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..7fae09502 --- /dev/null +++ b/docs/user-guide/neighborhood-filter.ipynb @@ -0,0 +1,377 @@ +{ + "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": "## 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\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` — 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 b961a5a37..7ebc9f01d 100644 --- a/test/core/test_dataarray.py +++ b/test/core/test_dataarray.py @@ -163,3 +163,146 @@ 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) + 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/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 177f888bb..15b57b6b8 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, @@ -28,6 +29,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 @@ -118,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 @@ -1989,6 +1992,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( @@ -2188,6 +2192,91 @@ 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. + + 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 ValueError( + 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}." + ) + + # 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 4331dd30d..b220cd00d 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.formatting_html import dataset_repr @@ -655,6 +656,63 @@ 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. + + 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 f7d5e828c..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" @@ -1776,7 +1776,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, @@ -1784,9 +1789,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 @@ -1876,7 +1878,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, @@ -1885,10 +1892,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..7dfe42e00 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,111 @@ 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``. + """ + + 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) + + # 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. + for i, idx in enumerate(neighbor_indices): + if len(idx): + destination_data[..., i] = func(data[..., idx], axis=-1) + + return destination_data