ENH: linalg: add batch support for matrix -> scalar funcs (scipy#22127)

* ENH: linalg.expm_cond: add batch support

* DOC: linalg: inject batch support note automatically

* ENH: linalg.issymmetric/ishermitian: add batch support

* MAINT: linalg: adjustments to address CI failures

Author: mdhaber

scipy/_lib/_util.py

@@ -1176,6 +1176,14 @@ def _dict_formatter(d, n=0, mplus=1, sorter=None):
     return s
 
 
+_batch_note = """
+The documentation is written assuming array arguments are of specified
+"core" shapes. However, array argument(s) of this function may have additional
+"batch" dimensions prepended to the core shape. In this case, the array is treated
+as a batch of lower-dimensional slices; see :ref:`linalg_batch` for details.
+"""
+
+
 def _apply_over_batch(*argdefs):
     """
     Factory for decorator that applies a function over batched arguments.
@@ -1266,5 +1274,9 @@ def wrapper(*args, **kwargs):
             # contributor to pass an `pack_result` callable to the decorator factory.
             return results[0] if len(results) == 1 else results
 
+        doc = FunctionDoc(wrapper)
+        doc['Extended Summary'].append(_batch_note.rstrip())
+        wrapper.__doc__ = str(doc).split("\n", 1)[1]  # remove signature
+
         return wrapper
     return decorator

scipy/linalg/_cythonized_array_utils.pyx

@@ -8,6 +8,7 @@ from scipy.linalg._cythonized_array_utils cimport (
     )
 from scipy.linalg.cython_lapack cimport sgetrf, dgetrf, cgetrf, zgetrf
 from libc.stdlib cimport malloc, free
+from scipy._lib._util import _apply_over_batch
 
 __all__ = ['bandwidth', 'issymmetric', 'ishermitian']
 
@@ -237,6 +238,7 @@ cdef inline (int, int) band_check_internal_noncontig(const np_numeric_t[:, :]A)
     return lower_band, upper_band
 
 
+@_apply_over_batch(('a', 2))
 @cython.embedsignature(True)
 def issymmetric(a, atol=None, rtol=None):
     """Check if a square 2D array is symmetric.
@@ -367,6 +369,7 @@ cdef inline bint is_sym_her_real_noncontig_internal(const np_numeric_t[:, :]A) n
     return True
 
 
+@_apply_over_batch(('a', 2))
 @cython.embedsignature(True)
 def ishermitian(a, atol=None, rtol=None):
     """Check if a square 2D array is Hermitian.

scipy/linalg/_decomp.py

@@ -124,11 +124,6 @@ def eig(a, b=None, left=False, right=True, overwrite_a=False,
 
     where ``.H`` is the Hermitian conjugation.
 
-    The documentation is written assuming array arguments are of specified
-    "core" shapes. However, array argument(s) of this function may have additional
-    "batch" dimensions prepended to the core shape. In this case, the array is treated
-    as a batch of lower-dimensional slices; see :ref:`linalg_batch` for details.
-
     Parameters
     ----------
     a : (M, M) array_like

scipy/linalg/_expm_frechet.py

@@ -1,6 +1,8 @@
 """Frechet derivative of the matrix exponential."""
 import numpy as np
 import scipy.linalg
+from scipy._lib._util import _apply_over_batch
+
 
 __all__ = ['expm_frechet', 'expm_cond']
 
@@ -351,6 +353,7 @@ def expm_frechet_kronform(A, method=None, check_finite=True):
     return np.vstack(cols).T
 
 
+@_apply_over_batch(('A', 2))
 def expm_cond(A, check_finite=True):
     """
     Relative condition number of the matrix exponential in the Frobenius norm.

scipy/linalg/tests/meson.build

@@ -1,6 +1,7 @@
 python_sources = [
   '__init__.py',
   'test_basic.py',
+  'test_batch.py',
   'test_blas.py',
   'test_cython_blas.py',
   'test_cython_lapack.py',

scipy/linalg/tests/test_batch.py

@@ -0,0 +1,56 @@
+import pytest
+import numpy as np
+from numpy.testing import assert_allclose
+from scipy import linalg
+
+
+real_floating = [np.float32, np.float64]
+complex_floating = [np.complex64, np.complex128]
+floating = real_floating + complex_floating
+
+
+def get_nearly_hermitian(shape, dtype, atol, rng):
+    # Generate a batch of nearly Hermitian matrices with specified
+    # `shape` and `dtype`. `atol` controls the level of noise in
+    # Hermitian-ness to by generated by `rng`.
+    A = rng.random(shape).astype(dtype)
+    At = np.conj(A.swapaxes(-1, -2))
+    noise = rng.standard_normal(size=A.shape).astype(dtype) * atol
+    return A + At + noise
+
+
+class TestMatrixInScalarOut:
+
+    def batch_test(self, fun, args=(), kwargs=None, dtype=np.float64,
+                   batch_shape=(5, 3), core_shape=(4, 4), seed=8342310302941288912051):
+        kwargs = {} if kwargs is None else kwargs
+        rng = np.random.default_rng(seed)
+        # test_expm_cond doesn't need symmetric/hermitian matrices, and
+        # test_issymmetric doesn't need hermitian matrices, but it doesn't hurt.
+        A = get_nearly_hermitian(batch_shape + core_shape, dtype, 3e-4, rng)
+
+        res = fun(A, *args, **kwargs)
+
+        for i in range(batch_shape[0]):
+            for j in range(batch_shape[1]):
+                ref = fun(A[i, j], *args, **kwargs)
+                assert_allclose(res[i, j], ref)
+
+        return res
+
+    @pytest.mark.parametrize('dtype', floating)
+    def test_expm_cond(self, dtype):
+        self.batch_test(linalg.expm_cond, dtype=dtype)
+
+    @pytest.mark.parametrize('dtype', floating)
+    def test_issymmetric(self, dtype):
+        res = self.batch_test(linalg.issymmetric, dtype=dtype, kwargs=dict(atol=1e-3))
+        assert not np.all(res)  # ensure test is not trivial: not all True or False;
+        assert np.any(res)      # also confirms that `atol` is passed to issymmetric
+
+    @pytest.mark.parametrize('dtype', floating)
+    def test_ishermitian(self, dtype):
+        res = self.batch_test(linalg.issymmetric, dtype=dtype, kwargs=dict(atol=1e-3))
+        assert not np.all(res)  # ensure test is not trivial: not all True or False;
+        assert np.any(res)      # also confirms that `atol` is passed to ishermitian
+