Apply function to points within circular neighborhood by ahijevyc · Pull Request #941 · UXARRAY/uxarray

ahijevyc · 2024-09-09T16:44:16Z

Apply a neighborhood filter within a circular radius r to a UxDataset or UxDataArray.

Issue #930

Overview

This is kind of like uxarray.UxDataArray.inverse_distance_weighted_remap , but the neighborhood is defined by distance, not a number of nearest neighbors. This is ideally suited for a variable resolution mesh, in which a constant of neighbors doesn't have a constant sized neighborhood. Another difference is that this neighborhood filter does not weight data by inverse distance.

Just like uxarray.UxDataArray.subset.bounding_circle this function uses ball_tree.query_radius to select grid elements in a circular neighborhood, but this function finds the neighborhood for all elements in grid, not just one center_coordinate.

The filter function func may be a user-defined function, but uses np.mean by default. It could be min, max, np.median. It can even use functions that require additional arguments, like np.percentile if you supply the argument(s) with functools.partial (see below)

Expected Usage

from functools import partial
import numpy as np
import uxarray

grid_path = "/glade/campaign/mmm/wmr/weiwang/cps/irma3/2020/tk707_conus/init.nc"
data_path = "/glade/campaign/mmm/wmr/weiwang/cps/irma3/mp6/tk707/diag.2017-09-07_09.00.00.nc"
uxds = uxarray.open_mfdataset(
    grid_path,
    data_path
)

# Trim domain
lon_bounds = (-74, -64)
lat_bounds = (18, 24)
uxda = uxds["refl10cm_max"].isel(Time=0).subset.bounding_box(lon_bounds, lat_bounds)

# this is how you use this function to smooth with 0.25-deg filter.
uxda_mean = uxda.neighborhood_filter(func=np.mean, r=0.25)


# this is another way to use this function with np.percentile
uxda_max = uxda.neighborhood_filter(func=partial(np.percentile, q=90), r=0.25)

(uxda.plot.rasterize() + uxda_mean.plot.rasterize() + uxda_max.plot.rasterize()).cols(1)

PR Checklist

General

An issue is linked created and linked
Add appropriate labels
Filled out Overview and Expected Usage (if applicable) sections

Testing

Adequate tests are created if there is new functionality
Tests cover all possible logical paths in your function
Tests are not too basic (such as simply calling a function and nothing else)

Documentation

Docstrings have been added to all new functions
Docstrings have updated with any function changes
Internal functions have a preceding underscore (_) and have been added to docs/internal_api/index.rst
User functions have been added to docs/user_api/index.rst

Examples

Any new notebook examples added to docs/examples/ folder
Clear the output of all cells before committing
New notebook files added to docs/examples.rst toctree
New notebook files added to new entry in docs/gallery.yml with appropriate thumbnail photo in docs/_static/thumbnails/

philipc2 · 2024-09-09T17:02:48Z

+        # face centers, edge centers or nodes.
+        tree = self.uxgrid.get_ball_tree(coordinates=data_mapping, reconstruct=True)


@aaronzedwick

We should probably fix this logic in get_ball_tree(), since we shouldn't need to manually set reconstruct=False

if self._ball_tree is None or reconstruct: self._ball_tree = BallTree( self, coordinates=coordinates, distance_metric=distance_metric, coordinate_system=coordinate_system, reconstruct=reconstruct, ) else: if coordinates != self._ball_tree._coordinates: self._ball_tree.coordinates = coordinates

The coordinates != self._ball_tree._coordinates check should be included in the first if

That makes sense. So, move the coordinates check to the if-clause like this?

if ( self._ball_tree is None or coordinates != self._ball_tree._coordinates or reconstruct ): self._ball_tree = BallTree( self, coordinates=coordinates, distance_metric=distance_metric, coordinate_system=coordinate_system, reconstruct=reconstruct, )

What if the coordinate_system is different? Would that also require a newly constructed tree?

Whatever logic is fixed in Grid.get_ball_tree should also be applied to Grid.get_kdtree.

checking coordinate system also (coordinate_system is not a hidden variable of _ball_tree; it has no underscore):

