def test_rv_linop_kroneckercov(): """Create a random variable with linear operator mean and Kronecker product covariance.""" def mv(v): return np.array([2 * v[0], 3 * v[1]]) A = linear_operators.LinearOperator(shape=(2, 2), matvec=mv) V = linear_operators.Kronecker(A, A) probability.RandomVariable(distribution=probability.Normal(mean=A, cov=V))
def test_sample(mean, cov): """Draw samples and check all sample dimensions.""" # TODO: check dimension of each realization in dist_sample dist = probability.Normal(mean=mean, cov=cov, random_state=1) dist_sample = dist.sample(size=5) if not np.isscalar(dist.mean()): ndims_rv = len(mean.shape) np.testing.assert_equal( dist_sample.shape[-ndims_rv:], mean.shape, err_msg="Realization shape does not match mean shape.")
def test_sample_zero_cov(mean, cov): """Draw sample from distribution with zero covariance and check whether it equals the mean.""" dist = probability.Normal(mean=mean, cov=0 * cov, random_state=1) dist_sample = dist.sample(size=1) assert_str = "Draw with covariance zero does not match mean." if isinstance(dist.mean(), linear_operators.LinearOperator): np.testing.assert_allclose(dist_sample, dist.mean().todense(), err_msg=assert_str) else: np.testing.assert_allclose(dist_sample, dist.mean(), err_msg=assert_str)
def test_matrixprior(matlinsolve): """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 b = np.random.rand(n, 1) # Prior distribution on A covA = linear_operators.SymmetricKronecker(A=np.eye(n), B=np.eye(n)) Ainv0 = probability.RandomVariable(distribution=probability.Normal(mean=np.eye(n), cov=covA)) x, Ahat, Ainvhat, info = matlinsolve(A=A, Ainv0=Ainv0, b=b) xnp = np.linalg.solve(A, b).ravel() np.testing.assert_allclose(x.mean(), xnp, rtol=1e-4, err_msg="Solution does not match np.linalg.solve.")
def test_symmetric_samples(): """Samples from a normal distribution with symmetric Kronecker covariance of two symmetric matrices are symmetric.""" np.random.seed(42) n = 3 A = np.random.uniform(size=(n, n)) A = 0.5 * (A + A.T) + n * np.eye(n) dist = probability.Normal(mean=np.eye(A.shape[0]), cov=linear_operators.SymmetricKronecker(A=A)) dist_sample = dist.sample(size=10) for i, B in enumerate(dist_sample): np.testing.assert_allclose( B, B.T, atol=1e-5, rtol=1e-5, err_msg= "Sample {} from symmetric Kronecker distribution is not symmetric." .format(i))
def test_posterior_distribution_parameters(matblinsolve, poisson_linear_system): """Compute the posterior parameters of the matrix-based probabilistic linear solvers directly and compare.""" # Initialization A, f = poisson_linear_system S = [] # search directions Y = [] # observations # Priors H0 = linear_operators.Identity(A.shape[0]) # inverse prior mean A0 = linear_operators.Identity(A.shape[0]) # prior mean WH0 = H0 # inverse prior Kronecker factor WA0 = A # prior Kronecker factor covH = linear_operators.SymmetricKronecker(WH0, WH0) covA = linear_operators.SymmetricKronecker(WA0, WA0) Ahat0 = probability.RandomVariable(distribution=probability.Normal(mean=A0, cov=covA)) Ainvhat0 = probability.RandomVariable(distribution=probability.Normal(mean=H0, cov=covH)) # Define callback function to obtain search directions def callback_postparams(xk, Ak, Ainvk, sk, yk, alphak, resid): S.append(sk) Y.append(yk) # Solve linear system u_solver, Ahat, Ainvhat, info = matblinsolve(A=A, b=f, A0=Ahat0, Ainv0=Ainvhat0, callback=callback_postparams, calibrate=False) # Create arrays from lists S = np.squeeze(np.array(S)).T Y = np.squeeze(np.array(Y)).T # E[A] and E[A^-1] def posterior_mean(A0, WA0, S, Y): """Compute posterior mean of the symmetric probabilistic linear solver.""" Delta = (Y - A0 @ S) U_T = np.linalg.solve(S.T @ (WA0 @ S), (WA0 @ S).T) U = U_T.T Ak = A0 + Delta @ U_T + U @ Delta.T - U @ S.T @ Delta @ U_T return Ak Ak = posterior_mean(A0.todense(), WA0, S, Y) Hk = posterior_mean(H0.todense(), WH0, Y, S) np.testing.assert_allclose(Ahat.mean().todense(), Ak, rtol=1e-5, err_msg="The matrix estimated by the probabilistic linear solver does not match the " + "directly computed one.") np.testing.assert_allclose(Ainvhat.mean().todense(), Hk, rtol=1e-5, err_msg="The inverse matrix estimated by the probabilistic linear solver does not" + "match the directly computed one.") # Cov[A] and Cov[A^-1] def posterior_cov_kronfac(WA0, S): """Compute the covariance symmetric Kronecker factor of the probabilistic linear solver.""" U_AT = np.linalg.solve(S.T @ (WA0 @ S), (WA0 @ S).T) covfac = WA0 @ (np.identity(np.shape(WA0)[0]) - S @ U_AT) return covfac A_covfac = posterior_cov_kronfac(WA0, S) H_covfac = posterior_cov_kronfac(WH0, Y) np.testing.assert_allclose(Ahat.cov().A.todense(), A_covfac, rtol=1e-5, err_msg="The covariance estimated by the probabilistic linear solver does not match the " + "directly computed one.") np.testing.assert_allclose(Ainvhat.cov().A.todense(), H_covfac, rtol=1e-5, err_msg="The covariance estimated by the probabilistic linear solver does not" + "match the directly computed one.")
def _create_output_randvars(self, S=None, Y=None, Phi=None, Psi=None): """Return output random variables x, A, Ainv from their means and covariances.""" _A_covfactor = self.A_covfactor _Ainv_covfactor = self.Ainv_covfactor # Set degrees of freedom based on uncertainty calibration in unexplored space if Phi is not None: def _mv(x): def _I_S_fun(x): return x - S @ np.linalg.solve(S.T @ S, S.T @ x) return _I_S_fun(Phi @ _I_S_fun(x)) I_S_Phi_I_S_op = linear_operators.LinearOperator( shape=self.A.shape, matvec=_mv) _A_covfactor = self.A_covfactor + I_S_Phi_I_S_op if Psi is not None: def _mv(x): def _I_Y_fun(x): return x - Y @ np.linalg.solve(Y.T @ Y, Y.T @ x) return _I_Y_fun(Psi @ _I_Y_fun(x)) I_Y_Psi_I_Y_op = linear_operators.LinearOperator( shape=self.A.shape, matvec=_mv) _Ainv_covfactor = self.Ainv_covfactor + I_Y_Psi_I_Y_op # Create output random variables A = probability.RandomVariable( shape=self.A_mean.shape, dtype=float, distribution=probability.Normal( mean=self.A_mean, cov=linear_operators.SymmetricKronecker(A=_A_covfactor))) cov_Ainv = linear_operators.SymmetricKronecker(A=_Ainv_covfactor) Ainv = probability.RandomVariable(shape=self.Ainv_mean.shape, dtype=float, distribution=probability.Normal( mean=self.Ainv_mean, cov=cov_Ainv)) # Induced distribution on x via Ainv # Exp = x = A^-1 b, Cov = 1/2 (W b'Wb + Wbb'W) Wb = _Ainv_covfactor @ self.b bWb = np.squeeze(Wb.T @ self.b) def _mv(x): return 0.5 * (bWb * _Ainv_covfactor @ x + Wb @ (Wb.T @ x)) cov_op = linear_operators.LinearOperator( shape=np.shape(_Ainv_covfactor), dtype=float, matvec=_mv, matmat=_mv) x = probability.RandomVariable(shape=(self.A_mean.shape[0], ), dtype=float, distribution=probability.Normal( mean=self.x.ravel(), cov=cov_op)) return x, A, Ainv
def _preprocess_linear_system(A, b, assume_A, A0=None, Ainv0=None, x0=None): """ Transform the linear system to linear operator and random variable form. Parameters ---------- A : array-like or LinearOperator or RandomVariable A square matrix, linear operator or random variable representing the prior belief over :math:`A`. b : array_like, shape=(n,) or (n, nrhs) Right-hand side vector or matrix in :math:`A x = b`. assume_A : str, default="sympos" Assumptions on the matrix, which can influence solver choice or behavior. The available options are ==================== ========= generic matrix ``gen`` symmetric ``sym`` positive definite ``pos`` symmetric pos. def. ``sympos`` ==================== ========= If ``A`` or ``Ainv`` are random variables, then the encoded assumptions in the distribution are used automatically. A0 : RandomVariable, shape=(n,n) Random variable representing the prior belief over the linear operator :math:`A`. Ainv0 : array-like or LinearOperator or RandomVariable, shape=(n,n) Optional. A square matrix, linear operator or random variable representing the prior belief over the inverse :math:`H=A^{-1}`. x0 : array-like, or RandomVariable, shape=(n,) or (n, nrhs) Optional. Prior belief for the solution of the linear system. Will be ignored if ``Ainv`` is given. Returns ------- A : RandomVariable, shape=(n,n) Prior belief over the linear operator :math:`A`. b : array-like, shape=(n,) or (n, nrhs) Right-hand-side of the linear system. A0 : RandomVariable, shape=(n,n) Prior belief over the linear operator :math:`A`. Ainv0 : RandomVariable, shape=(n,n) Prior belief over the linear operator inverse :math:`H=A^{-1}`. x : array-like or RandomVariable, shape=(n,) or (n, nrhs) Prior belief over the solution :math:`x` to the linear system. """ # Choose matrix based view if not clear from arguments if (Ainv0 is not None or A0 is not None) and x0 is not None: warnings.warn( "Cannot use prior information on both the matrix (inverse) and the solution. The latter will be ignored." ) x = None else: x = x0 # Check matrix assumptions if assume_A not in ["gen", "sym", "pos", "sympos"]: raise ValueError( '\'{}\' is not a recognized linear operator assumption.'.format( assume_A)) # Choose priors for A and Ainv if not specified, based on matrix assumptions in "assume_A" if assume_A == "sympos": # No priors specified if A0 is None and Ainv0 is None: dist = probability.Normal( mean=linear_operators.Identity(shape=A.shape[0]), cov=linear_operators.SymmetricKronecker( linear_operators.Identity(shape=A.shape[0]))) Ainv0 = probability.RandomVariable(distribution=dist) dist = probability.Normal( mean=linear_operators.Identity(shape=A.shape[0]), cov=linear_operators.SymmetricKronecker( linear_operators.Identity(shape=A.shape[0]))) A0 = probability.RandomVariable(distribution=dist) # Only prior on Ainv specified elif A0 is None and Ainv0 is not None: try: if isinstance(Ainv0, probability.RandomVariable): A0_mean = Ainv0.mean().inv() else: A0_mean = Ainv0.inv() except AttributeError: warnings.warn( message= "Prior specified only for Ainv. Inverting prior mean naively. " + "This operation is computationally costly! Specify an inverse prior (mean) instead." ) A0_mean = np.linalg.inv(Ainv0.mean()) except NotImplementedError: A0_mean = linear_operators.Identity(A.shape[0]) warnings.warn( message= "Prior specified only for Ainv. Automatic prior mean inversion not implemented, " + "falling back to standard normal prior.") # hereditary positive definiteness A0_covfactor = A dist = probability.Normal( mean=A0_mean, cov=linear_operators.SymmetricKronecker(A=A0_covfactor)) A0 = probability.RandomVariable(distribution=dist) # Only prior on A specified if A0 is not None and Ainv0 is None: try: if isinstance(A0, probability.RandomVariable): Ainv0_mean = A0.mean().inv() else: Ainv0_mean = A0.inv() except AttributeError: warnings.warn( message= "Prior specified only for Ainv. Inverting prior mean naively. " + "This operation is computationally costly! Specify an inverse prior (mean) instead." ) Ainv0_mean = np.linalg.inv(A0.mean()) except NotImplementedError: Ainv0_mean = linear_operators.Identity(A.shape[0]) warnings.warn( message="Prior specified only for Ainv. " + "Automatic prior mean inversion failed, falling back to standard normal prior." ) # (non-symmetric) posterior correspondence Ainv0_covfactor = Ainv0_mean dist = probability.Normal( mean=Ainv0_mean, cov=linear_operators.SymmetricKronecker(A=Ainv0_covfactor)) Ainv0 = probability.RandomVariable(distribution=dist) elif assume_A == "sym": raise NotImplementedError elif assume_A == "pos": raise NotImplementedError elif assume_A == "gen": # TODO: Implement case where only a pre-conditioner is given as Ainv0 # TODO: Automatic prior selection based on data scale, matrix trace, etc. raise NotImplementedError # Transform linear system to correct dimensions b = utils.as_colvec(b) # (n,) -> (n, 1) if x0 is not None: x = utils.as_colvec(x0) # (n,) -> (n, 1) assert (not (Ainv0 is None and x is None)), "Neither Ainv nor x are specified." return A, b, A0, Ainv0, x
def test_normal_pdf(mean, cov): """Evaluate pdf at random input.""" dist = probability.Normal(mean=mean, cov=cov) pass
def test_normal_instantiation(mean, cov): """Instantiation of a normal distribution with mixed mean and cov type.""" probability.Normal(mean=mean, cov=cov)
def test_normal_dimension_mismatch(mean, cov): """Instantiating a normal distribution with mismatched mean and covariance should result in a ValueError.""" with pytest.raises(ValueError): assert probability.Normal( mean=mean, cov=cov), "Mean and covariance mismatch in normal distribution."
A = linear_operators.LinearOperator(shape=(2, 2), matvec=mv) V = linear_operators.Kronecker(A, A) probability.RandomVariable(distribution=probability.Normal(mean=A, cov=V)) # Random variable arithmetic arrays2d = [ np.empty(2), np.zeros(2), np.array([np.inf, 1]), np.array([1, -2.5]) ] matrices2d = [np.array([[1, 2], [3, 2]]), np.array([[0, 0], [1.0, -4.3]])] linops2d = [linear_operators.MatrixMult(A=np.array([[1, 2], [4, 5]]))] randvars2d = [ probability.RandomVariable(distribution=probability.Normal( mean=np.array([1, 2]), cov=np.array([[2, 0], [0, 5]]))) ] randvars2x2 = [ probability.RandomVariable(shape=(2, 2), distribution=probability.Normal( mean=np.array([[-2, .3], [0, 1]]), cov=linear_operators.SymmetricKronecker( A=np.eye(2), B=np.ones((2, 2))))) ] @pytest.mark.parametrize("x,rv", list(itertools.product(arrays2d, randvars2d))) def test_rv_addition(x, rv): """Addition with random variables.""" z1 = x + rv z2 = rv + x