patrick-kidger
Patched edge-case in `LinearInterpolation`
Fixes an edge case where dfx.LinearInterpolation will return a nan value when input arrays have a repeated leading constant value. ie. the following example will return a nan value in current Diffrax:
import diffrax
import jax.numpy as jnp
ts = jnp.array([1., 1., 1., 2., 3.])
ys = jnp.array([2., 2., 2., 3., 4.])
interpolator = diffrax.LinearInterpolation(ts, ys)
nan_result = interpolator.evaluate(1.)
Note the interpolator will only get nan evaluated on the repeated value.
Right now this is deliberate behaviour in Diffrax. Having repeated times doesn't make a lot of sense in general?
I find either behavior reasonable. To me, multiple consistent values like this seems fine from the interpolation perspective, but maybe not from the diffeq perspective.
This type of data came up for me when using the Diffrax interpolators outside of standard diffeq solves.
Do note though that either way something like
ts = jnp.array([0., 1., 1., 1., 2., 3.])
ys = jnp.array([1., 2., 2., 2., 3., 4.])
will give a non-nan value for interpolator.evaluate(1.).
Ech. FWIW this is currently in tracked in #31. I'd like to (but will realistically never find time to) fix this up more carefully for every interpolation routine, but let's handle things ad-hoc for now.
I'd be happy to merge this, but I think you can make things computationally cheaper: relative to the original code, just replace fractional_part / diff_t with misc.linear_rescale(prev_t, fractional_part, next_t).
Actually, I think the numerics here are slightly wrong. From the docstring for linear_rescale: "Calculates (t - t0) / (t1 - t0)". But we want fractional_part / (t_next - t_prev).