ENH: linalg: enable N-D batch support in special matrix functions (scipy#21446)

---------

Co-authored-by: Jake Bowhay <[email protected]>

Author: mdhaber

scipy/linalg/_special_matrices.py

@@ -4,6 +4,7 @@
 import numpy as np
 from numpy.lib.stride_tricks import as_strided
 
+
 __all__ = ['toeplitz', 'circulant', 'hankel',
            'hadamard', 'leslie', 'kron', 'block_diag', 'companion',
            'helmert', 'hilbert', 'invhilbert', 'pascal', 'invpascal', 'dft',
@@ -16,7 +17,7 @@
 
 
 def toeplitz(c, r=None):
-    """
+    r"""
     Construct a Toeplitz matrix.
 
     The Toeplitz matrix has constant diagonals, with c as its first column
@@ -35,6 +36,12 @@ def toeplitz(c, r=None):
         ``[c[0], r[1:]]``.  Whatever the actual shape of `r`, it will be
         converted to a 1-D array.
 
+        .. warning::
+
+            Beginning in SciPy 1.17, multidimensional input will be treated as a batch,
+            not ``ravel``\ ed. To preserve the existing behavior, ``ravel`` aruments
+            before passing them to `toeplitz`.
+
     Returns
     -------
     A : (len(c), len(r)) ndarray
@@ -65,11 +72,19 @@ def toeplitz(c, r=None):
            [ 4.-1.j,  2.+3.j,  1.+0.j]])
 
     """
-    c = np.asarray(c).ravel()
+    c = np.asarray(c)
     if r is None:
         r = c.conjugate()
     else:
-        r = np.asarray(r).ravel()
+        r = np.asarray(r)
+
+    if c.ndim > 1 or r.ndim > 1:
+        msg = ("Beginning in SciPy 1.17, multidimensional input will be treated as a "
+               "batch, not `ravel`ed. To preserve the existing behavior and silence "
+               "this warning, `ravel` aruments before passing them to `toeplitz`.")
+        warnings.warn(msg, FutureWarning, stacklevel=2)
+
+    c, r = c.ravel(), r.ravel()
     # Form a 1-D array containing a reversed c followed by r[1:] that could be
     # strided to give us toeplitz matrix.
     vals = np.concatenate((c[::-1], r[1:]))
@@ -246,18 +261,19 @@ def leslie(f, s):
 
     Parameters
     ----------
-    f : (N,) array_like
+    f : (..., N,) array_like
         The "fecundity" coefficients.
-    s : (N-1,) array_like
-        The "survival" coefficients, has to be 1-D.  The length of `s`
-        must be one less than the length of `f`, and it must be at least 1.
+    s : (..., N-1,) array_like
+        The "survival" coefficients. The length of each slice of `s` (along the last
+        axis) must be one less than the length of `f`, and it must be at least 1.
 
     Returns
     -------
-    L : (N, N) ndarray
+    L : (..., N, N) ndarray
         The array is zero except for the first row,
         which is `f`, and the first sub-diagonal, which is `s`.
-        The data-type of the array will be the data-type of ``f[0]+s[0]``.
+        For 1-D input, the data-type of the array will be the data-type of
+        ``f[0]+s[0]``.
 
     Notes
     -----
@@ -270,6 +286,11 @@ def leslie(f, s):
     class, and the `n` - 1 "survival coefficients", which give the
     per-capita survival rate of each age class.
 
+    N-dimensional input are treated as a batches of coefficient arrays: each
+    slice along the last axis of the input arrays is a 1-D coefficient array,
+    and each slice along the last two dimensions of the output is the
+    corresponding Leslie matrix.
+
     References
     ----------
     .. [1] P. H. Leslie, On the use of matrices in certain population
@@ -290,18 +311,21 @@ def leslie(f, s):
     """
     f = np.atleast_1d(f)
     s = np.atleast_1d(s)
-    if f.ndim != 1:
-        raise ValueError("Incorrect shape for f.  f must be 1D")
-    if s.ndim != 1:
-        raise ValueError("Incorrect shape for s.  s must be 1D")
-    if f.size != s.size + 1:
-        raise ValueError("Incorrect lengths for f and s.  The length"
-                         " of s must be one less than the length of f.")
-    if s.size == 0:
+
+    if f.shape[-1] != s.shape[-1] + 1:
+        raise ValueError("Incorrect lengths for f and s. The length of s along "
+                         "the last axis must be one less than the length of f.")
+    if s.shape[-1] == 0:
         raise ValueError("The length of s must be at least 1.")
 
+    n = f.shape[-1]
+
+    if f.ndim > 1 or s.ndim > 1:
+        from scipy.stats._resampling import _vectorize_statistic
+        _leslie_nd = _vectorize_statistic(leslie)
+        return np.moveaxis(_leslie_nd(f, s, axis=-1), [0, 1], [-2, -1])
+
     tmp = f[0] + s[0]
-    n = f.size
     a = np.zeros((n, n), dtype=tmp.dtype)
     a[0] = f
     a[list(range(1, n)), list(range(0, n - 1))] = s
@@ -948,12 +972,16 @@ def fiedler(a):
 
     Parameters
     ----------
-    a : (n,) array_like
-        coefficient array
+    a : (..., n,) array_like
+        Coefficient array. N-dimensional arrays are treated as a batch:
+        each slice along the last axis is a 1-D coefficient array.
 
     Returns
     -------
-    F : (n, n) ndarray
+    F : (..., n, n) ndarray
+        Fiedler matrix. For batch input, each slice of shape ``(n, n)``
+        along the last two dimensions of the output corresponds with a
+        slice of shape ``(n,)`` along the last dimension of the input.
 
     See Also
     --------
@@ -1003,8 +1031,8 @@ def fiedler(a):
     """
     a = np.atleast_1d(a)
 
-    if a.ndim != 1:
-        raise ValueError("Input 'a' must be a 1D array.")
+    if a.ndim > 1:
+        return np.apply_along_axis(fiedler, -1, a)
 
     if a.size == 0:
         return np.array([], dtype=float)
@@ -1023,23 +1051,28 @@ def fiedler_companion(a):
 
     Parameters
     ----------
-    a : (N,) array_like
+    a : (..., N) array_like
         1-D array of polynomial coefficients in descending order with a nonzero
         leading coefficient. For ``N < 2``, an empty array is returned.
+        N-dimensional arrays are treated as a batch: each slice along the last
+        axis is a 1-D array of polynomial coefficients.
 
     Returns
     -------
-    c : (N-1, N-1) ndarray
-        Resulting companion matrix
+    c : (..., N-1, N-1) ndarray
+        Resulting companion matrix. For batch input, each slice of shape
+        ``(N-1, N-1)`` along the last two dimensions of the output corresponds
+        with a slice of shape ``(N,)`` along the last dimension of the input.
 
     See Also
     --------
     companion
 
     Notes
     -----
-    Similar to `companion` the leading coefficient should be nonzero. In the case
-    the leading coefficient is not 1, other coefficients are rescaled before
+    Similar to `companion`, each leading coefficient along the last axis of the
+    input should be nonzero.
+    If the leading coefficient is not 1, other coefficients are rescaled before
     the array generation. To avoid numerical issues, it is best to provide a
     monic polynomial.
 
@@ -1067,8 +1100,8 @@ def fiedler_companion(a):
     """
     a = np.atleast_1d(a)
 
-    if a.ndim != 1:
-        raise ValueError("Input 'a' must be a 1-D array.")
+    if a.ndim > 1:
+        return np.apply_along_axis(fiedler_companion, -1, a)
 
     if a.size <= 2:
         if a.size == 2:
@@ -1101,8 +1134,9 @@ def convolution_matrix(a, n, mode='full'):
 
     Parameters
     ----------
-    a : (m,) array_like
-        The 1-D array to convolve.
+    a : (..., m) array_like
+        The 1-D array to convolve. N-dimensional arrays are treated as a
+        batch: each slice along the last axis is a 1-D array to convolve.
     n : int
         The number of columns in the resulting matrix.  It gives the length
         of the input to be convolved with `a`.  This is analogous to the
@@ -1114,7 +1148,7 @@ def convolution_matrix(a, n, mode='full'):
 
     Returns
     -------
-    A : (k, n) ndarray
+    A : (..., k, n) ndarray
         The convolution matrix whose row count `k` depends on `mode`::
 
             =======  =========================
