This is an automated email from the ASF dual-hosted git repository.

rkk pushed a commit to branch tmp-stv
in repository https://gitbox.apache.org/repos/asf/sdap-nexus.git

commit e4bee9ade47f89eb6992e33e681f9237db47d7e1
Author: rileykk <[email protected]>
AuthorDate: Tue Jun 25 07:30:45 2024 -0700

    Tomo improvements
---
 analysis/webservice/algorithms/Tomogram.py   |  56 +++++++---
 analysis/webservice/algorithms/Tomogram3D.py | 149 +++++++++++++++++++++------
 2 files changed, 157 insertions(+), 48 deletions(-)

diff --git a/analysis/webservice/algorithms/Tomogram.py 
b/analysis/webservice/algorithms/Tomogram.py
index 2e01afd..9f3f779 100644
--- a/analysis/webservice/algorithms/Tomogram.py
+++ b/analysis/webservice/algorithms/Tomogram.py
@@ -15,12 +15,13 @@
 
 import logging
 from io import BytesIO
-from typing import Dict, Literal, Union, Tuple
-from collections import namedtuple
+from typing import Tuple
 
+import matplotlib as mpl
 import matplotlib.pyplot as plt
 import numpy as np
 import xarray as xr
+
 from webservice.NexusHandler import nexus_handler
 from webservice.algorithms.NexusCalcHandler import NexusCalcHandler
 from webservice.webmodel import NexusResults, NexusProcessingException, 
NoDataException
@@ -405,7 +406,13 @@ class LongitudeTomogramImpl(TomogramBaseClass):
             "type": "boolean",
             "description": "Display percentiles of cumulative returns over 
elevation. Requires both canopy_ds and "
                            "ground_ds be set."
-        }
+        },
+        "cmap": {
+            "name": "Color Map",
+            "type": "string",
+            "description": f"Color map to use. Will default to viridis if not 
provided or an unsupported map is "
+                           f"provided. Supported cmaps: {[k for k in 
sorted(mpl.colormaps.keys())]}"
+        },
     }
     singleton = True
 
@@ -454,12 +461,17 @@ class LongitudeTomogramImpl(TomogramBaseClass):
 
         peaks = compute_options.get_boolean_arg('peaks') and not percentiles
 
+        cmap = mpl.colormaps.get(
+            compute_options.get_argument('cmap', 'viridis'),
+            mpl.colormaps['viridis']
+        )
+
         return (ds, parameter, longitude, min_lat, max_lat, min_elevation, 
max_elevation, elevation_margin,
-                horizontal_margin, stride, peaks, ch_ds, g_ds, percentiles)
+                horizontal_margin, stride, peaks, ch_ds, g_ds, percentiles, 
cmap)
 
     def calc(self, compute_options, **args):
         (dataset, parameter, longitude, min_lat, max_lat, min_elevation, 
max_elevation, elevation_margin,
-         horizontal_margin, stride, calc_peaks, ch_ds, g_ds, percentiles) = 
self.parse_args(compute_options)
+         horizontal_margin, stride, calc_peaks, ch_ds, g_ds, percentiles, 
cmap) = self.parse_args(compute_options)
 
         slices = dict(
             lat=slice(min_lat, max_lat),
@@ -499,9 +511,9 @@ class LongitudeTomogramImpl(TomogramBaseClass):
         ds = ds.sel(longitude=longitude, method='nearest')
 
         if 'ch' in ds.data_vars:
-            ds = ds.where(ds.tomo.elevation <= ds.ch)
+            ds['tomo'] = ds['tomo'].where(ds.tomo.elevation <= ds.ch)
         if 'gh' in ds.data_vars:
-            ds = ds.where(ds.tomo.elevation >= ds.gh)
+            ds['tomo'] = ds['tomo'].where(ds.tomo.elevation >= ds.gh)
 
         lats = ds.latitude.to_numpy()
 
@@ -539,7 +551,8 @@ class LongitudeTomogramImpl(TomogramBaseClass):
             s={'longitude': longitude},
             coords=dict(x=lats, y=ds.elevation.to_numpy()),
             meta=dict(dataset=dataset),
-            style='db' if not percentiles else 'percentile'
+            style='db' if not percentiles else 'percentile',
+            cmap=cmap
         )
 
 
