diff --git a/benchmarks/mpas_ocean.py b/benchmarks/mpas_ocean.py index 659c19014..c49b0908f 100644 --- a/benchmarks/mpas_ocean.py +++ b/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) diff --git a/uxarray/core/zonal.py b/uxarray/core/zonal.py index 24a957c26..63279cb29 100644 --- a/uxarray/core/zonal.py +++ b/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)) diff --git a/uxarray/grid/grid.py b/uxarray/grid/grid.py index f7d5e828c..44556043a 100644 --- a/uxarray/grid/grid.py +++ b/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, diff --git a/uxarray/grid/intersections.py b/uxarray/grid/intersections.py index dfd408c51..d92cdfb74 100644 --- a/uxarray/grid/intersections.py +++ b/uxarray/grid/intersections.py @@ -1,7 +1,7 @@ import math import numpy as np -from numba import njit, prange +from numba import njit from uxarray.constants import ERROR_TOLERANCE, INT_DTYPE from uxarray.grid.arcs import _on_minor_arc_xyz, on_minor_arc @@ -19,10 +19,17 @@ # Edge screeners: fast O(n) passes used by Grid.get_edges_at_constant_* to # identify candidate edges before the expensive GCA intersection. "no_extreme" # means arc z-extrema along the great circle are not considered. +# +# These are deliberately plain NumPy (no @njit): they are memory-bound elementwise +# predicates where NumPy is ~2x faster than a Numba prange loop, -@njit(parallel=True, nogil=True, cache=True) -def constant_lat_intersections_no_extreme(lat, edge_node_z, n_edge): +def _flatnonzero(mask): + """Sorted indices where a 1-D boolean mask is True""" + return np.asarray(np.flatnonzero(mask)) + + +def constant_lat_intersections_no_extreme(lat, edge_node_z): """Determine which edges intersect a constant line of latitude on a sphere, without wrapping to the opposite longitude, with extremes along each great circle arc not considered. @@ -31,55 +38,30 @@ def constant_lat_intersections_no_extreme(lat, edge_node_z, n_edge): ---------- lat: Constant latitude value in degrees. - edge_node_x: - Array of shape (n_edge, 2) containing x-coordinates of the edge nodes. - edge_node_y: - Array of shape (n_edge, 2) containing y-coordinates of the edge nodes. - n_edge: - Total number of edges to check. + May be NumPy or dask array. + edge_node_z: + Array of shape (n_edge, 2) containing z-coordinates of the edge nodes. + May be NumPy or dask array. Returns ------- intersecting_edges: array of indices of edges that intersect the constant latitude. """ - lat = np.deg2rad(lat) - - intersecting_edges_mask = np.zeros(n_edge, dtype=np.int32) - - # Calculate the constant z-value for the given latitude - z_constant = np.sin(lat) - - # Iterate through each edge and check for intersections - for i in prange(n_edge): - # Check if the edge crosses the constant latitude or lies exactly on it - if edge_intersects_constant_lat_no_extreme(edge_node_z[i], z_constant): - intersecting_edges_mask[i] = 1 + z_constant = np.sin(np.deg2rad(lat)) - intersecting_edges = np.argwhere(intersecting_edges_mask) + d0 = edge_node_z[:, 0] - z_constant + d1 = edge_node_z[:, 1] - z_constant - return np.unique(intersecting_edges) - - -@njit(cache=True, nogil=True) -def edge_intersects_constant_lat_no_extreme(edge_node_z, z_constant): - """Helper to compute whether an edge intersects a line of constant latitude.""" - - # z coordinate of edge nodes - z0 = edge_node_z[0] - z1 = edge_node_z[1] + # Edge crosses the latitude (endpoints straddle it) or lies exactly on it. + intersecting = (d0 * d1 < 0.0) | ( + (np.abs(d0) < ERROR_TOLERANCE) & (np.abs(d1) < ERROR_TOLERANCE) + ) - if (z0 - z_constant) * (z1 - z_constant) < 0.0 or ( - abs(z0 - z_constant) < ERROR_TOLERANCE - and abs(z1 - z_constant) < ERROR_TOLERANCE - ): - return True - else: - return False + return _flatnonzero(intersecting) -@njit(parallel=True, nogil=True, cache=True) -def constant_lon_intersections_no_extreme(lon, edge_node_x, edge_node_y, n_edge): +def constant_lon_intersections_no_extreme(lon, edge_node_x, edge_node_y): """Determine which edges intersect a constant line of longitude on a sphere, without wrapping to the opposite longitude, with extremes along each great circle arc not considered. @@ -90,50 +72,42 @@ def constant_lon_intersections_no_extreme(lon, edge_node_x, edge_node_y, n_edge) Constant longitude value in degrees. edge_node_x: Array of shape (n_edge, 2) containing x-coordinates of the edge nodes. + May be NumPy or dask array. edge_node_y: Array of shape (n_edge, 2) containing y-coordinates of the edge nodes. - n_edge: - Total number of edges to check. + May be NumPy or dask array. Returns ------- intersecting_edges: array of indices of edges that intersect the constant longitude. """ - lon = np.deg2rad(lon) - - intersecting_edges_mask = np.zeros(n_edge, dtype=np.int32) - - # calculate the cos and sin of the constant longitude cos_lon = np.cos(lon) sin_lon = np.sin(lon) - for i in prange(n_edge): - # get the x and y coordinates of the edge's nodes - x0, x1 = edge_node_x[i, 0], edge_node_x[i, 1] - y0, y1 = edge_node_y[i, 0], edge_node_y[i, 1] + x0 = edge_node_x[:, 0] + x1 = edge_node_x[:, 1] + y0 = edge_node_y[:, 0] + y1 = edge_node_y[:, 1] - # calculate the dot products to determine on which side of the constant longitude the points lie - dot0 = x0 * sin_lon - y0 * cos_lon - dot1 = x1 * sin_lon - y1 * cos_lon + # Signed distance to the meridian plane for each endpoint. + dot0 = x0 * sin_lon - y0 * cos_lon + dot1 = x1 * sin_lon - y1 * cos_lon - # ensure that both points are not on the opposite longitude (180 degrees away) - if (x0 * cos_lon + y0 * sin_lon) < 0.0 or (x1 * cos_lon + y1 * sin_lon) < 0.0: - continue - - # check if the edge crosses the constant longitude or lies exactly on it - if dot0 * dot1 < 0.0 or ( - abs(dot0) < ERROR_TOLERANCE and abs(dot1) < ERROR_TOLERANCE - ): - intersecting_edges_mask[i] = 1 + # Discard edges with an endpoint on the opposite meridian (180 deg away). + not_opposite = (x0 * cos_lon + y0 * sin_lon >= 0.0) & ( + x1 * cos_lon + y1 * sin_lon >= 0.0 + ) - intersecting_edges = np.argwhere(intersecting_edges_mask) + # Edge crosses the longitude or lies exactly on it. + crosses = (dot0 * dot1 < 0.0) | ( + (np.abs(dot0) < ERROR_TOLERANCE) & (np.abs(dot1) < ERROR_TOLERANCE) + ) - return np.unique(intersecting_edges) + return _flatnonzero(not_opposite & crosses) -@njit(cache=True) def constant_lat_intersections_face_bounds(lat: float, face_bounds_lat: np.ndarray): """ Identify candidate faces that intersect with a given constant latitude line. @@ -153,15 +127,10 @@ def constant_lat_intersections_face_bounds(lat: float, face_bounds_lat: np.ndarr A 1D array of integers containing the indices of the faces that intersect the given latitude. """ - face_bounds_lat_min = face_bounds_lat[:, 0] - face_bounds_lat_max = face_bounds_lat[:, 1] - - within_bounds = (face_bounds_lat_min <= lat) & (face_bounds_lat_max >= lat) - candidate_faces = np.where(within_bounds)[0] - return candidate_faces + within_bounds = (face_bounds_lat[:, 0] <= lat) & (face_bounds_lat[:, 1] >= lat) + return _flatnonzero(within_bounds) -@njit(cache=True) def constant_lon_intersections_face_bounds(lon: float, face_bounds_lon: np.ndarray): """ Identify candidate faces that intersect with a given constant longitude line. @@ -183,25 +152,18 @@ def constant_lon_intersections_face_bounds(lon: float, face_bounds_lon: np.ndarr """ face_bounds_lon_min = face_bounds_lon[:, 0] face_bounds_lon_max = face_bounds_lon[:, 1] - n_face = face_bounds_lon.shape[0] - - candidate_faces = [] - for i in range(n_face): - cur_face_bounds_lon_min = face_bounds_lon_min[i] - cur_face_bounds_lon_max = face_bounds_lon_max[i] - if cur_face_bounds_lon_min < cur_face_bounds_lon_max: - if (lon >= cur_face_bounds_lon_min) and (lon <= cur_face_bounds_lon_max): - candidate_faces.append(i) - else: - # antimeridian case - if (lon >= cur_face_bounds_lon_min) or (lon <= cur_face_bounds_lon_max): - candidate_faces.append(i) + # Normal faces (min < max): lon inside [min, max]. + normal = face_bounds_lon_min < face_bounds_lon_max + in_normal = normal & (lon >= face_bounds_lon_min) & (lon <= face_bounds_lon_max) + # Antimeridian faces (min >= max): lon >= min OR lon <= max. + in_antimeridian = (~normal) & ( + (lon >= face_bounds_lon_min) | (lon <= face_bounds_lon_max) + ) - return np.array(candidate_faces, dtype=INT_DTYPE) + return _flatnonzero(in_normal | in_antimeridian).astype(INT_DTYPE, copy=False) -@njit(cache=True) def faces_within_lon_bounds(lons, face_bounds_lon): """ Identify candidate faces that lie within a specified longitudinal interval. @@ -225,51 +187,39 @@ def faces_within_lon_bounds(lons, face_bounds_lon): """ face_bounds_lon_min = face_bounds_lon[:, 0] face_bounds_lon_max = face_bounds_lon[:, 1] - n_face = face_bounds_lon.shape[0] min_lon, max_lon = lons # For example, a query of (160, -160) would cross the antimeridian antimeridian = min_lon > max_lon - candidate_faces = [] - for i in range(n_face): - cur_face_min = face_bounds_lon_min[i] - cur_face_max = face_bounds_lon_max[i] - - # Check if the face itself crosses the antimeridian - face_crosses_antimeridian = cur_face_min > cur_face_max - - if not antimeridian: - # Normal case: min_lon <= max_lon - # Face must be strictly contained within [min_lon, max_lon] - if not face_crosses_antimeridian: - if (cur_face_min >= min_lon) and (cur_face_max <= max_lon): - candidate_faces.append(i) - else: - # If face crosses antimeridian, it cannot be strictly contained - # in a non-antimeridian query interval - continue - else: - # Antimeridian case: interval crosses the -180/180 boundary - # The query interval is effectively [min_lon, 180] U [-180, max_lon] - - if face_crosses_antimeridian: - # If face crosses antimeridian, check if it's contained in the full query range - if (cur_face_min >= min_lon) and (cur_face_max <= max_lon): - candidate_faces.append(i) - else: - # For non-crossing faces, check if they're strictly contained in either part - contained_part1 = (cur_face_min >= min_lon) and (cur_face_max <= 180) - contained_part2 = (cur_face_min >= -180) and (cur_face_max <= max_lon) - - if contained_part1 or contained_part2: - candidate_faces.append(i) + # A face itself crosses the antimeridian when its stored min > max. + face_crosses = face_bounds_lon_min > face_bounds_lon_max + + if not antimeridian: + # Normal query interval [min_lon, max_lon]. Only non-crossing faces can + # be strictly contained; crossing faces are excluded. + mask = ( + (~face_crosses) + & (face_bounds_lon_min >= min_lon) + & (face_bounds_lon_max <= max_lon) + ) + else: + # Antimeridian query: effectively [min_lon, 180] U [-180, max_lon]. + crossing_contained = ( + face_crosses + & (face_bounds_lon_min >= min_lon) + & (face_bounds_lon_max <= max_lon) + ) + noncrossing_contained = (~face_crosses) & ( + ((face_bounds_lon_min >= min_lon) & (face_bounds_lon_max <= 180)) + | ((face_bounds_lon_min >= -180) & (face_bounds_lon_max <= max_lon)) + ) + mask = crossing_contained | noncrossing_contained - return np.array(candidate_faces, dtype=INT_DTYPE) + return _flatnonzero(mask).astype(INT_DTYPE, copy=False) -@njit(cache=True) def faces_within_lat_bounds(lats, face_bounds_lat): """ Identify candidate faces that lie within a specified latitudinal interval. @@ -297,8 +247,7 @@ def faces_within_lat_bounds(lats, face_bounds_lat): face_bounds_lat_max = face_bounds_lat[:, 1] within_bounds = (face_bounds_lat_max <= max_lat) & (face_bounds_lat_min >= min_lat) - candidate_faces = np.where(within_bounds)[0] - return candidate_faces + return _flatnonzero(within_bounds) @njit(cache=True, inline="always", error_model="numpy") diff --git a/uxarray/grid/utils.py b/uxarray/grid/utils.py index 9287bc325..a70b2b0aa 100644 --- a/uxarray/grid/utils.py +++ b/uxarray/grid/utils.py @@ -1,6 +1,6 @@ import numpy as np import xarray as xr -from numba import njit +from numba import njit, prange from uxarray.constants import INT_FILL_VALUE @@ -260,6 +260,60 @@ def _get_cartesian_face_edge_nodes_array( return face_edges_cartesian.reshape(n_face, n_max_face_edges, 2, 3) +@njit(cache=True, parallel=True) +def _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, +): + """Build the Cartesian edge array for a subset of faces. + + Parameters + ---------- + face_indices : np.ndarray + 1D array of face indices to build, shape ``(n_selected,)``. + face_node_connectivity : np.ndarray + Face-node connectivity, shape ``(n_face, n_max_face_edges)``. + n_nodes_per_face : np.ndarray + Number of non-fill nodes (== edges) for each face, shape ``(n_face,)``. + n_max_face_edges : int + Maximum number of edges for any face in the grid. + node_x, node_y, node_z : np.ndarray + Cartesian node coordinates, each shape ``(n_node,)``. + + Returns + ------- + np.ndarray + Array of shape ``(n_selected, n_max_face_edges, 2, 3)``. Unused edge + slots are padded with ``INT_FILL_VALUE`` cast to float + """ + n_selected = face_indices.shape[0] + out = np.full( + (n_selected, n_max_face_edges, 2, 3), INT_FILL_VALUE, dtype=np.float64 + ) + + for i in prange(n_selected): + f = face_indices[i] + n_edges = n_nodes_per_face[f] + for e in range(n_edges): + n_start = face_node_connectivity[f, e] + # Wrap the last edge back to the first node to close the polygon, + # mirroring the np.roll(-1) used by the whole-grid builder. + n_end = face_node_connectivity[f, (e + 1) % n_edges] + out[i, e, 0, 0] = node_x[n_start] + out[i, e, 0, 1] = node_y[n_start] + out[i, e, 0, 2] = node_z[n_start] + out[i, e, 1, 0] = node_x[n_end] + out[i, e, 1, 1] = node_y[n_end] + out[i, e, 1, 2] = node_z[n_end] + + return out + + def _get_lonlat_rad_face_edge_nodes_array( face_node_conn, n_face, n_max_face_edges, node_lon, node_lat ):