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 attempt_step(self, state: _odesolver_state.ODESolverState, dt: FloatLike):
"""Perturb the original stopping point.
Perform one perturbed step and project the solution back to the original
stopping point.
Parameters
----------
state
Current state of the ODE solver.
dt
Step-size.
Returns
-------
_odesolver_state.ODESolverState
New state.
"""
noisy_step = self.perturb_step(self.rng, dt)
new_state = self.solver.attempt_step(state, noisy_step)
scale = noisy_step / dt
self.scales.append(scale)
t_new = state.t + dt
state = _odesolver_state.ODESolverState(
ivp=state.ivp,
rv=new_state.rv,
t=t_new,
error_estimate=new_state.error_estimate,
reference_state=new_state.reference_state,
)
return state
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 initialize(self, ivp):
self.information_operator.incorporate_ode(ode=ivp)
self.measurement_model = self.approx_strategy(
self.information_operator
).as_transition()
initrv = self.init_routine(
ivp=ivp,
prior_process=self.prior_process,
)
state = _odesolver_state.ODESolverState(
ivp=ivp, t=ivp.t0, rv=initrv, error_estimate=None, reference_state=None
)
return state
def attempt_step(self, state, dt):
r"""Gaussian IVP filter step as nonlinear Kalman filtering with zero data.
It goes as follows:
1. The current state :math:`x(t) \sim \mathcal{N}(m(t), P(t))` is split into a
deterministic component and a noisy component,
.. math::
x(t) = m(t) + p(t), \quad p(t) \sim \mathcal{N}(0, P(t))
which is required for accurate calibration and error estimation: in order to
only work with the local error, ODE solvers often assume that the error at the
previous step is zero.
2. The deterministic component is propagated through dynamics model and
measurement model
.. math::
\hat{z}(t+\Delta t)\sim\mathcal{N}(H\Phi(\Delta t)m(t),H Q(\Delta t) H^\top)
which is a random variable that estimates the expected local defect
:math:`\dot y - f(y)`.
3. The counterpart of :math:`\hat{z}` in the (much more likely) interpretation
of an erronous previous state (recall that :math:`\hat z` was based on the
interpretation of an error-free previous state) is computed as,
.. math::
z(t + \Delta t) \sim \mathcal{N}(\mathbb{E}[\hat{z}(t + \Delta t)],
\mathbb{C}[\hat{z}(t+\Delta t)]+H\Phi(\Delta t)P(t)\Phi(\Delta t)^\top
H^\top),
which acknowledges the covariance :math:`P(t)` of the previous state.
:math:`\mathbb{E}` is the mean, and :math:`\mathbb{C}` is the covariance.
Both :math:`z(t + \Delta t)` and :math:`\hat z(t + \Delta t)` give rise to a
reasonable diffusion_model estimate. Which one to use is handled by the
``Diffusion`` attribute of the solver. At this point already we can compute a
local error estimate of the current step.
4. Depending on the diffusion model, there are two options now:
4.1. For a piecewise constant diffusion, the covariances are calibrated
locally -- that is, at each step. In this case we update the predicted
covariance and measured covariance with the most recent diffusion estimate.
4.2. For a constant diffusion, the calibration happens post hoc, and the
only step that is carried out here is an assembly of the full predicted
random variable (up to now, only its parts were available).
5. With the results of either 4.1. or 4.2. (which both return a predicted RV and
a measured RV), we finally compute the Kalman update and return the result.
Recall that the error estimate has been computed in the third step.
"""
# Read off system matrices; required for calibration / error estimation
# Use only where a full call to forward_*() would be too costly.
# We use the mathematical symbol `Phi` (and later, `H`), because this makes
# it easier to read for us.
discrete_dynamics = self.prior_process.transition.discretise(dt)
Phi = discrete_dynamics.transition_matrix
# Split the current RV into a deterministic part and a noisy part.
# This split is necessary for efficient calibration; see docstring.
error_free_state = state.rv.mean.copy()
noisy_component = randvars.Normal(
mean=np.zeros(state.rv.shape),
cov=state.rv.cov.copy(),
cov_cholesky=state.rv.cov_cholesky.copy(),
)
# Compute the measurements for the error-free component
pred_rv_error_free, _ = discrete_dynamics.forward_realization(
error_free_state, t=state.t
)
linearized_observation = self.measurement_model.linearize(
state.t + dt, pred_rv_error_free
)
meas_rv_error_free, _ = linearized_observation.forward_rv(
pred_rv_error_free, t=state.t + dt
)
H = linearized_observation.transition_matrix_fun(t=state.t + dt)
# Compute the measurements for the full components.
# Since the means of noise-free and noisy measurements coincide,
# we manually update only the covariance.
# The first two are only matrix square-roots and will be turned into proper
# Cholesky factors below.
pred_sqrtm = Phi @ noisy_component.cov_cholesky
meas_sqrtm = H @ pred_sqrtm
full_meas_cov_cholesky = utils.linalg.cholesky_update(
meas_rv_error_free.cov_cholesky, meas_sqrtm
)
full_meas_cov = full_meas_cov_cholesky @ full_meas_cov_cholesky.T
meas_rv = randvars.Normal(
mean=meas_rv_error_free.mean,
cov=full_meas_cov,
cov_cholesky=full_meas_cov_cholesky,
)
# Estimate local diffusion_model and error
local_diffusion = self.diffusion_model.estimate_locally(
meas_rv=meas_rv,
meas_rv_assuming_zero_previous_cov=meas_rv_error_free,
t=state.t + dt,
)
local_errors = np.sqrt(local_diffusion) * np.sqrt(
np.diag(meas_rv_error_free.cov)
)
proj = self.prior_process.transition.proj2coord(
coord=self._reference_coordinates
)
reference_values = np.abs(proj @ state.rv.mean)
# Overwrite the acceptance/rejection step here, because we need control over
# Appending or not appending the diffusion (and because the computations below
# are sufficiently costly such that skipping them here will have a positive
# impact).
internal_norm = self.steprule.errorest_to_norm(
errorest=local_errors,
reference_state=reference_values,
)
if not self.steprule.is_accepted(internal_norm):
t_new = state.t + dt
new_rv = randvars.Normal(
mean=state.rv.mean.copy(),
cov=state.rv.cov.copy(),
cov_cholesky=state.rv.cov_cholesky.copy(),
)
state = _odesolver_state.ODESolverState(
ivp=state.ivp,
rv=np.nan * new_rv,
t=t_new,
error_estimate=local_errors,
reference_state=reference_values,
)
return state
# Now we can be certain that the step will be accepted. In this case
# we update the diffusion model and continue the rest of the iteration.,
self.diffusion_model.update_in_place(local_diffusion, t=state.t)
if self._calibrate_at_each_step:
# With the updated diffusion, we need to re-compute the covariances of the
# predicted RV and measured RV.
# The resulting predicted and measured RV are overwritten herein.
full_pred_cov_cholesky = utils.linalg.cholesky_update(
np.sqrt(local_diffusion) * pred_rv_error_free.cov_cholesky, pred_sqrtm
)
full_pred_cov = full_pred_cov_cholesky @ full_pred_cov_cholesky.T
pred_rv = randvars.Normal(
mean=pred_rv_error_free.mean,
cov=full_pred_cov,
cov_cholesky=full_pred_cov_cholesky,
)
full_meas_cov_cholesky = utils.linalg.cholesky_update(
np.sqrt(local_diffusion) * meas_rv_error_free.cov_cholesky, meas_sqrtm
)
full_meas_cov = full_meas_cov_cholesky @ full_meas_cov_cholesky.T
meas_rv = randvars.Normal(
mean=meas_rv_error_free.mean,
cov=full_meas_cov,
cov_cholesky=full_meas_cov_cholesky,
)
else:
# Combine error-free and noisy-component predictions into a full prediction.
# This has not been assembled as a standalone random variable yet,
# but is needed for the update below.
# (The measurement has been updated already.)
full_pred_cov_cholesky = utils.linalg.cholesky_update(
pred_rv_error_free.cov_cholesky, pred_sqrtm
)
full_pred_cov = full_pred_cov_cholesky @ full_pred_cov_cholesky.T
pred_rv = randvars.Normal(
mean=pred_rv_error_free.mean,
cov=full_pred_cov,
cov_cholesky=full_pred_cov_cholesky,
)
# Gain needs manual catching up, too. Use it to compute the update
crosscov = full_pred_cov @ H.T
gain = scipy.linalg.cho_solve((meas_rv.cov_cholesky, True), crosscov.T).T
zero_data = np.zeros(meas_rv.mean.shape)
filt_rv, _ = self.measurement_model.backward_realization(
zero_data, pred_rv, rv_forwarded=meas_rv, gain=gain
)
t_new = state.t + dt
state = _odesolver_state.ODESolverState(
ivp=state.ivp,
rv=filt_rv,
t=t_new,
error_estimate=local_errors,
reference_state=reference_values,
)
return state