@@ -622,7 +635,13 @@ class LatitudeTomogramImpl(TomogramBaseClass):
             "type": "boolean",
             "description": "Display percentiles of cumulative returns over 
elevation. Requires both canopy_ds and "
                            "ground_ds be set."
-        }
+        },
+        "cmap": {
+            "name": "Color Map",
+            "type": "string",
+            "description": f"Color map to use. Will default to viridis if not 
provided or an unsupported map is "
+                           f"provided. Supported cmaps: {[k for k in 
sorted(mpl.colormaps.keys())]}"
+        },
     }
     singleton = True
 
@@ -671,12 +690,17 @@ class LatitudeTomogramImpl(TomogramBaseClass):
 
         peaks = compute_options.get_boolean_arg('peaks') and not percentiles
 
+        cmap = mpl.colormaps.get(
+            compute_options.get_argument('cmap', 'viridis'),
+            mpl.colormaps['viridis']
+        )
+
         return (ds, parameter, latitude, min_lon, max_lon, min_elevation, 
max_elevation, elevation_margin,
-                horizontal_margin, stride, peaks, ch_ds, g_ds, percentiles)
+                horizontal_margin, stride, peaks, ch_ds, g_ds, percentiles, 
cmap)
 
     def calc(self, compute_options, **args):
         (dataset, parameter, latitude, min_lon, max_lon, min_elevation, 
max_elevation, elevation_margin,
-         horizontal_margin, stride, calc_peaks, ch_ds, g_ds, percentiles) = 
self.parse_args(compute_options)
+         horizontal_margin, stride, calc_peaks, ch_ds, g_ds, percentiles, 
cmap) = self.parse_args(compute_options)
 
         slices = dict(
             lon=slice(min_lon, max_lon),
@@ -712,9 +736,9 @@ class LatitudeTomogramImpl(TomogramBaseClass):
         ds = ds.sel(latitude=latitude, method='nearest')
 
         if 'ch' in ds.data_vars:
-            ds = ds.where(ds.tomo.elevation <= ds.ch)
+            ds['tomo'] = ds['tomo'].where(ds.tomo.elevation <= ds.ch)
         if 'gh' in ds.data_vars:
-            ds = ds.where(ds.tomo.elevation >= ds.gh)
+            ds['tomo'] = ds['tomo'].where(ds.tomo.elevation >= ds.gh)
 
         lons = ds.longitude.to_numpy()
 
@@ -752,7 +776,8 @@ class LatitudeTomogramImpl(TomogramBaseClass):
             s={'latitude': latitude},
             coords=dict(x=lons, y=ds.elevation.to_numpy()),
             meta=dict(dataset=dataset),
-            style='db' if not percentiles else 'percentile'
+            style='db' if not percentiles else 'percentile',
+            cmap=cmap
         )
 
 
@@ -810,6 +835,7 @@ class ProfileTomoResults(NexusResults):
         self.__coords = coords
         self.__slice = s
         self.__style = style
+        self.__cmap = args.get('cmap', mpl.colormaps['viridis'])
 
     def toImage(self):
         rows, peaks = self.results()
@@ -830,7 +856,7 @@ class ProfileTomoResults(NexusResults):
             v = dict(vmin=-30, vmax=-10)
         else:
             v = dict(vmin=0, vmax=1)
-        plt.pcolormesh(self.__coords['x'], self.__coords['y'], rows, **v)
+        plt.pcolormesh(self.__coords['x'], self.__coords['y'], rows, 
cmap=self.__cmap, **v)
         plt.colorbar()
         plt.title(f'{self.meta()["dataset"]} tomogram slice\n{title_row}')
         plt.xlabel(xlabel)
diff --git a/analysis/webservice/algorithms/Tomogram3D.py 
b/analysis/webservice/algorithms/Tomogram3D.py
index 0324825..2436474 100644
--- a/analysis/webservice/algorithms/Tomogram3D.py
+++ b/analysis/webservice/algorithms/Tomogram3D.py
@@ -25,14 +25,29 @@ import matplotlib.pyplot as plt
 import numpy as np
 import pandas as pd
 import requests
-from mpl_toolkits.basemap import Basemap
 from PIL import Image
