diff --git a/docs/api.rst b/docs/api.rst index 2addaabd8..4666c0efa 100644 --- a/docs/api.rst +++ b/docs/api.rst @@ -194,6 +194,7 @@ Methods Grid.compute_face_areas Grid.construct_face_centers Grid.get_ball_tree + Grid.neighborhoods Grid.get_kd_tree Grid.get_spatial_hash Grid.get_faces_containing_point @@ -565,6 +566,30 @@ 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 + +Finding the neighbors is usually more expensive than reducing over them. To apply several +reductions at one radius, build the neighborhoods once and reduce over them repeatedly. + +.. autosummary:: + :toctree: generated/ + + Grid.neighborhoods + uxarray.grid.neighbors.Neighborhoods + uxarray.grid.neighbors.Neighborhoods.reduce + uxarray.grid.neighbors.Neighborhoods.n_neighbors + + Zonal Average ~~~~~~~~~~~~~ .. autosummary:: diff --git a/docs/user-guide/neighborhood-filter.ipynb b/docs/user-guide/neighborhood-filter.ipynb new file mode 100644 index 000000000..f2d7865b8 --- /dev/null +++ b/docs/user-guide/neighborhood-filter.ipynb @@ -0,0 +1,548 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "7fb27b941602401d91542211134fc71a", + "metadata": {}, + "source": [ + "# Neighborhood Filter\n", + "\n", + "A **neighborhood filter** replaces the value at each grid element with a\n", + "reduction of all grid elements whose centers fall within\n", + "a circular neighborhood of radius `r` degrees around that element.\n", + "\n", + "Unlike a fixed *k*-nearest-neighbor average, a radius-based filter imposes a\n", + "consistent spatial scale across the whole mesh—useful for variable-resolution grids\n", + "where the number of neighbors varies from region to region.\n", + "\n", + "**Supported element types:** face-centered, node-centered, and edge-centered data.\n", + "\n", + "**API at a glance:**\n", + "\n", + "| Object | Method |\n", + "|---|---|\n", + "| `UxDataArray` | `da.neighborhood_filter(\"mean\", r=5.0)` |\n", + "| `UxDataset` | `ds.neighborhood_filter(\"mean\", r=5.0)` |\n", + "\n", + "The returned object is always the same type as the input, with the same grid, dims,\n", + "coordinates, name, and attributes preserved.\n" + ] + }, + { + "cell_type": "markdown", + "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", + "\n", + "Calling `neighborhood_filter` with `\"mean\"` and a radius of 5° replaces\n", + "each face value with the mean of all face centers within 5° of that face's center.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7623eae2785240b9bd12b16a66d81610", + "metadata": {}, + "outputs": [], + "source": [ + "uxda_smooth = uxda.neighborhood_filter(\"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(\"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": [ + "## Other Reductions\n", + "\n", + "Pass the name of the reduction you want. The available names are `mean`,\n", + "`sum`, `min`, `max`, `median`, `ptp`, `std`, `var`, `quantile`, and\n", + "`percentile`. Reductions that take a parameter receive it as a keyword\n", + "argument: `q` for `quantile` (0–1) and `percentile` (0–100), `ddof` for\n", + "`std` and `var`.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "b43b363d81ae4b689946ece5c682cd59", + "metadata": {}, + "outputs": [], + "source": [ + "# 90th-percentile filter — highlights local maxima\n", + "uxda_p90 = uxda.neighborhood_filter(\"percentile\", r=5.0, q=90)\n", + "\n", + "# Maximum filter\n", + "uxda_max = uxda.neighborhood_filter(\"max\", r=5.0)\n", + "\n", + "# Median filter — robust to outliers\n", + "uxda_med = uxda.neighborhood_filter(\"median\", r=5.0)\n", + "\n", + "# Local spread, as a sample standard deviation\n", + "uxda_std = uxda.neighborhood_filter(\"std\", r=5.0, ddof=1)\n", + "\n", + "print(\"max filter max :\", uxda_max.values.max())\n", + "print(\"p90 filter max :\", uxda_p90.values.max())\n", + "print(\"median filter max:\", uxda_med.values.max())\n", + "print(\"std filter max :\", uxda_std.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": "28d3efd5258a48a79c179ea5c6759f01", + "metadata": {}, + "source": [ + "### Reductions Without a Name\n", + "\n", + "If you need something not in that list, pass a callable instead. It is applied as\n", + "`func(values, axis=-1)` to a block whose last axis is the neighborhood, once per\n", + "grid element, in Python — noticeably slower than a named reduction, so reach for it\n", + "only when no name fits.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3f9bc0b9dd2c44919cc8dcca39b469f8", + "metadata": {}, + "outputs": [], + "source": [ + "# a root-mean-square filter, which has no named equivalent\n", + "def rms(values, axis):\n", + " return np.sqrt(np.mean(values**2, axis=axis))\n", + "\n", + "\n", + "uxda_rms = uxda.neighborhood_filter(rms, r=5.0)\n", + "\n", + "# `functools.partial` also works, though `\"percentile\"` is the faster way here\n", + "uxda_p90_slow = uxda.neighborhood_filter(partial(np.percentile, q=90), r=5.0)\n", + "print(\n", + " \"partial matches the named reduction:\",\n", + " np.allclose(uxda_p90_slow.values, uxda_p90.values),\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "0e382214b5f147d187d36a2058b9c724", + "metadata": {}, + "source": [ + "## Reusing Neighborhoods Across Reductions\n", + "\n", + "Each call to `neighborhood_filter` searches the grid for the neighbors of every\n", + "element. That search usually costs far more than the reduction itself, so\n", + "applying several reductions at the same radius repeats the expensive part.\n", + "\n", + "`Grid.neighborhoods` does the search once and returns an object you can reduce\n", + "over as many times as you like.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5b09d5ef5b5e4bb6ab9b829b10b6a29f", + "metadata": {}, + "outputs": [], + "source": [ + "nb5 = uxda.uxgrid.neighborhoods(r=5.0)\n", + "nb5" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a50416e276a0479cbe66534ed1713a40", + "metadata": {}, + "outputs": [], + "source": [ + "# the search is already done; each of these only runs a reduction\n", + "smooth = nb5.reduce(uxda, \"mean\")\n", + "spread = nb5.reduce(uxda, \"std\")\n", + "p90 = nb5.reduce(uxda, \"percentile\", q=90)\n", + "\n", + "print(\n", + " \"identical to the one-shot call:\",\n", + " np.allclose(smooth.values, uxda.neighborhood_filter(\"mean\", r=5.0).values),\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "46a27a456b804aa2a380d5edf15a5daf", + "metadata": {}, + "source": [ + "`n_neighbors` reports how many elements fell inside each neighborhood. On a\n", + "variable-resolution mesh this varies by region, which is worth checking before\n", + "reading too much into a filtered field.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1944c39560714e6e80c856f20744a8e5", + "metadata": {}, + "outputs": [], + "source": [ + "counts = nb5.n_neighbors\n", + "print(\n", + " \"neighbors per face: min\",\n", + " int(counts.min()),\n", + " \" max\",\n", + " int(counts.max()),\n", + " \" mean\",\n", + " float(counts.mean()).__round__(1),\n", + ")" + ] + }, + { + "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(\"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(\"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(\"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(\"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(\"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(\"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(\"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(\"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", + "\n", + "See also:\n", + "\n", + "- {py:meth}`uxarray.UxDataArray.neighborhood_filter`\n", + "- {py:meth}`uxarray.UxDataset.neighborhood_filter`\n", + "- {py:meth}`uxarray.Grid.neighborhoods` — reusable neighborhoods\n", + "\n", + "Related methods that apply aggregations across different grid element types:\n", + "\n", + "- {py:meth}`uxarray.UxDataArray.topological_mean` — aggregate node→face, node→edge, etc.\n", + "- {py:meth}`uxarray.UxDataArray.zonal_mean` — latitude-band averages\n", + "- {py:meth}`uxarray.UxDataArray.azimuthal_mean` — rings of constant great-circle distance\n" + ] + } + ], + "metadata": { + "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 5932809dc..021709743 100644 --- a/test/core/test_dataarray.py +++ b/test/core/test_dataarray.py @@ -1,6 +1,7 @@ +import warnings import numpy as np import uxarray as ux -from uxarray.errors import DimensionError, GridInvalidError +from uxarray.errors import DataCenteringError, DimensionError, GridInvalidError from uxarray.grid.geometry import _build_polygon_shells, _build_corrected_polygon_shells from uxarray.core.dataset import UxDataset, UxDataArray import pytest @@ -166,6 +167,176 @@ def test_data_location(): 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_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) def test_uxgrid_None_is_invalid_in_uxdataarray(): """Ensures GridInvalidError gets raised if uxgrid=None when getting UxDataArray.uxgrid. Regression test for #1620. @@ -188,3 +359,214 @@ def test_uxgrid_None_is_invalid_in_uxdataarray(): # it also applies (for non-None non-Grid objects) during __init__: with pytest.raises(TypeError): ux.UxDataArray([4,5], dims=['n_face'], uxgrid="not a grid") + + """Tests for ``UxDataArray.neighborhood_filter``.""" + + @pytest.fixture + def vortex(self, gridpath, datasetpath): + """The ``psi`` field on outCSne30, which most of these tests filter.""" + uxds = ux.open_dataset( + gridpath("ugrid", "outCSne30", "outCSne30.ug"), + datasetpath("ugrid", "outCSne30", "outCSne30_vortex.nc"), + ) + return uxds["psi"] + + # Every named reduction, with the NumPy expression it must equal and the + # parameter it takes. The reference runs through the generic callable path, + # so this pins each compiled kernel against the loop it bypasses. + NAMED_REDUCTIONS = [ + ("mean", {}, lambda a, axis: np.mean(a, axis=axis)), + ("sum", {}, lambda a, axis: np.sum(a, axis=axis)), + ("min", {}, lambda a, axis: np.min(a, axis=axis)), + ("max", {}, lambda a, axis: np.max(a, axis=axis)), + ("median", {}, lambda a, axis: np.median(a, axis=axis)), + ("ptp", {}, lambda a, axis: np.ptp(a, axis=axis)), + ("std", {"ddof": 1}, lambda a, axis: np.std(a, axis=axis, ddof=1)), + ("var", {"ddof": 1}, lambda a, axis: np.var(a, axis=axis, ddof=1)), + ("quantile", {"q": 0.9}, lambda a, axis: np.quantile(a, 0.9, axis=axis)), + ("percentile", {"q": 90}, lambda a, axis: np.percentile(a, 90, axis=axis)), + ] + + @pytest.mark.parametrize("name,kwargs,reference", NAMED_REDUCTIONS) + def test_named_reduction_matches_numpy(self, name, kwargs, reference): + """Each compiled reduction must equal its NumPy expression, including + where NaN lands. + + The field is partly masked on purpose. NaN handling is the easy thing + to get wrong in a kernel: a hand-written ``if value > best`` loop skips + NaN where ``np.max`` propagates it, and numba's ``np.median`` + propagates only depending on where the NaN falls in its partition. The + extra leading dimension exercises the gufunc's broadcast loop. + """ + uxgrid = ux.Grid.from_healpix(zoom=2) + rng = np.random.default_rng(0) + values = rng.random((3, uxgrid.n_face)) + # mask a tenth of the faces, as a land/ocean mask would + values[:, rng.choice(uxgrid.n_face, uxgrid.n_face // 10, replace=False)] = np.nan + uxda = UxDataArray(values, dims=["time", "n_face"], uxgrid=uxgrid, name="masked") + + nb = uxgrid.neighborhoods(r=20.0) + got = nb.reduce(uxda, name, **kwargs).values + with warnings.catch_warnings(): + warnings.simplefilter("ignore", RuntimeWarning) # numpy all-NaN slices + expected = nb.reduce(uxda, reference).values + + assert np.isnan(got).any(), "expected a neighborhood to hit a masked value" + assert not np.isnan(got).all(), "expected some neighborhood to be clean" + np.testing.assert_array_equal(np.isnan(got), np.isnan(expected)) + finite = ~np.isnan(expected) + np.testing.assert_allclose(got[finite], expected[finite], rtol=1e-12) + + def test_callable_still_accepted(self, vortex): + """The original ``func=np.mean`` signature keeps working, and reaches + the same kernel the name does rather than dropping to the generic + loop.""" + from uxarray.grid.neighbors import _CALLABLE_ALIASES, _resolve_reduction + + for callable_func, name in _CALLABLE_ALIASES.items(): + assert ( + _resolve_reduction(callable_func, {})[0] + is _resolve_reduction(name, {})[0] + ), f"{callable_func} should reach the {name!r} kernel" + + np.testing.assert_allclose( + vortex.neighborhood_filter(np.mean, r=3.0).values, + vortex.neighborhood_filter("mean", r=3.0).values, + ) + + def test_callable_escape_hatch(self, vortex): + """A user's own function, with no compiled equivalent, still works on + the ``axis=-1`` contract. + + ``functools.partial`` is covered by test_custom_func_with_partial. + """ + + # a user's own function, with no NumPy equivalent at all + def rms(values, axis): + return np.sqrt(np.mean(values**2, axis=axis)) + + filtered = vortex.neighborhood_filter(rms, r=3.0) + assert filtered.shape == vortex.shape + assert np.all(filtered.values >= 0) + + @pytest.mark.parametrize( + "func,kwargs,error,match", + [ + ("meen", {}, ValueError, "Unknown reduction 'meen'"), + ("mean", {"q": 90}, TypeError, "unexpected keyword argument"), + ("quantile", {}, TypeError, "requires the 'q' keyword"), + ("quantile", {"q": 90}, ValueError, "between 0 and 1"), + (42, {}, TypeError, "name of a reduction or a callable"), + ], + ) + def test_invalid_reduction(self, func, kwargs, error, match): + """Naming a reduction makes bad input catchable up front, rather than + as a TypeError from inside the loop.""" + uxgrid = ux.Grid.from_healpix(zoom=1) + uxda = UxDataArray( + np.arange(uxgrid.n_face, dtype=float), dims=["n_face"], uxgrid=uxgrid + ) + with pytest.raises(error, match=match): + uxda.neighborhood_filter(func, r=1.0, **kwargs) + + def test_neighborhoods_reuse(self, vortex): + """A reused Neighborhoods must give the same answer as the one-shot + filter -- the point of holding onto it is that it costs one query.""" + nb = vortex.uxgrid.neighborhoods(r=4.0) + + assert (nb.r, nb.on, nb.grid_dim) == (4.0, "face centers", "n_face") + counts = nb.n_neighbors + assert counts.dims == ("n_face",) + # every element is its own neighbor, so no neighborhood is ever empty + assert counts.min() >= 1 + + for name, kwargs in [("mean", {}), ("percentile", {"q": 90})]: + np.testing.assert_allclose( + nb.reduce(vortex, name, **kwargs).values, + vortex.neighborhood_filter(name, r=4.0, **kwargs).values, + rtol=1e-12, + ) + + def test_neighborhoods_reject_wrong_data(self): + """Reducing data mapped elsewhere must fail loudly rather than index + into the wrong element set.""" + uxgrid = ux.Grid.from_healpix(zoom=1) + _ = uxgrid.n_node # populate node coords before the tree is built + nb = uxgrid.neighborhoods(r=30.0, on="face centers") + + node_data = UxDataArray( + np.arange(uxgrid.n_node, dtype=float), dims=["n_node"], uxgrid=uxgrid + ) + with pytest.raises(DataCenteringError, match="reduce over 'n_face'"): + nb.reduce(node_data, "mean") + + other = ux.Grid.from_healpix(zoom=2) + wrong_size = UxDataArray( + np.arange(other.n_face, dtype=float), dims=["n_face"], uxgrid=other + ) + with pytest.raises(DataCenteringError, match="different grid"): + nb.reduce(wrong_size, "mean") + + with pytest.raises(ValueError, match="Invalid `on`"): + uxgrid.neighborhoods(r=1.0, on="face_centers") + + def test_dask_input_stays_lazy(self, vortex): + """Lazy input stays lazy: the grid dimension is a core dimension, but + the others stay chunked and unevaluated.""" + eager = vortex.neighborhood_filter("mean", r=2.0) + + stacked = UxDataArray( + np.tile(vortex.values, (6, 1)), + dims=["time", "n_face"], + uxgrid=vortex.uxgrid, + name="psi", + ).chunk({"time": 2, "n_face": -1}) + + # chunking a non-grid dimension is the supported case: untouched, silent + with warnings.catch_warnings(): + warnings.simplefilter("error", UserWarning) + filtered = stacked.neighborhood_filter("mean", r=2.0) + + assert filtered.chunks is not None, "the filter should not force a compute" + assert filtered.chunksizes["time"] == (2, 2, 2) + assert isinstance(filtered, UxDataArray) + np.testing.assert_allclose( + filtered.compute().values, np.tile(eager.values, (6, 1)) + ) + + def test_grid_dim_chunks_are_collapsed_with_warning(self, vortex): + """A neighborhood may span the whole grid, so the grid dimension cannot + stay chunked. Collapsing it undoes a memory decision the user made, so + it is not done silently.""" + expected = vortex.neighborhood_filter("mean", r=2.0).values + + uxda = vortex.chunk({"n_face": 1000}) + assert len(uxda.chunksizes["n_face"]) > 1 + + with pytest.warns(UserWarning, match="Rechunking 'n_face'"): + filtered = uxda.neighborhood_filter("mean", r=2.0) + + assert filtered.chunksizes["n_face"] == (uxda.sizes["n_face"],) + np.testing.assert_allclose(filtered.compute().values, expected) + + def test_output_is_always_float64(self, vortex): + """float32 hits the kernel's float32 signature and integers have no + signature at all; both must come back as float64, as the generic path + does by writing into a float64 output.""" + uxgrid = ux.Grid.from_healpix(zoom=1) + integers = UxDataArray( + np.arange(uxgrid.n_face), dims=["n_face"], uxgrid=uxgrid, name="int_var" + ) + filtered = integers.neighborhood_filter("mean", r=0.0) + assert filtered.dtype == np.float64 + np.testing.assert_allclose(filtered.values, integers.values) + + as_float32 = UxDataArray( + vortex.values.astype(np.float32), dims=vortex.dims, uxgrid=vortex.uxgrid + ) + filtered32 = as_float32.neighborhood_filter("mean", r=2.0) + assert filtered32.dtype == np.float64 + np.testing.assert_allclose( + filtered32.values, vortex.neighborhood_filter("mean", r=2.0).values, + rtol=1e-6, + ) diff --git a/test/core/test_dataset.py b/test/core/test_dataset.py index 9b26d84d3..8c0557d22 100644 --- a/test/core/test_dataset.py +++ b/test/core/test_dataset.py @@ -168,6 +168,41 @@ def test_uxdataset_to_array(): 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) + def test_uxgrid_None_is_invalid_in_uxdataset(): """Ensures GridInvalidError gets raised if uxgrid=None when getting UxDataset.uxgrid. Regression test for #1620. @@ -190,3 +225,43 @@ def test_uxgrid_None_is_invalid_in_uxdataset(): # it also applies (for non-None non-Grid objects) during __init__: with pytest.raises(TypeError): ux.UxDataset({'arr1': xr.DataArray([4,5], dims=['n_face'])}, uxgrid=[1,2]) + def test_one_query_per_grid_location(self): + """Variables sharing a grid location must share one neighbor query. + + The query dominates the cost of a reduction, so rebuilding it per + variable would make a dataset filter scale with the number of + variables. Counting calls is the only way to see that from outside. + """ + from unittest.mock import patch + + import uxarray.grid.neighbors as neighbors + + uxgrid = ux.Grid.from_healpix(zoom=2) + # touch both locations first: a HEALPix grid cannot populate node + # coordinates lazily from inside the tree build + n_node, n_face = uxgrid.n_node, uxgrid.n_face + rng = np.random.default_rng(0) + uxds = UxDataset( + data_vars={ + "face_a": ("n_face", rng.random(n_face)), + "face_b": ("n_face", rng.random(n_face)), + "face_c": ("n_face", rng.random(n_face)), + "node_a": ("n_node", rng.random(n_node)), + }, + uxgrid=uxgrid, + ) + + real = neighbors._csr_neighbors + with patch.object(neighbors, "_csr_neighbors", side_effect=real) as spy: + filtered = uxds.neighborhood_filter("percentile", r=20.0, q=90) + + assert spy.call_count == 2, ( + f"expected one query per grid location (faces, nodes), got " + f"{spy.call_count}" + ) + # and the reduction, with its parameter, reached every variable + for name in ("face_a", "node_a"): + nt.assert_allclose( + filtered[name].values, + uxds[name].neighborhood_filter("percentile", r=20.0, q=90).values, + ) diff --git a/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 28fa22886..3de832632 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, @@ -34,6 +35,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 Neighborhoods from uxarray.grid.validation import _check_duplicate_nodes_indices from uxarray.io._healpix import get_zoom_from_cells from uxarray.plot.accessor import UxDataArrayPlotAccessor @@ -2199,6 +2201,107 @@ def get_dual(self): return uxda + def _neighborhood_location(self, caller: str) -> str: + """Grid location this data is mapped to, in ``Neighborhoods`` terms.""" + if self._face_centered(): + return "face centers" + if self._node_centered(): + return "nodes" + if self._edge_centered(): + return "edge centers" + raise DataCenteringError( + f"{caller} requires data mapped to nodes, edges, or faces, " + f"but the dimensions {self.dims!r} do not match any grid dimension " + f"{GRID_DIMS}." + ) + + def neighborhood_filter( + self, + func: str | Callable = "mean", + r: float = 1.0, + **kwargs, + ) -> UxDataArray: + """Apply a neighborhood filter, replacing the value at each grid + element with a reduction of all elements within a circular + neighborhood of radius ``r``. + + Parameters + ---------- + func : str or Callable, default="mean" + Name of the reduction to apply: "mean", "sum", "min", "max", + "median", "ptp", "std", "var", "quantile", or "percentile". Named + reductions run compiled. A callable is accepted as an escape hatch + for anything not in that list; see Notes. + r : float, default=1. + Radius of the neighborhood, in degrees. + **kwargs + Parameter for the named reduction: ``q`` for "quantile" (0-1) and + "percentile" (0-100), ``ddof`` for "std" and "var". + + Returns + ------- + uxda_filter : UxDataArray + Filtered data, as float64. + + Raises + ------ + DataCenteringError (subclass of ValueError) + If the data is not mapped to nodes, edges, or faces. + ValueError + If ``func`` names a reduction that does not exist. + TypeError + If ``func`` is a callable that does not accept an ``axis`` keyword + argument, or if a keyword argument does not apply to ``func``. + + 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``. + + A callable ``func`` is applied as ``func(values, axis=-1)`` over a block + whose last axis is the neighborhood, once per grid element, in Python. + That is considerably slower than a named reduction, so prefer a name + where one exists. + + Each call queries the grid for neighbors, which usually costs more than + the reduction itself. To apply several reductions at one radius, build + the neighborhoods once with :meth:`Grid.neighborhoods` and call + :meth:`Neighborhoods.reduce` on it. + + A neighborhood may span the whole grid, so the grid dimension cannot be + chunked; it is collapsed to a single chunk (with a warning) for + dask-backed data. The remaining dimensions stay chunked and lazy, so + chunk along ``time`` rather than the grid dimension. + + Examples + -------- + Apply a mean filter with a 5-degree radius: + + >>> import uxarray as ux + >>> uxds = ux.tutorial.open_dataset("outCSne30-vortex") + >>> uxda = uxds["psi"] + >>> smoothed = uxda.neighborhood_filter("mean", r=5.0) + + Reductions taking a parameter receive it as a keyword argument: + + >>> p90 = uxda.neighborhood_filter("percentile", r=5.0, q=90) + >>> spread = uxda.neighborhood_filter("std", r=5.0, ddof=1) + + See Also + -------- + Grid.neighborhoods : Reusable neighborhoods, for several reductions at one radius. + UxDataArray.topological_mean : Aggregate values across neighboring grid element types. + UxDataArray.zonal_mean : Average over latitude bands. + UxDataArray.azimuthal_mean : Average over rings of constant great-circle distance. + """ + neighborhoods = Neighborhoods( + self.uxgrid, + r=r, + on=self._neighborhood_location("neighborhood_filter"), + ) + return neighborhoods.reduce(self, func=func, **kwargs) + 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 42057a41c..a90d5b985 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,12 +13,14 @@ 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, GridInvalidError from uxarray.formatting_html import dataset_repr from uxarray.grid import Grid from uxarray.grid.dual import construct_dual +from uxarray.grid.neighbors import Neighborhoods from uxarray.grid.validation import _check_duplicate_nodes_indices from uxarray.io._healpix import get_zoom_from_cells from uxarray.plot.accessor import UxDatasetPlotAccessor @@ -677,6 +679,77 @@ def to_array( return UxDataArray(xarr, uxgrid=self._uxgrid) # _uxgrid not uxgrid; converting to UxDataArray is not a grid-aware method. + def neighborhood_filter( + self, + func: str | Callable = "mean", + r: float = 1.0, + **kwargs, + ) -> UxDataset: + """Apply a neighborhood filter, replacing the value at each grid + element of every data variable with a reduction of all elements within + a circular neighborhood of radius ``r``. + + Parameters are as for :meth:`UxDataArray.neighborhood_filter`, which + documents the available reductions and their keyword arguments. + + Returns + ------- + destination_uxds : UxDataset + Filtered dataset. + + Notes + ----- + Variables without a grid dimension are passed through unchanged. + + Variables mapped to the same grid location share one neighbor query, so + filtering a dataset costs one query per location present rather than + one per variable. + + Examples + -------- + Apply a mean filter to all grid-mapped variables in a dataset: + + >>> import uxarray as ux + >>> uxds = ux.tutorial.open_dataset("outCSne30-vortex") + >>> uxds_smooth = uxds.neighborhood_filter("mean", r=5.0) + + See Also + -------- + UxDataArray.neighborhood_filter : Filter a single data variable. + Grid.neighborhoods : Reusable neighborhoods, for several reductions at one radius. + UxDataArray.zonal_mean : Average over latitude bands. + UxDataArray.azimuthal_mean : Average over rings of constant great-circle distance. + """ + + destination_uxds = self._copy() + + # The neighbor query dominates the cost of a reduction, and it depends + # only on (grid, location, radius) -- not on the data. Variables mapped + # to the same location therefore share one query, built on first use. + neighborhoods: dict[str, Neighborhoods] = {} + + # Loop through UxDataArrays in UxDataset and apply the filter to every + # variable that is mapped to a grid element (node, edge, or face). + # Variables without a grid dimension are left unchanged. + 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 + + location = uxda._neighborhood_location("neighborhood_filter") + if location not in neighborhoods: + neighborhoods[location] = Neighborhoods(self.uxgrid, r=r, on=location) + + # Neighborhoods.reduce restores the input dimension order, so it is + # always preserved. + destination_uxds[var_name] = neighborhoods[location].reduce( + uxda, func=func, **kwargs + ) + + return destination_uxds + 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 6fa02b069..0d59c1554 100644 --- a/uxarray/grid/grid.py +++ b/uxarray/grid/grid.py @@ -60,6 +60,7 @@ from uxarray.grid.neighbors import ( BallTree, KDTree, + Neighborhoods, SpatialHash, _populate_edge_face_distances, _populate_edge_node_distances, @@ -1766,7 +1767,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" @@ -1783,7 +1784,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, @@ -1791,12 +1802,45 @@ 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 + def neighborhoods(self, r: float = 1.0, on: str = "face centers") -> Neighborhoods: + """Finds the grid elements within ``r`` degrees of every element of + ``on``, returning a reusable :class:`Neighborhoods`. + + The radius query behind this dominates the cost of a neighborhood + reduction, so building this once and reducing several times over it is + substantially cheaper than calling + :meth:`UxDataArray.neighborhood_filter` repeatedly, which rebuilds it + on every call. + + Parameters + ---------- + r : float, default=1. + Radius of the neighborhood, in degrees of great-circle distance. + on : str, default="face centers" + Grid location to center the neighborhoods on: "nodes", + "edge centers", or "face centers". + + Returns + ------- + Neighborhoods + + Examples + -------- + >>> import uxarray as ux + >>> uxds = ux.tutorial.open_dataset("outCSne30-vortex") # doctest: +SKIP + >>> nb = uxds.uxgrid.neighborhoods(r=5.0) # doctest: +SKIP + >>> smooth = nb.reduce(uxds["psi"], "mean") # doctest: +SKIP + >>> p90 = nb.reduce(uxds["psi"], "percentile", q=90) # doctest: +SKIP + + See Also + -------- + UxDataArray.neighborhood_filter : One-shot filter for a single reduction. + """ + return Neighborhoods(self, r=r, on=on) + def _get_scipy_kd_tree( self, coordinates: str | None = "face", reconstruct: bool = False ): @@ -1883,7 +1927,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, @@ -1892,10 +1944,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..83f78aa90 100644 --- a/uxarray/grid/neighbors.py +++ b/uxarray/grid/neighbors.py @@ -1,6 +1,9 @@ +import warnings +from typing import Callable, NamedTuple + import numpy as np import xarray as xr -from numba import njit +from numba import guvectorize, njit from numpy import deg2rad from uxarray.constants import ERROR_TOLERANCE, INT_DTYPE, INT_FILL_VALUE @@ -1129,3 +1132,540 @@ 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}" + ) + + +# A neighborhood reduction is a segmented reduction over a ragged (CSR-like) +# neighbor structure: elementwise in every dimension except the grid axis, +# which it reduces over. That is exactly a generalized ufunc signature, so the +# kernels below declare the grid axis as a core dimension. Two consequences +# fall out of stating it that way: +# +# * dask can parallelize over the remaining (chunked) dimensions on its own, +# so the filter stays lazy instead of materializing the whole array, and +# * the grid axis is a *core* dimension, so dask refuses to split it rather +# than silently handing a kernel a block the neighbor indices overrun. +# +# ``(n)`` is the source grid axis, ``(k)`` the flattened neighbor index array, +# and ``(m)`` the destination axis. Output is float64 regardless of input +# dtype, matching the behaviour of the generic path below. +_GUFUNC_SIGNATURES = [ + "void(float64[:], int64[:], int64[:], int64[:], float64, float64[:])", + "void(float32[:], int64[:], int64[:], int64[:], float64, float64[:])", +] +_GUFUNC_LAYOUT = "(n),(k),(m),(m),()->(m)" +_GUFUNC_KWARGS = {"nopython": True, "cache": True, "target": "parallel"} + + +def _make_kernel(reduce_fn): + """Builds a kernel that gathers each neighborhood, then calls + ``reduce_fn(window, param)`` on the 1-D result. + + ``reduce_fn`` must be numba-compilable, and must be defined in a real + source file for ``cache=True`` to find it. + """ + # A reducer shared between kernels arrives already compiled; numba rejects + # jitting a dispatcher twice. + if not hasattr(reduce_fn, "py_func"): + reduce_fn = njit(cache=True)(reduce_fn) + + @guvectorize(_GUFUNC_SIGNATURES, _GUFUNC_LAYOUT, **_GUFUNC_KWARGS) + def kernel(data, flat, starts, counts, param, out): + widest = 0 + for i in range(counts.shape[0]): + if counts[i] > widest: + widest = counts[i] + buffer = np.empty(widest, dtype=np.float64) + + for i in range(starts.shape[0]): + count = counts[i] + if count == 0: + out[i] = np.nan + continue + start = starts[i] + for j in range(count): + buffer[j] = data[flat[start + j]] + out[i] = reduce_fn(buffer[:count], param) + + return kernel + + +class _Reduction(NamedTuple): + """A named reduction, and the single scalar parameter it accepts (if any). + + Limiting reductions to one parameter is what keeps the gufunc layout above + down to one; it covers every reduction implemented here. + """ + + kernel: object + param: str | None = None + default: float = 0.0 + + +# Reductions with a compiled kernel, addressed by name. A name always takes the +# fast path, which is why the public API documents names rather than callables: +# dispatching on a function object cannot see through ``functools.partial``, so +# a parameterized reduction could never hit a kernel that way. +# +# Adding a reduction is one line here. Reducers take ``(window, param)``; +# those with no parameter ignore the second argument. Numba keys its cache by +# code object rather than qualified name, so the identically-named lambdas do +# not collide. +@njit(cache=True) +def _variance(window, ddof): + """Variance with a delta degrees of freedom. Numba's ``np.var`` takes no + ``ddof``, so the two-pass form is spelled out.""" + denominator = window.size - ddof + if denominator <= 0: + return np.nan + center = np.mean(window) + total = 0.0 + for value in window: + total += (value - center) ** 2 + return total / denominator + + +@njit(cache=True) +def _median(window, _): + # numba's ``np.median`` selects by partitioning, and whether a NaN survives + # that depends on where it lands -- so unlike numpy's, it propagates NaN + # only sometimes. This spelling short-circuits and allocates nothing: + # ``np.any(np.isnan(window))`` costs ~14% more, and routing through + # ``np.quantile``, which does propagate, costs 2.5x. + for value in window: + if np.isnan(value): + return np.nan + return np.median(window) + + +_quantile_kernel = _make_kernel(lambda window, q: np.quantile(window, q)) + +_REDUCTIONS = { + "mean": _Reduction(_make_kernel(lambda window, _: np.mean(window))), + "sum": _Reduction(_make_kernel(lambda window, _: np.sum(window))), + "min": _Reduction(_make_kernel(lambda window, _: np.min(window))), + "max": _Reduction(_make_kernel(lambda window, _: np.max(window))), + "ptp": _Reduction(_make_kernel(lambda window, _: np.max(window) - np.min(window))), + "median": _Reduction(_make_kernel(_median)), + "var": _Reduction(_make_kernel(_variance), param="ddof"), + "std": _Reduction( + _make_kernel(lambda window, ddof: np.sqrt(_variance(window, ddof))), + param="ddof", + ), + "quantile": _Reduction(_quantile_kernel, param="q"), + "percentile": _Reduction(_quantile_kernel, param="q"), +} + +# Callables accepted for backwards compatibility, so that code written against +# the original ``func=np.mean`` signature keeps the fast path instead of +# silently dropping to the generic loop. +_CALLABLE_ALIASES = { + np.mean: "mean", + np.sum: "sum", + np.max: "max", + np.amax: "max", + np.min: "min", + np.amin: "min", + np.median: "median", + np.std: "std", + np.var: "var", + np.ptp: "ptp", +} + + +def _resolve_reduction(func, kwargs): + """Maps ``func`` (a name or a callable) onto a kernel and its parameter. + + Returns ``(kernel, param_value)`` for a compiled reduction, or + ``(None, None)`` when ``func`` is a callable that has to go through the + generic loop. + """ + name = func if isinstance(func, str) else _CALLABLE_ALIASES.get(func) + + if name is None: + if not callable(func): + raise TypeError( + f"`func` must be the name of a reduction or a callable, but got " + f"{func!r}. Valid names: {', '.join(sorted(_REDUCTIONS))}." + ) + if kwargs: + raise TypeError( + f"Got unexpected keyword argument(s) {', '.join(sorted(kwargs))} " + f"for a callable `func`. Parameters are only supported for named " + f"reductions; use `functools.partial` to bind them to a callable." + ) + return None, None + + if name not in _REDUCTIONS: + raise ValueError( + f"Unknown reduction {name!r}. Expected one of: " + f"{', '.join(sorted(_REDUCTIONS))}." + ) + + reduction = _REDUCTIONS[name] + unexpected = set(kwargs) - ({reduction.param} if reduction.param else set()) + if unexpected: + raise TypeError( + f"Reduction {name!r} got unexpected keyword argument(s) " + f"{', '.join(sorted(unexpected))}." + + (f" It accepts {reduction.param!r}." if reduction.param else "") + ) + + if reduction.param is None: + # The kernel still takes a parameter; this one ignores it. + return reduction.kernel, 0.0 + + if reduction.param in kwargs: + value = float(kwargs[reduction.param]) + elif name in ("quantile", "percentile"): + raise TypeError( + f"Reduction {name!r} requires the {reduction.param!r} keyword argument." + ) + else: + value = reduction.default + + # `percentile` is `quantile` on a 0-100 scale; normalize so both share one + # kernel rather than compiling a near-duplicate. + if name == "percentile": + if not 0.0 <= value <= 100.0: + raise ValueError(f"`q` must be between 0 and 100, but got {value}.") + value /= 100.0 + elif name == "quantile" and not 0.0 <= value <= 1.0: + raise ValueError(f"`q` must be between 0 and 1, but got {value}.") + + return reduction.kernel, value + + +def _csr_neighbors(grid, data_mapping: str, r: float): + """Queries the neighborhood of every element and returns it in CSR form. + + ``query_radius`` returns a ragged sequence of index arrays, one per + element. Flattening it into ``(flat, starts, counts)`` gives the kernels a + layout they can walk without allocating per-neighborhood temporaries. + + Returns + ------- + flat : np.ndarray + Concatenated neighbor indices for every element. + starts : np.ndarray + Offset into ``flat`` at which each element's neighbors begin. + counts : np.ndarray + Number of neighbors of each element. + """ + # Request a spherical/haversine tree explicitly rather than relying on the + # defaults. Without this, a cartesian tree cached by an earlier call would + # be reused and ``r`` would be silently interpreted as a chord length + # instead of the great-circle degrees documented by the callers. + coordinate_system = "spherical" + tree = grid.get_ball_tree( + coordinates=data_mapping, + coordinate_system=coordinate_system, + distance_metric="haversine", + ) + + dest_coords = _get_element_coords(grid, data_mapping, coordinate_system) + neighbor_indices = tree.query_radius(dest_coords, r=r) + + # ``query_radius`` unwraps its result for a single query point, which a + # one-element grid would hit. + if isinstance(neighbor_indices, np.ndarray): + neighbor_indices = [neighbor_indices] + + counts = np.fromiter( + map(len, neighbor_indices), dtype=np.int64, count=len(neighbor_indices) + ) + starts = np.zeros(counts.size, dtype=np.int64) + np.cumsum(counts[:-1], out=starts[1:]) + flat = np.concatenate(neighbor_indices).astype(np.int64, copy=False) + + return flat, starts, counts + + +def _neighborhood_reduce(block, flat, starts, counts, func: Callable): + """Generic fallback: applies ``func`` to each neighborhood in turn. + + Used when ``func`` has no compiled kernel. ``block`` is a NumPy array with + the grid dimension last. + """ + destination_data = np.full(block.shape, np.nan) + + # The `axis` check lives outside the loop: whether `func` accepts the + # keyword cannot change between iterations, so validating it once is + # equivalent to validating it every time and leaves the loop body bare. + try: + for i in range(starts.shape[0]): + idx = flat[starts[i] : starts[i] + counts[i]] + # Apply func along the last (grid) axis only, so any extra leading + # dimensions (e.g. time) are preserved rather than being collapsed. + destination_data[..., i] = func(block[..., idx], axis=-1) + except TypeError as exc: + if "axis" not in str(exc): + raise + raise TypeError( + f"`func` must accept an `axis` keyword argument so that the " + f"reduction is applied over the neighborhood only, but " + f"{getattr(func, '__name__', func)!r} does not. Use a NumPy " + f"reduction such as `np.mean` or `np.median`, or wrap your " + f"function with `functools.partial` to supply `axis`." + ) from exc + + return destination_data + + +def _rechunk_grid_dim(uxda, grid_dim: str): + """Collapses the grid dimension to a single chunk, warning if that changes + the user's chunking. + + Neighborhoods are global — an element near a chunk boundary draws on + elements in other chunks — so the grid dimension cannot be chunked. This is + done explicitly rather than through ``allow_rechunk``, which would do it + silently and also disable ``apply_gufunc``'s other consistency checks. + """ + if uxda.chunks is None: + return uxda + + grid_chunks = uxda.chunksizes.get(grid_dim, ()) + if len(grid_chunks) <= 1: + return uxda + + warnings.warn( + f"Rechunking {grid_dim!r} from {len(grid_chunks)} chunks into one, as a " + f"neighborhood may span the whole grid. Each task will hold " + f"{uxda.sizes[grid_dim]} elements along {grid_dim!r}; chunk the " + f"non-grid dimensions instead to bound memory use.", + UserWarning, + stacklevel=3, + ) + + return uxda.chunk({grid_dim: -1}) + + +ELEMENT_DIMS = { + "nodes": "n_node", + "edge centers": "n_edge", + "face centers": "n_face", +} + + +class Neighborhoods: + """The set of grid elements within a radius ``r`` of every element of one + grid location, ready to be reduced over. + + Building this queries a ``BallTree`` once, which is by far the dominant + cost of a neighborhood reduction — typically far more than the reduction + itself. Holding onto the result lets several reductions, or several + variables, share that one query instead of repeating it. + + Parameters + ---------- + grid : Grid + Grid whose elements define the neighborhoods. + r : float, default=1. + Radius of the neighborhood, in degrees of great-circle distance. + on : str, default="face centers" + Grid location the neighborhoods are built around: "nodes", + "edge centers", or "face centers". + + Examples + -------- + >>> import uxarray as ux + >>> uxds = ux.tutorial.open_dataset("outCSne30-vortex") # doctest: +SKIP + >>> nb = uxds.uxgrid.neighborhoods(r=5.0) # doctest: +SKIP + >>> smooth = nb.reduce(uxds["psi"], "mean") # doctest: +SKIP + >>> spread = nb.reduce(uxds["psi"], "std") # doctest: +SKIP + + See Also + -------- + UxDataArray.neighborhood_filter : One-shot filter that builds this internally. + """ + + def __init__(self, grid, r: float = 1.0, on: str = "face centers"): + if on not in ELEMENT_DIMS: + raise ValueError( + f"Invalid `on`. Expected one of {', '.join(sorted(ELEMENT_DIMS))}, " + f"but received {on!r}." + ) + + self._grid = grid + self._r = float(r) + self._on = on + self._flat, self._starts, self._counts = _csr_neighbors(grid, on, self._r) + + @property + def grid(self): + """Grid the neighborhoods were built from.""" + return self._grid + + @property + def r(self) -> float: + """Neighborhood radius, in degrees.""" + return self._r + + @property + def on(self) -> str: + """Grid location the neighborhoods are centered on.""" + return self._on + + @property + def grid_dim(self) -> str: + """Name of the grid dimension this reduces over.""" + return ELEMENT_DIMS[self._on] + + @property + def n_neighbors(self) -> xr.DataArray: + """Number of elements in each neighborhood, itself a grid-mapped field. + + Useful for seeing how a fixed radius samples a variable-resolution + mesh, where the count varies by region. + """ + return xr.DataArray( + self._counts.copy(), + dims=[self.grid_dim], + name="n_neighbors", + attrs={"long_name": f"elements within {self._r} degrees"}, + ) + + def __repr__(self) -> str: + return ( + f"" + ) + + def reduce(self, uxda, func="mean", **kwargs): + """Reduces ``uxda`` over each neighborhood. + + Parameters + ---------- + uxda : UxDataArray + Data to reduce, mapped to the same grid location as ``on``. The + grid dimension may sit at any position. + func : str or Callable, default="mean" + Name of a compiled reduction — "mean", "sum", "min", "max", + "median", "ptp", "std", "var", "quantile", "percentile" — or a + callable taking an ``axis`` keyword (see Notes). + **kwargs + Parameter for the named reduction: ``q`` for "quantile" (0-1) and + "percentile" (0-100), ``ddof`` for "std" and "var". + + Returns + ------- + UxDataArray + Reduced data as float64, with the input's dimension order. Lazy if + the input was lazy. + + Notes + ----- + A callable is an escape hatch for reductions not implemented here. It + is applied as ``func(values, axis=-1)`` over a block whose last axis is + the neighborhood, once per element, in Python — considerably slower + than a named reduction. Named reductions run compiled. + """ + # Local import: uxarray.core.dataarray imports this module. + from uxarray.core.dataarray import UxDataArray + from uxarray.errors import DataCenteringError + + grid_dim = self.grid_dim + if grid_dim not in uxda.dims: + raise DataCenteringError( + f"These neighborhoods are built on {self._on!r} and reduce over " + f"{grid_dim!r}, but the data has dimensions {tuple(uxda.dims)!r}." + ) + if uxda.sizes[grid_dim] != self._counts.size: + raise DataCenteringError( + f"Data has {uxda.sizes[grid_dim]} elements along {grid_dim!r}, but " + f"these neighborhoods describe {self._counts.size}. The data is " + f"probably mapped to a different grid." + ) + + kernel, param = _resolve_reduction(func, kwargs) + + if kernel is None: + + def _apply(block): + return _neighborhood_reduce( + block, self._flat, self._starts, self._counts, func + ) + else: + + def _apply(block): + # The kernels are compiled for float32/float64 only; anything + # else (integer fields, say) is promoted, which the generic + # path does too by writing into a float64 output. + if block.dtype not in (np.float64, np.float32): + block = block.astype(np.float64) + return kernel(block, self._flat, self._starts, self._counts, param) + + work = _rechunk_grid_dim(uxda, grid_dim) + + # ``apply_ufunc`` moves the grid dimension last before calling + # ``_apply`` and, for dask-backed input, hands each chunk over as a + # materialized NumPy block. Indexing the array one destination element + # at a time would instead trigger one graph execution per grid element. + filtered = xr.apply_ufunc( + _apply, + work, + input_core_dims=[[grid_dim]], + output_core_dims=[[grid_dim]], + dask="parallelized", + output_dtypes=[np.float64], + keep_attrs=True, + ) + + # Core dimensions come back appended last, so restore the input order. + if filtered.dims != uxda.dims: + filtered = filtered.transpose(*uxda.dims) + + # ``apply_ufunc`` returns a plain xr.DataArray, dropping the subclass + # and its grid. Name, coords and attrs are carried through already. + return UxDataArray(filtered, uxgrid=getattr(uxda, "uxgrid", self._grid))