if ( self._ball_tree is None or coordinates != self._ball_tree._coordinates or coordinate_system != self._ball_tree.coordinate_system or reconstruct ): self._ball_tree = BallTree( self, coordinates=coordinates, distance_metric=distance_metric, coordinate_system=coordinate_system, reconstruct=reconstruct, )

philipc2

A few initial comments:

ahijevyc · 2024-10-02T17:18:58Z

+
+        coordinate_system = tree.coordinate_system
+
+        if coordinate_system == "spherical":
+            if data_mapping == "nodes":
+                lon, lat = (
+                    self.uxgrid.node_lon.values,
+                    self.uxgrid.node_lat.values,
+                )
+            elif data_mapping == "face centers":
+                lon, lat = (
+                    self.uxgrid.face_lon.values,
+                    self.uxgrid.face_lat.values,
+                )
+            elif data_mapping == "edge centers":
+                lon, lat = (
+                    self.uxgrid.edge_lon.values,
+                    self.uxgrid.edge_lat.values,
+                )
+            else:
+                raise ValueError(
+                    f"Invalid data_mapping. Expected 'nodes', 'edge centers', or 'face centers', "
+                    f"but received: {data_mapping}"
+                )
+
+            dest_coords = np.vstack((lon, lat)).T
+
+        elif coordinate_system == "cartesian":
+            if data_mapping == "nodes":
+                x, y, z = (
+                    self.uxgrid.node_x.values,
+                    self.uxgrid.node_y.values,
+                    self.uxgrid.node_z.values,
+                )
+            elif data_mapping == "face centers":
+                x, y, z = (
+                    self.uxgrid.face_x.values,
+                    self.uxgrid.face_y.values,
+                    self.uxgrid.face_z.values,
+                )
+            elif data_mapping == "edge centers":
+                x, y, z = (
+                    self.uxgrid.edge_x.values,
+                    self.uxgrid.edge_y.values,
+                    self.uxgrid.edge_z.values,
+                )
+            else:
+                raise ValueError(
+                    f"Invalid data_mapping. Expected 'nodes', 'edge centers', or 'face centers', "
+                    f"but received: {data_mapping}"
+                )


Use #974 's new remap.utils._remap_grid_parse instead of this code block.

Kept neighborhood and dual additions

philipc2 · 2025-03-07T22:23:17Z

HI @ahijevyc

Apologies for not getting to this PR earlier.

Looking at the implementation here, it looks great. It does however bring to light a possible need for us to consider a better, more streamlined, approach to handling these types of groupings and then applying some function on the result.

I mention this because of our Topological Aggregations. For this family of functions, we have distinct methods (i.e. topological_mean()), which looking back at doesn't seem like the preferred approach, especially if we plan to implement groupings like the neighborhood one and perhaps other spatial ones.

Very generally speaking, these functions essentially:

Group unstructured grid elements based on some condition/algorithm. Here we use the KD/BallTree to determine the candidate elements, while in the topological aggregations we use the connectivity information
Apply some function to the grouping (i.e. mean())
Store the results back on the unstructured grid element (node, edge, or face)

I wonder if this would be a good opportunity to extend the inherited .groupby() method from Xarray to support these spatial groupings.

I'm not sure of calling these approaches "kernels" is appropriate, but for the sake of this example, we could provide spatial kernels the user could pass into groupby() and then perform aggregations directly on the result. This feels much more in line with Xarray's design philosophy.

# radial neighborhood of r=0.25
uxds['t2m'].groupby(kernel=ux.BoundingCircle(r=0.25)).mean()

# 2 deg by 2 deg bounding box 
uxds['t2m'].groupby(kernel=ux.BoundingBox(dlon=2, dlat=2))

# group the nodes that surround each face and find the maximum
uxds['node_centered_var'].groupby(kernel=ux.FaceNode()).max()

# this is equivalent to the following in the current release
uxds['node_centered_var'].topological_max(destination='face')

I'll ping @aaronzedwick and @erogluorhan for their thoughts on this. I personally really like the design above and think that it aligns well with the overall design.

aaronzedwick · 2025-03-10T16:34:00Z

