Add Azimuthal Averaging

[jpan84]

Closes #1254

Compute azimuthal mean in radial rings around a central point.

Overview

Like uxarray.UxDataArray.zonal_mean(), this aggregates over faces to output the data with a new horizontal dimension radius. Following uxarray.UxDataArray.subset.bounding_circle, this function uses ball_tree.query_radius to select grid elements in a circular neighborhood from one center coordinate. For radii r=[r1, r2, r3, ...], face centers at distance d (given in great-circle degrees), where r2<d<=r3 will be used for the azimuthal mean at r3.

This is one solution for issue #1254

This may be tangentially relevant to issue #930.

Expected Usage

import uxarray as ux
import matplotlib.pyplot as plt

grid_path = '/glade/p/cesmdata/inputdata/share/scripgrids/ne120np4_pentagons_100310.nc'
data_path = '/glade/work/jpan/azim_test_data.nc'

uxds = ux.open_dataset(grid_path, data_path)
uxda = uxds['OMEGA500']
print(uxda)

azim_mean, hit_count = uxda.azimuthal_mean((183., -18.933), 5, 0.3)

print(azim_mean)
print(hit_count)
plt.plot(azim_mean.radius, azim_mean.isel(time=0))
plt.show()

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
• [N/A ] Docstrings have updated with any function changes
• [N/A] 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/

[jpan84]

sorry, this is my 1st time creating a new feature. I'm not sure how to properly link this to the issue or attach a "new feature" label.

[philipc2]

This seems similar to #941

[jpan84]

perhaps the case where a bin contains 0 hits (i.e., hit_count[ii] = 0) should be set to nan in the output data array. so before the continue statement:
means[ii, ...] = np.nan

[jpan84]

why are the outer radius and radius step arguments required to be integers rather than any non-negative scalar?

[philipc2]

why are the outer radius and radius step arguments required to be integers rather than any non-negative scalar?

Good catch, I've updated the type hints.

[Copilot]

Pull Request Overview

This PR adds an azimuthal_mean function to the UxDataArray class that computes averages along circles of constant great-circle distance from a central point. The implementation follows a similar pattern to the existing zonal_mean function and uses ball tree spatial queries to efficiently identify faces within radial bins.

• Adds azimuthal_mean method to compute radial averages around a specified center coordinate
• Implements spatial binning using KD-tree ball queries with radial ring exclusion logic
• Returns a UxDataset containing both the azimuthal means and hit counts for each radius

Reviewed Changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 4 comments.

File	Description
uxarray/core/dataarray.py	Adds the azimuthal_mean method implementation with radial binning logic
test/core/test_azimuthal.py	Adds basic tests for Gaussian and inverse Gaussian distributions
docs/api.rst	Updates API documentation to include the new azimuthal mean functionality

---

_{Tip: Customize your code reviews with copilot-instructions.md. Create the file or learn how to get started.}

[erogluorhan]

This seems similar to #941

Not really, based on my understanding of each. Geometrically, #941 suggests that a bounding circular neighborhood is obtained for each grid face (where the center coordinate varies per face center but the radius is constant) while this one looks into bounding circles centered on a fixed center coordinate but differs from each other with varying radii.

Do you agree with this?

[philipc2]

This seems similar to #941

Not really, based on my understanding of each. Geometrically, #941 suggests that a bounding circular neighborhood is obtained for each grid face (where the center coordinate varies per face center but the radius is constant) while this one looks into bounding circles centered on a fixed center coordinate but differs from each other with varying radii.

Do you agree with this?

I do agree. #941 applies an aggregation within all points that lie a given radius from a center coordinate.

This one incrementally increases the radius and applies the aggregation within those circles.

[erogluorhan]

The functionality seems great!

I think I am hesitant though about returning UxDataset out of such an aggregation function (it likewise applies to zonal_mean() as well; I think I overlooked it before): Shouldn't the logic be to create x number of scalars that are corresponding to the mean of each circular region, and not to worry about providing an USG topology? If so, shouldn't return types simply be np.ndarray, or something similar?

cc: @rajeeja

[philipc2]

The functionality seems great!

I think I am hesitant though about returning UxDataset out of such an aggregation function (it likewise applies to zonal_mean() as well; I think I overlooked it before): Shouldn't the logic be to create x number of scalars that are corresponding to the mean of each circular region, and not to worry about providing an USG topology? If so, shouldn't return types simply be np.ndarray, or something similar?

cc: @rajeeja

Good suggestion. Since there's no longer a valid grid associate with the output, the result should be an xarray.DataArray/Dataset (like with our cross-sections).

[erogluorhan]

Good suggestion. Since there's no longer a valid grid associate with the output, the result should be an xarray.DataArray/Dataset (like with our cross-sections).

Absolutely! Probably multiple xarray.DataArray's as needed rather than a xarray.Dataset, but not fully against the latter.

[philipc2]

The functionality seems great!
I think I am hesitant though about returning UxDataset out of such an aggregation function (it likewise applies to zonal_mean() as well; I think I overlooked it before): Shouldn't the logic be to create x number of scalars that are corresponding to the mean of each circular region, and not to worry about providing an USG topology? If so, shouldn't return types simply be np.ndarray, or something similar?
cc: @rajeeja

Good suggestion. Since there's no longer a valid grid associate with the output, the result should be an xarray.DataArray/Dataset (like with our cross-sections).

Absolutely! Probably multiple xarray.DataArray's as needed rather than a xarray.Dataset, but not fully against the latter.

Yep! Added a return_hit_counts parameter to return the hits as an optional DataArray as opposed to being stored in the Dataset

[erogluorhan]

Code looks great to me; thanks very much! No plans with an immediate notebook into the user guide with this PR? If the preference is to add that later, could you please create an issue regarding that, if none exists. Also, who is going to write that guide, @jpan84 maybe you?

[philipc2]

No plans with an immediate notebook into the user guide with this PR? If the preference is to add that later, could you please create an issue regarding that, if none exists. Also, who is going to write that guide, @jpan84 maybe you?

Any thoughts @jpan84? I'm happy to move forward with this PR with the simple docstring example, and create an issue to develop a notebook down the line.

[erogluorhan]

I am approving it to help make its way forward. Thanks a lot for this addition!

[rajeeja]

Good PR. Thanks this work and I agree that we should not return Uxdataset here as already handled. Couple of minor points:
Handle case where no faces are found at all np.all(hit_count == 0) also some parameter validations outer_radius <= 0 or radius_step <= 0, len(center_coord) != 2 , something for lon/lat deg..

[jpan84]

Thank you for moving forward with this PR—the tweaks to the docstring and functionality look great! Please do create an issue for the notebook.

[rajeeja]

Minor safeguards as pointed in comments

[Copilot]

Pull Request Overview

Copilot reviewed 3 out of 3 changed files in this pull request and generated 6 comments.

---

_{Tip: Customize your code reviews with copilot-instructions.md. Create the file or learn how to get started.}