Update ndimage tests and add jit to map_coordinates

src/lcm/ndimage.py

@@ -20,9 +20,10 @@
 from collections.abc import Sequence
 
 import jax.numpy as jnp
-from jax import Array, lax, util
+from jax import Array, jit, lax, util
 
 
+@jit
 def map_coordinates(
     input: Array,
     coordinates: Sequence[Array],

tests/test_ndimage.py

@@ -100,13 +100,3 @@ def f(step):
     # Gradient of f(step) is 2 * step
     assert_allclose(jax.grad(f)(0.5), 1.0)
     assert_allclose(jax.grad(f)(1.0), 2.0)
-
-
-def test_extrapolation():
-    x = jnp.arange(3.0)
-    c = [jnp.array([-2.0, -1.0, 5.0, 10.0])]
-
-    got = lcm_map_coordinates(x, c)
-    expected = c[0]
-
-    assert_array_equal(got, expected)

tests/test_ndimage_unit.py

@@ -1,14 +1,33 @@
 import jax.numpy as jnp
+import pytest
 from numpy.testing import assert_array_equal
 
 from lcm.ndimage import (
     _compute_indices_and_weights,
     _multiply_all,
     _round_half_away_from_zero,
     _sum_all,
+    map_coordinates,
 )
 
 
+def test_map_coordinates_wrong_input_dimensions():
+    values = jnp.arange(2)  # ndim = 1
+    coordinates = [jnp.array([0]), jnp.array([1])]  # len = 2
+    with pytest.raises(ValueError, match="coordinates must be a sequence of length"):
+        map_coordinates(values, coordinates)
+
+
+def test_map_coordinates_extrapolation():
+    x = jnp.arange(3.0)
+    c = [jnp.array([-2.0, -1.0, 5.0, 10.0])]
+
+    got = map_coordinates(x, c)
+    expected = c[0]
+
+    assert_array_equal(got, expected)
+
+
 def test_nonempty_sum():
     a = jnp.arange(3)