MAINT: signal: convert lp2{lp,hp, bp, bs} to be array api compatible

Author: ev-br

scipy/signal/_filter_design.py

@@ -5,7 +5,7 @@
 
 import numpy as np
 from numpy import (atleast_1d, poly, polyval, roots, real, asarray,
-                   resize, pi, absolute, sqrt, tan, log10,
+                   pi, absolute, sqrt, tan, log10,
                    arcsinh, sin, exp, cosh, arccosh, ceil, conjugate,
                    zeros, sinh, append, concatenate, prod, ones, full, array,
                    mintypecode)
@@ -17,6 +17,9 @@
 from scipy._lib._util import float_factorial
 from scipy.signal._arraytools import _validate_fs
 
+import scipy._lib.array_api_extra as xpx
+from scipy._lib._array_api import array_namespace, xp_promote, xp_size
+
 
 __all__ = ['findfreqs', 'freqs', 'freqz', 'tf2zpk', 'zpk2tf', 'normalize',
            'lp2lp', 'lp2hp', 'lp2bp', 'lp2bs', 'bilinear', 'iirdesign',
@@ -1676,7 +1679,7 @@ def idx_worst(p):
     return sos
 
 
-def _align_nums(nums):
+def _align_nums(nums, xp):
     """Aligns the shapes of multiple numerators.
 
     Given an array of numerator coefficient arrays [[a_1, a_2,...,
@@ -1701,19 +1704,19 @@ def _align_nums(nums):
         # The statement can throw a ValueError if one
         # of the numerators is a single digit and another
         # is array-like e.g. if nums = [5, [1, 2, 3]]
-        nums = asarray(nums)
+        nums = xp.asarray(nums)
 
-        if not np.issubdtype(nums.dtype, np.number):
+        if not xp.isdtype(nums.dtype, "numeric"):
             raise ValueError("dtype of numerator is non-numeric")
 
         return nums
 
     except ValueError:
-        nums = [np.atleast_1d(num) for num in nums]
-        max_width = max(num.size for num in nums)
+        nums = [xpx.atleast_nd(xp.asarray(num), ndim=1) for num in nums]
+        max_width = max(xp_size(num) for num in nums)
 
         # pre-allocate
-        aligned_nums = np.zeros((len(nums), max_width))
+        aligned_nums = xp.zeros((nums.shape[0], max_width))
 
         # Create numerators with padded zeros
         for index, num in enumerate(nums):
@@ -1722,6 +1725,26 @@ def _align_nums(nums):
         return aligned_nums
 
 
+def _trim_zeros(filt, trim='fb'):
+    # https://github.com/numpy/numpy/blob/v2.1.0/numpy/lib/_function_base_impl.py#L1874-L1925
+    first = 0
+    trim = trim.upper()
+    if 'F' in trim:
+        for i in filt:
+            if i != 0.:
+                break
+            else:
+                first = first + 1
+    last = filt.shape[0]
+    if 'B' in trim:
+        for i in filt[::-1]:
+            if i != 0.:
+                break
+            else:
+                last = last - 1
+    return filt[first:last]
+
+
 def normalize(b, a):
     """Normalize numerator/denominator of a continuous-time transfer function.
 
@@ -1778,30 +1801,33 @@ def normalize(b, a):
     Badly conditioned filter coefficients (numerator): the results may be meaningless
 
     """
-    num, den = b, a
+    xp = array_namespace(b, a)
 
-    den = np.asarray(den)
-    den = np.atleast_1d(den)
-    num = np.atleast_2d(_align_nums(num))
+    den = xp.asarray(a)
+    den = xpx.atleast_nd(den, ndim=1, xp=xp)
+
+    num = xp.asarray(b)
+    num = xpx.atleast_nd(_align_nums(num, xp), ndim=2, xp=xp)
 
     if den.ndim != 1:
         raise ValueError("Denominator polynomial must be rank-1 array.")
     if num.ndim > 2:
         raise ValueError("Numerator polynomial must be rank-1 or"
                          " rank-2 array.")
-    if np.all(den == 0):
+    if xp.all(den == 0):
         raise ValueError("Denominator must have at least on nonzero element.")
 
     # Trim leading zeros in denominator, leave at least one.
-    den = np.trim_zeros(den, 'f')
+    den = _trim_zeros(den, 'f')
 
     # Normalize transfer function
     num, den = num / den[0], den / den[0]
 
     # Count numerator columns that are all zero
     leading_zeros = 0
-    for col in num.T:
-        if np.allclose(col, 0, atol=1e-14):
+    for j in range(num.shape[-1]):
+        col = num[:, j]
+        if xp.all(xp.abs(col) <= 1e-14):
             leading_zeros += 1
         else:
             break
@@ -1879,22 +1905,49 @@ def lp2lp(b, a, wo=1.0):
     >>> plt.legend()
 
     """
-    a, b = map(atleast_1d, (a, b))
+    xp = array_namespace(a, b)
+    a, b = map(xp.asarray, (a, b))
+    a, b = xp_promote(a, b, force_floating=True, xp=xp)
+    a = xpx.atleast_nd(a, ndim=1, xp=xp)
+    b = xpx.atleast_nd(b, ndim=1, xp=xp)
+
     try:
         wo = float(wo)
     except TypeError:
         wo = float(wo[0])
-    d = len(a)
-    n = len(b)
+    d = a.shape[0]
+    n = b.shape[0]
     M = max((d, n))
-    pwo = pow(wo, np.arange(M - 1, -1, -1))
+    pwo = wo ** xp.arange(M - 1, -1, -1, dtype=xp.float64)
     start1 = max((n - d, 0))
     start2 = max((d - n, 0))
     b = b * pwo[start1] / pwo[start2:]
     a = a * pwo[start1] / pwo[start1:]
     return normalize(b, a)
 
 
+def _resize(a, new_shape, xp):
+    # https://github.com/numpy/numpy/blob/v2.2.4/numpy/_core/fromnumeric.py#L1535
+    a = xp.reshape(a, (-1,))
+
+    new_size = 1
+    for dim_length in new_shape:
+        new_size *= dim_length
+        if dim_length < 0:
+            raise ValueError(
+                'all elements of `new_shape` must be non-negative'
+            )
+
+    if xp_size(a) == 0 or new_size == 0:
+        # First case must zero fill. The second would have repeats == 0.
+        return xp.zeros_like(a, shape=new_shape)
+
+    repeats = -(-new_size // xp_size(a))  # ceil division
+    a = xp.concat((a,) * repeats)[:new_size]
+
+    return xp.reshape(a, new_shape)
+
+
 def lp2hp(b, a, wo=1.0):
     r"""
     Transform a lowpass filter prototype to a highpass filter.
@@ -1953,27 +2006,34 @@ def lp2hp(b, a, wo=1.0):
     >>> plt.legend()
 
     """
-    a, b = map(atleast_1d, (a, b))
+    xp = array_namespace(a, b)
+
+    a, b = map(xp.asarray, (a, b))
+    a, b = xp_promote(a, b, force_floating=True, xp=xp)
+    a = xpx.atleast_nd(a, ndim=1, xp=xp)
+    b = xpx.atleast_nd(b, ndim=1, xp=xp)
+
     try:
         wo = float(wo)
     except TypeError:
         wo = float(wo[0])
-    d = len(a)
-    n = len(b)
+    d = a.shape[0]
+    n = b.shape[0]
     if wo != 1:
-        pwo = pow(wo, np.arange(max((d, n))))
+        pwo = wo ** xp.arange(max((d, n)), dtype=xp.float64)
     else:
-        pwo = np.ones(max((d, n)), b.dtype.char)
+        pwo = xp.ones(max((d, n)), dtype=b.dtype)
     if d >= n:
-        outa = a[::-1] * pwo
-        outb = resize(b, (d,))
+        outa = xp.flip(a) * pwo
+        outb = xp.concat((xp.zeros(n, dtype=b.dtype), ))
+        outb = _resize(b, (d,), xp=xp)
         outb[n:] = 0.0
-        outb[:n] = b[::-1] * pwo[:n]
+        outb[:n] = xp.flip(b) * pwo[:n]
     else:
-        outb = b[::-1] * pwo
-        outa = resize(a, (n,))
+        outb = xp.flip(b) * pwo
+        outa = _resize(a, (n,), xp=xp)
         outa[d:] = 0.0
-        outa[:d] = a[::-1] * pwo[:d]
+        outa[:d] = xp.flip(a) * pwo[:d]
 
     return normalize(outb, outa)
 
@@ -2038,16 +2098,20 @@ def lp2bp(b, a, wo=1.0, bw=1.0):
     >>> plt.ylabel('Amplitude [dB]')
     >>> plt.legend()
     """
+    xp = array_namespace(a, b)
+
+    a, b = map(xp.asarray, (a, b))
+    a, b = xp_promote(a, b, force_floating=True, xp=xp)
+    a = xpx.atleast_nd(a, ndim=1, xp=xp)
+    b = xpx.atleast_nd(b, ndim=1, xp=xp)
 
-    a, b = map(atleast_1d, (a, b))
-    D = len(a) - 1
-    N = len(b) - 1
-    artype = mintypecode((a, b))
+    D = a.shape[0] - 1
+    N = b.shape[0] - 1
     ma = max([N, D])
     Np = N + ma
     Dp = D + ma
-    bprime = np.empty(Np + 1, artype)
-    aprime = np.empty(Dp + 1, artype)
+    bprime = xp.empty(Np + 1, dtype=b.dtype)
+    aprime = xp.empty(Dp + 1, dtype=a.dtype)
     wosq = wo * wo
     for j in range(Np + 1):
         val = 0.0
@@ -2126,15 +2190,20 @@ def lp2bs(b, a, wo=1.0, bw=1.0):
     >>> plt.ylabel('Amplitude [dB]')
     >>> plt.legend()
     """
-    a, b = map(atleast_1d, (a, b))
-    D = len(a) - 1
-    N = len(b) - 1
-    artype = mintypecode((a, b))
+    xp = array_namespace(a, b)
+
+    a, b = map(xp.asarray, (a, b))
+    a, b = xp_promote(a, b, force_floating=True, xp=xp)
+    a = xpx.atleast_nd(a, ndim=1, xp=xp)
+    b = xpx.atleast_nd(b, ndim=1, xp=xp)
+
+    D = a.shape[0] - 1
+    N = b.shape[0] - 1
     M = max([N, D])
     Np = M + M
     Dp = M + M
-    bprime = np.empty(Np + 1, artype)
-    aprime = np.empty(Dp + 1, artype)
+    bprime = xp.empty(Np + 1, dtype=b.dtype)
+    aprime = xp.empty(Dp + 1, dtype=a.dtype)
     wosq = wo * wo
     for j in range(Np + 1):
         val = 0.0

scipy/signal/tests/test_filter_design.py

@@ -1,3 +1,4 @@
+import math
 import warnings
 
 from itertools import product
@@ -28,6 +29,10 @@
 from scipy.signal._filter_design import (_cplxreal, _cplxpair, _norm_factor,
                                         _bessel_poly, _bessel_zeros)
 
+skip_xp_backends = pytest.mark.skip_xp_backends
+xfail_xp_backends = pytest.mark.xfail_xp_backends
+
+
 try:
     import mpmath
 except ImportError:
@@ -1351,43 +1356,50 @@ def test_errors(self):
 
 class TestLp2lp:
 
-    def test_basic(self):
-        b = [1]
-        a = [1, np.sqrt(2), 1]
+    def test_basic(self, xp):
+        b = xp.asarray([1])
+        a = xp.asarray([1, math.sqrt(2), 1])
         b_lp, a_lp = lp2lp(b, a, 0.38574256627112119)
-        assert_array_almost_equal(b_lp, [0.1488], decimal=4)
-        assert_array_almost_equal(a_lp, [1, 0.5455, 0.1488], decimal=4)
+        assert_array_almost_equal(b_lp, xp.asarray([0.1488]), decimal=4)
+        assert_array_almost_equal(a_lp, xp.asarray([1, 0.5455, 0.1488]), decimal=4)
 
 
 class TestLp2hp:
 
-    def test_basic(self):
-        b = [0.25059432325190018]
-        a = [1, 0.59724041654134863, 0.92834805757524175, 0.25059432325190018]
-        b_hp, a_hp = lp2hp(b, a, 2*np.pi*5000)
-        xp_assert_close(b_hp, [1.0, 0, 0, 0])
-        xp_assert_close(a_hp, [1, 1.1638e5, 2.3522e9, 1.2373e14], rtol=1e-4)
+    def test_basic(self, xp):
+        b = xp.asarray([0.25059432325190018])
+        a = xp.asarray(
+            [1, 0.59724041654134863, 0.92834805757524175, 0.25059432325190018]
+        )
+        b_hp, a_hp = lp2hp(b, a, 2*math.pi*5000)
+        xp_assert_close(b_hp, xp.asarray([1.0, 0, 0, 0]))
+        xp_assert_close(
+            a_hp, xp.asarray([1, 1.1638e5, 2.3522e9, 1.2373e14]), rtol=1e-4
+        )
 
 
 class TestLp2bp:
 
-    def test_basic(self):
-        b = [1]
-        a = [1, 2, 2, 1]
-        b_bp, a_bp = lp2bp(b, a, 2*np.pi*4000, 2*np.pi*2000)
-        xp_assert_close(b_bp, [1.9844e12, 0, 0, 0], rtol=1e-6)
-        xp_assert_close(a_bp, [1, 2.5133e4, 2.2108e9, 3.3735e13,
-                               1.3965e18, 1.0028e22, 2.5202e26], rtol=1e-4)
+    def test_basic(self, xp):
+        b = xp.asarray([1])
+        a = xp.asarray([1, 2, 2, 1])
+        b_bp, a_bp = lp2bp(b, a, 2*math.pi*4000, 2*math.pi*2000)
+        xp_assert_close(b_bp, xp.asarray([1.9844e12, 0, 0, 0]), rtol=1e-6)
+        xp_assert_close(
+            a_bp,
+            xp.asarray([1, 2.5133e4, 2.2108e9, 3.3735e13,
+                        1.3965e18, 1.0028e22, 2.5202e26]), rtol=1e-4
+        )
 
 
 class TestLp2bs:
 
-    def test_basic(self):
-        b = [1]
-        a = [1, 1]
+    def test_basic(self, xp):
+        b = xp.asarray([1])
+        a = xp.asarray([1, 1])
         b_bs, a_bs = lp2bs(b, a, 0.41722257286366754, 0.18460575326152251)
-        assert_array_almost_equal(b_bs, [1, 0, 0.17407], decimal=5)
-        assert_array_almost_equal(a_bs, [1, 0.18461, 0.17407], decimal=5)
+        assert_array_almost_equal(b_bs, xp.asarray([1, 0, 0.17407]), decimal=5)
+        assert_array_almost_equal(a_bs, xp.asarray([1, 0.18461, 0.17407]), decimal=5)
 
 
 class TestBilinear: