Better performance of the beam interpolation with faster domain checking and spline reuse

[tyler-a-cox]

Description

This PR improves the performance of UVBeam azimuth/zenith spline interpolation without changing the public API.

The changes:

• Replace the domain check’s repeated grid scans with nearest-neighbor lookup using np.searchsorted.
• For standard interpolating splines (s=0 with kx/ky options), construct the tensor-product spline coefficients once and evaluate all beam data slices together using NdBSpline.
• Preserve the existing RectBivariateSpline implementation as a fallback when spline reuse is requested, smoothing or other custom options are used, or NdBSpline is unavailable.
• Clip points outside the spline grid to preserve RectBivariateSpline’s existing boundary-value behavior.
• Preserve the existing spline-cache behavior when reuse_spline=True. The new fast path does not create a persistent cache.

Motivation and Context

Interpolating a UVBeam can be slow when evaluating large numbers of source positions. The previous implementation repeatedly scanned the coordinate grid during domain checking and constructed and evaluated separate bivariate splines for each beam data slice.

The new implementation addresses these costs independently:

1. The optimized domain checker finds the neighboring grid values with np.searchsorted, avoiding repeated scans of either the interpolation points or coordinate grid.
2. The optimized interpolation path shares tensor-product spline calculations across beam data slices and evaluates them together.

A representative local benchmark used 10 million interpolation points, a 91 × 360 beam grid, one frequency, and a complex E-field beam with two vector axes and two feeds:

Domain checker	Spline implementation	Runtime	End-to-end speedup
Previous grid scans	RectBivariateSpline	81.6 s	1.0×
Optimized with np.searchsorted	RectBivariateSpline	33.2 s	2.5×
Optimized with np.searchsorted	Batched NdBSpline	9.9 s	8.2×

The intermediate measurement changes only the domain checker, reducing the end-to-end runtime from 81.6 to 33.2 seconds. Changing only the spline implementation after that reduces the runtime by a further factor of approximately 3.4, from 33.2 to 9.9 seconds.

A warm run using the existing reuse_spline=True cache and the optimized domain checker took approximately 33.2 seconds in the same benchmark. This indicates that the additional improvement from the batched path comes primarily from faster spline evaluation rather than caching.

Testing

The added tests:

• Compare the optimized spline results with the preserved RectBivariateSpline implementation for both real power beams and complex E-field beams.
• Cover the default cubic interpolation order and custom kx/ky orders.
• Check agreement to an absolute tolerance of 1e-12, including points near the grid boundaries.
• Directly compare the optimized domain checker against the legacy implementation.
• Exercise both size-dependent paths from the legacy domain checker.
• Check covered points, azimuth and zenith failures, the exact coverage threshold, adjacent floating-point values, and matching exception messages.

The new interpolation and domain-check test set passes with 34 tests.

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
• Other — performance optimization

Checklist

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

Other:

• No public docstring changes are required because the API and documented behavior are unchanged.
• No README or tutorial changes are required because this is an internal performance optimization.
• I have added/updated tests to cover my changes.
• I have updated the CHANGELOG.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 100.00%. Comparing base (6f593ac) to head (15ca6d0).

Additional details and impacted files

@@            Coverage Diff            @@
##              main     #1716   +/-   ##
=========================================
  Coverage   100.00%   100.00%           
=========================================
  Files           71        71           
  Lines        23898     23912   +14     
=========================================
+ Hits         23898     23912   +14     

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

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[bhazelton]

This looks like a great improvement, thank you! It looks like NdBSpline became part of the public API in scipy 1.12? If that's right, I'd be comfortable bumping our minimum required version to that because it's been out for >2.5 years now (barring any unforeseen issues with package availability). It seems like that would significantly simplify the code.

We could try it out on a branch off of this one to make sure there's no issues in CI (the one thing I'm a little worried about is intel mac builds, some packages are not available for intel macs...)

