Better chunking strategies for constlat intersections and zonal routines.

benchmarks/mpas_ocean.py

@@ -229,3 +229,46 @@ def setup(self, resolution, *args, **kwargs):
     def time_zonal_average(self, resolution):
         lat_step = 10
         self.uxds['bottomDepth'].zonal_mean(lat=(-45, 45, lat_step))
+
+
+class ZonalAveragePeakMem:
+    """Peak memory of a cold-start non-conservative zonal-mean sweep."""
+
+    param_names = ["resolution"]
+    params = [["480km", "120km"]]
+
+    def setup_cache(self):
+        """Compile the njit kernels before anything is measured."""
+        for resolution in self.params[0]:
+            grid, data = file_path_dict[resolution]
+            uxds = ux.open_dataset(grid, data)
+            uxds.uxgrid.bounds
+            uxds[data_var].zonal_mean(lat=(-45, 45, 10))
+
+    def peakmem_zonal_average(self, resolution):
+        grid, data = file_path_dict[resolution]
+        uxds = ux.open_dataset(grid, data)
+        uxds.uxgrid.bounds
+        uxds[data_var].zonal_mean(lat=(-45, 45, 10))
+
+
+class CrossSectionsPeakMem:
+    """Peak memory of a cold-start constant-latitude cross-section sweep."""
+
+    param_names = ["resolution", "lat_step"]
+    params = [["480km", "120km"], [1, 2, 4]]
+
+    def setup_cache(self):
+        """Compile the njit kernels before anything is measured."""
+        for resolution in self.params[0]:
+            uxgrid = ux.open_grid(file_path_dict[resolution][0])
+            uxgrid.normalize_cartesian_coordinates()
+            uxgrid.bounds
+            uxgrid.cross_section.constant_latitude(0.0)
+
+    def peakmem_const_lat(self, resolution, lat_step):
+        uxgrid = ux.open_grid(file_path_dict[resolution][0])
+        uxgrid.normalize_cartesian_coordinates()
+        uxgrid.bounds
+        for lat in np.arange(-45, 45, lat_step):
+            uxgrid.cross_section.constant_latitude(lat)

uxarray/core/zonal.py

@@ -9,7 +9,7 @@
     get_number_of_intersections,
 )
 from uxarray.grid.utils import (
-    _get_cartesian_face_edge_nodes_array,
+    _get_cartesian_face_edge_nodes_array_subset,
     _small_angle_of_2_vectors,
 )
 
