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

Author: ev-br

scipy/signal/_filter_design.py

@@ -2794,7 +2794,7 @@ def _relative_degree(z, p):
     """
     Return relative degree of transfer function from zeros and poles
     """
-    degree = len(p) - len(z)
+    degree = p.shape[0] - z.shape[0]
     if degree < 0:
         raise ValueError("Improper transfer function. "
                          "Must have at least as many poles as zeros.")
@@ -2941,8 +2941,12 @@ def lp2lp_zpk(z, p, k, wo=1.0):
     >>> lp2lp_zpk(z, p, k, wo)
     (   array([2.8, 0.8]), array([2. , 5.2]), 0.8)
     """
-    z = atleast_1d(z)
-    p = atleast_1d(p)
+    xp = array_namespace(z, p)
+
+    z, p = map(xp.asarray, (z, p))
+    z = xpx.atleast_nd(z, ndim=1, xp=xp)
+    p = xpx.atleast_nd(p, ndim=1, xp=xp)
+
     wo = float(wo)  # Avoid int wraparound
 
     degree = _relative_degree(z, p)
@@ -3018,8 +3022,12 @@ def lp2hp_zpk(z, p, k, wo=1.0):
         array([-0.6 , -0.15]),
         8.5)
     """
-    z = atleast_1d(z)
-    p = atleast_1d(p)
+    xp = array_namespace(z, p)
+
+    z, p = map(xp.asarray, (z, p))
+    z = xpx.atleast_nd(z, ndim=1, xp=xp)
+    p = xpx.atleast_nd(p, ndim=1, xp=xp)
+
     wo = float(wo)
 
     degree = _relative_degree(z, p)
@@ -3030,10 +3038,10 @@ def lp2hp_zpk(z, p, k, wo=1.0):
     p_hp = wo / p
 
     # If lowpass had zeros at infinity, inverting moves them to origin.
-    z_hp = append(z_hp, zeros(degree))
+    z_hp = xp.concat((z_hp, xp.zeros(degree)))
 
     # Cancel out gain change caused by inversion
-    k_hp = k * real(prod(-z) / prod(-p))
+    k_hp = k * xp.real(xp.prod(-z) / xp.prod(-p))
 
     return z_hp, p_hp, k_hp
 
@@ -3104,8 +3112,12 @@ def lp2bp_zpk(z, p, k, wo=1.0, bw=1.0):
         array([1.04996339e+02+0.j, -1.60167736e-03+0.j,  3.66108003e-03+0.j,
                -2.39998398e+02+0.j]), 0.8)
     """
-    z = atleast_1d(z)
-    p = atleast_1d(p)
+    xp = array_namespace(z, p)
+
+    z, p = map(xp.asarray, (z, p))
+    z = xpx.atleast_nd(z, ndim=1, xp=xp)
+    p = xpx.atleast_nd(p, ndim=1, xp=xp)
+
     wo = float(wo)
     bw = float(bw)
 
@@ -3116,17 +3128,17 @@ def lp2bp_zpk(z, p, k, wo=1.0, bw=1.0):
     p_lp = p * bw/2
 
     # Square root needs to produce complex result, not NaN
-    z_lp = z_lp.astype(complex)
-    p_lp = p_lp.astype(complex)
+    z_lp = xp.astype(z_lp, xp.complex128)
+    p_lp = xp.astype(p_lp, xp.complex128)
 
     # Duplicate poles and zeros and shift from baseband to +wo and -wo
-    z_bp = concatenate((z_lp + sqrt(z_lp**2 - wo**2),
-                        z_lp - sqrt(z_lp**2 - wo**2)))
-    p_bp = concatenate((p_lp + sqrt(p_lp**2 - wo**2),
-                        p_lp - sqrt(p_lp**2 - wo**2)))
+    z_bp = xp.concat((z_lp + xp.sqrt(z_lp**2 - wo**2),
+                      z_lp - xp.sqrt(z_lp**2 - wo**2)))
+    p_bp = xp.concat((p_lp + xp.sqrt(p_lp**2 - wo**2),
+                      p_lp - xp.sqrt(p_lp**2 - wo**2)))
 
     # Move degree zeros to origin, leaving degree zeros at infinity for BPF
-    z_bp = append(z_bp, zeros(degree))
+    z_bp = xp.concat((z_bp, xp.zeros(degree)))
 
     # Cancel out gain change from frequency scaling
     k_bp = k * bw**degree
@@ -3199,8 +3211,12 @@ def lp2bs_zpk(z, p, k, wo=1.0, bw=1.0):
         array([14.2681928 +0.j, -0.02506281+0.j,  0.01752149+0.j, -9.97493719+0.j]),
         -12.857142857142858)
     """
-    z = atleast_1d(z)
-    p = atleast_1d(p)
+    xp = array_namespace(z, p)
+
+    z, p = map(xp.asarray, (z, p))
+    z = xpx.atleast_nd(z, ndim=1, xp=xp)
+    p = xpx.atleast_nd(p, ndim=1, xp=xp)
+
     wo = float(wo)
     bw = float(bw)
 
@@ -3211,21 +3227,21 @@ def lp2bs_zpk(z, p, k, wo=1.0, bw=1.0):
     p_hp = (bw/2) / p
 
     # Square root needs to produce complex result, not NaN
-    z_hp = z_hp.astype(complex)
-    p_hp = p_hp.astype(complex)
+    z_hp = xp.astype(z_hp, xp.complex128)
+    p_hp = xp.astype(p_hp, xp.complex128)
 
     # Duplicate poles and zeros and shift from baseband to +wo and -wo
-    z_bs = concatenate((z_hp + sqrt(z_hp**2 - wo**2),
-                        z_hp - sqrt(z_hp**2 - wo**2)))
-    p_bs = concatenate((p_hp + sqrt(p_hp**2 - wo**2),
-                        p_hp - sqrt(p_hp**2 - wo**2)))
+    z_bs = xp.concat((z_hp + xp.sqrt(z_hp**2 - wo**2),
+                      z_hp - xp.sqrt(z_hp**2 - wo**2)))
+    p_bs = xp.concat((p_hp + xp.sqrt(p_hp**2 - wo**2),
+                      p_hp - xp.sqrt(p_hp**2 - wo**2)))
 
     # Move any zeros that were at infinity to the center of the stopband
-    z_bs = append(z_bs, full(degree, +1j*wo))
-    z_bs = append(z_bs, full(degree, -1j*wo))
+    z_bs = xp.concat((z_bs, xp.full(degree, +1j*wo)))
+    z_bs = xp.concat((z_bs, xp.full(degree, -1j*wo)))
 
     # Cancel out gain change caused by inversion
-    k_bs = k * real(prod(-z) / prod(-p))
+    k_bs = k * xp.real(xp.prod(-z) / xp.prod(-p))
 
     return z_bs, p_bs, k_bs
 

scipy/signal/tests/test_filter_design.py

@@ -1,4 +1,5 @@
 import math
+import cmath
 import warnings
 
 from itertools import product
@@ -1491,25 +1492,42 @@ def test_fs_validation(self):
             bilinear(b, a, fs=None)
 
 
+def _sort_cmplx(arr, xp):
+    # xp.sort is undefined for complex dtypes. Here we only need some
+    # consistent way to sort a complex array, including equal magnitude elements.
+    arr = xp.asarray(arr)
+    if xp.isdtype(arr.dtype, 'complex floating'):
+        sorter = abs(arr) + xp.real(arr) + xp.imag(arr)**3
+    else:
+        sorter = arr
+
+    idxs = xp.argsort(sorter)
+    return arr[idxs]
+
+
 class TestLp2lp_zpk:
 
-    def test_basic(self):
-        z = []
-        p = [(-1+1j)/np.sqrt(2), (-1-1j)/np.sqrt(2)]
+    def test_basic(self, xp):
+        z = xp.asarray([])
+        p = xp.asarray([(-1+1j) / math.sqrt(2), (-1-1j) / math.sqrt(2)])
         k = 1
         z_lp, p_lp, k_lp = lp2lp_zpk(z, p, k, 5)
-        xp_assert_equal(z_lp, [])
-        xp_assert_close(sort(p_lp), sort(p)*5)
-        xp_assert_close(k_lp, 25.)
+        xp_assert_equal(z_lp, xp.asarray([]))
+        xp_assert_close(_sort_cmplx(p_lp, xp=xp), _sort_cmplx(p, xp=xp) * 5)
+        assert k_lp == 25.
 
         # Pseudo-Chebyshev with both poles and zeros
-        z = [-2j, +2j]
-        p = [-0.75, -0.5-0.5j, -0.5+0.5j]
+        z = xp.asarray([-2j, +2j])
+        p = xp.asarray([-0.75, -0.5-0.5j, -0.5+0.5j])
         k = 3
         z_lp, p_lp, k_lp = lp2lp_zpk(z, p, k, 20)
-        xp_assert_close(sort(z_lp), sort([-40j, +40j]))
-        xp_assert_close(sort(p_lp), sort([-15, -10-10j, -10+10j]))
-        xp_assert_close(k_lp, 60.)
+        xp_assert_close(
+            _sort_cmplx(z_lp, xp=xp), _sort_cmplx([-40j, +40j], xp=xp)
+        )
+        xp_assert_close(
+            _sort_cmplx(p_lp, xp=xp), _sort_cmplx([-15, -10-10j, -10+10j], xp=xp)
+        )
+        assert k_lp == 60.
 
     def test_fs_validation(self):
         z = [-2j, +2j]