[bhazelton]

@tyler-a-cox I made a branch (scipy1.12) where I increased the scipy requirement to 1.12 and it seems to run just fine on CI. We also talked about this on our telecon this morning and we agreed that moving to scipy 1.12 is fine. So I'd like to go that route which should help with the code complexity on this PR. Happy to chat on slack or zoom as well if you'd like.

[tyler-a-cox]

@bhazelton thanks for taking a look at this and for checking the versioning! I'll go ahead and bump the version requirement.

That said, dropping the RectBivariateSpline path entirely isn't quite equivalent to the NdBSpline path, since NdBSpline doesn't support all the same options, with smoothing being main difference. I looked into whether it could be folded in and I don't think it can. With s=0 the knots only depend on the coordinate grids, so every data slice shares the same ones and we can stack them all into a single NdBSpline. With s>0 FITPACK chooses them based on the data instead, so each slice gets its own set and can't share a spline with the others. That puts us back to one spline per slice, which is what the speedup was avoiding in the first place.

So I think it's either keep that bit of the RectBivariateSpline code around for the smoothing case, or drop the smoothing option. I don't know if anyone actually uses it (nothing I saw in the tests does), but I didn't want to remove it here without a discussion first. Curious which way you'd rather go!

[bhazelton]

This selection criteria is pretty hidden from the user. They might prefer to use this new approach but not realize that setting reuse_spline=True (which sounds like it will save time from the current docstring) will prevent it.

We can fix this in two ways:

• one is to add a bunch of info to the docs to help people pick which parameters and spline opts to use.
• another option would be to make the NdBSpline approach have an entirely separate function that users can choose to use. This is what we did with the map_coordinate option. This would still require some significant docs updates but it might be more transparent to a user. Given that this generally gives the same results, I'm fine with changing the default to be this if the spline_opts are compatible.

I'm not at all sure which of these is better, I don't want to duplicate a lot of code and I also want to make this both easy and clear for users with good defaults. I think there's not great guidance in the existing docs about the choice between az_za_simple and az_za_map_coordinates.

I am not aware that anyone uses any spline opts beyond kx and ky, but I agree that keeping around the full flexibility is nice.

[tyler-a-cox]

Thanks for taking a look! I agree, that behavior made things unclear for the user. I've made some changes based on your feedback.

I decided to keep the NdBSpline and RectBivariateSpline paths in the same function, and to make NdBSpline the default whenever only kx and ky are passed. If anything else is passed, it falls back to RectBivariateSpline. I also added spline caching to the NdBSpline path, so reuse_spline=True now works for both and no longer decides which one you get.

The one decision I wanted your input on is bumping the minimum scipy to 1.15. Even spline orders above 2 aren't supported before that, so kx=4 raises an error on 1.12. The other option is to stay at 1.12 and fall back to RectBivariateSpline for the orders that aren't supported, though that means carrying a special case for order 4. Curious to hear your thoughts.

[bhazelton]

I like the approach you took and I love the new docs, thank you!!

I'm ok with going to scipy 1.15 since all the CI works with it and it's > 1 year old. Much better than maintaining complicated work-arounds.

[bhazelton]

We just put out a new version to get a bug fix out to our users, so this branch needs to be rebased on main and all of this needs to end up in the unreleased section.

Also, please note the scipy dependency change, you can see other examples of that lower in this changelog.

[tyler-a-cox]

@bhazelton Okay, I just changed the CHANGELOG so I think this should be good to go! One of the tests is failing, but it passed before I made this last commit so I think it might just need to be re-run.

[bhazelton]

Ack! I'm so sorry! I merged another PR which made another changelog conflict. 🤦‍♀️ I should have merged yours first. Can you please rebase again?

[tyler-a-cox]

Ahhh gotcha, no problem! Just rebased and it looks like tests are passing

[bhazelton]

Thank you @tyler-a-cox ! This is fantastic.