Refactor UKF Components to use linearize()

src/probnum/diffeq/odefilter/_odefilter.py

@@ -189,9 +189,8 @@ def attempt_step(self, state, dt):
         # 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.
-        # The system matrix H of the measurement model can be accessed after the first forward_*,
-        # therefore we read it off further below.
         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.
@@ -204,15 +203,16 @@ def attempt_step(self, state, dt):
         )
 
         # Compute the measurements for the error-free component
-        pred_rv_error_free, _ = self.prior_process.transition.forward_realization(
-            error_free_state, t=state.t, dt=dt
+        pred_rv_error_free, _ = discrete_dynamics.forward_realization(
+            error_free_state, t=state.t
         )
-        meas_rv_error_free, _ = self.measurement_model.forward_rv(
-            pred_rv_error_free, t=state.t + dt
+        linearized_observation = self.measurement_model.linearize(
+            state.t + dt, pred_rv_error_free
         )
-        H = self.measurement_model.linearized_model.transition_matrix_fun(
-            t=state.t + dt
+        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,

src/probnum/filtsmooth/gaussian/_kalmanposterior.py

@@ -20,7 +20,6 @@
     approx.DiscreteUKFComponent,
     randprocs.markov.continuous.LinearSDE,
     approx.ContinuousEKFComponent,
-    approx.ContinuousUKFComponent,
 ]
 """Any linear(ized) transition can define an (approximate) Gauss-Markov prior."""
 

src/probnum/filtsmooth/gaussian/approx/__init__.py

@@ -1,26 +1,17 @@
 """Approximate Gaussian filtering and smoothing."""
 
-from ._extendedkalman import ContinuousEKFComponent, DiscreteEKFComponent, EKFComponent
-from ._unscentedkalman import ContinuousUKFComponent, DiscreteUKFComponent, UKFComponent
-from ._unscentedtransform import UnscentedTransform
+from ._extendedkalman import ContinuousEKFComponent, DiscreteEKFComponent
+from ._unscentedkalman import DiscreteUKFComponent
 
 # Public classes and functions. Order is reflected in documentation.
 __all__ = [
-    "EKFComponent",
     "ContinuousEKFComponent",
     "DiscreteEKFComponent",
-    "UKFComponent",
-    "ContinuousUKFComponent",
     "DiscreteUKFComponent",
-    "UnscentedTransform",
 ]
 
 # Set correct module paths (for superclasses).
 # Corrects links and module paths in documentation.
-EKFComponent.__module__ = "probnum.filtsmooth.gaussian.approx"
 ContinuousEKFComponent.__module__ = "probnum.filtsmooth.gaussian.approx"
 DiscreteEKFComponent.__module__ = "probnum.filtsmooth.gaussian.approx"
-UKFComponent.__module__ = "probnum.filtsmooth.gaussian.approx"
-ContinuousUKFComponent.__module__ = "probnum.filtsmooth.gaussian.approx"
 DiscreteUKFComponent.__module__ = "probnum.filtsmooth.gaussian.approx"
-UnscentedTransform.__module__ = "probnum.filtsmooth.gaussian.approx"

src/probnum/filtsmooth/gaussian/approx/_extendedkalman.py

@@ -1,130 +1,13 @@
-"""Gaussian filtering and smoothing based on making intractable quantities tractable
-through Taylor-method approximations, e.g. linearization."""
-
-import abc
-from typing import Dict, Tuple
+"""Extended Kalman filtering."""
 
 from probnum import randprocs, randvars
 
-
-class EKFComponent(abc.ABC):
-    """Interface for extended Kalman filtering components."""
-
-    def __init__(
-        self,
-        non_linear_model,
-    ) -> None:
-
-        self.non_linear_model = non_linear_model
-
-        # Will be constructed later
-        self.linearized_model = None
-
-    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)
-        )
-        self.linearized_model = self.linearize(at_this_rv=compute_jacobian_at)
-        return self.linearized_model.forward_realization(
-            realization=realization,
-            t=t,
-            dt=dt,
-            compute_gain=compute_gain,
-            _diffusion=_diffusion,
-        )
-
-    def forward_rv(
-        self,
-        rv,
-        t,
-        dt=None,
-        compute_gain=False,
-        _diffusion=1.0,
-        _linearise_at=None,
-    ) -> Tuple[randvars.Normal, Dict]:
-        """Approximate forward-propagation of a random variable."""
-
-        compute_jacobian_at = _linearise_at if _linearise_at is not None else rv
-        self.linearized_model = self.linearize(at_this_rv=compute_jacobian_at)
-        return self.linearized_model.forward_rv(
-            rv=rv,
-            t=t,
-            dt=dt,
-            compute_gain=compute_gain,
-            _diffusion=_diffusion,
-        )
-
-    def backward_realization(
-        self,
-        realization_obtained,
-        rv,
-        rv_forwarded=None,
-        gain=None,
-        t=None,
-        dt=None,
-        _diffusion=1.0,
-        _linearise_at=None,
-    ):
-        """Approximate backward-propagation of a realization of a random variable."""
-
-        return self._backward_realization_via_backward_rv(
-            realization_obtained,
-            rv=rv,
-            rv_forwarded=rv_forwarded,
-            gain=gain,
-            t=t,
-            dt=dt,
-            _diffusion=_diffusion,
-            _linearise_at=_linearise_at,
-        )
-
-    def backward_rv(
-        self,
-        rv_obtained,
-        rv,
-        rv_forwarded=None,
-        gain=None,
-        t=None,
-        dt=None,
-        _diffusion=1.0,
-        _linearise_at=None,
-    ):
-        """Approximate backward-propagation of a random variable."""
-
-        compute_jacobian_at = _linearise_at if _linearise_at is not None else rv
-        self.linearized_model = self.linearize(at_this_rv=compute_jacobian_at)
-        return self.linearized_model.backward_rv(
-            rv_obtained=rv_obtained,
-            rv=rv,
-            rv_forwarded=rv_forwarded,
-            gain=gain,
-            t=t,
-            dt=dt,
-            _diffusion=_diffusion,
-        )
-
-    @abc.abstractmethod
-    def linearize(
-        self, at_this_rv: randvars.RandomVariable
-    ) -> randprocs.markov.Transition:
-        """Linearize the transition and make it tractable."""
-        raise NotImplementedError
+from ._interface import _LinearizationInterface
 
 
 # Order of inheritance matters, because forward and backward
 # are defined in EKFComponent, and must not be inherited from SDE.
