Use jnp.interp

jax_cosmo/scipy/interpolate.py

@@ -11,34 +11,7 @@
 
 __all__ = ["interp"]
 
-
-@functools.partial(vmap, in_axes=(0, None, None))
-def interp(x, xp, fp):
-    """
-    Simple equivalent of np.interp that compute a linear interpolation.
-
-    We are not doing any checks, so make sure your query points are lying
-    inside the array.
-
-    TODO: Implement proper interpolation!
-
-    x, xp, fp need to be 1d arrays
-    """
-    # First we find the nearest neighbour
-    ind = np.argmin((x - xp) ** 2)
-
-    # Perform linear interpolation
-    ind = np.clip(ind, 1, len(xp) - 2)
-
-    xi = xp[ind]
-    # Figure out if we are on the right or the left of nearest
-    s = np.sign(np.clip(x, xp[1], xp[-2]) - xi).astype(np.int64)
-    a = (fp[ind + np.copysign(1, s).astype(np.int64)] - fp[ind]) / (
-        xp[ind + np.copysign(1, s).astype(np.int64)] - xp[ind]
-    )
-    b = fp[ind] - a * xp[ind]
-    return a * x + b
-
+interp = np.intern
 
 @register_pytree_node_class
 class InterpolatedUnivariateSpline(object):

tests/test_spline.py

@@ -1,6 +1,6 @@
 # This module tests the InterpolatedUnivariateSpline implementation against
 # SciPy
-from jax.config import config
+from jax import config
 
 config.update("jax_enable_x64", True)
 import jax.numpy as np