@@ -1125,6 +1159,10 @@ def convolution_matrix(a, n, mode='full'):
             'valid'  max(m, n) - min(m, n) + 1
             =======  =========================
 
+        For batch input, each slice of shape ``(k, n)`` along the last two
+        dimensions of the output corresponds with a slice of shape ``(m,)``
+        along the last dimension of the input.
+
     See Also
     --------
     toeplitz : Toeplitz matrix
@@ -1244,16 +1282,17 @@ def convolution_matrix(a, n, mode='full'):
         raise ValueError('n must be a positive integer.')
 
     a = np.asarray(a)
-    if a.ndim != 1:
-        raise ValueError('convolution_matrix expects a one-dimensional '
-                         'array as input')
+
     if a.size == 0:
         raise ValueError('len(a) must be at least 1.')
 
     if mode not in ('full', 'valid', 'same'):
         raise ValueError(
             "'mode' argument must be one of ('full', 'valid', 'same')")
 
+    if a.ndim > 1:
+        return np.apply_along_axis(lambda a: convolution_matrix(a, n, mode), -1, a)
+
     # create zero padded versions of the array
     az = np.pad(a, (0, n-1), 'constant')
     raz = np.pad(a[::-1], (0, n-1), 'constant')

scipy/linalg/tests/test_matmul_toeplitz.py

@@ -128,7 +128,8 @@ def do(self, x, c, r=None, check_finite=False, workers=None):
         if r is None:
             actual = matmul_toeplitz(c, x, check_finite, workers)
         else:
+            r = np.ravel(r)
             actual = matmul_toeplitz((c, r), x, check_finite)
-        desired = toeplitz(c, r) @ x
+        desired = toeplitz(np.ravel(c), r) @ x
         assert_allclose(actual, desired,
             rtol=self.tolerance, atol=self.tolerance)

scipy/linalg/tests/test_special_matrices.py

@@ -100,7 +100,6 @@ class TestLeslie:
 
     def test_bad_shapes(self):
         assert_raises(ValueError, leslie, [[1, 1], [2, 2]], [3, 4, 5])
-        assert_raises(ValueError, leslie, [3, 4, 5], [[1, 1], [2, 2]])
         assert_raises(ValueError, leslie, [1, 2], [1, 2])
         assert_raises(ValueError, leslie, [1], [])
 
@@ -585,11 +584,6 @@ def test_bad_n(self):
         with pytest.raises(ValueError, match='n must be a positive integer'):
             convolution_matrix([1, 2, 3], 0)
 
-    def test_bad_first_arg(self):
-        # first arg must be a 1d array, otherwise ValueError
-        with pytest.raises(ValueError, match='one-dimensional'):
-            convolution_matrix(1, 4)
-
     def test_empty_first_arg(self):
         # first arg must have at least one value
         with pytest.raises(ValueError, match=r'len\(a\)'):
@@ -615,3 +609,29 @@ def test_against_numpy_convolve(self, cpx, na, nv, mode):
             A = convolution_matrix(a, nv, mode)
         y2 = A @ v
         assert_array_almost_equal(y1, y2)
+
+
+@pytest.mark.fail_slow(5)  # `leslie` has an import in the function
+@pytest.mark.parametrize('f, args', [(companion, ()),
+                                     (convolution_matrix, (5, 'same')),
+                                     (fiedler, ()),
+                                     (fiedler_companion, ()),
+                                     (leslie, (np.arange(9),)),
+                                     (toeplitz, (np.arange(9),)),
+                                     ])
+def test_batch(f, args):
+    rng = np.random.default_rng(283592436523456)
+    batch_shape = (2, 3)
+    m = 10
+    A = rng.random(batch_shape + (m,))
+
+    if f in {toeplitz}:
+        message = "Beginning in SciPy 1.17, multidimensional input will be..."
+        with pytest.warns(FutureWarning, match=message):
+            f(A, *args)
+        return
+
+    res = f(A, *args)
+    ref = np.asarray([f(a, *args) for a in A.reshape(-1, m)])
+    ref = ref.reshape(A.shape[:-1] + ref.shape[-2:])
+    assert_allclose(res, ref)