From 6ff99ead89b1f19d4bb36147fc9a80aa61870870 Mon Sep 17 00:00:00 2001 From: Alexandre Georges Date: Wed, 5 Aug 2026 13:08:27 -0700 Subject: [PATCH 1/5] Refactor(integrate)!: atomized integration functions to single-timestep inputs --- suxarray/grid/integrate.py | 163 ++++++++++++++++++++++--------------- 1 file changed, 97 insertions(+), 66 deletions(-) diff --git a/suxarray/grid/integrate.py b/suxarray/grid/integrate.py index 8529c61..fc1e7b2 100644 --- a/suxarray/grid/integrate.py +++ b/suxarray/grid/integrate.py @@ -2,6 +2,33 @@ from .area import _integrate_nodal +def _coerce_single_timestep(arr, name): + """Return a single-timestep 2D view for (node, layer) kernels. + + Accepts: + - 2D arrays with shape (n_node, n_layer) + - 3D arrays with shape (1, n_node, n_layer), which are squeezed + + Raises for multi-timestep 3D inputs. + """ + if arr.ndim == 2: + return arr + + if arr.ndim == 3: + if arr.shape[0] == 1: + return arr[0] + raise ValueError( + f"{name} must contain exactly one timestep for this kernel. " + f"Got shape {arr.shape}. Select one timestep first (for example, " + f"isel(time=i)) and parallelize externally with xarray/dask." + ) + + raise ValueError( + f"{name} must be 2D (n_node, n_layer) or 3D with a single time " + f"step (1, n_node, n_layer). Got ndim={arr.ndim}." + ) + + def _earth_radius_at_latitude(latitude_deg): """Compute geocentric Earth radius (meters) at geodetic latitude. @@ -28,24 +55,24 @@ def _face_mean_from_nodal(values, connectivity): Parameters ---------- - values : ndarray, shape (n_time, n_node, n_interface) + values : ndarray, shape (n_node, n_interface) Nodal values to average. connectivity : ndarray, shape (n_face, n_max_face_nodes) Face-node connectivity with negative fill values. Returns ------- - ndarray, shape (n_time, n_interface, n_face) + ndarray, shape (n_interface, n_face) Face-averaged values. """ - n_time, _, n_interface = values.shape + n_node, n_interface = values.shape n_face = connectivity.shape[0] - face_mean = np.zeros((n_time, n_interface, n_face), dtype=np.float64) + face_mean = np.zeros((n_interface, n_face), dtype=np.float64) for face_i in range(n_face): node_ids = connectivity[face_i] node_ids = node_ids[node_ids >= 0] - face_mean[:, :, face_i] = values[:, node_ids, :].mean(axis=1) + face_mean[:, face_i] = values[node_ids, :].mean(axis=0) return face_mean @@ -55,21 +82,25 @@ def _calc_layers_thickness(z, bottom_indices): Calculate the thickness of each layer under a set of nodes. Parameters - z: ndarray, shape (n_time, n_node, n_layer) + z : ndarray, shape (n_node, n_layer) z-coordinates of nodes - bottom_indices: ndarray, shape (n_node,) + bottom_indices : ndarray, shape (n_node,) 0-based bottom layer index for each node Returns - ndarray, shape (n_time, n_layer-1, n_node) + ndarray, shape (n_node, n_layer-1) """ + if z.ndim != 2: + raise ValueError( + "z must be 2D (n_node, n_layer) for thickness calculations." + ) - thickness = z[:, :, 1:] - z[:, :, :-1] + thickness = z[:, 1:] - z[:, :-1] if np.any(bottom_indices > 0): - layer_idx = np.arange(thickness.shape[2]) + layer_idx = np.arange(thickness.shape[-1]) below_bottom = layer_idx[np.newaxis, :] < bottom_indices[:, np.newaxis] - thickness = np.where(below_bottom[np.newaxis, :, :], 0.0, thickness) + thickness = np.where(below_bottom, 0.0, thickness) return thickness @@ -85,8 +116,8 @@ def _calculate_prism_volumes(node_x, node_y, z, connectivity, bottom_indices): x coordinates of nodes node_y : ndarray, shape (n_node,) y coordinates of nodes - z : ndarray, shape (n_time, n_node, n_layer) - z-coordinates + z : ndarray, shape (n_node, n_layer) + z-coordinates for a single timestep. connectivity : ndarray, shape (n_face, n_node_per_face) Node indices for each face bottom_indices : ndarray, shape (n_element,) @@ -94,28 +125,24 @@ def _calculate_prism_volumes(node_x, node_y, z, connectivity, bottom_indices): Returns ------- - ndarray, shape (n_time, n_layer-1, n_face) - Volume of each prism layer under the selection of nodes. + ndarray, shape (n_layer-1, n_face). + Prism volumes per layer and face. """ # Enforce float64 for numba node_x = node_x.astype(np.float64) node_y = node_y.astype(np.float64) z = z.astype(np.float64) - # Extract dimensions - n_time, n_node, n_layer = z.shape - n_face = connectivity.shape[0] - n_interfaces = n_layer - 1 + z = _coerce_single_timestep(z, "z") + # Calculate thickness of layers under our given nodes thickness = _calc_layers_thickness(z, bottom_indices) + h = thickness.T - # Integrate thickness over elements per timestep - result = np.zeros((n_time, n_interfaces, n_face)) - for t_i in range(n_time): - h = thickness[t_i].T - result[t_i] = _integrate_nodal(node_x, node_y, h, connectivity) + # Integrate! + integrated = _integrate_nodal(node_x, node_y, h, connectivity) - return result + return integrated def calc_volumes_from_projected(sxda): @@ -126,26 +153,25 @@ def calc_volumes_from_projected(sxda): Parameters ---------- sxda : SxDataArray - SxDataArray with 3D z-coordinates and bottom indices + SxDataArray with z-coordinates and bottom indices. + This function enforces single-timestep behavior. Returns ------- - ndarray, shape (n_time, n_layer-1, n_face) + ndarray, shape (n_layer-1, n_face) Volume of each prism layer under the selection of nodes. """ grid = sxda.sxgrid + z = _coerce_single_timestep(grid.sgrid_info.zCoordinates.values, "z") return _calculate_prism_volumes( node_x=grid.node_x.values, node_y=grid.node_y.values, - z=grid.sgrid_info.zCoordinates.values, + z=z, connectivity=grid.face_node_connectivity.values, bottom_indices=grid.sgrid_info.bottom_index_node.values - 1, ) -# TODO: Remove time indexing loops in function and have them only calculate at -# one timestep - def _calculate_volumes_uxarray(sxda, radius_m=None): """Calculate per-face prism volumes using uxarray spherical face areas. @@ -155,7 +181,8 @@ def _calculate_volumes_uxarray(sxda, radius_m=None): Parameters ---------- sxda : SxDataArray - SxDataArray with 3D z-coordinates and bottom indices. + SxDataArray with z-coordinates and bottom indices. + This function enforces single-timestep behavior. radius_m : float, optional Physical sphere radius in meters used to scale uxarray face areas. If not provided, computes a local geocentric Earth radius from the @@ -163,19 +190,21 @@ def _calculate_volumes_uxarray(sxda, radius_m=None): Returns ------- - ndarray, shape (n_time, n_layer-1, n_face) + ndarray, shape (n_layer-1, n_face) Volume of each prism layer under each face element. """ grid = sxda.sxgrid - z = grid.sgrid_info.zCoordinates.values.astype(np.float64) + z = _coerce_single_timestep( + grid.sgrid_info.zCoordinates.values.astype(np.float64), "z" + ) connectivity = grid.face_node_connectivity.values - # Thickness at nodes: shape (n_time, n_node, n_interface) + # Thickness at nodes: shape (n_node, n_interface) bottom_indices = grid.sgrid_info.bottom_index_node.values - 1 thickness = _calc_layers_thickness(z, bottom_indices) - # Face-average nodal thickness: shape (n_time, n_interface, n_face) + # Face-average nodal thickness: shape (n_interface, n_face) # NOTE: This is an arithmetic node mean. For strict parity with the planar # path on quads, we may need to mirror its bilinear quadrature weighting. thickness_face = _face_mean_from_nodal(thickness, connectivity) @@ -210,9 +239,9 @@ def _calculate_volumes_uxarray(sxda, radius_m=None): radius_m = _earth_radius_at_latitude(mean_lat) else: radius_m = float(radius_m) - face_areas_m2 = np.asarray(face_areas_unit, dtype=np.float64) * (radius_m ** 2) + face_areas_m2 = np.asarray(face_areas_unit, dtype=np.float64) * radius_m**2 - return thickness_face * face_areas_m2[np.newaxis, np.newaxis, :] + return thickness_face * face_areas_m2[np.newaxis, :] # TODO: implement generic prism integration of variables (such as energy, @@ -231,58 +260,60 @@ def _calculate_mass_in_ocean(sxda_salinity, sxda_temperature=None): Parameters ---------- sxda_salinity : SxDataArray - Practical salinity (PSU), shape (n_time, n_node, n_layer) + Practical salinity (PSU), shape (n_node, n_layer) + or (1, n_node, n_layer). sxda_temperature : SxDataArray, optional - In-situ temperature (°C), shape (n_time, n_node, n_layer). + In-situ temperature (°C), shape (n_node, n_layer) + or (1, n_node, n_layer). Defaults to 15.0 °C if not provided. Returns ------- - ndarray, shape (n_time, n_layer-1, n_face) + ndarray, shape (n_layer-1, n_face) Water mass per prism layer per face element (kg). """ import gsw grid = sxda_salinity.sxgrid - # Pressure from z-coordinates: shape (n_time, n_node, n_layer) - z = grid.sgrid_info.zCoordinates.values - # Expand to (1, n_node, 1) to broadcast with z shape (n_time,n_node,n_layer) - lat = grid.node_lat.values[np.newaxis, :, np.newaxis] - lon = grid.node_lon.values[np.newaxis, :, np.newaxis] - pressure = gsw.p_from_z(z, lat) # shape (n_time, n_node, n_layer) + # Pressure from z-coordinates: shape (n_node, n_layer) + z = _coerce_single_timestep(grid.sgrid_info.zCoordinates.values, "z") + lat = grid.node_lat.values[:, np.newaxis] + lon = grid.node_lon.values[:, np.newaxis] + pressure = gsw.p_from_z(z, lat) # shape (n_node, n_layer) # Convert salinity to Absolute Salinity (g/kg) - SP = sxda_salinity.values - SA = gsw.SA_from_SP(SP, pressure, lon, lat) # shape (n_time,n_node,n_layer) + SP = _coerce_single_timestep(sxda_salinity.values, "salinity") + SA = gsw.SA_from_SP(SP, pressure, lon, lat) # shape (n_node, n_layer) # Convert temperature to Conservative Temperature - t = sxda_temperature.values if sxda_temperature is not None else 15.0 - CT = gsw.CT_from_t(SA, t, pressure) # shape (n_time, n_node, n_layer) + if sxda_temperature is not None: + t = _coerce_single_timestep(sxda_temperature.values, "temperature") + else: + t = 15.0 + CT = gsw.CT_from_t(SA, t, pressure) # shape (n_node, n_layer) - # density: shape (n_time, n_node, n_layer) + # density: shape (n_node, n_layer) rho = gsw.density.rho(SA, CT, pressure) - # Average density to layer intervals: shape (n_time, n_node, n_layer-1) - rho_layer = 0.5 * (rho[:, :, :-1] + rho[:, :, 1:]) + # Average density to layer intervals: shape (n_node, n_layer-1) + rho_layer = 0.5 * (rho[:, :-1] + rho[:, 1:]) - # Thickness per layer interval: shape (n_time, n_node, n_layer-1) + # Thickness per layer interval: shape (n_node, n_layer-1) bottom_indices = grid.sgrid_info.bottom_index_node.values - 1 + thickness = _calc_layers_thickness(z, bottom_indices) - # rho * h at each node per layer: shape (n_time, n_node, n_layer-1) + # rho * h at each node per layer: shape (n_node, n_layer-1) rho_h = rho_layer * thickness - - # Integrate rho*h over faces: shape (n_time, n_layer-1, n_face) connectivity = grid.face_node_connectivity.values node_x = grid.node_x.values.astype(np.float64) node_y = grid.node_y.values.astype(np.float64) - n_time, n_node, n_interfaces = rho_h.shape - n_face = connectivity.shape[0] - mass = np.zeros((n_time, n_interfaces, n_face)) - for t_i in range(n_time): - mass[t_i] = _integrate_nodal( - node_x, node_y, rho_h[t_i].T.astype(np.float64), connectivity - ) - return mass \ No newline at end of file + values = rho_h.T # (n_layer-1, n_node) + + integrated = _integrate_nodal( + node_x, node_y, values.astype(np.float64), connectivity + ) + + return integrated \ No newline at end of file From e3c5c57d364af852c3012ac4c46e20124f8fa698 Mon Sep 17 00:00:00 2001 From: Alexandre Georges Date: Wed, 5 Aug 2026 13:09:08 -0700 Subject: [PATCH 2/5] refactor: changed test_colume to single timestep to reflect atomic integrate functions (enforcing single timestep) --- tests/test_volume.py | 42 +++++++++++++----------------------------- 1 file changed, 13 insertions(+), 29 deletions(-) diff --git a/tests/test_volume.py b/tests/test_volume.py index 1b2c742..0792257 100644 --- a/tests/test_volume.py +++ b/tests/test_volume.py @@ -25,23 +25,14 @@ def test_calculate_prism_volume_synthetic_shaved_cells(): # 0-based node bottom indices. k > 0 masks interfaces [:k] at that node. bottom_indices = np.array([0, 1, 2, 0, 1], dtype=np.int64) - # z shape: (n_time=2, n_node=5, n_layer=4) -> 3 interfaces. + # z shape: (n_node=5, n_layer=4) -> 3 interfaces. z = np.array( [ - [ - [-3.0, -2.0, -1.0, 0.0], - [-3.0, -2.0, -1.0, 0.0], - [-3.0, -2.0, -1.0, 0.0], - [-3.0, -2.0, -1.0, 0.0], - [-3.0, -2.0, -1.0, 0.0], - ], - [ - [-6.0, -4.0, -2.0, 0.0], - [-6.0, -4.0, -2.0, 0.0], - [-6.0, -4.0, -2.0, 0.0], - [-6.0, -4.0, -2.0, 0.0], - [-6.0, -4.0, -2.0, 0.0], - ], + [-3.0, -2.0, -1.0, 0.0], + [-3.0, -2.0, -1.0, 0.0], + [-3.0, -2.0, -1.0, 0.0], + [-3.0, -2.0, -1.0, 0.0], + [-3.0, -2.0, -1.0, 0.0], ], dtype=np.float64, ) @@ -54,24 +45,17 @@ def test_calculate_prism_volume_synthetic_shaved_cells(): bottom_indices=bottom_indices, ) - # Expected shape: (n_time, n_interface, n_face) - assert volumes.shape == (2, 3, 3) + # Expected shape: (n_interface, n_face) + assert volumes.shape == (3, 3) # Make the geometric factor explicit for readers: # expected volume = triangle_area * face-mean(thickness after masking) triangle_area_m2 = 2.0 expected_face_mean_thickness = np.array( [ - [ - [1.0 / 3.0, 2.0 / 3.0, 2.0 / 3.0], - [2.0 / 3.0, 2.0 / 3.0, 1.0], - [1.0, 1.0, 1.0], - ], - [ - [2.0 / 3.0, 4.0 / 3.0, 4.0 / 3.0], - [4.0 / 3.0, 4.0 / 3.0, 2.0], - [2.0, 2.0, 2.0], - ], + [1.0 / 3.0, 2.0 / 3.0, 2.0 / 3.0], + [2.0 / 3.0, 2.0 / 3.0, 1.0], + [1.0, 1.0, 1.0], ], dtype=np.float64, ) @@ -80,5 +64,5 @@ def test_calculate_prism_volume_synthetic_shaved_cells(): np.testing.assert_allclose(volumes, expected, rtol=0.0, atol=1e-12) # Spot-check a couple of entries with pytest.approx for readability. - assert volumes[0, 0, 0] == pytest.approx(2.0 / 3.0) - assert volumes[1, 2, 2] == pytest.approx(4.0) \ No newline at end of file + assert volumes[0, 0] == pytest.approx(2.0 / 3.0) + assert volumes[2, 2] == pytest.approx(2.0) \ No newline at end of file From 09963e0c249305a63f69d182e077486b2e1b7d6c Mon Sep 17 00:00:00 2001 From: Alexandre Georges Date: Wed, 5 Aug 2026 14:03:09 -0700 Subject: [PATCH 3/5] feat: add projected_volume_ufunc_kernel for dask parallelization support --- suxarray/grid/integrate.py | 54 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 54 insertions(+) diff --git a/suxarray/grid/integrate.py b/suxarray/grid/integrate.py index fc1e7b2..8ed4b06 100644 --- a/suxarray/grid/integrate.py +++ b/suxarray/grid/integrate.py @@ -145,6 +145,46 @@ def _calculate_prism_volumes(node_x, node_y, z, connectivity, bottom_indices): return integrated +def projected_volume_ufunc_kernel( + z, + node_x, + node_y, + connectivity, + bottom_indices, +): + """Array-kernel companion for ``calc_volumes_from_projected``. + + This helper is designed for external orchestration with + ``xarray.apply_ufunc(..., dask='parallelized')`` where ``z`` is passed as + an ndarray block/chunk. + + Parameters + ---------- + z : ndarray, shape (n_node, n_layer) + Single-timestep z-coordinates. + node_x : ndarray, shape (n_node,) + x coordinates of nodes. + node_y : ndarray, shape (n_node,) + y coordinates of nodes. + connectivity : ndarray, shape (n_face, n_node_per_face) + Node indices for each face. + bottom_indices : ndarray, shape (n_node,) + 0-based bottom layer index for each node. + + Returns + ------- + ndarray, shape (n_layer-1, n_face) + Prism volumes per layer and face for one timestep. + """ + return _calculate_prism_volumes( + node_x=node_x, + node_y=node_y, + z=z, + connectivity=connectivity, + bottom_indices=bottom_indices, + ) + + def calc_volumes_from_projected(sxda): """ Calculate the volume of prism under a selection of nodes for each layer in @@ -160,7 +200,21 @@ def calc_volumes_from_projected(sxda): ------- ndarray, shape (n_layer-1, n_face) Volume of each prism layer under the selection of nodes. + + Notes + ----- + For external dask/xarray parallelization over time, use + ``calc_volumes_from_projected_kernel`` with ``xarray.apply_ufunc``. """ + if not hasattr(sxda, "sxgrid"): + raise TypeError( + "calc_volumes_from_projected expects an SxDataArray (with .sxgrid)." + "If you are using xarray.apply_ufunc(..., dask='parallelized'), " + "use calc_volumes_from_projected_kernel on z-coordinate array " + "blocks and pass grid arrays (node_x, node_y, connectivity, " + "bottom_indices) as static arguments." + ) + grid = sxda.sxgrid z = _coerce_single_timestep(grid.sgrid_info.zCoordinates.values, "z") From 2910a315261979e379d42ef559117dde911086d9 Mon Sep 17 00:00:00 2001 From: Alexandre Georges Date: Thu, 6 Aug 2026 17:21:00 -0700 Subject: [PATCH 4/5] feat: add integrate_prism_variable function for generic nodal data integration over water column --- suxarray/grid/integrate.py | 58 ++++++++++++++++++++++++++++++++++++-- 1 file changed, 55 insertions(+), 3 deletions(-) diff --git a/suxarray/grid/integrate.py b/suxarray/grid/integrate.py index 8ed4b06..eaeec51 100644 --- a/suxarray/grid/integrate.py +++ b/suxarray/grid/integrate.py @@ -185,6 +185,60 @@ def projected_volume_ufunc_kernel( ) +def integrate_prism_variable(sxda, return_per_layer=False): + """ + Integrate a variable under an element over the water column for nodal data. + + Parameters + ---------- + sxda : SxDataArray + SxDataArray with z-coordinates and bottom indices containing the + variable to integrate. Variable must be a concentration or density-like + quantity (e.g., chlorophyll, salt, energy) that can be integrated over + the water column to produce a total quantity per face element. + This function enforces single-timestep behavior. + return_per_layer : bool, optional + If True, return per-layer integrated values with shape + (n_layer-1, n_face). If False (default), return total integrated + values per face with shape (n_face). + + Returns + ------- + ndarray + If ``return_per_layer=False``: shape (n_face), integrated variable + per face element summed over all layers. + If ``return_per_layer=True``: shape (n_layer-1, n_face), integrated + variable per layer and face. + """ + + grid = sxda.sxgrid + + # Extract grid properties + node_x = grid.node_x.values.astype(np.float64) + node_y = grid.node_y.values.astype(np.float64) + connectivity = grid.face_node_connectivity.values + z = _coerce_single_timestep(grid.sgrid_info.zCoordinates.values, "z") + bottom_indices = grid.sgrid_info.bottom_index_node.values - 1 + thickness = _calc_layers_thickness(z, bottom_indices) + + # Pull variable to integrate from sxda + var_arr = _coerce_single_timestep(sxda.values, "variable") + + # Average variable to layer intervals: shape (n_node, n_layer-1) + var_layer = 0.5 * (var_arr[:, :-1] + var_arr[:, 1:]) + values = (var_layer * thickness).T + + # Integrate! + integrated = _integrate_nodal( + node_x, node_y, values.astype(np.float64), connectivity + ) + + if return_per_layer: + return integrated + + return integrated.sum(axis=0) + + def calc_volumes_from_projected(sxda): """ Calculate the volume of prism under a selection of nodes for each layer in @@ -226,6 +280,7 @@ def calc_volumes_from_projected(sxda): bottom_indices=grid.sgrid_info.bottom_index_node.values - 1, ) + def _calculate_volumes_uxarray(sxda, radius_m=None): """Calculate per-face prism volumes using uxarray spherical face areas. @@ -298,9 +353,6 @@ def _calculate_volumes_uxarray(sxda, radius_m=None): return thickness_face * face_areas_m2[np.newaxis, :] -# TODO: implement generic prism integration of variables (such as energy, -# chlorophyll or salt) to complement mass calculation function - # TODO: Test impact of gsw method on performance. def _calculate_mass_in_ocean(sxda_salinity, sxda_temperature=None): From 656f93df4f26a637349abb0036045c50fa7db210 Mon Sep 17 00:00:00 2001 From: Alexandre Georges Date: Thu, 6 Aug 2026 17:21:18 -0700 Subject: [PATCH 5/5] minor formatting fix --- tests/test_suxarray.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/tests/test_suxarray.py b/tests/test_suxarray.py index 7fc72cf..d7f266a 100644 --- a/tests/test_suxarray.py +++ b/tests/test_suxarray.py @@ -95,7 +95,8 @@ def test_find_nearest_node(grid_test): def test_subset_bounding_polygon(sxds_test_dask): """Test subsetting dataarray with a polygon""" - # polygon = Polygon([(0.0, 0.0), (1000.0, 0.0), (1000.0, 10400.0), (0.0, 10400.0)]) + # polygon = Polygon([(0.0, 0.0), (1000.0, 0.0), + # (1000.0, 10400.0), (0.0, 10400.0)]) polygon = Polygon( [ (569866.8953279585, 4213065.012236144), @@ -126,7 +127,8 @@ def test_subset_elements_bounding_polygon(sxds_element_test): def test_subset_bounding_box(sxds_test_dask): """Test find_element_at""" da_subset = sxds_test_dask["salinity"].subset.bounding_box_xy( - [569866.8953279585, 570131.3752643355], [4213065.012236144, 4213262.690668024] + [569866.8953279585, 570131.3752643355], + [4213065.012236144, 4213262.690668024] ) assert da_subset.uxgrid.n_node == 15 assert da_subset.uxgrid.n_face == 10