Apply function to points within circular neighborhood

[ahijevyc]

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]

@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

[ahijevyc]

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?

[ahijevyc]

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

[ahijevyc]

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:

[philipc2]

This implementation looks great! May we move the bulk of the logic into the uxarray.grid.neighbors module and call that helper from here?

We can keep the data-mapping checks here, and anything related to constructing and returining the final data array but the bulk of the computations would go inside a helper in the module mentioned above.

[ahijevyc]

I have to think about how to do that, but I am happy to defer to you.

[ahijevyc]

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