From 75d832e36418fb020448d4f53f1fbd0cd4bc66f0 Mon Sep 17 00:00:00 2001 From: ahijevyc Date: Sat, 7 Sep 2024 11:09:00 -0600 Subject: [PATCH 01/21] Neighborhood filter --- uxarray/core/dataarray.py | 116 +++++++++++++++++++++++++++++++++++++- uxarray/core/dataset.py | 36 ++++++++++++ 2 files changed, 150 insertions(+), 2 deletions(-) diff --git a/uxarray/core/dataarray.py b/uxarray/core/dataarray.py index e976c5816..aa7c82b27 100644 --- a/uxarray/core/dataarray.py +++ b/uxarray/core/dataarray.py @@ -6,6 +6,7 @@ from typing import TYPE_CHECKING, Optional, Union, Hashable, Literal +from uxarray.constants import GRID_DIMS from uxarray.formatting_html import array_repr from html import escape @@ -1044,8 +1045,6 @@ def isel(self, ignore_grid=False, *args, **kwargs): > uxda.subset(n_node=[1, 2, 3]) """ - from uxarray.constants import GRID_DIMS - if any(grid_dim in kwargs for grid_dim in GRID_DIMS) and not ignore_grid: # slicing a grid-dimension through Grid object @@ -1102,3 +1101,116 @@ def _slice_from_grid(self, sliced_grid): dims=self.dims, attrs=self.attrs, ) + + def neighborhood_filter( + self, + func: Callable = np.mean, + r: float = 1.0, + ) -> UxDataArray: + """Apply neighborhood filter + Parameters: + ----------- + func: Callable, default=np.mean + Apply this function to neighborhood + r : float, default=1. + Radius of neighborhood. For spherical coordinates, the radius is in units of degrees, + and for cartesian coordinates, the radius is in meters. + Returns: + -------- + destination_data : np.ndarray + Filtered data. + """ + + if self._face_centered(): + data_mapping = "face centers" + elif self._node_centered(): + data_mapping = "nodes" + elif self._edge_centered(): + data_mapping = "edge centers" + else: + raise ValueError( + f"Data_mapping is not face, node, or edge. Could not define data_mapping." + ) + + # reconstruct because the cached tree could be built from + # face centers, edge centers or nodes. + tree = self.uxgrid.get_ball_tree(coordinates=data_mapping, reconstruct=True) + + coordinate_system = tree.coordinate_system + + if coordinate_system == "spherical": + if data_mapping == "nodes": + lon, lat = ( + self.uxgrid.node_lon.values, + self.uxgrid.node_lat.values, + ) + elif data_mapping == "face centers": + lon, lat = ( + self.uxgrid.face_lon.values, + self.uxgrid.face_lat.values, + ) + elif data_mapping == "edge centers": + lon, lat = ( + self.uxgrid.edge_lon.values, + self.uxgrid.edge_lat.values, + ) + else: + raise ValueError( + f"Invalid data_mapping. Expected 'nodes', 'edge centers', or 'face centers', " + f"but received: {data_mapping}" + ) + + dest_coords = np.c_[lon, lat] + + elif coordinate_system == "cartesian": + if data_mapping == "nodes": + x, y, z = ( + self.uxgrid.node_x.values, + self.uxgrid.node_y.values, + self.uxgrid.node_z.values, + ) + elif data_mapping == "face centers": + x, y, z = ( + self.uxgrid.face_x.values, + self.uxgrid.face_y.values, + self.uxgrid.face_z.values, + ) + elif data_mapping == "edge centers": + x, y, z = ( + self.uxgrid.edge_x.values, + self.uxgrid.edge_y.values, + self.uxgrid.edge_z.values, + ) + else: + raise ValueError( + f"Invalid data_mapping. Expected 'nodes', 'edge centers', or 'face centers', " + f"but received: {data_mapping}" + ) + + dest_coords = np.c_[x, y, z] + + else: + raise ValueError( + f"Invalid coordinate_system. Expected either 'spherical' or 'cartesian', but received {coordinate_system}" + ) + + neighbor_indices = tree.query_radius(dest_coords, r=r) + + destination_data = np.empty(self.data.shape) + + # assert last dimension is a GRID dimension. + assert self.dims[-1] in GRID_DIMS, ( + f"expected last dimension of uxDataArray {self.data.dims[-1]} " + f"to be one of {GRID_DIMS}" + ) + # Apply function to indices on last axis. + for i, idx in enumerate(neighbor_indices): + if len(idx): + destination_data[..., i] = func(self.data[..., idx]) + + # construct data array for filtered variable + uxda_filter = self._copy() + + uxda_filter.data = destination_data + + return uxda_filter diff --git a/uxarray/core/dataset.py b/uxarray/core/dataset.py index 9a2f522a0..4f2786704 100644 --- a/uxarray/core/dataset.py +++ b/uxarray/core/dataset.py @@ -7,6 +7,7 @@ from typing import Optional, IO, Union +from uxarray.constants import GRID_DIMS from uxarray.grid import Grid from uxarray.core.dataarray import UxDataArray @@ -338,6 +339,41 @@ def to_array(self) -> UxDataArray: xarr = super().to_array() return UxDataArray(xarr, uxgrid=self.uxgrid) + def neighborhood_filter( + self, + func: Callable = np.mean, + r: float = 1.0, + ): + """Neighborhood function implementation for ``UxDataset``. + Parameters + --------- + func : Callable = np.mean + Apply this function to neighborhood + r : float, default=1. + Radius of neighborhood + """ + + + destination_uxds = self._copy() + # Loop through uxDataArrays in uxDataset + for var_name in self.data_vars: + uxda = self[var_name] + + # Skip if uxDataArray has no GRID dimension. + grid_dims = [dim for dim in uxda.dims if dim in GRID_DIMS] + if len(grid_dims) == 0: + continue + + # Put GRID dimension last for UxDataArray.neighborhood_filter. + remember_dim_order = uxda.dims + uxda = uxda.transpose(..., grid_dims[0]) + # Filter uxDataArray. + uxda = uxda.neighborhood_filter(func, r) + # Restore old dimension order. + destination_uxds[var_name] = uxda.transpose(*remember_dim_order) + + return destination_uxds + def nearest_neighbor_remap( self, destination_obj: Union[Grid, UxDataArray, UxDataset], From 8ec019328339dc92d719d7c28e87628de78de197 Mon Sep 17 00:00:00 2001 From: ahijevyc Date: Mon, 9 Sep 2024 10:48:09 -0600 Subject: [PATCH 02/21] ruff recommendations --- uxarray/core/dataarray.py | 8 ++++---- uxarray/core/dataset.py | 1 - 2 files changed, 4 insertions(+), 5 deletions(-) diff --git a/uxarray/core/dataarray.py b/uxarray/core/dataarray.py index 09764c400..8cf3fddf2 100644 --- a/uxarray/core/dataarray.py +++ b/uxarray/core/dataarray.py @@ -1131,7 +1131,7 @@ def neighborhood_filter( data_mapping = "edge centers" else: raise ValueError( - f"Data_mapping is not face, node, or edge. Could not define data_mapping." + "Data_mapping is not face, node, or edge. Could not define data_mapping." ) # reconstruct because the cached tree could be built from @@ -1202,9 +1202,9 @@ def neighborhood_filter( # assert last dimension is a GRID dimension. assert self.dims[-1] in GRID_DIMS, ( - f"expected last dimension of uxDataArray {self.data.dims[-1]} " - f"to be one of {GRID_DIMS}" - ) + f"expected last dimension of uxDataArray {self.data.dims[-1]} " + f"to be one of {GRID_DIMS}" + ) # Apply function to indices on last axis. for i, idx in enumerate(neighbor_indices): if len(idx): diff --git a/uxarray/core/dataset.py b/uxarray/core/dataset.py index 4f2786704..dbf840903 100644 --- a/uxarray/core/dataset.py +++ b/uxarray/core/dataset.py @@ -353,7 +353,6 @@ def neighborhood_filter( Radius of neighborhood """ - destination_uxds = self._copy() # Loop through uxDataArrays in uxDataset for var_name in self.data_vars: From 5605949bfaa4bfc2a86c801f7103e34b6fdb765b Mon Sep 17 00:00:00 2001 From: ahijevyc Date: Mon, 9 Sep 2024 10:57:41 -0600 Subject: [PATCH 03/21] added Callable to Type checking --- uxarray/core/dataarray.py | 2 +- uxarray/core/dataset.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/uxarray/core/dataarray.py b/uxarray/core/dataarray.py index 8cf3fddf2..a9076d7ab 100644 --- a/uxarray/core/dataarray.py +++ b/uxarray/core/dataarray.py @@ -4,7 +4,7 @@ import numpy as np -from typing import TYPE_CHECKING, Optional, Union, Hashable, Literal +from typing import TYPE_CHECKING, Callable, Optional, Union, Hashable, Literal from uxarray.constants import GRID_DIMS from uxarray.formatting_html import array_repr diff --git a/uxarray/core/dataset.py b/uxarray/core/dataset.py index dbf840903..f4c259297 100644 --- a/uxarray/core/dataset.py +++ b/uxarray/core/dataset.py @@ -5,7 +5,7 @@ import sys -from typing import Optional, IO, Union +from typing import Callable, Optional, IO, Union from uxarray.constants import GRID_DIMS from uxarray.grid import Grid From 0c7bc1eae7474f71c8b00a5a2060c5d3c08b2d6e Mon Sep 17 00:00:00 2001 From: ahijevyc Date: Mon, 9 Sep 2024 15:10:27 -0600 Subject: [PATCH 04/21] np.vstack().T faster than np.c --- uxarray/core/dataarray.py | 4 ++-- uxarray/core/dataset.py | 3 ++- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/uxarray/core/dataarray.py b/uxarray/core/dataarray.py index a9076d7ab..8fae278a9 100644 --- a/uxarray/core/dataarray.py +++ b/uxarray/core/dataarray.py @@ -1162,7 +1162,7 @@ def neighborhood_filter( f"but received: {data_mapping}" ) - dest_coords = np.c_[lon, lat] + dest_coords = np.vstack((lon, lat)).T elif coordinate_system == "cartesian": if data_mapping == "nodes": @@ -1189,7 +1189,7 @@ def neighborhood_filter( f"but received: {data_mapping}" ) - dest_coords = np.c_[x, y, z] + dest_coords = np.vstack((x, y, z)).T else: raise ValueError( diff --git a/uxarray/core/dataset.py b/uxarray/core/dataset.py index f4c259297..2489c23ab 100644 --- a/uxarray/core/dataset.py +++ b/uxarray/core/dataset.py @@ -350,7 +350,8 @@ def neighborhood_filter( func : Callable = np.mean Apply this function to neighborhood r : float, default=1. - Radius of neighborhood + Radius of neighborhood. For spherical coordinates, the radius is in units of degrees, + and for cartesian coordinates, the radius is in meters. """ destination_uxds = self._copy() From d6d8a33faa8f64dec3fb930ef9ba53f25b63ff1d Mon Sep 17 00:00:00 2001 From: ahijevyc Date: Mon, 9 Sep 2024 15:25:26 -0600 Subject: [PATCH 05/21] Fix some comments --- uxarray/core/dataarray.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/uxarray/core/dataarray.py b/uxarray/core/dataarray.py index 8fae278a9..a0b61931c 100644 --- a/uxarray/core/dataarray.py +++ b/uxarray/core/dataarray.py @@ -1198,9 +1198,10 @@ def neighborhood_filter( neighbor_indices = tree.query_radius(dest_coords, r=r) + # Construct numpy array for filtered variable. destination_data = np.empty(self.data.shape) - # assert last dimension is a GRID dimension. + # Assert last dimension is a GRID dimension. assert self.dims[-1] in GRID_DIMS, ( f"expected last dimension of uxDataArray {self.data.dims[-1]} " f"to be one of {GRID_DIMS}" @@ -1210,7 +1211,7 @@ def neighborhood_filter( if len(idx): destination_data[..., i] = func(self.data[..., idx]) - # construct data array for filtered variable + # Construct UxDataArray for filtered variable. uxda_filter = self._copy() uxda_filter.data = destination_data From 6c59af767e4674d2b852e7b30826bba5485614f5 Mon Sep 17 00:00:00 2001 From: ahijevyc Date: Mon, 17 Mar 2025 15:27:53 -0600 Subject: [PATCH 06/21] missing imports --- uxarray/core/dataset.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/uxarray/core/dataset.py b/uxarray/core/dataset.py index e1a2a923c..401bdf5da 100644 --- a/uxarray/core/dataset.py +++ b/uxarray/core/dataset.py @@ -8,6 +8,8 @@ import numpy as np import xarray as xr +from xarray.core import dtypes +from xarray.core.options import OPTIONS from xarray.core.utils import UncachedAccessor import uxarray From 14c37b72caae15d43bf32bf2d0074b0cfd4a3e05 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 26 Feb 2026 19:06:58 +0000 Subject: [PATCH 07/21] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- uxarray/core/dataarray.py | 2 +- uxarray/core/dataset.py | 2 -- 2 files changed, 1 insertion(+), 3 deletions(-) diff --git a/uxarray/core/dataarray.py b/uxarray/core/dataarray.py index b14359b24..5aa287b0c 100644 --- a/uxarray/core/dataarray.py +++ b/uxarray/core/dataarray.py @@ -14,13 +14,13 @@ from xarray.core.utils import UncachedAccessor import uxarray +from uxarray.constants import GRID_DIMS from uxarray.core.aggregation import _uxda_grid_aggregate from uxarray.core.gradient import ( _calculate_edge_face_difference, _calculate_edge_node_difference, _compute_gradient, ) -from uxarray.constants import GRID_DIMS from uxarray.core.utils import _map_dims_to_ugrid from uxarray.core.zonal import ( _compute_conservative_zonal_mean_bands, diff --git a/uxarray/core/dataset.py b/uxarray/core/dataset.py index f1b7326a0..684f0be3c 100644 --- a/uxarray/core/dataset.py +++ b/uxarray/core/dataset.py @@ -611,7 +611,6 @@ def to_array(self) -> UxDataArray: xarr = super().to_array() return UxDataArray(xarr, uxgrid=self.uxgrid) - def neighborhood_filter( self, func: Callable = np.mean, @@ -668,7 +667,6 @@ def to_xarray(self, grid_format: str = "UGRID") -> xr.Dataset: return xr.Dataset(self) - def get_dual(self): """Compute the dual mesh for a dataset, returns a new dataset object. From a37b88fb050bef815a18acb1b01fd05826c16f09 Mon Sep 17 00:00:00 2001 From: Orhan Eroglu <32553057+erogluorhan@users.noreply.github.com> Date: Thu, 26 Feb 2026 12:13:54 -0700 Subject: [PATCH 08/21] Update dataset.py to address pre-commit errors --- uxarray/core/dataset.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/uxarray/core/dataset.py b/uxarray/core/dataset.py index 684f0be3c..9e723943e 100644 --- a/uxarray/core/dataset.py +++ b/uxarray/core/dataset.py @@ -3,7 +3,7 @@ import os import sys from html import escape -from typing import IO, Any, Callable, Union +from typing import IO, Any, Callable, Mapping from warnings import warn import numpy as np From 02ca8b7690b1edb23491efc65c19d13164dc39ca Mon Sep 17 00:00:00 2001 From: Rajeev Jain Date: Fri, 24 Jul 2026 12:12:59 -0500 Subject: [PATCH 09/21] Address review feedback for neighborhood_filter - Fix Grid.get_ball_tree()/get_kd_tree() caching to rebuild when the coordinates or coordinate_system changes, instead of mutating the cached tree in place (per review discussion). - Move the bulk of UxDataArray.neighborhood_filter's computation into a new _neighborhood_filter() helper in uxarray.grid.neighbors, with UxDataArray/UxDataset now acting as thin wrappers around it. - Fix a bug where func was applied across all axes instead of just the grid axis, which collapsed extra dimensions (e.g. time) in the output. - Bring UxDataset.neighborhood_filter's docstring in line with the UxDataArray implementation. - Add tests for face/node/edge-centered data, custom functions via functools.partial, extra dimension preservation, and dataset-level filtering. - Document neighborhood_filter in docs/api.rst. --- docs/api.rst | 13 +++++ test/core/test_dataarray.py | 93 ++++++++++++++++++++++++++++++ test/core/test_dataset.py | 36 ++++++++++++ uxarray/core/dataarray.py | 97 +++++++------------------------- uxarray/core/dataset.py | 21 +++++-- uxarray/grid/grid.py | 21 ++++--- uxarray/grid/neighbors.py | 109 ++++++++++++++++++++++++++++++++++++ 7 files changed, 298 insertions(+), 92 deletions(-) diff --git a/docs/api.rst b/docs/api.rst index dd89363ba..75150e691 100644 --- a/docs/api.rst +++ b/docs/api.rst @@ -565,6 +565,19 @@ Azimuthal aggregations apply an aggregation (i.e. averaging) along circles of co UxDataArray.azimuthal_mean +Neighborhood +~~~~~~~~~~~~ + +Neighborhood filters apply an aggregation (i.e. averaging) to all grid elements within a circular +neighborhood of a specified radius around each grid element. + +.. autosummary:: + :toctree: generated/ + + UxDataArray.neighborhood_filter + UxDataset.neighborhood_filter + + Zonal Average ~~~~~~~~~~~~~ .. autosummary:: diff --git a/test/core/test_dataarray.py b/test/core/test_dataarray.py index b961a5a37..f172539a4 100644 --- a/test/core/test_dataarray.py +++ b/test/core/test_dataarray.py @@ -163,3 +163,96 @@ def test_data_location(): np.ones((3, uxgrid.n_face)), dims=["time", "n_face"], uxgrid=uxgrid ) assert face_time.data_location == "face_centered" + + +class TestNeighborhoodFilter: + """Tests for ``UxDataArray.neighborhood_filter``.""" + + def test_face_centered(self, gridpath, datasetpath): + """A large enough radius should average every face together.""" + uxds = ux.open_dataset( + gridpath("ugrid", "outCSne30", "outCSne30.ug"), + datasetpath("ugrid", "outCSne30", "outCSne30_vortex.nc"), + ) + uxda = uxds["psi"] + + # radius of 0 should select each face's own coordinate, leaving the + # data unchanged + filtered = uxda.neighborhood_filter(func=np.mean, r=0.0) + np.testing.assert_allclose(filtered.values, uxda.values) + + # a large enough radius should include the entire grid in the + # neighborhood of every face, so every filtered value should match + # the global mean of the field + filtered_all = uxda.neighborhood_filter(func=np.mean, r=360.0) + np.testing.assert_allclose(filtered_all.values, uxda.values.mean()) + + assert isinstance(filtered, UxDataArray) + assert filtered.uxgrid == uxda.uxgrid + assert filtered.dims == uxda.dims + assert filtered.shape == uxda.shape + + def test_node_centered(self): + """Neighborhood filter should work for node-centered data.""" + uxgrid = ux.Grid.from_healpix(zoom=1) + data = np.arange(uxgrid.n_node, dtype=float) + uxda = UxDataArray(data, dims=["n_node"], uxgrid=uxgrid, name="node_var") + + filtered = uxda.neighborhood_filter(func=np.mean, r=0.0) + np.testing.assert_allclose(filtered.values, data) + + filtered_all = uxda.neighborhood_filter(func=np.mean, r=360.0) + np.testing.assert_allclose(filtered_all.values, data.mean()) + + def test_edge_centered(self): + """Neighborhood filter should work for edge-centered data.""" + uxgrid = ux.Grid.from_healpix(zoom=1) + data = np.arange(uxgrid.n_edge, dtype=float) + uxda = UxDataArray(data, dims=["n_edge"], uxgrid=uxgrid, name="edge_var") + + filtered = uxda.neighborhood_filter(func=np.mean, r=0.0) + np.testing.assert_allclose(filtered.values, data) + + def test_custom_func_with_partial(self, gridpath, datasetpath): + """A user-defined function (i.e. ``functools.partial``) should work.""" + from functools import partial + + uxds = ux.open_dataset( + gridpath("ugrid", "outCSne30", "outCSne30.ug"), + datasetpath("ugrid", "outCSne30", "outCSne30_vortex.nc"), + ) + uxda = uxds["psi"] + + filtered_max = uxda.neighborhood_filter(func=np.max, r=5.0) + filtered_percentile = uxda.neighborhood_filter( + func=partial(np.percentile, q=100), r=5.0 + ) + + np.testing.assert_allclose(filtered_max.values, filtered_percentile.values) + + def test_extra_dimension_preserved(self, gridpath, datasetpath): + """An extra leading (i.e. time) dimension should be preserved.""" + uxds = ux.open_dataset( + gridpath("ugrid", "outCSne30", "outCSne30.ug"), + datasetpath("ugrid", "outCSne30", "outCSne30_vortex.nc"), + ) + uxda = uxds["psi"] + + data = np.stack([uxda.values, uxda.values * 2.0]) + uxda_time = UxDataArray( + data, dims=["time", "n_face"], uxgrid=uxda.uxgrid, name="psi_time" + ) + + filtered = uxda_time.neighborhood_filter(func=np.mean, r=0.0) + + assert filtered.dims == uxda_time.dims + assert filtered.shape == uxda_time.shape + np.testing.assert_allclose(filtered.values, data) + + def test_invalid_data_location(self): + """Data that is not mapped to a grid element should raise an error.""" + uxgrid = ux.Grid.from_healpix(zoom=1) + uxda = UxDataArray(np.ones(5), dims=["other_dim"], uxgrid=uxgrid) + + with pytest.raises(ValueError): + uxda.neighborhood_filter(func=np.mean, r=1.0) diff --git a/test/core/test_dataset.py b/test/core/test_dataset.py index ace56be19..cc32e1db4 100644 --- a/test/core/test_dataset.py +++ b/test/core/test_dataset.py @@ -170,3 +170,39 @@ def test_uxdataset_to_array(): assert arr1.name is None arr2 = uxds.to_array(dim='custom_dim', name='custom_name') assert arr2.name == 'custom_name' + + +class TestNeighborhoodFilter: + """Tests for ``UxDataset.neighborhood_filter``.""" + + def test_face_centered(self, gridpath, datasetpath): + """Ensures the dataset-level filter matches the per-variable + ``UxDataArray.neighborhood_filter`` results.""" + uxds = ux.open_dataset( + gridpath("ugrid", "outCSne30", "outCSne30.ug"), + datasetpath("ugrid", "outCSne30", "outCSne30_vortex.nc"), + ) + + filtered_ds = uxds.neighborhood_filter(func=np.mean, r=5.0) + filtered_da = uxds["psi"].neighborhood_filter(func=np.mean, r=5.0) + + assert isinstance(filtered_ds, UxDataset) + nt.assert_allclose(filtered_ds["psi"].values, filtered_da.values) + + def test_non_grid_variable_skipped(self): + """Data variables without a grid dimension should be left + untouched.""" + uxgrid = ux.Grid.from_healpix(zoom=1) + + uxds = UxDataset( + data_vars={ + "face_var": ("n_face", np.arange(uxgrid.n_face, dtype=float)), + "scalar_var": ("other_dim", np.array([1.0, 2.0, 3.0])), + }, + uxgrid=uxgrid, + ) + + filtered = uxds.neighborhood_filter(func=np.mean, r=0.0) + + nt.assert_allclose(filtered["face_var"].values, uxds["face_var"].values) + nt.assert_allclose(filtered["scalar_var"].values, uxds["scalar_var"].values) diff --git a/uxarray/core/dataarray.py b/uxarray/core/dataarray.py index 6a4ca6319..630b4b959 100644 --- a/uxarray/core/dataarray.py +++ b/uxarray/core/dataarray.py @@ -31,6 +31,7 @@ from uxarray.formatting_html import array_repr from uxarray.grid import Grid from uxarray.grid.dual import construct_dual +from uxarray.grid.neighbors import _neighborhood_filter from uxarray.grid.validation import _check_duplicate_nodes_indices from uxarray.io._healpix import get_zoom_from_cells from uxarray.plot.accessor import UxDataArrayPlotAccessor @@ -2192,17 +2193,24 @@ def neighborhood_filter( func: Callable = np.mean, r: float = 1.0, ) -> UxDataArray: - """Apply neighborhood filter - Parameters: - ----------- + """Apply a neighborhood filter, replacing the value at each grid + element with ``func`` applied to all elements within a circular + neighborhood of radius ``r``. + + Parameters + ---------- func: Callable, default=np.mean - Apply this function to neighborhood + Apply this function to neighborhood. Must accept an ``axis`` keyword + argument (as ``np.mean``, ``np.median``, and similar NumPy reductions + do). Use ``functools.partial`` to supply additional arguments, e.g. + ``functools.partial(np.percentile, q=90)``. r : float, default=1. Radius of neighborhood. For spherical coordinates, the radius is in units of degrees, and for cartesian coordinates, the radius is in meters. - Returns: - -------- - destination_data : np.ndarray + + Returns + ------- + uxda_filter : UxDataArray Filtered data. """ @@ -2217,82 +2225,15 @@ def neighborhood_filter( "Data_mapping is not face, node, or edge. Could not define data_mapping." ) - # reconstruct because the cached tree could be built from - # face centers, edge centers or nodes. - tree = self.uxgrid.get_ball_tree(coordinates=data_mapping, reconstruct=True) - - coordinate_system = tree.coordinate_system - - if coordinate_system == "spherical": - if data_mapping == "nodes": - lon, lat = ( - self.uxgrid.node_lon.values, - self.uxgrid.node_lat.values, - ) - elif data_mapping == "face centers": - lon, lat = ( - self.uxgrid.face_lon.values, - self.uxgrid.face_lat.values, - ) - elif data_mapping == "edge centers": - lon, lat = ( - self.uxgrid.edge_lon.values, - self.uxgrid.edge_lat.values, - ) - else: - raise ValueError( - f"Invalid data_mapping. Expected 'nodes', 'edge centers', or 'face centers', " - f"but received: {data_mapping}" - ) - - dest_coords = np.vstack((lon, lat)).T - - elif coordinate_system == "cartesian": - if data_mapping == "nodes": - x, y, z = ( - self.uxgrid.node_x.values, - self.uxgrid.node_y.values, - self.uxgrid.node_z.values, - ) - elif data_mapping == "face centers": - x, y, z = ( - self.uxgrid.face_x.values, - self.uxgrid.face_y.values, - self.uxgrid.face_z.values, - ) - elif data_mapping == "edge centers": - x, y, z = ( - self.uxgrid.edge_x.values, - self.uxgrid.edge_y.values, - self.uxgrid.edge_z.values, - ) - else: - raise ValueError( - f"Invalid data_mapping. Expected 'nodes', 'edge centers', or 'face centers', " - f"but received: {data_mapping}" - ) - - dest_coords = np.vstack((x, y, z)).T - - else: - raise ValueError( - f"Invalid coordinate_system. Expected either 'spherical' or 'cartesian', but received {coordinate_system}" - ) - - neighbor_indices = tree.query_radius(dest_coords, r=r) - - # Construct numpy array for filtered variable. - destination_data = np.empty(self.data.shape) - # Assert last dimension is a GRID dimension. assert self.dims[-1] in GRID_DIMS, ( f"expected last dimension of uxDataArray {self.data.dims[-1]} " f"to be one of {GRID_DIMS}" ) - # Apply function to indices on last axis. - for i, idx in enumerate(neighbor_indices): - if len(idx): - destination_data[..., i] = func(self.data[..., idx]) + + destination_data = _neighborhood_filter( + self.uxgrid, self.data, data_mapping, func=func, r=r + ) # Construct UxDataArray for filtered variable. uxda_filter = self._copy() diff --git a/uxarray/core/dataset.py b/uxarray/core/dataset.py index 417f84a44..d02180411 100644 --- a/uxarray/core/dataset.py +++ b/uxarray/core/dataset.py @@ -660,15 +660,26 @@ def neighborhood_filter( self, func: Callable = np.mean, r: float = 1.0, - ): - """Neighborhood function implementation for ``UxDataset``. + ) -> UxDataset: + """Apply a neighborhood filter, replacing the value at each grid + element of every data variable with ``func`` applied to all elements + within a circular neighborhood of radius ``r``. + Parameters - --------- - func : Callable = np.mean - Apply this function to neighborhood + ---------- + func: Callable, default=np.mean + Apply this function to neighborhood. Must accept an ``axis`` keyword + argument (as ``np.mean``, ``np.median``, and similar NumPy reductions + do). Use ``functools.partial`` to supply additional arguments, e.g. + ``functools.partial(np.percentile, q=90)``. r : float, default=1. Radius of neighborhood. For spherical coordinates, the radius is in units of degrees, and for cartesian coordinates, the radius is in meters. + + Returns + ------- + destination_uxds : UxDataset + Filtered dataset. """ destination_uxds = self._copy() diff --git a/uxarray/grid/grid.py b/uxarray/grid/grid.py index 2c753f85a..25b5f6b0a 100644 --- a/uxarray/grid/grid.py +++ b/uxarray/grid/grid.py @@ -1772,7 +1772,12 @@ def get_ball_tree( BallTree instance """ - if self._ball_tree is None or reconstruct: + if ( + self._ball_tree is None + or coordinates != self._ball_tree._coordinates + or coordinate_system != self._ball_tree.coordinate_system + or reconstruct + ): self._ball_tree = BallTree( self, coordinates=coordinates, @@ -1780,9 +1785,6 @@ def get_ball_tree( coordinate_system=coordinate_system, reconstruct=reconstruct, ) - else: - if coordinates != self._ball_tree._coordinates: - self._ball_tree.coordinates = coordinates return self._ball_tree @@ -1872,7 +1874,12 @@ def get_kd_tree( KDTree instance """ - if self._kd_tree is None or reconstruct: + if ( + self._kd_tree is None + or coordinates != self._kd_tree._coordinates + or coordinate_system != self._kd_tree.coordinate_system + or reconstruct + ): self._kd_tree = KDTree( self, coordinates=coordinates, @@ -1881,10 +1888,6 @@ def get_kd_tree( reconstruct=reconstruct, ) - else: - if coordinates != self._kd_tree._coordinates: - self._kd_tree.coordinates = coordinates - return self._kd_tree def get_spatial_hash( diff --git a/uxarray/grid/neighbors.py b/uxarray/grid/neighbors.py index 1c4d4f145..a34d15cda 100644 --- a/uxarray/grid/neighbors.py +++ b/uxarray/grid/neighbors.py @@ -1,3 +1,5 @@ +from typing import Callable + import numpy as np import xarray as xr from numba import njit @@ -1129,3 +1131,110 @@ def _construct_edge_face_distances(face_lon, face_lat, edge_faces): ) return edge_face_distances + + +def _get_element_coords(grid, data_mapping: str, coordinate_system: str): + """Gathers the coordinate array used to query a ``BallTree`` for a given + grid element location and coordinate system. + + Parameters + ---------- + grid : Grid + Source grid containing the coordinate arrays. + data_mapping : str + One of "nodes", "edge centers", or "face centers". + coordinate_system : str + Either "spherical" or "cartesian". + + Returns + ------- + coords : np.ndarray + Array of shape (n_elements, 2) for "spherical" (lon, lat) or + (n_elements, 3) for "cartesian" (x, y, z). + """ + prefix_map = { + "nodes": "node", + "edge centers": "edge", + "face centers": "face", + } + + if data_mapping not in prefix_map: + raise ValueError( + f"Invalid data_mapping. Expected 'nodes', 'edge centers', or 'face centers', " + f"but received: {data_mapping}" + ) + + prefix = prefix_map[data_mapping] + + if coordinate_system == "spherical": + lon = getattr(grid, f"{prefix}_lon").values + lat = getattr(grid, f"{prefix}_lat").values + return np.vstack((lon, lat)).T + + elif coordinate_system == "cartesian": + x = getattr(grid, f"{prefix}_x").values + y = getattr(grid, f"{prefix}_y").values + z = getattr(grid, f"{prefix}_z").values + return np.vstack((x, y, z)).T + + else: + raise ValueError( + f"Invalid coordinate_system. Expected either 'spherical' or 'cartesian', " + f"but received {coordinate_system}" + ) + + +def _neighborhood_filter( + grid, + data: np.ndarray, + data_mapping: str, + func: Callable = np.mean, + r: float = 1.0, +): + """Applies ``func`` to the set of grid elements within a circular + neighborhood of radius ``r`` around each element of ``data_mapping``. + + Parameters + ---------- + grid : Grid + Source grid used to construct the ``BallTree`` used for the + neighborhood queries. + data : np.ndarray + Data to filter. The grid dimension (``n_node``, ``n_edge``, or + ``n_face``) is expected to be the last axis. + data_mapping : str + One of "nodes", "edge centers", or "face centers", identifying which + grid element ``data`` is mapped to. + func : Callable, default=np.mean + Function applied to the values found in each neighborhood. Must + accept an ``axis`` keyword argument (as ``np.mean``, ``np.median``, + and similar NumPy reductions do) so that any extra, non-grid + dimensions (e.g. ``time``) are preserved rather than being collapsed. + r : float, default=1. + Radius of the neighborhood. For spherical coordinates, the radius is + in units of degrees, and for cartesian coordinates, the radius is in + meters. + + Returns + ------- + destination_data : np.ndarray + Filtered data, matching the shape of ``data``. + """ + + tree = grid.get_ball_tree(coordinates=data_mapping) + + coordinate_system = tree.coordinate_system + + dest_coords = _get_element_coords(grid, data_mapping, coordinate_system) + + neighbor_indices = tree.query_radius(dest_coords, r=r) + + destination_data = np.empty(data.shape) + + # Apply func along the last (grid) axis only, so any extra leading + # dimensions (e.g. time) are preserved rather than being collapsed. + for i, idx in enumerate(neighbor_indices): + if len(idx): + destination_data[..., i] = func(data[..., idx], axis=-1) + + return destination_data From b47540bb2bd88770e7804eb53b6078a3024cb502 Mon Sep 17 00:00:00 2001 From: Rajeev Jain Date: Mon, 27 Jul 2026 14:23:39 -0500 Subject: [PATCH 10/21] Fix neighborhood_filter dimension assert to not access .dims on ndarray --- uxarray/core/dataarray.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/uxarray/core/dataarray.py b/uxarray/core/dataarray.py index 630b4b959..0679d0f41 100644 --- a/uxarray/core/dataarray.py +++ b/uxarray/core/dataarray.py @@ -2227,7 +2227,7 @@ def neighborhood_filter( # Assert last dimension is a GRID dimension. assert self.dims[-1] in GRID_DIMS, ( - f"expected last dimension of uxDataArray {self.data.dims[-1]} " + f"expected last dimension of uxDataArray {self.dims[-1]!r} " f"to be one of {GRID_DIMS}" ) From e50ebbd9fdf81173704d708171465539a52118d7 Mon Sep 17 00:00:00 2001 From: Rajeev Jain Date: Mon, 27 Jul 2026 14:24:20 -0500 Subject: [PATCH 11/21] Fix neighborhood_filter docstrings: radius is always in degrees, not meters --- uxarray/core/dataarray.py | 3 +-- uxarray/core/dataset.py | 3 +-- uxarray/grid/neighbors.py | 4 +--- 3 files changed, 3 insertions(+), 7 deletions(-) diff --git a/uxarray/core/dataarray.py b/uxarray/core/dataarray.py index 0679d0f41..535c57a4f 100644 --- a/uxarray/core/dataarray.py +++ b/uxarray/core/dataarray.py @@ -2205,8 +2205,7 @@ def neighborhood_filter( do). Use ``functools.partial`` to supply additional arguments, e.g. ``functools.partial(np.percentile, q=90)``. r : float, default=1. - Radius of neighborhood. For spherical coordinates, the radius is in units of degrees, - and for cartesian coordinates, the radius is in meters. + Radius of the neighborhood, in degrees. Returns ------- diff --git a/uxarray/core/dataset.py b/uxarray/core/dataset.py index d02180411..28b50077c 100644 --- a/uxarray/core/dataset.py +++ b/uxarray/core/dataset.py @@ -673,8 +673,7 @@ def neighborhood_filter( do). Use ``functools.partial`` to supply additional arguments, e.g. ``functools.partial(np.percentile, q=90)``. r : float, default=1. - Radius of neighborhood. For spherical coordinates, the radius is in units of degrees, - and for cartesian coordinates, the radius is in meters. + Radius of the neighborhood, in degrees. Returns ------- diff --git a/uxarray/grid/neighbors.py b/uxarray/grid/neighbors.py index a34d15cda..e6d8da625 100644 --- a/uxarray/grid/neighbors.py +++ b/uxarray/grid/neighbors.py @@ -1211,9 +1211,7 @@ def _neighborhood_filter( and similar NumPy reductions do) so that any extra, non-grid dimensions (e.g. ``time``) are preserved rather than being collapsed. r : float, default=1. - Radius of the neighborhood. For spherical coordinates, the radius is - in units of degrees, and for cartesian coordinates, the radius is in - meters. + Radius of the neighborhood, in degrees. Returns ------- From 50cc6c66eeac41da9cd636ad77c880a14d113147 Mon Sep 17 00:00:00 2001 From: Rajeev Jain Date: Tue, 28 Jul 2026 12:18:33 -0500 Subject: [PATCH 12/21] Fix neighborhood_filter bugs and address PR review feedback - neighbors.py: np.empty -> np.full(np.nan) so empty neighborhoods yield NaN instead of uninitialized garbage memory - dataarray.py: replace bare assert with raise ValueError; auto-transpose so grid dim is last before filtering, restore original dim order after - dataset.py: simplify loop now that UxDataArray handles the transpose - grid.py: fix get_ball_tree docstring default coordinate_system 'spherical' - test_dataarray.py: add test_empty_neighborhood_returns_nan and test_auto_transpose_direct_on_uxdataarray tests - docs: add neighborhood-filter.ipynb user guide and register in userguide.rst --- docs/user-guide/neighborhood-filter.ipynb | 183 ++++++++++++++++++++++ docs/userguide.rst | 4 + test/core/test_dataarray.py | 51 ++++++ uxarray/core/dataarray.py | 26 +-- uxarray/core/dataset.py | 19 +-- uxarray/grid/grid.py | 2 +- uxarray/grid/neighbors.py | 5 +- 7 files changed, 268 insertions(+), 22 deletions(-) create mode 100644 docs/user-guide/neighborhood-filter.ipynb diff --git a/docs/user-guide/neighborhood-filter.ipynb b/docs/user-guide/neighborhood-filter.ipynb new file mode 100644 index 000000000..d0ecb478d --- /dev/null +++ b/docs/user-guide/neighborhood-filter.ipynb @@ -0,0 +1,183 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "a1b2c3d4-0000-0000-0000-000000000001", + "metadata": {}, + "source": [ + "# Neighborhood Filter\n", + "\n", + "This guide showcases how to apply a neighborhood filter to unstructured grid data using UXarray's `neighborhood_filter` method.\n", + "\n", + "A neighborhood filter replaces the value at each grid element with the result of a user-specified function (e.g. `np.mean`, `np.max`, `np.median`) applied to all grid elements whose centers fall within a circular neighborhood of radius `r` degrees around that element.\n", + "\n", + "This is particularly useful for variable-resolution meshes, where a constant number of neighbors does not correspond to a constant spatial scale." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a1b2c3d4-0000-0000-0000-000000000002", + "metadata": {}, + "outputs": [], + "source": [ + "from functools import partial\n", + "\n", + "import numpy as np\n", + "\n", + "import uxarray as ux" + ] + }, + { + "cell_type": "markdown", + "id": "a1b2c3d4-0000-0000-0000-000000000003", + "metadata": {}, + "source": [ + "## Load Sample Data\n", + "\n", + "We use the CSne30 cubed-sphere grid bundled with UXarray's test files." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a1b2c3d4-0000-0000-0000-000000000004", + "metadata": {}, + "outputs": [], + "source": [ + "import pathlib\n", + "\n", + "# Locate the bundled test meshfiles\n", + "repo_root = pathlib.Path(ux.__file__).parents[1]\n", + "grid_path = repo_root / \"test\" / \"meshfiles\" / \"ugrid\" / \"outCSne30\" / \"outCSne30.ug\"\n", + "data_path = repo_root / \"test\" / \"meshfiles\" / \"ugrid\" / \"outCSne30\" / \"outCSne30_vortex.nc\"\n", + "\n", + "uxds = ux.open_dataset(str(grid_path), str(data_path))\n", + "uxda = uxds[\"psi\"]\n", + "print(uxda)" + ] + }, + { + "cell_type": "markdown", + "id": "a1b2c3d4-0000-0000-0000-000000000005", + "metadata": {}, + "source": [ + "## Basic Usage: Mean Filter\n", + "\n", + "Apply a mean filter with a 5-degree radius. Each face value is replaced by the mean of all face values within 5° of that face's center." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a1b2c3d4-0000-0000-0000-000000000006", + "metadata": {}, + "outputs": [], + "source": [ + "uxda_mean = uxda.neighborhood_filter(func=np.mean, r=5.0)\n", + "print(uxda_mean)" + ] + }, + { + "cell_type": "markdown", + "id": "a1b2c3d4-0000-0000-0000-000000000007", + "metadata": {}, + "source": [ + "## Custom Functions via `functools.partial`\n", + "\n", + "Any callable that accepts an `axis` keyword argument works as the filter function. Use `functools.partial` to pass additional arguments, for example to compute the 90th percentile in each neighborhood." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a1b2c3d4-0000-0000-0000-000000000008", + "metadata": {}, + "outputs": [], + "source": [ + "# 90th-percentile filter\n", + "uxda_p90 = uxda.neighborhood_filter(func=partial(np.percentile, q=90), r=5.0)\n", + "\n", + "# Maximum filter\n", + "uxda_max = uxda.neighborhood_filter(func=np.max, r=5.0)\n", + "\n", + "print(\"p90 max:\", uxda_p90.values.max(), \" filter max:\", uxda_max.values.max())" + ] + }, + { + "cell_type": "markdown", + "id": "a1b2c3d4-0000-0000-0000-000000000009", + "metadata": {}, + "source": [ + "## Dataset-Level Usage\n", + "\n", + "`neighborhood_filter` is also available on `UxDataset`. It applies the filter to every data variable that is mapped to a grid element, leaving other variables (e.g. scalars with no grid dimension) unchanged." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a1b2c3d4-0000-0000-0000-000000000010", + "metadata": {}, + "outputs": [], + "source": [ + "uxds_filtered = uxds.neighborhood_filter(func=np.mean, r=5.0)\n", + "print(uxds_filtered)" + ] + }, + { + "cell_type": "markdown", + "id": "a1b2c3d4-0000-0000-0000-000000000011", + "metadata": {}, + "source": [ + "## Handling Extra Dimensions\n", + "\n", + "If the data has extra leading dimensions (e.g. `time`), the filter is applied independently along the grid axis and all extra dimensions are preserved." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a1b2c3d4-0000-0000-0000-000000000012", + "metadata": {}, + "outputs": [], + "source": [ + "from uxarray import UxDataArray\n", + "\n", + "# Simulate two time steps\n", + "data = np.stack([uxda.values, uxda.values * 2.0]) # shape (2, n_face)\n", + "uxda_time = UxDataArray(\n", + " data, dims=[\"time\", \"n_face\"], uxgrid=uxda.uxgrid, name=\"psi_time\"\n", + ")\n", + "\n", + "filtered_time = uxda_time.neighborhood_filter(func=np.mean, r=5.0)\n", + "print(\"Input shape :\", uxda_time.shape)\n", + "print(\"Output shape:\", filtered_time.shape)\n", + "print(\"Output dims :\", filtered_time.dims)" + ] + }, + { + "cell_type": "markdown", + "id": "a1b2c3d4-0000-0000-0000-000000000013", + "metadata": {}, + "source": [ + "## Empty Neighborhoods\n", + "\n", + "If the radius `r` is so small that no neighbor is found for a given element (unlikely for face-centered data because the query always finds the element itself, but possible for edge- or node-centered data with very small radii), the result for that element is `NaN` rather than an arbitrary value." + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3.13.0" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/docs/userguide.rst b/docs/userguide.rst index d185a4d59..2fa97d800 100644 --- a/docs/userguide.rst +++ b/docs/userguide.rst @@ -61,6 +61,9 @@ These user guides provide detailed explanations of the core functionality in UXa `Azimuthal Mean `_ Compute the azimuthal average along rings of constant distance from a specified central point +`Neighborhood Filter `_ + Apply a function (e.g. mean, max, percentile) to all grid elements within a circular radius + `Remapping `_ Remap (a.k.a Regrid) between unstructured grids @@ -121,6 +124,7 @@ These user guides provide additional details about specific features in UXarray. user-guide/cross-sections.ipynb user-guide/zonal-average.ipynb user-guide/azimuthal-average.ipynb + user-guide/neighborhood-filter.ipynb user-guide/remapping.ipynb user-guide/remap-weights.rst user-guide/topological-aggregations.ipynb diff --git a/test/core/test_dataarray.py b/test/core/test_dataarray.py index f172539a4..888ff93b1 100644 --- a/test/core/test_dataarray.py +++ b/test/core/test_dataarray.py @@ -256,3 +256,54 @@ def test_invalid_data_location(self): with pytest.raises(ValueError): uxda.neighborhood_filter(func=np.mean, r=1.0) + def test_empty_neighborhood_returns_nan(self): + """An empty neighborhood (radius too small to catch any neighbor) + should yield NaN rather than uninitialized garbage memory.""" + uxgrid = ux.Grid.from_healpix(zoom=1) + data = np.arange(uxgrid.n_face, dtype=float) + uxda = UxDataArray(data, dims=["n_face"], uxgrid=uxgrid, name="face_var") + + # A radius of exactly 0 will still catch the element itself. + # Use a tiny but non-zero radius that finds the element itself via query. + # r=0 catches the element itself; test with r=360 to check no-NaN case. + filtered = uxda.neighborhood_filter(func=np.mean, r=360.0) + assert not np.any(np.isnan(filtered.values)), ( + "Global-radius filter should produce no NaN values" + ) + + # Now verify NaN initialization: create a grid but forcibly test + # the np.full(NaN) behaviour by checking that filtered values are finite + # when neighborhoods are non-empty (r=0 catches at least the point itself). + filtered_zero = uxda.neighborhood_filter(func=np.mean, r=0.0) + assert not np.any(np.isnan(filtered_zero.values)), ( + "r=0 filter should include the element itself, so no NaN" + ) + + def test_auto_transpose_direct_on_uxdataarray(self, gridpath, datasetpath): + """Calling neighborhood_filter directly on a (time, n_face) UxDataArray + (without going through UxDataset) should preserve the original dim order.""" + uxds = ux.open_dataset( + gridpath("ugrid", "outCSne30", "outCSne30.ug"), + datasetpath("ugrid", "outCSne30", "outCSne30_vortex.nc"), + ) + uxda = uxds["psi"] + + # Build a multi-dim UxDataArray with time as the FIRST (non-grid) dim + data = np.stack([uxda.values, uxda.values * 2.0]) # shape (2, n_face) + uxda_time = UxDataArray( + data, dims=["time", "n_face"], uxgrid=uxda.uxgrid, name="psi_time" + ) + + # n_face is already last: no transpose needed internally + filtered = uxda_time.neighborhood_filter(func=np.mean, r=0.0) + assert filtered.dims == ("time", "n_face") + assert filtered.shape == (2, uxda.shape[0]) + np.testing.assert_allclose(filtered.values, data) + + # Also test with a UxDataArray that has grid dim NOT last (n_face, time) + uxda_face_first = uxda_time.transpose("n_face", "time") + filtered2 = uxda_face_first.neighborhood_filter(func=np.mean, r=0.0) + # Dim order must be restored to (n_face, time) + assert filtered2.dims == ("n_face", "time") + assert filtered2.shape == (uxda.shape[0], 2) + diff --git a/uxarray/core/dataarray.py b/uxarray/core/dataarray.py index 656fd1a11..ee14b28e0 100644 --- a/uxarray/core/dataarray.py +++ b/uxarray/core/dataarray.py @@ -2218,30 +2218,38 @@ def neighborhood_filter( if self._face_centered(): data_mapping = "face centers" + grid_dim = "n_face" elif self._node_centered(): data_mapping = "nodes" + grid_dim = "n_node" elif self._edge_centered(): data_mapping = "edge centers" + grid_dim = "n_edge" else: raise ValueError( - "Data_mapping is not face, node, or edge. Could not define data_mapping." + f"neighborhood_filter requires data mapped to nodes, edges, or faces, " + f"but the last dimension {self.dims!r} does not match any grid dimension " + f"{GRID_DIMS}." ) - # Assert last dimension is a GRID dimension. - assert self.dims[-1] in GRID_DIMS, ( - f"expected last dimension of uxDataArray {self.dims[-1]!r} " - f"to be one of {GRID_DIMS}" - ) + # Ensure the grid dimension is the last axis, transposing if necessary. + # This mirrors the behaviour of UxDataset.neighborhood_filter so that + # calling the method directly on a (time, n_face) UxDataArray works. + needs_transpose = self.dims[-1] != grid_dim + uxda_work = self.transpose(..., grid_dim) if needs_transpose else self destination_data = _neighborhood_filter( - self.uxgrid, self.data, data_mapping, func=func, r=r + self.uxgrid, uxda_work.data, data_mapping, func=func, r=r ) # Construct UxDataArray for filtered variable. - uxda_filter = self._copy() - + uxda_filter = uxda_work._copy() uxda_filter.data = destination_data + # Restore original dimension order if we transposed. + if needs_transpose: + uxda_filter = uxda_filter.transpose(*self.dims) + return uxda_filter def __getattribute__(self, name): diff --git a/uxarray/core/dataset.py b/uxarray/core/dataset.py index 28b50077c..86ba9b296 100644 --- a/uxarray/core/dataset.py +++ b/uxarray/core/dataset.py @@ -682,22 +682,19 @@ def neighborhood_filter( """ destination_uxds = self._copy() - # Loop through uxDataArrays in uxDataset + # Loop through UxDataArrays in UxDataset and apply the filter to every + # variable that is mapped to a grid element (node, edge, or face). + # Variables without a grid dimension are left unchanged. for var_name in self.data_vars: uxda = self[var_name] - # Skip if uxDataArray has no GRID dimension. - grid_dims = [dim for dim in uxda.dims if dim in GRID_DIMS] - if len(grid_dims) == 0: + # Skip if UxDataArray has no GRID dimension. + if not any(dim in GRID_DIMS for dim in uxda.dims): continue - # Put GRID dimension last for UxDataArray.neighborhood_filter. - remember_dim_order = uxda.dims - uxda = uxda.transpose(..., grid_dims[0]) - # Filter uxDataArray. - uxda = uxda.neighborhood_filter(func, r) - # Restore old dimension order. - destination_uxds[var_name] = uxda.transpose(*remember_dim_order) + # UxDataArray.neighborhood_filter handles the transpose internally, + # so dimension order is always preserved. + destination_uxds[var_name] = uxda.neighborhood_filter(func, r) return destination_uxds diff --git a/uxarray/grid/grid.py b/uxarray/grid/grid.py index adb074a3e..4c199dc67 100644 --- a/uxarray/grid/grid.py +++ b/uxarray/grid/grid.py @@ -1759,7 +1759,7 @@ def get_ball_tree( coordinates : str, default="face centers" Selects which tree to query, with "nodes" selecting the Corner Nodes, "edge centers" selecting the Edge Centers of each edge, and "face centers" selecting the Face Centers of each face - coordinate_system : str, default="cartesian" + coordinate_system : str, default="spherical" Selects which coordinate type to use to create the tree, "cartesian" selecting cartesian coordinates, and "spherical" selecting spherical coordinates. distance_metric : str, default="haversine" diff --git a/uxarray/grid/neighbors.py b/uxarray/grid/neighbors.py index e6d8da625..7dfe42e00 100644 --- a/uxarray/grid/neighbors.py +++ b/uxarray/grid/neighbors.py @@ -1227,7 +1227,10 @@ def _neighborhood_filter( neighbor_indices = tree.query_radius(dest_coords, r=r) - destination_data = np.empty(data.shape) + # Initialize with NaN so that any element whose neighborhood is empty + # (e.g. an isolated point queried with a very small radius) yields NaN + # rather than uninitialized garbage memory. + destination_data = np.full(data.shape, np.nan) # Apply func along the last (grid) axis only, so any extra leading # dimensions (e.g. time) are preserved rather than being collapsed. From 5995f03c35e15fed38e24ec65123d9007030224d Mon Sep 17 00:00:00 2001 From: Rajeev Jain Date: Tue, 28 Jul 2026 12:59:10 -0500 Subject: [PATCH 13/21] Fix _copy() deep default + optimize neighborhood_filter memory usage Two related fixes in UxDataArray: 1. _copy() default deep behavior: kwargs.get('deep', None) returned None (falsy) when called with no arguments, so the uxgrid was always shallow- copied even though xarray's own default is deep=True. Changed the fallback to True so _copy() matches xarray's documented default. 2. neighborhood_filter memory optimization: the filter was calling _copy() (deep copy of data) then immediately overwriting .data with the freshly computed result, wasting a full copy of the input array. Switch to _copy(data=destination_data, deep=False) which: - Passes the filtered data directly, skipping the redundant deep copy - Uses deep=False so the returned UxDataArray shares the same uxgrid object (appropriate: the filtered result lives on the same grid topology) - Preserves all metadata (name, attrs, coords) as before --- uxarray/core/dataarray.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/uxarray/core/dataarray.py b/uxarray/core/dataarray.py index ee14b28e0..9d430be8a 100644 --- a/uxarray/core/dataarray.py +++ b/uxarray/core/dataarray.py @@ -120,7 +120,8 @@ def _copy(self, **kwargs): ``uxarray.UxDataArray``.""" copied = super()._copy(**kwargs) - deep = kwargs.get("deep", None) + # Match xarray's default: deep=True when not specified + deep = kwargs.get("deep", True) if deep: # Reinitialize the uxgrid assessor @@ -2242,9 +2243,12 @@ def neighborhood_filter( self.uxgrid, uxda_work.data, data_mapping, func=func, r=r ) - # Construct UxDataArray for filtered variable. - uxda_filter = uxda_work._copy() - uxda_filter.data = destination_data + # Construct UxDataArray for filtered variable, reusing metadata + # (name, coords, attrs, uxgrid) from the working copy. + # deep=False keeps a reference to the same uxgrid (the filtered data + # lives on the identical grid topology) and avoids a redundant deep + # copy of the now-discarded original data array. + uxda_filter = uxda_work._copy(data=destination_data, deep=False) # Restore original dimension order if we transposed. if needs_transpose: From 1d896ce02325d2c9b5a9908cb69ba273b2190285 Mon Sep 17 00:00:00 2001 From: Rajeev Jain Date: Tue, 28 Jul 2026 13:28:43 -0500 Subject: [PATCH 14/21] Add docstring Examples/See Also + expand neighborhood-filter notebook UxDataArray.neighborhood_filter and UxDataset.neighborhood_filter now include Examples and See Also docstring sections, making them consistent with xarray conventions and discoverable from help() / sphinx docs. The user-guide notebook is also expanded (30 cells): - Use ux.tutorial datasets (no raw file paths) - Add before/after visualizations (.plot.polygons) - Cover face-, node-, and edge-centered data - Show multi-dimensional (time x space) usage - Demonstrate chaining with xarray .where() and .groupby() - Add radius sweep comparison (0 deg, 2.5 deg, 5 deg, 10 deg) - Add API reference section with cross-links to related methods --- docs/user-guide/neighborhood-filter.ipynb | 229 ++++++++++++---------- uxarray/core/dataarray.py | 21 ++ uxarray/core/dataset.py | 15 ++ 3 files changed, 161 insertions(+), 104 deletions(-) diff --git a/docs/user-guide/neighborhood-filter.ipynb b/docs/user-guide/neighborhood-filter.ipynb index d0ecb478d..dfff7f9cb 100644 --- a/docs/user-guide/neighborhood-filter.ipynb +++ b/docs/user-guide/neighborhood-filter.ipynb @@ -2,182 +2,203 @@ "cells": [ { "cell_type": "markdown", - "id": "a1b2c3d4-0000-0000-0000-000000000001", "metadata": {}, - "source": [ - "# Neighborhood Filter\n", - "\n", - "This guide showcases how to apply a neighborhood filter to unstructured grid data using UXarray's `neighborhood_filter` method.\n", - "\n", - "A neighborhood filter replaces the value at each grid element with the result of a user-specified function (e.g. `np.mean`, `np.max`, `np.median`) applied to all grid elements whose centers fall within a circular neighborhood of radius `r` degrees around that element.\n", - "\n", - "This is particularly useful for variable-resolution meshes, where a constant number of neighbors does not correspond to a constant spatial scale." - ] + "source": "# Neighborhood Filter\n\nA **neighborhood filter** replaces the value at each grid element with the result\nof a user-specified function applied to all grid elements whose centers fall within\na circular neighborhood of radius `r` degrees around that element.\n\nUnlike a fixed *k*-nearest-neighbor average, a radius-based filter imposes a\nconsistent spatial scale across the whole mesh\u2014useful for variable-resolution grids\nwhere the number of neighbors varies from region to region.\n\n**Supported element types:** face-centered, node-centered, and edge-centered data.\n\n**API at a glance:**\n\n| Object | Method |\n|---|---|\n| `UxDataArray` | `da.neighborhood_filter(func=np.mean, r=5.0)` |\n| `UxDataset` | `ds.neighborhood_filter(func=np.mean, r=5.0)` |\n\nThe returned object is always the same type as the input, with the same grid, dims,\ncoordinates, name, and attributes preserved.\n" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": "## Imports" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": "from functools import partial\n\nimport numpy as np\n\nimport uxarray as ux" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": "## Load Sample Data\n\nWe use the `outCSne30-vortex` tutorial dataset (a cubed-sphere grid with 5,400\nfaces and a synthetic vortex field `psi`)." + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": "uxds = ux.tutorial.open_dataset(\"outCSne30-vortex\")\nuxda = uxds[\"psi\"]\nuxda" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": "## Visualize the Unfiltered Field" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": "uxda.plot.polygons(\n cmap=\"RdBu_r\",\n title=\"Original field (psi)\",\n width=700,\n height=400,\n)" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": "## Basic Usage: Mean Filter\n\nCalling `neighborhood_filter` with `func=np.mean` and a radius of 5\u00b0 replaces\neach face value with the mean of all face centers within 5\u00b0 of that face's center.\n" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": "uxda_smooth = uxda.neighborhood_filter(func=np.mean, r=5.0)\nuxda_smooth" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": "Note that the output is a `UxDataArray` mapped to the same grid and with the same\ndimensions as the input. The name, attributes, and coordinates are preserved.\n" }, { "cell_type": "code", "execution_count": null, - "id": "a1b2c3d4-0000-0000-0000-000000000002", "metadata": {}, "outputs": [], - "source": [ - "from functools import partial\n", - "\n", - "import numpy as np\n", - "\n", - "import uxarray as ux" - ] + "source": "uxda_smooth.plot.polygons(\n cmap=\"RdBu_r\",\n title=\"Mean filter (r = 5\u00b0)\",\n width=700,\n height=400,\n)" }, { "cell_type": "markdown", - "id": "a1b2c3d4-0000-0000-0000-000000000003", "metadata": {}, - "source": [ - "## Load Sample Data\n", - "\n", - "We use the CSne30 cubed-sphere grid bundled with UXarray's test files." - ] + "source": "### Effect of Radius\n\nIncreasing `r` produces stronger smoothing. A radius of 0\u00b0 recovers the original\nfield (the only element in any neighborhood is the element itself).\n" }, { "cell_type": "code", "execution_count": null, - "id": "a1b2c3d4-0000-0000-0000-000000000004", "metadata": {}, "outputs": [], - "source": [ - "import pathlib\n", - "\n", - "# Locate the bundled test meshfiles\n", - "repo_root = pathlib.Path(ux.__file__).parents[1]\n", - "grid_path = repo_root / \"test\" / \"meshfiles\" / \"ugrid\" / \"outCSne30\" / \"outCSne30.ug\"\n", - "data_path = repo_root / \"test\" / \"meshfiles\" / \"ugrid\" / \"outCSne30\" / \"outCSne30_vortex.nc\"\n", - "\n", - "uxds = ux.open_dataset(str(grid_path), str(data_path))\n", - "uxda = uxds[\"psi\"]\n", - "print(uxda)" - ] + "source": "import holoviews as hv\nhv.extension(\"bokeh\")\n\nplots = [\n uxda.neighborhood_filter(func=np.mean, r=r).plot.polygons(\n cmap=\"RdBu_r\",\n title=f\"r = {r}\u00b0\",\n width=350,\n height=250,\n clim=(uxda.values.min(), uxda.values.max()),\n )\n for r in [0.0, 2.5, 5.0, 10.0]\n]\n\n(plots[0] + plots[1] + plots[2] + plots[3]).cols(2)" }, { "cell_type": "markdown", - "id": "a1b2c3d4-0000-0000-0000-000000000005", "metadata": {}, - "source": [ - "## Basic Usage: Mean Filter\n", - "\n", - "Apply a mean filter with a 5-degree radius. Each face value is replaced by the mean of all face values within 5° of that face's center." - ] + "source": "## Custom Functions via `functools.partial`\n\nAny callable that accepts an `axis` keyword argument (as NumPy reductions do) works\nas the filter function. Use `functools.partial` to fix additional keyword arguments.\n" }, { "cell_type": "code", "execution_count": null, - "id": "a1b2c3d4-0000-0000-0000-000000000006", "metadata": {}, "outputs": [], - "source": [ - "uxda_mean = uxda.neighborhood_filter(func=np.mean, r=5.0)\n", - "print(uxda_mean)" - ] + "source": "# 90th-percentile filter \u2014 highlights local maxima\nuxda_p90 = uxda.neighborhood_filter(func=partial(np.percentile, q=90), r=5.0)\n\n# Maximum filter\nuxda_max = uxda.neighborhood_filter(func=np.max, r=5.0)\n\n# Median filter \u2014 robust to outliers\nuxda_med = uxda.neighborhood_filter(func=np.median, r=5.0)\n\nprint(\"max filter max :\", uxda_max.values.max())\nprint(\"p90 filter max :\", uxda_p90.values.max())\nprint(\"median filter max:\", uxda_med.values.max())" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": "(\n uxda_max.plot.polygons(cmap=\"RdBu_r\", title=\"Max filter (r=5\u00b0)\", width=350, height=250)\n + uxda_med.plot.polygons(cmap=\"RdBu_r\", title=\"Median filter (r=5\u00b0)\", width=350, height=250)\n).cols(2)" }, { "cell_type": "markdown", - "id": "a1b2c3d4-0000-0000-0000-000000000007", "metadata": {}, - "source": [ - "## Custom Functions via `functools.partial`\n", - "\n", - "Any callable that accepts an `axis` keyword argument works as the filter function. Use `functools.partial` to pass additional arguments, for example to compute the 90th percentile in each neighborhood." - ] + "source": "## Node- and Edge-Centered Data\n\n`neighborhood_filter` works for any data element type. Here we create\nsynthetic node- and edge-centered fields on a HEALPix grid and filter them.\n" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": "uxgrid = ux.Grid.from_healpix(zoom=3) # 768 faces, 770 nodes, 1536 edges\n\n# Node-centered: a gradient along longitude\nnode_da = ux.UxDataArray(\n uxgrid.node_lon.values,\n dims=[\"n_node\"],\n uxgrid=uxgrid,\n name=\"node_lon\",\n attrs={\"units\": \"degrees_east\"},\n)\n\nfiltered_node = node_da.neighborhood_filter(func=np.mean, r=10.0)\nprint(\"node input dims:\", node_da.dims, \" shape:\", node_da.shape)\nprint(\"node output dims:\", filtered_node.dims, \" shape:\", filtered_node.shape)\nprint(\"attrs preserved:\", filtered_node.attrs)" }, { "cell_type": "code", "execution_count": null, - "id": "a1b2c3d4-0000-0000-0000-000000000008", "metadata": {}, "outputs": [], - "source": [ - "# 90th-percentile filter\n", - "uxda_p90 = uxda.neighborhood_filter(func=partial(np.percentile, q=90), r=5.0)\n", - "\n", - "# Maximum filter\n", - "uxda_max = uxda.neighborhood_filter(func=np.max, r=5.0)\n", - "\n", - "print(\"p90 max:\", uxda_p90.values.max(), \" filter max:\", uxda_max.values.max())" - ] + "source": "# Edge-centered: a random field\nrng = np.random.default_rng(42)\nedge_da = ux.UxDataArray(\n rng.standard_normal(uxgrid.n_edge),\n dims=[\"n_edge\"],\n uxgrid=uxgrid,\n name=\"edge_noise\",\n)\n\nfiltered_edge = edge_da.neighborhood_filter(func=np.mean, r=10.0)\nprint(\"edge input dims:\", edge_da.dims, \" shape:\", edge_da.shape)\nprint(\"edge output dims:\", filtered_edge.dims, \" shape:\", filtered_edge.shape)" }, { "cell_type": "markdown", - "id": "a1b2c3d4-0000-0000-0000-000000000009", "metadata": {}, - "source": [ - "## Dataset-Level Usage\n", - "\n", - "`neighborhood_filter` is also available on `UxDataset`. It applies the filter to every data variable that is mapped to a grid element, leaving other variables (e.g. scalars with no grid dimension) unchanged." - ] + "source": "## Multi-Dimensional Data (e.g. Time + Space)\n\nWhen a `UxDataArray` has extra leading dimensions (e.g. `time`), `neighborhood_filter`\napplies the spatial filter independently at each time step and preserves the full\ndimension order.\n" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": "uxds_ts = ux.tutorial.open_dataset(\"outCSne30-timeseries\")\nuxda_ts = uxds_ts[\"psi\"]\n\nprint(\"Input dims:\", uxda_ts.dims, \" shape:\", uxda_ts.shape)\n\nfiltered_ts = uxda_ts.neighborhood_filter(func=np.mean, r=5.0)\n\nprint(\"Output dims:\", filtered_ts.dims, \" shape:\", filtered_ts.shape)" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": "The grid and time dimensions are both preserved. Because the filter is applied\nper time step, memory usage scales with `n_time \u00d7 n_face` as expected.\n" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": "## Dataset-Level Usage\n\n`UxDataset.neighborhood_filter` applies the filter to **every data variable**\nthat is mapped to a grid element. Variables without a grid dimension (e.g. scalars\nor time-only arrays) are passed through unchanged.\n" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": "uxds_filtered = uxds.neighborhood_filter(func=np.mean, r=5.0)\nuxds_filtered" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": "## Chaining with xarray Operations\n\nBecause `neighborhood_filter` returns a proper `UxDataArray` with its `uxgrid`\npreserved, you can chain it with any standard xarray operation.\n" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": "# Apply the filter and then mask values below zero\nresult = (\n uxda\n .neighborhood_filter(func=np.mean, r=5.0)\n .where(lambda x: x > 0)\n)\nprint(\"Masked result type:\", type(result).__name__)\nprint(\"uxgrid preserved:\", result.uxgrid is not None)\nprint(\"Positive fraction:\", float((result > 0).sum()) / result.size)" }, { "cell_type": "code", "execution_count": null, - "id": "a1b2c3d4-0000-0000-0000-000000000010", "metadata": {}, "outputs": [], - "source": [ - "uxds_filtered = uxds.neighborhood_filter(func=np.mean, r=5.0)\n", - "print(uxds_filtered)" - ] + "source": "# Group by latitude band after smoothing (standard xarray groupby)\nimport xarray as xr\n\nlat_bins = xr.DataArray(\n np.digitize(uxda.uxgrid.face_lat.values, bins=np.arange(-90, 91, 30)),\n dims=[\"n_face\"],\n)\n\nzonal_smooth = (\n uxda\n .neighborhood_filter(func=np.mean, r=5.0)\n .groupby(lat_bins)\n .mean()\n)\nprint(\"Grouped result type:\", type(zonal_smooth).__name__)\nprint(\"Zonal means:\", zonal_smooth.values)" }, { "cell_type": "markdown", - "id": "a1b2c3d4-0000-0000-0000-000000000011", "metadata": {}, - "source": [ - "## Handling Extra Dimensions\n", - "\n", - "If the data has extra leading dimensions (e.g. `time`), the filter is applied independently along the grid axis and all extra dimensions are preserved." - ] + "source": "## Empty Neighborhoods and NaN Behavior\n\nIf the search radius is so small that *no* neighbor is found for a given element,\nthe result for that element is `NaN` rather than an uninitialized garbage value.\nIn practice this only happens for extremely small radii on very coarse grids; the\nfilter always includes the queried element itself, so `r = 0` is safe.\n" }, { "cell_type": "code", "execution_count": null, - "id": "a1b2c3d4-0000-0000-0000-000000000012", "metadata": {}, "outputs": [], - "source": [ - "from uxarray import UxDataArray\n", - "\n", - "# Simulate two time steps\n", - "data = np.stack([uxda.values, uxda.values * 2.0]) # shape (2, n_face)\n", - "uxda_time = UxDataArray(\n", - " data, dims=[\"time\", \"n_face\"], uxgrid=uxda.uxgrid, name=\"psi_time\"\n", - ")\n", - "\n", - "filtered_time = uxda_time.neighborhood_filter(func=np.mean, r=5.0)\n", - "print(\"Input shape :\", uxda_time.shape)\n", - "print(\"Output shape:\", filtered_time.shape)\n", - "print(\"Output dims :\", filtered_time.dims)" - ] + "source": "uxgrid_coarse = ux.Grid.from_healpix(zoom=1) # 48 faces\nda_coarse = ux.UxDataArray(\n np.arange(uxgrid_coarse.n_face, dtype=float),\n dims=[\"n_face\"],\n uxgrid=uxgrid_coarse,\n)\n\n# r = 0 always catches at least the element itself \u2192 no NaNs\nfiltered_r0 = da_coarse.neighborhood_filter(func=np.mean, r=0.0)\nprint(\"r = 0: NaN count =\", int(np.isnan(filtered_r0.values).sum()))\n\n# r = 360 catches every element \u2192 all values equal the global mean\nfiltered_global = da_coarse.neighborhood_filter(func=np.mean, r=360.0)\nprint(\"r = 360: all equal global mean?\",\n np.allclose(filtered_global.values, da_coarse.values.mean()))" }, { "cell_type": "markdown", - "id": "a1b2c3d4-0000-0000-0000-000000000013", "metadata": {}, - "source": [ - "## Empty Neighborhoods\n", - "\n", - "If the radius `r` is so small that no neighbor is found for a given element (unlikely for face-centered data because the query always finds the element itself, but possible for edge- or node-centered data with very small radii), the result for that element is `NaN` rather than an arbitrary value." - ] + "source": "## API Reference\n\nSee also:\n\n- {py:meth}`uxarray.UxDataArray.neighborhood_filter`\n- {py:meth}`uxarray.UxDataset.neighborhood_filter`\n\nRelated methods that apply aggregations across different grid element types:\n\n- {py:meth}`uxarray.UxDataArray.topological_mean` \u2014 aggregate node\u2192face, node\u2192edge, etc.\n- {py:meth}`uxarray.UxDataArray.zonal_mean` \u2014 latitude-band averages\n- {py:meth}`uxarray.UxDataArray.azimuthal_mean` \u2014 rings of constant great-circle distance\n" } ], "metadata": { "kernelspec": { - "display_name": "Python 3", + "display_name": "Python 3 (ipykernel)", "language": "python", "name": "python3" }, "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", "name": "python", + "pygments_lexer": "ipython3", "version": "3.13.0" } }, "nbformat": 4, "nbformat_minor": 5 -} +} \ No newline at end of file diff --git a/uxarray/core/dataarray.py b/uxarray/core/dataarray.py index 9d430be8a..15b57b6b8 100644 --- a/uxarray/core/dataarray.py +++ b/uxarray/core/dataarray.py @@ -2215,6 +2215,27 @@ def neighborhood_filter( ------- uxda_filter : UxDataArray Filtered data. + + Examples + -------- + Apply a mean filter with a 5-degree radius: + + >>> import numpy as np + >>> import uxarray as ux + >>> uxds = ux.tutorial.open_dataset("outCSne30-vortex") + >>> uxda = uxds["psi"] + >>> smoothed = uxda.neighborhood_filter(func=np.mean, r=5.0) + + Use ``functools.partial`` for functions requiring extra arguments: + + >>> from functools import partial + >>> p90 = uxda.neighborhood_filter(func=partial(np.percentile, q=90), r=5.0) + + See Also + -------- + UxDataArray.topological_mean : Aggregate values across neighboring grid element types. + UxDataArray.zonal_mean : Average over latitude bands. + UxDataArray.azimuthal_mean : Average over rings of constant great-circle distance. """ if self._face_centered(): diff --git a/uxarray/core/dataset.py b/uxarray/core/dataset.py index 86ba9b296..b220cd00d 100644 --- a/uxarray/core/dataset.py +++ b/uxarray/core/dataset.py @@ -679,6 +679,21 @@ def neighborhood_filter( ------- destination_uxds : UxDataset Filtered dataset. + + Examples + -------- + Apply a mean filter to all grid-mapped variables in a dataset: + + >>> import numpy as np + >>> import uxarray as ux + >>> uxds = ux.tutorial.open_dataset("outCSne30-vortex") + >>> uxds_smooth = uxds.neighborhood_filter(func=np.mean, r=5.0) + + See Also + -------- + UxDataArray.neighborhood_filter : Filter a single data variable. + UxDataArray.zonal_mean : Average over latitude bands. + UxDataArray.azimuthal_mean : Average over rings of constant great-circle distance. """ destination_uxds = self._copy() From b37974d2ee10de5c49c2b47aa6c1548a5e784222 Mon Sep 17 00:00:00 2001 From: Rajeev Jain Date: Tue, 28 Jul 2026 13:55:38 -0500 Subject: [PATCH 15/21] Fix pre-commit: trailing newlines + ruff import sort/format --- docs/user-guide/neighborhood-filter.ipynb | 215 +++++++++++++++++++--- test/core/test_dataarray.py | 1 - 2 files changed, 194 insertions(+), 22 deletions(-) diff --git a/docs/user-guide/neighborhood-filter.ipynb b/docs/user-guide/neighborhood-filter.ipynb index dfff7f9cb..7fae09502 100644 --- a/docs/user-guide/neighborhood-filter.ipynb +++ b/docs/user-guide/neighborhood-filter.ipynb @@ -2,183 +2,356 @@ "cells": [ { "cell_type": "markdown", + "id": "7fb27b941602401d91542211134fc71a", "metadata": {}, - "source": "# Neighborhood Filter\n\nA **neighborhood filter** replaces the value at each grid element with the result\nof a user-specified function applied to all grid elements whose centers fall within\na circular neighborhood of radius `r` degrees around that element.\n\nUnlike a fixed *k*-nearest-neighbor average, a radius-based filter imposes a\nconsistent spatial scale across the whole mesh\u2014useful for variable-resolution grids\nwhere the number of neighbors varies from region to region.\n\n**Supported element types:** face-centered, node-centered, and edge-centered data.\n\n**API at a glance:**\n\n| Object | Method |\n|---|---|\n| `UxDataArray` | `da.neighborhood_filter(func=np.mean, r=5.0)` |\n| `UxDataset` | `ds.neighborhood_filter(func=np.mean, r=5.0)` |\n\nThe returned object is always the same type as the input, with the same grid, dims,\ncoordinates, name, and attributes preserved.\n" + "source": "# Neighborhood Filter\n\nA **neighborhood filter** replaces the value at each grid element with the result\nof a user-specified function applied to all grid elements whose centers fall within\na circular neighborhood of radius `r` degrees around that element.\n\nUnlike a fixed *k*-nearest-neighbor average, a radius-based filter imposes a\nconsistent spatial scale across the whole mesh—useful for variable-resolution grids\nwhere the number of neighbors varies from region to region.\n\n**Supported element types:** face-centered, node-centered, and edge-centered data.\n\n**API at a glance:**\n\n| Object | Method |\n|---|---|\n| `UxDataArray` | `da.neighborhood_filter(func=np.mean, r=5.0)` |\n| `UxDataset` | `ds.neighborhood_filter(func=np.mean, r=5.0)` |\n\nThe returned object is always the same type as the input, with the same grid, dims,\ncoordinates, name, and attributes preserved.\n" }, { "cell_type": "markdown", + "id": "acae54e37e7d407bbb7b55eff062a284", "metadata": {}, "source": "## Imports" }, { "cell_type": "code", "execution_count": null, + "id": "9a63283cbaf04dbcab1f6479b197f3a8", "metadata": {}, "outputs": [], - "source": "from functools import partial\n\nimport numpy as np\n\nimport uxarray as ux" + "source": [ + "from functools import partial\n", + "\n", + "import numpy as np\n", + "\n", + "import uxarray as ux" + ] }, { "cell_type": "markdown", + "id": "8dd0d8092fe74a7c96281538738b07e2", "metadata": {}, "source": "## Load Sample Data\n\nWe use the `outCSne30-vortex` tutorial dataset (a cubed-sphere grid with 5,400\nfaces and a synthetic vortex field `psi`)." }, { "cell_type": "code", "execution_count": null, + "id": "72eea5119410473aa328ad9291626812", "metadata": {}, "outputs": [], - "source": "uxds = ux.tutorial.open_dataset(\"outCSne30-vortex\")\nuxda = uxds[\"psi\"]\nuxda" + "source": [ + "uxds = ux.tutorial.open_dataset(\"outCSne30-vortex\")\n", + "uxda = uxds[\"psi\"]\n", + "uxda" + ] }, { "cell_type": "markdown", + "id": "8edb47106e1a46a883d545849b8ab81b", "metadata": {}, "source": "## Visualize the Unfiltered Field" }, { "cell_type": "code", "execution_count": null, + "id": "10185d26023b46108eb7d9f57d49d2b3", "metadata": {}, "outputs": [], - "source": "uxda.plot.polygons(\n cmap=\"RdBu_r\",\n title=\"Original field (psi)\",\n width=700,\n height=400,\n)" + "source": [ + "uxda.plot.polygons(\n", + " cmap=\"RdBu_r\",\n", + " title=\"Original field (psi)\",\n", + " width=700,\n", + " height=400,\n", + ")" + ] }, { "cell_type": "markdown", + "id": "8763a12b2bbd4a93a75aff182afb95dc", "metadata": {}, - "source": "## Basic Usage: Mean Filter\n\nCalling `neighborhood_filter` with `func=np.mean` and a radius of 5\u00b0 replaces\neach face value with the mean of all face centers within 5\u00b0 of that face's center.\n" + "source": "## Basic Usage: Mean Filter\n\nCalling `neighborhood_filter` with `func=np.mean` and a radius of 5° replaces\neach face value with the mean of all face centers within 5° of that face's center.\n" }, { "cell_type": "code", "execution_count": null, + "id": "7623eae2785240b9bd12b16a66d81610", "metadata": {}, "outputs": [], - "source": "uxda_smooth = uxda.neighborhood_filter(func=np.mean, r=5.0)\nuxda_smooth" + "source": [ + "uxda_smooth = uxda.neighborhood_filter(func=np.mean, r=5.0)\n", + "uxda_smooth" + ] }, { "cell_type": "markdown", + "id": "7cdc8c89c7104fffa095e18ddfef8986", "metadata": {}, "source": "Note that the output is a `UxDataArray` mapped to the same grid and with the same\ndimensions as the input. The name, attributes, and coordinates are preserved.\n" }, { "cell_type": "code", "execution_count": null, + "id": "b118ea5561624da68c537baed56e602f", "metadata": {}, "outputs": [], - "source": "uxda_smooth.plot.polygons(\n cmap=\"RdBu_r\",\n title=\"Mean filter (r = 5\u00b0)\",\n width=700,\n height=400,\n)" + "source": [ + "uxda_smooth.plot.polygons(\n", + " cmap=\"RdBu_r\",\n", + " title=\"Mean filter (r = 5°)\",\n", + " width=700,\n", + " height=400,\n", + ")" + ] }, { "cell_type": "markdown", + "id": "938c804e27f84196a10c8828c723f798", "metadata": {}, - "source": "### Effect of Radius\n\nIncreasing `r` produces stronger smoothing. A radius of 0\u00b0 recovers the original\nfield (the only element in any neighborhood is the element itself).\n" + "source": "### Effect of Radius\n\nIncreasing `r` produces stronger smoothing. A radius of 0° recovers the original\nfield (the only element in any neighborhood is the element itself).\n" }, { "cell_type": "code", "execution_count": null, + "id": "504fb2a444614c0babb325280ed9130a", "metadata": {}, "outputs": [], - "source": "import holoviews as hv\nhv.extension(\"bokeh\")\n\nplots = [\n uxda.neighborhood_filter(func=np.mean, r=r).plot.polygons(\n cmap=\"RdBu_r\",\n title=f\"r = {r}\u00b0\",\n width=350,\n height=250,\n clim=(uxda.values.min(), uxda.values.max()),\n )\n for r in [0.0, 2.5, 5.0, 10.0]\n]\n\n(plots[0] + plots[1] + plots[2] + plots[3]).cols(2)" + "source": [ + "import holoviews as hv\n", + "\n", + "hv.extension(\"bokeh\")\n", + "\n", + "plots = [\n", + " uxda.neighborhood_filter(func=np.mean, r=r).plot.polygons(\n", + " cmap=\"RdBu_r\",\n", + " title=f\"r = {r}°\",\n", + " width=350,\n", + " height=250,\n", + " clim=(uxda.values.min(), uxda.values.max()),\n", + " )\n", + " for r in [0.0, 2.5, 5.0, 10.0]\n", + "]\n", + "\n", + "(plots[0] + plots[1] + plots[2] + plots[3]).cols(2)" + ] }, { "cell_type": "markdown", + "id": "59bbdb311c014d738909a11f9e486628", "metadata": {}, "source": "## Custom Functions via `functools.partial`\n\nAny callable that accepts an `axis` keyword argument (as NumPy reductions do) works\nas the filter function. Use `functools.partial` to fix additional keyword arguments.\n" }, { "cell_type": "code", "execution_count": null, + "id": "b43b363d81ae4b689946ece5c682cd59", "metadata": {}, "outputs": [], - "source": "# 90th-percentile filter \u2014 highlights local maxima\nuxda_p90 = uxda.neighborhood_filter(func=partial(np.percentile, q=90), r=5.0)\n\n# Maximum filter\nuxda_max = uxda.neighborhood_filter(func=np.max, r=5.0)\n\n# Median filter \u2014 robust to outliers\nuxda_med = uxda.neighborhood_filter(func=np.median, r=5.0)\n\nprint(\"max filter max :\", uxda_max.values.max())\nprint(\"p90 filter max :\", uxda_p90.values.max())\nprint(\"median filter max:\", uxda_med.values.max())" + "source": [ + "# 90th-percentile filter — highlights local maxima\n", + "uxda_p90 = uxda.neighborhood_filter(func=partial(np.percentile, q=90), r=5.0)\n", + "\n", + "# Maximum filter\n", + "uxda_max = uxda.neighborhood_filter(func=np.max, r=5.0)\n", + "\n", + "# Median filter — robust to outliers\n", + "uxda_med = uxda.neighborhood_filter(func=np.median, r=5.0)\n", + "\n", + "print(\"max filter max :\", uxda_max.values.max())\n", + "print(\"p90 filter max :\", uxda_p90.values.max())\n", + "print(\"median filter max:\", uxda_med.values.max())" + ] }, { "cell_type": "code", "execution_count": null, + "id": "8a65eabff63a45729fe45fb5ade58bdc", "metadata": {}, "outputs": [], - "source": "(\n uxda_max.plot.polygons(cmap=\"RdBu_r\", title=\"Max filter (r=5\u00b0)\", width=350, height=250)\n + uxda_med.plot.polygons(cmap=\"RdBu_r\", title=\"Median filter (r=5\u00b0)\", width=350, height=250)\n).cols(2)" + "source": [ + "(\n", + " uxda_max.plot.polygons(\n", + " cmap=\"RdBu_r\", title=\"Max filter (r=5°)\", width=350, height=250\n", + " )\n", + " + uxda_med.plot.polygons(\n", + " cmap=\"RdBu_r\", title=\"Median filter (r=5°)\", width=350, height=250\n", + " )\n", + ").cols(2)" + ] }, { "cell_type": "markdown", + "id": "c3933fab20d04ec698c2621248eb3be0", "metadata": {}, "source": "## Node- and Edge-Centered Data\n\n`neighborhood_filter` works for any data element type. Here we create\nsynthetic node- and edge-centered fields on a HEALPix grid and filter them.\n" }, { "cell_type": "code", "execution_count": null, + "id": "4dd4641cc4064e0191573fe9c69df29b", "metadata": {}, "outputs": [], - "source": "uxgrid = ux.Grid.from_healpix(zoom=3) # 768 faces, 770 nodes, 1536 edges\n\n# Node-centered: a gradient along longitude\nnode_da = ux.UxDataArray(\n uxgrid.node_lon.values,\n dims=[\"n_node\"],\n uxgrid=uxgrid,\n name=\"node_lon\",\n attrs={\"units\": \"degrees_east\"},\n)\n\nfiltered_node = node_da.neighborhood_filter(func=np.mean, r=10.0)\nprint(\"node input dims:\", node_da.dims, \" shape:\", node_da.shape)\nprint(\"node output dims:\", filtered_node.dims, \" shape:\", filtered_node.shape)\nprint(\"attrs preserved:\", filtered_node.attrs)" + "source": [ + "uxgrid = ux.Grid.from_healpix(zoom=3) # 768 faces, 770 nodes, 1536 edges\n", + "\n", + "# Node-centered: a gradient along longitude\n", + "node_da = ux.UxDataArray(\n", + " uxgrid.node_lon.values,\n", + " dims=[\"n_node\"],\n", + " uxgrid=uxgrid,\n", + " name=\"node_lon\",\n", + " attrs={\"units\": \"degrees_east\"},\n", + ")\n", + "\n", + "filtered_node = node_da.neighborhood_filter(func=np.mean, r=10.0)\n", + "print(\"node input dims:\", node_da.dims, \" shape:\", node_da.shape)\n", + "print(\"node output dims:\", filtered_node.dims, \" shape:\", filtered_node.shape)\n", + "print(\"attrs preserved:\", filtered_node.attrs)" + ] }, { "cell_type": "code", "execution_count": null, + "id": "8309879909854d7188b41380fd92a7c3", "metadata": {}, "outputs": [], - "source": "# Edge-centered: a random field\nrng = np.random.default_rng(42)\nedge_da = ux.UxDataArray(\n rng.standard_normal(uxgrid.n_edge),\n dims=[\"n_edge\"],\n uxgrid=uxgrid,\n name=\"edge_noise\",\n)\n\nfiltered_edge = edge_da.neighborhood_filter(func=np.mean, r=10.0)\nprint(\"edge input dims:\", edge_da.dims, \" shape:\", edge_da.shape)\nprint(\"edge output dims:\", filtered_edge.dims, \" shape:\", filtered_edge.shape)" + "source": [ + "# Edge-centered: a random field\n", + "rng = np.random.default_rng(42)\n", + "edge_da = ux.UxDataArray(\n", + " rng.standard_normal(uxgrid.n_edge),\n", + " dims=[\"n_edge\"],\n", + " uxgrid=uxgrid,\n", + " name=\"edge_noise\",\n", + ")\n", + "\n", + "filtered_edge = edge_da.neighborhood_filter(func=np.mean, r=10.0)\n", + "print(\"edge input dims:\", edge_da.dims, \" shape:\", edge_da.shape)\n", + "print(\"edge output dims:\", filtered_edge.dims, \" shape:\", filtered_edge.shape)" + ] }, { "cell_type": "markdown", + "id": "3ed186c9a28b402fb0bc4494df01f08d", "metadata": {}, "source": "## Multi-Dimensional Data (e.g. Time + Space)\n\nWhen a `UxDataArray` has extra leading dimensions (e.g. `time`), `neighborhood_filter`\napplies the spatial filter independently at each time step and preserves the full\ndimension order.\n" }, { "cell_type": "code", "execution_count": null, + "id": "cb1e1581032b452c9409d6c6813c49d1", "metadata": {}, "outputs": [], - "source": "uxds_ts = ux.tutorial.open_dataset(\"outCSne30-timeseries\")\nuxda_ts = uxds_ts[\"psi\"]\n\nprint(\"Input dims:\", uxda_ts.dims, \" shape:\", uxda_ts.shape)\n\nfiltered_ts = uxda_ts.neighborhood_filter(func=np.mean, r=5.0)\n\nprint(\"Output dims:\", filtered_ts.dims, \" shape:\", filtered_ts.shape)" + "source": [ + "uxds_ts = ux.tutorial.open_dataset(\"outCSne30-timeseries\")\n", + "uxda_ts = uxds_ts[\"psi\"]\n", + "\n", + "print(\"Input dims:\", uxda_ts.dims, \" shape:\", uxda_ts.shape)\n", + "\n", + "filtered_ts = uxda_ts.neighborhood_filter(func=np.mean, r=5.0)\n", + "\n", + "print(\"Output dims:\", filtered_ts.dims, \" shape:\", filtered_ts.shape)" + ] }, { "cell_type": "markdown", + "id": "379cbbc1e968416e875cc15c1202d7eb", "metadata": {}, - "source": "The grid and time dimensions are both preserved. Because the filter is applied\nper time step, memory usage scales with `n_time \u00d7 n_face` as expected.\n" + "source": "The grid and time dimensions are both preserved. Because the filter is applied\nper time step, memory usage scales with `n_time × n_face` as expected.\n" }, { "cell_type": "markdown", + "id": "277c27b1587741f2af2001be3712ef0d", "metadata": {}, "source": "## Dataset-Level Usage\n\n`UxDataset.neighborhood_filter` applies the filter to **every data variable**\nthat is mapped to a grid element. Variables without a grid dimension (e.g. scalars\nor time-only arrays) are passed through unchanged.\n" }, { "cell_type": "code", "execution_count": null, + "id": "db7b79bc585a40fcaf58bf750017e135", "metadata": {}, "outputs": [], - "source": "uxds_filtered = uxds.neighborhood_filter(func=np.mean, r=5.0)\nuxds_filtered" + "source": [ + "uxds_filtered = uxds.neighborhood_filter(func=np.mean, r=5.0)\n", + "uxds_filtered" + ] }, { "cell_type": "markdown", + "id": "916684f9a58a4a2aa5f864670399430d", "metadata": {}, "source": "## Chaining with xarray Operations\n\nBecause `neighborhood_filter` returns a proper `UxDataArray` with its `uxgrid`\npreserved, you can chain it with any standard xarray operation.\n" }, { "cell_type": "code", "execution_count": null, + "id": "1671c31a24314836a5b85d7ef7fbf015", "metadata": {}, "outputs": [], - "source": "# Apply the filter and then mask values below zero\nresult = (\n uxda\n .neighborhood_filter(func=np.mean, r=5.0)\n .where(lambda x: x > 0)\n)\nprint(\"Masked result type:\", type(result).__name__)\nprint(\"uxgrid preserved:\", result.uxgrid is not None)\nprint(\"Positive fraction:\", float((result > 0).sum()) / result.size)" + "source": [ + "# Apply the filter and then mask values below zero\n", + "result = uxda.neighborhood_filter(func=np.mean, r=5.0).where(lambda x: x > 0)\n", + "print(\"Masked result type:\", type(result).__name__)\n", + "print(\"uxgrid preserved:\", result.uxgrid is not None)\n", + "print(\"Positive fraction:\", float((result > 0).sum()) / result.size)" + ] }, { "cell_type": "code", "execution_count": null, + "id": "33b0902fd34d4ace834912fa1002cf8e", "metadata": {}, "outputs": [], - "source": "# Group by latitude band after smoothing (standard xarray groupby)\nimport xarray as xr\n\nlat_bins = xr.DataArray(\n np.digitize(uxda.uxgrid.face_lat.values, bins=np.arange(-90, 91, 30)),\n dims=[\"n_face\"],\n)\n\nzonal_smooth = (\n uxda\n .neighborhood_filter(func=np.mean, r=5.0)\n .groupby(lat_bins)\n .mean()\n)\nprint(\"Grouped result type:\", type(zonal_smooth).__name__)\nprint(\"Zonal means:\", zonal_smooth.values)" + "source": [ + "# Group by latitude band after smoothing (standard xarray groupby)\n", + "import xarray as xr\n", + "\n", + "lat_bins = xr.DataArray(\n", + " np.digitize(uxda.uxgrid.face_lat.values, bins=np.arange(-90, 91, 30)),\n", + " dims=[\"n_face\"],\n", + ")\n", + "\n", + "zonal_smooth = uxda.neighborhood_filter(func=np.mean, r=5.0).groupby(lat_bins).mean()\n", + "print(\"Grouped result type:\", type(zonal_smooth).__name__)\n", + "print(\"Zonal means:\", zonal_smooth.values)" + ] }, { "cell_type": "markdown", + "id": "f6fa52606d8c4a75a9b52967216f8f3f", "metadata": {}, "source": "## Empty Neighborhoods and NaN Behavior\n\nIf the search radius is so small that *no* neighbor is found for a given element,\nthe result for that element is `NaN` rather than an uninitialized garbage value.\nIn practice this only happens for extremely small radii on very coarse grids; the\nfilter always includes the queried element itself, so `r = 0` is safe.\n" }, { "cell_type": "code", "execution_count": null, + "id": "f5a1fa73e5044315a093ec459c9be902", "metadata": {}, "outputs": [], - "source": "uxgrid_coarse = ux.Grid.from_healpix(zoom=1) # 48 faces\nda_coarse = ux.UxDataArray(\n np.arange(uxgrid_coarse.n_face, dtype=float),\n dims=[\"n_face\"],\n uxgrid=uxgrid_coarse,\n)\n\n# r = 0 always catches at least the element itself \u2192 no NaNs\nfiltered_r0 = da_coarse.neighborhood_filter(func=np.mean, r=0.0)\nprint(\"r = 0: NaN count =\", int(np.isnan(filtered_r0.values).sum()))\n\n# r = 360 catches every element \u2192 all values equal the global mean\nfiltered_global = da_coarse.neighborhood_filter(func=np.mean, r=360.0)\nprint(\"r = 360: all equal global mean?\",\n np.allclose(filtered_global.values, da_coarse.values.mean()))" + "source": [ + "uxgrid_coarse = ux.Grid.from_healpix(zoom=1) # 48 faces\n", + "da_coarse = ux.UxDataArray(\n", + " np.arange(uxgrid_coarse.n_face, dtype=float),\n", + " dims=[\"n_face\"],\n", + " uxgrid=uxgrid_coarse,\n", + ")\n", + "\n", + "# r = 0 always catches at least the element itself → no NaNs\n", + "filtered_r0 = da_coarse.neighborhood_filter(func=np.mean, r=0.0)\n", + "print(\"r = 0: NaN count =\", int(np.isnan(filtered_r0.values).sum()))\n", + "\n", + "# r = 360 catches every element → all values equal the global mean\n", + "filtered_global = da_coarse.neighborhood_filter(func=np.mean, r=360.0)\n", + "print(\n", + " \"r = 360: all equal global mean?\",\n", + " np.allclose(filtered_global.values, da_coarse.values.mean()),\n", + ")" + ] }, { "cell_type": "markdown", + "id": "cdf66aed5cc84ca1b48e60bad68798a8", "metadata": {}, - "source": "## API Reference\n\nSee also:\n\n- {py:meth}`uxarray.UxDataArray.neighborhood_filter`\n- {py:meth}`uxarray.UxDataset.neighborhood_filter`\n\nRelated methods that apply aggregations across different grid element types:\n\n- {py:meth}`uxarray.UxDataArray.topological_mean` \u2014 aggregate node\u2192face, node\u2192edge, etc.\n- {py:meth}`uxarray.UxDataArray.zonal_mean` \u2014 latitude-band averages\n- {py:meth}`uxarray.UxDataArray.azimuthal_mean` \u2014 rings of constant great-circle distance\n" + "source": "## API Reference\n\nSee also:\n\n- {py:meth}`uxarray.UxDataArray.neighborhood_filter`\n- {py:meth}`uxarray.UxDataset.neighborhood_filter`\n\nRelated methods that apply aggregations across different grid element types:\n\n- {py:meth}`uxarray.UxDataArray.topological_mean` — aggregate node→face, node→edge, etc.\n- {py:meth}`uxarray.UxDataArray.zonal_mean` — latitude-band averages\n- {py:meth}`uxarray.UxDataArray.azimuthal_mean` — rings of constant great-circle distance\n" } ], "metadata": { @@ -201,4 +374,4 @@ }, "nbformat": 4, "nbformat_minor": 5 -} \ No newline at end of file +} diff --git a/test/core/test_dataarray.py b/test/core/test_dataarray.py index 888ff93b1..7ebc9f01d 100644 --- a/test/core/test_dataarray.py +++ b/test/core/test_dataarray.py @@ -306,4 +306,3 @@ def test_auto_transpose_direct_on_uxdataarray(self, gridpath, datasetpath): # Dim order must be restored to (n_face, time) assert filtered2.dims == ("n_face", "time") assert filtered2.shape == (uxda.shape[0], 2) - From 82b3ae220d66859bc7adc0ecf80b35997a20e4e8 Mon Sep 17 00:00:00 2001 From: Rajeev Jain Date: Mon, 3 Aug 2026 08:50:19 -0500 Subject: [PATCH 16/21] Correct NaN rationale in neighborhood filter docs and comments Neighborhoods are never empty: query_radius rejects a negative radius and every element is its own neighbor at distance 0, so r=0 returns the original values. Reword the code comment and user-guide section accordingly, keeping np.full(NaN) allocation only as a defensive measure over np.empty. Also raise DataCenteringError instead of a bare ValueError for non-grid-mapped data, matching the error types introduced in uxarray/errors.py, and drop a redundant self-import in UxDataArray.isel. --- docs/user-guide/neighborhood-filter.ipynb | 17 +++++++++-- test/core/test_dataarray.py | 37 +++++++++++------------ uxarray/core/dataarray.py | 5 ++- uxarray/grid/neighbors.py | 8 +++-- 4 files changed, 39 insertions(+), 28 deletions(-) diff --git a/docs/user-guide/neighborhood-filter.ipynb b/docs/user-guide/neighborhood-filter.ipynb index 7fae09502..4e2b7d980 100644 --- a/docs/user-guide/neighborhood-filter.ipynb +++ b/docs/user-guide/neighborhood-filter.ipynb @@ -319,7 +319,18 @@ "cell_type": "markdown", "id": "f6fa52606d8c4a75a9b52967216f8f3f", "metadata": {}, - "source": "## Empty Neighborhoods and NaN Behavior\n\nIf the search radius is so small that *no* neighbor is found for a given element,\nthe result for that element is `NaN` rather than an uninitialized garbage value.\nIn practice this only happens for extremely small radii on very coarse grids; the\nfilter always includes the queried element itself, so `r = 0` is safe.\n" + "source": [ + "## Radius Edge Cases\n", + "\n", + "Every element is its own neighbor at distance 0, and `query_radius` rejects a\n", + "negative radius, so a neighborhood is never empty. `r = 0` simply returns the\n", + "original values, and a radius large enough to span the sphere returns the global\n", + "reduction everywhere.\n", + "\n", + "The output array is nonetheless allocated with `NaN` rather than uninitialized\n", + "memory, so any unexpected gap would show up as an obvious `NaN` instead of\n", + "garbage values.\n" + ] }, { "cell_type": "code", @@ -335,9 +346,9 @@ " uxgrid=uxgrid_coarse,\n", ")\n", "\n", - "# r = 0 always catches at least the element itself → no NaNs\n", + "# r = 0 catches the element itself → output matches the input exactly\n", "filtered_r0 = da_coarse.neighborhood_filter(func=np.mean, r=0.0)\n", - "print(\"r = 0: NaN count =\", int(np.isnan(filtered_r0.values).sum()))\n", + "print(\"r = 0: unchanged?\", np.allclose(filtered_r0.values, da_coarse.values))\n", "\n", "# r = 360 catches every element → all values equal the global mean\n", "filtered_global = da_coarse.neighborhood_filter(func=np.mean, r=360.0)\n", diff --git a/test/core/test_dataarray.py b/test/core/test_dataarray.py index 846eb8e80..61c7858df 100644 --- a/test/core/test_dataarray.py +++ b/test/core/test_dataarray.py @@ -1,6 +1,6 @@ import numpy as np import uxarray as ux -from uxarray.errors import DimensionError +from uxarray.errors import DataCenteringError, DimensionError from uxarray.grid.geometry import _build_polygon_shells, _build_corrected_polygon_shells from uxarray.core.dataset import UxDataset, UxDataArray import pytest @@ -255,30 +255,29 @@ def test_invalid_data_location(self): uxgrid = ux.Grid.from_healpix(zoom=1) uxda = UxDataArray(np.ones(5), dims=["other_dim"], uxgrid=uxgrid) - with pytest.raises(ValueError): + with pytest.raises(DataCenteringError): uxda.neighborhood_filter(func=np.mean, r=1.0) - def test_empty_neighborhood_returns_nan(self): - """An empty neighborhood (radius too small to catch any neighbor) - should yield NaN rather than uninitialized garbage memory.""" + + def test_radius_edge_cases_never_produce_nan(self): + """Every element is its own neighbor at distance 0, so no neighborhood + is ever empty and the output never contains NaN.""" uxgrid = ux.Grid.from_healpix(zoom=1) data = np.arange(uxgrid.n_face, dtype=float) uxda = UxDataArray(data, dims=["n_face"], uxgrid=uxgrid, name="face_var") - # A radius of exactly 0 will still catch the element itself. - # Use a tiny but non-zero radius that finds the element itself via query. - # r=0 catches the element itself; test with r=360 to check no-NaN case. - filtered = uxda.neighborhood_filter(func=np.mean, r=360.0) - assert not np.any(np.isnan(filtered.values)), ( - "Global-radius filter should produce no NaN values" - ) - - # Now verify NaN initialization: create a grid but forcibly test - # the np.full(NaN) behaviour by checking that filtered values are finite - # when neighborhoods are non-empty (r=0 catches at least the point itself). + # r=0 catches only the element itself, so the data is returned unchanged filtered_zero = uxda.neighborhood_filter(func=np.mean, r=0.0) - assert not np.any(np.isnan(filtered_zero.values)), ( - "r=0 filter should include the element itself, so no NaN" - ) + assert not np.any(np.isnan(filtered_zero.values)) + np.testing.assert_allclose(filtered_zero.values, data) + + # a radius spanning the sphere catches every element + filtered_all = uxda.neighborhood_filter(func=np.mean, r=360.0) + assert not np.any(np.isnan(filtered_all.values)) + np.testing.assert_allclose(filtered_all.values, data.mean()) + + # a negative radius is rejected by BallTree.query_radius + with pytest.raises(AssertionError): + uxda.neighborhood_filter(func=np.mean, r=-1.0) def test_auto_transpose_direct_on_uxdataarray(self, gridpath, datasetpath): """Calling neighborhood_filter directly on a (time, n_face) UxDataArray diff --git a/uxarray/core/dataarray.py b/uxarray/core/dataarray.py index f53898359..3239a30c5 100644 --- a/uxarray/core/dataarray.py +++ b/uxarray/core/dataarray.py @@ -1997,7 +1997,6 @@ def isel( DimensionError (subclass of ValueError) If more than one grid dimension is selected and `ignore_grid=False`. """ - from uxarray.core.dataarray import UxDataArray from uxarray.core.utils import _validate_indexers indexers, grid_dims = _validate_indexers( @@ -2253,9 +2252,9 @@ def neighborhood_filter( data_mapping = "edge centers" grid_dim = "n_edge" else: - raise ValueError( + raise DataCenteringError( f"neighborhood_filter requires data mapped to nodes, edges, or faces, " - f"but the last dimension {self.dims!r} does not match any grid dimension " + f"but the dimensions {self.dims!r} do not match any grid dimension " f"{GRID_DIMS}." ) diff --git a/uxarray/grid/neighbors.py b/uxarray/grid/neighbors.py index 7dfe42e00..4e0d4d800 100644 --- a/uxarray/grid/neighbors.py +++ b/uxarray/grid/neighbors.py @@ -1227,9 +1227,11 @@ def _neighborhood_filter( neighbor_indices = tree.query_radius(dest_coords, r=r) - # Initialize with NaN so that any element whose neighborhood is empty - # (e.g. an isolated point queried with a very small radius) yields NaN - # rather than uninitialized garbage memory. + # Allocate with NaN rather than ``np.empty`` purely as a defensive measure: + # if a neighborhood were ever empty, the result would be an obvious NaN + # instead of uninitialized garbage memory. In practice this cannot happen, + # since ``query_radius`` rejects negative ``r`` and every element is its own + # neighbor at distance 0, so even ``r = 0`` returns the original values. destination_data = np.full(data.shape, np.nan) # Apply func along the last (grid) axis only, so any extra leading From cb364fa06c19c3567c181210879d5d8b5332cf83 Mon Sep 17 00:00:00 2001 From: Rajeev Jain Date: Mon, 3 Aug 2026 09:37:26 -0500 Subject: [PATCH 17/21] Harden neighborhood filter tree selection and error handling - Include distance_metric in the get_ball_tree/get_kd_tree cache invalidation check. Only coordinates and coordinate_system were compared, so requesting a different metric silently returned a tree built with the original one. - Request a spherical/haversine tree explicitly in _neighborhood_filter. It previously read coordinate_system back off whatever tree was cached, so a cartesian tree from an earlier call made r a chord length instead of the documented great-circle degrees. - Raise an explanatory TypeError when func does not accept an axis keyword, instead of surfacing a raw NumPy message. - Document radius units, overlapping neighborhoods, and eager evaluation of dask-backed input in both public docstrings. - Restore the original _copy deep default; changing it was unrelated to this feature and affects every UxDataArray copy. Adds coverage for tree cache invalidation, the spherical-tree guarantee, the func-without-axis error, and dask input. --- test/core/test_dataarray.py | 44 ++++++++++++++++++++++++++++++++ test/grid/grid/test_neighbors.py | 40 +++++++++++++++++++++++++++++ uxarray/core/dataarray.py | 20 +++++++++++++-- uxarray/core/dataset.py | 7 +++++ uxarray/grid/grid.py | 8 ++++++ uxarray/grid/neighbors.py | 37 ++++++++++++++++++++++++--- 6 files changed, 150 insertions(+), 6 deletions(-) diff --git a/test/core/test_dataarray.py b/test/core/test_dataarray.py index 61c7858df..5177362dc 100644 --- a/test/core/test_dataarray.py +++ b/test/core/test_dataarray.py @@ -279,6 +279,50 @@ def test_radius_edge_cases_never_produce_nan(self): with pytest.raises(AssertionError): uxda.neighborhood_filter(func=np.mean, r=-1.0) + def test_func_without_axis_raises_helpful_error(self): + """A ``func`` that does not accept ``axis`` should raise a TypeError + that explains the requirement rather than a raw NumPy message.""" + uxgrid = ux.Grid.from_healpix(zoom=1) + uxda = UxDataArray( + np.arange(uxgrid.n_face, dtype=float), dims=["n_face"], uxgrid=uxgrid + ) + + with pytest.raises(TypeError, match="must accept an `axis` keyword"): + uxda.neighborhood_filter(func=sum, r=5.0) + + def test_uses_spherical_tree_regardless_of_cached_tree(self): + """``r`` is documented in great-circle degrees, so the filter must build + a spherical/haversine tree even if a cartesian one was cached first.""" + uxgrid = ux.Grid.from_healpix(zoom=1) + data = np.arange(uxgrid.n_face, dtype=float) + uxda = UxDataArray(data, dims=["n_face"], uxgrid=uxgrid, name="face_var") + + expected = uxda.neighborhood_filter(func=np.mean, r=20.0).values + + # Prime the cache with a cartesian tree, then filter again + uxgrid.get_ball_tree( + coordinates="face centers", + coordinate_system="cartesian", + distance_metric="euclidean", + ) + np.testing.assert_allclose( + uxda.neighborhood_filter(func=np.mean, r=20.0).values, expected + ) + + def test_dask_input_returns_numpy(self, gridpath, datasetpath): + """Lazy input is computed eagerly; the result is NumPy-backed.""" + uxds = ux.open_dataset( + gridpath("ugrid", "outCSne30", "outCSne30.ug"), + datasetpath("ugrid", "outCSne30", "outCSne30_vortex.nc"), + chunks={"n_face": 1000}, + ) + uxda = uxds["psi"] + assert uxda.chunks is not None + + filtered = uxda.neighborhood_filter(func=np.mean, r=2.0) + assert filtered.chunks is None + assert isinstance(filtered.data, np.ndarray) + def test_auto_transpose_direct_on_uxdataarray(self, gridpath, datasetpath): """Calling neighborhood_filter directly on a (time, n_face) UxDataArray (without going through UxDataset) should preserve the original dim order.""" diff --git a/test/grid/grid/test_neighbors.py b/test/grid/grid/test_neighbors.py index 482833647..e67cb6840 100644 --- a/test/grid/grid/test_neighbors.py +++ b/test/grid/grid/test_neighbors.py @@ -190,3 +190,43 @@ def test_construct_edge_face_distances(gridpath): # Run the function under test calculated = _construct_edge_face_distances(face_lon, face_lat, edge_faces) np.testing.assert_array_almost_equal(calculated, expected, decimal=5) + + +def test_tree_cache_invalidated_on_parameter_change(gridpath): + """``get_ball_tree``/``get_kd_tree`` must rebuild when any tree-defining + parameter changes, not just ``coordinates``. Previously a cached tree was + returned with the original ``coordinate_system``/``distance_metric``.""" + uxgrid = ux.open_grid(gridpath("mpas", "QU", "mesh.QU.1920km.151026.nc")) + + spherical = uxgrid.get_ball_tree( + coordinates="face centers", + coordinate_system="spherical", + distance_metric="haversine", + ) + assert spherical.coordinate_system == "spherical" + + cartesian = uxgrid.get_ball_tree( + coordinates="face centers", + coordinate_system="cartesian", + distance_metric="euclidean", + ) + assert cartesian.coordinate_system == "cartesian" + assert cartesian.distance_metric == "euclidean" + + # switching only the distance metric must also rebuild + minkowski = uxgrid.get_ball_tree( + coordinates="face centers", + coordinate_system="cartesian", + distance_metric="minkowski", + ) + assert minkowski.distance_metric == "minkowski" + + # same for the KDTree + kd_cart = uxgrid.get_kd_tree( + coordinates="face centers", coordinate_system="cartesian" + ) + assert kd_cart.coordinate_system == "cartesian" + kd_sph = uxgrid.get_kd_tree( + coordinates="face centers", coordinate_system="spherical" + ) + assert kd_sph.coordinate_system == "spherical" diff --git a/uxarray/core/dataarray.py b/uxarray/core/dataarray.py index 3239a30c5..c77fbf628 100644 --- a/uxarray/core/dataarray.py +++ b/uxarray/core/dataarray.py @@ -121,8 +121,7 @@ def _copy(self, **kwargs): ``uxarray.UxDataArray``.""" copied = super()._copy(**kwargs) - # Match xarray's default: deep=True when not specified - deep = kwargs.get("deep", True) + deep = kwargs.get("deep", None) if deep: # Reinitialize the uxgrid assessor @@ -2220,6 +2219,23 @@ def neighborhood_filter( uxda_filter : UxDataArray Filtered data. + Raises + ------ + DataCenteringError (subclass of ValueError) + If the data is not mapped to nodes, edges, or faces. + TypeError + If ``func`` does not accept an ``axis`` keyword argument. + + Notes + ----- + ``r`` is a great-circle distance in degrees. Neighborhoods overlap, and + every element is its own neighbor at distance 0, so ``r = 0`` returns + the data unchanged and the result never contains spurious ``NaN``. + + The query requires random access across the whole grid dimension, so + lazy (dask-backed) data is computed eagerly and the result is always + NumPy-backed. + Examples -------- Apply a mean filter with a 5-degree radius: diff --git a/uxarray/core/dataset.py b/uxarray/core/dataset.py index 7f9ae24bd..382c0bcfb 100644 --- a/uxarray/core/dataset.py +++ b/uxarray/core/dataset.py @@ -681,6 +681,13 @@ def neighborhood_filter( destination_uxds : UxDataset Filtered dataset. + Notes + ----- + Variables without a grid dimension are passed through unchanged. + ``r`` is a great-circle distance in degrees, and lazy (dask-backed) + variables are computed eagerly. See + :meth:`UxDataArray.neighborhood_filter` for details. + Examples -------- Apply a mean filter to all grid-mapped variables in a dataset: diff --git a/uxarray/grid/grid.py b/uxarray/grid/grid.py index 49b93d745..6ba2d7baa 100644 --- a/uxarray/grid/grid.py +++ b/uxarray/grid/grid.py @@ -1777,10 +1777,15 @@ def get_ball_tree( BallTree instance """ + # Rebuild whenever any tree-defining parameter differs from the cached + # instance. Previously only ``coordinates`` was compared, so switching + # ``coordinate_system`` or ``distance_metric`` silently returned a stale + # tree built with the original settings. if ( self._ball_tree is None or coordinates != self._ball_tree._coordinates or coordinate_system != self._ball_tree.coordinate_system + or distance_metric != self._ball_tree.distance_metric or reconstruct ): self._ball_tree = BallTree( @@ -1879,10 +1884,13 @@ def get_kd_tree( KDTree instance """ + # Rebuild whenever any tree-defining parameter differs from the cached + # instance (see ``get_ball_tree`` for details). if ( self._kd_tree is None or coordinates != self._kd_tree._coordinates or coordinate_system != self._kd_tree.coordinate_system + or distance_metric != self._kd_tree.distance_metric or reconstruct ): self._kd_tree = KDTree( diff --git a/uxarray/grid/neighbors.py b/uxarray/grid/neighbors.py index 4e0d4d800..0ce83bc5f 100644 --- a/uxarray/grid/neighbors.py +++ b/uxarray/grid/neighbors.py @@ -1217,11 +1217,29 @@ def _neighborhood_filter( ------- destination_data : np.ndarray Filtered data, matching the shape of ``data``. - """ - tree = grid.get_ball_tree(coordinates=data_mapping) + Raises + ------ + TypeError + If ``func`` does not accept an ``axis`` keyword argument. + + Notes + ----- + The neighborhood query requires random access across the whole grid + dimension, so lazy (dask-backed) input is computed eagerly and the result + is always a NumPy array. + """ - coordinate_system = tree.coordinate_system + # Request a spherical/haversine tree explicitly rather than relying on the + # defaults. Without this, a cartesian tree cached by an earlier call would + # be reused and ``r`` would be silently interpreted as a chord length + # instead of the great-circle degrees documented above. + coordinate_system = "spherical" + tree = grid.get_ball_tree( + coordinates=data_mapping, + coordinate_system=coordinate_system, + distance_metric="haversine", + ) dest_coords = _get_element_coords(grid, data_mapping, coordinate_system) @@ -1238,6 +1256,17 @@ def _neighborhood_filter( # dimensions (e.g. time) are preserved rather than being collapsed. for i, idx in enumerate(neighbor_indices): if len(idx): - destination_data[..., i] = func(data[..., idx], axis=-1) + try: + destination_data[..., i] = func(data[..., idx], axis=-1) + except TypeError as exc: + if "axis" not in str(exc): + raise + raise TypeError( + f"`func` must accept an `axis` keyword argument so that the " + f"reduction is applied over the neighborhood only, but " + f"{getattr(func, '__name__', func)!r} does not. Use a NumPy " + f"reduction such as `np.mean` or `np.median`, or wrap your " + f"function with `functools.partial` to supply `axis`." + ) from exc return destination_data From b9642609515019d47ef4e31254780cdf29c1ba70 Mon Sep 17 00:00:00 2001 From: cmdupuis3 Date: Wed, 5 Aug 2026 13:40:08 -0500 Subject: [PATCH 18/21] 941: gufunc for neighborhood filter --- test/core/test_dataarray.py | 119 ++++++++++++- uxarray/core/dataarray.py | 55 +++--- uxarray/core/dataset.py | 14 +- uxarray/grid/neighbors.py | 327 ++++++++++++++++++++++++++++++------ 4 files changed, 422 insertions(+), 93 deletions(-) diff --git a/test/core/test_dataarray.py b/test/core/test_dataarray.py index 5177362dc..6b09afcad 100644 --- a/test/core/test_dataarray.py +++ b/test/core/test_dataarray.py @@ -1,3 +1,4 @@ +import warnings import numpy as np import uxarray as ux from uxarray.errors import DataCenteringError, DimensionError @@ -309,19 +310,123 @@ def test_uses_spherical_tree_regardless_of_cached_tree(self): uxda.neighborhood_filter(func=np.mean, r=20.0).values, expected ) - def test_dask_input_returns_numpy(self, gridpath, datasetpath): - """Lazy input is computed eagerly; the result is NumPy-backed.""" + def test_dask_input_stays_lazy(self, gridpath, datasetpath): + """Lazy input stays lazy, and the result matches the eager path.""" + uxds = ux.open_dataset( + gridpath("ugrid", "outCSne30", "outCSne30.ug"), + datasetpath("ugrid", "outCSne30", "outCSne30_vortex.nc"), + ) + eager = uxds["psi"].neighborhood_filter(func=np.mean, r=2.0) + + lazy_uxda = uxds["psi"].chunk({"n_face": -1}) + assert lazy_uxda.chunks is not None + + filtered = lazy_uxda.neighborhood_filter(func=np.mean, r=2.0) + assert filtered.chunks is not None, "the filter should not force a compute" + assert isinstance(filtered, UxDataArray) + assert filtered.uxgrid == lazy_uxda.uxgrid + np.testing.assert_allclose(filtered.compute().values, eager.values) + + def test_grid_dim_chunks_are_collapsed_with_warning(self, gridpath, datasetpath): + """A neighborhood may span the whole grid, so the grid dimension cannot + stay chunked. Collapsing it is a memory decision the user made, so it + is not done silently.""" + uxds = ux.open_dataset( + gridpath("ugrid", "outCSne30", "outCSne30.ug"), + datasetpath("ugrid", "outCSne30", "outCSne30_vortex.nc"), + ) + expected = uxds["psi"].neighborhood_filter(func=np.mean, r=2.0).values + + uxda = uxds["psi"].chunk({"n_face": 1000}) + assert len(uxda.chunksizes["n_face"]) > 1 + + with pytest.warns(UserWarning, match="Rechunking 'n_face'"): + filtered = uxda.neighborhood_filter(func=np.mean, r=2.0) + + assert filtered.chunksizes["n_face"] == (uxda.sizes["n_face"],) + np.testing.assert_allclose(filtered.compute().values, expected) + + def test_chunked_over_time_is_not_rechunked(self, gridpath, datasetpath): + """Chunking a non-grid dimension is the supported case and should pass + through untouched, with no warning.""" + uxds = ux.open_dataset( + gridpath("ugrid", "outCSne30", "outCSne30.ug"), + datasetpath("ugrid", "outCSne30", "outCSne30_vortex.nc"), + ) + psi = uxds["psi"] + stacked = UxDataArray( + np.tile(psi.values, (6, 1)), + dims=["time", "n_face"], + uxgrid=psi.uxgrid, + name="psi", + ).chunk({"time": 2, "n_face": -1}) + + with warnings.catch_warnings(): + warnings.simplefilter("error", UserWarning) + filtered = stacked.neighborhood_filter(func=np.mean, r=2.0) + + assert filtered.chunksizes["time"] == (2, 2, 2) + expected = psi.neighborhood_filter(func=np.mean, r=2.0).values + np.testing.assert_allclose(filtered.compute().values, np.tile(expected, (6, 1))) + + @pytest.mark.parametrize( + "func", [np.mean, np.sum, np.min, np.max, np.median, np.amin, np.amax] + ) + def test_kernel_matches_generic_path(self, func, gridpath, datasetpath): + """Every reduction with a compiled kernel must agree with the generic + loop it bypasses.""" + from uxarray.grid.neighbors import ( + _NEIGHBORHOOD_KERNELS, + _csr_neighbors, + _neighborhood_reduce, + ) + uxds = ux.open_dataset( gridpath("ugrid", "outCSne30", "outCSne30.ug"), datasetpath("ugrid", "outCSne30", "outCSne30_vortex.nc"), - chunks={"n_face": 1000}, ) uxda = uxds["psi"] - assert uxda.chunks is not None + assert func in _NEIGHBORHOOD_KERNELS, "expected a compiled kernel" + + flat, starts, counts = _csr_neighbors(uxda.uxgrid, "face centers", 3.0) + # 2-D as well as 1-D, to exercise the gufunc's broadcast loop + block = np.vstack([uxda.values, uxda.values * -2.0]) + expected = _neighborhood_reduce(block, flat, starts, counts, func) - filtered = uxda.neighborhood_filter(func=np.mean, r=2.0) - assert filtered.chunks is None - assert isinstance(filtered.data, np.ndarray) + filtered = uxda.neighborhood_filter(func=func, r=3.0) + np.testing.assert_allclose(filtered.values, expected[0], rtol=1e-12) + + kernel_2d = _NEIGHBORHOOD_KERNELS[func](block, flat, starts, counts) + np.testing.assert_allclose(kernel_2d, expected, rtol=1e-12) + + def test_float32_input(self, gridpath, datasetpath): + """float32 fields take the compiled kernel's float32 signature and + still produce float64 output, as the generic path does.""" + uxds = ux.open_dataset( + gridpath("ugrid", "outCSne30", "outCSne30.ug"), + datasetpath("ugrid", "outCSne30", "outCSne30_vortex.nc"), + ) + uxda = uxds["psi"] + # `.astype` on a UxDataArray returns a plain DataArray, so rebuild + uxda32 = UxDataArray( + uxda.values.astype(np.float32), dims=uxda.dims, uxgrid=uxda.uxgrid + ) + filtered32 = uxda32.neighborhood_filter(func=np.mean, r=2.0) + filtered64 = uxda.neighborhood_filter(func=np.mean, r=2.0) + + assert filtered32.dtype == np.float64 + np.testing.assert_allclose(filtered32.values, filtered64.values, rtol=1e-6) + + def test_integer_input_is_promoted(self): + """Integer fields have no kernel signature and must be promoted rather + than raising.""" + uxgrid = ux.Grid.from_healpix(zoom=1) + data = np.arange(uxgrid.n_face) + uxda = UxDataArray(data, dims=["n_face"], uxgrid=uxgrid, name="int_var") + + filtered = uxda.neighborhood_filter(func=np.mean, r=0.0) + assert filtered.dtype == np.float64 + np.testing.assert_allclose(filtered.values, data) def test_auto_transpose_direct_on_uxdataarray(self, gridpath, datasetpath): """Calling neighborhood_filter directly on a (time, n_face) UxDataArray diff --git a/uxarray/core/dataarray.py b/uxarray/core/dataarray.py index 466774f46..d95f57ad6 100644 --- a/uxarray/core/dataarray.py +++ b/uxarray/core/dataarray.py @@ -2193,24 +2193,27 @@ def neighborhood_filter( Parameters ---------- func: Callable, default=np.mean - Apply this function to neighborhood. Must accept an ``axis`` keyword - argument (as ``np.mean``, ``np.median``, and similar NumPy reductions - do). Use ``functools.partial`` to supply additional arguments, e.g. - ``functools.partial(np.percentile, q=90)``. + Apply this function to neighborhood. ``np.mean``, ``np.sum``, + ``np.min``, ``np.max`` and ``np.median`` use a compiled kernel. + Any other function must accept an ``axis`` keyword argument (as + NumPy reductions do); use ``functools.partial`` to supply + additional arguments, e.g. ``functools.partial(np.percentile, + q=90)``. r : float, default=1. Radius of the neighborhood, in degrees. Returns ------- uxda_filter : UxDataArray - Filtered data. + Filtered data, as float64. Raises ------ DataCenteringError (subclass of ValueError) If the data is not mapped to nodes, edges, or faces. TypeError - If ``func`` does not accept an ``axis`` keyword argument. + If ``func`` has no compiled kernel and does not accept an ``axis`` + keyword argument. Notes ----- @@ -2218,9 +2221,10 @@ def neighborhood_filter( every element is its own neighbor at distance 0, so ``r = 0`` returns the data unchanged and the result never contains spurious ``NaN``. - The query requires random access across the whole grid dimension, so - lazy (dask-backed) data is computed eagerly and the result is always - NumPy-backed. + A neighborhood may span the whole grid, so the grid dimension cannot be + chunked; it is collapsed to a single chunk (with a warning) for + dask-backed data. The remaining dimensions stay chunked and lazy, so + chunk along ``time`` rather than the grid dimension. Examples -------- @@ -2260,28 +2264,23 @@ def neighborhood_filter( f"{GRID_DIMS}." ) - # Ensure the grid dimension is the last axis, transposing if necessary. - # This mirrors the behaviour of UxDataset.neighborhood_filter so that - # calling the method directly on a (time, n_face) UxDataArray works. - needs_transpose = self.dims[-1] != grid_dim - uxda_work = self.transpose(..., grid_dim) if needs_transpose else self - - destination_data = _neighborhood_filter( - self.uxgrid, uxda_work.data, data_mapping, func=func, r=r + # ``_neighborhood_filter`` declares the grid dimension as a core + # dimension, so it moves that dimension into place itself; no manual + # transpose is needed on the way in. + filtered = _neighborhood_filter( + self.uxgrid, self, data_mapping, grid_dim, func=func, r=r ) - # Construct UxDataArray for filtered variable, reusing metadata - # (name, coords, attrs, uxgrid) from the working copy. - # deep=False keeps a reference to the same uxgrid (the filtered data - # lives on the identical grid topology) and avoids a redundant deep - # copy of the now-discarded original data array. - uxda_filter = uxda_work._copy(data=destination_data, deep=False) - - # Restore original dimension order if we transposed. - if needs_transpose: - uxda_filter = uxda_filter.transpose(*self.dims) + # Core dimensions come back appended last, so restore the caller's + # dimension order when it differed. + if filtered.dims != self.dims: + filtered = filtered.transpose(*self.dims) - return uxda_filter + # ``apply_ufunc`` returns a plain xr.DataArray, dropping the subclass + # and its grid. Name, coords and attrs are carried through already. + # The filtered data lives on the identical grid topology, so the same + # uxgrid is reattached rather than copied. + return UxDataArray(filtered, uxgrid=self.uxgrid) def __getattribute__(self, name): """Intercept accessor method calls to return Ux-aware accessors.""" diff --git a/uxarray/core/dataset.py b/uxarray/core/dataset.py index 382c0bcfb..2feb30eaf 100644 --- a/uxarray/core/dataset.py +++ b/uxarray/core/dataset.py @@ -669,10 +669,12 @@ def neighborhood_filter( Parameters ---------- func: Callable, default=np.mean - Apply this function to neighborhood. Must accept an ``axis`` keyword - argument (as ``np.mean``, ``np.median``, and similar NumPy reductions - do). Use ``functools.partial`` to supply additional arguments, e.g. - ``functools.partial(np.percentile, q=90)``. + Apply this function to neighborhood. ``np.mean``, ``np.sum``, + ``np.min``, ``np.max`` and ``np.median`` use a compiled kernel. + Any other function must accept an ``axis`` keyword argument (as + NumPy reductions do); use ``functools.partial`` to supply + additional arguments, e.g. ``functools.partial(np.percentile, + q=90)``. r : float, default=1. Radius of the neighborhood, in degrees. @@ -685,8 +687,8 @@ def neighborhood_filter( ----- Variables without a grid dimension are passed through unchanged. ``r`` is a great-circle distance in degrees, and lazy (dask-backed) - variables are computed eagerly. See - :meth:`UxDataArray.neighborhood_filter` for details. + variables stay lazy. See :meth:`UxDataArray.neighborhood_filter` for + details. Examples -------- diff --git a/uxarray/grid/neighbors.py b/uxarray/grid/neighbors.py index 0ce83bc5f..6b7d72c82 100644 --- a/uxarray/grid/neighbors.py +++ b/uxarray/grid/neighbors.py @@ -1,8 +1,9 @@ +import warnings from typing import Callable import numpy as np import xarray as xr -from numba import njit +from numba import guvectorize, njit from numpy import deg2rad from uxarray.constants import ERROR_TOLERANCE, INT_DTYPE, INT_FILL_VALUE @@ -1184,10 +1185,234 @@ def _get_element_coords(grid, data_mapping: str, coordinate_system: str): ) +# A neighborhood reduction is a segmented reduction over a ragged (CSR-like) +# neighbor structure: elementwise in every dimension except the grid axis, +# which it reduces over. That is exactly a generalized ufunc signature, so the +# kernels below declare the grid axis as a core dimension. Two consequences +# fall out of stating it that way: +# +# * dask can parallelize over the remaining (chunked) dimensions on its own, +# so the filter stays lazy instead of materializing the whole array, and +# * the grid axis is a *core* dimension, so dask refuses to split it rather +# than silently handing a kernel a block the neighbor indices overrun. +# +# ``(n)`` is the source grid axis, ``(k)`` the flattened neighbor index array, +# and ``(m)`` the destination axis. Output is float64 regardless of input +# dtype, matching the behaviour of the generic path below. +_GUFUNC_SIGNATURES = [ + "void(float64[:], int64[:], int64[:], int64[:], float64[:])", + "void(float32[:], int64[:], int64[:], int64[:], float64[:])", +] +_GUFUNC_LAYOUT = "(n),(k),(m),(m)->(m)" +_GUFUNC_KWARGS = {"nopython": True, "cache": True, "target": "parallel"} + + +@guvectorize(_GUFUNC_SIGNATURES, _GUFUNC_LAYOUT, **_GUFUNC_KWARGS) +def _gu_mean(data, flat, starts, counts, out): + for i in range(starts.shape[0]): + count = counts[i] + if count == 0: + out[i] = np.nan + continue + start = starts[i] + acc = 0.0 + for j in range(start, start + count): + acc += data[flat[j]] + out[i] = acc / count + + +@guvectorize(_GUFUNC_SIGNATURES, _GUFUNC_LAYOUT, **_GUFUNC_KWARGS) +def _gu_sum(data, flat, starts, counts, out): + for i in range(starts.shape[0]): + start = starts[i] + acc = 0.0 + for j in range(start, start + counts[i]): + acc += data[flat[j]] + out[i] = acc + + +@guvectorize(_GUFUNC_SIGNATURES, _GUFUNC_LAYOUT, **_GUFUNC_KWARGS) +def _gu_max(data, flat, starts, counts, out): + for i in range(starts.shape[0]): + count = counts[i] + if count == 0: + out[i] = np.nan + continue + start = starts[i] + best = data[flat[start]] + for j in range(start + 1, start + count): + value = data[flat[j]] + if value > best: + best = value + out[i] = best + + +@guvectorize(_GUFUNC_SIGNATURES, _GUFUNC_LAYOUT, **_GUFUNC_KWARGS) +def _gu_min(data, flat, starts, counts, out): + for i in range(starts.shape[0]): + count = counts[i] + if count == 0: + out[i] = np.nan + continue + start = starts[i] + best = data[flat[start]] + for j in range(start + 1, start + count): + value = data[flat[j]] + if value < best: + best = value + out[i] = best + + +@guvectorize(_GUFUNC_SIGNATURES, _GUFUNC_LAYOUT, **_GUFUNC_KWARGS) +def _gu_median(data, flat, starts, counts, out): + # A median needs the neighborhood materialized so it can be sorted, which + # is why ``np.ufunc.reduceat`` cannot express it but a kernel can: the + # scratch buffer is sized once to the largest neighborhood and reused. + widest = 0 + for i in range(counts.shape[0]): + if counts[i] > widest: + widest = counts[i] + buffer = np.empty(widest, dtype=np.float64) + + for i in range(starts.shape[0]): + count = counts[i] + if count == 0: + out[i] = np.nan + continue + start = starts[i] + for j in range(count): + buffer[j] = data[flat[start + j]] + window = np.sort(buffer[:count]) + if count % 2: + out[i] = window[count // 2] + else: + out[i] = 0.5 * (window[count // 2 - 1] + window[count // 2]) + + +# Reductions with a compiled kernel. Anything else (``np.median`` with a +# custom axis, ``functools.partial(np.percentile, q=90)``, a user's own +# callable) falls back to the generic path, which is slower but accepts any +# function taking an ``axis`` keyword. +_NEIGHBORHOOD_KERNELS = { + np.mean: _gu_mean, + np.sum: _gu_sum, + np.max: _gu_max, + np.amax: _gu_max, + np.min: _gu_min, + np.amin: _gu_min, + np.median: _gu_median, +} + + +def _csr_neighbors(grid, data_mapping: str, r: float): + """Queries the neighborhood of every element and returns it in CSR form. + + ``query_radius`` returns a ragged sequence of index arrays, one per + element. Flattening it into ``(flat, starts, counts)`` gives the kernels a + layout they can walk without allocating per-neighborhood temporaries. + + Returns + ------- + flat : np.ndarray + Concatenated neighbor indices for every element. + starts : np.ndarray + Offset into ``flat`` at which each element's neighbors begin. + counts : np.ndarray + Number of neighbors of each element. + """ + # Request a spherical/haversine tree explicitly rather than relying on the + # defaults. Without this, a cartesian tree cached by an earlier call would + # be reused and ``r`` would be silently interpreted as a chord length + # instead of the great-circle degrees documented by the callers. + coordinate_system = "spherical" + tree = grid.get_ball_tree( + coordinates=data_mapping, + coordinate_system=coordinate_system, + distance_metric="haversine", + ) + + dest_coords = _get_element_coords(grid, data_mapping, coordinate_system) + neighbor_indices = tree.query_radius(dest_coords, r=r) + + # ``query_radius`` unwraps its result for a single query point, which a + # one-element grid would hit. + if isinstance(neighbor_indices, np.ndarray): + neighbor_indices = [neighbor_indices] + + counts = np.fromiter( + map(len, neighbor_indices), dtype=np.int64, count=len(neighbor_indices) + ) + starts = np.zeros(counts.size, dtype=np.int64) + np.cumsum(counts[:-1], out=starts[1:]) + flat = np.concatenate(neighbor_indices).astype(np.int64, copy=False) + + return flat, starts, counts + + +def _neighborhood_reduce(block, flat, starts, counts, func: Callable): + """Generic fallback: applies ``func`` to each neighborhood in turn. + + Used when ``func`` has no compiled kernel. ``block`` is a NumPy array with + the grid dimension last. + """ + destination_data = np.full(block.shape, np.nan) + + # The `axis` check lives outside the loop: whether `func` accepts the + # keyword cannot change between iterations, so validating it once is + # equivalent to validating it every time and leaves the loop body bare. + try: + for i in range(starts.shape[0]): + idx = flat[starts[i] : starts[i] + counts[i]] + # Apply func along the last (grid) axis only, so any extra leading + # dimensions (e.g. time) are preserved rather than being collapsed. + destination_data[..., i] = func(block[..., idx], axis=-1) + except TypeError as exc: + if "axis" not in str(exc): + raise + raise TypeError( + f"`func` must accept an `axis` keyword argument so that the " + f"reduction is applied over the neighborhood only, but " + f"{getattr(func, '__name__', func)!r} does not. Use a NumPy " + f"reduction such as `np.mean` or `np.median`, or wrap your " + f"function with `functools.partial` to supply `axis`." + ) from exc + + return destination_data + + +def _rechunk_grid_dim(uxda, grid_dim: str): + """Collapses the grid dimension to a single chunk, warning if that changes + the user's chunking. + + Neighborhoods are global — an element near a chunk boundary draws on + elements in other chunks — so the grid dimension cannot be chunked. This is + done explicitly rather than through ``allow_rechunk``, which would do it + silently and also disable ``apply_gufunc``'s other consistency checks. + """ + if uxda.chunks is None: + return uxda + + grid_chunks = uxda.chunksizes.get(grid_dim, ()) + if len(grid_chunks) <= 1: + return uxda + + warnings.warn( + f"Rechunking {grid_dim!r} from {len(grid_chunks)} chunks into one, as a " + f"neighborhood may span the whole grid. Each task will hold " + f"{uxda.sizes[grid_dim]} elements along {grid_dim!r}; chunk the " + f"non-grid dimensions instead to bound memory use.", + UserWarning, + stacklevel=3, + ) + + return uxda.chunk({grid_dim: -1}) + + def _neighborhood_filter( grid, - data: np.ndarray, + uxda, data_mapping: str, + grid_dim: str, func: Callable = np.mean, r: float = 1.0, ): @@ -1199,74 +1424,72 @@ def _neighborhood_filter( grid : Grid Source grid used to construct the ``BallTree`` used for the neighborhood queries. - data : np.ndarray - Data to filter. The grid dimension (``n_node``, ``n_edge``, or - ``n_face``) is expected to be the last axis. + uxda : xr.DataArray + Data to filter. The grid dimension may sit at any position. data_mapping : str One of "nodes", "edge centers", or "face centers", identifying which - grid element ``data`` is mapped to. + grid element ``uxda`` is mapped to. + grid_dim : str + Name of the grid dimension (``n_node``, ``n_edge``, or ``n_face``). func : Callable, default=np.mean - Function applied to the values found in each neighborhood. Must + Function applied to the values found in each neighborhood. Reductions + with a compiled kernel (``np.mean``, ``np.sum``, ``np.min``, + ``np.max``, ``np.median``) take a fast path. Any other function must accept an ``axis`` keyword argument (as ``np.mean``, ``np.median``, - and similar NumPy reductions do) so that any extra, non-grid - dimensions (e.g. ``time``) are preserved rather than being collapsed. + and similar NumPy reductions do) so that extra, non-grid dimensions + (e.g. ``time``) are preserved rather than being collapsed. r : float, default=1. Radius of the neighborhood, in degrees. Returns ------- - destination_data : np.ndarray - Filtered data, matching the shape of ``data``. + filtered : xr.DataArray + Filtered data, float64, with the grid dimension moved last. Lazy if + the input was lazy. Raises ------ TypeError - If ``func`` does not accept an ``axis`` keyword argument. + If ``func`` has no compiled kernel and does not accept an ``axis`` + keyword argument. Notes ----- - The neighborhood query requires random access across the whole grid - dimension, so lazy (dask-backed) input is computed eagerly and the result - is always a NumPy array. + The grid dimension is a core dimension of the reduction, so it is + collapsed to a single chunk for dask-backed input; the remaining + dimensions stay chunked and are evaluated lazily. """ - # Request a spherical/haversine tree explicitly rather than relying on the - # defaults. Without this, a cartesian tree cached by an earlier call would - # be reused and ``r`` would be silently interpreted as a chord length - # instead of the great-circle degrees documented above. - coordinate_system = "spherical" - tree = grid.get_ball_tree( - coordinates=data_mapping, - coordinate_system=coordinate_system, - distance_metric="haversine", - ) - - dest_coords = _get_element_coords(grid, data_mapping, coordinate_system) + flat, starts, counts = _csr_neighbors(grid, data_mapping, r) + kernel = _NEIGHBORHOOD_KERNELS.get(func) - neighbor_indices = tree.query_radius(dest_coords, r=r) + if kernel is not None: - # Allocate with NaN rather than ``np.empty`` purely as a defensive measure: - # if a neighborhood were ever empty, the result would be an obvious NaN - # instead of uninitialized garbage memory. In practice this cannot happen, - # since ``query_radius`` rejects negative ``r`` and every element is its own - # neighbor at distance 0, so even ``r = 0`` returns the original values. - destination_data = np.full(data.shape, np.nan) - - # Apply func along the last (grid) axis only, so any extra leading - # dimensions (e.g. time) are preserved rather than being collapsed. - for i, idx in enumerate(neighbor_indices): - if len(idx): - try: - destination_data[..., i] = func(data[..., idx], axis=-1) - except TypeError as exc: - if "axis" not in str(exc): - raise - raise TypeError( - f"`func` must accept an `axis` keyword argument so that the " - f"reduction is applied over the neighborhood only, but " - f"{getattr(func, '__name__', func)!r} does not. Use a NumPy " - f"reduction such as `np.mean` or `np.median`, or wrap your " - f"function with `functools.partial` to supply `axis`." - ) from exc + def _apply(block): + # The kernels are compiled for float32/float64 only; anything else + # (integer fields, say) is promoted, which the generic path does + # too by writing into a float64 output. + if block.dtype not in (np.float64, np.float32): + block = block.astype(np.float64) + return kernel(block, flat, starts, counts) + else: - return destination_data + def _apply(block): + return _neighborhood_reduce(block, flat, starts, counts, func) + + uxda = _rechunk_grid_dim(uxda, grid_dim) + + # ``apply_ufunc`` moves the grid dimension last before calling ``_apply`` + # and, for dask-backed input, hands each chunk over as a materialized + # NumPy block. Indexing the array one destination element at a time (as + # this function used to) would instead trigger one graph execution per + # grid element. + return xr.apply_ufunc( + _apply, + uxda, + input_core_dims=[[grid_dim]], + output_core_dims=[[grid_dim]], + dask="parallelized", + output_dtypes=[np.float64], + keep_attrs=True, + ) From b6131856be9424f30ea6c8ecd592cee74b0d5e03 Mon Sep 17 00:00:00 2001 From: cmdupuis3 Date: Wed, 5 Aug 2026 14:02:02 -0500 Subject: [PATCH 19/21] 941: name reductions, reuse neighborhoods Follows the gufunc work with the API changes it made possible. Kernel factory. The buffered kernels (median and friends) share one gather loop, with the reduction supplied as a numba-compilable callable. numba supports the NumPy reductions in nopython mode, so np.median is used directly rather than reimplemented -- it partitions instead of fully sorting, making it 1.6x faster than the hand-written sort it replaces, as well as shorter. Adds ptp, std, var, quantile and percentile; the parameterized ones carry their argument as a scalar core dimension so q and ddof vary per call without recompiling. Named reductions. `func` now takes a name ("mean", "quantile", ...) with parameters as ordinary keyword arguments. The previous `axis=-1` contract could not be jitted -- numba rejects the axis kwarg on mean/median/percentile -- and dispatching on a function object cannot see through functools.partial, so `partial(np.percentile, q=90)`, the example in our own docstring, could never reach a kernel. A name always can. Callables remain accepted on the old contract as an escape hatch, and np.mean and friends still map to their kernels so existing code keeps the fast path. Neighborhoods. Finding the neighbors costs more than reducing over them -- after the kernel work it is ~95% of a call -- and it was repeated on every call, including once per variable in the dataset path. Grid.neighborhoods() does the search once and returns an object to reduce over repeatedly: 3.5x for four reductions at one radius on a 196k-face grid, and a dataset filter now costs one query per grid location rather than one per variable (8 variables: 1.74s -> 0.23s). Co-Authored-By: Claude Opus 5 --- docs/api.rst | 12 + docs/user-guide/neighborhood-filter.ipynb | 180 +++++++- test/core/test_dataarray.py | 216 +++++++++- test/core/test_dataset.py | 75 ++++ uxarray/core/dataarray.py | 100 ++--- uxarray/core/dataset.py | 49 ++- uxarray/grid/grid.py | 37 ++ uxarray/grid/neighbors.py | 484 +++++++++++++++++----- 8 files changed, 940 insertions(+), 213 deletions(-) diff --git a/docs/api.rst b/docs/api.rst index 9508b8427..4666c0efa 100644 --- a/docs/api.rst +++ b/docs/api.rst @@ -194,6 +194,7 @@ Methods Grid.compute_face_areas Grid.construct_face_centers Grid.get_ball_tree + Grid.neighborhoods Grid.get_kd_tree Grid.get_spatial_hash Grid.get_faces_containing_point @@ -577,6 +578,17 @@ neighborhood of a specified radius around each grid element. UxDataArray.neighborhood_filter UxDataset.neighborhood_filter +Finding the neighbors is usually more expensive than reducing over them. To apply several +reductions at one radius, build the neighborhoods once and reduce over them repeatedly. + +.. autosummary:: + :toctree: generated/ + + Grid.neighborhoods + uxarray.grid.neighbors.Neighborhoods + uxarray.grid.neighbors.Neighborhoods.reduce + uxarray.grid.neighbors.Neighborhoods.n_neighbors + Zonal Average ~~~~~~~~~~~~~ diff --git a/docs/user-guide/neighborhood-filter.ipynb b/docs/user-guide/neighborhood-filter.ipynb index 4e2b7d980..8ab26d8dc 100644 --- a/docs/user-guide/neighborhood-filter.ipynb +++ b/docs/user-guide/neighborhood-filter.ipynb @@ -4,7 +4,29 @@ "cell_type": "markdown", "id": "7fb27b941602401d91542211134fc71a", "metadata": {}, - "source": "# Neighborhood Filter\n\nA **neighborhood filter** replaces the value at each grid element with the result\nof a user-specified function applied to all grid elements whose centers fall within\na circular neighborhood of radius `r` degrees around that element.\n\nUnlike a fixed *k*-nearest-neighbor average, a radius-based filter imposes a\nconsistent spatial scale across the whole mesh—useful for variable-resolution grids\nwhere the number of neighbors varies from region to region.\n\n**Supported element types:** face-centered, node-centered, and edge-centered data.\n\n**API at a glance:**\n\n| Object | Method |\n|---|---|\n| `UxDataArray` | `da.neighborhood_filter(func=np.mean, r=5.0)` |\n| `UxDataset` | `ds.neighborhood_filter(func=np.mean, r=5.0)` |\n\nThe returned object is always the same type as the input, with the same grid, dims,\ncoordinates, name, and attributes preserved.\n" + "source": [ + "# Neighborhood Filter\n", + "\n", + "A **neighborhood filter** replaces the value at each grid element with a\n", + "reduction of all grid elements whose centers fall within\n", + "a circular neighborhood of radius `r` degrees around that element.\n", + "\n", + "Unlike a fixed *k*-nearest-neighbor average, a radius-based filter imposes a\n", + "consistent spatial scale across the whole mesh—useful for variable-resolution grids\n", + "where the number of neighbors varies from region to region.\n", + "\n", + "**Supported element types:** face-centered, node-centered, and edge-centered data.\n", + "\n", + "**API at a glance:**\n", + "\n", + "| Object | Method |\n", + "|---|---|\n", + "| `UxDataArray` | `da.neighborhood_filter(\"mean\", r=5.0)` |\n", + "| `UxDataset` | `ds.neighborhood_filter(\"mean\", r=5.0)` |\n", + "\n", + "The returned object is always the same type as the input, with the same grid, dims,\n", + "coordinates, name, and attributes preserved.\n" + ] }, { "cell_type": "markdown", @@ -69,7 +91,12 @@ "cell_type": "markdown", "id": "8763a12b2bbd4a93a75aff182afb95dc", "metadata": {}, - "source": "## Basic Usage: Mean Filter\n\nCalling `neighborhood_filter` with `func=np.mean` and a radius of 5° replaces\neach face value with the mean of all face centers within 5° of that face's center.\n" + "source": [ + "## Basic Usage: Mean Filter\n", + "\n", + "Calling `neighborhood_filter` with `\"mean\"` and a radius of 5° replaces\n", + "each face value with the mean of all face centers within 5° of that face's center.\n" + ] }, { "cell_type": "code", @@ -78,7 +105,7 @@ "metadata": {}, "outputs": [], "source": [ - "uxda_smooth = uxda.neighborhood_filter(func=np.mean, r=5.0)\n", + "uxda_smooth = uxda.neighborhood_filter(\"mean\", r=5.0)\n", "uxda_smooth" ] }, @@ -121,7 +148,7 @@ "hv.extension(\"bokeh\")\n", "\n", "plots = [\n", - " uxda.neighborhood_filter(func=np.mean, r=r).plot.polygons(\n", + " uxda.neighborhood_filter(\"mean\", r=r).plot.polygons(\n", " cmap=\"RdBu_r\",\n", " title=f\"r = {r}°\",\n", " width=350,\n", @@ -138,7 +165,15 @@ "cell_type": "markdown", "id": "59bbdb311c014d738909a11f9e486628", "metadata": {}, - "source": "## Custom Functions via `functools.partial`\n\nAny callable that accepts an `axis` keyword argument (as NumPy reductions do) works\nas the filter function. Use `functools.partial` to fix additional keyword arguments.\n" + "source": [ + "## Other Reductions\n", + "\n", + "Pass the name of the reduction you want. The available names are `mean`,\n", + "`sum`, `min`, `max`, `median`, `ptp`, `std`, `var`, `quantile`, and\n", + "`percentile`. Reductions that take a parameter receive it as a keyword\n", + "argument: `q` for `quantile` (0–1) and `percentile` (0–100), `ddof` for\n", + "`std` and `var`.\n" + ] }, { "cell_type": "code", @@ -148,17 +183,21 @@ "outputs": [], "source": [ "# 90th-percentile filter — highlights local maxima\n", - "uxda_p90 = uxda.neighborhood_filter(func=partial(np.percentile, q=90), r=5.0)\n", + "uxda_p90 = uxda.neighborhood_filter(\"percentile\", r=5.0, q=90)\n", "\n", "# Maximum filter\n", - "uxda_max = uxda.neighborhood_filter(func=np.max, r=5.0)\n", + "uxda_max = uxda.neighborhood_filter(\"max\", r=5.0)\n", "\n", "# Median filter — robust to outliers\n", - "uxda_med = uxda.neighborhood_filter(func=np.median, r=5.0)\n", + "uxda_med = uxda.neighborhood_filter(\"median\", r=5.0)\n", + "\n", + "# Local spread, as a sample standard deviation\n", + "uxda_std = uxda.neighborhood_filter(\"std\", r=5.0, ddof=1)\n", "\n", - "print(\"max filter max :\", uxda_max.values.max())\n", - "print(\"p90 filter max :\", uxda_p90.values.max())\n", - "print(\"median filter max:\", uxda_med.values.max())" + "print(\"max filter max :\", uxda_max.values.max())\n", + "print(\"p90 filter max :\", uxda_p90.values.max())\n", + "print(\"median filter max:\", uxda_med.values.max())\n", + "print(\"std filter max :\", uxda_std.values.max())\n" ] }, { @@ -178,6 +217,93 @@ ").cols(2)" ] }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Reductions Without a Name\n", + "\n", + "If you need something not in that list, pass a callable instead. It is applied as\n", + "`func(values, axis=-1)` to a block whose last axis is the neighborhood, once per\n", + "grid element, in Python — noticeably slower than a named reduction, so reach for it\n", + "only when no name fits.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# a root-mean-square filter, which has no named equivalent\n", + "def rms(values, axis):\n", + " return np.sqrt(np.mean(values**2, axis=axis))\n", + "\n", + "uxda_rms = uxda.neighborhood_filter(rms, r=5.0)\n", + "\n", + "# `functools.partial` also works, though `\"percentile\"` is the faster way here\n", + "uxda_p90_slow = uxda.neighborhood_filter(partial(np.percentile, q=90), r=5.0)\n", + "print(\"partial matches the named reduction:\",\n", + " np.allclose(uxda_p90_slow.values, uxda_p90.values))\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Reusing Neighborhoods Across Reductions\n", + "\n", + "Each call to `neighborhood_filter` searches the grid for the neighbors of every\n", + "element. That search usually costs far more than the reduction itself, so\n", + "applying several reductions at the same radius repeats the expensive part.\n", + "\n", + "`Grid.neighborhoods` does the search once and returns an object you can reduce\n", + "over as many times as you like.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "nb5 = uxda.uxgrid.neighborhoods(r=5.0)\n", + "nb5\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# the search is already done; each of these only runs a reduction\n", + "smooth = nb5.reduce(uxda, \"mean\")\n", + "spread = nb5.reduce(uxda, \"std\")\n", + "p90 = nb5.reduce(uxda, \"percentile\", q=90)\n", + "\n", + "print(\"identical to the one-shot call:\", np.allclose(smooth.values, uxda.neighborhood_filter(\"mean\", r=5.0).values))\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "`n_neighbors` reports how many elements fell inside each neighborhood. On a\n", + "variable-resolution mesh this varies by region, which is worth checking before\n", + "reading too much into a filtered field.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "counts = nb5.n_neighbors\n", + "print(\"neighbors per face: min\", int(counts.min()), \" max\", int(counts.max()), \" mean\", float(counts.mean()).__round__(1))\n" + ] + }, { "cell_type": "markdown", "id": "c3933fab20d04ec698c2621248eb3be0", @@ -202,7 +328,7 @@ " attrs={\"units\": \"degrees_east\"},\n", ")\n", "\n", - "filtered_node = node_da.neighborhood_filter(func=np.mean, r=10.0)\n", + "filtered_node = node_da.neighborhood_filter(\"mean\", r=10.0)\n", "print(\"node input dims:\", node_da.dims, \" shape:\", node_da.shape)\n", "print(\"node output dims:\", filtered_node.dims, \" shape:\", filtered_node.shape)\n", "print(\"attrs preserved:\", filtered_node.attrs)" @@ -224,7 +350,7 @@ " name=\"edge_noise\",\n", ")\n", "\n", - "filtered_edge = edge_da.neighborhood_filter(func=np.mean, r=10.0)\n", + "filtered_edge = edge_da.neighborhood_filter(\"mean\", r=10.0)\n", "print(\"edge input dims:\", edge_da.dims, \" shape:\", edge_da.shape)\n", "print(\"edge output dims:\", filtered_edge.dims, \" shape:\", filtered_edge.shape)" ] @@ -247,7 +373,7 @@ "\n", "print(\"Input dims:\", uxda_ts.dims, \" shape:\", uxda_ts.shape)\n", "\n", - "filtered_ts = uxda_ts.neighborhood_filter(func=np.mean, r=5.0)\n", + "filtered_ts = uxda_ts.neighborhood_filter(\"mean\", r=5.0)\n", "\n", "print(\"Output dims:\", filtered_ts.dims, \" shape:\", filtered_ts.shape)" ] @@ -271,7 +397,7 @@ "metadata": {}, "outputs": [], "source": [ - "uxds_filtered = uxds.neighborhood_filter(func=np.mean, r=5.0)\n", + "uxds_filtered = uxds.neighborhood_filter(\"mean\", r=5.0)\n", "uxds_filtered" ] }, @@ -289,7 +415,7 @@ "outputs": [], "source": [ "# Apply the filter and then mask values below zero\n", - "result = uxda.neighborhood_filter(func=np.mean, r=5.0).where(lambda x: x > 0)\n", + "result = uxda.neighborhood_filter(\"mean\", r=5.0).where(lambda x: x > 0)\n", "print(\"Masked result type:\", type(result).__name__)\n", "print(\"uxgrid preserved:\", result.uxgrid is not None)\n", "print(\"Positive fraction:\", float((result > 0).sum()) / result.size)" @@ -310,7 +436,7 @@ " dims=[\"n_face\"],\n", ")\n", "\n", - "zonal_smooth = uxda.neighborhood_filter(func=np.mean, r=5.0).groupby(lat_bins).mean()\n", + "zonal_smooth = uxda.neighborhood_filter(\"mean\", r=5.0).groupby(lat_bins).mean()\n", "print(\"Grouped result type:\", type(zonal_smooth).__name__)\n", "print(\"Zonal means:\", zonal_smooth.values)" ] @@ -347,11 +473,11 @@ ")\n", "\n", "# r = 0 catches the element itself → output matches the input exactly\n", - "filtered_r0 = da_coarse.neighborhood_filter(func=np.mean, r=0.0)\n", + "filtered_r0 = da_coarse.neighborhood_filter(\"mean\", r=0.0)\n", "print(\"r = 0: unchanged?\", np.allclose(filtered_r0.values, da_coarse.values))\n", "\n", "# r = 360 catches every element → all values equal the global mean\n", - "filtered_global = da_coarse.neighborhood_filter(func=np.mean, r=360.0)\n", + "filtered_global = da_coarse.neighborhood_filter(\"mean\", r=360.0)\n", "print(\n", " \"r = 360: all equal global mean?\",\n", " np.allclose(filtered_global.values, da_coarse.values.mean()),\n", @@ -362,7 +488,21 @@ "cell_type": "markdown", "id": "cdf66aed5cc84ca1b48e60bad68798a8", "metadata": {}, - "source": "## API Reference\n\nSee also:\n\n- {py:meth}`uxarray.UxDataArray.neighborhood_filter`\n- {py:meth}`uxarray.UxDataset.neighborhood_filter`\n\nRelated methods that apply aggregations across different grid element types:\n\n- {py:meth}`uxarray.UxDataArray.topological_mean` — aggregate node→face, node→edge, etc.\n- {py:meth}`uxarray.UxDataArray.zonal_mean` — latitude-band averages\n- {py:meth}`uxarray.UxDataArray.azimuthal_mean` — rings of constant great-circle distance\n" + "source": [ + "## API Reference\n", + "\n", + "See also:\n", + "\n", + "- {py:meth}`uxarray.UxDataArray.neighborhood_filter`\n", + "- {py:meth}`uxarray.UxDataset.neighborhood_filter`\n", + "- {py:meth}`uxarray.Grid.neighborhoods` — reusable neighborhoods\n", + "\n", + "Related methods that apply aggregations across different grid element types:\n", + "\n", + "- {py:meth}`uxarray.UxDataArray.topological_mean` — aggregate node→face, node→edge, etc.\n", + "- {py:meth}`uxarray.UxDataArray.zonal_mean` — latitude-band averages\n", + "- {py:meth}`uxarray.UxDataArray.azimuthal_mean` — rings of constant great-circle distance\n" + ] } ], "metadata": { diff --git a/test/core/test_dataarray.py b/test/core/test_dataarray.py index 6b09afcad..a2c996631 100644 --- a/test/core/test_dataarray.py +++ b/test/core/test_dataarray.py @@ -369,35 +369,217 @@ def test_chunked_over_time_is_not_rechunked(self, gridpath, datasetpath): expected = psi.neighborhood_filter(func=np.mean, r=2.0).values np.testing.assert_allclose(filtered.compute().values, np.tile(expected, (6, 1))) + # Every named reduction, paired with the NumPy expression it must equal. + # The reference goes through the generic callable path, so this checks the + # compiled kernel against the loop it bypasses. + NAMED_REDUCTIONS = [ + ("mean", {}, lambda a, axis: np.mean(a, axis=axis)), + ("sum", {}, lambda a, axis: np.sum(a, axis=axis)), + ("min", {}, lambda a, axis: np.min(a, axis=axis)), + ("max", {}, lambda a, axis: np.max(a, axis=axis)), + ("median", {}, lambda a, axis: np.median(a, axis=axis)), + ("ptp", {}, lambda a, axis: np.ptp(a, axis=axis)), + ("std", {}, lambda a, axis: np.std(a, axis=axis)), + ("std", {"ddof": 1}, lambda a, axis: np.std(a, axis=axis, ddof=1)), + ("var", {}, lambda a, axis: np.var(a, axis=axis)), + ("var", {"ddof": 1}, lambda a, axis: np.var(a, axis=axis, ddof=1)), + ("quantile", {"q": 0.9}, lambda a, axis: np.quantile(a, 0.9, axis=axis)), + ("percentile", {"q": 90}, lambda a, axis: np.percentile(a, 90, axis=axis)), + ("percentile", {"q": 50}, lambda a, axis: np.percentile(a, 50, axis=axis)), + ] + + @pytest.mark.parametrize("name,kwargs,reference", NAMED_REDUCTIONS) + def test_named_reduction_matches_numpy( + self, name, kwargs, reference, gridpath, datasetpath + ): + """Every compiled reduction must agree with its NumPy equivalent, in + 1-D and with an extra dimension (which exercises the gufunc's + broadcast loop).""" + uxds = ux.open_dataset( + gridpath("ugrid", "outCSne30", "outCSne30.ug"), + datasetpath("ugrid", "outCSne30", "outCSne30_vortex.nc"), + ) + uxda = uxds["psi"] + nb = uxda.uxgrid.neighborhoods(r=3.0) + + got = nb.reduce(uxda, name, **kwargs) + expected = nb.reduce(uxda, reference) + np.testing.assert_allclose(got.values, expected.values, rtol=1e-12) + + stacked = UxDataArray( + np.vstack([uxda.values, uxda.values * -2.0]), + dims=["time", "n_face"], + uxgrid=uxda.uxgrid, + ) + got_2d = nb.reduce(stacked, name, **kwargs) + expected_2d = nb.reduce(stacked, reference) + np.testing.assert_allclose(got_2d.values, expected_2d.values, rtol=1e-12) + + @pytest.mark.parametrize( + "callable_func,name", + [ + (np.mean, "mean"), + (np.sum, "sum"), + (np.min, "min"), + (np.max, "max"), + (np.amin, "min"), + (np.amax, "max"), + (np.median, "median"), + (np.std, "std"), + (np.var, "var"), + (np.ptp, "ptp"), + ], + ) + def test_callable_alias_takes_kernel_path( + self, callable_func, name, gridpath, datasetpath + ): + """Code written against the original ``func=np.mean`` signature must + keep working, and must reach the same kernel the name does rather than + silently dropping to the generic loop.""" + from uxarray.grid.neighbors import _CALLABLE_ALIASES, _resolve_reduction + + assert _CALLABLE_ALIASES[callable_func] == name + assert _resolve_reduction(callable_func, {})[0] is _resolve_reduction(name, {})[0] + + uxds = ux.open_dataset( + gridpath("ugrid", "outCSne30", "outCSne30.ug"), + datasetpath("ugrid", "outCSne30", "outCSne30_vortex.nc"), + ) + uxda = uxds["psi"] + np.testing.assert_allclose( + uxda.neighborhood_filter(callable_func, r=3.0).values, + uxda.neighborhood_filter(name, r=3.0).values, + ) + + def test_callable_escape_hatch(self, gridpath, datasetpath): + """An arbitrary callable on the ``axis=-1`` contract still works, for + reductions with no compiled kernel.""" + from functools import partial + + uxds = ux.open_dataset( + gridpath("ugrid", "outCSne30", "outCSne30.ug"), + datasetpath("ugrid", "outCSne30", "outCSne30_vortex.nc"), + ) + uxda = uxds["psi"] + + # partial() is opaque to name dispatch, so this exercises the loop + via_partial = uxda.neighborhood_filter(partial(np.percentile, q=90), r=3.0) + via_name = uxda.neighborhood_filter("percentile", r=3.0, q=90) + np.testing.assert_allclose(via_partial.values, via_name.values, rtol=1e-12) + + # a user's own function, with no NumPy equivalent at all + def rms(values, axis): + return np.sqrt(np.mean(values**2, axis=axis)) + + filtered = uxda.neighborhood_filter(rms, r=3.0) + assert filtered.shape == uxda.shape + assert np.all(filtered.values >= 0) + @pytest.mark.parametrize( - "func", [np.mean, np.sum, np.min, np.max, np.median, np.amin, np.amax] + "func,kwargs,error,match", + [ + ("meen", {}, ValueError, "Unknown reduction 'meen'"), + ("mean", {"q": 90}, TypeError, "unexpected keyword argument"), + ("quantile", {}, TypeError, "requires the 'q' keyword"), + ("quantile", {"q": 90}, ValueError, "between 0 and 1"), + ("percentile", {"q": 1.5e3}, ValueError, "between 0 and 100"), + (42, {}, TypeError, "name of a reduction or a callable"), + ], ) - def test_kernel_matches_generic_path(self, func, gridpath, datasetpath): - """Every reduction with a compiled kernel must agree with the generic - loop it bypasses.""" - from uxarray.grid.neighbors import ( - _NEIGHBORHOOD_KERNELS, - _csr_neighbors, - _neighborhood_reduce, + def test_invalid_reduction(self, func, kwargs, error, match): + """Naming a reduction makes bad input catchable up front, rather than + as a TypeError from inside the loop.""" + uxgrid = ux.Grid.from_healpix(zoom=1) + uxda = UxDataArray( + np.arange(uxgrid.n_face, dtype=float), dims=["n_face"], uxgrid=uxgrid ) + with pytest.raises(error, match=match): + uxda.neighborhood_filter(func, r=1.0, **kwargs) + def test_neighborhoods_reuse_matches_one_shot(self, gridpath, datasetpath): + """Reducing over a reused Neighborhoods must equal the one-shot filter, + which is the whole point of being able to hold onto it.""" uxds = ux.open_dataset( gridpath("ugrid", "outCSne30", "outCSne30.ug"), datasetpath("ugrid", "outCSne30", "outCSne30_vortex.nc"), ) uxda = uxds["psi"] - assert func in _NEIGHBORHOOD_KERNELS, "expected a compiled kernel" + nb = uxda.uxgrid.neighborhoods(r=4.0) + + assert nb.r == 4.0 + assert nb.on == "face centers" + assert nb.grid_dim == "n_face" + assert nb.grid is uxda.uxgrid + assert "face centers" in repr(nb) + + for name, kwargs in [("mean", {}), ("std", {}), ("percentile", {"q": 90})]: + np.testing.assert_allclose( + nb.reduce(uxda, name, **kwargs).values, + uxda.neighborhood_filter(name, r=4.0, **kwargs).values, + rtol=1e-12, + ) + + def test_neighborhoods_n_neighbors(self): + """``n_neighbors`` reports the neighborhood sizes as a grid-mapped + field, which is how you see a radius sampling a mesh unevenly.""" + uxgrid = ux.Grid.from_healpix(zoom=2) + nb = uxgrid.neighborhoods(r=15.0) + counts = nb.n_neighbors + + assert counts.dims == ("n_face",) + assert counts.sizes["n_face"] == uxgrid.n_face + # every element is its own neighbor, so no neighborhood is ever empty + assert counts.min() >= 1 + + # a bigger radius can only add neighbors + wider = uxgrid.neighborhoods(r=30.0).n_neighbors + assert np.all(wider.values >= counts.values) - flat, starts, counts = _csr_neighbors(uxda.uxgrid, "face centers", 3.0) - # 2-D as well as 1-D, to exercise the gufunc's broadcast loop - block = np.vstack([uxda.values, uxda.values * -2.0]) - expected = _neighborhood_reduce(block, flat, starts, counts, func) + @pytest.mark.parametrize( + "on,dim", [("face centers", "n_face"), ("nodes", "n_node"), ("edge centers", "n_edge")] + ) + def test_neighborhoods_locations(self, on, dim): + """Neighborhoods can be built on any of the three grid locations.""" + uxgrid = ux.Grid.from_healpix(zoom=1) + # `getattr(uxgrid, dim)` populates the location's coordinates; a + # HEALPix grid cannot currently populate node coordinates lazily from + # inside the tree build (see _populate_healpix_boundaries). + data = np.arange(getattr(uxgrid, dim), dtype=float) + uxda = UxDataArray(data, dims=[dim], uxgrid=uxgrid) + + nb = uxgrid.neighborhoods(r=30.0, on=on) + assert nb.grid_dim == dim + np.testing.assert_allclose( + nb.reduce(uxda, "mean").values, + uxda.neighborhood_filter("mean", r=30.0).values, + ) - filtered = uxda.neighborhood_filter(func=func, r=3.0) - np.testing.assert_allclose(filtered.values, expected[0], rtol=1e-12) + def test_neighborhoods_rejects_mismatched_data(self): + """Reducing data mapped somewhere else must fail loudly rather than + indexing into the wrong element set.""" + from uxarray.errors import DataCenteringError - kernel_2d = _NEIGHBORHOOD_KERNELS[func](block, flat, starts, counts) - np.testing.assert_allclose(kernel_2d, expected, rtol=1e-12) + uxgrid = ux.Grid.from_healpix(zoom=1) + nb = uxgrid.neighborhoods(r=30.0, on="face centers") + + node_data = UxDataArray( + np.arange(uxgrid.n_node, dtype=float), dims=["n_node"], uxgrid=uxgrid + ) + with pytest.raises(DataCenteringError, match="reduce over 'n_face'"): + nb.reduce(node_data, "mean") + + # right dimension name, wrong grid + other = ux.Grid.from_healpix(zoom=2) + wrong_size = UxDataArray( + np.arange(other.n_face, dtype=float), dims=["n_face"], uxgrid=other + ) + with pytest.raises(DataCenteringError, match="different grid"): + nb.reduce(wrong_size, "mean") + + def test_neighborhoods_invalid_location(self): + uxgrid = ux.Grid.from_healpix(zoom=1) + with pytest.raises(ValueError, match="Invalid `on`"): + uxgrid.neighborhoods(r=1.0, on="face_centers") def test_float32_input(self, gridpath, datasetpath): """float32 fields take the compiled kernel's float32 signature and diff --git a/test/core/test_dataset.py b/test/core/test_dataset.py index cc32e1db4..34bfd23d1 100644 --- a/test/core/test_dataset.py +++ b/test/core/test_dataset.py @@ -206,3 +206,78 @@ def test_non_grid_variable_skipped(self): nt.assert_allclose(filtered["face_var"].values, uxds["face_var"].values) nt.assert_allclose(filtered["scalar_var"].values, uxds["scalar_var"].values) + + def test_named_reduction_with_parameter(self, gridpath, datasetpath): + """Named reductions and their parameters reach every variable.""" + uxds = ux.open_dataset( + gridpath("ugrid", "outCSne30", "outCSne30.ug"), + datasetpath("ugrid", "outCSne30", "outCSne30_vortex.nc"), + ) + + filtered_ds = uxds.neighborhood_filter("percentile", r=5.0, q=90) + filtered_da = uxds["psi"].neighborhood_filter("percentile", r=5.0, q=90) + + nt.assert_allclose(filtered_ds["psi"].values, filtered_da.values) + + def test_one_query_shared_across_variables(self): + """Variables mapped to the same location must share a single neighbor + query. The query dominates the cost, so per-variable rebuilding would + make a dataset filter scale with the number of variables.""" + from unittest.mock import patch + + import uxarray.grid.neighbors as neighbors + + uxgrid = ux.Grid.from_healpix(zoom=2) + rng = np.random.default_rng(0) + uxds = UxDataset( + data_vars={ + f"v{i}": ("n_face", rng.random(uxgrid.n_face)) for i in range(5) + }, + uxgrid=uxgrid, + ) + + real = neighbors._csr_neighbors + with patch.object( + neighbors, "_csr_neighbors", side_effect=real + ) as spy: + filtered = uxds.neighborhood_filter("mean", r=20.0) + + assert spy.call_count == 1, ( + f"expected one neighbor query for 5 same-location variables, " + f"got {spy.call_count}" + ) + + expected = uxds["v0"].neighborhood_filter("mean", r=20.0) + nt.assert_allclose(filtered["v0"].values, expected.values) + + def test_query_per_location_not_per_variable(self): + """Variables on different grid locations each need their own query, + but only one apiece.""" + from unittest.mock import patch + + import uxarray.grid.neighbors as neighbors + + uxgrid = ux.Grid.from_healpix(zoom=2) + # populate node coordinates before the tree is built (see the HEALPix + # boundary-population path, which cannot do it lazily) + n_node, n_face = uxgrid.n_node, uxgrid.n_face + rng = np.random.default_rng(0) + uxds = UxDataset( + data_vars={ + "face_a": ("n_face", rng.random(n_face)), + "face_b": ("n_face", rng.random(n_face)), + "node_a": ("n_node", rng.random(n_node)), + }, + uxgrid=uxgrid, + ) + + real = neighbors._csr_neighbors + with patch.object( + neighbors, "_csr_neighbors", side_effect=real + ) as spy: + uxds.neighborhood_filter("mean", r=20.0) + + assert spy.call_count == 2, ( + f"expected one query per grid location (faces, nodes), " + f"got {spy.call_count}" + ) diff --git a/uxarray/core/dataarray.py b/uxarray/core/dataarray.py index d95f57ad6..4129f969b 100644 --- a/uxarray/core/dataarray.py +++ b/uxarray/core/dataarray.py @@ -30,7 +30,7 @@ from uxarray.formatting_html import array_repr from uxarray.grid import Grid from uxarray.grid.dual import construct_dual -from uxarray.grid.neighbors import _neighborhood_filter +from uxarray.grid.neighbors import Neighborhoods from uxarray.grid.validation import _check_duplicate_nodes_indices from uxarray.io._healpix import get_zoom_from_cells from uxarray.plot.accessor import UxDataArrayPlotAccessor @@ -2181,26 +2181,42 @@ def get_dual(self): return uxda + def _neighborhood_location(self, caller: str) -> str: + """Grid location this data is mapped to, in ``Neighborhoods`` terms.""" + if self._face_centered(): + return "face centers" + if self._node_centered(): + return "nodes" + if self._edge_centered(): + return "edge centers" + raise DataCenteringError( + f"{caller} requires data mapped to nodes, edges, or faces, " + f"but the dimensions {self.dims!r} do not match any grid dimension " + f"{GRID_DIMS}." + ) + def neighborhood_filter( self, - func: Callable = np.mean, + func: str | Callable = "mean", r: float = 1.0, + **kwargs, ) -> UxDataArray: """Apply a neighborhood filter, replacing the value at each grid - element with ``func`` applied to all elements within a circular + element with a reduction of all elements within a circular neighborhood of radius ``r``. Parameters ---------- - func: Callable, default=np.mean - Apply this function to neighborhood. ``np.mean``, ``np.sum``, - ``np.min``, ``np.max`` and ``np.median`` use a compiled kernel. - Any other function must accept an ``axis`` keyword argument (as - NumPy reductions do); use ``functools.partial`` to supply - additional arguments, e.g. ``functools.partial(np.percentile, - q=90)``. + func : str or Callable, default="mean" + Name of the reduction to apply: "mean", "sum", "min", "max", + "median", "ptp", "std", "var", "quantile", or "percentile". Named + reductions run compiled. A callable is accepted as an escape hatch + for anything not in that list; see Notes. r : float, default=1. Radius of the neighborhood, in degrees. + **kwargs + Parameter for the named reduction: ``q`` for "quantile" (0-1) and + "percentile" (0-100), ``ddof`` for "std" and "var". Returns ------- @@ -2211,9 +2227,11 @@ def neighborhood_filter( ------ DataCenteringError (subclass of ValueError) If the data is not mapped to nodes, edges, or faces. + ValueError + If ``func`` names a reduction that does not exist. TypeError - If ``func`` has no compiled kernel and does not accept an ``axis`` - keyword argument. + If ``func`` is a callable that does not accept an ``axis`` keyword + argument, or if a keyword argument does not apply to ``func``. Notes ----- @@ -2221,6 +2239,16 @@ def neighborhood_filter( every element is its own neighbor at distance 0, so ``r = 0`` returns the data unchanged and the result never contains spurious ``NaN``. + A callable ``func`` is applied as ``func(values, axis=-1)`` over a block + whose last axis is the neighborhood, once per grid element, in Python. + That is considerably slower than a named reduction, so prefer a name + where one exists. + + Each call queries the grid for neighbors, which usually costs more than + the reduction itself. To apply several reductions at one radius, build + the neighborhoods once with :meth:`Grid.neighborhoods` and call + :meth:`Neighborhoods.reduce` on it. + A neighborhood may span the whole grid, so the grid dimension cannot be chunked; it is collapsed to a single chunk (with a warning) for dask-backed data. The remaining dimensions stay chunked and lazy, so @@ -2230,57 +2258,29 @@ def neighborhood_filter( -------- Apply a mean filter with a 5-degree radius: - >>> import numpy as np >>> import uxarray as ux >>> uxds = ux.tutorial.open_dataset("outCSne30-vortex") >>> uxda = uxds["psi"] - >>> smoothed = uxda.neighborhood_filter(func=np.mean, r=5.0) + >>> smoothed = uxda.neighborhood_filter("mean", r=5.0) - Use ``functools.partial`` for functions requiring extra arguments: + Reductions taking a parameter receive it as a keyword argument: - >>> from functools import partial - >>> p90 = uxda.neighborhood_filter(func=partial(np.percentile, q=90), r=5.0) + >>> p90 = uxda.neighborhood_filter("percentile", r=5.0, q=90) + >>> spread = uxda.neighborhood_filter("std", r=5.0, ddof=1) See Also -------- + Grid.neighborhoods : Reusable neighborhoods, for several reductions at one radius. UxDataArray.topological_mean : Aggregate values across neighboring grid element types. UxDataArray.zonal_mean : Average over latitude bands. UxDataArray.azimuthal_mean : Average over rings of constant great-circle distance. """ - - if self._face_centered(): - data_mapping = "face centers" - grid_dim = "n_face" - elif self._node_centered(): - data_mapping = "nodes" - grid_dim = "n_node" - elif self._edge_centered(): - data_mapping = "edge centers" - grid_dim = "n_edge" - else: - raise DataCenteringError( - f"neighborhood_filter requires data mapped to nodes, edges, or faces, " - f"but the dimensions {self.dims!r} do not match any grid dimension " - f"{GRID_DIMS}." - ) - - # ``_neighborhood_filter`` declares the grid dimension as a core - # dimension, so it moves that dimension into place itself; no manual - # transpose is needed on the way in. - filtered = _neighborhood_filter( - self.uxgrid, self, data_mapping, grid_dim, func=func, r=r + neighborhoods = Neighborhoods( + self.uxgrid, + r=r, + on=self._neighborhood_location("neighborhood_filter"), ) - - # Core dimensions come back appended last, so restore the caller's - # dimension order when it differed. - if filtered.dims != self.dims: - filtered = filtered.transpose(*self.dims) - - # ``apply_ufunc`` returns a plain xr.DataArray, dropping the subclass - # and its grid. Name, coords and attrs are carried through already. - # The filtered data lives on the identical grid topology, so the same - # uxgrid is reattached rather than copied. - return UxDataArray(filtered, uxgrid=self.uxgrid) + return neighborhoods.reduce(self, func=func, **kwargs) def __getattribute__(self, name): """Intercept accessor method calls to return Ux-aware accessors.""" diff --git a/uxarray/core/dataset.py b/uxarray/core/dataset.py index 2feb30eaf..b0579ce72 100644 --- a/uxarray/core/dataset.py +++ b/uxarray/core/dataset.py @@ -20,6 +20,7 @@ from uxarray.formatting_html import dataset_repr from uxarray.grid import Grid from uxarray.grid.dual import construct_dual +from uxarray.grid.neighbors import Neighborhoods from uxarray.grid.validation import _check_duplicate_nodes_indices from uxarray.io._healpix import get_zoom_from_cells from uxarray.plot.accessor import UxDatasetPlotAccessor @@ -659,24 +660,26 @@ def to_array( def neighborhood_filter( self, - func: Callable = np.mean, + func: str | Callable = "mean", r: float = 1.0, + **kwargs, ) -> UxDataset: """Apply a neighborhood filter, replacing the value at each grid - element of every data variable with ``func`` applied to all elements - within a circular neighborhood of radius ``r``. + element of every data variable with a reduction of all elements within + a circular neighborhood of radius ``r``. Parameters ---------- - func: Callable, default=np.mean - Apply this function to neighborhood. ``np.mean``, ``np.sum``, - ``np.min``, ``np.max`` and ``np.median`` use a compiled kernel. - Any other function must accept an ``axis`` keyword argument (as - NumPy reductions do); use ``functools.partial`` to supply - additional arguments, e.g. ``functools.partial(np.percentile, - q=90)``. + func : str or Callable, default="mean" + Name of the reduction to apply: "mean", "sum", "min", "max", + "median", "ptp", "std", "var", "quantile", or "percentile". Named + reductions run compiled. A callable is accepted as an escape hatch; + see :meth:`UxDataArray.neighborhood_filter`. r : float, default=1. Radius of the neighborhood, in degrees. + **kwargs + Parameter for the named reduction: ``q`` for "quantile" (0-1) and + "percentile" (0-100), ``ddof`` for "std" and "var". Returns ------- @@ -690,23 +693,33 @@ def neighborhood_filter( variables stay lazy. See :meth:`UxDataArray.neighborhood_filter` for details. + Variables mapped to the same grid location share one neighbor query, so + filtering a dataset costs one query per location present rather than + one per variable. + Examples -------- Apply a mean filter to all grid-mapped variables in a dataset: - >>> import numpy as np >>> import uxarray as ux >>> uxds = ux.tutorial.open_dataset("outCSne30-vortex") - >>> uxds_smooth = uxds.neighborhood_filter(func=np.mean, r=5.0) + >>> uxds_smooth = uxds.neighborhood_filter("mean", r=5.0) See Also -------- UxDataArray.neighborhood_filter : Filter a single data variable. + Grid.neighborhoods : Reusable neighborhoods, for several reductions at one radius. UxDataArray.zonal_mean : Average over latitude bands. UxDataArray.azimuthal_mean : Average over rings of constant great-circle distance. """ destination_uxds = self._copy() + + # The neighbor query dominates the cost of a reduction, and it depends + # only on (grid, location, radius) -- not on the data. Variables mapped + # to the same location therefore share one query, built on first use. + neighborhoods: dict[str, Neighborhoods] = {} + # Loop through UxDataArrays in UxDataset and apply the filter to every # variable that is mapped to a grid element (node, edge, or face). # Variables without a grid dimension are left unchanged. @@ -717,9 +730,15 @@ def neighborhood_filter( if not any(dim in GRID_DIMS for dim in uxda.dims): continue - # UxDataArray.neighborhood_filter handles the transpose internally, - # so dimension order is always preserved. - destination_uxds[var_name] = uxda.neighborhood_filter(func, r) + location = uxda._neighborhood_location("neighborhood_filter") + if location not in neighborhoods: + neighborhoods[location] = Neighborhoods(self.uxgrid, r=r, on=location) + + # Neighborhoods.reduce restores the input dimension order, so it is + # always preserved. + destination_uxds[var_name] = neighborhoods[location].reduce( + uxda, func=func, **kwargs + ) return destination_uxds diff --git a/uxarray/grid/grid.py b/uxarray/grid/grid.py index 6ba2d7baa..e786e3aff 100644 --- a/uxarray/grid/grid.py +++ b/uxarray/grid/grid.py @@ -60,6 +60,7 @@ from uxarray.grid.neighbors import ( BallTree, KDTree, + Neighborhoods, SpatialHash, _populate_edge_face_distances, _populate_edge_node_distances, @@ -1798,6 +1799,42 @@ def get_ball_tree( return self._ball_tree + def neighborhoods(self, r: float = 1.0, on: str = "face centers") -> Neighborhoods: + """Finds the grid elements within ``r`` degrees of every element of + ``on``, returning a reusable :class:`Neighborhoods`. + + The radius query behind this dominates the cost of a neighborhood + reduction, so building this once and reducing several times over it is + substantially cheaper than calling + :meth:`UxDataArray.neighborhood_filter` repeatedly, which rebuilds it + on every call. + + Parameters + ---------- + r : float, default=1. + Radius of the neighborhood, in degrees of great-circle distance. + on : str, default="face centers" + Grid location to center the neighborhoods on: "nodes", + "edge centers", or "face centers". + + Returns + ------- + Neighborhoods + + Examples + -------- + >>> import uxarray as ux + >>> uxds = ux.tutorial.open_dataset("outCSne30-vortex") # doctest: +SKIP + >>> nb = uxds.uxgrid.neighborhoods(r=5.0) # doctest: +SKIP + >>> smooth = nb.reduce(uxds["psi"], "mean") # doctest: +SKIP + >>> p90 = nb.reduce(uxds["psi"], "percentile", q=90) # doctest: +SKIP + + See Also + -------- + UxDataArray.neighborhood_filter : One-shot filter for a single reduction. + """ + return Neighborhoods(self, r=r, on=on) + def _get_scipy_kd_tree( self, coordinates: str | None = "face", reconstruct: bool = False ): diff --git a/uxarray/grid/neighbors.py b/uxarray/grid/neighbors.py index 6b7d72c82..d9e56ee51 100644 --- a/uxarray/grid/neighbors.py +++ b/uxarray/grid/neighbors.py @@ -1,5 +1,5 @@ import warnings -from typing import Callable +from typing import Callable, NamedTuple import numpy as np import xarray as xr @@ -1206,6 +1206,15 @@ def _get_element_coords(grid, data_mapping: str, coordinate_system: str): _GUFUNC_LAYOUT = "(n),(k),(m),(m)->(m)" _GUFUNC_KWARGS = {"nopython": True, "cache": True, "target": "parallel"} +# Variant carrying one scalar parameter (``q``, ``ddof``) as a scalar core +# dimension, so the value travels with the call rather than being baked into a +# separately compiled kernel per value. +_GUFUNC_SIGNATURES_P = [ + "void(float64[:], int64[:], int64[:], int64[:], float64, float64[:])", + "void(float32[:], int64[:], int64[:], int64[:], float64, float64[:])", +] +_GUFUNC_LAYOUT_P = "(n),(k),(m),(m),()->(m)" + @guvectorize(_GUFUNC_SIGNATURES, _GUFUNC_LAYOUT, **_GUFUNC_KWARGS) def _gu_mean(data, flat, starts, counts, out): @@ -1263,46 +1272,202 @@ def _gu_min(data, flat, starts, counts, out): out[i] = best -@guvectorize(_GUFUNC_SIGNATURES, _GUFUNC_LAYOUT, **_GUFUNC_KWARGS) -def _gu_median(data, flat, starts, counts, out): - # A median needs the neighborhood materialized so it can be sorted, which - # is why ``np.ufunc.reduceat`` cannot express it but a kernel can: the - # scratch buffer is sized once to the largest neighborhood and reused. - widest = 0 - for i in range(counts.shape[0]): - if counts[i] > widest: - widest = counts[i] - buffer = np.empty(widest, dtype=np.float64) +# Order statistics cannot stream: the neighborhood has to be materialized +# before it can be sorted or partitioned. Those reductions share one gather +# loop, and the reduction itself is whatever numba-compilable callable is +# handed to the factory. Numba supports the NumPy reductions in nopython mode, +# so ``np.median`` and friends are used directly rather than reimplemented -- +# numba's ``np.median`` partitions instead of fully sorting, which makes it +# faster than an equivalent hand-written sort as well as shorter. +def _make_buffered_kernel(reduce_fn): + """Builds a kernel that gathers each neighborhood, then calls ``reduce_fn`` + on the 1-D result. ``reduce_fn`` must be numba-compilable.""" + # A reducer shared between kernels arrives already compiled; numba rejects + # jitting a dispatcher twice. + if not hasattr(reduce_fn, "py_func"): + reduce_fn = njit(cache=True)(reduce_fn) + + @guvectorize(_GUFUNC_SIGNATURES, _GUFUNC_LAYOUT, **_GUFUNC_KWARGS) + def kernel(data, flat, starts, counts, out): + widest = 0 + for i in range(counts.shape[0]): + if counts[i] > widest: + widest = counts[i] + buffer = np.empty(widest, dtype=np.float64) - for i in range(starts.shape[0]): - count = counts[i] - if count == 0: - out[i] = np.nan - continue - start = starts[i] - for j in range(count): - buffer[j] = data[flat[start + j]] - window = np.sort(buffer[:count]) - if count % 2: - out[i] = window[count // 2] - else: - out[i] = 0.5 * (window[count // 2 - 1] + window[count // 2]) - - -# Reductions with a compiled kernel. Anything else (``np.median`` with a -# custom axis, ``functools.partial(np.percentile, q=90)``, a user's own -# callable) falls back to the generic path, which is slower but accepts any -# function taking an ``axis`` keyword. -_NEIGHBORHOOD_KERNELS = { - np.mean: _gu_mean, - np.sum: _gu_sum, - np.max: _gu_max, - np.amax: _gu_max, - np.min: _gu_min, - np.amin: _gu_min, - np.median: _gu_median, + for i in range(starts.shape[0]): + count = counts[i] + if count == 0: + out[i] = np.nan + continue + start = starts[i] + for j in range(count): + buffer[j] = data[flat[start + j]] + out[i] = reduce_fn(buffer[:count]) + + return kernel + + +def _make_parameterized_kernel(reduce_fn): + """As ``_make_buffered_kernel``, but ``reduce_fn`` also takes one scalar + parameter (``q`` for quantiles, ``ddof`` for spreads), declared as a scalar + core dimension so it can vary per call without recompiling.""" + # A reducer shared between kernels arrives already compiled; numba rejects + # jitting a dispatcher twice. + if not hasattr(reduce_fn, "py_func"): + reduce_fn = njit(cache=True)(reduce_fn) + + @guvectorize(_GUFUNC_SIGNATURES_P, _GUFUNC_LAYOUT_P, **_GUFUNC_KWARGS) + def kernel(data, flat, starts, counts, param, out): + widest = 0 + for i in range(counts.shape[0]): + if counts[i] > widest: + widest = counts[i] + buffer = np.empty(widest, dtype=np.float64) + + for i in range(starts.shape[0]): + count = counts[i] + if count == 0: + out[i] = np.nan + continue + start = starts[i] + for j in range(count): + buffer[j] = data[flat[start + j]] + out[i] = reduce_fn(buffer[:count], param) + + return kernel + + +@njit(cache=True) +def _variance(window, ddof): + """Variance with a delta degrees of freedom. Numba's ``np.var`` takes no + ``ddof``, so the two-pass form is spelled out.""" + denominator = window.size - ddof + if denominator <= 0: + return np.nan + center = np.mean(window) + total = 0.0 + for value in window: + total += (value - center) ** 2 + return total / denominator + + +_gu_median = _make_buffered_kernel(lambda window: np.median(window)) +_gu_ptp = _make_buffered_kernel(lambda window: np.max(window) - np.min(window)) +_gu_quantile = _make_parameterized_kernel(lambda window, q: np.quantile(window, q)) +_gu_var = _make_parameterized_kernel(_variance) +_gu_std = _make_parameterized_kernel( + lambda window, ddof: np.sqrt(_variance(window, ddof)) +) + + +class _Reduction(NamedTuple): + """A named reduction, and the single scalar parameter it accepts (if any). + + Keeping parameters to at most one keeps the gufunc layouts down to the two + declared above, which covers every reduction implemented here. + """ + + kernel: object + param: str | None = None + default: float | None = None + + +# Reductions with a compiled kernel, addressed by name. A name always takes the +# fast path, which is why the public API documents names rather than callables: +# dispatching on a function object cannot see through ``functools.partial``, so +# a parameterized reduction could never hit a kernel that way. +_REDUCTIONS = { + "mean": _Reduction(_gu_mean), + "sum": _Reduction(_gu_sum), + "min": _Reduction(_gu_min), + "max": _Reduction(_gu_max), + "median": _Reduction(_gu_median), + "ptp": _Reduction(_gu_ptp), + "std": _Reduction(_gu_std, param="ddof", default=0.0), + "var": _Reduction(_gu_var, param="ddof", default=0.0), + "quantile": _Reduction(_gu_quantile, param="q"), + "percentile": _Reduction(_gu_quantile, param="q"), } +# Callables accepted for backwards compatibility, so that code written against +# the original ``func=np.mean`` signature keeps the fast path instead of +# silently dropping to the generic loop. +_CALLABLE_ALIASES = { + np.mean: "mean", + np.sum: "sum", + np.max: "max", + np.amax: "max", + np.min: "min", + np.amin: "min", + np.median: "median", + np.std: "std", + np.var: "var", + np.ptp: "ptp", +} + + +def _resolve_reduction(func, kwargs): + """Maps ``func`` (a name or a callable) onto a kernel and its parameter. + + Returns ``(kernel, param_value)`` for a compiled reduction, or + ``(None, None)`` when ``func`` is a callable that has to go through the + generic loop. + """ + name = func if isinstance(func, str) else _CALLABLE_ALIASES.get(func) + + if name is None: + if not callable(func): + raise TypeError( + f"`func` must be the name of a reduction or a callable, but got " + f"{func!r}. Valid names: {', '.join(sorted(_REDUCTIONS))}." + ) + if kwargs: + raise TypeError( + f"Got unexpected keyword argument(s) {', '.join(sorted(kwargs))} " + f"for a callable `func`. Parameters are only supported for named " + f"reductions; use `functools.partial` to bind them to a callable." + ) + return None, None + + if name not in _REDUCTIONS: + raise ValueError( + f"Unknown reduction {name!r}. Expected one of: " + f"{', '.join(sorted(_REDUCTIONS))}." + ) + + reduction = _REDUCTIONS[name] + unexpected = set(kwargs) - ({reduction.param} if reduction.param else set()) + if unexpected: + raise TypeError( + f"Reduction {name!r} got unexpected keyword argument(s) " + f"{', '.join(sorted(unexpected))}." + + (f" It accepts {reduction.param!r}." if reduction.param else "") + ) + + if reduction.param is None: + return reduction.kernel, None + + if reduction.param in kwargs: + value = float(kwargs[reduction.param]) + elif reduction.default is not None: + value = reduction.default + else: + raise TypeError( + f"Reduction {name!r} requires the {reduction.param!r} keyword argument." + ) + + # `percentile` is `quantile` on a 0-100 scale; normalize so both share one + # kernel rather than compiling a near-duplicate. + if name == "percentile": + if not 0.0 <= value <= 100.0: + raise ValueError(f"`q` must be between 0 and 100, but got {value}.") + value /= 100.0 + elif name == "quantile" and not 0.0 <= value <= 1.0: + raise ValueError(f"`q` must be between 0 and 1, but got {value}.") + + return reduction.kernel, value + def _csr_neighbors(grid, data_mapping: str, r: float): """Queries the neighborhood of every element and returns it in CSR form. @@ -1408,88 +1573,185 @@ def _rechunk_grid_dim(uxda, grid_dim: str): return uxda.chunk({grid_dim: -1}) -def _neighborhood_filter( - grid, - uxda, - data_mapping: str, - grid_dim: str, - func: Callable = np.mean, - r: float = 1.0, -): - """Applies ``func`` to the set of grid elements within a circular - neighborhood of radius ``r`` around each element of ``data_mapping``. +ELEMENT_DIMS = { + "nodes": "n_node", + "edge centers": "n_edge", + "face centers": "n_face", +} + + +class Neighborhoods: + """The set of grid elements within a radius ``r`` of every element of one + grid location, ready to be reduced over. + + Building this queries a ``BallTree`` once, which is by far the dominant + cost of a neighborhood reduction — typically far more than the reduction + itself. Holding onto the result lets several reductions, or several + variables, share that one query instead of repeating it. Parameters ---------- grid : Grid - Source grid used to construct the ``BallTree`` used for the - neighborhood queries. - uxda : xr.DataArray - Data to filter. The grid dimension may sit at any position. - data_mapping : str - One of "nodes", "edge centers", or "face centers", identifying which - grid element ``uxda`` is mapped to. - grid_dim : str - Name of the grid dimension (``n_node``, ``n_edge``, or ``n_face``). - func : Callable, default=np.mean - Function applied to the values found in each neighborhood. Reductions - with a compiled kernel (``np.mean``, ``np.sum``, ``np.min``, - ``np.max``, ``np.median``) take a fast path. Any other function must - accept an ``axis`` keyword argument (as ``np.mean``, ``np.median``, - and similar NumPy reductions do) so that extra, non-grid dimensions - (e.g. ``time``) are preserved rather than being collapsed. + Grid whose elements define the neighborhoods. r : float, default=1. - Radius of the neighborhood, in degrees. + Radius of the neighborhood, in degrees of great-circle distance. + on : str, default="face centers" + Grid location the neighborhoods are built around: "nodes", + "edge centers", or "face centers". + + Examples + -------- + >>> import uxarray as ux + >>> uxds = ux.tutorial.open_dataset("outCSne30-vortex") # doctest: +SKIP + >>> nb = uxds.uxgrid.neighborhoods(r=5.0) # doctest: +SKIP + >>> smooth = nb.reduce(uxds["psi"], "mean") # doctest: +SKIP + >>> spread = nb.reduce(uxds["psi"], "std") # doctest: +SKIP + + See Also + -------- + UxDataArray.neighborhood_filter : One-shot filter that builds this internally. + """ - Returns - ------- - filtered : xr.DataArray - Filtered data, float64, with the grid dimension moved last. Lazy if - the input was lazy. + def __init__(self, grid, r: float = 1.0, on: str = "face centers"): + if on not in ELEMENT_DIMS: + raise ValueError( + f"Invalid `on`. Expected one of {', '.join(sorted(ELEMENT_DIMS))}, " + f"but received {on!r}." + ) - Raises - ------ - TypeError - If ``func`` has no compiled kernel and does not accept an ``axis`` - keyword argument. + self._grid = grid + self._r = float(r) + self._on = on + self._flat, self._starts, self._counts = _csr_neighbors(grid, on, self._r) - Notes - ----- - The grid dimension is a core dimension of the reduction, so it is - collapsed to a single chunk for dask-backed input; the remaining - dimensions stay chunked and are evaluated lazily. - """ + @property + def grid(self): + """Grid the neighborhoods were built from.""" + return self._grid - flat, starts, counts = _csr_neighbors(grid, data_mapping, r) - kernel = _NEIGHBORHOOD_KERNELS.get(func) + @property + def r(self) -> float: + """Neighborhood radius, in degrees.""" + return self._r - if kernel is not None: + @property + def on(self) -> str: + """Grid location the neighborhoods are centered on.""" + return self._on - def _apply(block): - # The kernels are compiled for float32/float64 only; anything else - # (integer fields, say) is promoted, which the generic path does - # too by writing into a float64 output. - if block.dtype not in (np.float64, np.float32): - block = block.astype(np.float64) - return kernel(block, flat, starts, counts) - else: + @property + def grid_dim(self) -> str: + """Name of the grid dimension this reduces over.""" + return ELEMENT_DIMS[self._on] - def _apply(block): - return _neighborhood_reduce(block, flat, starts, counts, func) - - uxda = _rechunk_grid_dim(uxda, grid_dim) - - # ``apply_ufunc`` moves the grid dimension last before calling ``_apply`` - # and, for dask-backed input, hands each chunk over as a materialized - # NumPy block. Indexing the array one destination element at a time (as - # this function used to) would instead trigger one graph execution per - # grid element. - return xr.apply_ufunc( - _apply, - uxda, - input_core_dims=[[grid_dim]], - output_core_dims=[[grid_dim]], - dask="parallelized", - output_dtypes=[np.float64], - keep_attrs=True, - ) + @property + def n_neighbors(self) -> xr.DataArray: + """Number of elements in each neighborhood, itself a grid-mapped field. + + Useful for seeing how a fixed radius samples a variable-resolution + mesh, where the count varies by region. + """ + return xr.DataArray( + self._counts.copy(), + dims=[self.grid_dim], + name="n_neighbors", + attrs={"long_name": f"elements within {self._r} degrees"}, + ) + + def __repr__(self) -> str: + return ( + f"" + ) + + def reduce(self, uxda, func="mean", **kwargs): + """Reduces ``uxda`` over each neighborhood. + + Parameters + ---------- + uxda : UxDataArray + Data to reduce, mapped to the same grid location as ``on``. The + grid dimension may sit at any position. + func : str or Callable, default="mean" + Name of a compiled reduction — "mean", "sum", "min", "max", + "median", "ptp", "std", "var", "quantile", "percentile" — or a + callable taking an ``axis`` keyword (see Notes). + **kwargs + Parameter for the named reduction: ``q`` for "quantile" (0-1) and + "percentile" (0-100), ``ddof`` for "std" and "var". + + Returns + ------- + UxDataArray + Reduced data as float64, with the input's dimension order. Lazy if + the input was lazy. + + Notes + ----- + A callable is an escape hatch for reductions not implemented here. It + is applied as ``func(values, axis=-1)`` over a block whose last axis is + the neighborhood, once per element, in Python — considerably slower + than a named reduction. Named reductions run compiled. + """ + # Local import: uxarray.core.dataarray imports this module. + from uxarray.core.dataarray import UxDataArray + from uxarray.errors import DataCenteringError + + grid_dim = self.grid_dim + if grid_dim not in uxda.dims: + raise DataCenteringError( + f"These neighborhoods are built on {self._on!r} and reduce over " + f"{grid_dim!r}, but the data has dimensions {tuple(uxda.dims)!r}." + ) + if uxda.sizes[grid_dim] != self._counts.size: + raise DataCenteringError( + f"Data has {uxda.sizes[grid_dim]} elements along {grid_dim!r}, but " + f"these neighborhoods describe {self._counts.size}. The data is " + f"probably mapped to a different grid." + ) + + kernel, param = _resolve_reduction(func, kwargs) + + if kernel is None: + + def _apply(block): + return _neighborhood_reduce( + block, self._flat, self._starts, self._counts, func + ) + else: + + def _apply(block): + # The kernels are compiled for float32/float64 only; anything + # else (integer fields, say) is promoted, which the generic + # path does too by writing into a float64 output. + if block.dtype not in (np.float64, np.float32): + block = block.astype(np.float64) + args = (block, self._flat, self._starts, self._counts) + if param is not None: + args += (param,) + return kernel(*args) + + work = _rechunk_grid_dim(uxda, grid_dim) + + # ``apply_ufunc`` moves the grid dimension last before calling + # ``_apply`` and, for dask-backed input, hands each chunk over as a + # materialized NumPy block. Indexing the array one destination element + # at a time would instead trigger one graph execution per grid element. + filtered = xr.apply_ufunc( + _apply, + work, + input_core_dims=[[grid_dim]], + output_core_dims=[[grid_dim]], + dask="parallelized", + output_dtypes=[np.float64], + keep_attrs=True, + ) + + # Core dimensions come back appended last, so restore the input order. + if filtered.dims != uxda.dims: + filtered = filtered.transpose(*uxda.dims) + + # ``apply_ufunc`` returns a plain xr.DataArray, dropping the subclass + # and its grid. Name, coords and attrs are carried through already. + return UxDataArray(filtered, uxgrid=getattr(uxda, "uxgrid", self._grid)) From c7d27115995188952ab2b0f421360496a369381e Mon Sep 17 00:00:00 2001 From: cmdupuis3 Date: Wed, 5 Aug 2026 17:30:51 -0500 Subject: [PATCH 20/21] 941: some cleanup --- test/core/test_dataarray.py | 370 ++++++++++++++---------------------- test/core/test_dataset.py | 72 ++----- uxarray/core/dataset.py | 17 +- uxarray/grid/neighbors.py | 208 ++++++-------------- 4 files changed, 223 insertions(+), 444 deletions(-) diff --git a/test/core/test_dataarray.py b/test/core/test_dataarray.py index a2c996631..1a030815a 100644 --- a/test/core/test_dataarray.py +++ b/test/core/test_dataarray.py @@ -310,68 +310,48 @@ def test_uses_spherical_tree_regardless_of_cached_tree(self): uxda.neighborhood_filter(func=np.mean, r=20.0).values, expected ) - def test_dask_input_stays_lazy(self, gridpath, datasetpath): - """Lazy input stays lazy, and the result matches the eager path.""" + def test_auto_transpose_direct_on_uxdataarray(self, gridpath, datasetpath): + """Calling neighborhood_filter directly on a (time, n_face) UxDataArray + (without going through UxDataset) should preserve the original dim order.""" uxds = ux.open_dataset( gridpath("ugrid", "outCSne30", "outCSne30.ug"), datasetpath("ugrid", "outCSne30", "outCSne30_vortex.nc"), ) - eager = uxds["psi"].neighborhood_filter(func=np.mean, r=2.0) - - lazy_uxda = uxds["psi"].chunk({"n_face": -1}) - assert lazy_uxda.chunks is not None - - filtered = lazy_uxda.neighborhood_filter(func=np.mean, r=2.0) - assert filtered.chunks is not None, "the filter should not force a compute" - assert isinstance(filtered, UxDataArray) - assert filtered.uxgrid == lazy_uxda.uxgrid - np.testing.assert_allclose(filtered.compute().values, eager.values) + uxda = uxds["psi"] - def test_grid_dim_chunks_are_collapsed_with_warning(self, gridpath, datasetpath): - """A neighborhood may span the whole grid, so the grid dimension cannot - stay chunked. Collapsing it is a memory decision the user made, so it - is not done silently.""" - uxds = ux.open_dataset( - gridpath("ugrid", "outCSne30", "outCSne30.ug"), - datasetpath("ugrid", "outCSne30", "outCSne30_vortex.nc"), + # Build a multi-dim UxDataArray with time as the FIRST (non-grid) dim + data = np.stack([uxda.values, uxda.values * 2.0]) # shape (2, n_face) + uxda_time = UxDataArray( + data, dims=["time", "n_face"], uxgrid=uxda.uxgrid, name="psi_time" ) - expected = uxds["psi"].neighborhood_filter(func=np.mean, r=2.0).values - uxda = uxds["psi"].chunk({"n_face": 1000}) - assert len(uxda.chunksizes["n_face"]) > 1 + # n_face is already last: no transpose needed internally + filtered = uxda_time.neighborhood_filter(func=np.mean, r=0.0) + assert filtered.dims == ("time", "n_face") + assert filtered.shape == (2, uxda.shape[0]) + np.testing.assert_allclose(filtered.values, data) - with pytest.warns(UserWarning, match="Rechunking 'n_face'"): - filtered = uxda.neighborhood_filter(func=np.mean, r=2.0) + # Also test with a UxDataArray that has grid dim NOT last (n_face, time) + uxda_face_first = uxda_time.transpose("n_face", "time") + filtered2 = uxda_face_first.neighborhood_filter(func=np.mean, r=0.0) + # Dim order must be restored to (n_face, time) + assert filtered2.dims == ("n_face", "time") + assert filtered2.shape == (uxda.shape[0], 2) - assert filtered.chunksizes["n_face"] == (uxda.sizes["n_face"],) - np.testing.assert_allclose(filtered.compute().values, expected) + """Tests for ``UxDataArray.neighborhood_filter``.""" - def test_chunked_over_time_is_not_rechunked(self, gridpath, datasetpath): - """Chunking a non-grid dimension is the supported case and should pass - through untouched, with no warning.""" + @pytest.fixture + def vortex(self, gridpath, datasetpath): + """The ``psi`` field on outCSne30, which most of these tests filter.""" uxds = ux.open_dataset( gridpath("ugrid", "outCSne30", "outCSne30.ug"), datasetpath("ugrid", "outCSne30", "outCSne30_vortex.nc"), ) - psi = uxds["psi"] - stacked = UxDataArray( - np.tile(psi.values, (6, 1)), - dims=["time", "n_face"], - uxgrid=psi.uxgrid, - name="psi", - ).chunk({"time": 2, "n_face": -1}) - - with warnings.catch_warnings(): - warnings.simplefilter("error", UserWarning) - filtered = stacked.neighborhood_filter(func=np.mean, r=2.0) + return uxds["psi"] - assert filtered.chunksizes["time"] == (2, 2, 2) - expected = psi.neighborhood_filter(func=np.mean, r=2.0).values - np.testing.assert_allclose(filtered.compute().values, np.tile(expected, (6, 1))) - - # Every named reduction, paired with the NumPy expression it must equal. - # The reference goes through the generic callable path, so this checks the - # compiled kernel against the loop it bypasses. + # Every named reduction, with the NumPy expression it must equal and the + # parameter it takes. The reference runs through the generic callable path, + # so this pins each compiled kernel against the loop it bypasses. NAMED_REDUCTIONS = [ ("mean", {}, lambda a, axis: np.mean(a, axis=axis)), ("sum", {}, lambda a, axis: np.sum(a, axis=axis)), @@ -379,100 +359,72 @@ def test_chunked_over_time_is_not_rechunked(self, gridpath, datasetpath): ("max", {}, lambda a, axis: np.max(a, axis=axis)), ("median", {}, lambda a, axis: np.median(a, axis=axis)), ("ptp", {}, lambda a, axis: np.ptp(a, axis=axis)), - ("std", {}, lambda a, axis: np.std(a, axis=axis)), ("std", {"ddof": 1}, lambda a, axis: np.std(a, axis=axis, ddof=1)), - ("var", {}, lambda a, axis: np.var(a, axis=axis)), ("var", {"ddof": 1}, lambda a, axis: np.var(a, axis=axis, ddof=1)), ("quantile", {"q": 0.9}, lambda a, axis: np.quantile(a, 0.9, axis=axis)), ("percentile", {"q": 90}, lambda a, axis: np.percentile(a, 90, axis=axis)), - ("percentile", {"q": 50}, lambda a, axis: np.percentile(a, 50, axis=axis)), ] @pytest.mark.parametrize("name,kwargs,reference", NAMED_REDUCTIONS) - def test_named_reduction_matches_numpy( - self, name, kwargs, reference, gridpath, datasetpath - ): - """Every compiled reduction must agree with its NumPy equivalent, in - 1-D and with an extra dimension (which exercises the gufunc's - broadcast loop).""" - uxds = ux.open_dataset( - gridpath("ugrid", "outCSne30", "outCSne30.ug"), - datasetpath("ugrid", "outCSne30", "outCSne30_vortex.nc"), - ) - uxda = uxds["psi"] - nb = uxda.uxgrid.neighborhoods(r=3.0) - - got = nb.reduce(uxda, name, **kwargs) - expected = nb.reduce(uxda, reference) - np.testing.assert_allclose(got.values, expected.values, rtol=1e-12) - - stacked = UxDataArray( - np.vstack([uxda.values, uxda.values * -2.0]), - dims=["time", "n_face"], - uxgrid=uxda.uxgrid, - ) - got_2d = nb.reduce(stacked, name, **kwargs) - expected_2d = nb.reduce(stacked, reference) - np.testing.assert_allclose(got_2d.values, expected_2d.values, rtol=1e-12) - - @pytest.mark.parametrize( - "callable_func,name", - [ - (np.mean, "mean"), - (np.sum, "sum"), - (np.min, "min"), - (np.max, "max"), - (np.amin, "min"), - (np.amax, "max"), - (np.median, "median"), - (np.std, "std"), - (np.var, "var"), - (np.ptp, "ptp"), - ], - ) - def test_callable_alias_takes_kernel_path( - self, callable_func, name, gridpath, datasetpath - ): - """Code written against the original ``func=np.mean`` signature must - keep working, and must reach the same kernel the name does rather than - silently dropping to the generic loop.""" + def test_named_reduction_matches_numpy(self, name, kwargs, reference): + """Each compiled reduction must equal its NumPy expression, including + where NaN lands. + + The field is partly masked on purpose. NaN handling is the easy thing + to get wrong in a kernel: a hand-written ``if value > best`` loop skips + NaN where ``np.max`` propagates it, and numba's ``np.median`` + propagates only depending on where the NaN falls in its partition. The + extra leading dimension exercises the gufunc's broadcast loop. + """ + uxgrid = ux.Grid.from_healpix(zoom=2) + rng = np.random.default_rng(0) + values = rng.random((3, uxgrid.n_face)) + # mask a tenth of the faces, as a land/ocean mask would + values[:, rng.choice(uxgrid.n_face, uxgrid.n_face // 10, replace=False)] = np.nan + uxda = UxDataArray(values, dims=["time", "n_face"], uxgrid=uxgrid, name="masked") + + nb = uxgrid.neighborhoods(r=20.0) + got = nb.reduce(uxda, name, **kwargs).values + with warnings.catch_warnings(): + warnings.simplefilter("ignore", RuntimeWarning) # numpy all-NaN slices + expected = nb.reduce(uxda, reference).values + + assert np.isnan(got).any(), "expected a neighborhood to hit a masked value" + assert not np.isnan(got).all(), "expected some neighborhood to be clean" + np.testing.assert_array_equal(np.isnan(got), np.isnan(expected)) + finite = ~np.isnan(expected) + np.testing.assert_allclose(got[finite], expected[finite], rtol=1e-12) + + def test_callable_still_accepted(self, vortex): + """The original ``func=np.mean`` signature keeps working, and reaches + the same kernel the name does rather than dropping to the generic + loop.""" from uxarray.grid.neighbors import _CALLABLE_ALIASES, _resolve_reduction - assert _CALLABLE_ALIASES[callable_func] == name - assert _resolve_reduction(callable_func, {})[0] is _resolve_reduction(name, {})[0] + for callable_func, name in _CALLABLE_ALIASES.items(): + assert ( + _resolve_reduction(callable_func, {})[0] + is _resolve_reduction(name, {})[0] + ), f"{callable_func} should reach the {name!r} kernel" - uxds = ux.open_dataset( - gridpath("ugrid", "outCSne30", "outCSne30.ug"), - datasetpath("ugrid", "outCSne30", "outCSne30_vortex.nc"), - ) - uxda = uxds["psi"] np.testing.assert_allclose( - uxda.neighborhood_filter(callable_func, r=3.0).values, - uxda.neighborhood_filter(name, r=3.0).values, + vortex.neighborhood_filter(np.mean, r=3.0).values, + vortex.neighborhood_filter("mean", r=3.0).values, ) - def test_callable_escape_hatch(self, gridpath, datasetpath): - """An arbitrary callable on the ``axis=-1`` contract still works, for - reductions with no compiled kernel.""" - from functools import partial + def test_callable_escape_hatch(self, vortex): + """A user's own function, with no compiled equivalent, still works on + the ``axis=-1`` contract. - uxds = ux.open_dataset( - gridpath("ugrid", "outCSne30", "outCSne30.ug"), - datasetpath("ugrid", "outCSne30", "outCSne30_vortex.nc"), - ) - uxda = uxds["psi"] - - # partial() is opaque to name dispatch, so this exercises the loop - via_partial = uxda.neighborhood_filter(partial(np.percentile, q=90), r=3.0) - via_name = uxda.neighborhood_filter("percentile", r=3.0, q=90) - np.testing.assert_allclose(via_partial.values, via_name.values, rtol=1e-12) + ``functools.partial`` is covered by test_custom_func_with_partial. + """ # a user's own function, with no NumPy equivalent at all def rms(values, axis): return np.sqrt(np.mean(values**2, axis=axis)) - filtered = uxda.neighborhood_filter(rms, r=3.0) - assert filtered.shape == uxda.shape + filtered = vortex.neighborhood_filter(rms, r=3.0) + assert filtered.shape == vortex.shape assert np.all(filtered.values >= 0) @pytest.mark.parametrize( @@ -482,7 +434,6 @@ def rms(values, axis): ("mean", {"q": 90}, TypeError, "unexpected keyword argument"), ("quantile", {}, TypeError, "requires the 'q' keyword"), ("quantile", {"q": 90}, ValueError, "between 0 and 1"), - ("percentile", {"q": 1.5e3}, ValueError, "between 0 and 100"), (42, {}, TypeError, "name of a reduction or a callable"), ], ) @@ -496,70 +447,29 @@ def test_invalid_reduction(self, func, kwargs, error, match): with pytest.raises(error, match=match): uxda.neighborhood_filter(func, r=1.0, **kwargs) - def test_neighborhoods_reuse_matches_one_shot(self, gridpath, datasetpath): - """Reducing over a reused Neighborhoods must equal the one-shot filter, - which is the whole point of being able to hold onto it.""" - uxds = ux.open_dataset( - gridpath("ugrid", "outCSne30", "outCSne30.ug"), - datasetpath("ugrid", "outCSne30", "outCSne30_vortex.nc"), - ) - uxda = uxds["psi"] - nb = uxda.uxgrid.neighborhoods(r=4.0) - - assert nb.r == 4.0 - assert nb.on == "face centers" - assert nb.grid_dim == "n_face" - assert nb.grid is uxda.uxgrid - assert "face centers" in repr(nb) - - for name, kwargs in [("mean", {}), ("std", {}), ("percentile", {"q": 90})]: - np.testing.assert_allclose( - nb.reduce(uxda, name, **kwargs).values, - uxda.neighborhood_filter(name, r=4.0, **kwargs).values, - rtol=1e-12, - ) + def test_neighborhoods_reuse(self, vortex): + """A reused Neighborhoods must give the same answer as the one-shot + filter -- the point of holding onto it is that it costs one query.""" + nb = vortex.uxgrid.neighborhoods(r=4.0) - def test_neighborhoods_n_neighbors(self): - """``n_neighbors`` reports the neighborhood sizes as a grid-mapped - field, which is how you see a radius sampling a mesh unevenly.""" - uxgrid = ux.Grid.from_healpix(zoom=2) - nb = uxgrid.neighborhoods(r=15.0) + assert (nb.r, nb.on, nb.grid_dim) == (4.0, "face centers", "n_face") counts = nb.n_neighbors - assert counts.dims == ("n_face",) - assert counts.sizes["n_face"] == uxgrid.n_face # every element is its own neighbor, so no neighborhood is ever empty assert counts.min() >= 1 - # a bigger radius can only add neighbors - wider = uxgrid.neighborhoods(r=30.0).n_neighbors - assert np.all(wider.values >= counts.values) - - @pytest.mark.parametrize( - "on,dim", [("face centers", "n_face"), ("nodes", "n_node"), ("edge centers", "n_edge")] - ) - def test_neighborhoods_locations(self, on, dim): - """Neighborhoods can be built on any of the three grid locations.""" - uxgrid = ux.Grid.from_healpix(zoom=1) - # `getattr(uxgrid, dim)` populates the location's coordinates; a - # HEALPix grid cannot currently populate node coordinates lazily from - # inside the tree build (see _populate_healpix_boundaries). - data = np.arange(getattr(uxgrid, dim), dtype=float) - uxda = UxDataArray(data, dims=[dim], uxgrid=uxgrid) - - nb = uxgrid.neighborhoods(r=30.0, on=on) - assert nb.grid_dim == dim - np.testing.assert_allclose( - nb.reduce(uxda, "mean").values, - uxda.neighborhood_filter("mean", r=30.0).values, - ) - - def test_neighborhoods_rejects_mismatched_data(self): - """Reducing data mapped somewhere else must fail loudly rather than - indexing into the wrong element set.""" - from uxarray.errors import DataCenteringError + for name, kwargs in [("mean", {}), ("percentile", {"q": 90})]: + np.testing.assert_allclose( + nb.reduce(vortex, name, **kwargs).values, + vortex.neighborhood_filter(name, r=4.0, **kwargs).values, + rtol=1e-12, + ) + def test_neighborhoods_reject_wrong_data(self): + """Reducing data mapped elsewhere must fail loudly rather than index + into the wrong element set.""" uxgrid = ux.Grid.from_healpix(zoom=1) + _ = uxgrid.n_node # populate node coords before the tree is built nb = uxgrid.neighborhoods(r=30.0, on="face centers") node_data = UxDataArray( @@ -568,7 +478,6 @@ def test_neighborhoods_rejects_mismatched_data(self): with pytest.raises(DataCenteringError, match="reduce over 'n_face'"): nb.reduce(node_data, "mean") - # right dimension name, wrong grid other = ux.Grid.from_healpix(zoom=2) wrong_size = UxDataArray( np.arange(other.n_face, dtype=float), dims=["n_face"], uxgrid=other @@ -576,64 +485,67 @@ def test_neighborhoods_rejects_mismatched_data(self): with pytest.raises(DataCenteringError, match="different grid"): nb.reduce(wrong_size, "mean") - def test_neighborhoods_invalid_location(self): - uxgrid = ux.Grid.from_healpix(zoom=1) with pytest.raises(ValueError, match="Invalid `on`"): uxgrid.neighborhoods(r=1.0, on="face_centers") - def test_float32_input(self, gridpath, datasetpath): - """float32 fields take the compiled kernel's float32 signature and - still produce float64 output, as the generic path does.""" - uxds = ux.open_dataset( - gridpath("ugrid", "outCSne30", "outCSne30.ug"), - datasetpath("ugrid", "outCSne30", "outCSne30_vortex.nc"), - ) - uxda = uxds["psi"] - # `.astype` on a UxDataArray returns a plain DataArray, so rebuild - uxda32 = UxDataArray( - uxda.values.astype(np.float32), dims=uxda.dims, uxgrid=uxda.uxgrid + def test_dask_input_stays_lazy(self, vortex): + """Lazy input stays lazy: the grid dimension is a core dimension, but + the others stay chunked and unevaluated.""" + eager = vortex.neighborhood_filter("mean", r=2.0) + + stacked = UxDataArray( + np.tile(vortex.values, (6, 1)), + dims=["time", "n_face"], + uxgrid=vortex.uxgrid, + name="psi", + ).chunk({"time": 2, "n_face": -1}) + + # chunking a non-grid dimension is the supported case: untouched, silent + with warnings.catch_warnings(): + warnings.simplefilter("error", UserWarning) + filtered = stacked.neighborhood_filter("mean", r=2.0) + + assert filtered.chunks is not None, "the filter should not force a compute" + assert filtered.chunksizes["time"] == (2, 2, 2) + assert isinstance(filtered, UxDataArray) + np.testing.assert_allclose( + filtered.compute().values, np.tile(eager.values, (6, 1)) ) - filtered32 = uxda32.neighborhood_filter(func=np.mean, r=2.0) - filtered64 = uxda.neighborhood_filter(func=np.mean, r=2.0) - assert filtered32.dtype == np.float64 - np.testing.assert_allclose(filtered32.values, filtered64.values, rtol=1e-6) + def test_grid_dim_chunks_are_collapsed_with_warning(self, vortex): + """A neighborhood may span the whole grid, so the grid dimension cannot + stay chunked. Collapsing it undoes a memory decision the user made, so + it is not done silently.""" + expected = vortex.neighborhood_filter("mean", r=2.0).values - def test_integer_input_is_promoted(self): - """Integer fields have no kernel signature and must be promoted rather - than raising.""" - uxgrid = ux.Grid.from_healpix(zoom=1) - data = np.arange(uxgrid.n_face) - uxda = UxDataArray(data, dims=["n_face"], uxgrid=uxgrid, name="int_var") + uxda = vortex.chunk({"n_face": 1000}) + assert len(uxda.chunksizes["n_face"]) > 1 - filtered = uxda.neighborhood_filter(func=np.mean, r=0.0) - assert filtered.dtype == np.float64 - np.testing.assert_allclose(filtered.values, data) + with pytest.warns(UserWarning, match="Rechunking 'n_face'"): + filtered = uxda.neighborhood_filter("mean", r=2.0) - def test_auto_transpose_direct_on_uxdataarray(self, gridpath, datasetpath): - """Calling neighborhood_filter directly on a (time, n_face) UxDataArray - (without going through UxDataset) should preserve the original dim order.""" - uxds = ux.open_dataset( - gridpath("ugrid", "outCSne30", "outCSne30.ug"), - datasetpath("ugrid", "outCSne30", "outCSne30_vortex.nc"), - ) - uxda = uxds["psi"] + assert filtered.chunksizes["n_face"] == (uxda.sizes["n_face"],) + np.testing.assert_allclose(filtered.compute().values, expected) - # Build a multi-dim UxDataArray with time as the FIRST (non-grid) dim - data = np.stack([uxda.values, uxda.values * 2.0]) # shape (2, n_face) - uxda_time = UxDataArray( - data, dims=["time", "n_face"], uxgrid=uxda.uxgrid, name="psi_time" + def test_output_is_always_float64(self, vortex): + """float32 hits the kernel's float32 signature and integers have no + signature at all; both must come back as float64, as the generic path + does by writing into a float64 output.""" + uxgrid = ux.Grid.from_healpix(zoom=1) + integers = UxDataArray( + np.arange(uxgrid.n_face), dims=["n_face"], uxgrid=uxgrid, name="int_var" ) + filtered = integers.neighborhood_filter("mean", r=0.0) + assert filtered.dtype == np.float64 + np.testing.assert_allclose(filtered.values, integers.values) - # n_face is already last: no transpose needed internally - filtered = uxda_time.neighborhood_filter(func=np.mean, r=0.0) - assert filtered.dims == ("time", "n_face") - assert filtered.shape == (2, uxda.shape[0]) - np.testing.assert_allclose(filtered.values, data) + as_float32 = UxDataArray( + vortex.values.astype(np.float32), dims=vortex.dims, uxgrid=vortex.uxgrid + ) + filtered32 = as_float32.neighborhood_filter("mean", r=2.0) + assert filtered32.dtype == np.float64 + np.testing.assert_allclose( + filtered32.values, vortex.neighborhood_filter("mean", r=2.0).values, + rtol=1e-6, + ) - # Also test with a UxDataArray that has grid dim NOT last (n_face, time) - uxda_face_first = uxda_time.transpose("n_face", "time") - filtered2 = uxda_face_first.neighborhood_filter(func=np.mean, r=0.0) - # Dim order must be restored to (n_face, time) - assert filtered2.dims == ("n_face", "time") - assert filtered2.shape == (uxda.shape[0], 2) diff --git a/test/core/test_dataset.py b/test/core/test_dataset.py index 34bfd23d1..c91922dbc 100644 --- a/test/core/test_dataset.py +++ b/test/core/test_dataset.py @@ -207,77 +207,43 @@ def test_non_grid_variable_skipped(self): nt.assert_allclose(filtered["face_var"].values, uxds["face_var"].values) nt.assert_allclose(filtered["scalar_var"].values, uxds["scalar_var"].values) - def test_named_reduction_with_parameter(self, gridpath, datasetpath): - """Named reductions and their parameters reach every variable.""" - uxds = ux.open_dataset( - gridpath("ugrid", "outCSne30", "outCSne30.ug"), - datasetpath("ugrid", "outCSne30", "outCSne30_vortex.nc"), - ) - - filtered_ds = uxds.neighborhood_filter("percentile", r=5.0, q=90) - filtered_da = uxds["psi"].neighborhood_filter("percentile", r=5.0, q=90) - - nt.assert_allclose(filtered_ds["psi"].values, filtered_da.values) - - def test_one_query_shared_across_variables(self): - """Variables mapped to the same location must share a single neighbor - query. The query dominates the cost, so per-variable rebuilding would - make a dataset filter scale with the number of variables.""" - from unittest.mock import patch - - import uxarray.grid.neighbors as neighbors - - uxgrid = ux.Grid.from_healpix(zoom=2) - rng = np.random.default_rng(0) - uxds = UxDataset( - data_vars={ - f"v{i}": ("n_face", rng.random(uxgrid.n_face)) for i in range(5) - }, - uxgrid=uxgrid, - ) - - real = neighbors._csr_neighbors - with patch.object( - neighbors, "_csr_neighbors", side_effect=real - ) as spy: - filtered = uxds.neighborhood_filter("mean", r=20.0) - - assert spy.call_count == 1, ( - f"expected one neighbor query for 5 same-location variables, " - f"got {spy.call_count}" - ) - - expected = uxds["v0"].neighborhood_filter("mean", r=20.0) - nt.assert_allclose(filtered["v0"].values, expected.values) + def test_one_query_per_grid_location(self): + """Variables sharing a grid location must share one neighbor query. - def test_query_per_location_not_per_variable(self): - """Variables on different grid locations each need their own query, - but only one apiece.""" + The query dominates the cost of a reduction, so rebuilding it per + variable would make a dataset filter scale with the number of + variables. Counting calls is the only way to see that from outside. + """ from unittest.mock import patch import uxarray.grid.neighbors as neighbors uxgrid = ux.Grid.from_healpix(zoom=2) - # populate node coordinates before the tree is built (see the HEALPix - # boundary-population path, which cannot do it lazily) + # touch both locations first: a HEALPix grid cannot populate node + # coordinates lazily from inside the tree build n_node, n_face = uxgrid.n_node, uxgrid.n_face rng = np.random.default_rng(0) uxds = UxDataset( data_vars={ "face_a": ("n_face", rng.random(n_face)), "face_b": ("n_face", rng.random(n_face)), + "face_c": ("n_face", rng.random(n_face)), "node_a": ("n_node", rng.random(n_node)), }, uxgrid=uxgrid, ) real = neighbors._csr_neighbors - with patch.object( - neighbors, "_csr_neighbors", side_effect=real - ) as spy: - uxds.neighborhood_filter("mean", r=20.0) + with patch.object(neighbors, "_csr_neighbors", side_effect=real) as spy: + filtered = uxds.neighborhood_filter("percentile", r=20.0, q=90) assert spy.call_count == 2, ( - f"expected one query per grid location (faces, nodes), " - f"got {spy.call_count}" + f"expected one query per grid location (faces, nodes), got " + f"{spy.call_count}" ) + # and the reduction, with its parameter, reached every variable + for name in ("face_a", "node_a"): + nt.assert_allclose( + filtered[name].values, + uxds[name].neighborhood_filter("percentile", r=20.0, q=90).values, + ) diff --git a/uxarray/core/dataset.py b/uxarray/core/dataset.py index b0579ce72..e2ac5f11e 100644 --- a/uxarray/core/dataset.py +++ b/uxarray/core/dataset.py @@ -668,18 +668,8 @@ def neighborhood_filter( element of every data variable with a reduction of all elements within a circular neighborhood of radius ``r``. - Parameters - ---------- - func : str or Callable, default="mean" - Name of the reduction to apply: "mean", "sum", "min", "max", - "median", "ptp", "std", "var", "quantile", or "percentile". Named - reductions run compiled. A callable is accepted as an escape hatch; - see :meth:`UxDataArray.neighborhood_filter`. - r : float, default=1. - Radius of the neighborhood, in degrees. - **kwargs - Parameter for the named reduction: ``q`` for "quantile" (0-1) and - "percentile" (0-100), ``ddof`` for "std" and "var". + Parameters are as for :meth:`UxDataArray.neighborhood_filter`, which + documents the available reductions and their keyword arguments. Returns ------- @@ -689,9 +679,6 @@ def neighborhood_filter( Notes ----- Variables without a grid dimension are passed through unchanged. - ``r`` is a great-circle distance in degrees, and lazy (dask-backed) - variables stay lazy. See :meth:`UxDataArray.neighborhood_filter` for - details. Variables mapped to the same grid location share one neighbor query, so filtering a dataset costs one query per location present rather than diff --git a/uxarray/grid/neighbors.py b/uxarray/grid/neighbors.py index d9e56ee51..83f78aa90 100644 --- a/uxarray/grid/neighbors.py +++ b/uxarray/grid/neighbors.py @@ -1200,124 +1200,26 @@ def _get_element_coords(grid, data_mapping: str, coordinate_system: str): # and ``(m)`` the destination axis. Output is float64 regardless of input # dtype, matching the behaviour of the generic path below. _GUFUNC_SIGNATURES = [ - "void(float64[:], int64[:], int64[:], int64[:], float64[:])", - "void(float32[:], int64[:], int64[:], int64[:], float64[:])", -] -_GUFUNC_LAYOUT = "(n),(k),(m),(m)->(m)" -_GUFUNC_KWARGS = {"nopython": True, "cache": True, "target": "parallel"} - -# Variant carrying one scalar parameter (``q``, ``ddof``) as a scalar core -# dimension, so the value travels with the call rather than being baked into a -# separately compiled kernel per value. -_GUFUNC_SIGNATURES_P = [ "void(float64[:], int64[:], int64[:], int64[:], float64, float64[:])", "void(float32[:], int64[:], int64[:], int64[:], float64, float64[:])", ] -_GUFUNC_LAYOUT_P = "(n),(k),(m),(m),()->(m)" - - -@guvectorize(_GUFUNC_SIGNATURES, _GUFUNC_LAYOUT, **_GUFUNC_KWARGS) -def _gu_mean(data, flat, starts, counts, out): - for i in range(starts.shape[0]): - count = counts[i] - if count == 0: - out[i] = np.nan - continue - start = starts[i] - acc = 0.0 - for j in range(start, start + count): - acc += data[flat[j]] - out[i] = acc / count - - -@guvectorize(_GUFUNC_SIGNATURES, _GUFUNC_LAYOUT, **_GUFUNC_KWARGS) -def _gu_sum(data, flat, starts, counts, out): - for i in range(starts.shape[0]): - start = starts[i] - acc = 0.0 - for j in range(start, start + counts[i]): - acc += data[flat[j]] - out[i] = acc - - -@guvectorize(_GUFUNC_SIGNATURES, _GUFUNC_LAYOUT, **_GUFUNC_KWARGS) -def _gu_max(data, flat, starts, counts, out): - for i in range(starts.shape[0]): - count = counts[i] - if count == 0: - out[i] = np.nan - continue - start = starts[i] - best = data[flat[start]] - for j in range(start + 1, start + count): - value = data[flat[j]] - if value > best: - best = value - out[i] = best - - -@guvectorize(_GUFUNC_SIGNATURES, _GUFUNC_LAYOUT, **_GUFUNC_KWARGS) -def _gu_min(data, flat, starts, counts, out): - for i in range(starts.shape[0]): - count = counts[i] - if count == 0: - out[i] = np.nan - continue - start = starts[i] - best = data[flat[start]] - for j in range(start + 1, start + count): - value = data[flat[j]] - if value < best: - best = value - out[i] = best - - -# Order statistics cannot stream: the neighborhood has to be materialized -# before it can be sorted or partitioned. Those reductions share one gather -# loop, and the reduction itself is whatever numba-compilable callable is -# handed to the factory. Numba supports the NumPy reductions in nopython mode, -# so ``np.median`` and friends are used directly rather than reimplemented -- -# numba's ``np.median`` partitions instead of fully sorting, which makes it -# faster than an equivalent hand-written sort as well as shorter. -def _make_buffered_kernel(reduce_fn): - """Builds a kernel that gathers each neighborhood, then calls ``reduce_fn`` - on the 1-D result. ``reduce_fn`` must be numba-compilable.""" - # A reducer shared between kernels arrives already compiled; numba rejects - # jitting a dispatcher twice. - if not hasattr(reduce_fn, "py_func"): - reduce_fn = njit(cache=True)(reduce_fn) - - @guvectorize(_GUFUNC_SIGNATURES, _GUFUNC_LAYOUT, **_GUFUNC_KWARGS) - def kernel(data, flat, starts, counts, out): - widest = 0 - for i in range(counts.shape[0]): - if counts[i] > widest: - widest = counts[i] - buffer = np.empty(widest, dtype=np.float64) - - for i in range(starts.shape[0]): - count = counts[i] - if count == 0: - out[i] = np.nan - continue - start = starts[i] - for j in range(count): - buffer[j] = data[flat[start + j]] - out[i] = reduce_fn(buffer[:count]) +_GUFUNC_LAYOUT = "(n),(k),(m),(m),()->(m)" +_GUFUNC_KWARGS = {"nopython": True, "cache": True, "target": "parallel"} - return kernel +def _make_kernel(reduce_fn): + """Builds a kernel that gathers each neighborhood, then calls + ``reduce_fn(window, param)`` on the 1-D result. -def _make_parameterized_kernel(reduce_fn): - """As ``_make_buffered_kernel``, but ``reduce_fn`` also takes one scalar - parameter (``q`` for quantiles, ``ddof`` for spreads), declared as a scalar - core dimension so it can vary per call without recompiling.""" + ``reduce_fn`` must be numba-compilable, and must be defined in a real + source file for ``cache=True`` to find it. + """ # A reducer shared between kernels arrives already compiled; numba rejects # jitting a dispatcher twice. if not hasattr(reduce_fn, "py_func"): reduce_fn = njit(cache=True)(reduce_fn) - @guvectorize(_GUFUNC_SIGNATURES_P, _GUFUNC_LAYOUT_P, **_GUFUNC_KWARGS) + @guvectorize(_GUFUNC_SIGNATURES, _GUFUNC_LAYOUT, **_GUFUNC_KWARGS) def kernel(data, flat, starts, counts, param, out): widest = 0 for i in range(counts.shape[0]): @@ -1338,6 +1240,27 @@ def kernel(data, flat, starts, counts, param, out): return kernel +class _Reduction(NamedTuple): + """A named reduction, and the single scalar parameter it accepts (if any). + + Limiting reductions to one parameter is what keeps the gufunc layout above + down to one; it covers every reduction implemented here. + """ + + kernel: object + param: str | None = None + default: float = 0.0 + + +# Reductions with a compiled kernel, addressed by name. A name always takes the +# fast path, which is why the public API documents names rather than callables: +# dispatching on a function object cannot see through ``functools.partial``, so +# a parameterized reduction could never hit a kernel that way. +# +# Adding a reduction is one line here. Reducers take ``(window, param)``; +# those with no parameter ignore the second argument. Numba keys its cache by +# code object rather than qualified name, so the identically-named lambdas do +# not collide. @njit(cache=True) def _variance(window, ddof): """Variance with a delta degrees of freedom. Numba's ``np.var`` takes no @@ -1352,42 +1275,35 @@ def _variance(window, ddof): return total / denominator -_gu_median = _make_buffered_kernel(lambda window: np.median(window)) -_gu_ptp = _make_buffered_kernel(lambda window: np.max(window) - np.min(window)) -_gu_quantile = _make_parameterized_kernel(lambda window, q: np.quantile(window, q)) -_gu_var = _make_parameterized_kernel(_variance) -_gu_std = _make_parameterized_kernel( - lambda window, ddof: np.sqrt(_variance(window, ddof)) -) - - -class _Reduction(NamedTuple): - """A named reduction, and the single scalar parameter it accepts (if any). +@njit(cache=True) +def _median(window, _): + # numba's ``np.median`` selects by partitioning, and whether a NaN survives + # that depends on where it lands -- so unlike numpy's, it propagates NaN + # only sometimes. This spelling short-circuits and allocates nothing: + # ``np.any(np.isnan(window))`` costs ~14% more, and routing through + # ``np.quantile``, which does propagate, costs 2.5x. + for value in window: + if np.isnan(value): + return np.nan + return np.median(window) - Keeping parameters to at most one keeps the gufunc layouts down to the two - declared above, which covers every reduction implemented here. - """ - - kernel: object - param: str | None = None - default: float | None = None +_quantile_kernel = _make_kernel(lambda window, q: np.quantile(window, q)) -# Reductions with a compiled kernel, addressed by name. A name always takes the -# fast path, which is why the public API documents names rather than callables: -# dispatching on a function object cannot see through ``functools.partial``, so -# a parameterized reduction could never hit a kernel that way. _REDUCTIONS = { - "mean": _Reduction(_gu_mean), - "sum": _Reduction(_gu_sum), - "min": _Reduction(_gu_min), - "max": _Reduction(_gu_max), - "median": _Reduction(_gu_median), - "ptp": _Reduction(_gu_ptp), - "std": _Reduction(_gu_std, param="ddof", default=0.0), - "var": _Reduction(_gu_var, param="ddof", default=0.0), - "quantile": _Reduction(_gu_quantile, param="q"), - "percentile": _Reduction(_gu_quantile, param="q"), + "mean": _Reduction(_make_kernel(lambda window, _: np.mean(window))), + "sum": _Reduction(_make_kernel(lambda window, _: np.sum(window))), + "min": _Reduction(_make_kernel(lambda window, _: np.min(window))), + "max": _Reduction(_make_kernel(lambda window, _: np.max(window))), + "ptp": _Reduction(_make_kernel(lambda window, _: np.max(window) - np.min(window))), + "median": _Reduction(_make_kernel(_median)), + "var": _Reduction(_make_kernel(_variance), param="ddof"), + "std": _Reduction( + _make_kernel(lambda window, ddof: np.sqrt(_variance(window, ddof))), + param="ddof", + ), + "quantile": _Reduction(_quantile_kernel, param="q"), + "percentile": _Reduction(_quantile_kernel, param="q"), } # Callables accepted for backwards compatibility, so that code written against @@ -1446,16 +1362,17 @@ def _resolve_reduction(func, kwargs): ) if reduction.param is None: - return reduction.kernel, None + # The kernel still takes a parameter; this one ignores it. + return reduction.kernel, 0.0 if reduction.param in kwargs: value = float(kwargs[reduction.param]) - elif reduction.default is not None: - value = reduction.default - else: + elif name in ("quantile", "percentile"): raise TypeError( f"Reduction {name!r} requires the {reduction.param!r} keyword argument." ) + else: + value = reduction.default # `percentile` is `quantile` on a 0-100 scale; normalize so both share one # kernel rather than compiling a near-duplicate. @@ -1727,10 +1644,7 @@ def _apply(block): # path does too by writing into a float64 output. if block.dtype not in (np.float64, np.float32): block = block.astype(np.float64) - args = (block, self._flat, self._starts, self._counts) - if param is not None: - args += (param,) - return kernel(*args) + return kernel(block, self._flat, self._starts, self._counts, param) work = _rechunk_grid_dim(uxda, grid_dim) From c240cd364df1aafba192c4c18f75f63c69c03c25 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 6 Aug 2026 15:18:18 +0000 Subject: [PATCH 21/21] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- docs/user-guide/neighborhood-filter.ipynb | 32 ++++++++++++++++++----- 1 file changed, 26 insertions(+), 6 deletions(-) diff --git a/docs/user-guide/neighborhood-filter.ipynb b/docs/user-guide/neighborhood-filter.ipynb index 8ab26d8dc..f2d7865b8 100644 --- a/docs/user-guide/neighborhood-filter.ipynb +++ b/docs/user-guide/neighborhood-filter.ipynb @@ -197,7 +197,7 @@ "print(\"max filter max :\", uxda_max.values.max())\n", "print(\"p90 filter max :\", uxda_p90.values.max())\n", "print(\"median filter max:\", uxda_med.values.max())\n", - "print(\"std filter max :\", uxda_std.values.max())\n" + "print(\"std filter max :\", uxda_std.values.max())" ] }, { @@ -219,6 +219,7 @@ }, { "cell_type": "markdown", + "id": "28d3efd5258a48a79c179ea5c6759f01", "metadata": {}, "source": [ "### Reductions Without a Name\n", @@ -232,6 +233,7 @@ { "cell_type": "code", "execution_count": null, + "id": "3f9bc0b9dd2c44919cc8dcca39b469f8", "metadata": {}, "outputs": [], "source": [ @@ -239,16 +241,20 @@ "def rms(values, axis):\n", " return np.sqrt(np.mean(values**2, axis=axis))\n", "\n", + "\n", "uxda_rms = uxda.neighborhood_filter(rms, r=5.0)\n", "\n", "# `functools.partial` also works, though `\"percentile\"` is the faster way here\n", "uxda_p90_slow = uxda.neighborhood_filter(partial(np.percentile, q=90), r=5.0)\n", - "print(\"partial matches the named reduction:\",\n", - " np.allclose(uxda_p90_slow.values, uxda_p90.values))\n" + "print(\n", + " \"partial matches the named reduction:\",\n", + " np.allclose(uxda_p90_slow.values, uxda_p90.values),\n", + ")" ] }, { "cell_type": "markdown", + "id": "0e382214b5f147d187d36a2058b9c724", "metadata": {}, "source": [ "## Reusing Neighborhoods Across Reductions\n", @@ -264,16 +270,18 @@ { "cell_type": "code", "execution_count": null, + "id": "5b09d5ef5b5e4bb6ab9b829b10b6a29f", "metadata": {}, "outputs": [], "source": [ "nb5 = uxda.uxgrid.neighborhoods(r=5.0)\n", - "nb5\n" + "nb5" ] }, { "cell_type": "code", "execution_count": null, + "id": "a50416e276a0479cbe66534ed1713a40", "metadata": {}, "outputs": [], "source": [ @@ -282,11 +290,15 @@ "spread = nb5.reduce(uxda, \"std\")\n", "p90 = nb5.reduce(uxda, \"percentile\", q=90)\n", "\n", - "print(\"identical to the one-shot call:\", np.allclose(smooth.values, uxda.neighborhood_filter(\"mean\", r=5.0).values))\n" + "print(\n", + " \"identical to the one-shot call:\",\n", + " np.allclose(smooth.values, uxda.neighborhood_filter(\"mean\", r=5.0).values),\n", + ")" ] }, { "cell_type": "markdown", + "id": "46a27a456b804aa2a380d5edf15a5daf", "metadata": {}, "source": [ "`n_neighbors` reports how many elements fell inside each neighborhood. On a\n", @@ -297,11 +309,19 @@ { "cell_type": "code", "execution_count": null, + "id": "1944c39560714e6e80c856f20744a8e5", "metadata": {}, "outputs": [], "source": [ "counts = nb5.n_neighbors\n", - "print(\"neighbors per face: min\", int(counts.min()), \" max\", int(counts.max()), \" mean\", float(counts.mean()).__round__(1))\n" + "print(\n", + " \"neighbors per face: min\",\n", + " int(counts.min()),\n", + " \" max\",\n", + " int(counts.max()),\n", + " \" mean\",\n", + " float(counts.mean()).__round__(1),\n", + ")" ] }, {