numba / numba-scipy

numba_scipy extends Numba to make it aware of SciPy
https://numba-scipy.readthedocs.io/en/latest/
BSD 2-Clause "Simplified" License
250 stars 33 forks source link

triangular solver #91

Open samotto1 opened 3 years ago

samotto1 commented 3 years ago

Feature request

Hi,

I would like to be able to solve symmetric positive-definite linear systems in Numba using Cholesky factorization. While the Cholesky factorization is implemented, there is no triangular solver! It would be wonderful if a triangular solver could be implemented so that I can take advantage of the Cholesky function. Thanks.

Sam

gmarkall commented 3 years ago

Thanks for the request! Do you have an executable example of the function you'd like to work please? Or can you perhaps point towards the particular functions / libraries you'd like to see supported?

stuartarchibald commented 3 years ago

I don't think a back solver is part of NumPy: https://numpy.org/doc/stable/reference/routines.linalg.html, it does exist in SciPy though https://docs.scipy.org/doc/scipy/reference/generated/scipy.linalg.solve_triangular.html#scipy.linalg.solve_triangular, another option is to go via https://docs.scipy.org/doc/scipy/reference/generated/scipy.linalg.cho_factor.html#scipy.linalg.cho_factor and https://docs.scipy.org/doc/scipy/reference/generated/scipy.linalg.cho_solve.html#scipy.linalg.cho_solve, or just https://docs.scipy.org/doc/scipy/reference/generated/scipy.linalg.solve.html#scipy.linalg.solve which will select an optimal solve method if you tell the routine that your system is indeed S.P.D. I don't think this is something Numba should implement as it is out of scope. May be something for https://github.com/numba/numba-scipy, eventually.

If you really want to bind to LAPACK routines to do this from within JIT compiled code then there's an example of doing so in this issue https://github.com/numba/numba/issues/5301, it's not particularly easy but works.

gmarkall commented 3 years ago

Closing this issue as no further info provided recently - please feel free to open with more specific information about the request.

agramfort commented 1 year ago

I have the same request. Here is a gist to replicate the need

import numpy as np
from scipy import linalg
from numba import njit

# Build 3 random symmetric positive matrices
np.random.seed(0)
A = np.random.randn(3, 30, 30)
A = 0.5 * (A @ np.swapaxes(A, -1, -2))
B = np.random.randn(3, 30, 10)
AinvB = np.linalg.solve(A, B)

@njit
def cholesky_solve(A, B):
    # Compute the Cholesky decomposition
    L = np.linalg.cholesky(A)
    AinvB = np.empty((3, 30, 10))
    for i in range(3):
        AinvB[i, :, :] = linalg.cho_solve(
            (L[i], True), B[i]
        )
    return AinvB

np.testing.assert_allclose(cholesky_solve(A, B), np.linalg.solve(A, B))

it would be great to see this working. Thanks

stuartarchibald commented 1 year ago

I transferred this in from numba/numba (was issue 6394) as it's a request for supporting a feature in scipy.

agramfort commented 1 year ago

@gmarkall @stuartarchibald any chance this can be looked into soon? thanks