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))