From 6b099e4e4ad873b4703beb648035184ae8a8ed1f Mon Sep 17 00:00:00 2001 From: Berk Geveci Date: Thu, 17 Sep 2026 11:35:57 -0400 Subject: [PATCH] feat(dycore): read EAM native spectral-element grids, and make formats pluggable QuickView handled one grid: the ne*pg2 physics grid, whose SCRIP connectivity gives every cell its own four corners and whose values are cell averages. The dynamical core runs on GLL nodes shared between spectral elements, and its values are nodal point samples -- the opposite shape. Supporting both means the format, not the application, has to decide what the pipeline looks like. pipeline.py now has a DataPath base with one subclass per format, chosen by sniffing the connectivity file (element_corners vs corner_lat): Pg2Path cell data reader > CenterMeridian > Extract > Project > Surface DycorePath point data reader > CenterMeridian > Extract > Project > Surface EAMDycoreSource is the tier-1 reader: one bilinear quad per (np-1)^2 subcell, vertices *are* the GLL nodes, so no topology is reconstructed and no point is invented. It lays the sphere flat itself, duplicating nodes at the date line and the poles rather than clipping, which keeps every cell whole and every value exact. HOMME winds element_corners inward, so the winding is reversed at read time; a pole node's longitude is arbitrary and must take no part in the seam test, or float noise flings a corner a whole turn out of the map. The map's longitude window is now movable: the left edge sits anywhere in [-180, 180] and the right edge a turn east of it. One rule shared by every filter -- cut at c = input_origin + (origin - input_origin) % 360, translate the halves by origin+360-c and origin-c -- and at origin -180 with a [0, 360) input it reduces to the previous hardcoded behaviour. EAMCenterMeridian needs to be told its *input's* window because the two readers emit different ones. The projection is re-centred before PROJ sees it: PROJ normalises longitude into [-180, 180) before applying lon_0, which would send the window's right edge to the left rim and draw cells and coastlines across the map. Association is plumbed through view_panel, compute.extract_avgs and app.py, and ViewManager.drop_views() discards views on a format switch, since each binds its mapper to the tail of one specific pipeline. Performance: EAMCenterMeridian passes through untouched when the input already sits in the requested window (the dycore default), and on a real rebuild hands the cached points and cells objects downstream, because EAMProject keys its cache on point identity and EAMExtract on their modified time. EAMExtract keeps its cache for nodal data through a point pedigree map -- it removes whole cells and never interpolates, so the ids stay exact. Roughly 7 ms per pipeline pass at ne30 with 20 variables loaded, matching pg2. Verified against ParaView 6.0.1 (VTK 9.5.2) and ParaView master (VTK 9.7). --- src/e3sm_quickview/app.py | 30 +- src/e3sm_quickview/components/toolbars.py | 28 +- src/e3sm_quickview/pipeline.py | 275 ++++++++- src/e3sm_quickview/plugins/eam_projection.py | 329 ++++++++--- src/e3sm_quickview/plugins/eam_reader.py | 590 +++++++++++++++++-- src/e3sm_quickview/utils/compute.py | 16 +- src/e3sm_quickview/view_manager.py | 12 + src/e3sm_quickview/view_panel.py | 11 +- 8 files changed, 1129 insertions(+), 162 deletions(-) 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):