def test_all_backward_realization_same_no_cache(self, some_normal_rv1,
some_normal_rv2,
diffusion):
"""Assert all implementations give the same output -- no gain or forwarded RV
passed."""
out_classic, _ = self.transition._backward_rv_classic(
randvars.Constant(some_normal_rv1.mean),
some_normal_rv2,
t=0.0,
_diffusion=diffusion,
)
out_sqrt, _ = self.transition._backward_rv_sqrt(
randvars.Constant(some_normal_rv1.mean),
some_normal_rv2,
t=0.0,
_diffusion=diffusion,
)
out_joseph, _ = self.transition._backward_rv_joseph(
randvars.Constant(some_normal_rv1.mean),
some_normal_rv2,
t=0.0,
_diffusion=diffusion,
)
# Classic -- sqrt
np.testing.assert_allclose(out_classic.mean, out_sqrt.mean)
np.testing.assert_allclose(out_classic.cov, out_sqrt.cov)
# Joseph -- sqrt
np.testing.assert_allclose(out_joseph.mean, out_sqrt.mean)
np.testing.assert_allclose(out_joseph.cov, out_sqrt.cov)
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_cov_linops(self):
# Ignore scalar support, because in this case, LinOps make no sense.
for supp in self.supports[1:]:
with self.subTest():
with config(matrix_free=True):
rv_lo = randvars.Constant(support=supp)
assert isinstance(rv_lo.cov, linops.LinearOperator)
with config(matrix_free=False):
rv_de = randvars.Constant(support=supp)
assert isinstance(rv_de.cov, np.ndarray)
def test_induced_solution_belief(rng: np.random.Generator):
"""Test whether a consistent belief over the solution is inferred from a belief over
the inverse."""
n = 5
A = randvars.Constant(random_spd_matrix(dim=n, rng=rng))
Ainv = randvars.Normal(
mean=linops.Scaling(factors=1 / np.diag(A.mean)),
cov=linops.SymmetricKronecker(linops.Identity(n)),
)
b = randvars.Constant(rng.normal(size=(n, 1)))
prior = LinearSystemBelief(A=A, Ainv=Ainv, x=None, b=b)
x_infer = Ainv @ b
np.testing.assert_allclose(prior.x.mean, x_infer.mean)
np.testing.assert_allclose(prior.x.cov.todense(), x_infer.cov.todense())
def test_dense_output(solvers, y, start_point, stop_point):
testsolver, scipysolver = solvers
# step has to be performed before dense-output can be computed
scipysolver.step()
# perform step of the same size
testsolver.step(
scipysolver.t_old,
scipysolver.t,
randvars.Constant(scipysolver.y_old),
)
testsolver_dense = testsolver.dense_output()
scipy_dense = scipysolver._dense_output_impl()
np.testing.assert_allclose(
testsolver_dense(scipysolver.t_old),
scipy_dense(scipysolver.t_old),
atol=1e-14,
rtol=1e-14,
)
np.testing.assert_allclose(
testsolver_dense(scipysolver.t),
scipy_dense(scipysolver.t),
atol=1e-14,
rtol=1e-14,
)
np.testing.assert_allclose(
testsolver_dense((scipysolver.t_old + scipysolver.t) / 2),
scipy_dense((scipysolver.t_old + scipysolver.t) / 2),
atol=1e-14,
rtol=1e-14,
)
def initialize(self, ivp):
"""Return t0 and y0 (for the solver, which might be different to ivp.y0) and
initialize the solver. Reset the solver when solving the ODE multiple times,
i.e. explicitly setting y_old, t, y and f to the respective initial values,
otherwise those are wrong when running the solver twice.
Returns
-------
self.ivp.t0: float
initial time point
self.ivp.initrv: randvars.RandomVariable
initial random variable
"""
self.solver = self.solver_type(ivp.f, ivp.t0, ivp.y0, ivp.tmax)
self.ivp = ivp
self.interpolants = []
self.solver.y_old = None
self.solver.t = self.ivp.t0
self.solver.y = self.ivp.y0
self.solver.f = self.solver.fun(self.solver.t, self.solver.y)
state = _odesolver_state.ODESolverState(
ivp=ivp,
rv=randvars.Constant(self.ivp.y0),
t=self.ivp.t0,
error_estimate=None,
reference_state=None,
)
return state
def test_conjugate_actions(policy: policies.ConjugateGradientPolicy,
state: LinearSolverState):
"""Tests whether actions generated by the policy are A-conjugate via a naive CG
implementation."""
A = state.problem.A
for _ in range(A.shape[1]):
# Action
s = policy(state)
state.action = s
# Observation
y = A @ s
state.observation = y
# Residual
r = state.residual
# Step size
alpha = np.linalg.norm(r)**2 / np.inner(s, y)
# Solution update
x = state.belief.x.mean
state.belief.x = randvars.Constant(x + alpha * s)
state.next_step()
actions = np.array(state.actions[:-1]).T
innerprods = actions.T @ A @ actions
np.testing.assert_allclose(innerprods,
np.diag(np.diag(innerprods)),
atol=1e-7,
rtol=1e7)
def __call__(
self, solver_state: "probnum.linalg.solvers.LinearSolverState"
) -> LinearSystemBelief:
proj_resid = solver_state.observation
# Compute gain and covariance update
action_A = solver_state.action.T @ solver_state.problem.A
cov_xy = solver_state.belief.x.cov @ action_A.T
gram = action_A @ cov_xy + self.noise_var
gram_pinv = 1.0 / gram if gram > 0.0 else 0.0
gain = cov_xy * gram_pinv
cov_update = np.outer(gain, cov_xy)
x = randvars.Normal(
mean=solver_state.belief.x.mean + gain * proj_resid,
cov=solver_state.belief.x.cov - cov_update,
)
if solver_state.belief.Ainv is None:
Ainv = randvars.Constant(cov_update)
else:
Ainv = solver_state.belief.Ainv + cov_update
return LinearSystemBelief(x=x,
A=solver_state.belief.A,
Ainv=Ainv,
b=solver_state.belief.b)
def test_non_randvar_arguments_raises_type_error():
A = np.eye(5)
Ainv = np.eye(5)
x = np.ones((5, 1))
b = np.ones((5, 1))
with pytest.raises(TypeError):
LinearSystemBelief(x=x)
with pytest.raises(TypeError):
LinearSystemBelief(Ainv=Ainv)
with pytest.raises(TypeError):
LinearSystemBelief(x=randvars.Constant(x), A=A)
with pytest.raises(TypeError):
LinearSystemBelief(x=randvars.Constant(x), b=b)
def posterior(stepsize, timespan):
"""Kalman smoothing posterior."""
initrv = randvars.Constant(20 * np.ones(2))
ivp = diffeq.ode.lotkavolterra(timespan, initrv)
f = ivp.rhs
t0, tmax = ivp.timespan
y0 = ivp.initrv.mean
return diffeq.probsolve_ivp(f, t0, tmax, y0, step=stepsize, adaptive=False)
def test_step_variables(solvers, y, start_point, stop_point):
testsolver, scipysolver = solvers
solver_y_new, solver_error_estimation = testsolver.step(
start_point, stop_point, randvars.Constant(y))
y_new, f_new = rk.rk_step(
scipysolver.fun,
start_point,
y,
scipysolver.f,
stop_point - start_point,
scipysolver.A,
scipysolver.B,
scipysolver.C,
scipysolver.K,
)
# error estimation is correct
scipy_error_estimation = scipysolver._estimate_error(
scipysolver.K, stop_point - start_point)
np.testing.assert_allclose(solver_error_estimation,
scipy_error_estimation,
atol=1e-14,
rtol=1e-14)
# locations are correct
np.testing.assert_allclose(testsolver.solver.t_old,
start_point,
atol=1e-14,
rtol=1e-14)
np.testing.assert_allclose(testsolver.solver.t,
stop_point,
atol=1e-14,
rtol=1e-14)
np.testing.assert_allclose(
testsolver.solver.h_previous,
stop_point - start_point,
atol=1e-14,
rtol=1e-14,
)
# evaluations are correct
np.testing.assert_allclose(testsolver.solver.y_old.mean,
y,
atol=1e-14,
rtol=1e-14)
np.testing.assert_allclose(testsolver.solver.y,
y_new,
atol=1e-14,
rtol=1e-14)
np.testing.assert_allclose(testsolver.solver.h_abs,
stop_point - start_point,
atol=1e-14,
rtol=1e-14)
np.testing.assert_allclose(testsolver.solver.f,
f_new,
atol=1e-14,
rtol=1e-14)
def test_all_backward_realization_same_with_cache(self, some_normal_rv1,
some_normal_rv2,
diffusion):
"""Assert all implementations give the same output -- gain and forwarded RV
passed."""
rv_forward, info = self.transition.forward_rv(some_normal_rv2,
0.0,
compute_gain=True,
_diffusion=diffusion)
gain = info["gain"]
out_classic, _ = self.transition._backward_rv_classic(
randvars.Constant(some_normal_rv1.mean),
some_normal_rv2,
rv_forwarded=rv_forward,
gain=gain,
t=0.0,
_diffusion=diffusion,
)
out_sqrt, _ = self.transition._backward_rv_sqrt(
randvars.Constant(some_normal_rv1.mean),
some_normal_rv2,
rv_forwarded=rv_forward,
gain=gain,
t=0.0,
_diffusion=diffusion,
)
out_joseph, _ = self.transition._backward_rv_joseph(
randvars.Constant(some_normal_rv1.mean),
some_normal_rv2,
rv_forwarded=rv_forward,
gain=gain,
t=0.0,
_diffusion=diffusion,
)
# Classic -- sqrt
np.testing.assert_allclose(out_classic.mean, out_sqrt.mean)
np.testing.assert_allclose(out_classic.cov, out_sqrt.cov)
# Joseph -- sqrt
np.testing.assert_allclose(out_joseph.mean, out_sqrt.mean)
np.testing.assert_allclose(out_joseph.cov, out_sqrt.cov)
def test_dimension_mismatch_raises_value_error():
"""Test whether mismatched components result in a ValueError."""
m, n, nrhs = 5, 3, 2
A = randvars.Normal(mean=np.ones((m, n)), cov=np.eye(m * n))
Ainv = A
x = randvars.Normal(mean=np.zeros((n, nrhs)), cov=np.eye(n * nrhs))
b = randvars.Constant(np.ones((m, nrhs)))
# A does not match b
with pytest.raises(ValueError):
LinearSystemBelief(A=A,
Ainv=Ainv,
x=x,
b=randvars.Constant(np.ones((m + 1, nrhs))))
# A does not match x
with pytest.raises(ValueError):
LinearSystemBelief(
A=A,
Ainv=Ainv,
x=randvars.Normal(mean=np.zeros((n + 1, nrhs)),
cov=np.eye((n + 1) * nrhs)),
b=b,
)
# x does not match b
with pytest.raises(ValueError):
LinearSystemBelief(
A=A,
Ainv=Ainv,
x=randvars.Normal(mean=np.zeros((n, nrhs + 1)),
cov=np.eye(n * (nrhs + 1))),
b=b,
)
# A does not match Ainv
with pytest.raises(ValueError):
LinearSystemBelief(
A=A,
Ainv=randvars.Normal(mean=np.ones((m + 1, n)),
cov=np.eye((m + 1) * n)),
x=x,
b=b,
)
def test_generate_shapes(times, test_ndim):
"""Output shapes are as expected."""
mocktrans = MockTransition(dim=test_ndim)
initrv = randvars.Constant(np.random.rand(test_ndim))
states, obs = pnss.generate_samples(mocktrans, mocktrans, initrv, times)
assert states.shape[0] == len(times)
assert states.shape[1] == test_ndim
assert obs.shape[0] == len(times)
assert obs.shape[1] == test_ndim
def attempt_step(self, state: _odesolver_state.ODESolverState,
dt: FloatLike):
"""Perform one ODE-step from start to stop and set variables to the
corresponding values.
To specify start and stop directly, rk_step() and not _step_impl() is used.
Parameters
----------
state
Current state of the ODE solver.
dt
Step-size.
Returns
-------
_odesolver_state.ODESolverState
New state.
"""
y_new, f_new = rk.rk_step(
self.solver.fun,
state.t,
state.rv.mean,
self.solver.f,
dt,
self.solver.A,
self.solver.B,
self.solver.C,
self.solver.K,
)
# Unnormalized error estimation is used as the error estimation is normalized in
# solve().
error_estimation = self.solver._estimate_error(self.solver.K, dt)
y_new_as_rv = randvars.Constant(y_new)
new_state = _odesolver_state.ODESolverState(
ivp=state.ivp,
rv=y_new_as_rv,
t=state.t + dt,
error_estimate=error_estimation,
reference_state=state.rv.mean,
)
# Update the solver settings. This part is copied from scipy's _step_impl().
self.solver.h_previous = dt
self.solver.y_old = state.rv.mean
self.solver.t_old = state.t
self.solver.t = state.t + dt
self.solver.y = y_new
self.solver.h_abs = dt
self.solver.f = f_new
return new_state
def test_step_execution(solvers):
testsolver, scipysolver = solvers
scipysolver.step()
# perform step of the same size
random_var, error_est = testsolver.step(
scipysolver.t_old,
scipysolver.t,
randvars.Constant(scipysolver.y_old),
)
np.testing.assert_allclose(scipysolver.y, random_var.mean)
def __call__(self, t: DenseOutputLocationArgType) -> DenseOutputValueType:
"""Evaluate the time-continuous solution at time t.
Parameters
----------
t
Location / time at which to evaluate the continuous ODE solution.
Returns
-------
randvars.RandomVariable or randvars._RandomVariableList
Estimate of the states at time ``t`` based on a fourth order polynomial.
"""
states = self.scipy_solution(t).T
if np.isscalar(t):
solution_as_rv = randvars.Constant(states)
else:
solution_as_rv = randvars._RandomVariableList(
[randvars.Constant(state) for state in states])
return solution_as_rv
def test_perfect_information(solver: ProbabilisticLinearSolver,
problem: problems.LinearSystem, ncols: int):
"""Test whether a solver given perfect information converges instantly."""
# Construct prior belief with perfect information
belief = beliefs.LinearSystemBelief(
x=randvars.Normal(mean=problem.solution,
cov=linops.Scaling(factors=0.0,
shape=(ncols, ncols))),
A=randvars.Constant(problem.A),
Ainv=randvars.Constant(np.linalg.inv(problem.A @ np.eye(ncols))),
)
# Run solver
belief, solver_state = solver.solve(prior=belief,
problem=problem,
rng=np.random.default_rng(1))
# Check for instant convergence
assert solver_state.step == 0
np.testing.assert_allclose(belief.x.mean, problem.solution)
def step(self, start: FloatArgType, stop: FloatArgType, current: randvars,
**kwargs):
"""Perform one ODE-step from start to stop and set variables to the
corresponding values.
To specify start and stop directly, rk_step() and not _step_impl() is used.
Parameters
----------
start : float
starting location of the step
stop : float
stopping location of the step
current : :obj:`list` of :obj:`RandomVariable`
current state of the ODE.
Returns
-------
random_var : randvars.RandomVariable
Estimated states of the discrete-time solution.
error_estimation : float
estimated error after having performed the step.
"""
y = current.mean
dt = stop - start
y_new, f_new = rk.rk_step(
self.solver.fun,
start,
y,
self.solver.f,
dt,
self.solver.A,
self.solver.B,
self.solver.C,
self.solver.K,
)
# Unnormalized error estimation is used as the error estimation is normalized in
# solve().
error_estimation = self.solver._estimate_error(self.solver.K, dt)
y_new_as_rv = randvars.Constant(y_new)
# Update the solver settings. This part is copied from scipy's _step_impl().
self.solver.h_previous = dt
self.solver.y_old = current
self.solver.t_old = start
self.solver.t = stop
self.solver.y = y_new
self.solver.h_abs = dt
self.solver.f = f_new
return y_new_as_rv, error_estimation
def test_step_variables(solvers, y, start_point, stop_point):
testsolver, scipysolver, ode = solvers
teststate = diffeq.ODESolverState(
ivp=ode,
rv=randvars.Constant(y),
t=start_point,
error_estimate=None,
reference_state=None,
)
testsolver.initialize(ode)
solver_y_new = testsolver.attempt_step(teststate, dt=stop_point - start_point)
y_new, f_new = rk.rk_step(
scipysolver.fun,
start_point,
y,
scipysolver.f,
stop_point - start_point,
scipysolver.A,
scipysolver.B,
scipysolver.C,
scipysolver.K,
)
# error estimation is correct
scipy_error_estimation = scipysolver._estimate_error(
scipysolver.K, stop_point - start_point
)
np.testing.assert_allclose(
solver_y_new.error_estimate, scipy_error_estimation, atol=1e-13, rtol=1e-13
)
# locations are correct
np.testing.assert_allclose(
testsolver.solver.t_old, start_point, atol=1e-13, rtol=1e-13
)
np.testing.assert_allclose(testsolver.solver.t, stop_point, atol=1e-13, rtol=1e-13)
np.testing.assert_allclose(
testsolver.solver.h_previous,
stop_point - start_point,
atol=1e-13,
rtol=1e-13,
)
# evaluations are correct
np.testing.assert_allclose(testsolver.solver.y_old, y, atol=1e-13, rtol=1e-13)
np.testing.assert_allclose(testsolver.solver.y, y_new, atol=1e-13, rtol=1e-13)
np.testing.assert_allclose(
testsolver.solver.h_abs, stop_point - start_point, atol=1e-13, rtol=1e-13
)
np.testing.assert_allclose(testsolver.solver.f, f_new, atol=1e-13, rtol=1e-13)
def test_non_two_dimensional_raises_value_error():
"""Test whether specifying higher-dimensional random variables raise a
ValueError."""
A = randvars.Constant(np.eye(5))
Ainv = randvars.Constant(np.eye(5))
x = randvars.Constant(np.ones((5, 1)))
b = randvars.Constant(np.ones((5, 1)))
# A.ndim == 3
with pytest.raises(ValueError):
LinearSystemBelief(A=A[:, None], Ainv=Ainv, x=x, b=b)
# Ainv.ndim == 3
with pytest.raises(ValueError):
LinearSystemBelief(A=A, Ainv=Ainv[:, None], x=x, b=b)
# x.ndim == 3
with pytest.raises(ValueError):
LinearSystemBelief(A=A, Ainv=Ainv, x=x[:, None], b=b)
# b.ndim == 3
with pytest.raises(ValueError):
LinearSystemBelief(A=A, Ainv=Ainv, x=x, b=b[:, None])
def test_sample_shapes(self):
"""Test whether samples have the correct shapes."""
for supp in self.supports:
for sample_size in [1, (), 10, (4, ), (3, 2)]:
with self.subTest():
s = randvars.Constant(support=supp).sample(
size=sample_size)
if sample_size == ():
self.assertEqual(np.shape(supp), np.shape(s))
elif isinstance(sample_size, tuple):
self.assertEqual(sample_size + np.shape(supp),
np.shape(s))
else:
self.assertEqual(tuple([sample_size, *np.shape(supp)]),
np.shape(s))
def test_generate_shapes(times, test_ndim, rng):
"""Output shapes are as expected."""
mocktrans = MockTransition(dim=test_ndim)
initrv = randvars.Constant(np.random.rand(test_ndim))
proc = randprocs.markov.MarkovProcess(
initarg=times[0], initrv=initrv, transition=mocktrans
)
states, obs = randprocs.markov.utils.generate_artificial_measurements(
rng, prior_process=proc, measmod=mocktrans, times=times
)
assert states.shape[0] == len(times)
assert states.shape[1] == test_ndim
assert obs.shape[0] == len(times)
assert obs.shape[1] == test_ndim
def test_step_execution(solvers):
testsolver, scipysolver, ode = solvers
scipysolver.step()
# perform step of the same size
teststate = diffeq.ODESolverState(
ivp=ode,
rv=randvars.Constant(scipysolver.y_old),
t=scipysolver.t_old,
error_estimate=None,
reference_state=None,
)
testsolver.initialize(ode)
dt = scipysolver.t - scipysolver.t_old
new_state = testsolver.attempt_step(teststate, dt)
np.testing.assert_allclose(scipysolver.y, new_state.rv.mean)
def from_callable(
cls,
input_dim: IntLike,
output_dim: IntLike,
transition_fun: Callable[[FloatLike, ArrayLike], ArrayLike],
transition_fun_jacobian: Callable[[FloatLike, ArrayLike], ArrayLike],
):
"""Turn a callable into a deterministic transition."""
return cls(
input_dim=input_dim,
output_dim=output_dim,
transition_fun=transition_fun,
transition_fun_jacobian=transition_fun_jacobian,
noise_fun=lambda t: randvars.Constant(np.zeros(output_dim)),
)
def __init__(self, solver: rk.RungeKutta):
self.solver = solver
self.interpolants = None
# ProbNum ODESolver needs an ivp
ivp = diffeq.IVP(
timespan=[self.solver.t, self.solver.t_bound],
initrv=randvars.Constant(self.solver.y),
rhs=self.solver._fun,
)
# Dopri853 as implemented in SciPy computes the dense output differently.
if isinstance(solver, rk.DOP853):
raise TypeError(
"Dense output interpolation of DOP853 is currently not supported. Choose a different RK-method."
)
super().__init__(ivp=ivp, order=solver.order)
def interpolate(
self,
t: FloatLike,
previous_index: Optional[FloatLike] = None,
next_index: Optional[FloatLike] = None,
):
# For the first state, no interpolation has to be performed.
if t == self.locations[0]:
return self.states[0]
if t == self.locations[-1]:
return self.states[-1]
interpolant = self.interpolants[previous_index]
relative_time = (
t - self.locations[previous_index]) * self.scales[previous_index]
previous_time = self.locations[previous_index]
evaluation = interpolant(previous_time + relative_time)
res_as_rv = randvars.Constant(evaluation)
return res_as_rv
def forward_realization(
self,
realization,
t,
dt=None,
compute_gain=False,
_diffusion=1.0,
_linearise_at=None,
) -> Tuple[randvars.Normal, Dict]:
"""Approximate forward-propagation of a realization of a random variable."""
compute_jacobian_at = (_linearise_at if _linearise_at is not None else
randvars.Constant(realization))
linearized_model = self.linearize(t=t, at_this_rv=compute_jacobian_at)
return linearized_model.forward_realization(
realization=realization,
t=t,
dt=dt,
compute_gain=compute_gain,
_diffusion=_diffusion,
)
def from_linop(
cls,
transition_matrix: LinearOperatorLike,
noise_mean: ArrayLike,
forward_implementation="classic",
backward_implementation="classic",
):
"""Turn a linear operator (or numpy array) into a noise-free transition."""
# Currently, this is only a numpy array.
# In the future, once linops are more widely adopted here, this will become
# a linop.
if transition_matrix.ndim != 2:
raise ValueError
return cls(
transition_matrix=transition_matrix,
noise=randvars.Constant(noise_mean),
forward_implementation=forward_implementation,
backward_implementation=backward_implementation,
)
def test_dense_output(solvers):
testsolver, scipysolver, ode = solvers
# perform steps of the same size
testsolver.initialize(ode)
scipysolver.step()
teststate = diffeq.ODESolverState(
ivp=ode,
rv=randvars.Constant(scipysolver.y_old),
t=scipysolver.t_old,
error_estimate=None,
reference_state=None,
)
state = testsolver.attempt_step(
state=teststate, dt=scipysolver.t - scipysolver.t_old
)
# sanity check: the steps are the same
# (this is contained in a different test already, but if this one
# does not work, the dense output test below is meaningless)
np.testing.assert_allclose(scipysolver.y, state.rv.mean)
testsolver_dense = testsolver.dense_output()
scipy_dense = scipysolver._dense_output_impl()
t_old = scipysolver.t_old
t = scipysolver.t
t_mid = (t_old + t) / 2.0
for time in [t_old, t, t_mid]:
test_dense = testsolver_dense(time)
ref_dense = scipy_dense(time)
np.testing.assert_allclose(
test_dense,
ref_dense,
atol=1e-13,
rtol=1e-13,
)
