Apply function to points within circular neighborhood

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::

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)

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)

uxarray/core/dataarray.py

@@ -2,7 +2,7 @@
 
 import warnings
 from html import escape
-from typing import TYPE_CHECKING, Any, Hashable, Literal, Mapping, Optional
+from typing import TYPE_CHECKING, Any, Callable, Hashable, Literal, Mapping, Optional
 from warnings import warn
 
 import cartopy.crs as ccrs
@@ -14,6 +14,7 @@
 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,
@@ -30,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
@@ -1986,6 +1988,7 @@ def isel(
         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(
@@ -2185,6 +2188,60 @@ def get_dual(self):
 
         return uxda
 
+    def neighborhood_filter(
+        self,
+        func: Callable = np.mean,
+        r: float = 1.0,
+    ) -> UxDataArray:
+        """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. 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
+        -------
+        uxda_filter : UxDataArray
+            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(
+                "Data_mapping is not face, node, or edge. Could not define data_mapping."
+            )
+
+        # 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}"
+        )
+
+        destination_data = _neighborhood_filter(
+            self.uxgrid, self.data, data_mapping, func=func, r=r
+        )
+
+        # Construct UxDataArray for filtered variable.
+        uxda_filter = self._copy()
+
+        uxda_filter.data = destination_data
+
+        return uxda_filter
+
     def __getattribute__(self, name):
         """Intercept accessor method calls to return Ux-aware accessors."""
         # Lazy import to avoid circular imports

uxarray/core/dataset.py

@@ -3,7 +3,7 @@
 import os
 import sys
 from html import escape
-from typing import IO, Any, Hashable, Mapping
+from typing import IO, Any, Callable, Hashable, Mapping
 from warnings import warn
 
 import numpy as np
@@ -13,6 +13,7 @@
 from xarray.core.utils import UncachedAccessor
 
 import uxarray
+from uxarray.constants import GRID_DIMS
 from uxarray.core.dataarray import UxDataArray
 from uxarray.core.utils import _map_dims_to_ugrid, _open_dataset_with_fallback
 from uxarray.formatting_html import dataset_repr
@@ -655,6 +656,52 @@ def to_array(
         xarr = super().to_array(dim=dim, name=name)
         return UxDataArray(xarr, uxgrid=self.uxgrid)
 
+    def neighborhood_filter(
+        self,
+        func: Callable = np.mean,
+        r: float = 1.0,
+    ) -> 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, 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()
+        # 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 to_xarray(self, grid_format: str = "UGRID") -> xr.Dataset:
         """
         Converts a ``ux.UXDataset`` to a ``xr.Dataset``.

uxarray/grid/grid.py

@@ -1772,17 +1772,19 @@ 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,
                 distance_metric=distance_metric,
                 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(