Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

Hi - sorry we missed this! Can you please add some tests for the new functionality? Existing tests for this function can be found here: https://github.com/google/jax/blob/b4ec72deaefa1094891829cd0b78efb5cc19c28b/tests/scipy_ndimage_test.py#L79

Hey Jake thanks for picking this up!

So I've locally updated the tests to cover the third order interpolation however I've implemented a version with a different boundary condition, scipy uses the 'not-a-knot' condition as far as I can tell (they are not explicit about the particular algorithm they use) whereas I've implemented 'natural' boundary conditions. This gives machine precision matching results at the inner parts of the array, but differing results at the boundaries (shown below).

The existing tests use arrays of shapes (5,), (3, 4), (3, 4, 5) meaning that the tests comparing to scipy fail and can't just compare the central portion as the arrays are too small. I'm wondering what the way forward is here. I haven't been able to solve this problem for the 'not-a-knot' boundary conditions, which may not even be the correct one and I think will require an entirely different implementation 😬. I'm thinking as a solution the tests for third order could use larger arrays and only compare the inner regions (which isn't exactly ideal, but works) and specifying in the docs about using different boundary conditions to scipy.

Let me know your thoughts!

Did you see the discussion in #5687?

The existing linear interpolation scheme is somewhat inconsistent with scipy but we would like to fix this. In particular see the proposed solution in #5687 (comment)

Is your PR consistent with one of SciPy (new) supported boundary modes?
https://docs.scipy.org/doc/scipy/tutorial/ndimage.html#interpolation-boundary-handling

Yes I'm aware of the current issues with the boundary mode and have been following these, my PR does not address anything to do with this. The issue you linked is to do with the boundary mode, how the output coordinates are mapped outside of the range of the input coordinates, which is different to the boundary conditions I'm referring to in my above comments.

The boundary conditions are a set of extra conditions at the boundary of the arrays that are needed to solve for the spline coefficients, of which there are a few although scipy doesn't provide options for this. The plot in the above comment is defined entirely within the input coordinates so there are no effects from the boundary mode.

There should be no issues with any future changes made to the boundary modes since this just changes how the output coordinates are mapped, which are then subsequently fed into the interpolation algorithm.

I just wanted to mention that it would be great to have this merged, even if it does not match the scipy boundary conditions. Figuring out what those boundary conditions even are and then implementing them would be challenging. And having some form of cubic interpolation for map_ccordinates is better than having none :)

I am happy to help out with making this happen if you need it. Let me know what needs to be done!

@JamesAllingham If you need something now you can build my jax fork from source until we get this though https://github.com/LouisDesdoigts/jax. Also feel free to send me your email and I can add you to a thread I have trying to sort this out! Any help is greatly appreciated!

Just checking on this - this is a dependency for some science we are currently doing.

I really appreciate you putting this together, but I don't think it's a good idea for JAX's cubic spline code to diverge further from SciPy. The JAX team does not have the expertise on the numerics of splines to get this code right.

I would suggest either:

1. Work on getting these features into SciPy. Once they have been merged over there, we would be happy to consider your PR to mirror the features in JAX.
2. Release a separate library with JAX splines (JAX is open source, after all). We would be happy to include a link in the relevant section of the JAX docs.

Thanks for this – I think given the discussion in #18137, this functionality is probably out-of-scope for jax core. Sorry we weren't able to accept your contribution here!