HI @ahijevyc

Apologies for not getting to this PR earlier.

Looking at the implementation here, it looks great. It does however bring to light a possible need for us to consider a better, more streamlined, approach to handling these types of groupings and then applying some function on the result.

I mention this because of our Topological Aggregations. For this family of functions, we have distinct methods (i.e. topological_mean()), which looking back at doesn't seem like the preferred approach, especially if we plan to implement groupings like the neighborhood one and perhaps other spatial ones.

Very generally speaking, these functions essentially:

Group unstructured grid elements based on some condition/algorithm. Here we use the KD/BallTree to determine the candidate elements, while in the topological aggregations we use the connectivity information

Apply some function to the grouping (i.e. mean())

Store the results back on the unstructured grid element (node, edge, or face)

I wonder if this would be a good opportunity to extend the inherited .groupby() method from Xarray to support these spatial groupings.

I'm not sure of calling these approaches "kernels" is appropriate, but for the sake of this example, we could provide spatial kernels the user could pass into groupby() and then perform aggregations directly on the result. This feels much more in line with Xarray's design philosophy.
# radial neighborhood of r=0.25
uxds['t2m'].groupby(kernel=ux.BoundingCircle(r=0.25)).mean()

# 2 deg by 2 deg bounding box 
uxds['t2m'].groupby(kernel=ux.BoundingBox(dlon=2, dlat=2))

# group the nodes that surround each face and find the maximum
uxds['node_centered_var'].groupby(kernel=ux.FaceNode()).max()

# this is equivalent to the following in the current release
uxds['node_centered_var'].topological_max(destination='face')
I'll ping @aaronzedwick and @erogluorhan for their thoughts on this. I personally really like the design above and think that it aligns well with the overall design.

That is interesting. You suggesting changing the way we do aggregations entirely? Then this would affect the reduction PR I am working on then. Perhaps this PR could implement that change if you wish. I am fine with this, if you want to, it sounds like it would be intuitive.

philipc2 · 2025-03-10T17:53:25Z

You suggesting changing the way we do aggregations entirely?

We should be changing the API for it, since the current one is not sustainable if we plan to implement many times of spatial grouping methods.

The underlying functionality would remain mostly unchanged.

Perhaps this PR could implement that change if you wish.

If we do decide to do this, I'll open up a separate PR starting with the topological aggregations.

aaronzedwick · 2025-03-10T20:03:00Z

You suggesting changing the way we do aggregations entirely?

We should be changing the API for it, since the current one is not sustainable if we plan to implement many times of spatial grouping methods.

The underlying functionality would remain mostly unchanged.

Perhaps this PR could implement that change if you wish.

If we do decide to do this, I'll open up a separate PR starting with the topological aggregations.

So would the reductions PR be obsolete?

philipc2 · 2025-03-10T20:05:53Z

You suggesting changing the way we do aggregations entirely?

We should be changing the API for it, since the current one is not sustainable if we plan to implement many times of spatial grouping methods.
The underlying functionality would remain mostly unchanged.

Perhaps this PR could implement that change if you wish.

If we do decide to do this, I'll open up a separate PR starting with the topological aggregations.

So would the reductions PR be obsolete?

No. The underlying implementation would remain the same, since we would still need those implemented.

This would just provide a different interface for it, with a more "Xarray-like" interface.

aaronzedwick · 2025-03-10T20:15:35Z

You suggesting changing the way we do aggregations entirely?

We should be changing the API for it, since the current one is not sustainable if we plan to implement many times of spatial grouping methods.
The underlying functionality would remain mostly unchanged.

Perhaps this PR could implement that change if you wish.

If we do decide to do this, I'll open up a separate PR starting with the topological aggregations.

So would the reductions PR be obsolete?

No. The underlying implementation would remain the same, since we would still need those implemented.

This would just provide a different interface for it, with a more "Xarray-like" interface.

Ah, okay, I see. That makes sense, thanks for the clarification!

erogluorhan · 2026-02-26T19:06:29Z

pre-commit.ci autofix

for more information, see https://pre-commit.ci

