# SPDX-License-Identifier: BSD-3-Clause
# Copyright (c) 2026 Antoine COLLET
"""Wrap the sparse cholesky factorization from"""
import warnings
from typing import Tuple, TypeVar, overload
import numpy as np
import scipy as sp
from covmats._types import ArrayLike, NDArrayFloat, NDArrayInt
# Define a type variable for the input/output type
T = TypeVar("T", NDArrayFloat, sp.sparse.csc_array)
_fact_ds = (
r":math:`\mathbf{LDL}^{\mathrm{T}} = \mathbf{PAP}^{\mathrm{T}}` factorization"
)
[docs]
class SparseCholeskyFactor(sp.sparse.linalg.LinearOperator):
r"""
Sparse Cholesky factor of a matrix :math:`\mathbf{A}`.
.. _factref:
.. admonition:: Factorization
It relies on the :math:`\mathbf{LDL}^{\mathrm{T}} =`
:math:`\mathbf{PAP}^{\mathrm{T}}` factorization where
- :math:`\mathbf{L}` is the sparse lower triangle with unit diagonal.
- :math:`\mathbf{D}` is the sparse diagonal matrix with diagonal entries.
- :math:`\mathbf{P}` is the rows permutation vector.
Note
----
Since this code is released under BSD 3-Clause License, it is not possible to
include sparse factors provided by
`scikit-sparse <https://github.com/scikit-sparse/scikit-sparse>`_,
as it depends on external libraries with GPL licenses, such as
`SuiteSparse <https://github.com/DrTimothyAldenDavis/SuiteSparse?tab=License-1-ov-file>`_.
The main goal of this implementation is therefore to mimic useful features while
preserving the BSD 3-Clause License compatibility.
Examples
--------
>>> from your_module import SparseCholeskyFactor
>>> L, D, P = ... # sparse matrices
>>> factor = SparseCholeskyFactor(L, D, P)
>>> factor.matvec(x)
array([...])
"""
__slots__ = ["_P", "_Pt", "_D", "_invD", "_sqrtD", "_sqrtinvD", "_L", "_n_pts"]
[docs]
def __init__(
self, L: sp.sparse.sparray, D: sp.sparse.sparray, P: ArrayLike
) -> None:
r"""
Initialize the instance.
Parameters
----------
L : sp.sparse.sparray
Sparse lower triangle :math:`\mathbf{L}` of
the :ref:`factorization <factref>`.
D : sp.sparse.sparray
Sparse diagonal matrix $\mathbf{{D}}$ of the :ref:`factorization <factref>`.
P : ArrayLike
1D vector $\mathbf{{P}}$ storing rows permutations of the
:ref:`factorization <factref>`.
Raises
------
ValueError
If the dimensions of `L`, `D` and `P` do not match.
"""
# Make sure that is it 1D and integer
self.P = np.asarray(P, dtype=np.int64).ravel()
self.D = D
self.L = L
# Assert shapes are ok.
if self.D.shape != self.shape or self.P.size != self.n_pts:
raise ValueError(
f"`L` has shape {self.shape},"
f" `D` has shape {D.shape}, and `L` has shape {self.P.shape}!\n"
f"To be consistent, `D` and `P` are expected with shape {self.shape}"
f" and ({self.n_pts},) respectively."
)
@staticmethod
def _validate_lower_triangle(A: sp.sparse.sparray, name: str) -> sp.sparse.sparray:
"""Assert that the input matrix A is lower triangular."""
try:
assert_allclose_sparse(A, sp.sparse.tril(A))
except AssertionError as e:
message = f"The input `{name}` must be a lower-triangular sparse array."
raise ValueError(message) from e
return A
@property
def L(self) -> sp.sparse.csc_array:
r"""
Sparse lower triangle :math:`\mathbf{L}` with unit diagonal of
the :ref:`factorization <factref>`.
"""
return self._L
@L.setter
def L(self, L: sp.sparse.csc_array) -> None:
r"""
Set the sparse lower triangle of
the :ref:`factorization <factref>`
"""
self._L = self._validate_lower_triangle(L, "L")
# Extract the diagonal
try:
np.testing.assert_allclose(L.diagonal(), np.ones(L.shape[0]))
except AssertionError as e:
raise ValueError("The diagonal elements of `L` are assumed to be 1!") from e
@property
def D(self) -> sp.sparse.csc_array:
r"""
Sparse diagonal matrix :math:`\mathbf{D}` of
the :ref:`factorization <factref>`.
"""
return self._D
@D.setter
def D(self, D: sp.sparse.csc_array) -> None:
r"""
Set the sparse diagonal matrix :math:`\mathbf{D}` of
the :ref:`factorization <factref>`.
"""
nnv = np.count_nonzero(D.diagonal() < 0.0)
if nnv != 0:
raise ValueError(
f"{nnv} negative values have been detected in `D` which "
"means the factorized matrix is indefinite and non invertible."
)
nnv = np.count_nonzero(D.diagonal() == 0.0)
if nnv != 0:
raise ValueError(
f"{nnv} null values have been detected in `D` which "
"means the factorized matrix is singular and non invertible."
)
self._D = D
# Update the pseudo inverse as well
self._invD = sp.sparse.diags(1.0 / self.D.diagonal())
# Store the square roots for whitening and colorizing transformations
self._sqrtD = np.sqrt(self.D)
self._sqrtinvD = np.sqrt(self.invD)
@property
def invD(self) -> sp.sparse.csc_array:
r"""
Inverse of :math:`\mathbf{D}` in
the :ref:`factorization <factref>`.
Notes
-----
The property is read-only and updated when setting `D`.
"""
return self._invD
@property
def sqrtinvD(self) -> sp.sparse.csc_array:
r"""
Squared root inverse of :math:`\mathbf{D}` in
the :ref:`factorization <factref>`.
Notes
-----
The property is read-only and updated when setting `D`.
"""
return self._sqrtinvD
@property
def sqrtD(self) -> sp.sparse.csc_array:
r"""
Squared root of :math:`\mathbf{D}` in
the :ref:`factorization <factref>`.
Notes
-----
The property is read-only and updated when setting `D`.
"""
return self._sqrtD
@property
def P(self) -> NDArrayInt:
r"""
1D vector :math:`\mathbf{P}` storing rows permutations of the
:ref:`factorization <factref>`.
"""
return self._P
@P.setter
def P(self, P: ArrayLike) -> None:
r"""
Set the 1D vector :math:`\mathbf{P}` storing rows permutations of the
:ref:`factorization <factref>`.
"""
self._P = np.asarray(P, dtype=np.int64).ravel()
# Update the back pivot
self._Pt: NDArrayInt = np.argsort(self._P)
@property
def Pt(self) -> NDArrayInt:
r"""
1D vector :math:`\mathbf{P}` storing rows back-permutations of the
:ref:`factorization <factref>`.
"""
return self._Pt
@overload
def apply_P(self, x: NDArrayFloat) -> NDArrayFloat: ...
@overload
def apply_P(self, x: sp.sparse.csc_array) -> sp.sparse.csc_array: ...
[docs]
def apply_P(self, x: T) -> T:
"""Apply the permutation of rows."""
return x[self.P]
@overload
def apply_Pt(self, x: NDArrayFloat) -> NDArrayFloat: ...
@overload
def apply_Pt(self, x: sp.sparse.csc_array) -> sp.sparse.csc_array: ...
[docs]
def apply_Pt(self, x: T) -> T:
"""Apply the reverse permutation of rows."""
return x[self.Pt]
[docs]
def solve(self, b: NDArrayFloat) -> NDArrayFloat:
r"""
Return :math:`\mathbf{x} = \mathbf{A}^{-1} \mathbf{b}.`
# L @ D @ L' = P @ A @ P'
# A x = b
# P' L @ D @ L' P x = b
# x = P' L^-T @ D^{-1} @ L^{-1} P b
Parameters
----------
b: NDArrayFloat
Column vector with shape (n,) or ensemble matrix with shape (n, ne).
Returns
-------
NDArrayFloat
_description_
"""
return self.apply_Pt(
sp.sparse.linalg.spsolve_triangular(
self.L.T,
self.invD
@ sp.sparse.linalg.spsolve_triangular(
self.L, self.apply_P(b), lower=True, unit_diagonal=True
),
lower=False,
unit_diagonal=True,
)
)
def _matvec(self, x: NDArrayFloat) -> NDArrayFloat:
"""Return A @ x."""
L = self.apply_Pt(self.L)
return L @ (self.D @ (L.T @ x))
@property
def log_pdet(self) -> float:
"""
Log of the pseudo-determinant of the covariance matrix.
Note
----
In linear algebra and statistics, the pseudo-determinant [1] is the product of
all non-zero eigenvalues of a square matrix. It coincides with the regular
determinant when the matrix is non-singular.
"""
with warnings.catch_warnings():
warnings.simplefilter("ignore", category=RuntimeWarning)
return np.log(np.prod(self.D.diagonal()))
@property
def n_pts(self) -> int:
"""Number of points in the domain (n). This is read-only."""
return self.L.shape[0]
@property
def shape(self) -> Tuple[int, int]:
"""
Shape of the square matrix (n, n). Read-only.
"""
return self.L.shape[:2]
@property
def mat(self) -> sp.sparse.csc_array:
r"""
Return :math:`\mathbf{A}` as a dense array.
Note
----
Alias for :py:meth:`SparseCholeskyFactor.todense()`. This is read-only.
"""
# Permute and rebuild
L = self.apply_Pt(self.L)
return L @ self.D @ L.T
[docs]
def todense(self) -> NDArrayFloat:
r"""Return :math:`\mathbf{A}` as a dense array."""
# Permute and rebuild
return self.mat.toarray()
[docs]
def inv(self) -> NDArrayFloat:
r"""Return the inverse :math:`\mathbf{A}^{-1}` as a dense array."""
return self.solve(np.eye(self.n_pts))
[docs]
def get_diagonal(self) -> NDArrayFloat:
r"""
Return the diagonal entries of the matrix :math:`\mathbf{A}`.
It as shape (n,).
"""
return (self.L @ self.D @ self.L.T).diagonal()
[docs]
def get_invdiagonal(self) -> NDArrayFloat:
r"""
Return the diagonal entries of the inverse matrix :math:`\mathbf{A}^{-1}`.
It as shape (n,).
"""
cov_mat_diag = np.zeros(self.n_pts)
v = np.zeros(self.n_pts)
for i in range(self.n_pts):
v[i - 1] = 0.0
v[i] = 1.0
cov_mat_diag[i] = self.solve(v)[i]
return cov_mat_diag
[docs]
def colorize(self, x: NDArrayFloat) -> NDArrayFloat:
"""
Perform a colorizing transformation on data.
"Colorizing" ("color" as in "colored noise", in which different
frequencies may have different magnitudes) transforms a set of
uncorrelated random variables into a new set of random variables with
the desired covariance. When a coloring transform is applied to a
sample of points distributed according to a multivariate normal
distribution with identity covariance and zero mean, the covariance of
the transformed sample is approximately the covariance matrix used
in the coloring transform
:cite:p:`wikipediaWhiteningTransformation2025,novakGeneralizationColoringLinear2019`.
Parameters
----------
x : array_like
An array of points. The last dimension must correspond with the
dimensionality of the space, i.e., the number of columns in the
covariance matrix.
Returns
-------
x_ : array_like
The transformed array of points.
Note
----
We want to solve z.T = x @ K.T, where A = K @ K^{T}
We use the cholesky factorization LDL' = PA'AP'
with P' = P^{-1} the permutation that makes the decomposition unique.
So LD^{1/2} = PA' and A = D^{1/2}L'P
Finally z = (P' L D^{1/2} x')'
References
----------
.. bibliography::
:list: enumerated
:filter: False
wikipediaWhiteningTransformation2025
novakGeneralizationColoringLinear2019
Examples
--------
>>> import numpy as np
>>> import covmats
>>> rng = np.random.default_rng(1638083107694713882823079058616272161)
>>> n = 3
>>> A = rng.random(size=(n, n))
>>> cov_array = A @ A.T # make matrix symmetric positive definite
>>> cholesky = np.linalg.cholesky(cov_array)
>>> cov_object = covmats.CovViaCholesky(cholesky)
>>> x = rng.multivariate_normal(np.zeros(n), np.eye(n), size=(10000))
>>> x_ = cov_object.colorize(x)
>>> cov_data = np.cov(x_, rowvar=False)
>>> np.allclose(cov_data, cov_array, rtol=3e-2)
True
"""
return self.apply_Pt(self.L @ self.sqrtD @ x.T).T
[docs]
def colorize_inv(self, x: NDArrayFloat) -> NDArrayFloat:
r"""
Perform a colorizing transformation on data as in
:py:meth:`SparseCholeskyFactor:colorize` but for the inverse
:math:`\mathbf{A}^{-1}`.
"""
return self.apply_Pt(
sp.sparse.linalg.spsolve_triangular(
self.L.T, self.sqrtinvD @ x.T, unit_diagonal=True, lower=False
)
).T
[docs]
def whiten(self, x: NDArrayFloat) -> NDArrayFloat:
"""
Perform a whitening transformation on data.
"Whitening" ("white" as in "white noise", in which each frequency has
equal magnitude) transforms a set of random variables into a new set of
random variables with unit-diagonal covariance. When a whitening
transform is applied to a sample of points distributed according to
a multivariate normal distribution with zero mean, the covariance of
the transformed sample is approximately the identity matrix
:cite:p:`wikipediaWhiteningTransformation2025,
novakGeneralizationColoringLinear2019`.
Parameters
----------
x : array_like
An array of points. The last dimension must correspond with the
dimensionality of the space, i.e., the number of columns in the
covariance matrix.
Returns
-------
x_ : array_like
The transformed array of points.
Note
----
TODO.
We use the cholesky factorization LDL' = PA'AP'
We want to solve x = z @ A.T => z = x A^{-T}
with P' = P^{-1} the permutation that makes the decomposition unique.
So LD^{1/2} = PA' and A = D^{1/2}L'P
Finally z = (D^{-1/2} L^{-1} P x')'
References
----------
.. bibliography::
:list: enumerated
:filter: False
novakGeneralizationColoringLinear2019
wikipediaWhiteningTransformation2025
Examples
--------
>>> import numpy as np
>>> import covmats
>>> rng = np.random.default_rng()
>>> n = 3
>>> A = rng.random(size=(n, n))
>>> cov_array = A @ A.T # make matrix symmetric positive definite
>>> precision = np.linalg.inv(cov_array)
>>> cov_object = covmats.CovViaPrecisionCholesky(
.. sp.linalg.cholesky(precision, lower=True))
>>> x = rng.multivariate_normal(np.zeros(n), cov_array, size=(10000))
>>> x_ = cov_object.whiten(x)
>>> np.cov(x_, rowvar=False) # near-identity covariance
array([[0.97862122, 0.00893147, 0.02430451],
[0.00893147, 0.96719062, 0.02201312],
[0.02430451, 0.02201312, 0.99206881]])
"""
return (
self.sqrtinvD
[docs]
@ sp.sparse.linalg.spsolve_triangular(
self.L, self.apply_P(x.T), unit_diagonal=True, lower=True
)
).T
def whiten_inv(self, x: NDArrayFloat) -> NDArrayFloat:
r"""
Perform a whitening transformation on data as in
:py:meth:`SparseCholeskyFactor:whiten` but for the inverse
:math:`\mathbf{A}^{-1}`.
"""
return self.sqrtD @ (self.L.T @ self.apply_P(x.T)).T
def assert_allclose_sparse(
A: sp.sparse.sparray, B: sp.sparse.sparray, atol: float = 1e-8, rtol: float = 1e-8
) -> None:
"""
Assert that two sparse matrices or arrays are almost equal.
Parameters
----------
A : sp.sparse.sparray
First sparse matrix.
B : sp.sparse.sparray
Second sparse matrix.
atol : float, optional
Absolute tolerance for the equality test, by default 1e-8.
rtol : float, optional
Relative tolerance for the equality test, by default 1e-8.
"""
# If you want to check matrix shapes as well
assert np.array_equal(A.shape, B.shape)
r1, c1, v1 = sp.sparse.find(A)
r2, c2, v2 = sp.sparse.find(B)
np.testing.assert_equal(r1, r2)
np.testing.assert_equal(c1, c2)
np.testing.assert_allclose(v1, v2, atol=atol, rtol=rtol)
[docs]
def get_SPD_sparse_n11_example(seed: int) -> sp.sparse.csc_array:
"""
Create a symmetric positive definite matrix of shape(11, 11).
Parameters
----------
seed: int
Random seed.
"""
N = 11
rows = np.array([5, 6, 2, 7, 9, 10, 5, 9, 7, 10, 8, 9, 10, 9, 10, 10])
cols = np.array([0, 0, 1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 6, 7, 7, 9])
rng = np.random.default_rng(seed)
vals = rng.random(len(rows), dtype=np.float64)
L = sp.sparse.coo_array((vals, (rows, cols)), shape=(N, N))
A = L + L.T # make it symmetric
A.setdiag(2.0) # make it strongly positive definite
A = A.tocsc()
return A
[docs]
def get_SPD_sparse_example(
n: int, seed: int, diag_mean: float = 10.0
) -> sp.sparse.csc_array:
"""
Create a symmetric positive definite matrix.
Parameters
----------
n : int
Number of elements.
seed : int
Random seed.
Returns
-------
sp.sparse.cscarray
CSC sparse SPD.
"""
# Initialize a sparse matrix
A = sp.sparse.lil_matrix((n, n))
# Fill the diagonal with variances (non-zero values)
np.random.seed(42) # for reproducibility
variances = np.random.uniform(-1.5, 1.5, n) + diag_mean
for i in range(n):
A[i, i] = variances[i]
# Add some random non-zero covariances
for i in range(n):
for j in range(i + 1, n):
if np.random.rand() < 0.1: # 10% chance of non-zero covariance
cov = np.random.uniform(-0.5, 0.5)
A[i, j] = cov
A[j, i] = cov # symmetry
return A.tocsc()