Support regulariser with multi model particle filter

stonesoup/regulariser/particle.py

@@ -4,10 +4,11 @@
 from typing import Sequence, Callable
 
 from .base import Regulariser
+from ..base import Property
 from ..functions import cholesky_eps
-from ..types.state import ParticleState
 from ..models.transition import TransitionModel
-from ..base import Property
+from ..predictor.particle import MultiModelPredictor
+from ..types.state import ParticleState
 
 
 class MCMCRegulariser(Regulariser):
@@ -43,7 +44,22 @@
             "object and return an array-like object of logical indices. "
     )
 
-    def regularise(self, prior, posterior):
+    def _transition_prior(self, prior, posterior, **kwargs):
+        hypothesis = posterior.hypothesis[0] if isinstance(posterior.hypothesis, Sequence) \
+            else posterior.hypothesis
+        transition_model = hypothesis.prediction.transition_model
+        if not transition_model:
+            transition_model = self.transition_model
+
+        if transition_model:
+            transitioned_prior = copy.copy(prior)
+            transitioned_prior.state_vector = transition_model.function(
+                prior, noise=False, time_interval=posterior.timestamp - prior.timestamp)
+            return transitioned_prior
+        else:
+            return prior
+
+    def regularise(self, prior, posterior, **kwargs):
         """Regularise the particles
 
         Parameters
@@ -65,19 +81,8 @@
         if not isinstance(prior, ParticleState):
             raise TypeError('Only ParticleState type is supported!')
 
-        regularised_particles = copy.copy(posterior)
-        moved_particles = copy.copy(posterior)
-        transitioned_prior = copy.copy(prior)
-
         hypotheses = posterior.hypothesis if isinstance(posterior.hypothesis, Sequence) \
             else [posterior.hypothesis]
-
-        transition_model = hypotheses[0].prediction.transition_model or self.transition_model
-        if transition_model is not None:
-            time_interval = posterior.timestamp - prior.timestamp
-            transitioned_prior.state_vector = \
-                transition_model.function(prior, noise=False, time_interval=time_interval)
-
         detections = {hypothesis.measurement for hypothesis in hypotheses if hypothesis}
 
         if detections:
@@ -91,6 +96,7 @@
             covar_est = posterior.covar
 
             # move particles
+            moved_particles = copy.copy(posterior)
             moved_particles.state_vector = moved_particles.state_vector + \
                 hopt * cholesky_eps(covar_est) @ np.random.randn(ndim, nparticles)
 
@@ -100,6 +106,7 @@
                 moved_particles.state_vector[:, part_indx] = posterior.state_vector[:, part_indx]
 
             # Evaluate likelihoods
+            transitioned_prior = self._transition_prior(prior, posterior, **kwargs)
             part_diff = moved_particles.state_vector - transitioned_prior.state_vector
             move_likelihood = multivariate_normal.logpdf(part_diff.T,
                                                          cov=covar_est)
@@ -129,6 +136,43 @@
             selector = uniform.rvs(size=nparticles)
             index = alpha > selector
 
-            regularised_particles.state_vector[:, index] = moved_particles.state_vector[:, index]
+            posterior = copy.copy(posterior)
+            posterior.state_vector = copy.copy(posterior.state_vector)
+            posterior.state_vector[:, index] = moved_particles.state_vector[:, index]
+
+        return posterior
+
+
+class MultiModelMCMCRegulariser(MCMCRegulariser):
+    """MCMC Regulariser for :class:`~.MultiModelParticleState`
 