erogluorhan · 2026-02-26T19:11:03Z

pre-commit.ci autofix

rajeeja · 2026-07-08T20:13:23Z

@ahijevyc are you still looking into this? Let me know if not, I can takeover..

ahijevyc · 2026-07-08T21:09:44Z

@ahijevyc are you still looking into this? Let me know if not, I can takeover..
I appreciate that. No I am not still looking at this. It would take a while for me to spin up again.

erogluorhan · 2026-07-17T17:16:39Z

@ahijevyc your last message was a typo I think as it doesn't show any response. Please let us know if you are still looking into it. Otherwise, let @rajeeja take this one since this is an important functionality.

ahijevyc · 2026-07-17T17:33:31Z

@erogluorhan @rajeeja I forgot that I assigned myself to this! Please take it over, Rajeev, and thank you for that.

…_filter # Conflicts: # uxarray/core/dataset.py

- Fix Grid.get_ball_tree()/get_kd_tree() caching to rebuild when the coordinates or coordinate_system changes, instead of mutating the cached tree in place (per review discussion). - Move the bulk of UxDataArray.neighborhood_filter's computation into a new _neighborhood_filter() helper in uxarray.grid.neighbors, with UxDataArray/UxDataset now acting as thin wrappers around it. - Fix a bug where func was applied across all axes instead of just the grid axis, which collapsed extra dimensions (e.g. time) in the output. - Bring UxDataset.neighborhood_filter's docstring in line with the UxDataArray implementation. - Add tests for face/node/edge-centered data, custom functions via functools.partial, extra dimension preservation, and dataset-level filtering. - Document neighborhood_filter in docs/api.rst.

…meters

…ter' into ahijevyc/neighborhood_filter

- neighbors.py: np.empty -> np.full(np.nan) so empty neighborhoods yield NaN instead of uninitialized garbage memory - dataarray.py: replace bare assert with raise ValueError; auto-transpose so grid dim is last before filtering, restore original dim order after - dataset.py: simplify loop now that UxDataArray handles the transpose - grid.py: fix get_ball_tree docstring default coordinate_system 'spherical' - test_dataarray.py: add test_empty_neighborhood_returns_nan and test_auto_transpose_direct_on_uxdataarray tests - docs: add neighborhood-filter.ipynb user guide and register in userguide.rst

review-notebook-app · 2026-07-28T17:18:49Z

Check out this pull request on

See visual diffs & provide feedback on Jupyter Notebooks.

Powered by ReviewNB

Two related fixes in UxDataArray: 1. _copy() default deep behavior: kwargs.get('deep', None) returned None (falsy) when called with no arguments, so the uxgrid was always shallow- copied even though xarray's own default is deep=True. Changed the fallback to True so _copy() matches xarray's documented default. 2. neighborhood_filter memory optimization: the filter was calling _copy() (deep copy of data) then immediately overwriting .data with the freshly computed result, wasting a full copy of the input array. Switch to _copy(data=destination_data, deep=False) which: - Passes the filtered data directly, skipping the redundant deep copy - Uses deep=False so the returned UxDataArray shares the same uxgrid object (appropriate: the filtered result lives on the same grid topology) - Preserves all metadata (name, attrs, coords) as before

UxDataArray.neighborhood_filter and UxDataset.neighborhood_filter now include Examples and See Also docstring sections, making them consistent with xarray conventions and discoverable from help() / sphinx docs. The user-guide notebook is also expanded (30 cells): - Use ux.tutorial datasets (no raw file paths) - Add before/after visualizations (.plot.polygons) - Cover face-, node-, and edge-centered data - Show multi-dimensional (time x space) usage - Demonstrate chaining with xarray .where() and .groupby() - Add radius sweep comparison (0 deg, 2.5 deg, 5 deg, 10 deg) - Add API reference section with cross-links to related methods

rajeeja · 2026-07-28T18:29:40Z

@ahijevyc — since GitHub blocks adding the PR author as a reviewer, could you take a look when you get a chance? Key changes since your last review round:

Bug fixes:

