diff --git a/src/e3sm_quickview/app.py b/src/e3sm_quickview/app.py index 334fa30..255bb54 100644 --- a/src/e3sm_quickview/app.py +++ b/src/e3sm_quickview/app.py @@ -493,6 +493,11 @@ async def data_loading_open(self, simulation, connectivity): conn_file=connectivity, ) + # A different format means a different pipeline; existing views are + # still wired to the old one. + if self.source.path_changed: + self.view_manager.drop_views() + self.file_browser.loading_completed(self.source.valid) if self.source.valid: @@ -687,9 +692,30 @@ def _on_slicing_change(self, var, ind_var, **_): geom_filter.Update() data = geom_filter.GetOutput() self.state.fields_avgs = compute.extract_avgs( - data, self.selected_variable_names + data, self.selected_variable_names, self.source.association ) + @change("longitude_origin") + def _on_longitude_origin(self, longitude_origin, **_): + """Rotate the map so its left edge sits at the chosen longitude.""" + if not self.source.valid: + return + + origin = float(longitude_origin) + with perf.timed("longitude_origin.total"): + self.source.SetLongitudeOrigin(origin) + # The crop is expressed in the map's own longitudes, so rotating the + # map invalidates the previous selection -- a stale one would sit + # partly outside the window, or straddle its seam, which the box + # crop used for the continents and the graticule cannot represent. + with self.state as s: + s.crop_longitude = [origin, origin + 360] + s.crop_longitude_min = origin + s.crop_longitude_max = origin + 360 + self.view_manager.update_color_range() + self.view_manager.reset_camera() + self.view_manager.render() + @change( "variables_loaded", "crop_longitude", @@ -725,7 +751,7 @@ def _on_downstream_change( geom_filter.Update() data = geom_filter.GetOutput() self.state.fields_avgs = compute.extract_avgs( - data, self.selected_variable_names + data, self.selected_variable_names, self.source.association ) def toggle_toolbar(self, toolbar_name=None): diff --git a/src/e3sm_quickview/components/toolbars.py b/src/e3sm_quickview/components/toolbars.py index ff1da34..5a2fc83 100644 --- a/src/e3sm_quickview/components/toolbars.py +++ b/src/e3sm_quickview/components/toolbars.py @@ -324,8 +324,8 @@ def __init__(self): ) v3.VRangeSlider( v_model=("crop_longitude", [-180, 180]), - min=-180, - max=180, + min=("longitude_origin",), + max=("longitude_origin + 360",), step=1, density="compact", hide_details=True, @@ -407,11 +407,31 @@ def __init__(self): density="compact", hide_details=True, ) + + with v3.VCol(): + with v3.VRow(classes="mx-2 my-0"): + v3.VLabel( + "Map origin", + classes="text-subtitle-2", + ) + v3.VSpacer() + v3.VLabel( + "{{ longitude_origin }} to {{ longitude_origin + 360 }}", + classes="text-body-2", + ) + v3.VSlider( + v_model=("longitude_origin", -180), + min=-180, + max=180, + step=1, + density="compact", + hide_details=True, + ) with v3.VRow(classes="ma-0 pl-6 pr-2 align-center ga-4", v_else=True): v3.VNumberInput( label="Longitude (min)", v_model=("crop_longitude_min", -180), - min=[-180], + min=("longitude_origin",), max=("crop_longitude_max", 180), step=[1], hide_details=True, @@ -424,7 +444,7 @@ def __init__(self): label="Longitude (max)", v_model=("crop_longitude_max", 180), min=("crop_longitude_min", -180), - max=[180], + max=("longitude_origin + 360",), step=[1], hide_details=True, density="comfortable", diff --git a/src/e3sm_quickview/pipeline.py b/src/e3sm_quickview/pipeline.py index 30f72e7..7ee6ed9 100644 --- a/src/e3sm_quickview/pipeline.py +++ b/src/e3sm_quickview/pipeline.py @@ -7,6 +7,25 @@ from vtkmodules.vtkRenderingCore import vtkActor, vtkPolyDataMapper +def longitude_window(lon_range, origin): + """Express a geographic longitude range inside the window [origin, origin+360). + + Everything downstream -- the data, the continents, the graticule -- has to + agree on which 360-degree window it is working in, or the overlays drift off + the map when the origin is rotated. + """ + low, high = lon_range + width = (high - low) % 360.0 + if width == 0.0: + # A full turn has no meaningful left edge of its own -- it is the window. + return [origin, origin + 360.0] + low_in_window = origin + (low - origin) % 360.0 + # A selection may not run off the end of its own window: the box crop behind + # the continents and the graticule cannot wrap, so it would silently lose + # whatever fell past the seam. + return [low_in_window, min(low_in_window + width, origin + 360.0)] + + def load_plugins(): try: plugin_dir = Path(__file__).with_name("plugins") @@ -48,6 +67,7 @@ def __init__(self, projection="Mollweide"): LongitudeRange=[-180.0, 180.0], LatitudeRange=[-90.0, 90.0], ) + self._longitude_origin = -180.0 self.proj = simple.EAMProject( Input=self._crop, Projection=projection, @@ -70,6 +90,16 @@ def crop(self, longitude_min_max, latitude_min_max): self._crop.LongitudeRange = longitude_min_max self._crop.LatitudeRange = latitude_min_max + @property + def longitude_origin(self): + return self._longitude_origin + + @longitude_origin.setter + def longitude_origin(self, origin): + self._longitude_origin = origin + self._crop.LongitudeOrigin = origin + self.proj.LongitudeOrigin = origin + @property def projection(self): return self._projection @@ -87,6 +117,7 @@ class GridLines: def __init__(self, projection="Mollweide"): self._projection = projection self.grid_lines = simple.EAMGridLines() + self._longitude_origin = -180.0 self.proj = simple.EAMProject( Input=self.grid_lines, Projection=projection, @@ -110,6 +141,15 @@ def crop(self, longitude_min_max, latitude_min_max): self.grid_lines.LongitudeRange = longitude_min_max self.grid_lines.LatitudeRange = latitude_min_max + @property + def longitude_origin(self): + return self._longitude_origin + + @longitude_origin.setter + def longitude_origin(self, origin): + self._longitude_origin = origin + self.proj.LongitudeOrigin = origin + @property def projection(self): return self._projection @@ -124,7 +164,45 @@ def update(self): self.mapper.Update() -class DataReader: +def connectivity_kind(conn_file): + """Sniff a connectivity file and name the format it describes. + + Returns "dycore" for a HOMME np4/GLL grid file (nodal values on shared + spectral-element nodes), "pg2" for a SCRIP physics grid (cell values on + unshared corners), or None if neither signature is present. + """ + try: + import netCDF4 + + with netCDF4.Dataset(conn_file) as ds: + names = set(ds.variables) + if "element_corners" in names: + return "dycore" + if any("corner_lat" in name for name in names): + return "pg2" + except Exception as e: + print(f"Could not inspect connectivity file {conn_file}: {e}") + + return None + + +class DataPath: + """One format's pipeline, from its reader to the surface the views render. + + Each format brings its own chain of filters and its own attribute + association, so the rest of the application can stay format-agnostic: it + talks to whichever path is active through this interface and reads + ``association`` when it needs to know where the variables live. + + Subclasses provide ``_build_pipeline``, which must set ``reader``, ``proj`` + and ``geometry``. + """ + + kind = None + label = None + #: Where this format's variables land -- "cell" or "point". + association = "cell" + def __init__(self, projection="Mollweide"): self._file_connection = None self._file_mesh = None @@ -136,23 +214,7 @@ def __init__(self, projection="Mollweide"): self._dimensions = None self._slicing = defaultdict(int) - # Pipeline - self.reader = simple.EAMSliceDataReader() - self.center_meridian = simple.EAMCenterMeridian( - Input=self.reader, - Meridian=0, - ) - self._crop = simple.EAMExtract( - Input=self.center_meridian, - LongitudeRange=[-180, 180], - LatitudeRange=[-90, 90], - ) - self.proj = simple.EAMProject( # noqa: F821 - Input=self._crop, - Projection=projection, - Translate=0, - ) - self.geometry = simple.ExtractSurface(Input=self.proj) + self._build_pipeline(projection) self.vtk_geometry = self.geometry.GetClientSideObject() # Add observer to @@ -160,6 +222,9 @@ def __init__(self, projection="Mollweide"): vtk_obj.AddObserver("ErrorEvent", self._observer) vtk_obj.GetExecutive().AddObserver("ErrorEvent", self._observer) + def _build_pipeline(self, projection): + raise NotImplementedError + @property def valid(self): return self._valid and not self._observer.error_occurred @@ -226,6 +291,9 @@ def projection(self, value): self._projection = value self.proj.Projection = value + def set_variables(self, names): + self.reader.Variables = list(set([*names, "lat", "lon"])) + def update_slicing(self, dimension, value): current_value = self._slicing.get(dimension, 0) @@ -242,20 +310,143 @@ def update(self, time=0.0): self.geometry.UpdatePipeline(time) + def crop(self, longitude_min_max, latitude_min_max): + """Restrict the rendered region. Formats that cannot crop ignore this.""" + + def set_longitude_origin(self, origin): + """Place the left edge of the map; the right edge is a turn further east.""" + self.proj.LongitudeOrigin = origin + + +class Pg2Path(DataPath): + """ne*pg2 physics grid: cell values on a SCRIP corner mesh. + + Cells are cut at the central meridian, so coverage stays exact, and the + per-tick cost of re-cutting is avoided by remapping cell values through + the clip's PedigreeIds. + """ + + kind = "pg2" + label = "EAM physics grid (pg2)" + association = "cell" + + def _build_pipeline(self, projection): + self.reader = simple.EAMSliceDataReader() + self.center_meridian = simple.EAMCenterMeridian( + Input=self.reader, + Meridian=0, + ) + self._crop = simple.EAMExtract( + Input=self.center_meridian, + LongitudeRange=[-180, 180], + LatitudeRange=[-90, 90], + ) + self.proj = simple.EAMProject( # noqa: F821 + Input=self._crop, + Projection=projection, + Translate=0, + ) + self.geometry = simple.ExtractSurface(Input=self.proj) + def crop(self, longitude_min_max, latitude_min_max): self._crop.LongitudeRange = longitude_min_max self._crop.LatitudeRange = latitude_min_max + def set_longitude_origin(self, origin): + # The clip already rearranges the halves; it just needs the new seam. + super().set_longitude_origin(origin) + self.center_meridian.LongitudeOrigin = origin + + +class DycorePath(DataPath): + """ne*np4 dynamical core grid: nodal values on shared GLL nodes. + + The reader lays the sphere flat itself, duplicating nodes at the date line + and the poles instead of clipping, so cells stay whole and no value is + interpolated. That leaves nothing for EAMCenterMeridian to do -- the mesh + already tiles [-180, 180] exactly -- so the path goes straight to the + projection. + """ + + kind = "dycore" + label = "EAM dynamical core (np4)" + association = "point" + + def _build_pipeline(self, projection): + # The reader stays in its natural [-180, 180) window, where the GLL + # nodes land exactly on the seam and its node-duplication split tiles + # the map with no overhang and no gap. Rotating to any other origin is + # left to the clip: shifting whole cells would throw a polar cell -- + # up to 90 degrees wide -- past the edge, and a projection wraps that + # onto the far side of the map. + self.reader = simple.EAMDycoreReader() # noqa: F821 + self.center_meridian = simple.EAMCenterMeridian( # noqa: F821 + Input=self.reader, + Meridian=0, + InputLongitudeOrigin=-180, + ) + # EAMExtract is not a clip -- it hides whole cells and removes them -- + # so it never interpolates and the nodal values survive untouched. + self._crop = simple.EAMExtract( # noqa: F821 + Input=self.center_meridian, + LongitudeRange=[-180, 180], + LatitudeRange=[-90, 90], + ) + self.proj = simple.EAMProject( # noqa: F821 + Input=self._crop, + Projection=projection, + Translate=0, + ) + self.geometry = simple.ExtractSurface(Input=self.proj) + + def crop(self, longitude_min_max, latitude_min_max): + self._crop.LongitudeRange = longitude_min_max + self._crop.LatitudeRange = latitude_min_max + + def set_longitude_origin(self, origin): + # The reader keeps its exact window; the clip does the rotation. + super().set_longitude_origin(origin) + self.center_meridian.LongitudeOrigin = origin + + +#: Every format the application can open, in detection order. +DATA_PATHS = (Pg2Path, DycorePath) + +#: Kept so existing callers that expect the pg2 pipeline keep working. +DataReader = Pg2Path + class EAMVisSource: def __init__(self): self.projection = "Mollweide" load_plugins() - self.data_reader = DataReader(self.projection) + self._paths = {} + # pg2 is the default so that opening a physics-grid file behaves + # exactly as it did before formats became selectable. + self.data_reader = self._path_for(Pg2Path) self.continent = Continent(self.projection) self.grid_lines = GridLines(self.projection) self.views = {} + #: True when the last Update() switched to a different format. + self.path_changed = False + #: Left edge of the map; the right edge sits 360 degrees east of it. + self.longitude_origin = -180.0 + self._crop_lon = [-180.0, 180.0] + self._crop_lat = [-90.0, 90.0] + + def _path_for(self, path_cls): + """Return this format's pipeline, building it the first time it is used.""" + path = self._paths.get(path_cls.kind) + if path is None: + path = path_cls(self.projection) + self._paths[path_cls.kind] = path + return path + + @property + def association(self): + """Where the active format's variables live -- "cell" or "point".""" + return self.data_reader.association @property def valid(self): @@ -276,9 +467,33 @@ def ApplyClipping(self, cliplong, cliplat): if not self.valid: return - self.data_reader.crop(cliplong, cliplat) - self.continent.crop(cliplong, cliplat) - self.grid_lines.crop(cliplong, cliplat) + self._crop_lon = list(cliplong) + self._crop_lat = list(cliplat) + # The crop arrives in geographic longitude; everything downstream works + # in the current window, so move it there once, here. + window = longitude_window(cliplong, self.longitude_origin) + self.data_reader.crop(window, cliplat) + self.continent.crop(window, cliplat) + self.grid_lines.crop(window, cliplat) + + def SetLongitudeOrigin(self, origin): + """Rotate the map so its left edge is at `origin` degrees longitude.""" + if self.longitude_origin == origin: + return + + self.longitude_origin = origin + # Every format keeps its own pipeline, so they all have to be told -- + # not just the active one, or switching format would lose the window. + for path in self._paths.values(): + path.set_longitude_origin(origin) + self.continent.longitude_origin = origin + self.grid_lines.longitude_origin = origin + + # Re-express the crop in the new window and refresh the overlays. + self.ApplyClipping(self._crop_lon, self._crop_lat) + self.UpdatePipeline() + self.continent.update() + self.grid_lines.update() def UpdateProjection(self, proj): if not self.valid: @@ -300,6 +515,20 @@ def UpdateSlicing(self, dimension, slice): self.data_reader.update_slicing(dimension, slice) def Update(self, data_file, conn_file): # force_reload + kind = connectivity_kind(conn_file) + path_cls = next( + (cls for cls in DATA_PATHS if cls.kind == kind), + Pg2Path, + ) + path = self._path_for(path_cls) + + # Views bind to the tail of a specific pipeline, so a format switch has + # to be visible to the caller -- it invalidates every existing view. + self.path_changed = path is not self.data_reader + if self.path_changed: + self.data_reader = path + path.projection = self.projection + if self.data_reader.load(data_file, conn_file): self.views["atmosphere_data"] = self.data_reader.vtk_geometry self.views["continents"] = self.continent.proj @@ -312,7 +541,7 @@ def LoadVariables(self, vars): if not self.valid: return - self.data_reader.reader.Variables = list(set([*vars, "lat", "lon"])) + self.data_reader.set_variables(vars) def Clip(self, plane=None): self.grid_lines.mapper.RemoveAllClippingPlanes() diff --git a/src/e3sm_quickview/plugins/eam_projection.py b/src/e3sm_quickview/plugins/eam_projection.py index e282dd8..0ce792e 100644 --- a/src/e3sm_quickview/plugins/eam_projection.py +++ b/src/e3sm_quickview/plugins/eam_projection.py @@ -133,50 +133,61 @@ def ProcessPoint(point, radius): return [x, y, z] -def add_cell_arrays(inData, outData, cached_output): +def _translated(dataset, shift): + """dataset moved `shift` degrees in longitude (returned as-is when shift is 0).""" + if shift == 0.0: + return dataset + transform = vtkTransform() + transform.Translate(shift, 0, 0) + transform_filter = vtkTransformFilter() + transform_filter.SetInputData(dataset) + transform_filter.SetTransform(transform) + transform_filter.Update() + return transform_filter.GetOutput() + + +def _longitude_window(origin, input_origin=0.0): + """Cut meridian and the two shifts that move [input_origin, +360) to [origin, +360). + + Everything below the cut is translated by one turn relative to everything + above it, and a whole-turn offset then places the seam exactly at `origin`. + With an input running [0, 360) and origin = -180 this reduces to the + historical behaviour: cut at 180, right half shifted by -360. """ - Adds arrays not modified in inData to outData. - New arrays (or arrays modified) values are set using the PedigreeIds - because the number of values in the new array (just read from the file) - is different than the number of values in the arrays already processed - through the pipeline. + cut = input_origin + (origin - input_origin) % 360.0 + return cut, origin + 360.0 - cut, origin - cut + + +def _remap_arrays(in_attrs, cached_attrs, out_attrs, pedigree_vtk, label): + """Rebuild out_attrs from in_attrs, permuted through a pedigree map. + + The number of values in a freshly read array differs from the number that + came out of the pipeline, so values are gathered through the pedigree + permutation recorded when the geometry was last built. A single fancy-index gather does this. An earlier version walked the permutation as a list of monotonic run slices, assuming the runs were thousands of entries long. Measured against the permutations this pipeline - actually produces — mean run 55-110 — that loop is 8-16x *slower* than one - numpy gather, because the per-run Python overhead dominates. + actually produces — mean run 13 for point ids, 55-110 for cell ids — that + loop is 8-66x *slower* than one numpy gather, because the per-run Python + overhead dominates. """ - pedigreeIds = cached_output.cell_data["PedigreeIds"] - if pedigreeIds is None: - print_error("Error: no PedigreeIds array") - return - - pedigree_vtk = cached_output.GetCellData().GetArray("PedigreeIds") pid_np = numpy_support.vtk_to_numpy(pedigree_vtk) - - cached_cell_data = cached_output.GetCellData() - in_cell_data = inData.GetCellData() - outData.ShallowCopy(cached_output) - out_cell_data = outData.GetCellData() - - out_cell_data.Initialize() - for i in range(in_cell_data.GetNumberOfArrays()): - in_array = in_cell_data.GetArray(i) - cached_array = cached_cell_data.GetArray(in_array.GetName()) + n_tuples = pedigree_vtk.GetNumberOfTuples() + out_attrs.Initialize() + for i in range(in_attrs.GetNumberOfArrays()): + in_array = in_attrs.GetArray(i) + cached_array = cached_attrs.GetArray(in_array.GetName()) if cached_array and cached_array.GetMTime() >= in_array.GetMTime(): # This scalar has been seen before — reuse cached copy. - out_cell_data.AddArray(cached_array) + out_attrs.AddArray(cached_array) else: - with _perf.timed(f"add_cell_arrays.pedigree_copy.{in_array.GetName()}"): - array0 = cached_cell_data.GetArray(0) - n_comp = array0.GetNumberOfComponents() - n_tuples = array0.GetNumberOfTuples() + with _perf.timed(f"{label}.pedigree_copy.{in_array.GetName()}"): out_array = in_array.NewInstance() - out_array.SetNumberOfComponents(n_comp) + out_array.SetNumberOfComponents(in_array.GetNumberOfComponents()) out_array.SetNumberOfTuples(n_tuples) out_array.SetName(in_array.GetName()) - out_cell_data.AddArray(out_array) + out_attrs.AddArray(out_array) in_np = numpy_support.vtk_to_numpy(in_array) out_np = numpy_support.vtk_to_numpy(out_array) @@ -184,6 +195,58 @@ def add_cell_arrays(inData, outData, cached_output): out_array.Modified() +def add_cell_arrays(inData, outData, cached_output): + """Refresh cell arrays only — for filters that interpolate point data. + + A clip creates new points by interpolation, so an output point has no + single source point to gather from and the pedigree trick cannot work + for point data. Cells are only ever kept or dropped, so they can. + """ + pedigree_vtk = cached_output.GetCellData().GetArray("PedigreeIds") + if pedigree_vtk is None: + print_error("Error: no PedigreeIds array") + return + + outData.ShallowCopy(cached_output) + _remap_arrays( + inData.GetCellData(), + cached_output.GetCellData(), + outData.GetCellData(), + pedigree_vtk, + "add_cell_arrays", + ) + + +def add_cell_and_point_arrays(inData, outData, cached_output): + """Refresh cell *and* point arrays through their respective pedigree maps. + + Usable only where the filter subsets whole cells and never interpolates — + then every output point is a copy of an input point, so its pedigree id is + an exact gather index. EAMExtract qualifies; a clip does not. + """ + outData.ShallowCopy(cached_output) + + cell_pedigree = cached_output.GetCellData().GetArray("PedigreeIds") + if cell_pedigree is not None and inData.GetCellData().GetNumberOfArrays(): + _remap_arrays( + inData.GetCellData(), + cached_output.GetCellData(), + outData.GetCellData(), + cell_pedigree, + "add_cell_arrays", + ) + + point_pedigree = cached_output.GetPointData().GetArray("PointPedigreeIds") + if point_pedigree is not None and inData.GetPointData().GetNumberOfArrays(): + _remap_arrays( + inData.GetPointData(), + cached_output.GetPointData(), + outData.GetPointData(), + point_pedigree, + "add_point_arrays", + ) + + @smproxy.filter() @smproperty.input(name="Input") @smdomain.datatype( @@ -230,16 +293,10 @@ def RequestData(self, request, inInfo, outInfo): else: outData.DeepCopy(inData) - inPoints = numpy_support.vtk_to_numpy(inData.GetPoints().GetData()) + inPoints = inData.points pRadius = (self.radius + 1) if self.isData else self.radius outPoints = np.array(list(map(lambda x: ProcessPoint(x, pRadius), inPoints))) - vtk_coords = vtkPoints() - vtk_coords.SetData( - numpy_support.numpy_to_vtk( - outPoints, deep=True, array_type=vtkConstants.VTK_FLOAT - ) - ) - outData.SetPoints(vtk_coords) + outData.points = outPoints return 1 @@ -366,6 +423,12 @@ def RequestData(self, request, inInfo, outInfo): + + Left edge of the map; the projection is centred half a turn east of it. + """ ) class EAMProject(VTKPythonAlgorithmBase): @@ -382,6 +445,8 @@ def __init__(self): self._cached_input_points = None self._cached_key = None + self.longitude_origin = -180.0 + def _invalidate_cache(self): self.cached_points = None self._cached_input_points = None @@ -399,6 +464,14 @@ def SetProjection(self, project): self._invalidate_cache() self.Modified() + def SetLongitudeOrigin(self, origin): + """Left edge of the map. The projection is centred half a turn east of + it, so a rotated window still maps onto the middle of the figure.""" + if self.longitude_origin != origin: + self.longitude_origin = origin + self._invalidate_cache() + self.Modified() + def RequestData(self, request, inInfo, outInfo): with _perf.timed("project.RequestData"): inData = self.GetInputData(inInfo, 0, 0) @@ -423,6 +496,9 @@ def RequestData(self, request, inInfo, outInfo): out_points_vtk = vtkPoints() out_points_vtk.DeepCopy(outData.GetPoints()) outData.SetPoints(out_points_vtk) + # Go through numpy_support rather than the pythonic + # `.points`: VTK 9.7 returns a vtkPoints subclass there, + # where earlier versions handed back a numpy array. out_points_np = numpy_support.vtk_to_numpy( outData.GetPoints().GetData() ) @@ -455,6 +531,18 @@ def RequestData(self, request, inInfo, outInfo): # Should not reach here, but return without transformation return 1 + # Re-centre on the middle of the window here rather + # than through PROJ's lon_0. PROJ normalises its + # input into [-180, 180) *before* subtracting lon_0, + # which sends the window's right edge to the left + # rim -- drawing coastlines and cells straight + # across the map. The data is already confined to + # the window, so the offset lands in range on its + # own and PROJ never has to wrap anything. + x = np.clip( + x - (self.longitude_origin + 180.0), -180.0, 180.0 + ) + xformer = Transformer.from_proj( latlon, proj, always_xy=True ) @@ -494,6 +582,12 @@ def RequestData(self, request, inInfo, outInfo): number_of_elements="2" default_values="-180 180"> + + Left edge of the map; the right edge is 360 degrees east of it. + 180 or min > max: + # Ranges arrive in whichever 360-degree window the map is using, so + # only the ordering and the width are meaningful here. + if min > max or (max - min) > 360.0: print_error( f"SetLongitudeRange called with invalid parameters: {min=}, {max=}" ) @@ -625,7 +723,8 @@ def RequestData(self, request, inInfo, outInfo): with _perf.timed("extract.RequestData"): inData = self.GetInputData(inInfo, 0, 0) outData = self.GetOutputData(outInfo, 0) - if self.lon_range == [-180.0, 180.0] and self.lat_range == [-90.0, 90.0]: + spans_full_turn = (self.lon_range[1] - self.lon_range[0]) >= 359.999 + if spans_full_turn and self.lat_range == [-90.0, 90.0]: outData.ShallowCopy(inData) # Only invalidate the shared points when transitioning *out* of a # cropped state — the original code did it unconditionally, which @@ -661,14 +760,17 @@ def RequestData(self, request, inInfo, outInfo): self.GetMTime(), inData.GetPoints().GetMTime(), cell_centers.GetMTime() ): with _perf.timed("extract.cache_hit"): - outData.ShallowCopy(self._cached_output) - add_cell_arrays(inData, outData, self._cached_output) + add_cell_and_point_arrays(inData, outData, self._cached_output) else: with _perf.timed("extract.rebuild_trim"): # add PedigreeIds generate_ids = vtkGenerateIds() generate_ids.SetInputData(inData) - generate_ids.PointIdsOff() + # Point ids as well: RemoveGhostCells only ever drops whole + # cells, so a surviving point keeps an exact source index + # and nodal formats can be refreshed from the cache too. + generate_ids.PointIdsOn() + generate_ids.SetPointIdsArrayName("PointPedigreeIds") generate_ids.SetCellIdsArrayName("PedigreeIds") generate_ids.Update() outData.ShallowCopy(generate_ids.GetOutput()) @@ -690,9 +792,13 @@ def RequestData(self, request, inInfo, outInfo): # add HIDDENCELL based on ranges with _perf.timed("extract.ghost_mask"): + # Compare longitudes as offsets from lon_min taken + # modulo a turn, so the test is independent of which + # 360-degree window the data happens to live in and + # still works for a range that spans the seam. + lon_offset = (cc[:, 0] - lon_min) % 360.0 outside_mask = ( - (cc[:, 0] < lon_min) - | (cc[:, 0] > lon_max) + (lon_offset > ((lon_max - lon_min) % 360.0 or 360.0)) | (cc[:, 1] < lat_min) | (cc[:, 1] > lat_max) ) @@ -734,6 +840,19 @@ def RequestData(self, request, inInfo, outInfo): - 20: Often used to center Europe and Africa. + + + Left edge of the map; the right edge is 360 degrees east of it. + + + Left edge of the window the input already uses. + """ ) @smdomain.datatype( @@ -752,6 +871,7 @@ def __init__(self): ) # common values: self._center_meridian = 0 + self._input_origin = 0.0 self._cached_output = None def SetMeridian(self, meridian_): @@ -768,6 +888,34 @@ def SetMeridian(self, meridian_): self._center_meridian = meridian_ self.Modified() + def SetLongitudeOrigin(self, origin): + """Left edge of the map; the right edge is 360 degrees further east.""" + if origin < -180 or origin > 180: + print_error( + f"SetLongitudeOrigin called with parameter outside [-180, 180]: {origin}" + ) + return + meridian = origin + 180.0 + if self._center_meridian != meridian: + self._center_meridian = meridian + self._cached_output = None + self.Modified() + + def GetLongitudeOrigin(self): + return self._center_meridian - 180.0 + + def SetInputLongitudeOrigin(self, origin): + """Left edge of the window the *input* already uses. + + The pg2 reader emits [0, 360); the dycore reader emits [-180, 180). + Without this the cut lands outside the data and the rotation silently + does nothing. + """ + if self._input_origin != origin: + self._input_origin = origin + self._cached_output = None + self.Modified() + def GetMeridian(self): """ Returns the central meridian @@ -779,13 +927,29 @@ def RequestData(self, request, inInfo, outInfo): inData = self.GetInputData(inInfo, 0, 0) outData = self.GetOutputData(outInfo, 0) - if ( + + # Nothing to do when the input already sits in the requested + # window -- the dycore reader's default case. Clipping here would + # be a no-op that still rebuilds the points every pass, which also + # costs EAMProject its cache downstream. + origin = self._center_meridian - 180.0 + if (origin - self._input_origin) % 360.0 == 0.0: + with _perf.timed("center_meridian.passthrough"): + outData.ShallowCopy(inData) + return 1 + # A clip makes new points by interpolation, so an output point has + # no single source to gather from and the pedigree cache cannot + # refresh point data. Nodal formats therefore re-clip every pass; + # it costs a few milliseconds and is always correct. + has_point_arrays = inData.GetPointData().GetNumberOfArrays() > 0 + geometry_cached = bool( self._cached_output and self._cached_output.GetPoints().GetMTime() >= inData.GetPoints().GetMTime() and self._cached_output.GetCells().GetMTime() >= inData.GetCells().GetMTime() - ): + ) + if geometry_cached and not has_point_arrays: with _perf.timed("center_meridian.cache_hit"): add_cell_arrays(inData, outData, self._cached_output) else: @@ -795,9 +959,11 @@ def RequestData(self, request, inInfo, outInfo): generate_ids.PointIdsOff() generate_ids.SetCellIdsArrayName("PedigreeIds") - cut_meridian = self._center_meridian + 180 + cut, shift_low, shift_high = _longitude_window( + self._center_meridian - 180.0, self._input_origin + ) plane = vtkPlane() - plane.SetOrigin([cut_meridian, 0.0, 0.0]) + plane.SetOrigin([cut, 0.0, 0.0]) plane.SetNormal([-1, 0, 0]) # vtkClipPolyData hangs clipL = vtkTableBasedClipDataSet() @@ -813,21 +979,40 @@ def RequestData(self, request, inInfo, outInfo): with _perf.timed("center_meridian.clip_right"): clipR.Update() - transFunc = vtkTransform() - transFunc.Translate(-360, 0, 0) - transform = vtkTransformFilter() - transform.SetInputData(clipR.GetOutput()) - transform.SetTransform(transFunc) with _perf.timed("center_meridian.transform"): - transform.Update() + halves = [ + _translated(clipL.GetOutput(), shift_low), + _translated(clipR.GetOutput(), shift_high), + ] append = vtkAppendFilter() - append.AddInputData(clipL.GetOutput()) - append.AddInputData(transform.GetOutput()) + for half in halves: + append.AddInputData(half) with _perf.timed("center_meridian.append"): append.Update() outData.ShallowCopy(append.GetOutput()) - # previous _cached_output is available for garbage collection - self._cached_output = outData.NewInstance() - self._cached_output.ShallowCopy(outData) + + # The clip is deterministic, so when only the values + # changed the geometry it just produced is identical to the + # cached one. Hand the *same* points and cells objects + # downstream: EAMProject keys its cache on the identity of + # the incoming points, and EAMExtract on their modified + # time, so fresh copies would make both rebuild for nothing. + if ( + geometry_cached + and self._cached_output.GetNumberOfPoints() + == outData.GetNumberOfPoints() + and self._cached_output.GetNumberOfCells() + == outData.GetNumberOfCells() + ): + with _perf.timed("center_meridian.reuse_geometry"): + outData.SetPoints(self._cached_output.GetPoints()) + outData.SetCells( + _cell_types_array(self._cached_output), + self._cached_output.GetCells(), + ) + else: + # previous _cached_output is available for garbage collection + self._cached_output = outData.NewInstance() + self._cached_output.ShallowCopy(outData) return 1 diff --git a/src/e3sm_quickview/plugins/eam_reader.py b/src/e3sm_quickview/plugins/eam_reader.py index 5646fdf..7d5d5cf 100644 --- a/src/e3sm_quickview/plugins/eam_reader.py +++ b/src/e3sm_quickview/plugins/eam_reader.py @@ -197,62 +197,18 @@ def _markmodified(*args, **kwars): return _markmodified -@smproxy.reader( - name="EAMSliceSource", - label="EAM Slice Data Reader", - extensions="nc", - file_description="NETCDF files for EAM", -) -@smproperty.xml("""""") -@smproperty.xml( - """ - - - Specify the NetCDF data file name. - - """ -) -@smproperty.xml( - """ - - - Specify the NetCDF connecticity file name. - - """ -) -@smproperty.xml( - """ - - JSON representing dimension slices (e.g. {"lev": 0, "ilev": 1}) - - """ -) -@smproperty.xml( +class _EAMReaderBase(VTKPythonAlgorithmBase): + """Shared plumbing for the EAM readers: files, variable metadata, + dimension slicing and timesteps. + + This is deliberately left undecorated. ParaView's decorators replace the + class they wrap with a function, so a decorated reader cannot be used as + a base class -- the shared code has to live here instead. + + The geometry and data-placement methods below are the pg2 physics-grid + behaviour; EAMDycoreSource overrides them. """ - - - - If True, the points of the dataset will be float, otherwise they will be float or double depending - on the type of corner_lat and corner_lon variables in the connectivity file. - - - """ -) -class EAMSliceSource(VTKPythonAlgorithmBase): + def __init__(self): VTKPythonAlgorithmBase.__init__( self, nInputPorts=0, nOutputPorts=1, outputType="vtkUnstructuredGrid" @@ -824,3 +780,527 @@ def _RequestDataImpl(self, request, inInfo, outInfo): output.ShallowCopy(self._output) return 1 + + +@smproxy.reader( + name="EAMSliceSource", + label="EAM Slice Data Reader", + extensions="nc", + file_description="NETCDF files for EAM", +) +@smproperty.xml("""""") +@smproperty.xml( + """ + + + Specify the NetCDF data file name. + + """ +) +@smproperty.xml( + """ + + + Specify the NetCDF connecticity file name. + + """ +) +@smproperty.xml( + """ + + JSON representing dimension slices (e.g. {"lev": 0, "ilev": 1}) + + """ +) +@smproperty.xml( + """ + + + + If True, the points of the dataset will be float, otherwise they will be float or double depending + on the type of corner_lat and corner_lon variables in the connectivity file. + + + """ +) +class EAMSliceSource(_EAMReaderBase): + """ne*pg2 physics grid: cell values on an unshared SCRIP corner mesh.""" + + # ParaView builds a proxy's XML from the methods in the class's own + # __dict__, so these have to be declared on each decorated reader rather + # than inherited from the shared base. + @smproperty.doublevector( + name="TimestepValues", information_only="1", si_class="vtkSITimeStepsProperty" + ) + def GetTimestepValues(self): + return self._timeSteps + + @smproperty.dataarrayselection(name="Variables") + def GetSurfaceVariables(self): + return self._variable_selection + + +# --------------------------------------------------------------------------- +# Dycore reader: native spectral-element (ne*np4 / GLL) grids +# --------------------------------------------------------------------------- + + +@smproxy.reader( + name="EAMDycoreSource", + label="EAM Dycore Reader", + extensions="nc", + file_description="NETCDF files for the EAM dynamical core (np4/GLL)", +) +@smproperty.xml("""""") +@smproperty.xml( + """ + + + Specify the NetCDF data file name. + + """ +) +@smproperty.xml( + """ + + + Specify the HOMME np4 grid file (lat/lon + element_corners). + + """ +) +@smproperty.xml( + """ + + JSON representing dimension slices (e.g. {"lev": 0, "time": 1}) + + """ +) +@smproperty.xml( + """ + + + + If True, the points of the dataset will be float, otherwise double. + + + """ +) +@smproperty.xml( + """ + + + Left edge of the map in degrees. -180 places the seam on the date + line, where the ne*np4 element boundaries fall, so the split cells + tile the map with no overhang and no gap. + + + """ +) +class EAMDycoreSource(_EAMReaderBase): + """Read EAM/CAM-SE native spectral-element (ne*np4) output. + + The dynamical core runs on Gauss-Lobatto-Legendre quadrature nodes: 4x4 + nodes per spectral element, with edge and corner nodes *shared* between + neighbouring elements. Values are nodal point samples, not cell averages -- + the opposite of the ne*pg2 physics grid that ``EAMSliceSource`` reads. + + This is the "tier 1" representation: each element is split into + (np-1)^2 = 9 bilinear quads whose vertices are the GLL nodes themselves. + That connectivity is read straight from the grid file's ``element_corners`` + array, so no topology is reconstructed and no points are invented -- the + subdivision vertices *are* the data locations. Variables are therefore + attached as **point data**. + + Grid file (e.g. ne30np4_latlon.nc, written by HOMME2META.ncl): + lat(ncol), lon(ncol) GLL node positions, degrees + element_corners(ncorners, ncells) 1-based, element-major + + The reader lays the sphere flat itself, duplicating nodes at the date line + and at the poles rather than clipping, so every cell stays whole and no + value is ever interpolated. Output is in [lon_origin, lon_origin+360) and + feeds EAMProject directly. + """ + + def __init__(self): + super().__init__() + self._lon_origin = -180.0 + # GLL grid state + self._gll_lat = None # (ncol,) degrees + self._gll_lon = None # (ncol,) degrees + self._cell_verts = None # (ncells, 4) indices into ncol + self._node_source = None # (npoints,) -> ncol index, for gathering + self._winding_flipped = False + + # ParaView builds a proxy's XML from the methods in the class's own + # __dict__, so these have to be declared here rather than inherited. + @smproperty.doublevector( + name="TimestepValues", information_only="1", si_class="vtkSITimeStepsProperty" + ) + def GetTimestepValues(self): + return self._timeSteps + + @smproperty.dataarrayselection(name="Variables") + def GetSurfaceVariables(self): + return self._variable_selection + + # -- properties ---------------------------------------------------- + + def SetLongitudeOrigin(self, origin): + if self._lon_origin != origin: + self._lon_origin = origin + self._clear_geometry_cache() + self._dirty = True + self.Modified() + + def GetNodeSource(self): + """Map from output point id to GLL node id (None before execution).""" + return self._node_source + + # -- overrides ----------------------------------------------------- + + def _clear_geometry_cache(self): + super()._clear_geometry_cache() + self._gll_lat = None + self._gll_lon = None + self._cell_verts = None + self._node_source = None + + def _identify_horizontal_dimension(self, meshdata, vardata): + """Identify the GLL node dimension (ncol) and match it in the data file. + + The base class takes the first dimension of the connectivity file, + which happens to be ``ncol`` for ne30np4_latlon.nc but is not something + to rely on -- take it from the ``lat`` variable instead. + """ + if self._horizontal_dim and self._data_horizontal_dim: + return + + if "lat" not in meshdata.variables: + print_error("Grid file has no 'lat' variable; not an np4 grid file") + return + + self._horizontal_dim = meshdata.variables["lat"].dimensions[0] + n_nodes = meshdata.dimensions[self._horizontal_dim].size + + # Prefer a same-named dimension in the data file, else match by size. + dim = vardata.dimensions.get(self._horizontal_dim) + if dim is not None and dim.size == n_nodes: + self._data_horizontal_dim = self._horizontal_dim + return + + for dim_name, dim_obj in vardata.dimensions.items(): + if dim_obj.size == n_nodes: + self._data_horizontal_dim = dim_name + return + + print_error( + f"Could not match GLL node count {n_nodes} to any dimension in the data file" + ) + + def _read_grid(self, meshdata): + """Read GLL node positions and the 9-subcell connectivity.""" + lat = np.asarray(meshdata["lat"][:]).reshape(-1).astype(np.float64) + lon = np.asarray(meshdata["lon"][:]).reshape(-1).astype(np.float64) + + if "element_corners" not in meshdata.variables: + print_error("Grid file has no 'element_corners'; not an np4 grid file") + return False + + # element_corners is (ncorners, ncells), 1-based -> (ncells, 4), 0-based + ec = np.asarray(meshdata["element_corners"][:]).astype(np.int64) + verts = np.ascontiguousarray(ec.T) - 1 + + if verts.min() < 0 or verts.max() >= len(lat): + print_error( + f"element_corners indexes outside [0, {len(lat)}) after converting " + "from 1-based; grid file may use a different convention" + ) + return False + + self._gll_lat = lat + self._gll_lon = lon + self._cell_verts = verts + return True + + def _orient_outward(self): + """Reverse subcell winding if HOMME's corner order faces normals inward. + + A quad's normal is taken from the cross product of its diagonals and + compared with the outward radial direction at its centroid. HOMME winds + ``element_corners`` inward uniformly, so one whole-mesh test settles it. + """ + lat_r = np.radians(self._gll_lat) + lon_r = np.radians(self._gll_lon) + cos_lat = np.cos(lat_r) + xyz = np.column_stack( + [cos_lat * np.cos(lon_r), cos_lat * np.sin(lon_r), np.sin(lat_r)] + ) + + v = self._cell_verts + normals = np.cross(xyz[v[:, 2]] - xyz[v[:, 0]], xyz[v[:, 3]] - xyz[v[:, 1]]) + outward = (normals * xyz[v].mean(axis=1)).sum(axis=1) + n_inward = int((outward < 0).sum()) + + if n_inward == v.shape[0]: + self._cell_verts = np.ascontiguousarray(v[:, ::-1]) + self._winding_flipped = True + elif n_inward: + print_warning( + f"subcell winding is not consistent ({n_inward} of {v.shape[0]} " + "cells wind inward); leaving orientation as found" + ) + + def _latlon_mesh(self): + """Lay the sphere flat, splitting the date line and the poles. + + Two degeneracies have to be dealt with, and both are fixed by + duplicating nodes rather than by clipping -- so every cell stays whole + and no value is interpolated. + + *Date line.* A cell whose corners fall either side of the seam would + stretch across the whole map. Its low-side corners get a duplicate + shifted +360 degrees, leaving the cell whole at the right-hand edge. + + *Poles.* Longitude is undefined at a pole, so the file stores an + arbitrary one (0). Every cell touching a pole node is dragged toward + that meridian. Each such cell gets its own copy of the pole node at the + mean longitude of its other corners, which turns the pole from a single + point into the top edge of the map. + + Sets ``self._node_source``, mapping each output point back to its GLL + node so field arrays can be gathered with a single fancy index. + """ + lon_origin = self._lon_origin + lon = lon_origin + np.mod(self._gll_lon - lon_origin, 360.0) + lat = self._gll_lat + v = self._cell_verts + n0 = len(lat) + + extra_lon = [] + extra_src = [] + new_verts = v.copy() + + def add(orig, lon_value): + extra_lon.append(float(lon_value)) + extra_src.append(int(orig)) + return n0 + len(extra_src) - 1 + + def lon_of(idx): + return lon[idx] if idx < n0 else extra_lon[idx - n0] + + pole = np.where(np.abs(np.abs(lat) - 90.0) < 1e-9)[0] + is_pole_node = np.zeros(n0, dtype=bool) + is_pole_node[pole] = True + + # 1. date line + # + # A pole node's longitude is arbitrary (the file stores 0), so it must + # take no part in deciding whether a cell straddles the seam -- it is + # replaced below anyway. Including it makes a polar cell's span read as + # a full half-turn, and float noise then tips the comparison over and + # flings a legitimate corner a whole turn out of the map. + cl = np.where(is_pole_node[v], np.nan, lon[v]) + real_span = np.nanmax(cl, axis=1) - np.nanmin(cl, axis=1) + seam = real_span > 180.0 + midline = lon_origin + 180.0 + shifted = {} + for c in np.where(seam)[0]: + for k in range(4): + n = int(v[c, k]) + if is_pole_node[n]: + continue + if lon[n] < midline: + if n not in shifted: + shifted[n] = add(n, lon[n] + 360.0) + new_verts[c, k] = shifted[n] + + # 2. poles + if pole.size: + orig_all = np.concatenate( + [ + np.arange(n0, dtype=np.int64), + np.array(extra_src, dtype=np.int64) + if extra_src + else np.empty(0, dtype=np.int64), + ] + ) + is_pole = np.isin(orig_all, pole) + for c in np.where(is_pole[new_verts].any(axis=1))[0]: + for k in range(4): + idx = int(new_verts[c, k]) + orig = idx if idx < n0 else extra_src[idx - n0] + if orig in pole: + others = [ + lon_of(int(new_verts[c, j])) for j in range(4) if j != k + ] + new_verts[c, k] = add(orig, np.mean(others)) + + if extra_src: + src_extra = np.array(extra_src, dtype=np.int64) + source = np.concatenate([np.arange(n0, dtype=np.int64), src_extra]) + out_lon = np.concatenate([lon, np.array(extra_lon, dtype=np.float64)]) + out_lat = np.concatenate([lat, lat[src_extra]]) + else: + source = np.arange(n0, dtype=np.int64) + out_lon, out_lat = lon, lat + + self._node_source = source + return out_lon, out_lat, new_verts + + def _build_geometry(self, meshdata): + """Build and cache the tier-1 subdivided-element mesh.""" + if self._cached_points is not None: + return + + if not self._read_grid(meshdata): + return + + self._orient_outward() + lon, lat, verts = self._latlon_mesh() + + n_cells = verts.shape[0] + self._cached_ncells2D = n_cells + + points_type = np.float32 if self._ForceFloatPoints else np.float64 + coords = np.empty((len(lon), 3), dtype=points_type) + coords[:, 0] = lon + coords[:, 1] = lat + coords[:, 2] = 0.0 + + vtk_coords = vtkPoints() + vtk_coords.SetData(dsa.numpyTovtkDataArray(coords)) + self._cached_points = vtk_coords + + cellTypes = np.empty(n_cells, dtype=np.uint8) + cellTypes.fill(vtkConstants.VTK_QUAD) + self._cached_cell_types = numpy_support.numpy_to_vtk( + num_array=cellTypes.ravel(), + deep=True, + array_type=vtkConstants.VTK_UNSIGNED_CHAR, + ) + + offsets = np.arange(0, (4 * n_cells) + 1, 4, dtype=np.int64) + self._cached_offsets = numpy_support.numpy_to_vtk( + num_array=offsets.ravel(), deep=True, array_type=vtkConstants.VTK_ID_TYPE + ) + + self._cached_cells = numpy_support.numpy_to_vtk( + num_array=np.ascontiguousarray(verts).ravel(), + deep=True, + array_type=vtkConstants.VTK_ID_TYPE, + ) + + def _RequestDataImpl(self, request, inInfo, outInfo): + if ( + self._ConnFileName is None + or self._ConnFileName == "None" + or self._DataFileName is None + or self._DataFileName == "None" + ): + print_error( + "Either one or both, the data file or connectivity file, are not provided!" + ) + return 0 + if not _has_deps: + print_error("Required Python module 'netCDF4' or 'numpy' missing!") + return 0 + + meshdata = self._get_mesh_dataset() + vardata = self._get_var_dataset() + + self._identify_horizontal_dimension(meshdata, vardata) + if not self._horizontal_dim or not self._data_horizontal_dim: + print_error("Could not identify required dimensions from files") + return 0 + + self._build_geometry(meshdata) + if self._cached_points is None: + print_error("Could not build geometry from the np4 grid file") + return 0 + + output_mesh = dsa.WrapDataObject(self._output) + + if self._dirty: + self._output = vtkUnstructuredGrid() + output_mesh = dsa.WrapDataObject(self._output) + output_mesh.SetPoints(self._cached_points) + cellArray = vtkCellArray() + cellArray.SetData(self._cached_offsets, self._cached_cells) + output_mesh.VTKObject.SetCells(self._cached_cell_types, cellArray) + self._dirty = False + + # Values are nodal, so they are gathered onto the split point set + # through the node source map rather than used directly. + source = self._node_source + + to_remove = set() + for i in range(output_mesh.PointData.GetNumberOfArrays()): + to_remove.add(output_mesh.PointData.GetArrayName(i)) + + changed_dims = self._changed_dims + for name, varmeta in self._variables.items(): + if self._variable_selection.ArrayIsEnabled(name): + if output_mesh.PointData.HasArray(name): + to_remove.remove(name) + if changed_dims and not changed_dims.intersection( + varmeta.dimensions + ): + continue + data = self._load_variable(vardata, varmeta) + if data.size != len(source) and data.size == len(self._gll_lat): + data = np.ascontiguousarray(data[source]) + output_mesh.PointData.append(data, name) + + self._changed_dims = set() + + # CAM-SE files carry area(ncol): the GLL quadrature weight for each + # node. That is the correct weight for averaging nodal values, and it + # has to sit in PointData beside them to line up. + area_var_name = "area" + if self._areavar and not output_mesh.PointData.HasArray(area_var_name): + data = self._get_cached_area(vardata) + if data is not None and data.size == len(self._gll_lat): + output_mesh.PointData.append( + np.ascontiguousarray(data[source]), area_var_name + ) + if area_var_name in to_remove: + to_remove.remove(area_var_name) + + for var_name in to_remove: + output_mesh.PointData.RemoveArray(var_name) + + output = vtkUnstructuredGrid.GetData(outInfo, 0) + output.ShallowCopy(self._output) + + return 1 diff --git a/src/e3sm_quickview/utils/compute.py b/src/e3sm_quickview/utils/compute.py index d504adf..03f2676 100644 --- a/src/e3sm_quickview/utils/compute.py +++ b/src/e3sm_quickview/utils/compute.py @@ -35,11 +35,21 @@ def calculate_weighted_average( return float(np.mean(data)) -def extract_avgs(vtk_data, array_names): +def extract_avgs(vtk_data, array_names, association="cell"): + """Average each named array, weighted by "area" when the format supplies it. + + ``association`` says where the active format keeps its variables: "cell" + for the pg2 physics grid, "point" for the np4 dynamical core. The weights + have to come from the same attribute set as the values, or they would not + line up. + """ results = {} - area_array = vtk_data.GetCellData().GetArray("area") + attributes = ( + vtk_data.GetPointData() if association == "point" else vtk_data.GetCellData() + ) + area_array = attributes.GetArray("area") for name in array_names: - vtk_array = vtk_data.GetCellData().GetArray(name) + vtk_array = attributes.GetArray(name) if vtk_array is None: results[name] = np.nan continue diff --git a/src/e3sm_quickview/view_manager.py b/src/e3sm_quickview/view_manager.py index fa28744..319510b 100644 --- a/src/e3sm_quickview/view_manager.py +++ b/src/e3sm_quickview/view_manager.py @@ -306,6 +306,18 @@ def update_color_range(self): for view in list(self._var2view.values()): view.colormap.update_color_range() # colormaps module + def drop_views(self): + """Discard every view so the next layout rebinds to the active pipeline. + + A view's mapper is connected to the tail of whichever data path was + active when the view was built, so switching format leaves it pointing + at the wrong pipeline. compute_layout() rebuilds what is needed. + """ + for renderer in list(self._render_window.GetRenderers()): + self._render_window.RemoveRenderer(renderer) + self._var2view.clear() + self.layout_dirty = True + def get_view(self, variable_name, variable_type): view = self._var2view.get(variable_name) if view is None: diff --git a/src/e3sm_quickview/view_panel.py b/src/e3sm_quickview/view_panel.py index 9b065bc..8ee301e 100644 --- a/src/e3sm_quickview/view_panel.py +++ b/src/e3sm_quickview/view_panel.py @@ -64,7 +64,9 @@ def __init__(self, server, source, variable_name, variable_type, camera): server, mapper=self.mapper, data_array_fn=lambda: self.data_array, - ).set_data_array(variable_name, lambda: self.data_array, "cell") + ).set_data_array( + variable_name, lambda: self.data_array, source.data_reader.association + ) self.colormap.watch(["mapper_change"], lambda *_: self.render()) # GUI @@ -100,8 +102,11 @@ def render(self): @property def data_array(self): - self.source.data_reader.vtk_geometry.Update() - ds = self.source.data_reader.vtk_geometry.GetOutput() + data_reader = self.source.data_reader + data_reader.vtk_geometry.Update() + ds = data_reader.vtk_geometry.GetOutput() + if data_reader.association == "point": + return ds.GetPointData().GetArray(self.variable_name) return ds.GetCellData().GetArray(self.variable_name) def _build_ui(self):