Faster interpolation in UVBeam.interp

[tyler-a-cox]

Adds new option to UVBeam.interp that improves the interpolation speed of UVBeam._interp_az_za_rect_spline

Description

This PR adds additional functionality to UVBeam.interp that allows users to choose between two interpolators (scipy.interpolate.RectBivariateSpline and scipy.ndimage.map_coordinates) for interpolation along the spatial axis when running UVBeam._interp_az_za_rect_spline. Users can select between the two interpolators by using the newly added spatial_interp_func keyword argument in UVBeam.interp, where the default is RectBivariateSpline which is currently the interpolator used in UVBeam._interp_az_za_rect_spline.

Motivation and Context

When running visibility simulations with hera_sim, we noticed that beam interpolation was one of the most expensive steps. The current interpolator scipy.interpolate.RectBivariateSpline is quite flexible, but can be quite a bit slower than other gridded interpolators in scipy. I ran this branch on HERA's Vivaldi beam file using spatial_interp_func="map_coordinates" and found a ~3x improvement in the interpolation time compared to when spatial_interp_func="RectBivariateSpline".

Types of changes

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to change)
• Documentation change (documentation changes only)
• Version change
• Build or continuous integration change

Checklist:

• I have read the contribution guide.
• My code follows the code style of this project.

New feature checklist:

• I have added or updated the docstrings associated with my feature using the numpy docstring format.
• I have updated the tutorial to highlight my new feature (if appropriate).
• I have added tests to cover my new feature.
• All new and existing tests pass.
• I have updated the CHANGELOG.

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 99.92%. Comparing base (6b22c6a) to head (3720f46).
Report is 21 commits behind head on main.

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #1464   +/-   ##
=======================================
  Coverage   99.92%   99.92%           
=======================================
  Files          61       61           
  Lines       21319    21337   +18     
=======================================
+ Hits        21304    21322   +18     
  Misses         15       15           

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[steven-murray]

Thanks @tyler-a-cox! One thought: is RegularGridInterpolator exactly the same as RectBivariateSpline (when then grid happens to be regular)? If so, we could either just replace the original RectBivariateSpline with the new regular interpolator, or, if we suspect that we might get irregular grids from time to time, at least make the default dynamic, i.e. choose the regular grid if the grid is regular, otherwise choose the bivariate spline.

[tyler-a-cox]

Thanks for the reply, @steven-murray! As we discussed offline, I've swapped the RectangularGridInterpolator with map_coordinates which produces identical results to RectBivariateSpline for linear interpolation. For higher order interpolation, the results are largely identical, but there are some slightly differences at the edges of the interpolation range that I think are the result of differences in smoothing the interpolators do when using higher order interpolants. I would be interested in building in the infrastructure to use a different interpolator depending on if the interpolation grid is regular or not, but might be slightly concerned about doing so if the results aren't identical. I'm also slightly concerned about the fact that the beam.interp allows users to provide a dictionary of parameters to pass to the interpolation function. The arguments between these functions are slightly different, and it's unclear to me how we would handle that. I'd be interested to hear your thoughts!

[steven-murray]

Thanks @tyler-a-cox. I think this is probably the way to go. We don't want to do anything API-breaking here, so it does have to be a keyword argument rather than choosing a default dynamically, because as you say the spline_opts are different between the two.

[steven-murray]

Thanks @tyler-a-cox this all looks good to me. The only thing to still discuss I think is whether we need the check for regularity, given that as @bhazelton mentioned the data should always be regularly spaced.

[bhazelton]

This is looking really good, thanks @tyler-a-cox! I just had a couple comments on docstrings. It looks like there are a couple of lines of missed coverage, but they're similar to lines being covered for the other interp method, so maybe a judiciously placed pytest parametrize call would take care of them pretty easily?

[bhazelton]

We discussed maybe changing the name of this option (with deprecation!) to include something related to RectBivariateSpline. I think maybe it's best to push that to another PR, just wanted to mention it here.

[bhazelton]

moved this to an issue: #1473

[bhazelton]

Looks good to me now.

[bhazelton]

moved this to an issue: #1473