Add ability to replay tracks

stonesoup/feeder/__init__.py

@@ -4,6 +4,6 @@
 framework, and feed them into the tracking algorithms. These can then
 optionally be used to drop detections, deliver detections out of sequence,
 synchronise out of sequence detections, etc. """
-from .base import Feeder, DetectionFeeder, GroundTruthFeeder
+from .base import Feeder, DetectionFeeder, GroundTruthFeeder, TrackFeeder
 
-__all__ = ['Feeder', 'DetectionFeeder', 'GroundTruthFeeder']
+__all__ = ['Feeder', 'DetectionFeeder', 'GroundTruthFeeder', 'TrackFeeder']

stonesoup/feeder/base.py

@@ -2,7 +2,7 @@
 from abc import abstractmethod
 
 from ..base import Property
-from ..reader import Reader, DetectionReader, GroundTruthReader
+from ..reader import Reader, DetectionReader, GroundTruthReader, TrackReader
 from ..buffered_generator import BufferedGenerator
 
 
@@ -43,3 +43,15 @@ class GroundTruthFeeder(Feeder, GroundTruthReader):
     @BufferedGenerator.generator_method
     def groundtruth_paths_gen(self):
         yield from self.data_gen()
+
+
+class TrackFeeder(Feeder, TrackReader):
+    """Track feeder base class
+
+    Feeder consumes and outputs :class:`.Track` data and can be used to modify the sequence,
+    duplicate or drop data.
+    """
+
+    @BufferedGenerator.generator_method
+    def tracks_gen(self):
+        yield from self.data_gen()

stonesoup/feeder/tests/test_track.py

@@ -3,7 +3,7 @@
 
 from ...types.state import GaussianState
 from ...types.track import Track
-from ..track import Tracks2GaussianDetectionFeeder
+from ..track import Tracks2GaussianDetectionFeeder, ReplayTrackFeeder
 
 t1 = Track(GaussianState([1, 1, 1, 1], np.diag([2, 2, 2, 2]), timestamp=2))
 t2 = Track([GaussianState([1, 1, 1, 1], np.diag([2, 2, 2, 2]), timestamp=1),
@@ -13,6 +13,8 @@
             GaussianState([3, 1], np.diag([2, 2]), timestamp=2)])
 t4 = Track(GaussianState([1, 0, 1, 0, 1, 0], np.diag([2, 2, 2, 2, 2, 2]), timestamp=2))
 
+times = [0, 1, 2]
+
 
 @pytest.mark.parametrize(
     "tracks",
@@ -53,3 +55,22 @@ def test_Track2GaussianDetectionFeeder(tracks):
                               if detection.metadata['track_id'] == track.id))
         assert np.all(detection.state_vector == track.state_vector)
         assert np.all(detection.covar == track.covar)
+
+
+@pytest.mark.parametrize(("tracks", "times"),
+                         [([t1], times), ([t1], None),
+                          ([t1, t2], times), ([t1, t2], None),
+                          ([t2, t3], times), ([t2, t3], None),
+                          ([t1, t2, t3, t4], times), ([t1, t2, t3, t4], None)])
+def test_ReplayTrackFeeder(tracks, times):
+    feeder = ReplayTrackFeeder(reader=tracks, times=times)
+    feeder_times = []
+    feeder_tracks = set()
+    for new_time, new_tracks in feeder:
+        feeder_times.append(new_time)
+        feeder_tracks |= new_tracks
+    if times is not None:
+        assert times == feeder_times
+    assert len(tracks) == len(feeder_tracks)
+    assert (sorted([len(track) for track in tracks]) ==
+            sorted([len(track) for track in feeder_tracks]))

stonesoup/feeder/track.py

@@ -1,6 +1,9 @@
+from collections.abc import Collection
+import datetime
 import numpy as np
 
-from . import DetectionFeeder
+from . import DetectionFeeder, TrackFeeder
+from ..base import Property
 from ..buffered_generator import BufferedGenerator
 from ..models.measurement.linear import LinearGaussian
 from ..types.detection import GaussianDetection, Detection
@@ -42,3 +45,51 @@ def data_gen(self):
                     raise TypeError(f"track is of type {type(track)}. Expected Track or Detection")
 
             yield time, detections
+
+
+class ReplayTrackFeeder(TrackFeeder):
+    """
+    Feeder outputs Track objects from an input of tracks.
+    This allows an already produced set of tracks to be used as a reader.
+
+    At each timestep, the states of each track that existed at that point are output.
+    Any tracks which have ended are removed
+    """
+
+    reader: Collection[Track] = Property(doc="A collection of tracks to be replayed")
+    times: list[datetime.datetime] = Property(
+        default=None, doc="The timesteps at which the tracks should be replayed. "
+        "The default `None` will use all timestamps")
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        if self.times is None:
+            times = {state.timestamp for track in self.reader for state in track}
+            self.times = sorted(times)
+
+    @BufferedGenerator.generator_method
+    def data_gen(self):
+        output_tracks = {}
+        last_time = None
+
+        for time in self.times:
+            for track in self.reader:
+                if not track or track[0].timestamp > time:
+                    continue
+                if last_time is not None and last_time >= track.timestamp:
+                    if track in output_tracks:
+                        del output_tracks[track]
+                    continue
+                track_states = []
+                track_metadatas = []
+                for state, metadata in zip(track.states, track.metadatas):
+                    if state.timestamp > time:
+                        continue
+                    track_states.append(state)
+                    track_metadatas.append(metadata)
+                current_track = output_tracks.setdefault(
+                    track, Track(id=track.id, init_metadata=track.init_metadata))
+                current_track.states = track_states
+                current_track.metadatas = track_metadatas
+            last_time = time
+            yield time, set(output_tracks.values())

stonesoup/reader/__init__.py

@@ -1,5 +1,5 @@
 """Reader classes are used for getting data into the framework."""
-from .base import Reader, DetectionReader, GroundTruthReader, SensorDataReader
+from .base import Reader, DetectionReader, GroundTruthReader, SensorDataReader, TrackReader
 
 __all__ = [
-    'Reader', 'DetectionReader', 'GroundTruthReader', 'SensorDataReader']
+    'Reader', 'DetectionReader', 'GroundTruthReader', 'SensorDataReader', 'TrackReader']