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