#!/usr/bin/env python
# -*- coding: utf-8 -*-
import numpy as np
import scipy.linalg
[docs]def chol_solve(A, b):
L = np.linalg.cholesky(A)
x = scipy.linalg.solve_triangular(
L,
scipy.linalg.solve_triangular(
L, b,
lower=True,
check_finite=False),
trans='C',
lower=True,
check_finite=False
)
return x
[docs]class CholSolver:
def __init__(self, A):
self.L = np.linalg.cholesky(A)
def __matmul__(self, b):
return scipy.linalg.solve_triangular(
self.L,
scipy.linalg.solve_triangular(
self.L, b,
lower=True,
check_finite=False),
trans='C',
lower=True,
check_finite=False
)
[docs] def todense(self):
return self @ np.eye(*self.L.shape)
[docs] def diagonal(self):
return self.todense().diagonal()