@@ -1525,23 +1543,27 @@ def test_fs_validation(self):
 
 class TestLp2hp_zpk:
 
-    def test_basic(self):
-        z = []
-        p = [(-1+1j)/np.sqrt(2), (-1-1j)/np.sqrt(2)]
+    def test_basic(self, xp):
+        z = xp.asarray([])
+        p = xp.asarray([(-1+1j) / math.sqrt(2), (-1-1j) / math.sqrt(2)])
         k = 1
 
         z_hp, p_hp, k_hp = lp2hp_zpk(z, p, k, 5)
-        xp_assert_equal(z_hp, np.asarray([0.0, 0.0]))
-        xp_assert_close(sort(p_hp), sort(p)*5)
-        xp_assert_close(k_hp, 1.0)
+        xp_assert_equal(z_hp, xp.asarray([0.0, 0.0]))
+        xp_assert_close(_sort_cmplx(p_hp, xp=xp), _sort_cmplx(p, xp=xp) * 5)
+        assert math.isclose(k_hp, 1.0)
 
-        z = [-2j, +2j]
-        p = [-0.75, -0.5-0.5j, -0.5+0.5j]
+        z = xp.asarray([-2j, +2j])
+        p = xp.asarray([-0.75, -0.5-0.5j, -0.5+0.5j])
         k = 3
         z_hp, p_hp, k_hp = lp2hp_zpk(z, p, k, 6)
-        xp_assert_close(sort(z_hp), sort([-3j, 0, +3j]))
-        xp_assert_close(sort(p_hp), sort([-8, -6-6j, -6+6j]))
-        xp_assert_close(k_hp, 32.0)
+        xp_assert_close(
+            _sort_cmplx(z_hp, xp=xp), _sort_cmplx([-3j, 0, +3j], xp=xp)
+        )
+        xp_assert_close(
+            _sort_cmplx(p_hp, xp=xp), _sort_cmplx([-8, -6-6j, -6+6j], xp=xp)
+        )
+        assert k_hp == 32.0
 
 
 class TestLp2bp_zpk:
@@ -1563,25 +1585,31 @@ def test_basic(self):
 
 class TestLp2bs_zpk:
 
-    def test_basic(self):
-        z = [-2j, +2j]
-        p = [-0.75, -0.5-0.5j, -0.5+0.5j]
+    def test_basic(self, xp):
+        z = xp.asarray([-2j, +2j])
+        p = xp.asarray([-0.75, -0.5-0.5j, -0.5+0.5j])
         k = 3
 
         z_bs, p_bs, k_bs = lp2bs_zpk(z, p, k, 35, 12)
 
-        xp_assert_close(sort(z_bs), sort([+35j, -35j,
-                                          +3j+sqrt(1234)*1j,
-                                          -3j+sqrt(1234)*1j,
-                                          +3j-sqrt(1234)*1j,
-                                          -3j-sqrt(1234)*1j]))
-        xp_assert_close(sort(p_bs), sort([+3j*sqrt(129) - 8,
-                                          -3j*sqrt(129) - 8,
-                                          (-6 + 6j) - sqrt(-1225 - 72j),
-                                          (-6 - 6j) - sqrt(-1225 + 72j),
-                                          (-6 + 6j) + sqrt(-1225 - 72j),
-                                          (-6 - 6j) + sqrt(-1225 + 72j), ]))
-        xp_assert_close(k_bs, 32.0)
+        xp_assert_close(
+            _sort_cmplx(z_bs, xp=xp),
+            _sort_cmplx([+35j, -35j,
+                         +3j + math.sqrt(1234)*1j,
+                         -3j + math.sqrt(1234)*1j,
+                         +3j - math.sqrt(1234)*1j,
+                         -3j - math.sqrt(1234)*1j], xp=xp)
+        )
+        xp_assert_close(
+            _sort_cmplx(p_bs, xp=xp),
+            _sort_cmplx([+3j*math.sqrt(129) - 8,
+                         -3j*math.sqrt(129) - 8,
+                         (-6 + 6j) - cmath.sqrt(-1225 - 72j),
+                         (-6 - 6j) - cmath.sqrt(-1225 + 72j),
+                         (-6 + 6j) + cmath.sqrt(-1225 - 72j),
+                         (-6 - 6j) + cmath.sqrt(-1225 + 72j), ], xp=xp)
+        )
+        assert math.isclose(k_bs, 32.0)
 
 
 class TestBilinear_zpk: