feat: Add sparse.interp function for numba backend (#903)

* Add sparse.interp function for numba backend
* Add suggestions from review
* Add shortcut for dense data to spare.interp
* Move test to test_coo.py
* Adjust test to compare sparse.interp to numpy.interp
* Add test case for complex data
* Use random test data
* Fix failing pipeline tests
* Fix doctest

Author: peanutfun

sparse/numba_backend/__init__.py

@@ -105,6 +105,7 @@
     full,
     full_like,
     imag,
+    interp,
     isinf,
     isnan,
     matmul,
@@ -252,6 +253,7 @@
     "int32",
     "int64",
     "int8",
+    "interp",
     "isfinite",
     "isinf",
     "isnan",

sparse/numba_backend/_common.py

@@ -3250,3 +3250,88 @@ def diff(x, axis=-1, n=1, prepend=None, append=None):
     for _ in range(n):
         result = result[(slice(None),) * axis + (slice(1, None),)] - result[(slice(None),) * axis + (slice(None, -1),)]
     return result
+
+
+def interp(x, xp, fp, left=None, right=None, period=None):
+    """
+    An implementation of ``numpy.interp`` for sparse arrays.
+
+    Thanks to the function dispatch of numpy, this enables interpolation on sparse arrays
+    using the numpy universal function. This function effectively wraps ``np.interp`` by
+    calling it on the array data and the fill value. See the numpy documentation for
+    details on the parameters.
+
+    Parameters
+    ----------
+    x : SparseArray
+        The x-coordinates at which to evaluate the interpolated values.
+    xp : 1-D sequence or SparseArray
+        The x-coordinates of the data points.
+    fp : 1-D sequence or SparseArray
+        The y-coordinates of the data points, same length as ``xp``.
+    left : float or complex, optional
+        Value to return for ``x < xp[0]``, default is ``fp[0]``.
+    right : float or complex, optional
+        Value to return for ``x > xp[-1]``, default is ``fp[-1]``.
+    period : None or float, optional
+        A period for the x-coordinates.
+
+    Returns
+    -------
+    out : SparseArray
+        The interpolated values, same shape as x.
+
+    See Also
+    --------
+    https://numpy.org/doc/stable/reference/generated/numpy.interp.html
+
+    Examples
+    --------
+    When interpolating a sparse array, its data and the fill value are interpolated. The
+    returned array is pruned. Therefore, the fill value and the number of nonzero
+    elements might change.
+
+    >>> import numpy as np
+    >>> xp = [1, 2, 3]
+    >>> fp = [3, 2, 0]
+    >>> y = np.interp(sparse.COO.from_numpy(np.array([0, 1, 1.5, 2.72, 3.14])), xp, fp)
+    >>> y.todense()
+    array([3.  , 3.  , 2.5 , 0.56, 0.  ])
+    >>> y.fill_value
+    np.float64(3.0)
+    >>> y.nnz
+    3
+    """
+    from ._compressed import GCXS
+    from ._coo import COO
+    from ._dok import DOK
+
+    # Densify sparse interpolants
+    if isinstance(xp, SparseArray):
+        xp = xp.todense()
+    if isinstance(fp, SparseArray):
+        fp = fp.todense()
+
+    def interp_func(xx):
+        return np.interp(xx, xp, fp, left=left, right=right, period=period)
+
+    # Shortcut for dense arrays
+    if not isinstance(x, SparseArray):
+        return interp_func(x)
+
+    # Define output type
+    out_kwargs = {}
+    out_type = COO
+    if isinstance(x, GCXS):
+        out_type = GCXS
+        out_kwargs["compressed_axes"] = x.compressed_axes
+    elif isinstance(x, DOK):
+        out_type = DOK
+
+    # Perform interpolation on sparse object
+    arr = as_coo(x)
+    data = interp_func(arr.data)
+    fill_value = interp_func(arr.fill_value)
+    return COO(data=data, coords=arr.coords, shape=arr.shape, fill_value=fill_value, prune=True).asformat(
+        out_type, **out_kwargs
+    )

sparse/numba_backend/tests/test_coo.py

@@ -2070,3 +2070,44 @@ def test_diff_invalid_type():
     a = np.arange(6).reshape(2, 3)
     with pytest.raises(TypeError, match="must be a SparseArray"):
         sparse.diff(a)
+
+
+class TestInterp:
+    xp = [-1, 0, 1]
+    fp = [3, 2, 0]
+
+    @pytest.fixture(params=["coo", "dok", "gcxs", "dense"])
+    def x(self, request):
+        arr = sparse.random((10, 10, 10), fill_value=0)
+        if request.param == "dense":
+            return arr.todense()
+        return arr.asformat(request.param)
+
+    @pytest.mark.parametrize(
+        "xp",
+        [
+            xp,
+            sparse.COO.from_numpy(np.array(xp)),
+        ],
+    )
+    @pytest.mark.parametrize(
+        "fp",
+        [
+            fp,
+            sparse.COO.from_numpy(np.array(fp)),
+            sparse.COO.from_numpy(np.array(fp) + 1j),
+        ],
+    )
+    def test_interp(self, x, xp, fp):
+        def to_dense(arr):
+            if isinstance(x, sparse.SparseArray):
+                return arr.todense()
+            return arr
+
+        actual = sparse.interp(x, xp, fp)
+        expected = np.interp(to_dense(x), xp, fp)
+
+        assert isinstance(actual, type(x))
+        if isinstance(x, sparse.SparseArray):
+            assert actual.fill_value == fp[1]
+        np.testing.assert_array_equal(to_dense(actual), expected)

sparse/numba_backend/tests/test_namespace.py

@@ -83,6 +83,7 @@ def test_namespace():
         "int32",
         "int64",
         "int8",
+        "interp",
         "isfinite",
         "isinf",
         "isnan",