-class ContinuousEKFComponent(EKFComponent, randprocs.markov.continuous.SDE):
+class ContinuousEKFComponent(_LinearizationInterface, randprocs.markov.continuous.SDE):
     """Continuous-time extended Kalman filter transition.
 
     Parameters
@@ -163,15 +46,15 @@ def __init__(
             dispersion_function=non_linear_model.dispersion_function,
             drift_jacobian=non_linear_model.drift_jacobian,
         )
-        EKFComponent.__init__(self, non_linear_model=non_linear_model)
+        _LinearizationInterface.__init__(self, non_linear_model=non_linear_model)
 
         self.mde_atol = mde_atol
         self.mde_rtol = mde_rtol
         self.mde_solver = mde_solver
 
         self.forward_implementation = forward_implementation
 
-    def linearize(self, at_this_rv: randvars.Normal):
+    def linearize(self, t, at_this_rv: randvars.Normal):
         """Linearize the drift function with a first order Taylor expansion."""
 
         g = self.non_linear_model.drift_function
@@ -202,7 +85,9 @@ def dispersion_matrix_function(t):
         )
 
 
-class DiscreteEKFComponent(EKFComponent, randprocs.markov.discrete.NonlinearGaussian):
+class DiscreteEKFComponent(
+    _LinearizationInterface, randprocs.markov.discrete.NonlinearGaussian
+):
     """Discrete extended Kalman filter transition."""
 
     def __init__(
@@ -220,12 +105,12 @@ def __init__(
             noise_fun=non_linear_model.noise_fun,
             transition_fun_jacobian=non_linear_model.transition_fun_jacobian,
         )
-        EKFComponent.__init__(self, non_linear_model=non_linear_model)
+        _LinearizationInterface.__init__(self, non_linear_model=non_linear_model)
 
         self.forward_implementation = forward_implementation
         self.backward_implementation = backward_implementation
 
-    def linearize(self, at_this_rv: randvars.Normal):
+    def linearize(self, t, at_this_rv: randvars.Normal):
         """Linearize the dynamics function with a first order Taylor expansion."""
 
         g = self.non_linear_model.transition_fun

src/probnum/filtsmooth/gaussian/approx/_interface.py

@@ -0,0 +1,115 @@
+"""Temporary interface."""
+
+from typing import Dict, Tuple
+
+from probnum import randprocs, randvars
+
+
+class _LinearizationInterface:
+    """Interface for extended Kalman filtering components."""
+
+    def __init__(
+        self,
+        non_linear_model,
+    ) -> None:
+
+        self.non_linear_model = non_linear_model
+
+    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 forward_rv(
+        self,
+        rv,
+        t,
+        dt=None,
+        compute_gain=False,
+        _diffusion=1.0,
+        _linearise_at=None,
+    ) -> Tuple[randvars.Normal, Dict]:
+        """Approximate forward-propagation of a random variable."""
+
+        compute_jacobian_at = _linearise_at if _linearise_at is not None else rv
+        linearized_model = self.linearize(t=t, at_this_rv=compute_jacobian_at)
+        return linearized_model.forward_rv(
+            rv=rv,
+            t=t,
+            dt=dt,
+            compute_gain=compute_gain,
+            _diffusion=_diffusion,
+        )
+
+    def backward_realization(
+        self,
+        realization_obtained,
+        rv,
+        rv_forwarded=None,
+        gain=None,
+        t=None,
+        dt=None,
+        _diffusion=1.0,
+        _linearise_at=None,
+    ):
+        """Approximate backward-propagation of a realization of a random variable."""
+        return self._backward_realization_via_backward_rv(
+            realization_obtained,
+            rv=rv,
+            rv_forwarded=rv_forwarded,
+            gain=gain,
+            t=t,
+            dt=dt,
+            _diffusion=_diffusion,
+            _linearise_at=_linearise_at,
+        )
+
+    def backward_rv(
+        self,
+        rv_obtained,
+        rv,
+        rv_forwarded=None,
+        gain=None,
+        t=None,
+        dt=None,
+        _diffusion=1.0,
+        _linearise_at=None,
+    ):
+        """Approximate backward-propagation of a random variable."""
+
+        compute_jacobian_at = _linearise_at if _linearise_at is not None else rv
+        linearized_model = self.linearize(t=t, at_this_rv=compute_jacobian_at)
+        return linearized_model.backward_rv(
+            rv_obtained=rv_obtained,
+            rv=rv,
+            rv_forwarded=rv_forwarded,
+            gain=gain,
+            t=t,
+            dt=dt,
+            _diffusion=_diffusion,
+        )
+
+    def linearize(
+        self, t, at_this_rv: randvars.RandomVariable
+    ) -> randprocs.markov.Transition:
+        """Linearize the transition and make it tractable."""
+        raise NotImplementedError