diff --git a/benchmarks/mpas_ocean.py b/benchmarks/mpas_ocean.py index cd6177f47..adf679a3a 100644 --- a/benchmarks/mpas_ocean.py +++ b/benchmarks/mpas_ocean.py @@ -1,10 +1,12 @@ import os import urllib.request +import warnings from pathlib import Path import numpy as np import uxarray as ux +from uxarray.grid.neighbors import Neighborhood, _get_element_coords from .helpers._memsize import grid_nbytes from .helpers._peakmem import numba_threads, peak_allocated @@ -348,3 +350,178 @@ def peakmem_const_lat(self, resolution, lat_step): uxgrid.bounds for lat in np.arange(-45, 45, lat_step): uxgrid.cross_section.constant_latitude(lat) + + +class NeighborhoodBuild(DatasetBenchmark): + """Construction cost of a ``Neighborhood``, split into its three stages. + + ``Neighborhood`` claims the neighbor query costs more than any reduction run + on it. ``r`` is a great-circle radius in degrees, against + a mean element spacing of roughly 4.3 degrees at 480km and 1.1 at 120km, so + the smallest radius here is near self-only on the coarser mesh. + """ + + param_names = DatasetBenchmark.param_names + ['r'] + params = DatasetBenchmark.params + [[1.0, 5.0, 15.0]] + + def setup(self, resolution, r): + super().setup(resolution) + self.uxgrid = self.uxds.uxgrid + # Build the coordinates, and the njit paths behind them, here -- so the + # timings below are the query rather than lat/lon construction. + self.coords = _get_element_coords(self.uxgrid, "face centers", "spherical") + self.tree = self.uxgrid.get_ball_tree(coordinates="face centers", + coordinate_system="spherical", + distance_metric="haversine") + + def time_query_radius(self, resolution, r): + self.tree.query_radius(self.coords, r=r) + + def time_build(self, resolution, r): + """Query plus CSR flatten; ``setup`` has already cached the tree.""" + Neighborhood(self.uxgrid, r=r, on="face centers") + + def track_nbytes_neighbors(self, resolution, r): + """Size of the CSR structure a ``Neighborhood`` holds onto.""" + nb = Neighborhood(self.uxgrid, r=r, on="face centers") + return nb._flat.nbytes + nb._starts.nbytes + nb._counts.nbytes + + track_nbytes_neighbors.unit = "bytes" + + def track_peakmem_build(self, resolution, r): + """Transient high-water allocation of building a ``Neighborhood``.""" + return peak_allocated( + lambda: Neighborhood(self.uxgrid, r=r, on="face centers")) + + track_peakmem_build.unit = "bytes" + + def track_mean_neighbors(self, resolution, r): + """Mean neighborhood size -- the ``k`` behind every timing here.""" + nb = Neighborhood(self.uxgrid, r=r, on="face centers") + return round(float(nb.n_neighbors.mean()), 2) + + track_mean_neighbors.unit = "elements" + + +class NeighborhoodReduce(DatasetBenchmark): + """Reduction cost, and what reusing one neighbor query saves. + + One reduction, measured three ways: on a neighborhood built in ``setup``, + which is the compiled kernel alone; through ``UxDataArray.neighborhood``, + which pays for a query per call; and through ``UxDataset.neighborhood``, + which shares one query per grid location across every variable. The first + two bracket how much of a call is the query, and the third says whether + ``DatasetNeighborhood`` actually shares one. + + ``mean`` is linear in the neighborhood, while ``median`` partitions it. + """ + + param_names = DatasetBenchmark.param_names + ['reduction'] + params = DatasetBenchmark.params + [['mean', 'median']] + + radius = 15.0 + + @staticmethod + def _run(neighborhood, reduction): + """Calls ``reduction`` on an already-bound neighborhood.""" + if reduction == 'percentile': + return neighborhood.percentile(90) + if reduction == 'std': + return neighborhood.std(ddof=1) + return getattr(neighborhood, reduction)() + + def setup(self, resolution, reduction): + super().setup(resolution) + uxgrid = self.uxds.uxgrid + + # There is one compiled kernel per reduction, so warm the one under + # test on the coarsest grid. ``cache=True`` is an on-disk cache and asv + # builds a fresh environment per commit, so the first call still + # compiles. + grid, data = file_path_dict[self.params[0][0]] + warmup = ux.open_dataset(grid, data)[data_var].neighborhood(r=1.0) + _ = self._run(warmup, reduction) + + # A second face-centered variable, so the dataset case has something to + # share a query with, and one variable at each of the other two + # locations, so it has to build more than one. + self.uxds['depth_squared'] = self.uxds[data_var] ** 2 + self.uxds['node_var'] = ux.UxDataArray( + np.ones(uxgrid.n_node), dims=('n_node',), uxgrid=uxgrid) + self.uxds['edge_var'] = ux.UxDataArray( + np.ones(uxgrid.n_edge), dims=('n_edge',), uxgrid=uxgrid) + # Edge coordinates pull in edge_node_connectivity; build all three + # locations now so the first timed call is not the one that pays. + _, _, _ = uxgrid.node_lon, uxgrid.edge_lon, uxgrid.face_lon + + self.nb = self.uxds[data_var].neighborhood(r=self.radius) + + def time_reduce(self, resolution, reduction): + """The kernel alone: the query was paid for in ``setup``.""" + self._run(self.nb, reduction) + + def time_neighborhood_reduce(self, resolution, reduction): + """A query per call, which is what reuse is meant to avoid.""" + self._run(self.uxds[data_var].neighborhood(r=self.radius), reduction) + + def time_dataset_reduce(self, resolution, reduction): + """Four variables across three grid locations, sharing three queries.""" + self._run(self.uxds.neighborhood(r=self.radius), reduction) + + def track_peakmem_reduce(self, resolution, reduction): + """Transient allocation of the kernel, with the query already paid for. + + Held to one numba thread: the kernels are ``target="parallel"``, and + tracing serializes them on tracemalloc's allocator lock. + """ + with numba_threads(1): + return peak_allocated(lambda: self._run(self.nb, reduction)) + + track_peakmem_reduce.unit = "bytes" + + +class NeighborhoodDask(DatasetBenchmark): + """A reduction over lazy input, chunked three ways.""" + + param_names = DatasetBenchmark.param_names + ['chunking'] + params = DatasetBenchmark.params + [['numpy', 'time_chunks', 'grid_chunks']] + + n_time = 12 + radius = 5.0 + + def setup(self, resolution, chunking): + super().setup(resolution) + grid, data = file_path_dict[self.params[0][0]] + _ = ux.open_dataset(grid, data)[data_var].neighborhood(r=1.0).mean() + + base = self.uxds[data_var].values + stacked = np.broadcast_to(base, (self.n_time,) + base.shape).copy() + uxda = ux.UxDataArray(stacked, dims=('time', 'n_face'), + uxgrid=self.uxds.uxgrid) + if chunking == 'time_chunks': + uxda = uxda.chunk({'time': 1}) + elif chunking == 'grid_chunks': + uxda = uxda.chunk({'time': 1, + 'n_face': uxda.sizes['n_face'] // 4}) + + # Built here, so these measure the reduction and the graph it runs + # through rather than the query. + self.nb = uxda.neighborhood(r=self.radius) + + # One reduction here too, to warm the dask graph path -- and, for + # 'grid_chunks', to let the rechunk warning through exactly once... + _ = self.nb.mean().compute() + + # ...then silence the repeats. + warnings.filterwarnings('ignore', category=UserWarning, + message='Rechunking') + + def time_mean(self, resolution, chunking): + _ = self.nb.mean().compute() + + def track_peakmem_mean(self, resolution, chunking): + """High-water allocation of the reduction, tree query excluded.""" + with numba_threads(1): + return peak_allocated(lambda: self.nb.mean().compute()) + + track_peakmem_mean.unit = "bytes" diff --git a/docs/api.rst b/docs/api.rst index 68b70bf9d..e2a3d729a 100644 --- a/docs/api.rst +++ b/docs/api.rst @@ -197,6 +197,7 @@ Methods Grid.compute_face_node_angles Grid.construct_face_centers Grid.get_ball_tree + Grid.neighborhood Grid.get_kd_tree Grid.get_spatial_hash Grid.get_faces_containing_point @@ -568,6 +569,56 @@ Azimuthal aggregations apply an aggregation (i.e. averaging) along circles of co UxDataArray.azimuthal_mean +Neighborhood +~~~~~~~~~~~~ + +Neighborhood reductions apply an aggregation (i.e. averaging) to all grid elements within a +circular neighborhood of a specified radius around each grid element, as in a smoothing filter. +Grouping the data comes first, then a reduction over the groups:: + + uxda.neighborhood(r=5.0).mean() + uxda.neighborhood(r=5.0).percentile(90) + +.. autosummary:: + :toctree: generated/ + + UxDataArray.neighborhood + UxDataset.neighborhood + +Each reduction is a method of the object those return. ``reduce`` is the escape hatch for +anything without a method of its own; it takes a callable, and runs in Python rather than +compiled. + +.. autosummary:: + :toctree: generated/ + + uxarray.grid.neighbors.DataArrayNeighborhood + uxarray.grid.neighbors.DataArrayNeighborhood.mean + uxarray.grid.neighbors.DataArrayNeighborhood.sum + uxarray.grid.neighbors.DataArrayNeighborhood.min + uxarray.grid.neighbors.DataArrayNeighborhood.max + uxarray.grid.neighbors.DataArrayNeighborhood.ptp + uxarray.grid.neighbors.DataArrayNeighborhood.median + uxarray.grid.neighbors.DataArrayNeighborhood.std + uxarray.grid.neighbors.DataArrayNeighborhood.var + uxarray.grid.neighbors.DataArrayNeighborhood.quantile + uxarray.grid.neighbors.DataArrayNeighborhood.percentile + uxarray.grid.neighbors.DataArrayNeighborhood.reduce + uxarray.grid.neighbors.DataArrayNeighborhood.n_neighbors + uxarray.grid.neighbors.DatasetNeighborhood + +Finding the neighbors is usually more expensive than reducing over them. Reductions on one of the +objects above already share a single query. To share one across several variables too, build the +neighborhood from the grid instead; its reduction methods then take the data as an argument. + +.. autosummary:: + :toctree: generated/ + + Grid.neighborhood + uxarray.grid.neighbors.Neighborhood + uxarray.grid.neighbors.Neighborhood.reduce + uxarray.grid.neighbors.Neighborhood.n_neighbors + .. _zonal-average: Zonal Average diff --git a/docs/user-guide/neighborhood-filter.ipynb b/docs/user-guide/neighborhood-filter.ipynb new file mode 100644 index 000000000..9f662dbe4 --- /dev/null +++ b/docs/user-guide/neighborhood-filter.ipynb @@ -0,0 +1,576 @@ +{ + "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(r=5.0).mean()` |\n", + "| `UxDataset` | `ds.neighborhood(r=5.0).mean()` |\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` with a radius of 5° groups each face with every face\n", + "center within 5° of it; `mean` then reduces each of those groups.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7623eae2785240b9bd12b16a66d81610", + "metadata": {}, + "outputs": [], + "source": [ + "uxda_smooth = uxda.neighborhood(r=5.0).mean()\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(r=r)\n", + " .mean()\n", + " .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", + "Call the reduction you want as a method. The available ones are `mean`,\n", + "`sum`, `min`, `max`, `median`, `ptp`, `std`, `var`, `quantile`, and\n", + "`percentile`, each running a compiled kernel. Those taking a parameter\n", + "declare it in their own signature: `q` for `quantile` (0–1) and\n", + "`percentile` (0–100), `ddof` for `std` and `var`.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "b43b363d81ae4b689946ece5c682cd59", + "metadata": {}, + "outputs": [], + "source": [ + "# One grouping, reduced four different ways\n", + "nb = uxda.neighborhood(r=5.0)\n", + "\n", + "# 90th-percentile filter — highlights local maxima\n", + "uxda_p90 = nb.percentile(90)\n", + "\n", + "# Maximum filter\n", + "uxda_max = nb.max()\n", + "\n", + "# Median filter — robust to outliers\n", + "uxda_med = nb.median()\n", + "\n", + "# Local spread, as a sample standard deviation\n", + "uxda_std = nb.std(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 Method\n", + "\n", + "If you need something not in that list, hand it to `reduce` as a callable. It is\n", + "applied as `func(values, axis=-1)` to a block whose last axis is the neighborhood,\n", + "once per grid element, in Python — noticeably slower than the compiled methods, so\n", + "reach for it only when none of them 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(r=5.0).reduce(rms)\n", + "\n", + "# `functools.partial` also works, though `.percentile()` is the faster way here\n", + "uxda_p90_slow = uxda.neighborhood(r=5.0).reduce(partial(np.percentile, q=90))\n", + "print(\n", + " \"partial matches the compiled reduction:\",\n", + " np.allclose(uxda_p90_slow.values, uxda_p90.values),\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "0e382214b5f147d187d36a2058b9c724", + "metadata": {}, + "source": [ + "## Reusing a Neighborhood Across Reductions\n", + "\n", + "Each call to `neighborhood` searches the grid for the neighbors of every element.\n", + "That search usually costs far more than the reduction itself, so holding onto the\n", + "object and reducing it several times — as the cell above does — already avoids\n", + "repeating the expensive part.\n", + "\n", + "`Grid.neighborhood` goes one step further and shares that search across *variables*.\n", + "It is not bound to any data, so its reduction methods take the data as an argument.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5b09d5ef5b5e4bb6ab9b829b10b6a29f", + "metadata": {}, + "outputs": [], + "source": [ + "nb5 = uxda.uxgrid.neighborhood(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.mean(uxda)\n", + "spread = nb5.std(uxda)\n", + "p90 = nb5.percentile(uxda, 90)\n", + "\n", + "print(\n", + " \"identical to the data-bound call:\",\n", + " np.allclose(smooth.values, uxda.neighborhood(r=5.0).mean().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` works for any data element type. Here we create synthetic node-\n", + "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(r=10.0).mean()\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(r=10.0).mean()\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", + "\n", + "When a `UxDataArray` has extra leading dimensions (e.g. `time`), a neighborhood\n", + "reduction applies independently at each time step and preserves the full dimension\n", + "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(r=5.0).mean()\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` carries the same reductions, applying each to **every\n", + "data variable** that is mapped to a grid element. Variables without a grid dimension\n", + "(e.g. scalars or time-only arrays) are passed through unchanged.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "db7b79bc585a40fcaf58bf750017e135", + "metadata": {}, + "outputs": [], + "source": [ + "uxds_filtered = uxds.neighborhood(r=5.0).mean()\n", + "uxds_filtered" + ] + }, + { + "cell_type": "markdown", + "id": "916684f9a58a4a2aa5f864670399430d", + "metadata": {}, + "source": [ + "## Chaining with xarray Operations\n", + "\n", + "Because every reduction returns a proper `UxDataArray` with its `uxgrid`\n", + "preserved, 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(r=5.0).mean().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(r=5.0).mean().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(r=0.0).mean()\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(r=360.0).mean()\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`\n", + "- {py:meth}`uxarray.UxDataset.neighborhood`\n", + "- {py:meth}`uxarray.Grid.neighborhood` — a neighborhood shared across variables\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 aac16e3ab..813aa983f 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..0372a6742 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,361 @@ def test_data_location(): assert face_time.data_location == "face_centered" +class TestNeighborhood: + """Tests for ``UxDataArray.neighborhood`` and the reductions on it.""" + + @pytest.fixture + def vortex(self, gridpath, datasetpath): + """The ``psi`` field on outCSne30, which most of these tests reduce.""" + uxds = ux.open_dataset( + gridpath("ugrid", "outCSne30", "outCSne30.ug"), + datasetpath("ugrid", "outCSne30", "outCSne30_vortex.nc"), + ) + return uxds["psi"] + + def test_face_centered(self, vortex): + """A large enough radius should average every face together.""" + # radius of 0 should select each face's own coordinate, leaving the + # data unchanged + filtered = vortex.neighborhood(r=0.0).mean() + np.testing.assert_allclose(filtered.values, vortex.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 = vortex.neighborhood(r=360.0).mean() + np.testing.assert_allclose(filtered_all.values, vortex.values.mean()) + + assert isinstance(filtered, UxDataArray) + assert filtered.uxgrid == vortex.uxgrid + assert filtered.dims == vortex.dims + assert filtered.shape == vortex.shape + + def test_node_centered(self): + """Neighborhood reductions 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") + + np.testing.assert_allclose(uxda.neighborhood(r=0.0).mean().values, data) + np.testing.assert_allclose( + uxda.neighborhood(r=360.0).mean().values, data.mean() + ) + + def test_edge_centered(self): + """Neighborhood reductions 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") + + np.testing.assert_allclose(uxda.neighborhood(r=0.0).mean().values, data) + + def test_bound_object_reports_its_neighborhood(self, vortex): + """The object returned by ``neighborhood`` should describe the query it + holds, and expose it for reuse elsewhere.""" + nb = vortex.neighborhood(r=4.0) + + assert (nb.r, nb.on, nb.grid_dim) == (4.0, "face centers", "n_face") + assert nb.grid is vortex.uxgrid + assert nb.neighborhood.on == "face centers" + + 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 + + def test_extra_dimension_preserved(self, vortex): + """An extra leading (i.e. time) dimension should be preserved.""" + data = np.stack([vortex.values, vortex.values * 2.0]) + uxda_time = UxDataArray( + data, dims=["time", "n_face"], uxgrid=vortex.uxgrid, name="psi_time" + ) + + filtered = uxda_time.neighborhood(r=0.0).mean() + + 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(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(r=0.0).mean() + 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(r=360.0).mean() + 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(r=-1.0) + + def test_uses_spherical_tree_regardless_of_cached_tree(self): + """``r`` is documented in great-circle degrees, so the query 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(r=20.0).mean().values + + # Prime the cache with a cartesian tree, then reduce again + uxgrid.get_ball_tree( + coordinates="face centers", + coordinate_system="cartesian", + distance_metric="euclidean", + ) + np.testing.assert_allclose( + uxda.neighborhood(r=20.0).mean().values, expected + ) + + def test_auto_transpose_direct_on_uxdataarray(self, vortex): + """Reducing a (time, n_face) UxDataArray directly (without going + through UxDataset) should preserve the original dim order.""" + # Build a multi-dim UxDataArray with time as the FIRST (non-grid) dim + data = np.stack([vortex.values, vortex.values * 2.0]) # shape (2, n_face) + uxda_time = UxDataArray( + data, dims=["time", "n_face"], uxgrid=vortex.uxgrid, name="psi_time" + ) + + # n_face is already last: no transpose needed internally + filtered = uxda_time.neighborhood(r=0.0).mean() + assert filtered.dims == ("time", "n_face") + assert filtered.shape == (2, vortex.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(r=0.0).mean() + # Dim order must be restored to (n_face, time) + assert filtered2.dims == ("n_face", "time") + assert filtered2.shape == (vortex.shape[0], 2) + + # Every compiled reduction, with the NumPy expression it must equal and the + # arguments it takes. The reference runs through the generic callable path, + # so this pins each compiled kernel against the loop it bypasses. + COMPILED_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", COMPILED_REDUCTIONS) + def test_compiled_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 = uxda.neighborhood(r=20.0) + got = getattr(nb, name)(**kwargs).values + with warnings.catch_warnings(): + warnings.simplefilter("ignore", RuntimeWarning) # numpy all-NaN slices + expected = nb.reduce(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_escape_hatch(self, vortex): + """A user's own function, with no compiled equivalent, still works on + the ``axis=-1`` contract.""" + + # 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(r=3.0).reduce(rms) + assert filtered.shape == vortex.shape + assert np.all(filtered.values >= 0) + + def test_reduce_accepts_partial(self, vortex): + """``functools.partial`` binds a reduction's extra arguments on the + callable path, as it did before the compiled methods existed.""" + from functools import partial + + nb = vortex.neighborhood(r=5.0) + np.testing.assert_allclose( + nb.max().values, + nb.reduce(partial(np.percentile, q=100)).values, + ) + + 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(r=5.0).reduce(sum) + + def test_misspelled_reduction_is_an_attribute_error(self, vortex): + """A reduction that does not exist is not spelled at all, so it fails + as a missing attribute rather than at run time.""" + with pytest.raises(AttributeError, match="meen"): + vortex.neighborhood(r=1.0).meen() + + @pytest.mark.parametrize( + "name,kwargs,error,match", + [ + ("mean", {"q": 90}, TypeError, "unexpected keyword argument"), + ("quantile", {}, TypeError, "required positional argument"), + ("quantile", {"q": 90}, ValueError, "between 0 and 1"), + ("percentile", {"q": 101}, ValueError, "between 0 and 100"), + ], + ) + def test_invalid_reduction_arguments(self, name, kwargs, error, match): + """A reduction's parameters are its own signature, so bad input is + caught 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): + getattr(uxda.neighborhood(r=1.0), name)(**kwargs) + + def test_neighborhood_reuse(self, vortex): + """A neighborhood built from the grid must give the same answer as one + bound to the data -- the point of holding onto it is that several + variables cost one query.""" + nb = vortex.uxgrid.neighborhood(r=4.0) + + assert (nb.r, nb.on, nb.grid_dim) == (4.0, "face centers", "n_face") + + np.testing.assert_allclose( + nb.mean(vortex).values, + vortex.neighborhood(r=4.0).mean().values, + rtol=1e-12, + ) + np.testing.assert_allclose( + nb.percentile(vortex, 90).values, + vortex.neighborhood(r=4.0).percentile(90).values, + rtol=1e-12, + ) + + def test_neighborhood_rejects_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.neighborhood(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="reduces over 'n_face'"): + nb.mean(node_data) + + 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.mean(wrong_size) + + with pytest.raises(ValueError, match="Invalid `on`"): + uxgrid.neighborhood(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(r=2.0).mean() + + 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(r=2.0).mean() + + assert filtered.chunks is not None, "the reduction 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(r=2.0).mean().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(r=2.0).mean() + + 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(r=0.0).mean() + 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(r=2.0).mean() + assert filtered32.dtype == np.float64 + np.testing.assert_allclose( + filtered32.values, vortex.neighborhood(r=2.0).mean().values, + rtol=1e-6, + ) + + def test_uxgrid_None_is_invalid_in_uxdataarray(): """Ensures GridInvalidError gets raised if uxgrid=None when getting UxDataArray.uxgrid. Regression test for #1620. diff --git a/test/core/test_dataset.py b/test/core/test_dataset.py index 9b26d84d3..d69f86ba7 100644 --- a/test/core/test_dataset.py +++ b/test/core/test_dataset.py @@ -168,6 +168,123 @@ def test_uxdataset_to_array(): assert arr2.name == 'custom_name' +class TestNeighborhood: + """Tests for ``UxDataset.neighborhood`` and the reductions on it.""" + + def test_face_centered(self, gridpath, datasetpath): + """Ensures the dataset-level reduction matches the per-variable + ``UxDataArray.neighborhood`` results.""" + uxds = ux.open_dataset( + gridpath("ugrid", "outCSne30", "outCSne30.ug"), + datasetpath("ugrid", "outCSne30", "outCSne30_vortex.nc"), + ) + + filtered_ds = uxds.neighborhood(r=5.0).mean() + filtered_da = uxds["psi"].neighborhood(r=5.0).mean() + + 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(r=0.0).mean() + + nt.assert_allclose(filtered["face_var"].values, uxds["face_var"].values) + nt.assert_allclose(filtered["scalar_var"].values, uxds["scalar_var"].values) + + 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 reduction 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(r=20.0).percentile(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(r=20.0).percentile(90).values, + ) + + def test_one_query_reused_across_reductions(self): + """A DatasetNeighborhood holds its queries, so a second reduction on + the same object must not rebuild them.""" + from unittest.mock import patch + + import uxarray.grid.neighbors as neighbors + + uxgrid = ux.Grid.from_healpix(zoom=2) + rng = np.random.default_rng(0) + uxds = UxDataset( + data_vars={"face_a": ("n_face", rng.random(uxgrid.n_face))}, + uxgrid=uxgrid, + ) + + nb = uxds.neighborhood(r=20.0) + real = neighbors._csr_neighbors + with patch.object(neighbors, "_csr_neighbors", side_effect=real) as spy: + smooth, spread = nb.mean(), nb.std(ddof=1) + + assert spy.call_count == 1, ( + f"expected the query to be built once and reused, got {spy.call_count}" + ) + assert smooth["face_a"].shape == spread["face_a"].shape + + def test_callable_escape_hatch(self): + """``reduce`` applies a user's own function to every grid-mapped + variable.""" + uxgrid = ux.Grid.from_healpix(zoom=1) + rng = np.random.default_rng(0) + uxds = UxDataset( + data_vars={"face_a": ("n_face", rng.random(uxgrid.n_face))}, + uxgrid=uxgrid, + ) + + def rms(values, axis): + return np.sqrt(np.mean(values**2, axis=axis)) + + filtered = uxds.neighborhood(r=20.0).reduce(rms) + assert np.all(filtered["face_a"].values >= 0) + + def test_uxgrid_None_is_invalid_in_uxdataset(): """Ensures GridInvalidError gets raised if uxgrid=None when getting UxDataset.uxgrid. Regression test for #1620. 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 b62b7f1c9..ae2fea6cd 100644 --- a/uxarray/core/dataarray.py +++ b/uxarray/core/dataarray.py @@ -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 DataArrayNeighborhood, Neighborhood from uxarray.grid.validation import _check_duplicate_nodes_indices from uxarray.io._healpix import get_zoom_from_cells from uxarray.plot.accessor import UxDataArrayPlotAccessor @@ -2209,6 +2211,100 @@ def get_dual(self): return uxda + def _neighborhood_location(self, caller: str) -> str: + """Grid location this data is mapped to, in ``Neighborhood`` 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(self, r: float = 1.0) -> DataArrayNeighborhood: + """Groups this data by the elements within ``r`` degrees of each grid + element, to be reduced over by a method of the returned + :class:`DataArrayNeighborhood`. + + Each reduction replaces the value at every grid element with a + reduction of all elements within a circular neighborhood of radius + ``r``, as in a smoothing filter. + + Parameters + ---------- + r : float, default=1. + Radius of the neighborhood, in degrees. + + Returns + ------- + DataArrayNeighborhood + Bound to this data, so its reduction methods take only the + parameters of the reduction: ``mean()``, ``sum()``, ``min()``, + ``max()``, ``median()``, ``ptp()``, ``std(ddof)``, ``var(ddof)``, + ``quantile(q)``, ``percentile(q)``, or ``reduce(func)`` for + anything else. Each returns a ``UxDataArray`` of float64. + + Raises + ------ + DataCenteringError (subclass of ValueError) + If the data is not mapped to nodes, edges, or faces. + + Notes + ----- + ``r`` is a great-circle distance in degrees. An element's neighborhood + overlaps those of the elements around it, and every element is its own + neighbor at distance 0, so ``r = 0`` returns the data unchanged and the + result never contains spurious ``NaN``. + + Building this queries the grid for neighbors, which usually costs more + than the reduction itself. That query is what the returned object holds + on to, so several reductions at one radius should share one call rather + than repeat it. To share it across variables too, build the + neighborhood from the grid instead, with :meth:`Grid.neighborhood`. + + 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(r=5.0).mean() + + Reductions taking a parameter receive it as a keyword argument: + + >>> p90 = uxda.neighborhood(r=5.0).percentile(90) + >>> spread = uxda.neighborhood(r=5.0).std(ddof=1) + + Several reductions at one radius share the neighbor query: + + >>> nb = uxda.neighborhood(r=5.0) + >>> smoothed, spread = nb.mean(), nb.std() + + See Also + -------- + DataArrayNeighborhood : The reductions available on the returned object. + Grid.neighborhood : Neighborhood shared across several variables. + 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. + """ + neighborhood = Neighborhood( + self.uxgrid, + r=r, + on=self._neighborhood_location("neighborhood"), + ) + return DataArrayNeighborhood(neighborhood, self) + 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 129466589..2b4c8fca7 100644 --- a/uxarray/core/dataset.py +++ b/uxarray/core/dataset.py @@ -18,6 +18,7 @@ from uxarray.formatting_html import dataset_repr from uxarray.grid import Grid from uxarray.grid.dual import construct_dual +from uxarray.grid.neighbors import DatasetNeighborhood from uxarray.grid.validation import _check_duplicate_nodes_indices from uxarray.io._healpix import get_zoom_from_cells from uxarray.plot.accessor import UxDatasetPlotAccessor @@ -676,6 +677,48 @@ def to_array( return UxDataArray(xarr, uxgrid=self._uxgrid) # _uxgrid not uxgrid; converting to UxDataArray is not a grid-aware method. + def neighborhood(self, r: float = 1.0) -> DatasetNeighborhood: + """Groups every grid-mapped data variable by the elements within ``r`` + degrees of each grid element, to be reduced over by a method of the + returned :class:`DatasetNeighborhood`. + + Parameters + ---------- + r : float, default=1. + Radius of the neighborhood, in degrees. + + Returns + ------- + DatasetNeighborhood + Carrying the same reductions as :meth:`UxDataArray.neighborhood`, + applied to every data variable at once. Each returns a + ``UxDataset``. + + Notes + ----- + Variables without a grid dimension are passed through unchanged. + + Variables mapped to the same grid location share one neighbor query, so + reducing 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(r=5.0).mean() + + See Also + -------- + UxDataArray.neighborhood : Reduce a single data variable. + Grid.neighborhood : Neighborhood for one grid location, without data. + UxDataArray.zonal_mean : Average over latitude bands. + UxDataArray.azimuthal_mean : Average over rings of constant great-circle distance. + """ + return DatasetNeighborhood(self, r=r) + 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 711b30389..40ca7710f 100644 --- a/uxarray/grid/grid.py +++ b/uxarray/grid/grid.py @@ -61,6 +61,7 @@ from uxarray.grid.neighbors import ( BallTree, KDTree, + Neighborhood, SpatialHash, _populate_edge_face_distances, _populate_edge_node_distances, @@ -1790,7 +1791,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" @@ -1807,7 +1808,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, @@ -1815,12 +1826,50 @@ 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 neighborhood(self, r: float = 1.0, on: str = "face centers") -> Neighborhood: + """Finds the grid elements within ``r`` degrees of every element of + ``on``, returning a reusable :class:`Neighborhood`. + + 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` + repeatedly, which rebuilds it on every call. + + Unlike :meth:`UxDataArray.neighborhood`, the result is not bound to + any data, so its reduction methods take the data to reduce as an + argument. That is what lets several variables share one query. + + 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 neighborhood on: "nodes", + "edge centers", or "face centers". + + Returns + ------- + Neighborhood + + Examples + -------- + >>> import uxarray as ux + >>> uxds = ux.tutorial.open_dataset("outCSne30-vortex") # doctest: +SKIP + >>> nb = uxds.uxgrid.neighborhood(r=5.0) # doctest: +SKIP + >>> smooth = nb.mean(uxds["psi"]) # doctest: +SKIP + >>> p90 = nb.percentile(uxds["psi"], q=90) # doctest: +SKIP + + See Also + -------- + Neighborhood : The reductions available on the returned object. + UxDataArray.neighborhood : Neighborhood bound to a single variable. + UxDataset.neighborhood : Neighborhood across every variable in a dataset. + """ + return Neighborhood(self, r=r, on=on) + def _get_scipy_kd_tree( self, coordinates: str | None = "face", reconstruct: bool = False ): @@ -1907,7 +1956,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, @@ -1916,10 +1973,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 195bf4138..03a4c7cd5 100644 --- a/uxarray/grid/neighbors.py +++ b/uxarray/grid/neighbors.py @@ -1,9 +1,17 @@ +import warnings +from typing import Callable + 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 +from uxarray.constants import ( + ERROR_TOLERANCE, + GRID_DIMS, + INT_DTYPE, + INT_FILL_VALUE, +) from uxarray.errors import DimensionError @@ -1130,3 +1138,728 @@ 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 + + +# Reducers take ``(window, param)``; those without a parameter ignore the +# second argument. Numba keys its cache by code object rather than qualified +# name, so the identically-named lambdas below 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) + + +# One compiled kernel per reduction. The methods on ``Neighborhood`` below name +# these directly, so there is no dispatch table between the public API and the +# gufuncs: a reduction is reachable only if a method exists for it, and a method +# can only reach the kernel it names. ``Neighborhood`` is the only class that +# names them -- the data-bound classes reach a kernel by naming the +# ``Neighborhood`` method for it, so there is one place per reduction where its +# kernel and parameter are chosen. +_MEAN_KERNEL = _make_kernel(lambda window, _: np.mean(window)) +_SUM_KERNEL = _make_kernel(lambda window, _: np.sum(window)) +_MIN_KERNEL = _make_kernel(lambda window, _: np.min(window)) +_MAX_KERNEL = _make_kernel(lambda window, _: np.max(window)) +_PTP_KERNEL = _make_kernel(lambda window, _: np.max(window) - np.min(window)) +_MEDIAN_KERNEL = _make_kernel(_median) +_VAR_KERNEL = _make_kernel(_variance) +_STD_KERNEL = _make_kernel(lambda window, ddof: np.sqrt(_variance(window, ddof))) +# ``percentile`` is ``quantile`` on a 0-100 scale, so both methods rescale onto +# this one kernel rather than compiling a near-duplicate. +_QUANTILE_KERNEL = _make_kernel(lambda window, q: np.quantile(window, q)) + + +def _as_quantile(q, scale: float): + """Validates ``q`` on a 0-``scale`` scale and returns it as a 0-1 fraction.""" + value = float(q) + if not 0.0 <= value <= scale: + raise ValueError(f"`q` must be between 0 and {scale:g}, but got {q!r}.") + return value / scale + + +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. + + A neighborhood is 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 Neighborhood: + """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. + + Each reduction is a method: :meth:`mean`, :meth:`std`, :meth:`percentile` + and so on run a compiled kernel, and :meth:`reduce` takes an arbitrary + callable for anything without one. Every one takes the data to reduce as + its first argument, so several variables can share one query. + + Parameters + ---------- + grid : Grid + Grid whose elements define the neighborhood. + r : float, default=1. + Radius of the neighborhood, in degrees of great-circle distance. + on : str, default="face centers" + Grid location the neighborhood is 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.neighborhood(r=5.0) # doctest: +SKIP + >>> smooth = nb.mean(uxds["psi"]) # doctest: +SKIP + >>> spread = nb.std(uxds["psi"], ddof=1) # doctest: +SKIP + + See Also + -------- + UxDataArray.neighborhood : Neighborhood around a single variable's elements. + UxDataset.neighborhood : Neighborhood around every variable in a dataset. + """ + + 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 neighborhood was 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 neighborhood is 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 mean(self, uxda): + """Mean of each neighborhood.""" + return self._apply_kernel(uxda, _MEAN_KERNEL, 0.0) + + def sum(self, uxda): + """Sum of each neighborhood.""" + return self._apply_kernel(uxda, _SUM_KERNEL, 0.0) + + def min(self, uxda): + """Smallest value in each neighborhood.""" + return self._apply_kernel(uxda, _MIN_KERNEL, 0.0) + + def max(self, uxda): + """Largest value in each neighborhood.""" + return self._apply_kernel(uxda, _MAX_KERNEL, 0.0) + + def ptp(self, uxda): + """Peak-to-peak spread (``max - min``) of each neighborhood.""" + return self._apply_kernel(uxda, _PTP_KERNEL, 0.0) + + def median(self, uxda): + """Median of each neighborhood.""" + return self._apply_kernel(uxda, _MEDIAN_KERNEL, 0.0) + + def var(self, uxda, ddof: int = 0): + """Variance of each neighborhood, with ``ddof`` delta degrees of + freedom.""" + return self._apply_kernel(uxda, _VAR_KERNEL, float(ddof)) + + def std(self, uxda, ddof: int = 0): + """Standard deviation of each neighborhood, with ``ddof`` delta degrees + of freedom.""" + return self._apply_kernel(uxda, _STD_KERNEL, float(ddof)) + + def quantile(self, uxda, q: float): + """Quantile ``q`` (between 0 and 1) of each neighborhood.""" + return self._apply_kernel(uxda, _QUANTILE_KERNEL, _as_quantile(q, 1.0)) + + def percentile(self, uxda, q: float): + """Percentile ``q`` (between 0 and 100) of each neighborhood.""" + return self._apply_kernel(uxda, _QUANTILE_KERNEL, _as_quantile(q, 100.0)) + + def reduce(self, uxda, func: Callable): + """Reduces each neighborhood with an arbitrary callable. + + This is the escape hatch for reductions with no method of their own. + ``func`` is applied as ``func(values, axis=-1)`` over a block whose last + axis is the neighborhood, once per grid element, in Python — which is + considerably slower than the compiled methods above. Prefer a method + where one exists. + + Parameters + ---------- + uxda : UxDataArray + Data to reduce. + func : Callable + Reduction to apply. Must accept an ``axis`` keyword argument. Use + ``functools.partial`` to bind any further arguments. + + Returns + ------- + UxDataArray + Reduced data as float64, with the input's dimension order. Lazy if + the input was lazy. + + Examples + -------- + >>> from scipy.stats import skew # doctest: +SKIP + >>> nb.reduce(uxds["psi"], skew) # doctest: +SKIP + """ + + def run(block, arrays): + return _neighborhood_reduce(block, *arrays, func) + + return self._apply(uxda, run) + + def _apply_kernel(self, uxda, kernel, param: float): + """Runs a compiled ``kernel`` over every neighborhood.""" + + def run(block, arrays): + # 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, *arrays, param) + + return self._apply(uxda, run) + + def _apply(self, uxda, run): + """Validates ``uxda`` against this neighborhood and maps ``run`` over + it, one NumPy block at a time with the grid dimension last.""" + # 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"This neighborhood is built on {self._on!r} and reduces 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"this neighborhood describes {self._counts.size}. The data is " + f"probably mapped to a different grid." + ) + + arrays = (self._flat, self._starts, self._counts) + + def _apply(block): + return run(block, arrays) + + 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)) + + +class _BoundNeighborhoodReductions: + """The reduction vocabulary of a :class:`Neighborhood` whose data is already + supplied, spelled once for every class that carries it. + + A reduction has to be an attribute of ``Neighborhood``, which takes its data + as an argument, and of the classes below, which already hold theirs. Those + two signatures cannot share a definition, so the vocabulary is spelled twice + in all -- but only twice, and only one of the two chooses a kernel. + + This is the half that does not. Each method here hands ``_map`` the + ``Neighborhood`` method it stands for, as the function object itself rather + than a name to look up later, and subclasses implement ``_map`` to say + *which data* to call it on. That is the only thing that differs between a + neighborhood bound to one variable and one bound to a whole dataset, so + everything else is written once. + + Passing the method rather than a name is what keeps the two halves honest: + the reference is resolved when this class is created, so a reduction that + ``Neighborhood`` does not define cannot be spelled here at all, and a bound + method cannot reach a different kernel, or a different ``ddof``, than the + unbound one it names. + """ + + def _map(self, reduction: Callable, *args, **kwargs): + """Calls ``reduction(neighborhood, uxda, *args, **kwargs)`` on the data + this is bound to, returning the result in the same container. + + Subclasses must implement this; it is the only thing they need to. + """ + raise NotImplementedError + + def mean(self): + """Mean of each neighborhood.""" + return self._map(Neighborhood.mean) + + def sum(self): + """Sum of each neighborhood.""" + return self._map(Neighborhood.sum) + + def min(self): + """Smallest value in each neighborhood.""" + return self._map(Neighborhood.min) + + def max(self): + """Largest value in each neighborhood.""" + return self._map(Neighborhood.max) + + def ptp(self): + """Peak-to-peak spread (``max - min``) of each neighborhood.""" + return self._map(Neighborhood.ptp) + + def median(self): + """Median of each neighborhood.""" + return self._map(Neighborhood.median) + + def var(self, ddof: int = 0): + """Variance of each neighborhood, with ``ddof`` delta degrees of + freedom.""" + return self._map(Neighborhood.var, ddof=ddof) + + def std(self, ddof: int = 0): + """Standard deviation of each neighborhood, with ``ddof`` delta degrees + of freedom.""" + return self._map(Neighborhood.std, ddof=ddof) + + def quantile(self, q: float): + """Quantile ``q`` (between 0 and 1) of each neighborhood.""" + return self._map(Neighborhood.quantile, q) + + def percentile(self, q: float): + """Percentile ``q`` (between 0 and 100) of each neighborhood.""" + return self._map(Neighborhood.percentile, q) + + def reduce(self, func: Callable): + """Reduces each neighborhood with an arbitrary callable. + + See :meth:`Neighborhood.reduce`, which this supplies the data to. + """ + return self._map(Neighborhood.reduce, func) + + +class DataArrayNeighborhood(_BoundNeighborhoodReductions): + """Neighborhood of radius ``r`` around the elements of one data variable, + ready to be reduced over. + + Carries the same reduction methods as :class:`Neighborhood`, with the data + already supplied, so they take only the parameters of the reduction itself. + The underlying neighbor query is built once and reused by every reduction + called on this object. + + Returned by :meth:`UxDataArray.neighborhood`; not constructed directly. + + Examples + -------- + >>> import uxarray as ux + >>> uxds = ux.tutorial.open_dataset("outCSne30-vortex") # doctest: +SKIP + >>> nb = uxds["psi"].neighborhood(r=5.0) # doctest: +SKIP + >>> smooth, spread = nb.mean(), nb.std(ddof=1) # doctest: +SKIP + + See Also + -------- + Neighborhood : The unbound neighborhood this wraps. + """ + + def __init__(self, neighborhood: Neighborhood, uxda): + self._neighborhood = neighborhood + self._uxda = uxda + + @property + def neighborhood(self) -> Neighborhood: + """The unbound :class:`Neighborhood` these reductions run on. + + Reusable across other variables at the same grid location, without + repeating the neighbor query. + """ + return self._neighborhood + + @property + def grid(self): + """Grid the neighborhood was built from.""" + return self._neighborhood.grid + + @property + def r(self) -> float: + """Neighborhood radius, in degrees.""" + return self._neighborhood.r + + @property + def on(self) -> str: + """Grid location the neighborhood is centered on.""" + return self._neighborhood.on + + @property + def grid_dim(self) -> str: + """Name of the grid dimension this reduces over.""" + return self._neighborhood.grid_dim + + @property + def n_neighbors(self) -> xr.DataArray: + """Number of elements in each neighborhood.""" + return self._neighborhood.n_neighbors + + def __repr__(self) -> str: + return ( + f"" + ) + + def _map(self, reduction: Callable, *args, **kwargs): + """Applies ``reduction`` to the one variable this is bound to.""" + return reduction(self._neighborhood, self._uxda, *args, **kwargs) + + +class DatasetNeighborhood(_BoundNeighborhoodReductions): + """Neighborhood of radius ``r`` around every grid-mapped variable of a + dataset, ready to be reduced over. + + Carries the same reduction methods as :class:`Neighborhood`, applying each + to every data variable at once. Variables without a grid dimension pass + through unchanged, and variables mapped to the same grid location share one + neighbor query, so a reduction costs one query per location present rather + than one per variable. + + Returned by :meth:`UxDataset.neighborhood`; not constructed directly. + + Examples + -------- + >>> import uxarray as ux + >>> uxds = ux.tutorial.open_dataset("outCSne30-vortex") # doctest: +SKIP + >>> smooth = uxds.neighborhood(r=5.0).mean() # doctest: +SKIP + + See Also + -------- + Neighborhood : The single-location neighborhood this builds on. + """ + + def __init__(self, uxds, r: float = 1.0): + self._uxds = uxds + self._r = float(r) + # Built on first use, keyed by grid location, since the query depends + # only on (grid, location, radius) -- not on the data. + self._by_location: dict[str, Neighborhood] = {} + + @property + def r(self) -> float: + """Neighborhood radius, in degrees.""" + return self._r + + def __repr__(self) -> str: + return ( + f"" + ) + + def _map(self, reduction: Callable, *args, **kwargs): + """Applies ``reduction`` to every grid-mapped variable, sharing a + neighbor query between variables at the same grid location.""" + destination_uxds = self._uxds._copy() + + for var_name in self._uxds.data_vars: + uxda = self._uxds[var_name] + + # Skip variables that are not mapped to a grid element. + if not any(dim in GRID_DIMS for dim in uxda.dims): + continue + + location = uxda._neighborhood_location("neighborhood") + if location not in self._by_location: + self._by_location[location] = Neighborhood( + self._uxds.uxgrid, r=self._r, on=location + ) + + # The Neighborhood methods restore the input dimension order, so + # it is always preserved. + destination_uxds[var_name] = reduction( + self._by_location[location], uxda, *args, **kwargs + ) + + return destination_uxds