-        return regularised_particles
+    This is a version of the :class:`~.MCMCRegulariser` that supports case where multiple
+    models are used i.e. with :class:`~.MultiModelParticleUpdater`.
+    """
+    transition_model = None
+    transition_models: Sequence[TransitionModel] = Property(
+        doc="Transition models used to for particle transition, selected by model index on "
+            "particle. Models dimensions can be subset of the overall state space, by "
+            "using :attr:`model_mappings`."
+    )
+    model_mappings: Sequence[Sequence[int]] = Property(
+        doc="Sequence of mappings associated with each transition model. This enables mapping "
+            "between model and state space, enabling use of models that may have different "
+            "dimensions (e.g. velocity or acceleration). Parts of the state that aren't mapped "
+            "are set to zero.")
+
+    def _transition_prior(self, prior, posterior, **kwargs):
+        transitioned_prior = copy.copy(prior)
+        transitioned_prior.state_vector = copy.copy(prior.state_vector)
+        for model_index, transition_model in enumerate(self.transition_models):
+            current_model_indices = prior.dynamic_model == model_index
+            current_model_count = np.count_nonzero(current_model_indices)
+            if current_model_count == 0:
+                continue
+
+            new_state_vector = MultiModelPredictor.apply_model(
+                prior[current_model_indices], transition_model, posterior.timestamp,
+                self.model_mappings[model_index], noise=False)
+
+            transitioned_prior.state_vector[:, current_model_indices] = new_state_vector
+        return transitioned_prior

stonesoup/regulariser/tests/test_particle.py

@@ -1,16 +1,22 @@
-import numpy as np
 import datetime
+
+import numpy as np
 import pytest
 
-from ...types.state import ParticleState
-from ...types.particle import Particle
-from ...types.hypothesis import SingleHypothesis
-from ...types.prediction import ParticleStatePrediction, ParticleMeasurementPrediction
+from ..particle import MCMCRegulariser, MultiModelMCMCRegulariser
 from ...models.measurement.linear import LinearGaussian
 from ...models.transition.linear import CombinedLinearGaussianTransitionModel, ConstantVelocity
+from ...predictor.particle import MultiModelPredictor
+from ...resampler.particle import SystematicResampler
 from ...types.detection import Detection
+from ...types.hypothesis import SingleHypothesis
+from ...types.particle import Particle, MultiModelParticle
+from ...types.prediction import ParticleStatePrediction, ParticleMeasurementPrediction
+from ...types.state import ParticleState, MultiModelParticleState
 from ...types.update import Update, ParticleStateUpdate
-from ..particle import MCMCRegulariser
+from ...updater.particle import MultiModelParticleUpdater
+from ...updater.tests.test_multi_model_particle import (  # noqa: F401
+    dynamic_model_list, position_mappings, transition_matrix, resampler)
 
 
 def dummy_constraint_function(particles):
@@ -159,3 +165,51 @@ def test_no_measurement():
     assert any(new_particles.weight == state_update.weight)
     # Check that the timestamp is the same
     assert new_particles.timestamp == state_update.timestamp
+
+
+def test_multi_model_regulariser(
+        dynamic_model_list, position_mappings, transition_matrix):  # noqa: F811
+
+    # Initialise particles
+    particle1 = MultiModelParticle(
+        state_vector=[1, 1, -0.5, 1, 1, -0.5],
+        weight=1/3000,
+        dynamic_model=0)
+    particle2 = MultiModelParticle(
+        state_vector=[1, 1, 0.5, 1, 1, 0.5],
+        weight=1/3000,
+        dynamic_model=1)
+    particle3 = MultiModelParticle(
+        state_vector=[1, 1, 0.5, 1, 1, 0.5],
+        weight=1/3000,
+        dynamic_model=2)
+
+    particles = [particle1, particle2, particle3] * 1000
+    timestamp = datetime.datetime.now()
+
+    particle_state = MultiModelParticleState(
+        None, particle_list=particles, timestamp=timestamp)
+
+    predictor = MultiModelPredictor(
+        dynamic_model_list, transition_matrix, position_mappings
+    )
+
+    timestamp += datetime.timedelta(seconds=5)
+    prediction = predictor.predict(particle_state, timestamp)
+
+    measurement_model = LinearGaussian(6, [0, 3], np.diag([1, 1]))
+    updater = MultiModelParticleUpdater(measurement_model, predictor, SystematicResampler())
+
+    # Detection close to where known turn rate model would place particles
+    detection = Detection([[0.5, 7.]], timestamp, measurement_model)
+
+    update = updater.update(hypothesis=SingleHypothesis(prediction, detection))
+
+    regulariser = MultiModelMCMCRegulariser(
+        dynamic_model_list,
+        position_mappings)
+
+    new_particles = regulariser.regularise(particle_state, update)
+
+    assert not np.array_equal(update.state_vector, new_particles.state_vector)
+    assert np.array_equal(update.log_weight, new_particles.log_weight)

stonesoup/updater/particle.py

@@ -299,11 +299,25 @@
             + measurement_model.logpdf(hypothesis.measurement, update, **kwargs) \
             + np.log(transition_matrix[update.parent.dynamic_model, update.dynamic_model])
 
+        # Apply constraints if defined
+        if self.constraint_func is not None:
+            part_indx = self.constraint_func(update)
+            update.log_weight[part_indx] = -1*np.inf
+
         # Normalise the weights
         update.log_weight -= logsumexp(update.log_weight)
 
-        if self.resampler:
-            update = self.resampler.resample(update)
+        # Resample
+        resample_flag = True
+        if self.resampler is not None:
+            resampled_state = self.resampler.resample(update)
+            if resampled_state == update:
+                resample_flag = False
+            update = resampled_state
+
+        if self.regulariser is not None and resample_flag:
+            update = self.regulariser.regularise(update.parent, update)
+
         return update
 
 
@@ -335,7 +349,6 @@
 
         update = Update.from_state(
             hypothesis.prediction,
-            log_weight=copy.copy(hypothesis.prediction.log_weight),
             hypothesis=hypothesis,
             timestamp=hypothesis.measurement.timestamp,
         )
@@ -347,11 +360,25 @@
                                                                          update,
                                                                          **kwargs)
 
+        # Apply constraints if defined
+        if self.constraint_func is not None:
+            part_indx = self.constraint_func(update)
+            update.log_weight[part_indx] = -1*np.inf
+
         # Normalise the weights
         update.log_weight -= logsumexp(update.log_weight)
 
-        if self.resampler:
-            update = self.resampler.resample(update)
+        # Resample
+        resample_flag = True
+        if self.resampler is not None:
+            resampled_state = self.resampler.resample(update)
+            if resampled_state == update:
+                resample_flag = False
+            update = resampled_state
+
+        if self.regulariser is not None and resample_flag:
+            update = self.regulariser.regularise(update.parent, update)
+
         return update
 
     @staticmethod
@@ -498,6 +525,14 @@
             # Normalise weights
             updated_state.log_weight -= logsumexp(updated_state.log_weight)
 
+        # Apply constraints if defined
+        if self.constraint_func is not None:
+            part_indx = self.constraint_func(updated_state)
+            updated_state.log_weight[part_indx] = -1*np.inf
+            if not any(hypotheses):
+                updated_state.log_weight = copy.copy(updated_state.log_weight)
+            updated_state.log_weight -= logsumexp(updated_state.log_weight)
+
         # Resampling
         if self.resampler is not None:
             updated_state = self.resampler.resample(updated_state,