Repository: f0uriest/interpax

Issue:
evaluating splines with scalars returns 1D arrays when extrapolate=True
```python
import interpax
import jax.numpy as jnp

for extrapolate in [None, True, False]:
    spl = interpax.CubicSpline(jnp.linspace(0, 1, 11), jnp.arange(11), check=False, extrapolate=extrapolate)
    print(spl(1/2).shape)
```
prints `()`, `()`, and `(1,)`.

Discussion:
@f0uriest: Thanks for the clear repro! this should be fixed with #88 @zachjweiner: Thanks for the quick fix (and the library)!

Solution patch:
diff --git a/interpax/_ppoly.py b/interpax/_ppoly.py
index 0654053..d50baeb 100644
--- a/interpax/_ppoly.py
+++ b/interpax/_ppoly.py
@@ -242,8 +242,9 @@ class PPoly(eqx.Module):
         y = y.reshape(x_shape + self.c.shape[2:])
 
         if not extrapolate:
-            mask = jnp.logical_or(x > self.x[-1], x < self.x[0])
-            y = jnp.where(mask, jnp.nan, y.T).T
+            # x became 1d after flatten, so reshape this back to original x shape
+            mask = jnp.logical_or(x > self.x[-1], x < self.x[0]).reshape(x_shape)
+            y = jnp.where(mask.T, jnp.nan, y.T).T
 
         if self.axis != 0:
             # transpose to move the calculated values to the interpolation axis
@@ -384,9 +385,7 @@ class PPoly(eqx.Module):
 
             def falsefun():
                 return (
-                    jnp.zeros(self.c.shape[2:])
-                    if self.c.shape[2:]
-                    else jnp.array([0.0])
+                    jnp.zeros(self.c.shape[2:]) if self.c.shape[2:] else jnp.array(0.0)
                 )
 
             out = jax.lax.cond(n_periods > 0, truefun, falsefun)


Test patch:
diff --git a/tests/test_scipy.py b/tests/test_scipy.py
index c1304f1..13dcc4d 100644
--- a/tests/test_scipy.py
+++ b/tests/test_scipy.py
@@ -540,6 +540,7 @@ class TestPPoly:
         # 1 - x**2
         c = np.array([[-1, 0, 1]]).T
         x = np.array([0, 1])
+        x1 = 1 / 2
 
         for extrapolate in [True, False, None]:
             pp = PPoly(c, x, extrapolate=extrapolate)
@@ -555,6 +556,11 @@ class TestPPoly:
                 assert_(not np.isnan(pp_i([-0.1, 1.1])).any())
                 assert_(not np.isnan(pp_d([-0.1, 1.1])).any())
 
+            # extra test for gh#85
+            assert pp(x1).shape == ()
+            assert pp_d(x1).shape == ()
+            assert pp_i(x1).shape == ()
+
 
 def _ppoly_eval_1(c, x, xps):
     """Evaluate piecewise polynomial manually."""