def _setup(
self,
test_ndim,
spdmat1,
spdmat2,
):
with config(matrix_free=True):
self.noise_fun = lambda t: randvars.Normal(
mean=np.arange(test_ndim), cov=linops.aslinop(spdmat2))
self.transition_matrix_fun = lambda t: linops.aslinop(spdmat1)
self.transition = randprocs.markov.discrete.LinearGaussian(
input_dim=test_ndim,
output_dim=test_ndim,
transition_matrix_fun=self.transition_matrix_fun,
noise_fun=self.noise_fun,
forward_implementation="classic",
backward_implementation="classic",
)
self.sqrt_transition = randprocs.markov.discrete.LinearGaussian(
input_dim=test_ndim,
output_dim=test_ndim,
transition_matrix_fun=self.transition_matrix_fun,
noise_fun=self.noise_fun,
forward_implementation="sqrt",
backward_implementation="sqrt",
)
self.transition_fun = lambda t, x: self.transition_matrix_fun(t
) @ x
self.transition_fun_jacobian = lambda t, x: self.transition_matrix_fun(
t)
def _setup(
self,
test_ndim,
spdmat1,
spdmat2,
):
with config(matrix_free=True):
self.G = lambda t: linops.aslinop(spdmat1)
self.S = lambda t: linops.aslinop(spdmat2)
self.v = lambda t: np.arange(test_ndim)
self.transition = randprocs.markov.discrete.LinearGaussian(
test_ndim,
test_ndim,
self.G,
self.v,
self.S,
forward_implementation="classic",
backward_implementation="classic",
)
self.sqrt_transition = randprocs.markov.discrete.LinearGaussian(
test_ndim,
test_ndim,
self.G,
self.v,
self.S,
forward_implementation="sqrt",
backward_implementation="sqrt",
)
self.g = lambda t, x: self.G(t) @ x + self.v(t)
self.dg = lambda t, x: self.G(t)
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 = gain @ covfactor_Ms.T
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 _matrix_based_update(self, matrix: randvars.Normal, action: np.ndarray,
observ: np.ndarray) -> randvars.Normal:
"""Matrix-based inference update for linear information."""
if not isinstance(matrix.cov, linops.Kronecker):
raise ValueError(
f"Covariance must have Kronecker structure, but is '{type(matrix.cov).__name__}'."
)
pred = matrix.mean @ action
resid = observ - pred
covfactor_Ms = matrix.cov.B @ 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, :]
) # residual and gain are flipped due to matrix vectorization
return randvars.Normal(
mean=matrix.mean + resid_gain,
cov=linops.Kronecker(A=matrix.cov.A,
B=matrix.cov.B - covfactor_update),
)
def test_linop_construction(self):
"""Create linear operators via various construction methods."""
# Custom linear operator
linops.LinearOperator(shape=(2, 2), matvec=self.mv)
# Scipy linear operator
scipy_linop = scipy.sparse.linalg.LinearOperator(shape=(2, 2), matvec=self.mv)
linops.aslinop(scipy_linop)
def equivalent_discretisation_preconditioned(self):
"""Discretised IN THE PRECONDITIONED SPACE.
The preconditioned state transition is the flipped Pascal matrix.
The preconditioned process noise covariance is the flipped Hilbert matrix.
The shift is always zero.
Reference: https://arxiv.org/abs/2012.10106
"""
state_transition_1d = np.flip(
scipy.linalg.pascal(self.num_derivatives + 1, kind="lower", exact=False)
)
if config.matrix_free:
state_transition = linops.Kronecker(
A=linops.Identity(self.wiener_process_dimension),
B=linops.aslinop(state_transition_1d),
)
else:
state_transition = np.kron(
np.eye(self.wiener_process_dimension), state_transition_1d
)
process_noise_1d = np.flip(scipy.linalg.hilbert(self.num_derivatives + 1))
if config.matrix_free:
process_noise = linops.Kronecker(
A=linops.Identity(self.wiener_process_dimension),
B=linops.aslinop(process_noise_1d),
)
else:
process_noise = np.kron(
np.eye(self.wiener_process_dimension), process_noise_1d
)
empty_shift = np.zeros(
self.wiener_process_dimension * (self.num_derivatives + 1)
)
process_noise_cholesky_1d = np.linalg.cholesky(process_noise_1d)
if config.matrix_free:
process_noise_cholesky = linops.Kronecker(
A=linops.Identity(self.wiener_process_dimension),
B=linops.aslinop(process_noise_cholesky_1d),
)
else:
process_noise_cholesky = np.kron(
np.eye(self.wiener_process_dimension), process_noise_cholesky_1d
)
return discrete.LTIGaussian(
state_trans_mat=state_transition,
shift_vec=empty_shift,
proc_noise_cov_mat=process_noise,
proc_noise_cov_cholesky=process_noise_cholesky,
forward_implementation=self.forward_implementation,
backward_implementation=self.backward_implementation,
)
def _derivwise2coordwise_projmat(self) -> np.ndarray:
r"""Projection matrix to change the ordering of the state representation in an
:class:`Integrator` from coordinate-wise to derivative-wise representation.
A coordinate-wise ordering is
.. math:: (y_1, \dot y_1, \ddot y_1, y_2, \dot y_2, ..., y_d^{(\nu)})
and a derivative-wise ordering is
.. math:: (y_1, y_2, ..., y_d, \dot y_1, \dot y_2, ...,
\dot y_d, \ddot y_1, ..., y_d^{(\nu)}).
Default representation in an :class:`Integrator` is coordinate-wise ordering,
but sometimes, derivative-wise ordering is more convenient.
See Also
--------
:attr:`Integrator._convert_coordwise_to_derivwise`
:attr:`Integrator._convert_derivwise_to_coordwise`
"""
dim = (self.num_derivatives + 1) * self.wiener_process_dimension
projmat = np.zeros((dim, dim))
E = np.eye(dim)
for q in range(self.num_derivatives + 1):
projmat[q::(self.num_derivatives +
1)] = E[q * self.wiener_process_dimension:(q + 1) *
self.wiener_process_dimension]
if config.matrix_free:
projmat = linops.aslinop(projmat)
return projmat
def case_state_converged(rng: np.random.Generator, ):
"""State of a linear solver, which has converged at initialization."""
belief = linalg.solvers.beliefs.LinearSystemBelief(
A=randvars.Constant(linsys.A),
Ainv=randvars.Constant(linops.aslinop(linsys.A).inv().todense()),
x=randvars.Constant(linsys.solution),
b=randvars.Constant(linsys.b),
)
state = linalg.solvers.LinearSolverState(problem=linsys, prior=belief)
return state
def test_backward_realization(self, some_normal_rv1, some_normal_rv2):
with config(matrix_free=True):
array_cov_rv = some_normal_rv2
linop_cov_rv = randvars.Normal(array_cov_rv.mean.copy(),
linops.aslinop(array_cov_rv.cov))
with pytest.warns(RuntimeWarning):
self.transition.backward_realization(some_normal_rv1.mean,
array_cov_rv)
out, _ = self.transition.backward_realization(
some_normal_rv1.mean, linop_cov_rv)
assert isinstance(out, randvars.Normal)
def test_backward_rv_classic(self, some_normal_rv1, some_normal_rv2):
array_cov_rv1 = some_normal_rv1
linop_cov_rv1 = randvars.Normal(array_cov_rv1.mean.copy(),
linops.aslinop(array_cov_rv1.cov))
array_cov_rv2 = some_normal_rv2
linop_cov_rv2 = randvars.Normal(array_cov_rv2.mean.copy(),
linops.aslinop(array_cov_rv2.cov))
with config(matrix_free=True):
with pytest.warns(RuntimeWarning):
self.transition.backward_rv(array_cov_rv1, array_cov_rv2)
with pytest.warns(RuntimeWarning):
self.transition.backward_rv(linop_cov_rv1, array_cov_rv2)
with pytest.warns(RuntimeWarning):
self.transition.backward_rv(array_cov_rv1, linop_cov_rv2)
out, _ = self.transition.backward_rv(linop_cov_rv1, linop_cov_rv2)
assert isinstance(out, randvars.Normal)
assert isinstance(out.cov, linops.LinearOperator)
assert isinstance(out.cov_cholesky, linops.LinearOperator)
with pytest.raises(NotImplementedError):
self.sqrt_transition.backward_rv(array_cov_rv1, array_cov_rv2)
with pytest.raises(NotImplementedError):
self.sqrt_transition.backward_rv(linop_cov_rv1, linop_cov_rv2)
def test_forward_rv(self, some_normal_rv1):
array_cov_rv = some_normal_rv1
linop_cov_rv = randvars.Normal(array_cov_rv.mean.copy(),
linops.aslinop(array_cov_rv.cov))
with config(matrix_free=True):
with pytest.warns(RuntimeWarning):
self.transition.forward_rv(array_cov_rv, 0.0)
out, _ = self.transition.forward_rv(linop_cov_rv, 0.0)
assert isinstance(out, randvars.Normal)
assert isinstance(out.cov, linops.LinearOperator)
assert isinstance(out.cov_cholesky, linops.LinearOperator)
with pytest.raises(NotImplementedError):
self.sqrt_transition.forward_rv(array_cov_rv, 0.0)
with pytest.raises(NotImplementedError):
self.sqrt_transition.forward_rv(linop_cov_rv, 0.0)
def test_precompute_cov_cholesky_with_linops(self):
mean, cov = self.params
rv = randvars.Normal(mean, linops.aslinop(cov))
with self.subTest("No Cholesky precomputed"):
self.assertFalse(rv.cov_cholesky_is_precomputed)
with self.subTest("Damping factor check"):
with config(matrix_free=True):
rv.precompute_cov_cholesky(damping_factor=10.0)
self.assertIsInstance(rv.cov_cholesky, linops.LinearOperator)
self.assertAllClose(
rv.cov_cholesky.todense(),
np.linalg.cholesky(cov + 10.0 * np.eye(rv.cov.shape[0])),
)
with self.subTest("Cholesky is precomputed"):
self.assertTrue(rv.cov_cholesky_is_precomputed)
def solve(self,
callback=None,
maxiter=None,
atol=None,
rtol=None,
calibration=None):
"""Solve the linear system :math:`Ax=b`.
Parameters
----------
callback : function, optional
User-supplied function called after each iteration of the linear solver. It
is called as ``callback(xk, Ak, Ainvk, sk, yk, alphak, resid)`` and can be
used to return quantities from the iteration. Note that depending on the
function supplied, this can slow down the solver.
maxiter : int
Maximum number of iterations
atol : float
Absolute residual tolerance. Stops if
:math:`\\min(\\lVert r_i \\rVert, \\sqrt{\\operatorname{tr}(\\operatorname{Cov}(x))}) \\leq \\text{atol}`.
rtol : float
Relative residual tolerance. Stops if
:math:`\\min(\\lVert r_i \\rVert, \\sqrt{\\operatorname{tr}(\\operatorname{Cov}(x))}) \\leq \\text{rtol} \\lVert b \\rVert`.
calibration : str or float, default=False
If supplied calibrates the output via the given procedure or uncertainty
scale. Available calibration procedures / choices are
==================================== ================
No calibration None
Provided scale float
Most recent Rayleigh quotient ``adhoc``
Running (weighted) mean ``weightedmean``
GP regression for kernel matrices ``gpkern``
==================================== ================
Returns
-------
x : RandomVariable, shape=(n,) or (n, nrhs)
Approximate solution :math:`x` to the linear system. Shape of the return
matches the shape of ``b``.
A : RandomVariable, shape=(n,n)
Posterior belief over the linear operator.calibrate
Ainv : RandomVariable, shape=(n,n)
Posterior belief over the linear operator inverse :math:`H=A^{-1}`.
info : dict
Information on convergence of the solver.
""" # pylint: disable=line-too-long
# Initialization
self.iter_ = 0
resid = self.A @ self.x_mean - self.b
# Initialize uncertainty calibration
phi = None
psi = None
if calibration is None:
pass
elif (calibration is not None
or calibration is not False) and not self.is_calib_covclass:
warnings.warn(
message=
"Cannot use calibration without a compatible covariance class."
)
elif isinstance(calibration, str) and self.is_calib_covclass:
pass
elif self.is_calib_covclass:
if calibration < 0:
raise ValueError("Calibration scale must be non-negative.")
elif calibration == 0.0:
pass
else:
phi = calibration
psi = 1 / calibration
# Trace of solution covariance
_trace_Ainv_covfactor_update = 0
self.trace_Ainv_covfactor = self._compute_trace_Ainv_covfactor0(
Y=None, unc_scale=psi)
# Create output random variables
x, A, Ainv = self._get_output_randvars(Y_list=self.obs_list,
sy_list=self.sy,
phi=phi,
psi=psi)
# Iteration with stopping criteria
while True:
# Check convergence
_has_converged, _conv_crit = self.has_converged(iter=self.iter_,
maxiter=maxiter,
resid=resid,
atol=atol,
rtol=rtol)
if _has_converged:
break
# Compute search direction (with implicit reorthogonalization) via policy
search_dir = -self.Ainv_mean @ resid
self.search_dir_list.append(search_dir)
# Perform action and observe
obs = self.A @ search_dir
self.obs_list.append(obs)
# Compute step size
sy = search_dir.T @ obs
step_size = -np.squeeze((search_dir.T @ resid) / sy)
self.sy.append(sy)
# Step and residual update
self.x_mean = self.x_mean + step_size * search_dir
resid = resid + step_size * obs
# (Symmetric) mean and covariance updates
Vs = self.A_covfactor @ search_dir
delta_A = obs - self.A_mean @ search_dir
u_A = Vs / (search_dir.T @ Vs)
v_A = delta_A - 0.5 * (search_dir.T @ delta_A) * u_A
Wy = self.Ainv_covfactor @ obs
delta_Ainv = search_dir - self.Ainv_mean @ obs
yWy = np.squeeze(obs.T @ Wy)
u_Ainv = Wy / yWy
v_Ainv = delta_Ainv - 0.5 * (obs.T @ delta_Ainv) * u_Ainv
# Rank 2 mean updates (+= uv' + vu')
self.A_mean = linops.aslinop(self.A_mean) + self._mean_update(
u=u_A, v=v_A)
self.Ainv_mean = linops.aslinop(
self.Ainv_mean) + self._mean_update(u=u_Ainv, v=v_Ainv)
# Rank 1 covariance Kronecker factor update (-= u_A(Vs)' and -= u_Ainv(Wy)')
if self.iter_ == 0:
self._A_covfactor_update_term = self._covariance_update(u=u_A,
Ws=Vs)
self._Ainv_covfactor_update_term = self._covariance_update(
u=u_Ainv, Ws=Wy)
else:
self._A_covfactor_update_term = (
self._A_covfactor_update_term +
self._covariance_update(u=u_A, Ws=Vs))
self._Ainv_covfactor_update_term = (
self._Ainv_covfactor_update_term +
self._covariance_update(u=u_Ainv, Ws=Wy))
self.A_covfactor = (linops.aslinop(self.A_covfactor0) -
self._A_covfactor_update_term)
self.Ainv_covfactor = (linops.aslinop(self.Ainv_covfactor0) -
self._Ainv_covfactor_update_term)
# Calibrate uncertainty based on Rayleigh quotient
if isinstance(calibration, str) and self.is_calib_covclass:
phi, psi = self._calibrate_uncertainty(
S=np.hstack(self.search_dir_list),
sy=np.vstack(self.sy).ravel(),
method=calibration,
)
# Update trace of solution covariance: tr(Cov(Hb))
_trace_Ainv_covfactor_update += 1 / yWy * np.squeeze(Wy.T @ Wy)
self.trace_Ainv_covfactor = np.real_if_close(
self._compute_trace_Ainv_covfactor0(Y=np.hstack(self.obs_list),
unc_scale=psi) -
_trace_Ainv_covfactor_update).item()
# Create output random variables
x, A, Ainv = self._get_output_randvars(Y_list=self.obs_list,
sy_list=self.sy,
phi=phi,
psi=psi)
# Callback function used to extract quantities from iteration
if callback is not None:
callback(
xk=x,
Ak=A,
Ainvk=Ainv,
sk=search_dir,
yk=obs,
alphak=step_size,
resid=resid,
)
# Iteration increment
self.iter_ += 1
# Log information on solution
info = {
"iter": self.iter_,
"maxiter": maxiter,
"resid_l2norm": np.linalg.norm(resid, ord=2),
"trace_sol_cov": self.trace_sol_cov,
"conv_crit": _conv_crit,
"rel_cond":
None, # TODO: matrix condition from solver (see scipy solvers)
}
return x, A, Ainv, info
def _dense_cov_cholesky_as_linop(
self, damping_factor: FloatArgType
) -> linops.LinearOperator:
return linops.aslinop(self.dense_cov_cholesky(damping_factor=damping_factor))
