Implementation of Stochastic Integration Filter

stonesoup/functions/__init__.py

@@ -3,6 +3,7 @@
 import warnings
 
 import numpy as np
+from scipy.stats import ortho_group
 
 from ..types.numeric import Probability
 from ..types.array import StateVector, StateVectors, CovarianceMatrix
@@ -955,3 +956,148 @@
     cross_covar = cross_covar.view(CovarianceMatrix)
 
     return mean, covar, cross_covar, cubature_points_t
+
+
+def stochasticCubatureRulePoints(nx, order):
+    """Stochstic cubature rule points
+
+    computation of cubature points and weights for the stochastic integration
+
+    Parameters
+    ==========
+    integer
+        number of points
+
+    Returns
+    =======
+    numpy.ndarray dim x nx
+        Matrix of sigma points
+    """
+
+    if order == 1:
+        X = np.random.randn(nx, 1)
+        SCRSigmaPoints = np.concatenate((X, -X), axis=1)
+        weights = np.array([0.5, 0.5])
+    elif order == 3:
+        CRSigmaPoints = np.concatenate(
+            (np.zeros((nx, 1)), np.eye(nx), -np.eye(nx)), axis=1
+        )
+        rho = np.sqrt(np.random.chisquare(nx + 2))
+        Q = ortho_group.rvs(nx)
+        SCRSigmaPoints = Q * rho @ CRSigmaPoints
+        weights = np.insert(0.5 * np.ones(2 * nx) / rho**2, 0, (1 - nx / rho**2))
+
+    elif order == 5:
+        # generating random values
+        r = np.sqrt(np.random.chisquare(2 * nx + 7))
+
+        q = np.random.beta(nx + 2, 3 / 2)
+
+        rho = r * np.sin(np.arcsin(q) / 2)
+        delta = r * np.cos(np.arcsin(q) / 2)
+
+        # calculating weights
+        c1up = nx + 2 - delta**2
+        c1do = rho**2 * (rho**2 - delta**2)
+        c2up = nx + 2 - rho**2
+        c2do = delta**2 * (delta**2 - rho**2)
+        cdo = 2 * (nx + 1) ** 2 * (nx + 2)
+        c3 = (7 - nx) * nx**2
+        c4 = 4 * (nx - 1) ** 2
+        coef1 = c1up * c3 / cdo / c1do
+        coef2 = c2up * c3 / cdo / c2do
+        coef3 = c1up * c4 / cdo / c1do
+        coef4 = c2up * c4 / cdo / c2do
+
+        pom = np.concatenate(
+            (
+                np.ones(2 * nx + 2) * coef1,
+                np.ones(2 * nx + 2) * coef2,
+                np.ones(nx * (nx + 1)) * coef3,
+                np.ones(nx * (nx + 1)) * coef4,
+            ),
+            axis=0,
+        )
+        weights = np.insert(
+            pom, 0, (1 - nx * (rho**2 + delta**2 - nx - 2) / (rho**2 * delta**2))
+        )
+
+        # Calculating sigma points
+        Q = ortho_group.rvs(nx)
+        v = np.zeros((nx, nx + 1))
+        i_vals, j_vals = np.triu_indices(nx + 1, k=1)
+        v[i_vals, i_vals] = np.sqrt((nx+1) * (nx-i_vals) / (nx * (nx-i_vals+1)))
+        v[i_vals, j_vals] = -np.sqrt((nx+1) / ((nx-i_vals) * nx * (nx-i_vals+1)))
+        v = Q @ v
+        i_vals, j_vals = np.tril_indices(nx + 1, k=-1)
+        comb_v = v[:, i_vals] + v[:, j_vals]
+        y = comb_v / np.linalg.norm(comb_v, axis=0)
+
+        SCRSigmaPoints = np.block(
+            [
+                np.zeros((nx, 1)),
+                -rho * v,
+                rho * v,
+                -delta * v,
+                +delta * v,
+                -rho * y,
+                rho * y,
+                -delta * y,
+                delta * y,
+            ]
+        )
+    else:
+        raise ValueError("This order of SIF is not supported")
+
+    return (SCRSigmaPoints, weights)
+
+
+def cubPointsAndTransfer(nx, order, sqrtCov, mean, transFunct, state):
+    r""" Calculates cubature points for stochastic integration filter and
+    puts them through given function (measurement/dynamics)
+
+    Parameters
+    ==========
+    nx : integer
+        dimension
+    order : integer
+        order of SIF rule
+    sqrtCov : array of float 64 [nx x nx]
+        square root of covariance
+    mean : array of int 64 [nx x 1]
+        mean
+    transFunct : function
+        function to transfer state vectors
+    state : state
+        whole state
+
+    Returns
+    =======
+    points : numpy.ndarray nx x number of points (based on order and dim)
+        cubature points
+    w : numpy.ndarray number of points x
+        weights
+    trsfPoints : numpy.ndarray nx x number of points (based on order and dim)
+        cubature transformed points
+    """
+
+    # -- cubature points and weights computation (for standard normal PDF)
+    SCRSigmaPoints, w = stochasticCubatureRulePoints(nx, order)
+
+    # -- points transformation for given filtering mean and covariance matrix
+    points = sqrtCov @ SCRSigmaPoints + mean
+
+    # -- points transformation through the function
+    sigma_points_states = []
+    for point in points.T:
+        state_copy = copy.copy(state)
+        state_copy.state_vector = StateVector(point)
+        sigma_points_states.append(state_copy)
+        trsfPoints = StateVectors(
+            [
+                transFunct(sigma_points_state)
+                for sigma_points_state in sigma_points_states
+            ]
+        )
+
+    return points, w, trsfPoints

