diff --git a/docs/api.rst b/docs/api.rst index 2addaabd8..9508b8427 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/docs/user-guide/neighborhood-filter.ipynb b/docs/user-guide/neighborhood-filter.ipynb new file mode 100644 index 000000000..4e2b7d980 --- /dev/null +++ b/docs/user-guide/neighborhood-filter.ipynb @@ -0,0 +1,388 @@ +{ + "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—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", + "\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\")\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", + ")" + ] + }, + { + "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" + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7623eae2785240b9bd12b16a66d81610", + "metadata": {}, + "outputs": [], + "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°)\",\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° 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\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 — 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(\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\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\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\")\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 × 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)\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\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)\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": [ + "## 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", + "execution_count": null, + "id": "f5a1fa73e5044315a093ec459c9be902", + "metadata": {}, + "outputs": [], + "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 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: 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", + "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` — 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": { + "kernelspec": { + "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 +} 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 3ad3376ae..5177362dc 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 @@ -164,3 +164,189 @@ 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(DataCenteringError): + uxda.neighborhood_filter(func=np.mean, r=1.0) + + 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") + + # 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)) + 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_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.""" + 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/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/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 8f41a1627..c77fbf628 100644 --- a/uxarray/core/dataarray.py +++ b/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 numpy as np @@ -12,6 +12,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, @@ -29,6 +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.validation import _check_duplicate_nodes_indices from uxarray.io._healpix import get_zoom_from_cells from uxarray.plot.accessor import UxDataArrayPlotAccessor @@ -2193,6 +2195,108 @@ 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 the neighborhood, in degrees. + + Returns + ------- + 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: + + >>> 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(): + 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}." + ) + + # 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 + ) + + # 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) + + return uxda_filter + def __getattribute__(self, name): """Intercept accessor method calls to return Ux-aware accessors.""" # Lazy import to avoid circular imports diff --git a/uxarray/core/dataset.py b/uxarray/core/dataset.py index 25a16a13a..382c0bcfb 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, 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.errors import DimensionError @@ -656,6 +657,70 @@ 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 the neighborhood, in degrees. + + Returns + ------- + 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: + + >>> 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() + # 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. + 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) + + return destination_uxds + def to_xarray(self, grid_format: str = "UGRID") -> xr.Dataset: """ Converts a ``ux.UXDataset`` to a ``xr.Dataset``. diff --git a/uxarray/grid/grid.py b/uxarray/grid/grid.py index 27c65e183..6ba2d7baa 100644 --- a/uxarray/grid/grid.py +++ b/uxarray/grid/grid.py @@ -1760,7 +1760,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" @@ -1777,7 +1777,17 @@ def get_ball_tree( BallTree instance """ - if self._ball_tree is None or reconstruct: + # 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( self, coordinates=coordinates, @@ -1785,9 +1795,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 @@ -1877,7 +1884,15 @@ def get_kd_tree( KDTree instance """ - if self._kd_tree is None or reconstruct: + # 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( self, coordinates=coordinates, @@ -1886,10 +1901,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..0ce83bc5f 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,142 @@ 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, in degrees. + + Returns + ------- + destination_data : np.ndarray + Filtered data, matching the shape of ``data``. + + 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. + """ + + # 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) + + neighbor_indices = tree.query_radius(dest_coords, r=r) + + # 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 + + return destination_data