np.empty → np.full(…, np.nan) in neighbors.py — empty neighborhoods now correctly return NaN instead of arbitrary memory
Auto-transpose for non-last grid dim — (time, n_face) arrays now work without the caller having to manually transpose
raise ValueError instead of bare assert for dimension mismatch (cleaner user-facing error)
_copy(data=…, deep=False) in the filter path — skips a redundant deep copy of the source data, shares the same uxgrid reference

API compliance / xarray parity:

Investigated whether neighborhood_filter should use groupby: conclusion is no. Neighborhoods are overlapping spatial balls — a single element belongs to many neighborhoods. xarray's groupby requires non-overlapping partitions. The direct-method pattern is consistent with azimuthal_mean and zonal_mean.
Philip has been removed as a reviewer; @erogluorhan has been added.

Docs:

Expanded user-guide notebook (30 cells): covers face-, node-, and edge-centered data; multi-dim (time × space); xarray chaining (.where(), .groupby()); radius sweep comparison; ux.tutorial datasets throughout (no hardcoded file paths)
Added Examples and See Also sections to both UxDataArray.neighborhood_filter and UxDataset.neighborhood_filter docstrings

Happy to discuss further or adjust anything before merge!

Philip is no longer on the project. Dismissing stale review — all concerns have been addressed in subsequent commits.

ahijevyc · 2026-07-29T16:59:24Z

If the search radius is so small that no neighbor is found for a given element,
the result for that element is NaN rather than an uninitialized garbage value.
In practice this only happens for extremely small radii on very coarse grids; the
filter always includes the queried element itself, so r = 0 is safe.

I don't understand this comment. Even with a very coarse grid and an extremely small radius, I don't see how a NaN could be returned. After all, the filter always includes the queried element itself.

ahijevyc

Thanks for carrying this across the finish line! I just had a comment or two about the assertion that NaN could be returned for a coarse grid and small radius. I don't think this could occur in practice, as the filter should at least return the values at the original grid points even if the radius is zero. And grid.neighbors.BallTree.query_radius already makes sure the radius is not negative.

ahijevyc · 2026-07-29T17:13:12Z

+    neighbor_indices = tree.query_radius(dest_coords, r=r)
+
+    # Initialize with NaN so that any element whose neighborhood is empty
+    # (e.g. an isolated point queried with a very small radius) yields NaN


Initializing with NaN is a good defensive idea. But I don't think the reasoning about an isolated point queried with a very small radius is strictly correct. The filter should return the original values even if the radius is zero.

ahijevyc added 2 commits September 7, 2024 11:09

Neighborhood filter

75d832e

Merge remote-tracking branch 'upstream/main' into neighborhood_filter

b850bc4

ahijevyc added the new feature New feature or request label Sep 9, 2024

ahijevyc self-assigned this Sep 9, 2024

ahijevyc mentioned this pull request Sep 9, 2024

Add neighborhood filter to remap #923

Closed

14 tasks

ahijevyc added 2 commits September 9, 2024 10:48

ruff recommendations

8ec0193

added Callable to Type checking

5605949

philipc2 linked an issue Sep 9, 2024 that may be closed by this pull request

Apply a neighborhood filter with radius r to all elements of UxDataArray #930

Open

philipc2 reviewed Sep 9, 2024

View reviewed changes

Merge branch 'main' into ahijevyc/neighborhood_filter

70ba961

philipc2 previously requested changes Sep 9, 2024

View reviewed changes

Comment thread uxarray/core/dataarray.py Outdated

Comment thread uxarray/core/dataarray.py

Comment thread uxarray/core/dataset.py

ahijevyc and others added 4 commits September 9, 2024 15:10

np.vstack().T faster than np.c

0c7bc1e

Fix some comments

d6d8a33

Merge branch 'main' into ahijevyc/neighborhood_filter

47b9cda

Merge branch 'main' into ahijevyc/neighborhood_filter

f75db0d

ahijevyc commented Oct 2, 2024

View reviewed changes

Merge branch 'main' into ahijevyc/neighborhood_filter

bddd2fa

Kept neighborhood and dual additions

ahijevyc and others added 4 commits March 17, 2025 14:12