@@ -40,14 +40,15 @@ def _compute_non_conservative_zonal_mean(uxda, latitudes, use_robust_weights=Fal
         # Create a NumPy array for storing results
         result = np.full(shape, np.nan, dtype=result_dtype)
 
-    faces_edge_nodes_xyz = _get_cartesian_face_edge_nodes_array(
-        uxgrid.face_node_connectivity.values,
-        uxgrid.n_face,
-        uxgrid.n_max_face_nodes,
-        uxgrid.node_x.values,
-        uxgrid.node_y.values,
-        uxgrid.node_z.values,
-    )
+    # Grid arrays are read once here and passed to the per-latitude subset
+    # builder. The whole-grid (n_face, n_max, 2, 3) edge array is never
+    # materialized; only the faces intersecting each latitude are built, so peak
+    # memory scales with the candidate count (~1% of n_face) instead of n_face.
+    face_node_connectivity = uxgrid.face_node_connectivity.values
+    n_max_face_edges = uxgrid.n_max_face_nodes
+    node_x = uxgrid.node_x.values
+    node_y = uxgrid.node_y.values
+    node_z = uxgrid.node_z.values
 
     bounds = uxgrid.bounds.values
 
@@ -63,7 +64,16 @@ def _compute_non_conservative_zonal_mean(uxda, latitudes, use_robust_weights=Fal
 
         z = np.sin(np.deg2rad(lat))
 
-        fe = faces_edge_nodes_xyz[face_indices]
+        fe = _get_cartesian_face_edge_nodes_array_subset(
+            face_indices,
+            face_node_connectivity,
+            n_nodes_per_face,
+            n_max_face_edges,
+            node_x,
+            node_y,
+            node_z,
+        )
+
         nn = n_nodes_per_face[face_indices]
         b = bounds[face_indices]
 
@@ -258,14 +268,13 @@ def _compute_face_band_weights(uxgrid, bands):
             f"bands must be monotonic non-decreasing; got diff(bands)={np.diff(bands)}"
         )
 
-    faces_edge_nodes_xyz = _get_cartesian_face_edge_nodes_array(
-        uxgrid.face_node_connectivity.values,
-        uxgrid.n_face,
-        uxgrid.n_max_face_nodes,
-        uxgrid.node_x.values,
-        uxgrid.node_y.values,
-        uxgrid.node_z.values,
-    )
+    # Read grid arrays once; the per-band partial-face edge subsets are built
+    # on demand below rather than materializing the whole-grid edge array.
+    face_node_connectivity = uxgrid.face_node_connectivity.values
+    n_max_face_edges = uxgrid.n_max_face_nodes
+    node_x = uxgrid.node_x.values
+    node_y = uxgrid.node_y.values
+    node_z = uxgrid.node_z.values
     n_nodes_per_face = uxgrid.n_nodes_per_face.values
     face_bounds_lat = uxgrid.face_bounds_lat.values
     face_areas = uxgrid.face_areas.values
@@ -301,11 +310,22 @@ def _compute_face_band_weights(uxgrid, bands):
 
         partial = all_overlapping[~fc_mask]
         partial_pos = np.nonzero(~fc_mask)[0]
-        for pos, f in zip(partial_pos, partial):
-            nedge = n_nodes_per_face[f]
-            weights[pos] = _compute_band_overlap_area(
-                faces_edge_nodes_xyz[f, :nedge], zmin, zmax
+        if partial.size:
+            # Build edges for only the partially-overlapping faces of this band.
+            fe_partial = _get_cartesian_face_edge_nodes_array_subset(
+                partial,
+                face_node_connectivity,
+                n_nodes_per_face,
+                n_max_face_edges,
+                node_x,
+                node_y,
+                node_z,
             )
+            for k, (pos, f) in enumerate(zip(partial_pos, partial)):
+                nedge = n_nodes_per_face[f]
+                weights[pos] = _compute_band_overlap_area(
+                    fe_partial[k, :nedge], zmin, zmax
+                )
 
         per_band.append((all_overlapping.astype(np.int64), weights))
 

uxarray/grid/grid.py

@@ -2576,9 +2576,14 @@ def get_edges_at_constant_latitude(self, lat: float, use_face_bounds: bool = Fal
                 "is not yet supported."
             )
         else:
-            edges = constant_lat_intersections_no_extreme(
-                lat, self.edge_node_z.values, self.n_edge
-            )
+            # Gather per-edge z-coords positionally, mirroring the longitude
+            # sibling. A concrete connectivity indexer against node_z.data keeps
+            # node coords lazy on a chunked grid (dask indexed by a numpy array)
+            # and — unlike xarray indexing with a dask connectivity — does not
+            # raise. The screener then reduces the gathered array to candidates.
+            edge_nodes = self.edge_node_connectivity.values
+            edge_node_z = self.node_z.data[edge_nodes.ravel()].reshape(edge_nodes.shape)
+            edges = constant_lat_intersections_no_extreme(lat, edge_node_z)
 
         return edges.squeeze()
 
@@ -2608,7 +2613,7 @@ def get_faces_at_constant_latitude(
 
         faces = constant_lat_intersections_face_bounds(
             lat=lat,
-            face_bounds_lat=self.face_bounds_lat.values,
+            face_bounds_lat=self.face_bounds_lat.data,
         )
         return faces
 
@@ -2643,11 +2648,14 @@ def get_edges_at_constant_longitude(
                 "is not yet supported."
             )
         else:
-            edge_node_x = self.node_x[self.edge_node_connectivity].values
-            edge_node_y = self.node_y[self.edge_node_connectivity].values
-            edges = constant_lon_intersections_no_extreme(
-                lon, edge_node_x, edge_node_y, self.n_edge
-            )
+            # Positional gather of edge endpoint coords: a concrete connectivity
+            # indexer against node_[xy].data keeps node coords lazy on a chunked
+            # grid and does not raise (xarray vindex rejects a dask indexer).
+            edge_nodes = self.edge_node_connectivity.values
+            flat = edge_nodes.ravel()
+            edge_node_x = self.node_x.data[flat].reshape(edge_nodes.shape)
+            edge_node_y = self.node_y.data[flat].reshape(edge_nodes.shape)
+            edges = constant_lon_intersections_no_extreme(lon, edge_node_x, edge_node_y)
             return edges.squeeze()
 
     def get_faces_at_constant_longitude(self, lon: float):
@@ -2671,7 +2679,7 @@ def get_faces_at_constant_longitude(self, lon: float):
                 f"Longitude must be between -180 and 180 degrees. Received {lon}"
             )
 
-        faces = constant_lon_intersections_face_bounds(lon, self.face_bounds_lon.values)
+        faces = constant_lon_intersections_face_bounds(lon, self.face_bounds_lon.data)
         return faces
 
     def get_faces_between_longitudes(self, lons: tuple[float, float]):
@@ -2688,7 +2696,7 @@ def get_faces_between_longitudes(self, lons: tuple[float, float]):
             An array of face indices that are strictly between two lines of constant longitude.
 
         """
-        return faces_within_lon_bounds(lons, self.face_bounds_lon.values)
+        return faces_within_lon_bounds(lons, self.face_bounds_lon.data)
 
     def get_faces_between_latitudes(self, lats: tuple[float, float]):
         """Identifies the indices of faces that are strictly between two lines of constant latitude.
@@ -2704,7 +2712,7 @@ def get_faces_between_latitudes(self, lats: tuple[float, float]):
             An array of face indices that are strictly between two lines of constant latitude.
 
         """
-        return faces_within_lat_bounds(lats, self.face_bounds_lat.values)
+        return faces_within_lat_bounds(lats, self.face_bounds_lat.data)
 
     def get_faces_containing_point(
         self,