ENH: add max dist to NearestNDInterpolator

[harshilkamdar]

Reference issue

This does not close a particular issue as far as I know, but does speed up nearest neighbor interpolation by a siginificant amount in cases where a max distance arg is passed.

What does this implement/fix?

This allows for far quicker queries and nearest neighbor interpolation for large datasets where prior information about when/where to interpolate is known. In some 2D geospatial examples with tens of millions of points, this leads to a speedup of anywhere from 2-30x.

Additional information

n/a. This is my first PR here -- please let me know if I did anything wrong.

[ev-br]

Looks very reasonable!

My main comment is about naming: I think it would be better to

1. keep the name of the parameter consistent with KDTree, and
2. maybe think about exposing other KDTree.query parameters where it makes sense.

The second part here is optional of course.

[ev-br]

One other thing (sorry for the piecemeal review, only noticed this being next to **tree_options): you are adding these new arguments to the constructor and only use them at the call site. Is there a need to store workers and distance_upper_bound on the instance? Maybe better mirror KDTree and only supply them in __call__.

[harshilkamdar]

No worries at all. My rationale was looking to interpolate.griddata in the future and making it more flexible.

Right now, griddata does not support passing any options to the underlying interpolator (nearest, linear, or cubic). I'm hoping to change that in a new PR and figured it would be cleaner to only expose and use class-level options and keeping __call__ clean for now. Linear doesn't have interesting options, but CloughTocher2DInterpolator does.

This is not a super strong reason, so I'm happy to go with whatever you prefer.

[ev-br]

Right. So delineating between __init__ and __call__ optional arguments would require two sets of dicts as inputs to griddata (unpacked or not), and those would be method-dependent. Not nice.
OTOH, you current approach is that everything is passed in a constructor and __call__ signature only has data points and nothing else. Which is a reasonable invariant, so let's roll with it I'd say. WDYT @j-bowhay ?

I'd then ask for a small tweak:

• please add a leading underscore to things you store on the instance (as in self._workers is private, while self.workers is a part of public API).
• [EDITED IN] could you please stress that tree_options are passed to the underlying KDTree constructor.

In fact, it's not immediately clear from the docstring that this whole interpolator is for unstructured data and is KDTree-based. Would be great if you could update that in a follow-up PR.

[j-bowhay]

I appreciate the argument around future work on griddata but it seems cumbersome that I would have to rebuild my interpolator just because I want to change the workers / max dist. Could we not add an options dict argument to griddata (as we do for minimise and friends) then each method uses this where needed (whether that be __init__ or __call__)

[harshilkamdar]

@j-bowhay, @ev-br: thank you both for your careful feedback. I have updated the PR based on both your comments in the latest commit - ff1d52f. It now takes in a query_options dict and exposes some more cKDTree functionality. I've also added a couple more tests.

In fact, it's not immediately clear from the docstring that this whole interpolator is for unstructured data and is KDTree-based. Would be great if you could update that in a follow-up PR.

I will do this in a separate documentation PR.

[j-bowhay]

Sorry I wasn't quite clear, I was proposing the griddata take some kind of option dict to pass to the underlying interpolator and a signature like NearestNDInterpolator.__call__(*args, distance_upper_bound=np.inf, etc) but perhaps lets see what @ev-br thinks so we don't have to make lots of back and forth changes

[ev-br]

To be clear: I'm fine with either approach. I'm sold on that the inconvenience of lumping together constructor and call arguments is rather minor and a cleaner griddata signature outweights it; I can also live with a optimize-like approach. As long as we are not adding public attributes to interpolators, I'm OK and defer to @j-bowhay for the final API specification.

[harshilkamdar]

Sounds good - sorry about the added confusion. @j-bowhay let me know which option you prefer and I can implement. Either is fine on my end and should likely be minimal effort

[j-bowhay]

How about an interface like this?

[harshilkamdar]

Thanks both for your patience. @j-bowhay - I like it. I have committed and updated tests + docs and have re-requested review.

[ev-br]

CI failures look related, could you please take a look