Merge branch 'main' into ahijevyc/neighborhood_filter

a0b6361

missing imports

6c59af7

Merge branch 'main' into ahijevyc/neighborhood_filter

67d0c11

Merge branch 'main' into ahijevyc/neighborhood_filter

8979745

philipc2 mentioned this pull request May 14, 2025

Add Azimuthal Averaging #1255

Merged

9 tasks

ahijevyc added 2 commits May 28, 2025 11:47

Merge branch 'main' into ahijevyc/neighborhood_filter

a8875cd

Merge branch 'UXARRAY:main' into ahijevyc/neighborhood_filter

f4af498

[pre-commit.ci] auto fixes from pre-commit.com hooks

14c37b7

for more information, see https://pre-commit.ci

Update dataset.py to address pre-commit errors

a37b88f

rajeeja moved this to 👀 In review in UXarray Development Jul 8, 2026

rajeeja added this to UXarray Development Jul 8, 2026

ahijevyc assigned rajeeja Jul 17, 2026

rajeeja and others added 7 commits July 24, 2026 06:25

Merge remote-tracking branch 'origin/main' into ahijevyc/neighborhood…

47662c1

…_filter # Conflicts: # uxarray/core/dataset.py

Merge branch 'main' into ahijevyc/neighborhood_filter

42c7872

Fix neighborhood_filter dimension assert to not access .dims on ndarray

b47540b

Fix neighborhood_filter docstrings: radius is always in degrees, not …

e50ebbd

…meters

Merge remote-tracking branch 'ahijevyc-fork/ahijevyc/neighborhood_fil…

68fd51f

…ter' into ahijevyc/neighborhood_filter

rajeeja requested a review from erogluorhan July 28, 2026 17:19

rajeeja added 2 commits July 28, 2026 12:59

rajeeja added 2 commits July 28, 2026 13:54

Merge branch 'main' into ahijevyc/neighborhood_filter

29f1873

Fix pre-commit: trailing newlines + ruff import sort/format

b37974d

Merge branch 'UXARRAY:main' into ahijevyc/neighborhood_filter

ed35c01

ahijevyc commented Jul 29, 2026

View reviewed changes

		# face centers, edge centers or nodes.
		tree = self.uxgrid.get_ball_tree(coordinates=data_mapping, reconstruct=True)

Uh oh!

Conversation

ahijevyc commented Sep 9, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Overview

Expected Usage

PR Checklist

Uh oh!

philipc2 Sep 9, 2024

Choose a reason for hiding this comment

Uh oh!

ahijevyc Sep 9, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ahijevyc Sep 9, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ahijevyc Sep 10, 2024

Choose a reason for hiding this comment

Uh oh!

philipc2 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

ahijevyc Oct 2, 2024

Choose a reason for hiding this comment

Uh oh!

philipc2 commented Mar 7, 2025

Uh oh!

aaronzedwick commented Mar 10, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

philipc2 commented Mar 10, 2025

Uh oh!

aaronzedwick commented Mar 10, 2025

Uh oh!

philipc2 commented Mar 10, 2025

Uh oh!

aaronzedwick commented Mar 10, 2025

Uh oh!

erogluorhan commented Feb 26, 2026

Uh oh!

erogluorhan commented Feb 26, 2026

Uh oh!

rajeeja commented Jul 8, 2026

Uh oh!

ahijevyc commented Jul 8, 2026

Uh oh!

erogluorhan commented Jul 17, 2026

Uh oh!

ahijevyc commented Jul 17, 2026

Uh oh!

review-notebook-app Bot commented Jul 28, 2026

Uh oh!

rajeeja commented Jul 28, 2026

Uh oh!

ahijevyc Jul 29, 2026

Choose a reason for hiding this comment

Uh oh!

ahijevyc left a comment

Choose a reason for hiding this comment

Uh oh!

ahijevyc Jul 29, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

ahijevyc commented Sep 9, 2024 •

edited

Loading

ahijevyc Sep 9, 2024 •

edited

Loading

ahijevyc Sep 9, 2024 •

edited

Loading

aaronzedwick commented Mar 10, 2025 •

edited

Loading