Merge pull request scipy#23547 from ev-br/solve_sy_he

`linalg.solve`: move "sym"/"her" solvers to C

Author: tylerjereddy

benchmarks/benchmarks/linalg.py

@@ -102,6 +102,36 @@ def time_svd(self, size, contig, module):
     )
 
 
+class BatchedSolveBench(Benchmark):
+    params = [
+        [(100, 10, 10), (100, 20, 20), (100, 100)],
+        ["gen", "pos", "sym"],
+        ["scipy", "numpy"]
+    ]
+    param_names = ["shape", "structure" ,"module"]
+
+    def setup(self, shape, structure, module):
+        a = random(shape)
+        # larger diagonal ensures non-singularity:
+        for i in range(shape[-1]):
+            a[..., i, i] = 10*(.1+a[..., i, i])
+
+        if structure == "pos":
+            self.a = a @ a.mT
+        elif structure == "sym":
+            self.a = a + a.mT
+        else:
+            self.a = a
+
+        self.b = random([a.shape[-1]])
+
+    def time_solve(self, shape, structure, module):
+        if module == 'numpy':
+            nl.solve(self.a, self.b)
+        else:
+            sl.solve(self.a, self.b, assume_a=structure)
+
+
 class Norm(Benchmark):
     params = [
         [(20, 20), (100, 100), (1000, 1000), (20, 1000), (1000, 20)],

scipy/linalg/_basic.py

@@ -209,7 +209,7 @@ def solve(a, b, lower=False, overwrite_a=False,
            [ 5. , -4.5]])
     """
     if assume_a in [
-        'sym', 'her', 'symmetric', 'hermitian', 'diagonal', 'tridiagonal', 'banded'
+        'diagonal', 'tridiagonal', 'banded'
     ]:
         # TODO: handle these structures in this function
         return solve0(
@@ -225,6 +225,8 @@ def solve(a, b, lower=False, overwrite_a=False,
         'upper triangular': 21,
         'lower triangular': 22,
         'pos' : 101, 'positive definite': 101,
+        'sym' : 201, 'symmetric': 201,
+        'her' : 211, 'hermitian': 211,
     }.get(assume_a, 'unknown')
     if structure == 'unknown':
         raise ValueError(f'{assume_a} is not a recognized matrix structure')
@@ -1349,10 +1351,12 @@ def inv(a, overwrite_a=False, check_finite=True, assume_a=None, lower=False):
      upper triangular               'upper triangular'
      lower triangular               'lower triangular'
      symmetric positive definite    'pos'
+     symmetric                      'sym'
+     Hermitian                      'her'
     =============================  ================================
 
-    For the 'pos' option, only the triangle of the input matrix specified in
-    the `lower` argument is used, and the other triangle is not referenced.
+    For the 'pos', 'sym' and 'her' options, only the specified triangle of the input
+    matrix is used, and the other triangle is not referenced.
 
     Array argument(s) of this function may have additional
     "batch" dimensions prepended to the core shape. In this case, the array is treated
@@ -1432,14 +1436,16 @@ def inv(a, overwrite_a=False, check_finite=True, assume_a=None, lower=False):
         overwrite_a = True
         a1 = a1.copy()
 
-    # keep the numbers in sync with C
+    # keep the numbers in sync with C at `linalg/src/_common_array_utils.hh`
     structure = {
         None: -1,
         'general': 0,
         # 'diagonal': 11,
         'upper triangular': 21,
         'lower triangular': 22,
         'pos' : 101,
+        'sym' : 201,
+        'her' : 211,
     }[assume_a]
 
     # a1 is well behaved, invert it.

scipy/linalg/src/_batched_linalg_module.cc

@@ -6,7 +6,6 @@
 #include "_common_array_utils.hh"
 
 
-#define PYERR(errobj,message) {PyErr_SetString(errobj,message); return NULL;}
 static PyObject* _linalg_inv_error;
 
 
@@ -45,7 +44,8 @@ _linalg_inv(PyObject* Py_UNUSED(dummy), PyObject* args) {
     npy_intp* shape = PyArray_SHAPE(ap_Am);      // Array shape
     npy_intp n = shape[ndim - 1];                // Slice size
     if (n != shape[ndim - 2]) {
-        PYERR(PyExc_ValueError, "Last two dimensions of the input must be the same.")
+        PyErr_SetString(PyExc_ValueError, "Last two dimensions of the input must be the same.");
+        return NULL;
     }
 
     overwrite_a = 0; // TODO: enable it
@@ -79,13 +79,15 @@ _linalg_inv(PyObject* Py_UNUSED(dummy), PyObject* args) {
             info = _inverse<npy_complex128>(ap_Am, (npy_complex128 *)buf, structure, lower, overwrite_a, vec_status);
             break;
         default:
-            PYERR(PyExc_RuntimeError, "Unknown array type.")
+            PyErr_SetString(PyExc_RuntimeError, "Unknown array type.");
+            return NULL;
     }
 
     if (info < 0) {
         // Either OOM or internal LAPACK error.
         Py_DECREF(ap_Ainv);
-        PYERR(PyExc_RuntimeError, "Memory error in scipy.linalg.inv.")
+        PyErr_SetString(PyExc_RuntimeError, "Memory error in scipy.linalg.inv.");
+        return NULL;
     }
     PyObject *ret_lst = convert_vec_status(vec_status);
 
@@ -127,7 +129,8 @@ _linalg_solve(PyObject* Py_UNUSED(dummy), PyObject* args) {
     int ndim = PyArray_NDIM(ap_Am);
     npy_intp* shape = PyArray_SHAPE(ap_Am);
     if ((ndim < 2) || (shape[ndim - 1] != shape[ndim - 2])) {
-        PYERR(PyExc_ValueError, "Last two dimensions of `a` must be the same.")
+        PyErr_SetString(PyExc_ValueError, "Last two dimensions of `a` must be the same.");
+        return NULL;
     }
 
     // At the python call site, 
@@ -171,13 +174,15 @@ _linalg_solve(PyObject* Py_UNUSED(dummy), PyObject* args) {
             info = _solve<npy_complex128>(ap_Am, ap_b, (npy_complex128 *)buf, structure, lower, transposed, overwrite_a, vec_status);
             break;
         default:
-            PYERR(PyExc_RuntimeError, "Unknown array type.")
+            PyErr_SetString(PyExc_RuntimeError, "Unknown array type.");
+            return NULL;
     }
 
     if (info < 0) {
-        // Either OOM or internal LAPACK error.
+        // Either OOM error or requiested lwork too large.
         Py_DECREF(ap_x);
-        PYERR(PyExc_RuntimeError, "Memory error in scipy.linalg.solve.")
+        PyErr_SetString(PyExc_MemoryError, "Memory error in scipy.linalg.solve.");
+        return NULL;
     }
     PyObject *ret_lst = convert_vec_status(vec_status);