def test_posterior_mean_CG_equivalency(self):
"""The probabilistic linear solver(s) should recover CG iterates as a posterior
mean for specific covariances."""
# Linear system
A, b = self.poisson_linear_system
# Callback function to return CG iterates
cg_iterates = []
def callback_iterates_CG(xk):
cg_iterates.append(
np.eye(np.shape(A)[0]) @ xk
) # identity hack to actually save different iterations
# Solve linear system
# Initial guess as chosen by PLS: x0 = Ainv.mean @ b
x0 = b
# Conjugate gradient method
xhat_cg, info_cg = scipy.sparse.linalg.cg(
A=A, b=b, x0=x0, tol=10 ** -6, callback=callback_iterates_CG
)
cg_iters_arr = np.array([x0] + cg_iterates)
# Matrix priors (encoding weak symmetric posterior correspondence)
Ainv0 = rvs.Normal(
mean=linops.Identity(A.shape[1]),
cov=linops.SymmetricKronecker(A=linops.Identity(A.shape[1])),
)
A0 = rvs.Normal(
mean=linops.Identity(A.shape[1]),
cov=linops.SymmetricKronecker(A),
)
for kwargs in [{"assume_A": "sympos", "rtol": 10 ** -6}]:
with self.subTest():
# Define callback function to obtain search directions
pls_iterates = []
# pylint: disable=cell-var-from-loop
def callback_iterates_PLS(
xk, Ak, Ainvk, sk, yk, alphak, resid, **kwargs
):
pls_iterates.append(xk.mean)
# Probabilistic linear solver
xhat_pls, _, _, info_pls = linalg.problinsolve(
A=A,
b=b,
Ainv0=Ainv0,
A0=A0,
callback=callback_iterates_PLS,
**kwargs
)
pls_iters_arr = np.array([x0] + pls_iterates)
self.assertAllClose(xhat_pls.mean, xhat_cg, rtol=10 ** -12)
self.assertAllClose(pls_iters_arr, cg_iters_arr, rtol=10 ** -12)
def test_symmkronecker_transpose(self):
"""Kronecker product transpose property: (A (x) B)^T = A^T (x) B^T."""
for A, B in self.symmkronecker_matrices:
with self.subTest():
W = linops.SymmetricKronecker(A=A, B=B)
V = linops.SymmetricKronecker(A=A.T, B=B.T)
self.assertAllClose(W.T.todense(), V.todense())
def test_symmkronecker_commutation(self):
"""Symmetric Kronecker products fulfill A (x)_s B = B (x)_s A"""
for A, B in self.symmkronecker_matrices:
with self.subTest():
W = linops.SymmetricKronecker(A=A, B=B)
V = linops.SymmetricKronecker(A=B, B=A)
self.assertAllClose(W.todense(), V.todense())
def setUp(self):
"""Resources for tests."""
# Seed
np.random.seed(seed=42)
# Parameters
m = 7
n = 3
self.constants = [-1, -2.4, 0, 200, np.pi]
sparsemat = scipy.sparse.rand(m=m, n=n, density=0.1, random_state=1)
self.normal_params = [
# Univariate
(-1.0, 3.0),
(1, 3),
# Multivariate
(np.random.uniform(size=10), np.eye(10)),
(np.random.uniform(size=10), random_spd_matrix(10)),
# Matrixvariate
(
np.random.uniform(size=(2, 2)),
linops.SymmetricKronecker(
A=np.array([[1.0, 2.0], [2.0, 1.0]]),
B=np.array([[5.0, -1.0], [-1.0, 10.0]]),
).todense(),
),
# Operatorvariate
(
np.array([1.0, -5.0]),
linops.Matrix(A=np.array([[2.0, 1.0], [1.0, -0.1]])),
),
(
linops.Matrix(A=np.array([[0.0, -5.0]])),
linops.Identity(shape=(2, 2)),
),
(
np.array([[1.0, 2.0], [-3.0, -0.4], [4.0, 1.0]]),
linops.Kronecker(A=np.eye(3), B=5 * np.eye(2)),
),
(
linops.Matrix(A=sparsemat.todense()),
linops.Kronecker(0.1 * linops.Identity(m), linops.Identity(n)),
),
(
linops.Matrix(A=np.random.uniform(size=(2, 2))),
linops.SymmetricKronecker(
A=np.array([[1.0, 2.0], [2.0, 1.0]]),
B=np.array([[5.0, -1.0], [-1.0, 10.0]]),
),
),
# Symmetric Kronecker Identical Factors
(
linops.Identity(shape=25),
linops.SymmetricKronecker(A=linops.Identity(25)),
),
]
def get_randvar(rv_name):
"""Return a random variable for a given distribution name."""
# Distribution Means and Covariances
mean_0d = np.random.rand()
mean_1d = np.random.rand(5)
mean_2d_mat = SPD_MATRIX_5x5
mean_2d_linop = linops.MatrixMult(SPD_MATRIX_5x5)
cov_0d = np.random.rand() + 10**-12
cov_1d = SPD_MATRIX_5x5
cov_2d_kron = linops.Kronecker(A=SPD_MATRIX_5x5, B=SPD_MATRIX_5x5)
cov_2d_symkron = linops.SymmetricKronecker(A=SPD_MATRIX_5x5)
if rv_name == "univar_normal":
randvar = rvs.Normal(mean=mean_0d, cov=cov_0d)
elif rv_name == "multivar_normal":
randvar = rvs.Normal(mean=mean_1d, cov=cov_1d)
elif rv_name == "matrixvar_normal":
randvar = rvs.Normal(mean=mean_2d_mat, cov=cov_2d_kron)
elif rv_name == "symmatrixvar_normal":
randvar = rvs.Normal(mean=mean_2d_mat, cov=cov_2d_symkron)
elif rv_name == "operatorvar_normal":
randvar = rvs.Normal(mean=mean_2d_linop, cov=cov_2d_symkron)
else:
raise ValueError("Random variable not found.")
return randvar
def setUp(self):
"""Resources for tests."""
# Random Seed
np.random.seed(42)
# Scalars, arrays and operators
self.scalars = [0, int(1), 0.1, -4.2, np.nan, np.inf]
self.arrays = [np.random.normal(size=[5, 4]), np.array([[3, 4], [1, 5]])]
def mv(v):
return np.array([2 * v[0], v[0] + 3 * v[1]])
self.mv = mv
self.ops = [
linops.MatrixMult(np.array([[-1.5, 3], [0, -230]])),
linops.LinearOperator(shape=(2, 2), matvec=mv),
linops.Identity(shape=4),
linops.Kronecker(
A=linops.MatrixMult(np.array([[2, -3.5], [12, 6.5]])),
B=linops.Identity(shape=3),
),
linops.SymmetricKronecker(
A=linops.MatrixMult(np.array([[1, -2], [-2.2, 5]])),
B=linops.MatrixMult(np.array([[1, -3], [0, -0.5]])),
),
]
def _symmetric_matrix_based_update(self, matrix: randvars.Normal,
action: np.ndarray,
observ: np.ndarray) -> randvars.Normal:
"""Symmetric matrix-based inference update for linear information."""
if not isinstance(matrix.cov, linops.SymmetricKronecker):
raise ValueError(
f"Covariance must have symmetric Kronecker structure, but is '{type(matrix.cov).__name__}'."
)
pred = matrix.mean @ action
resid = observ - pred
covfactor_Ms = matrix.cov.A @ action
gram = action.T @ covfactor_Ms
gram_pinv = 1.0 / gram if gram > 0.0 else 0.0
gain = covfactor_Ms * gram_pinv
covfactor_update = linops.aslinop(gain[:, None]) @ linops.aslinop(
covfactor_Ms[None, :])
resid_gain = linops.aslinop(resid[:, None]) @ linops.aslinop(
gain[None, :])
return randvars.Normal(
mean=matrix.mean + resid_gain + resid_gain.T -
linops.aslinop(gain[:, None]) @ linops.aslinop(
(action.T @ resid_gain)[None, :]),
cov=linops.SymmetricKronecker(A=matrix.cov.A - covfactor_update),
)
def test_matrixprior(self):
"""Solve random linear system with a matrix-based linear solver."""
np.random.seed(1)
# Linear system
n = 10
A = np.random.rand(n, n)
A = A.dot(
A.T) + n * np.eye(n) # Symmetrize and make diagonally dominant
x_true = np.random.normal(size=(n, ))
b = A @ x_true
# Prior distributions on A
covA = linops.SymmetricKronecker(A=np.eye(n))
Ainv0 = rvs.Normal(mean=np.eye(n), cov=covA)
for matblinsolve in self.matblinsolvers:
with self.subTest():
x, Ahat, Ainvhat, info = matblinsolve(A=A, Ainv0=Ainv0, b=b)
self.assertAllClose(
x.mean,
x_true,
rtol=1e-6,
atol=1e-6,
msg=
"Solution for matrixvariate prior does not match true solution.",
)
def setUp(self) -> None:
"""Scalars, arrays, linear operators and random variables for tests."""
# Seed
np.random.seed(42)
# Random variable instantiation
self.scalars = [0, int(1), 0.1, np.nan, np.inf]
self.arrays = [np.empty(2), np.zeros(4), np.array([]), np.array([1, 2])]
# Random variable arithmetic
self.arrays2d = [
np.empty(2),
np.zeros(2),
np.array([np.inf, 1]),
np.array([1, -2.5]),
]
self.matrices2d = [np.array([[1, 2], [3, 2]]), np.array([[0, 0], [1.0, -4.3]])]
self.linops2d = [linops.Matrix(A=np.array([[1, 2], [4, 5]]))]
self.randvars2d = [
randvars.Normal(mean=np.array([1, 2]), cov=np.array([[2, 0], [0, 5]]))
]
self.randvars2x2 = [
randvars.Normal(
mean=np.array([[-2, 0.3], [0, 1]]),
cov=linops.SymmetricKronecker(A=np.eye(2), B=np.ones((2, 2))),
),
]
self.scipyrvs = [
scipy.stats.bernoulli(0.75),
scipy.stats.norm(4, 2.4),
scipy.stats.multivariate_normal(np.random.randn(10), np.eye(10)),
scipy.stats.gamma(0.74),
scipy.stats.dirichlet(alpha=np.array([0.1, 0.1, 0.2, 0.3])),
]
def symmetric_matrixvariate_normal(
shape: ShapeArgType, precompute_cov_cholesky: bool, rng: np.random.Generator
) -> randvars.Normal:
rv = randvars.Normal(
mean=random_spd_matrix(dim=shape[0], rng=rng),
cov=linops.SymmetricKronecker(A=random_spd_matrix(dim=shape[0], rng=rng)),
)
if precompute_cov_cholesky:
rv.precompute_cov_cholesky()
return rv
def case_state_symmetric_matrix_based(rng: np.random.Generator, ):
"""State of a symmetric matrix-based linear solver."""
prior = linalg.solvers.beliefs.LinearSystemBelief(
A=randvars.Normal(
mean=linops.Matrix(linsys.A),
cov=linops.SymmetricKronecker(A=linops.Identity(n)),
),
x=(Ainv @ b[:, None]).reshape((n, )),
Ainv=randvars.Normal(
mean=linops.Identity(n),
cov=linops.SymmetricKronecker(A=linops.Identity(n)),
),
b=b,
)
state = linalg.solvers.LinearSolverState(problem=linsys, prior=prior)
state.action = rng.standard_normal(size=state.problem.A.shape[1])
state.observation = rng.standard_normal(size=state.problem.A.shape[1])
return state
def test_symmkronecker_todense_symmetric(self):
"""Dense matrix from symmetric Kronecker product of two symmetric matrices must be symmetric."""
C = np.array([[5, 1], [1, 10]])
D = np.array([[-2, 0.1], [0.1, 8]])
Ws = linops.SymmetricKronecker(A=C, B=C)
Ws_dense = Ws.todense()
self.assertArrayEqual(
Ws_dense,
Ws_dense.T,
msg=
"Symmetric Kronecker product of symmetric matrices is not symmetric.",
)
def _symmetric_kronecker_identical_factors_cov_cholesky(
self,
damping_factor: Optional[FloatArgType] = COV_CHOLESKY_DAMPING,
) -> linops.SymmetricKronecker:
assert (isinstance(self._cov, linops.SymmetricKronecker)
and self._cov.identical_factors)
A = self._cov.A.todense()
return linops.SymmetricKronecker(A=scipy.linalg.cholesky(
A + damping_factor * np.eye(A.shape[0], dtype=self.dtype),
lower=True,
), )
def _symmetric_kronecker_identical_factors_cov_cholesky(
self,
) -> linops.SymmetricKronecker:
assert isinstance(self._cov, linops.SymmetricKronecker) and self._cov._ABequal
A = self._cov.A.todense()
return linops.SymmetricKronecker(
A=scipy.linalg.cholesky(
A + COV_CHOLESKY_DAMPING * np.eye(A.shape[0], dtype=self.dtype),
lower=True,
),
dtype=self.dtype,
)
def test_induced_solution_belief(rng: np.random.Generator):
"""Test whether a consistent belief over the solution is inferred from a belief over
the inverse."""
n = 5
A = randvars.Constant(random_spd_matrix(dim=n, rng=rng))
Ainv = randvars.Normal(
mean=linops.Scaling(factors=1 / np.diag(A.mean)),
cov=linops.SymmetricKronecker(linops.Identity(n)),
)
b = randvars.Constant(rng.normal(size=(n, 1)))
prior = LinearSystemBelief(A=A, Ainv=Ainv, x=None, b=b)
x_infer = Ainv @ b
np.testing.assert_allclose(prior.x.mean, x_infer.mean)
np.testing.assert_allclose(prior.x.cov.todense(), x_infer.cov.todense())
def dense_matrix_based_update(matrix: randvars.Normal, action: np.ndarray,
observ: np.ndarray):
pred = matrix.mean @ action
resid = observ - pred
covfactor_Ms = matrix.cov.A @ action
gram = action.T @ covfactor_Ms
gram_pinv = 1.0 / gram if gram > 0.0 else 0.0
gain = covfactor_Ms * gram_pinv
covfactor_update = gain @ covfactor_Ms.T
resid_gain = np.outer(resid, gain)
return randvars.Normal(
mean=matrix.mean + resid_gain + resid_gain.T -
np.outer(gain, action.T @ resid_gain),
cov=linops.SymmetricKronecker(A=matrix.cov.A - covfactor_update),
)
def test_symmetric_samples(self):
"""Samples from a normal distribution with symmetric Kronecker kernels of two
symmetric matrices are symmetric."""
n = 3
A = self.rng.uniform(size=(n, n))
A = 0.5 * (A + A.T) + n * np.eye(n)
rv = randvars.Normal(
mean=np.eye(A.shape[0]),
cov=linops.SymmetricKronecker(A=A),
)
rv = rv.sample(rng=self.rng, size=10)
for i, B in enumerate(rv):
self.assertAllClose(
B,
B.T,
atol=1e-5,
rtol=1e-5,
msg=
"Sample {} from symmetric Kronecker distribution is not symmetric."
.format(i),
)
import numpy as np
from pytest_cases import case
from probnum import linalg, linops, randvars
from probnum.problems.zoo.linalg import random_linear_system, random_spd_matrix
# Problem
n = 10
linsys = random_linear_system(
rng=np.random.default_rng(42), matrix=random_spd_matrix, dim=n
)
# Prior
Ainv = randvars.Normal(
mean=linops.Identity(n), cov=linops.SymmetricKronecker(linops.Identity(n))
)
b = randvars.Constant(linsys.b)
prior = linalg.solvers.beliefs.LinearSystemBelief(
A=randvars.Constant(linsys.A),
Ainv=Ainv,
x=(Ainv @ b[:, None]).reshape(
(n,)
), # TODO: This can be replaced by Ainv @ b once https://github.com/probabilistic-numerics/probnum/issues/456 is fixed
b=b,
)
@case(tags=["initial"])
def case_initial_state(
rng: np.random.Generator,
def _get_output_randvars(self, Y_list, sy_list, phi=None, psi=None):
"""Return output random variables x, A, Ainv from their means and
covariances."""
if self.iter_ > 0:
# Observations and inner products in A-space between actions
Y = np.hstack(Y_list)
sy = np.vstack(sy_list).ravel()
# Posterior covariance factors
if self.is_calib_covclass and (not phi is None) and (
not psi is None):
# Ensure prior covariance class only acts in span(S) like A
@linops.LinearOperator.broadcast_matvec
def _matmul(x):
# First term of calibration covariance class: AS(S'AS)^{-1}S'A
return (Y * sy**-1) @ (Y.T @ x.ravel())
_A_covfactor0 = linops.LinearOperator(
shape=(self.n, self.n),
dtype=np.result_type(Y, sy),
matmul=_matmul,
)
@linops.LinearOperator.broadcast_matvec
def _matmul(x):
# Term in covariance class: A_0^{-1}Y(Y'A_0^{-1}Y)^{-1}Y'A_0^{-1}
# TODO: for efficiency ensure that we dont have to compute
# (Y.T Y)^{-1} two times! For a scalar mean this is the same as in
# the null space projection
YAinv0Y_inv_YAinv0x = np.linalg.solve(
Y.T @ (self.Ainv_mean0 @ Y),
Y.T @ (self.Ainv_mean0 @ x))
return self.Ainv_mean0 @ (Y @ YAinv0Y_inv_YAinv0x)
_Ainv_covfactor0 = linops.LinearOperator(
shape=(self.n, self.n),
dtype=np.result_type(Y, self.Ainv_mean0),
matmul=_matmul,
)
# Set degrees of freedom based on uncertainty calibration in unexplored
# space
(
calibration_term_A,
calibration_term_Ainv,
) = self._get_calibration_covariance_update_terms(phi=phi,
psi=psi)
_A_covfactor = (_A_covfactor0 - self._A_covfactor_update_term +
calibration_term_A)
_Ainv_covfactor = (_Ainv_covfactor0 -
self._Ainv_covfactor_update_term +
calibration_term_Ainv)
else:
# No calibration
_A_covfactor = self.A_covfactor
_Ainv_covfactor = self.Ainv_covfactor
else:
# Converged before making any observations
_A_covfactor = self.A_covfactor0
_Ainv_covfactor = self.Ainv_covfactor0
# Create output random variables
A = randvars.Normal(mean=self.A_mean,
cov=linops.SymmetricKronecker(A=_A_covfactor))
Ainv = randvars.Normal(
mean=self.Ainv_mean,
cov=linops.SymmetricKronecker(A=_Ainv_covfactor),
)
# Induced distribution on x via Ainv
# Exp(x) = Ainv b, Cov(x) = 1/2 (W b'Wb + Wbb'W)
Wb = _Ainv_covfactor @ self.b
bWb = np.squeeze(Wb.T @ self.b)
def _matmul(x):
return 0.5 * (bWb * _Ainv_covfactor @ x + Wb @ (Wb.T @ x))
cov_op = linops.LinearOperator(
shape=(self.n, self.n),
dtype=np.result_type(Wb.dtype, bWb.dtype),
matmul=_matmul,
)
x = randvars.Normal(mean=self.x_mean.ravel(), cov=cov_op)
# Compute trace of solution covariance: tr(Cov(x))
self.trace_sol_cov = np.real_if_close(
self._compute_trace_solution_covariance(bWb=bWb, Wb=Wb)).item()
return x, A, Ainv