+from matplotlib.colors import LinearSegmentedColormap, XKCD_COLORS
+from mpl_toolkits.basemap import Basemap
+from xarray import Dataset, apply_ufunc
+
 from webservice.NexusHandler import nexus_handler
 from webservice.algorithms.NexusCalcHandler import NexusCalcHandler
 from webservice.webmodel import NexusResults, NexusProcessingException, 
NoDataException
-from xarray import Dataset, apply_ufunc
 
 logger = logging.getLogger(__name__)
+mpl.colormaps.register(
+    LinearSegmentedColormap.from_list(
+        'forest',
+        [
+            XKCD_COLORS['xkcd:dirt brown'],
+            # (0.0, 1.0, 0.0),
+            XKCD_COLORS['xkcd:dark forest green'],
+            XKCD_COLORS['xkcd:forest green'],
+            XKCD_COLORS['xkcd:forest'],
+            XKCD_COLORS['xkcd:light forest green'],
+            XKCD_COLORS['xkcd:leaf'],
+        ])
+)
 
 
 @nexus_handler
@@ -126,7 +141,21 @@ class Tomogram3D(NexusCalcHandler):
             "type": "boolean",
             "description": "Display percentiles of cumulative returns over 
elevation. Requires both canopy_ds and "
                            "ground_ds be set."
-        }
+        },
+        "renderRH": {
+            "name": "Render RH",
+            "type": "int",
+            "description": "With \'elevPercentiles\', this parameter is an 
integer between 0 and 100, which if set, "
+                           "will filter out all data whose cumulative return 
values is greater than the parameter "
+                           "value. It will also apply the color mapping to the 
voxel elevation rather than value, "
+                           "rendering a surface of that RH value."
+        },
+        "cmap": {
+            "name": "Color Map",
+            "type": "string",
+            "description": f"Color map to use. Will default to viridis if not 
provided or an unsupported map is "
+                           f"provided. Supported cmaps: {[k for k in 
sorted(mpl.colormaps.keys())]}"
+        },
     }
     singleton = True
 
@@ -214,6 +243,14 @@ class Tomogram3D(NexusCalcHandler):
                 reason='\'elevPercentiles\' argument requires \'canopy_ds\' 
and \'ground_ds\' be set'
             )
 
+        render_rh = compute_options.get_int_arg('renderRH', default=None)
+
+        if render_rh is not None and not 0 <= render_rh <= 100:
+            raise NexusProcessingException(
+                code=400,
+                reason="'renderRH' must be between 0 and 100"
+            )
+
         filter_arg = compute_options.get_argument('filter') if not percentiles 
else None
 
         if filter_arg is not None:
@@ -251,9 +288,14 @@ class Tomogram3D(NexusCalcHandler):
 
         hide_basemap = compute_options.get_boolean_arg('hideBasemap')
 
+        cmap = mpl.colormaps.get(
+            compute_options.get_argument('cmap', 'viridis'),
+            mpl.colormaps['viridis']
+        )
+
         return (ds, parameter_s, bounding_poly, min_elevation, max_elevation,
                 (orbit_elev, orbit_step, frame_duration, view_azim, 
view_elev), filter_arg, render_type, hide_basemap,
-                ch_ds, g_ds, percentiles, output)
+                ch_ds, g_ds, percentiles, output, render_rh, cmap)
 
     def do_subset(self, ds, parameter, bounds):
         tile_service = self._get_tile_service()
@@ -326,7 +368,7 @@ class Tomogram3D(NexusCalcHandler):
 
     def calc(self, computeOptions, **args):
         (ds, parameter, bounding_poly, min_elevation, max_elevation, 
render_params, filter_arg, render_type,
-         hide_basemap, ch_ds, g_ds, percentiles, output) = 
(self.parse_args(computeOptions))
+         hide_basemap, ch_ds, g_ds, percentiles, output, render_rh, cmap) = 
(self.parse_args(computeOptions))
 
         min_lat = bounding_poly.bounds[1]
         max_lat = bounding_poly.bounds[3]
@@ -366,12 +408,23 @@ class Tomogram3D(NexusCalcHandler):
                 ]
             )
         else:
