def test_cov_cholesky_cov_cholesky_not_passed(self): """No cov_cholesky is passed in init. In this case, the "is_precomputed" flag is False, a cov_cholesky is computed on demand, but can also be computed manually with any damping factor. """ mean, cov = self.params rv = rvs.Normal(mean, cov) with self.subTest("No Cholesky precomputed"): self.assertFalse(rv.cov_cholesky_is_precomputed) with self.subTest("Cholesky factor is computed correctly"): # The default damping factor 1e-12 does not mess up this test self.assertAllClose(rv.cov_cholesky, np.sqrt(rv.cov)) with self.subTest("Cholesky is precomputed"): self.assertTrue(rv.cov_cholesky_is_precomputed)
def test_cov_cholesky_cov_cholesky_not_passed(self): """No cov_cholesky is passed in init. In this case, the "is_precomputed" flag is False, a cov_cholesky is computed on demand, but can also be computed manually with any damping factor. """ rv = rvs.Normal(mean=np.random.uniform(size=(2, 2)), cov=random_spd_matrix(4)) with self.subTest("No Cholesky precomputed"): self.assertFalse(rv.cov_cholesky_is_precomputed) with self.subTest("Cholesky factor is computed correctly"): # The default damping factor 1e-12 does not mess up this test self.assertAllClose(rv.cov_cholesky, np.linalg.cholesky(rv.cov)) with self.subTest("Cholesky is precomputed"): self.assertTrue(rv.cov_cholesky_is_precomputed)
def test_reshape_newaxis(self): dist = rvs.Normal(*self.params) for ndim in range(1, 3): for method in ["newaxis", "reshape"]: with self.subTest(): newshape = tuple([1] * ndim) if method == "newaxis": if ndim == 1: newdist = dist[np.newaxis] elif ndim == 2: newdist = dist[np.newaxis, np.newaxis] elif method == "reshape": newdist = dist.reshape(newshape) self.assertEqual(newdist.shape, newshape) self.assertEqual(np.squeeze(newdist.mean), dist.mean) self.assertEqual(np.squeeze(newdist.cov), dist.cov)
def test_transform_of_gaussian_exact(self): sigpts = self.ut.sigma_points(pnrv.Normal(self.mean, self.covar)) ndim_meas = self.ndim + 1 # != self.ndim is important transmtrx = np.random.rand(ndim_meas, self.ndim) meascov = 0 * np.random.rand(ndim_meas, ndim_meas) proppts = self.ut.propagate(None, sigpts, lambda t, x: transmtrx @ x) mest, cest, ccest = self.ut.estimate_statistics( proppts, sigpts, meascov, self.mean) diff_mean = np.linalg.norm(mest - transmtrx @ self.mean) diff_covar = np.linalg.norm(cest - transmtrx @ self.covar @ transmtrx.T) diff_crosscovar = np.linalg.norm(ccest - self.covar @ transmtrx.T) self.assertLess(diff_mean / np.linalg.norm(transmtrx @ self.mean), 1e-11) self.assertLess( diff_covar / np.linalg.norm(transmtrx @ self.covar @ transmtrx.T), 1e-11) self.assertLess( diff_crosscovar / np.linalg.norm(self.covar @ transmtrx.T), 1e-11)
def _forward_rv_sqrt( self, rv, t, compute_gain=False, _diffusion=1.0 ) -> (pnrv.RandomVariable, typing.Dict): H = self.state_trans_mat_fun(t) SR = self.proc_noise_cov_cholesky_fun(t) shift = self.shift_vec_fun(t) new_mean = H @ rv.mean + shift new_cov_cholesky = cholesky_update( H @ rv.cov_cholesky, np.sqrt(_diffusion) * SR ) new_cov = new_cov_cholesky @ new_cov_cholesky.T crosscov = rv.cov @ H.T info = {"crosscov": crosscov} if compute_gain: gain = crosscov @ np.linalg.inv(new_cov) info["gain"] = gain return pnrv.Normal(new_mean, cov=new_cov, cov_cholesky=new_cov_cholesky), info
def _initialdistribution(ivp, prior): """ Conditions initialdistribution :math:`\\mathcal{N}(0, P P^\\top)` on the initial values :math:`(x_0, f(t_0, x_0), ...)` using as many available derivatives as possible. Note that the projection matrices :math:`H_0` and :math:`H_1` become :math:`H_0 P^{-1}` and :math:`H_1 P^{-1}`. """ if not issubclass(type(ivp.initrv), rvs.Normal): if not issubclass(type(ivp.initrv), rvs.Dirac): raise RuntimeError("Initial distribution not Normal nor Dirac") x0 = ivp.initialdistribution.mean dx0 = ivp.rhs(ivp.t0, x0) ddx0 = _ddx(ivp.t0, x0, ivp) dddx0 = _dddx(ivp.t0, x0, ivp) h0 = prior.proj2coord(coord=0) h1 = prior.proj2coord(coord=1) precond = prior.preconditioner # initcov = np.eye(*(precond @ precond.T).shape) initcov = precond @ precond.T if prior.ordint == 1: projmat = np.hstack((h0.T, h1.T)).T data = np.hstack((x0, dx0)) _size = 2 elif prior.ordint == 2: # try only jacobian h2 = prior.proj2coord(coord=2) projmat = np.hstack((h0.T, h1.T, h2.T)).T data = np.hstack((x0, dx0, ddx0)) _size = 3 else: # try jacobian and hessian h2 = prior.proj2coord(coord=2) h3 = prior.proj2coord(coord=3) projmat = np.hstack((h0.T, h1.T, h2.T, h3.T)).T data = np.hstack((x0, dx0, ddx0, dddx0)) _size = 4 largecov = np.kron(np.eye(_size), ivp.initialdistribution.cov) s = projmat @ initcov @ projmat.T + largecov crosscov = initcov @ projmat.T newmean = crosscov @ np.linalg.solve(s, data) newcov = initcov - (crosscov @ np.linalg.solve(s.T, crosscov.T)).T return rvs.Normal(newmean, newcov)
def stochastic_policy( fun: Callable[[FloatArgType], FloatArgType], fun_params0: pn.RandomVariable, random_state: RandomStateArgType = None, ) -> float: """ Policy returning a random action. Parameters ---------- fun : One-dimensional objective function. fun_params0 : Belief over the parameters of the quadratic objective. random_state : Random state of the policy. If None (or np.random), the global np.random state is used. If integer, it is used to seed the local :class:`~numpy.random.RandomState` instance. """ return rvs.Normal(mean=0.0, cov=1.0, random_state=random_state).sample()
def forward_rv(self, rv, t, compute_gain=False, _diffusion=1.0, _linearise_at=None, **kwargs) -> typing.Tuple[pnrv.Normal, typing.Dict]: compute_sigmapts_at = _linearise_at if _linearise_at is not None else rv self.sigma_points = self.assemble_sigma_points( at_this_rv=compute_sigmapts_at) proppts = self.ut.propagate(t, self.sigma_points, self.non_linear_model.state_trans_fun) meascov = _diffusion * self.non_linear_model.proc_noise_cov_mat_fun(t) mean, cov, crosscov = self.ut.estimate_statistics( proppts, self.sigma_points, meascov, rv.mean) info = {"crosscov": crosscov} if compute_gain: gain = crosscov @ np.linalg.inv(cov) info["gain"] = gain return pnrv.Normal(mean, cov), info
def test_reshape(self): rv = rvs.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 gaussian_belief_update( fun_params0: pn.RandomVariable, action: FloatArgType, observation: FloatArgType, noise_cov: Union[np.ndarray, linops.LinearOperator], ) -> pn.RandomVariable: """ Update the belief over the parameters with an observation. Parameters ---------- fun_params0 : Belief over the parameters of the quadratic objective. action : Action of the probabilistic quadratic optimizer. observation : Observation of the problem corresponding to the given `action`. noise_cov : *shape=(3, 3)* -- Covariance of the noise on the parameters of the quadratic objective given by the assumed observation model. """ # Feature vector x = np.asarray(action).reshape(1, -1) Phi = np.vstack((0.5 * x**2, x, np.ones_like(x))) # Mean and covariance mu = fun_params0.mean Sigma = fun_params0.cov # Gram matrix gram = Phi.T @ (Sigma + noise_cov) @ Phi # Posterior Mean m = mu + Sigma @ Phi @ np.linalg.solve(gram, observation - Phi.T @ mu) # Posterior Covariance S = Sigma - Sigma @ Phi @ np.linalg.solve(gram, Phi.T @ Sigma) return rvs.Normal(mean=m, cov=S)
def test_cov_cholesky_cov_cholesky_passed(self): """A value for cov_cholesky is passed in init. In this case, the "is_precomputed" flag is True, the cov_cholesky returns the argument that has been passed, but (p)recomputing overwrites the argument with a new factor. """ mean, cov = self.params # This is purposely not the correct Cholesky factor for test reasons cov_cholesky = np.random.rand(*cov.shape) rv = rvs.Normal(mean, cov, cov_cholesky=cov_cholesky) with self.subTest("Cholesky precomputed"): self.assertTrue(rv.cov_cholesky_is_precomputed) with self.subTest("Returns correct cov_cholesky"): self.assertAllClose(rv.cov_cholesky, cov_cholesky) with self.subTest("self.precompute raises exception"): with self.assertRaises(Exception): rv.precompute_cov_cholesky()
def explore_exploit_policy( fun: Callable[[FloatArgType], FloatArgType], fun_params0: pn.RandomVariable, random_state: RandomStateArgType = None, ) -> float: """ Policy exploring around the estimate of the minimum based on the certainty about the parameters. Parameters ---------- fun : One-dimensional objective function. fun_params0 : Belief over the parameters of the quadratic objective. random_state : Random state of the policy. If None (or np.random), the global np.random state is used. If integer, it is used to seed the local :class:`~numpy.random.RandomState` instance. """ a0, b0, c0 = fun_params0 return (-b0.mean / a0.mean + rvs.Normal( 0, np.trace(fun_params0.cov), random_state=random_state).sample())
def test_symmetric_samples(self): """Samples from a normal distribution with symmetric Kronecker kernels 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) rv = rvs.Normal( mean=np.eye(A.shape[0]), cov=linops.SymmetricKronecker(A=A), random_state=1, ) rv = rv.sample(size=10) for i, B in enumerate(rv): self.assertAllClose( B, B.T, atol=1e-5, rtol=1e-5, msg= "Sample {} from symmetric Kronecker distribution is not symmetric." .format(i), )
def _initialdistribution(ivp, prior): """Conditions initialdistribution :math:`\\mathcal{N}(0, P P^\\top)` on the initial values :math:`(x_0, f(t_0, x_0), ...)` using as many available derivatives as possible. Note that the projection matrices :math:`H_0` and :math:`H_1` become :math:`H_0 P^{-1}` and :math:`H_1 P^{-1}`. """ if not issubclass(type(ivp.initrv), pnrv.Normal): if not issubclass(type(ivp.initrv), pnrv.Constant): raise RuntimeError("Initial distribution not Normal or Constant") x0 = ivp.initialdistribution.mean dx0 = ivp.rhs(ivp.t0, x0) ddx0 = _ddx(ivp.t0, x0, ivp) h0 = prior.proj2coord(coord=0) h1 = prior.proj2coord(coord=1) initcov = np.eye(len(h0.T)) if prior.ordint == 1: projmat = np.hstack((h0.T, h1.T)).T data = np.hstack((x0, dx0)) else: h2 = prior.proj2coord(coord=2) projmat = np.hstack((h0.T, h1.T, h2.T)).T data = np.hstack((x0, dx0, ddx0)) s = projmat @ initcov @ projmat.T crosscov = initcov @ projmat.T # @ np.linalg.inv(s) newmean = crosscov @ np.linalg.solve(s, data) newcov = initcov - crosscov @ np.linalg.solve(s, crosscov.T) lu, d, _ = scipy.linalg.ldl(newcov, lower=True) chol = lu @ np.sqrt(d) return pnrv.Normal(newmean, newcov, cov_cholesky=chol)
def test_same_backward_outputs(both_transitions, diffusion): trans1, trans2 = both_transitions real = 1 + 0.1 * np.random.rand(trans1.dimension) real2 = 1 + 0.1 * np.random.rand(trans1.dimension) cov = random_spd_matrix(trans1.dimension) rv = pnrv.Normal(real2, cov) out_1, info1 = trans1.backward_realization(real, rv, t=0.0, dt=0.5, compute_gain=True, _diffusion=diffusion) out_2, info2 = trans2.backward_realization(real, rv, t=0.0, dt=0.5, compute_gain=True, _diffusion=diffusion) np.testing.assert_allclose(out_1.mean, out_2.mean) np.testing.assert_allclose(out_1.cov, out_2.cov) # Both dicts are empty? assert not info1 assert not info2
def test_transition_rv(self): out_rv, _ = self.nl.transition_rv( rvs.Normal(np.ones(self.nl.dimension), np.eye(self.nl.dimension)), start=None, ) self.assertIsInstance(out_rv, rvs.RandomVariable)
def test_normal_pdf(self): """Evaluate pdf at random input.""" for mean, cov in self.normal_params: with self.subTest(): dist = rvs.Normal(mean=mean, cov=cov) pass
def _backward_rv_sqrt( self, rv_obtained, rv, rv_forwarded=None, gain=None, t=None, _diffusion=1.0, ): """See Section 4.1f of: ``https://www.sciencedirect.com/science/article/abs/pii/S0005109805001810``. """ # forwarded_rv is ignored in square-root smoothing. # Smoothing updates need the gain, but # filtering updates "compute their own". # Thus, if we are doing smoothing (|cov_obtained|>0) an the gain is not provided, # make an extra prediction to compute the gain. if gain is None: if np.linalg.norm(rv_obtained.cov) > 0: _, info = self.forward_rv(rv, t=t, compute_gain=True, _diffusion=_diffusion) gain = info["gain"] else: gain = np.zeros((len(rv.mean), len(rv_obtained.mean))) state_trans = self.state_trans_mat_fun(t) proc_noise_chol = np.sqrt( _diffusion) * self.proc_noise_cov_cholesky_fun(t) shift = self.shift_vec_fun(t) chol_past = rv.cov_cholesky chol_obtained = rv_obtained.cov_cholesky output_dim = self.output_dim input_dim = self.input_dim zeros_bottomleft = np.zeros((output_dim, output_dim)) zeros_middleright = np.zeros((output_dim, input_dim)) blockmat = np.block([ [chol_past.T @ state_trans.T, chol_past.T], [proc_noise_chol.T, zeros_middleright], [zeros_bottomleft, chol_obtained.T @ gain.T], ]) big_triu = np.linalg.qr(blockmat, mode="r") new_chol_triu = big_triu[output_dim:(output_dim + input_dim), output_dim:(output_dim + input_dim)] # If no initial gain was provided, compute it from the QR-results # This is required in the Kalman update, where, other than in the smoothing update, # no initial gain was provided. # Recall that above, gain was set to zero in this setting. if np.linalg.norm(gain) == 0.0: gain = big_triu[:output_dim, output_dim:].T @ np.linalg.inv( big_triu[:output_dim, :output_dim].T) new_mean = rv.mean + gain @ (rv_obtained.mean - state_trans @ rv.mean - shift) new_cov_cholesky = tril_to_positive_tril(new_chol_triu.T) new_cov = new_cov_cholesky @ new_cov_cholesky.T return pnrv.Normal(new_mean, new_cov, cov_cholesky=new_cov_cholesky), {}
def _apply_precon(precon, rv): new_mean = precon @ rv.mean new_cov = precon @ rv.cov @ precon.T new_cov_choleky = precon @ rv.cov_cholesky # precon is diagonal, so this is valid return pnrv.Normal(new_mean, new_cov, cov_cholesky=new_cov_choleky)
def some_normal_rv4(test_ndim, spdmat4): return pnrv.Normal( mean=np.random.rand(test_ndim), cov=spdmat4, cov_cholesky=np.linalg.cholesky(spdmat4), )
def test_normal_instantiation(self): """Instantiation of a normal distribution with mixed mean and cov type.""" for mean, cov in self.normal_params: with self.subTest(): rvs.Normal(mean=mean, cov=cov)
def initialize_odefilter_with_rk(f, y0, t0, prior, initrv, df=None, h0=1e-2, method="DOP853"): r"""Initialize an ODE filter by fitting the prior process to a few steps of an approximate ODE solution computed with Scipy's RK. It goes as follows: 1. The ODE integration problem is set up on the interval ``[t0, t0 + (2*order+1)*h0]`` and solved with a call to ``scipy.integrate.solve_ivp``. The solver is uses adaptive steps with ``atol=rtol=1e-12``, but is forced to pass through the events ``(t0, t0+h0, t0 + 2*h0, ..., t0 + (2*order+1)*h0)``. The result is a vector of time points and states, with at least ``(2*order+1)``. Potentially, the adaptive steps selected many more steps, but because of the events, fewer steps cannot have happened. 2. A prescribed prior is fitted to the first ``(2*order+1)`` (t, y) pairs of the solution. ``order`` is the order of the prior. 3. The value of the resulting posterior at time ``t=t0`` is an estimate of the state and all its derivatives. The resulting marginal standard deviations estimate the error. This random variable is returned. Parameters ---------- f ODE vector field. y0 Initial value. t0 Initial time point. prior Prior distribution used for the ODE solver. For instance an integrated Brownian motion prior (``IBM``). initrv Initial random variable. df Jacobian of the ODE vector field. Optional. If specified, more components of the result will be exact. h0 Maximum step-size to use for computing the approximate ODE solution. The smaller, the more accurate, but also, the smaller, the less stable. The best value here depends on the ODE problem, and probably the chosen method. Optional. Default is ``1e-2``. method Which solver to use. This is communicated as a string that is compatible with ``scipy.integrate.solve_ivp(..., method=method)``. Optional. Default is `DOP853`. Returns ------- Normal Estimated (improved) initial random variable. Compatible with the specified prior. Examples -------- >>> from dataclasses import astuple >>> from probnum.random_variables import Normal >>> from probnum.statespace import IBM >>> from probnum.problems.zoo.diffeq import vanderpol Compute the initial values of the van-der-Pol problem as follows >>> f, t0, tmax, y0, df, *_ = astuple(vanderpol()) >>> print(y0) [2. 0.] >>> prior = IBM(ordint=3, spatialdim=2) >>> initrv = Normal(mean=np.zeros(prior.dimension), cov=np.eye(prior.dimension)) >>> improved_initrv = initialize_odefilter_with_rk(f, y0, t0, prior=prior, initrv=initrv, df=df) >>> print(prior.proj2coord(0) @ improved_initrv.mean) [2. 0.] >>> print(np.round(improved_initrv.mean, 1)) [ 2. 0. -2. 58.2 0. -2. 60. -1745.7] >>> print(np.round(np.log10(improved_initrv.std), 1)) [-13.8 -11.3 -9. -1.5 -13.8 -11.3 -9. -1.5] """ y0 = np.asarray(y0) ode_dim = y0.shape[0] if y0.ndim > 0 else 1 order = prior.ordint proj_to_y = prior.proj2coord(0) zeros_shift = np.zeros(ode_dim) zeros_cov = np.zeros((ode_dim, ode_dim)) measmod = pnss.DiscreteLTIGaussian( proj_to_y, zeros_shift, zeros_cov, proc_noise_cov_cholesky=zeros_cov, forward_implementation="sqrt", backward_implementation="sqrt", ) # order + 1 would suffice in theory, 2*order + 1 is for good measure # (the "+1" is a safety factor for order=1) num_steps = 2 * order + 1 t_eval = np.arange(t0, t0 + (num_steps + 1) * h0, h0) sol = sci.solve_ivp( f, (t0, t0 + (num_steps + 1) * h0), y0=y0, atol=1e-12, rtol=1e-12, t_eval=t_eval, method=method, ) ts = sol.t[:num_steps] ys = sol.y[:, :num_steps].T initmean = initrv.mean.copy() initmean[0::(order + 1)] = y0 initmean[1::(order + 1)] = f(t0, y0) initcov_diag = np.diag(initrv.cov).copy() initcov_diag[0::(order + 1)] = SMALL_VALUE initcov_diag[1::(order + 1)] = SMALL_VALUE if df is not None: if order > 1: initmean[2::(order + 1)] = df(t0, y0) @ f(t0, y0) initcov_diag[2::(order + 1)] = SMALL_VALUE initcov = np.diag(initcov_diag) initcov_cholesky = np.diag(np.sqrt(initcov_diag)) initrv = pnrv.Normal(initmean, initcov, cov_cholesky=initcov_cholesky) kalman = pnfs.Kalman(prior, measmod, initrv) out = kalman.filtsmooth(ys, ts) estimated_initrv = out.state_rvs[0] return estimated_initrv
def test_correct_instantiation(self): """Test whether different variants of the normal distribution are instances of Normal.""" for mean, cov in self.normal_params: with self.subTest(): dist = rvs.Normal(mean=mean, cov=cov) self.assertIsInstance(dist, rvs.Normal)
def initialize_odefilter_with_taylormode(f, y0, t0, prior, initrv): """Initialize an ODE filter with Taylor-mode automatic differentiation. This requires JAX. For an explanation of what happens ``under the hood``, see [1]_. References ---------- .. [1] Krämer, N. and Hennig, P., Stable implementation of probabilistic ODE solvers, *arXiv:2012.10106*, 2020. The implementation is inspired by the implementation in https://github.com/jacobjinkelly/easy-neural-ode/blob/master/latent_ode.py Parameters ---------- f ODE vector field. y0 Initial value. t0 Initial time point. prior Prior distribution used for the ODE solver. For instance an integrated Brownian motion prior (``IBM``). initrv Initial random variable. Returns ------- Normal Estimated initial random variable. Compatible with the specified prior. Examples -------- >>> import sys, pytest >>> if sys.platform.startswith('win'): ... pytest.skip('this doctest does not work on Windows') >>> from dataclasses import astuple >>> from probnum.random_variables import Normal >>> from probnum.problems.zoo.diffeq import threebody_jax, vanderpol_jax >>> from probnum.statespace import IBM Compute the initial values of the restricted three-body problem as follows >>> f, t0, tmax, y0, df, *_ = astuple(threebody_jax()) >>> print(y0) [ 0.994 0. 0. -2.00158511] >>> prior = IBM(ordint=3, spatialdim=4) >>> initrv = Normal(mean=np.zeros(prior.dimension), cov=np.eye(prior.dimension)) >>> improved_initrv = initialize_odefilter_with_taylormode(f, y0, t0, prior, initrv) >>> print(prior.proj2coord(0) @ improved_initrv.mean) [ 0.994 0. 0. -2.00158511] >>> print(improved_initrv.mean) [ 9.94000000e-01 0.00000000e+00 -3.15543023e+02 0.00000000e+00 0.00000000e+00 -2.00158511e+00 0.00000000e+00 9.99720945e+04 0.00000000e+00 -3.15543023e+02 0.00000000e+00 6.39028111e+07 -2.00158511e+00 0.00000000e+00 9.99720945e+04 0.00000000e+00] Compute the initial values of the van-der-Pol oscillator as follows >>> f, t0, tmax, y0, df, *_ = astuple(vanderpol_jax()) >>> print(y0) [2. 0.] >>> prior = IBM(ordint=3, spatialdim=2) >>> initrv = Normal(mean=np.zeros(prior.dimension), cov=np.eye(prior.dimension)) >>> improved_initrv = initialize_odefilter_with_taylormode(f, y0, t0, prior, initrv) >>> print(prior.proj2coord(0) @ improved_initrv.mean) [2. 0.] >>> print(improved_initrv.mean) [ 2. 0. -2. 60. 0. -2. 60. -1798.] >>> print(improved_initrv.std) [0. 0. 0. 0. 0. 0. 0. 0.] """ try: import jax.numpy as jnp from jax.config import config from jax.experimental.jet import jet config.update("jax_enable_x64", True) except ImportError as err: raise ImportError( "Cannot perform Taylor-mode initialisation without optional " "dependencies jax and jaxlib. Try installing them via `pip install jax jaxlib`." ) from err order = prior.ordint def total_derivative(z_t): """Total derivative.""" z, t = jnp.reshape(z_t[:-1], z_shape), z_t[-1] dz = jnp.ravel(f(t, z)) dt = jnp.array([1.0]) dz_t = jnp.concatenate((dz, dt)) return dz_t z_shape = y0.shape z_t = jnp.concatenate((jnp.ravel(y0), jnp.array([t0]))) derivs = [] derivs.extend(y0) if order == 0: all_derivs = pnss.Integrator._convert_derivwise_to_coordwise( np.asarray(jnp.array(derivs)), ordint=0, spatialdim=len(y0)) return pnrv.Normal( np.asarray(all_derivs), cov=np.asarray(jnp.diag(jnp.zeros(len(derivs)))), cov_cholesky=np.asarray(jnp.diag(jnp.zeros(len(derivs)))), ) (dy0, [*yns]) = jet(total_derivative, (z_t, ), ((jnp.ones_like(z_t), ), )) derivs.extend(dy0[:-1]) if order == 1: all_derivs = pnss.Integrator._convert_derivwise_to_coordwise( np.asarray(jnp.array(derivs)), ordint=1, spatialdim=len(y0)) return pnrv.Normal( np.asarray(all_derivs), cov=np.asarray(jnp.diag(jnp.zeros(len(derivs)))), cov_cholesky=np.asarray(jnp.diag(jnp.zeros(len(derivs)))), ) for _ in range(1, order): (dy0, [*yns]) = jet(total_derivative, (z_t, ), ((dy0, *yns), )) derivs.extend(yns[-2][:-1]) all_derivs = pnss.Integrator._convert_derivwise_to_coordwise( jnp.array(derivs), ordint=order, spatialdim=len(y0)) return pnrv.Normal( np.asarray(all_derivs), cov=np.asarray(jnp.diag(jnp.zeros(len(derivs)))), cov_cholesky=np.asarray(jnp.diag(jnp.zeros(len(derivs)))), )
def test_sigpts_shape(self): sigpts = self.ut.sigma_points(pnrv.Normal(self.mean, self.covar)) self.assertEqual(sigpts.ndim, 2) self.assertEqual(sigpts.shape[0], 2 * self.ndim + 1) self.assertEqual(sigpts.shape[1], self.ndim)
def condition_state_on_rv(attained_rv, forwarded_rv, rv, gain): new_mean = rv.mean + gain @ (attained_rv.mean - forwarded_rv.mean) new_cov = rv.cov + gain @ (attained_rv.cov - forwarded_rv.cov) @ gain.T return pnrv.Normal(new_mean, new_cov)
def test_propagate_shape(self): sigpts = self.ut.sigma_points(pnrv.Normal(self.mean, self.covar)) propagated = self.ut.propagate(None, sigpts, lambda t, x: np.sin(x)) self.assertEqual(propagated.ndim, 2) self.assertEqual(propagated.shape[0], 2 * self.ndim + 1) self.assertEqual(propagated.shape[1], self.ndim)