Repository: google/jax

Issue:
scipy.signal.detrend implementation in JAX
HI, 

Would it be possible to implement scipy.signal.detrend in JAX?  I need to compute the gradient of a function which involves de-trending a signal computed from input parameters.

Thanks,
York


Discussion:
Here's a quick implementation (not comprehensively tested yet):
```python
import jax.numpy as jnp
import numpy as np

def detrend(data, axis=-1, type='linear', bp=0):
  if type not in ['constant', 'linear', 'c', 'l']:
    raise ValueError("Trend type must be 'linear' or 'constant'.")
  data = jnp.asarray(data)
  if type in ['constant', 'c']:
    return data - jnp.mean(data, axis, keepdims=True)
  else:
    N = data.shape[axis]
    # bp is static, so we use np operations to avoid pushing to device.
    bp = np.sort(np.unique(np.r_[0, bp, N]))
    if np.any(bp > N):
      raise ValueError("Breakpoints must be less than length of data along given axis.")
    data = jnp.moveaxis(data, axis, 0)
    shape = data.shape
    data = data.reshape(N, -1)
    for m in range(len(bp) - 1):
      Npts = bp[m + 1] - bp[m]
      A = jnp.vstack([jnp.ones(Npts), jnp.arange(1, Npts + 1) / Npts]).T
      sl = slice(bp[m], bp[m + 1])
      coef, resids, rank, s = jnp.linalg.lstsq(A, data[sl])
      data = data.at[sl].add(-jnp.dot(A, coef))
    return jnp.moveaxis(data.reshape(shape), 0, axis)
```
I can work on getting it into ``jax.scipy.signal`` this week.
Thanks for the quick reply.  I tested the function with my data and it's very close to the result from using scipy.signal.detrend. It's a bit slower than the original scipy version, but I suppose wrapping it in jit will speed it up for repeated evaluations.

Resolving patch:
diff --git a/jax/scipy/signal.py b/jax/scipy/signal.py
--- a/jax/scipy/signal.py
+++ b/jax/scipy/signal.py
@@ -15,8 +15,11 @@
 import scipy.signal as osp_signal
 import warnings
 
+import numpy as np
+
 from .. import lax
 from ..numpy import lax_numpy as jnp
+from ..numpy import linalg
 from ..numpy.lax_numpy import _promote_dtypes_inexact
 from ..numpy._util import _wraps
 
@@ -104,3 +107,33 @@ def correlate2d(in1, in2, mode='full', boundary='fill', fillvalue=0,
   if jnp.ndim(in1) != 2 or jnp.ndim(in2) != 2:
     raise ValueError("correlate2d() only supports {ndim}-dimensional inputs.")
   return _convolve_nd(in1[::-1, ::-1], in2, mode, precision=precision)[::-1, ::-1]
+
+
+@_wraps(osp_signal.detrend)
+def detrend(data, axis=-1, type='linear', bp=0, overwrite_data=None):
+  if overwrite_data is not None:
+    raise NotImplementedError("overwrite_data argument not implemented.")
+  if type not in ['constant', 'linear']:
+    raise ValueError("Trend type must be 'linear' or 'constant'.")
+  data, = _promote_dtypes_inexact(jnp.asarray(data))
+  if type == 'constant':
+    return data - data.mean(axis, keepdims=True)
+  else:
+    N = data.shape[axis]
+    # bp is static, so we use np operations to avoid pushing to device.
+    bp = np.sort(np.unique(np.r_[0, bp, N]))
+    if bp[0] < 0 or bp[-1] > N:
+      raise ValueError("Breakpoints must be non-negative and less than length of data along given axis.")
+    data = jnp.moveaxis(data, axis, 0)
+    shape = data.shape
+    data = data.reshape(N, -1)
+    for m in range(len(bp) - 1):
+      Npts = bp[m + 1] - bp[m]
+      A = jnp.vstack([
+        jnp.ones(Npts, dtype=data.dtype),
+        jnp.arange(1, Npts + 1, dtype=data.dtype) / Npts
+      ]).T
+      sl = slice(bp[m], bp[m + 1])
+      coef, *_ = linalg.lstsq(A, data[sl])
+      data = data.at[sl].add(-jnp.matmul(A, coef, precision=lax.Precision.HIGHEST))
+    return jnp.moveaxis(data.reshape(shape), 0, axis)


Test patch:
diff --git a/tests/scipy_signal_test.py b/tests/scipy_signal_test.py
--- a/tests/scipy_signal_test.py
+++ b/tests/scipy_signal_test.py
@@ -44,7 +44,7 @@ class LaxBackedScipySignalTests(jtu.JaxTestCase):
   """Tests for LAX-backed scipy.stats implementations"""
 
   @parameterized.named_parameters(jtu.cases_from_list(
-      {"testcase_name": "_op={}_xshape=[{}]_yshape=[{}]_mode={}".format(
+      {"testcase_name": "_op={}_xshape={}_yshape={}_mode={}".format(
           op,
           jtu.format_shape_dtype_string(xshape, dtype),
           jtu.format_shape_dtype_string(yshape, dtype),
@@ -67,7 +67,7 @@ def testConvolutions(self, xshape, yshape, dtype, mode, jsp_op, osp_op):
     self._CompileAndCheck(jsp_fun, args_maker)
 
   @parameterized.named_parameters(jtu.cases_from_list(
-      {"testcase_name": "op={}_xshape=[{}]_yshape=[{}]_mode={}".format(
+      {"testcase_name": "op={}_xshape={}_yshape={}_mode={}".format(
           op,
           jtu.format_shape_dtype_string(xshape, dtype),
           jtu.format_shape_dtype_string(yshape, dtype),
@@ -89,6 +89,24 @@ def testConvolutions2D(self, xshape, yshape, dtype, mode, jsp_op, osp_op):
     self._CheckAgainstNumpy(osp_fun, jsp_fun, args_maker, check_dtypes=False, tol=tol)
     self._CompileAndCheck(jsp_fun, args_maker)
 
+  @parameterized.named_parameters(jtu.cases_from_list(
+      {"testcase_name": "_shape={}_axis={}_type={}_bp={}".format(
+          jtu.format_shape_dtype_string(shape, dtype), axis, type, bp),
+       "shape": shape, "dtype": dtype, "axis": axis, "type": type, "bp": bp}
+      for shape in [(5,), (4, 5), (3, 4, 5)]
+      for dtype in default_dtypes
+      for axis in [0, -1]
+      for type in ['constant', 'linear']
+      for bp in [0, [0, 2]]))
+  def testDetrend(self, shape, dtype, axis, type, bp):
+    rng = jtu.rand_default(self.rng())
+    args_maker = lambda: [rng(shape, dtype)]
+    osp_fun = partial(osp_signal.detrend, axis=axis, type=type, bp=bp)
+    jsp_fun = partial(jsp_signal.detrend, axis=axis, type=type, bp=bp)
+    tol = {onp.float32: 1e-5, onp.float64: 1e-12}
+    self._CheckAgainstNumpy(osp_fun, jsp_fun, args_maker, tol=tol)
+    self._CompileAndCheck(jsp_fun, args_maker, rtol=tol, atol=tol)
+
 
 if __name__ == "__main__":
     absltest.main()