-            df = self.__mask_by_elev_map(data, ch_ds, g_ds, bounds.copy(), 
parameter, percentiles, output)
+            df = self.__mask_by_elev_map(
+                data, ch_ds, g_ds, bounds.copy(), parameter, percentiles, 
output, render_rh is not None
+            )
 
         logger.info(f'Result repr:\n{df}')
 
         bounds = (bounds['lat'].start, bounds['lon'].start, 
bounds['lat'].stop, bounds['lon'].stop)
 
+        if percentiles:
+            if render_rh is None or output == 'NETCDF':
+                style = 'percentile'
+            else:
+                df = df[df['tomo_value'] > 1 - (render_rh / 100)]
+                style = 'percentile_elevation'
+        else:
+            style = 'db'
+
         return Tomogram3DResults(
             df,
             render_params,
@@ -379,37 +432,49 @@ class Tomogram3D(NexusCalcHandler):
             render_type,
             bounds=bounds,
             hide_basemap=hide_basemap,
-            style='db' if not percentiles else 'percentile'
+            style=style,
+            cmap=cmap
         )
 
-    def __mask_by_elev_map(self, data_points, ch_ds, g_ds, bounds, parameter, 
percentiles, output):
+    def __mask_by_elev_map(self, data_points, ch_ds, g_ds, bounds, parameter, 
percentiles, output, invert):
         from .Tomogram import TomogramBaseClass as Tomo
 
-        # Sample data_points to determine difference in elevation between 
slices
+        tries = 5
 
-        sample = data_points[random.randint(0, len(data_points) - 1)]
+        # Sample data_points to determine difference in elevation between 
slices
+        while tries >= 0:
+            # There's a bug where some samples yield non-uniform slices for 
what DEFINITELY was uniform data at ingest
+            # Trying to get around it by just resampling for now, but this 
will definitely need further investigation
+            tries -= 1
+            sample = data_points[random.randint(0, len(data_points) - 1)]
 
-        sample = [d['elevation'] for d in data_points if
-                  d['latitude'] == sample['latitude'] and d['longitude'] == 
sample['longitude']]
+            sample = [d['elevation'] for d in data_points if
+                      d['latitude'] == sample['latitude'] and d['longitude'] 
== sample['longitude']]
 
-        sample.sort()
-        sample = [sample[i+1] - sample[i] for i in range(len(sample) - 1)]
+            sample.sort()
+            sample = [sample[i + 1] - sample[i] for i in range(len(sample) - 
1)]
 
-        if min(sample) != max(sample):
-            raise NexusProcessingException(
-                code=500,
-                reason='There appears to be non-uniform elevation slices in 
the subset.'
-            )
+            if min(sample) == max(sample):
+                break
+            elif tries < 0:
+                raise NexusProcessingException(
+                    code=500,
+                    reason=f'There appears to be non-uniform elevation slices 
in the subset: '
+                           f'{sample}, {min(sample)}, {max(sample)}, 
{len(sample)}'
+                )
+            else:
+                logger.warning(f'Possible bad elevation sampling. Trying 
again. '
+                               f'{min(sample)}, {max(sample)}, {len(sample)}')
 
         stride = sample[0]
 
         # Bin and grid the subset data to an xarray Dataset
-        elevations = np.arange(bounds['elevation'].start, 
bounds['elevation'].stop + stride, stride/2)
+        elevations = np.arange(bounds['elevation'].start, 
bounds['elevation'].stop + stride, stride / 2)
         bounds['elevation'] = slice(None, None)  # Remove elevation from 
params as these are 2d datasets
 
         ds = Tomo.data_subset_to_ds_with_elevation(
             Tomo.bin_subset_elevations_to_range(
-                data_points, elevations, stride/2
+                data_points, elevations, stride / 2
             )
         )[3]
 
@@ -428,13 +493,16 @@ class Tomogram3D(NexusCalcHandler):
 
         # Trim the data
         if 'ch' in ds.data_vars:
-            ds = ds.where(ds.tomo.elevation <= ds.ch)
+            ds['tomo'] = ds['tomo'].where(ds.tomo.elevation <= ds.ch)
         if 'gh' in ds.data_vars:
-            ds = ds.where(ds.tomo.elevation >= ds.gh)
+            ds['tomo'] = ds['tomo'].where(ds.tomo.elevation >= ds.gh)
 
         if percentiles:
             tomo = ds['tomo'].cumsum(dim='elevation', 
skipna=True).where(ds['tomo'].notnull())
             ds['tomo'] = tomo / tomo.max(dim='elevation', skipna=True)
