From ddeee1bc0d85b092e730b2b630cfe6f3cd69141e Mon Sep 17 00:00:00 2001
From: Emily Hunter <59794482+ekhunter123@users.noreply.github.com>
Date: Thu, 28 Jul 2022 09:30:40 -0400
Subject: [PATCH] Apply suggestions from code review

Generalize the update and predicted measurement, rather than using the specific Gaussian state subclasses.

Co-authored-by: Steven Hiscocks <sdhiscocks@dstl.gov.uk>
---
 stonesoup/feeder/track.py     |  3 ++-
 stonesoup/updater/chernoff.py | 15 +++++++--------
 2 files changed, 9 insertions(+), 9 deletions(-)

diff --git a/stonesoup/feeder/track.py b/stonesoup/feeder/track.py
index 95e385638..b088c3c8f 100644
--- a/stonesoup/feeder/track.py
+++ b/stonesoup/feeder/track.py
@@ -26,7 +26,8 @@ def data_gen(self):
                 detections.append(
                     GaussianDetection.from_state(
                         track.state,
-                        measurement_model=LinearGaussian(dim, range(dim), np.asarray(track.covar)),
+                        measurement_model=LinearGaussian(
+                            dim, list(range(dim)), np.asarray(track.covar)),
                         target_type=GaussianDetection)
                 )
 
diff --git a/stonesoup/updater/chernoff.py b/stonesoup/updater/chernoff.py
index d04b98b3f..577f40359 100644
--- a/stonesoup/updater/chernoff.py
+++ b/stonesoup/updater/chernoff.py
@@ -3,8 +3,8 @@
 
 from ..base import Property
 from .base import Updater
-from ..types.prediction import GaussianMeasurementPrediction
-from ..types.update import GaussianStateUpdate
+from ..types.prediction import MeasurementPrediction
+from ..types.update import Update
 
 
 class ChernoffUpdater(Updater):
@@ -86,9 +86,9 @@ def predict_measurement(self, predicted_state, measurement_model=None,  **kwargs
         meas_cross_cov = predicted_state.covar
 
         # Combine everything into a GaussianMeasurementPrediction object
-        return GaussianMeasurementPrediction(predicted_meas, innov_covar,
-                                             predicted_state.timestamp,
-                                             cross_covar=meas_cross_cov)
+        return MeasurementPrediction.from_state(predicted_state, predicted_meas, innov_covar,
+                                                predicted_state.timestamp,
+                                                cross_covar=meas_cross_cov)
 
     def update(self, hypothesis, force_symmetric_covariance=False, **kwargs):
         '''
@@ -127,6 +127,5 @@ def update(self, hypothesis, force_symmetric_covariance=False, **kwargs):
                 (posterior_covariance + posterior_covariance.T)/2
 
         # Return the updated state
-        return GaussianStateUpdate(posterior_mean, posterior_covariance,
-                                   hypothesis,
-                                   hypothesis.measurement.timestamp)
+        return Update.from_state(hypothesis.prediction, posterior_mean, posterior_covariance,
+                                 hypothesis, hypothesis.measurement.timestamp)