ENH: linalg: Add sy/hetrs LAPACK wrappers

Author: ilayn

scipy/linalg/flapack_sym_herm.pyf.src

@@ -302,6 +302,46 @@ end subroutine <prefix2c>heevd_lwork
    end subroutine <prefix>sytrf_lwork
 
 
+   subroutine <prefix>sytrs(n,nrhs,a,lda,ipiv,b,ldb,info,lower)
+   
+   ! Solve A * X = B for symmetric A matrix after calling ?sytrf
+   
+     callstatement (*f2py_func)((lower?"L":"U"),&n,&nrhs,a,&lda,ipiv,b,&ldb,&info)
+     callprotoargument char*,F_INT*,F_INT*,<ctype>*,F_INT*,F_INT*,<ctype>*,F_INT*,F_INT*
+     
+     <ftype> intent(in), dimension(lda, n), check((lda >= n) && (n >= 0)) :: a
+     <ftype> intent(in,out,copy,out=x),dimension(ldb, nrhs) :: b
+     integer optional,intent(in),check(lower==0||lower==1) :: lower = 0
+     integer intent(in),dimension(n),depend(n) :: ipiv
+     integer intent(hide), depend(a),check(n >= 0) :: n = shape(a, 1)
+     integer intent(hide), depend(b) :: nrhs = shape(b, 1)
+     integer intent(hide), depend(a) :: lda = MAX(1, shape(a, 0))
+     integer intent(hide), depend(b,n),check(ldb >= n) :: ldb = MAX(1, shape(b, 0))
+     integer intent(out):: info
+
+   end subroutine <prefix>sytrs
+
+
+   subroutine <prefix2c>hetrs(n,nrhs,a,lda,ipiv,b,ldb,info,lower)
+   
+   ! Solve A * X = B for hermitian A matrix after calling ?hetrf
+   
+     callstatement (*f2py_func)((lower?"L":"U"),&n,&nrhs,a,&lda,ipiv,b,&ldb,&info)
+     callprotoargument char*,F_INT*,F_INT*,<ctype2c>*,F_INT*,F_INT*,<ctype2c>*,F_INT*,F_INT*
+     
+     <ftype2c> intent(in), dimension(lda, n), check((lda >= n) && (n >= 0)) :: a
+     <ftype2c> intent(in,out,copy,out=x),dimension(ldb, nrhs) :: b
+     integer optional,intent(in),check(lower==0||lower==1) :: lower = 0
+     integer intent(in),dimension(n),depend(n) :: ipiv
+     integer intent(hide), depend(a),check(n >= 0) :: n = shape(a, 1)
+     integer intent(hide), depend(b) :: nrhs = shape(b, 1)
+     integer intent(hide), depend(a) :: lda = MAX(1, shape(a, 0))
+     integer intent(hide), depend(b,n),check(ldb >= n) :: ldb = MAX(1, shape(b, 0))
+     integer intent(out):: info
+
+   end subroutine <prefix2c>hetrs
+   
+   
    subroutine <prefix>sysv(n,nrhs,a,lda,ipiv,b,ldb,work,lwork,info,lower)
 
    ! Solve A * X = B for symmetric A matrix

scipy/linalg/lapack.py

@@ -345,6 +345,9 @@
    chetrf_lwork
    zhetrf_lwork
 
+   chetrs
+   zhetrs
+
    chfrk
    zhfrk
 
@@ -660,6 +663,11 @@
    csytrf_lwork
    zsytrf_lwork
 
+   ssytrs
+   dsytrs
+   csytrs
+   zsytrs
+
    stbtrs
    dtbtrs
    ctbtrs

scipy/linalg/tests/test_lapack.py

@@ -3413,3 +3413,45 @@ def test_tgsyl(dtype, trans, ijob):
                         err_msg='lhs1 and rhs1 do not match')
         assert_allclose(lhs2, rhs2, atol=atol, rtol=0.,
                         err_msg='lhs2 and rhs2 do not match')
+
+
+@pytest.mark.parametrize('dtype', DTYPES)
+@pytest.mark.parametrize('lower', (0, 1))
+def test_sytrs(dtype, lower):
+    rng = np.random.default_rng(1723059677121834)
+    n, nrhs = 20, 5
+    if dtype in COMPLEX_DTYPES:
+        A = (rng.uniform(size=(n, n)) + rng.uniform(size=(n, n))*1j).astype(dtype)
+    else:
+        A = rng.uniform(size=(n, n)).astype(dtype)
+
+    A = A + A.T
+    b = rng.uniform(size=(n, nrhs)).astype(dtype)
+    sytrf, sytrf_lwork, sytrs = get_lapack_funcs(['sytrf', 'sytrf_lwork', 'sytrs'],
+                                                 dtype=dtype)
+    lwork = sytrf_lwork(n, lower=lower)
+    ldu, ipiv, info = sytrf(A, lwork=lwork)
+    assert info == 0
+    x, info = sytrs(a=ldu, ipiv=ipiv, b=b)
+    assert info == 0
+    eps = np.finfo(dtype).eps
+    assert_allclose(A@x, b, atol=100*n*eps)
+
+
+@pytest.mark.parametrize('dtype', COMPLEX_DTYPES)
+@pytest.mark.parametrize('lower', (0, 1))
+def test_hetrs(dtype, lower):
+    rng = np.random.default_rng(1723059677121834)
+    n, nrhs = 20, 5
+    A = (rng.uniform(size=(n, n)) + rng.uniform(size=(n, n))*1j).astype(dtype)
+    A = A + A.conj().T
+    b = np.random.rand(n, nrhs).astype(dtype)
+    hetrf, hetrf_lwork, hetrs = get_lapack_funcs(['hetrf', 'hetrf_lwork', 'hetrs'],
+                                                 dtype=dtype)
+    lwork = hetrf_lwork(n, lower=lower)
+    ldu, ipiv, info = hetrf(A, lwork=lwork)
+    assert info == 0
+    x, info = hetrs(a=ldu, ipiv=ipiv, b=b)
+    assert info == 0
+    eps = np.finfo(dtype).eps
+    assert_allclose(A @ x, b, atol=100*n*eps)