[harshilkamdar]

Sorry, dumb bug. Fixed now.

[ev-br]

CI is still unhappy :-(.
Might be faster to iterate locally: the incantation is $ python dev.py test -s interpolate or -t path/to/test/file.

[harshilkamdar]

@ev-br - good call.

Through the griddata test failures, I've found that both NearestNDInterpolator and griddata are silently actually supporting two things that I wasn't originally accounting for - the case where __call__ gets an n-D xi argument and the case where the y-values for the interpolation are also multidimensional.

I've come up with hopefully an acceptable solution to account for these two cases and updated docs & comments for NearestNDInterpolator to reflect this. I couldn't get scipy to build locally for some annoying reasons but have tried to make sure that all tests in test_ndgriddata.py pass by checking manually. If there end up being failures, I can fix them quickly.

[ev-br]

OK, this is getting close! I like the flatten-compute-restore logic, it makes a lot of sense.

I've left several comments about 1) minor simplifications of implementation details and 2) clarity of comments/documentation.

The story about stacks of arrays is always confusing, so let's try to clarify it here. So, there are three arrays: x, y and xi.

Let's see if I got it right:

• x is always 2D
• y is 1D or a stack of 1D, where the first dimension is the same as the second (or first?) dimension of x, and trailing dimensions represent the stack of 1D arrays.
• xi is normally a 1D array of dimension ndim or a stack of 1D arrays, where the stacking is along the leading dimension (this is what broadcastable word hints at inhttps://docs.scipy.org/doc/scipy/reference/generated/scipy.interpolate.NearestNDInterpolator.call.html#scipy.interpolate.NearestNDInterpolator.call ).

I think it would be best to explain this (or a correct version of it) to clearly separate what is the number of dimensions, what is the number of data points and how stacked dimensions come into play.
Two general considerations to keep in mind:

• the docs for NearestNDInterpolator, LinearNDInterpolator and CloughtTocher2DInterpolator should look relatively consistent. If one of them supports trailing/leading dimensions and the other does not, OK, so be it. But going from nearest to linear interpolator in a simple case should be seamless, and so should be the docs.
• Maybe it'd be useful to take a look at how RegularGridInterpolator documents essentially the same issue. Ideally, regular and scattered interpolants look similar.

Finally, since the documentation story is not strictly speaking related to your original changes, it would be perfectly fine to undo the docstring changes and keep the documentation update for a follow-up PR (which would be very welcome indeed).

[ev-br]

n*m* is confusing: what does the trailing asterisk stand for?

[harshilkamdar]

Yeah, I'm not being careful here in my words. Have given the description another go.

[ev-br]

let's not bring griddata here. The rest of the comment is also confusing (at least to me, and this is not the first time I'm seeing the story about leading/trailing dimensions).

I'm not sure what your m, n, k and l refer to TBH. Maybe take a look at how RegularGridInterpolator documents essentially the same story, possible trailing dimensions in y and xi:

https://docs.scipy.org/doc/scipy/reference/generated/scipy.interpolate.RegularGridInterpolator.__call__.html#scipy.interpolate.RegularGridInterpolator.__call__

I'm not saying RGI does it perfectly (it is not), all I'm saying you maybe can come up with a better formulation for both (likely in a follow-up PR).

[harshilkamdar]

Yeah, this is confusing and griddata is not relevant - updated.

[harshilkamdar]

@ev-br - thanks for the thorough comments. I think you're right that I'm confusing nD/2D and trailing dimensions. The pointer to RegularGridInterpolator was very helpful.

I have considered most of your comments and hopefully addressed bulk of it in the new commits. Hopefully this is a more accurate version of what's going on. Any remaining doc fix, I will leave to a separate PR for now.

[ev-br]

All my comments have been addressed, I think it is a great addition, so +1 from me.

Unless there are further comments, am going to hit the green button later this week.

[ev-br]

Am going to merge now, based on two core dev approvals. Thanks Jake for the review, thank you @harshilkamdar for the enhancement and congrats with what I believe is your first SciPy commit. Keep them coming!