+
+            if invert:
+                ds['tomo'] = 1 - ds['tomo']
         else:
             ds['tomo'] = apply_ufunc(lambda a: 10 * np.log10(a), ds.tomo)
 
@@ -466,6 +534,7 @@ class Tomogram3DResults(NexusResults):
         self.render_type = render_type
         self.style = style
         self.hide_basemap = 'hide_basemap' in args and args['hide_basemap']
+        self.cmap = args.get('cmap', mpl.colormaps['viridis'])
 
     def results(self):
         r: pd.DataFrame = NexusResults.results(self)
@@ -487,7 +556,7 @@ class Tomogram3DResults(NexusResults):
     def __common(self):
         xyz = self.results()[['lon', 'lat', 'elevation']].values
 
-        fig = plt.figure(figsize=(10,7))
+        fig = plt.figure(figsize=(10, 7))
         return xyz, (fig, fig.add_subplot(111, projection='3d'))
 
     @staticmethod
@@ -519,7 +588,7 @@ class Tomogram3DResults(NexusResults):
         if not self.hide_basemap:
             min_lat, min_lon, max_lat, max_lon = self.bounds
 
-            m = Basemap(llcrnrlon=min_lon, llcrnrlat=min_lat, 
urcrnrlat=max_lat, urcrnrlon=max_lon,)
+            m = Basemap(llcrnrlon=min_lon, llcrnrlat=min_lat, 
urcrnrlat=max_lat, urcrnrlon=max_lon, )
 
             basemap_size = 512
 
@@ -566,21 +635,28 @@ class Tomogram3DResults(NexusResults):
 
             if self.style == 'db':
                 v = dict(vmin=-30, vmax=-10)
-            else:
+                cbl = 'Tomogram (dB)'
+            elif self.style == 'percentile':
                 v = dict(vmin=0, vmax=1)
+                cbl = 'Tomogram Cumulative value (%/100)'
+            else:
+                elevs = results[['elevation']].values
+                v = dict(vmin=np.nanmin(elevs), vmax=np.nanmax(elevs))
+                cbl = 'Voxel elevation'
 
             s = ax.scatter(
                 xyz[:, 0], xyz[:, 1], xyz[:, 2],
                 marker=',',
-                c=results[['tomo_value']].values,
+                c=results[['tomo_value']].values if self.style != 
'percentile_elevation'
+                else results[['elevation']].values,
                 zdir='z',
                 depthshade=True,
-                cmap=mpl.colormaps['viridis'],
+                cmap=self.cmap,
                 **v
             )
 
             cbar = fig.colorbar(s, ax=ax)
-            cbar.set_label('Tomogram (dB)')
+            cbar.set_label(cbl)
         else:
             logger.info('Collecting data by lat/lon')
             lats = np.unique(results['lat'].values)
@@ -698,21 +774,28 @@ class Tomogram3DResults(NexusResults):
 
             if self.style == 'db':
                 v = dict(vmin=-30, vmax=-10)
-            else:
+                cbl = 'Tomogram (dB)'
+            elif self.style == 'percentile':
                 v = dict(vmin=0, vmax=1)
+                cbl = 'Tomogram Cumulative value (%/100)'
+            else:
+                elevs = results[['elevation']].values
+                v = dict(vmin=np.nanmin(elevs), vmax=np.nanmax(elevs))
+                cbl = 'Voxel elevation'
 
             s = ax.scatter(
                 xyz[:, 0], xyz[:, 1], xyz[:, 2],
                 marker=',',
-                c=results[['tomo_value']].values,
+                c=results[['tomo_value']].values if self.style != 
'percentile_elevation'
+                else results[['elevation']].values,
                 zdir='z',
                 depthshade=True,
-                cmap=mpl.colormaps['viridis'],
+                cmap=self.cmap,
                 **v
             )
 
             cbar = fig.colorbar(s, ax=ax)
-            cbar.set_label('Tomogram (dB)')
+            cbar.set_label(cbl)
         else:
             logger.info('Collecting data by lat/lon')
             lats = np.unique(results['lat'].values)

Reply via email to