def test_kronecker_transpose(self):
"""Kronecker product transpose property: (A (x) B)^T = A^T (x) B^T."""
for A, B in self.kronecker_matrices:
with self.subTest():
W = linops.Kronecker(A=A, B=B)
V = linops.Kronecker(A=A.T, B=B.T)
self.assertAllClose(W.T.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 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 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 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 proj2coord(self, coord: int) -> np.ndarray:
"""Projection matrix to :math:`i` th coordinates.
Computes the matrix
.. math:: H_i = \\left[ I_d \\otimes e_i \\right] P^{-1},
where :math:`e_i` is the :math:`i` th unit vector,
that projects to the :math:`i` th coordinate of a vector.
If the ODE is multidimensional, it projects to **each** of the
:math:`i` th coordinates of each ODE dimension.
Parameters
----------
coord : int
Coordinate index :math:`i` which to project to.
Expected to be in range :math:`0 \\leq i \\leq q + 1`.
Returns
-------
np.ndarray, shape=(d, d*(q+1))
Projection matrix :math:`H_i`.
"""
projvec1d = np.eye(self.num_derivatives + 1)[:, coord]
projmat1d = projvec1d.reshape((1, self.num_derivatives + 1))
if config.matrix_free:
return linops.Kronecker(
linops.Identity(self.wiener_process_dimension), projmat1d)
return np.kron(np.eye(self.wiener_process_dimension), projmat1d)
def _matmul_normal_constant(norm_rv: _Normal, constant_rv: _Constant) -> _Normal:
if norm_rv.ndim == 1 or (norm_rv.ndim == 2 and norm_rv.shape[0] == 1):
return _Normal(
mean=norm_rv.mean @ constant_rv.support,
cov=constant_rv.support.T @ (norm_rv.cov @ constant_rv.support),
random_state=_utils.derive_random_seed(
norm_rv.random_state, constant_rv.random_state
),
)
elif norm_rv.ndim == 2 and norm_rv.shape[0] > 1:
cov_update = _linear_operators.Kronecker(
_linear_operators.Identity(constant_rv.shape[0]), constant_rv.support
)
return _Normal(
mean=norm_rv.mean @ constant_rv.support,
cov=cov_update.T @ (norm_rv.cov @ cov_update),
random_state=_utils.derive_random_seed(
norm_rv.random_state, constant_rv.random_state
),
)
else:
raise TypeError(
"Currently, matrix multiplication is only supported for vector- and "
"matrix-variate Gaussians."
)
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 _matmul_normal_constant(norm_rv: _Normal,
constant_rv: _Constant) -> _Normal:
if norm_rv.ndim == 1 or (norm_rv.ndim == 2 and norm_rv.shape[0] == 1):
if norm_rv.cov_cholesky_is_precomputed:
cov_cholesky = _utils.linalg.cholesky_update(
constant_rv.support.T @ norm_rv.cov_cholesky)
else:
cov_cholesky = None
return _Normal(
mean=norm_rv.mean @ constant_rv.support,
cov=constant_rv.support.T @ (norm_rv.cov @ constant_rv.support),
cov_cholesky=cov_cholesky,
random_state=_utils.derive_random_seed(norm_rv.random_state,
constant_rv.random_state),
)
elif norm_rv.ndim == 2 and norm_rv.shape[0] > 1:
# This part does not do the Cholesky update,
# because of performance configurations: currently, there is no way of switching
# the Cholesky updates off, which might affect (large, potentially sparse) covariance matrices
# of matrix-variate Normal RVs. See Issue #335.
cov_update = _linear_operators.Kronecker(
_linear_operators.Identity(constant_rv.shape[0]),
constant_rv.support)
return _Normal(
mean=norm_rv.mean @ constant_rv.support,
cov=cov_update.T @ (norm_rv.cov @ cov_update),
random_state=_utils.derive_random_seed(norm_rv.random_state,
constant_rv.random_state),
)
else:
raise TypeError(
"Currently, matrix multiplication is only supported for vector- and "
"matrix-variate Gaussians.")
def _drift_matrix(self):
drift_matrix_1d = np.diag(np.ones(self.num_derivatives), 1)
if config.matrix_free:
return linops.Kronecker(
A=linops.Identity(self.wiener_process_dimension),
B=linops.Matrix(A=drift_matrix_1d),
)
return np.kron(np.eye(self.wiener_process_dimension), drift_matrix_1d)
def test_rv_linop_kroneckercov(self):
"""Create a rv with a normal distribution with linear operator mean and Kronecker product kernels."""
def mv(v):
return np.array([2 * v[0], 3 * v[1]])
A = linops.LinearOperator(shape=(2, 2), matvec=mv)
V = linops.Kronecker(A, A)
rvs.Normal(mean=A, cov=V)
def test_kronecker_explicit(self):
"""Test the Kronecker operator against explicit matrix representations."""
for A, B in self.kronecker_matrices:
with self.subTest():
W = linops.Kronecker(A=A, B=B)
AkronB = np.kron(A, B)
self.assertAllClose(W.todense(), AkronB)
def __call__(self, step):
scaling_vector = np.abs(step) ** self.powers / self.scales
if config.matrix_free:
return linops.Kronecker(
A=linops.Identity(self.dimension),
B=linops.Scaling(factors=scaling_vector),
)
return np.kron(np.eye(self.dimension), np.diag(scaling_vector))
def _dispersion_matrix(self):
dispersion_matrix_1d = np.zeros(self.num_derivatives + 1)
dispersion_matrix_1d[-1] = 1.0 # Unit diffusion
if config.matrix_free:
return linops.Kronecker(
A=linops.Identity(self.wiener_process_dimension),
B=linops.Matrix(A=dispersion_matrix_1d.reshape(-1, 1)),
)
return np.kron(np.eye(self.wiener_process_dimension), dispersion_matrix_1d).T
def case_state_matrix_based(rng: np.random.Generator, ):
"""State of a matrix-based linear solver."""
prior = linalg.solvers.beliefs.LinearSystemBelief(
A=randvars.Normal(
mean=linops.Matrix(linsys.A),
cov=linops.Kronecker(A=linops.Identity(n), B=linops.Identity(n)),
),
x=(Ainv @ b[:, None]).reshape((n, )),
Ainv=randvars.Normal(
mean=linops.Identity(n),
cov=linops.Kronecker(A=linops.Identity(n), B=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_rv_linop_kroneckercov(self):
"""Create a rv with a normal distribution with linear operator mean and
Kronecker product kernels."""
@linops.LinearOperator.broadcast_matvec
def _matmul(v):
return np.array([2 * v[0], 3 * v[1]])
A = linops.LinearOperator(shape=(2, 2), dtype=np.double, matmul=_matmul)
V = linops.Kronecker(A, A)
randvars.Normal(mean=A, cov=V)
def matrixvariate_normal(shape: ShapeLike, precompute_cov_cholesky: bool,
rng: np.random.Generator) -> randvars.Normal:
rv = randvars.Normal(
mean=rng.normal(size=shape),
cov=linops.Kronecker(
A=random_spd_matrix(dim=shape[0], rng=rng),
B=random_spd_matrix(dim=shape[1], rng=rng),
),
)
if precompute_cov_cholesky:
rv.precompute_cov_cholesky()
return rv
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.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 = np.outer(gain, covfactor_Ms)
return randvars.Normal(
mean=matrix.mean + np.outer(resid, gain),
cov=linops.Kronecker(A=matrix.cov.A,
B=matrix.cov.B - covfactor_update),
)
def _kronecker_cov_cholesky(self) -> linops.Kronecker:
assert isinstance(self._cov, linops.Kronecker)
A = self._cov.A.todense()
B = self._cov.B.todense()
return linops.Kronecker(
A=scipy.linalg.cholesky(
A + COV_CHOLESKY_DAMPING * np.eye(A.shape[0], dtype=self.dtype),
lower=True,
),
B=scipy.linalg.cholesky(
B + COV_CHOLESKY_DAMPING * np.eye(B.shape[0], dtype=self.dtype),
lower=True,
),
dtype=self.dtype,
)
def _kronecker_cov_cholesky(
self,
damping_factor: Optional[FloatArgType] = COV_CHOLESKY_DAMPING
) -> linops.Kronecker:
assert isinstance(self._cov, linops.Kronecker)
A = self._cov.A.todense()
B = self._cov.B.todense()
return linops.Kronecker(
A=scipy.linalg.cholesky(
A + damping_factor * np.eye(A.shape[0], dtype=self.dtype),
lower=True,
),
B=scipy.linalg.cholesky(
B + damping_factor * np.eye(B.shape[0], dtype=self.dtype),
lower=True,
),
)
def _kronecker_cov_cholesky(
self,
damping_factor: FloatLike,
) -> linops.Kronecker:
assert isinstance(self.cov, linops.Kronecker)
A = self.cov.A.todense()
B = self.cov.B.todense()
return linops.Kronecker(
A=scipy.linalg.cholesky(
A + damping_factor * np.eye(A.shape[0], dtype=self.dtype),
lower=True,
),
B=scipy.linalg.cholesky(
B + damping_factor * np.eye(B.shape[0], dtype=self.dtype),
lower=True,
),
)
def _matmul_normal_constant(norm_rv: _Normal, constant_rv: _Constant) -> _Normal:
"""Normal random variable multiplied with a vector or matrix.
Computes the distribution of the random variable :math:`Y = XA`, where :math:`X`
is a matrix- or multi-variate normal random variable and :math:`A` a constant.
"""
if norm_rv.ndim == 1 or (norm_rv.ndim == 2 and norm_rv.shape[0] == 1):
if norm_rv.cov_cholesky_is_precomputed:
cov_cholesky = _utils.linalg.cholesky_update(
constant_rv.support.T @ norm_rv.cov_cholesky
)
else:
cov_cholesky = None
mean = norm_rv.mean @ constant_rv.support
cov = constant_rv.support.T @ (norm_rv.cov @ constant_rv.support)
if cov.shape == () and mean.shape == (1,):
cov = cov.reshape((1, 1))
return _Normal(mean=mean, cov=cov, cov_cholesky=cov_cholesky)
# This part does not do the Cholesky update,
# because of performance configurations: currently, there is no way of switching
# the Cholesky updates off, which might affect (large, potentially sparse)
# covariance matrices of matrix-variate Normal RVs. See Issue #335.
if constant_rv.support.ndim == 1:
constant_rv_support = constant_rv.support[:, None]
else:
constant_rv_support = constant_rv.support
cov_update = _linear_operators.Kronecker(
_linear_operators.Identity(norm_rv.shape[0]), constant_rv_support.T
)
# Cov(rvec(XA)) = Cov((I (x) A.T)rvec(X)) = (I (x) A.T)Cov(rvec(X))(I (x) A.T).T
return _Normal(
mean=norm_rv.mean @ constant_rv.support,
cov=cov_update @ (norm_rv.cov @ cov_update.T),
)
def test_reshape(self):
rv = randvars.Normal(
mean=np.random.uniform(size=(4, 3)),
cov=linops.Kronecker(A=random_spd_matrix(4),
B=random_spd_matrix(3)).todense(),
)
newshape = (2, 6)
reshaped_rv = rv.reshape(newshape)
self.assertArrayEqual(reshaped_rv.mean, rv.mean.reshape(newshape))
self.assertArrayEqual(reshaped_rv.cov, rv.cov)
# Test sampling
rv.random_state = 42
dist_sample = rv.sample(size=5)
reshaped_rv.random_state = 42
dist_reshape_sample = reshaped_rv.sample(size=5)
self.assertArrayEqual(dist_reshape_sample,
dist_sample.reshape((-1, ) + newshape))
def _matmul_constant_normal(constant_rv: _Constant, norm_rv: _Normal) -> _Normal:
"""Matrix-multiplication with a normal random variable.
Computes the distribution of the random variable :math:`Y = AX`, where :math:`X` is
a matrix- or multi-variate normal random variable and :math:`A` a constant.
"""
if norm_rv.ndim == 1 or (norm_rv.ndim == 2 and norm_rv.shape[1] == 1):
if norm_rv.cov_cholesky_is_precomputed:
cov_cholesky = _utils.linalg.cholesky_update(
constant_rv.support @ norm_rv.cov_cholesky
)
else:
cov_cholesky = None
return _Normal(
mean=constant_rv.support @ norm_rv.mean,
cov=constant_rv.support @ (norm_rv.cov @ constant_rv.support.T),
cov_cholesky=cov_cholesky,
)
# This part does not do the Cholesky update,
# because of performance configurations: currently, there is no way of switching
# the Cholesky updates off, which might affect (large, potentially sparse)
# covariance matrices of matrix-variate Normal RVs. See Issue #335.
if constant_rv.support.ndim == 1:
constant_rv_support = constant_rv.support[None, :]
else:
constant_rv_support = constant_rv.support
cov_update = _linear_operators.Kronecker(
constant_rv_support,
_linear_operators.Identity(norm_rv.shape[1]),
)
# Cov(rvec(AX)) = Cov((A (x) I)rvec(X)) = (A (x) I)Cov(rvec(X))(A (x) I).T
return _Normal(
mean=constant_rv.support @ norm_rv.mean,
cov=cov_update @ (norm_rv.cov @ cov_update.T),
)
def test_reshape(self):
rv = randvars.Normal(
mean=np.random.uniform(size=(4, 3)),
cov=linops.Kronecker(
A=random_spd_matrix(rng=self.rng, dim=4),
B=random_spd_matrix(rng=self.rng, dim=3),
).todense(),
)
newshape = (2, 6)
reshaped_rv = rv.reshape(newshape)
self.assertArrayEqual(reshaped_rv.mean, rv.mean.reshape(newshape))
self.assertArrayEqual(reshaped_rv.cov, rv.cov)
# Test sampling
fixed_rng = np.random.default_rng(seed=self.seed)
dist_sample = rv.sample(rng=fixed_rng, size=5)
fixed_rng = np.random.default_rng(seed=self.seed)
dist_reshape_sample = reshaped_rv.sample(rng=fixed_rng, size=5)
self.assertArrayEqual(dist_reshape_sample,
dist_sample.reshape((-1, ) + newshape))
