From 75d832e36418fb020448d4f53f1fbd0cd4bc66f0 Mon Sep 17 00:00:00 2001 From: ahijevyc Date: Sat, 7 Sep 2024 11:09:00 -0600 Subject: [PATCH 1/9] 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 2/9] 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 3/9] 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 4/9] 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 5/9] 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 6/9] 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 7/9] [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 8/9] 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 9/9] 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