Apply function to points within circular neighborhood

docs/api.rst

@@ -565,6 +565,19 @@ Azimuthal aggregations apply an aggregation (i.e. averaging) along circles of co
    UxDataArray.azimuthal_mean
 
 
+Neighborhood
+~~~~~~~~~~~~
+
+Neighborhood filters apply an aggregation (i.e. averaging) to all grid elements within a circular
+neighborhood of a specified radius around each grid element.
+
+.. autosummary::
+   :toctree: generated/
+
+   UxDataArray.neighborhood_filter
+   UxDataset.neighborhood_filter
+
+
 Zonal Average
 ~~~~~~~~~~~~~
 .. autosummary::

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
+}