stonesoup/functions/tests/test_functions.py

@@ -7,7 +7,7 @@
 from .. import (
     cholesky_eps, jacobian, gm_reduce_single, mod_bearing, mod_elevation, gauss2sigma,
     rotx, roty, rotz, cart2sphere, cart2angles, pol2cart, sphere2cart, dotproduct, gm_sample,
-    gauss2cubature, cubature2gauss, cubature_transform)
+    gauss2cubature, cubature2gauss, cubature_transform, stochasticCubatureRulePoints)
 from ...types.array import StateVector, StateVectors, Matrix, CovarianceMatrix
 from ...types.state import State, GaussianState
 
@@ -380,3 +380,29 @@ def identity_function(inpu):
     assert np.allclose(outcovar, instate.covar)
     assert np.allclose(mean, instate.state_vector)
     assert np.allclose(covar, instate.covar)
+
+
+@pytest.mark.parametrize(
+    "order, nx",
+    [
+        (3, 3),
+        (5, 4),
+        (1, 2)
+    ]
+)
+def test_stochastic_integration(order, nx):
+    points, weights = stochasticCubatureRulePoints(nx, order)
+    # Mean
+    assert np.allclose(np.average(points, weights=weights, axis=1),
+                       0, atol=1e-5)
+    # Weights
+    assert np.isclose(np.sum(weights), 1, atol=1e-5)
+    if order != 1:  # For order 1 it does not make sense to check variance of points
+        # Covariance
+        var = ((weights * points) @ points.T)
+        # Check if diagonal elements are close to 1
+        diagonal_elements = np.diag(var)
+        assert np.allclose(diagonal_elements, 1, atol=1e-5)
+        # Check if off-diagonal elements are close to 0
+        off_diagonal_elements = var[~np.eye(nx, dtype=bool)]
+        assert np.allclose(off_diagonal_elements, 0, atol=1e-5)

stonesoup/predictor/kalman.py

@@ -3,6 +3,7 @@
 import numpy as np
 import scipy.linalg as la
 
+from ..types.array import CovarianceMatrix, StateVector
 from .base import Predictor
 from ._utils import predict_lru_cache
 from ..base import Property
@@ -12,12 +13,13 @@
 from ..models.transition.linear import LinearGaussianTransitionModel
 from ..models.control import ControlModel
 from ..models.control.linear import LinearControlModel
-from ..functions import gauss2sigma, unscented_transform, cubature_transform
+from ..functions import (gauss2sigma, unscented_transform, cubature_transform,
+                         cubPointsAndTransfer)
 
 
 class KalmanPredictor(Predictor):
     r"""A predictor class which forms the basis for the family of Kalman
-    predictors. This class also serves as the (specific) Kalman Filter
+    predictors. This class also servas the (specific) Kalman Filter
     :class:`~.Predictor` class. Here transition and control models must be linear:
 
     .. math::
@@ -563,3 +565,157 @@ def predict(self, prior, timestamp=None, control_input=None, **kwargs):
 
         return Prediction.from_state(prior, x_pred, p_pred, timestamp=timestamp,
                                      transition_model=self.transition_model, prior=prior)
+
+
+class StochasticIntegrationPredictor(KalmanPredictor):
+    """Stochastic Integration Kalman Filter class
+
+    The state prediction of nonlinear models is accomplished by the stochastic
+    integration approximation.
+    """
+
+    transition_model: TransitionModel = Property(doc="The transition model to be used.")
+    control_model: ControlModel = Property(
+        default=None,
+        doc="The control model to be used. Default `None` where the predictor "
+        "will create a zero-effect linear :class:`~.ControlModel`.",
+    )
+    Nmax: float = Property(default=10, doc="maximal number of iterations of SIR")
+    Nmin: float = Property(
+        default=5,
+        doc="minimal number of iterations of stochastic integration rule (SIR)",
+    )
+    Eps: float = Property(default=5e-3, doc="allowed threshold for integration error")
+    SIorder: float = Property(
+        default=5, doc="order of SIR (orders 1, 3, 5 are currently supported)"
+    )
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        self._time_interval = None
+
+    def _transition_and_control_function(self, prior_state, **kwargs):
+        r"""Returns the result of applying the transition and control functions
+        for the unscented transform
+
+        Parameters
+        ----------
+        prior_state_vector : :class:`~.State`
+            Prior state vector
+        **kwargs : various, optional
+            These are passed to :class:`~.TransitionModel.function`
+
+        Returns
+        -------
+        : :class:`numpy.ndarray`
+            The combined, noiseless, effect of applying the transition and
+            control
+        """
+
+        return self.transition_model.function(
+            prior_state, **kwargs
+        ) + self.control_model.function(**kwargs)
+
+    @predict_lru_cache()
+    def predict(self, prior, timestamp=None, control_input=None, **kwargs):
+        r"""The stochastic integration version of the predict step
+
+        Parameters
+        ----------
+        prior : :class:`~.State`
+            Prior state, :math:`\mathbf{x}_{k-1}`
+        timestamp : :class:`datetime.datetime`
+            Time to transit to (:math:`k`)
+        control_input: :class:`~.State`
+            Control input vector, :math:`\mathbf{u}_k`
+        **kwargs : various, optional
+            These are passed to :meth:`~.TransitionModel.covar`
+
+        Returns
+        -------
+        : :class:`~.GaussianStatePrediction`
+            The predicted state :math:`\mathbf{x}_{k|k-1}` and the predicted
+            state covariance :math:`P_{k|k-1}`
+        """
+
+        nx = np.size(prior.mean, 0)
+        Sp = np.linalg.cholesky(prior.covar)
+        Ix = np.zeros((nx, 1))
+        Vx = np.zeros((nx, 1))
+        IPx = np.zeros((nx, nx))
+        VPx = np.zeros((nx, nx))
+        # Iteration Count
+        N = 0
+
+        # Get the prediction interval
+        predict_over_interval = self._predict_over_interval(prior, timestamp)
+
+        # This ensures that function passed to unscented transform has the
+        # correct time interval
+        transition_and_control_function = partial(
+            self._transition_and_control_function,
+            control_input=control_input,
+            time_interval=predict_over_interval,
+        )
+
+        # - SIR recursion for state predictive moments computation
+        # -- until either max iterations or threshold is reached
+        while N < self.Nmin or np.all([N < self.Nmax,
+                                       np.any([(np.linalg.norm(Vx) > self.Eps),]),]):
+            N += 1
+
+            # -- cubature points and weights computation (for standard normal PDF)
+            # -- points transformation for given filtering mean and covariance matrix
+            xpoints, w, fpoints = cubPointsAndTransfer(nx, self.SIorder, Sp,
+                                                       prior.mean,
+                                                       transition_and_control_function,
+                                                       prior)
+            # Stochastic integration rule for predictive state mean and
+            # covariance matrix
+            SumRx = np.average(fpoints, axis=1, weights=w)
+
+            # --- update mean Ix
+            Dx = (SumRx - Ix) / N
+            Ix += Dx
+            Vx = (N - 2) * Vx / N + Dx ** 2
+
+        xp = Ix
+
+        N = 0
+        # - SIR recursion for state predictive moments computation
+        # -- until max iterations are reached or threshold is reached
+        while N < self.Nmin or np.all([N < self.Nmax, np.any([(np.linalg.norm(VPx) > self.Eps)])]):
+            N += 1
+            # -- cubature points and weights computation (for standard normal PDF)
+            # -- points transformation for given filtering mean and covariance matrix
+            xpoints, w, fpoints = cubPointsAndTransfer(nx, self.SIorder, Sp,
+                                                       prior.mean,
+                                                       transition_and_control_function,
+                                                       prior)
+
+            # -- stochastic integration rule for predictive state mean and covariance
+            #    matrix
+            fpoints_diff = fpoints - xp
+            SumRPx = fpoints_diff @ np.diag(w) @ fpoints_diff.T
+
+            # --- update covariance matrix IPx
+            DPx = (SumRPx - IPx) / N
+            IPx += DPx
+            VPx = (N - 2) * VPx / N + DPx ** 2
+
+        Q = self.transition_model.covar(time_interval=predict_over_interval, **kwargs)
+
+        Pp = IPx + Q + np.diag(Vx.ravel())
+        Pp = (Pp + Pp.T) / 2
+        Pp = Pp.astype(np.float64)
+
+        # and return a Gaussian state based on these parameters
+        return Prediction.from_state(
+            prior,
+            xp.view(StateVector),
+            Pp.view(CovarianceMatrix),
+            timestamp=timestamp,
+            transition_model=self.transition_model,
+            prior=prior,
+        )