feat: DataStream.__add__ and Acquistion._merge_data_stream_lists (#1634)

* feat: DataStream.__add__ and Acquistion._merge_data_stream_lists

* test: merging now merges data streams!

* tests: coverage for un-mergeable streams

* chore: overzealous codespell

* refactor: allow configurable overlap range, defaults to 2 minutes

* chore: lint

* refactor: merge_streams instead of two different lists

* tests: match new _merge_data_streams input type

Author: dbirman

pyproject.toml

@@ -103,4 +103,4 @@ fail-under = 100
 
 [tool.codespell]
 skip = '.git,*.pdf,*.svg'
-ignore-words-list = 'nd,assertIn'
+ignore-words-list = 'nd,assertIn,DeviceC'

src/aind_data_schema/core/acquisition.py

@@ -1,6 +1,7 @@
 """Schema describing data acquisition metadata and configurations"""
 
 from decimal import Decimal
+import logging
 from typing import Annotated, List, Literal, Optional
 
 from aind_data_schema_models.modalities import Modality
@@ -188,6 +189,54 @@ def check_connections(self):
 
         return self
 
+    @classmethod
+    def overlapping(cls, stream1: "DataStream", stream2: "DataStream", overlap_s: int) -> bool:
+        """Check if two DataStream objects have overlapping start and end times"""
+        start_diff = abs((stream1.stream_start_time - stream2.stream_start_time).total_seconds())
+        end_diff = abs((stream1.stream_end_time - stream2.stream_end_time).total_seconds())
+        return start_diff <= overlap_s and end_diff <= overlap_s
+
+    def __add__(self, other: "DataStream", overlap_s: int = 120) -> "DataStream":
+        """Combine two DataStream objects"""
+
+        if not DataStream.overlapping(self, other, overlap_s=overlap_s):
+            raise ValueError("Cannot combine DataStreams with non-overlapping start and end times.")
+
+        min_start_time = min(self.stream_start_time, other.stream_start_time)
+        max_end_time = max(self.stream_end_time, other.stream_end_time)
+
+        # Combine modalities
+        modalities = self.modalities + other.modalities
+        modalities = remove_duplicates(modalities)
+
+        # Combine active devices
+        active_devices = self.active_devices + other.active_devices
+        len_orig_devices = len(active_devices)
+        active_devices = remove_duplicates(active_devices)
+        if len(active_devices) < len_orig_devices:
+            logging.warning(
+                "Duplicate active devices were removed. Only DAQ devices should be shared in overlapped " "DataStreams."
+            )
+
+        # Combine configurations
+        configurations = self.configurations + other.configurations
+
+        # Combine connections
+        connections = self.connections + other.connections
+
+        # Combine notes
+        notes = merge_notes(self.notes, other.notes)
+
+        return DataStream(
+            stream_start_time=min_start_time,
+            stream_end_time=max_end_time,
+            modalities=modalities,
+            active_devices=active_devices,
+            configurations=configurations,
+            connections=connections,
+            notes=notes,
+        )
+
 
 class StimulusEpoch(DataModel):
     """All stimuli being presented to the subject. starting and stopping at approximately the
@@ -361,6 +410,37 @@ def specimen_required(self):
 
         return self
 
+    @classmethod
+    def _merge_data_streams(
+        cls, streams: List[DataStream], overlap_s: int = 120
+    ) -> List[DataStream]:
+        """Merge two lists of data streams"""
+        groups = []
+        visited = set()
+        for i in range(len(streams)):
+            if i in visited:
+                continue
+            group = [streams[i]]
+            visited.add(i)
+            for j in range(i + 1, len(streams)):
+                if j not in visited and DataStream.overlapping(streams[i], streams[j], overlap_s=overlap_s):
+                    group.append(streams[j])
+                    visited.add(j)
+            groups.append(group)
+
+        # Construct the final set of streams, including merged streams where applicable
+        merged_streams = []
+        for group in groups:
+            if len(group) == 1:
+                merged_streams.append(group[0])
+            else:
+                merged_stream = group[0]
+                for stream in group[1:]:
+                    merged_stream = merged_stream + stream
+                merged_streams.append(merged_stream)
+
+        return merged_streams
+
     def __add__(self, other: "Acquisition") -> "Acquisition":
         """Combine two Acquisition objects"""
 
@@ -400,7 +480,7 @@ def __add__(self, other: "Acquisition") -> "Acquisition":
         ethics_review_id = merge_optional_list(self.ethics_review_id, other.ethics_review_id)
         calibrations = self.calibrations + other.calibrations
         maintenance = self.maintenance + other.maintenance
-        data_streams = self.data_streams + other.data_streams
+        data_streams = Acquisition._merge_data_streams(self.data_streams + other.data_streams)
         stimulus_epochs = self.stimulus_epochs + other.stimulus_epochs
 
         # Remove duplicates

tests/test_acquisition.py

@@ -392,6 +392,166 @@ def extract_union_types(annotation):
         # If we get here, all device config subclasses are properly covered
         self.assertTrue(True, "All DeviceConfig subclasses are properly covered in discriminated unions")
 
+    def test_datastream_add_basic(self):
+        """Test combining two DataStream objects"""
+        acq = ephys_acquisition.model_copy()
+
+        stream1 = acq.data_streams[0].model_copy()
+        stream2 = acq.data_streams[0].model_copy()
+
+        stream1.active_devices = ["Device1", "Device2"]
+        stream2.active_devices = ["Device3", "Device4"]
+
+        combined_stream = stream1 + stream2
+
+        self.assertIsNotNone(combined_stream)
+        self.assertEqual(combined_stream.stream_start_time, stream1.stream_start_time)
+        self.assertEqual(combined_stream.stream_end_time, stream1.stream_end_time)
+        self.assertEqual(len(combined_stream.modalities), 1)
+        self.assertEqual(len(combined_stream.active_devices), 4)
+        self.assertEqual(len(combined_stream.configurations), 4)
+
+        # Also check that an error is raised if the streams cannot be combined
+
+        stream2.stream_end_time = stream2.stream_end_time.replace(year=2100)
+
+        with self.assertRaises(ValueError):
+            _ = stream1 + stream2
+
+    def test_datastream_add_combines_notes(self):
+        """Test that notes are properly merged when combining DataStreams"""
+        acq = ephys_acquisition.model_copy()
+
+        stream1 = acq.data_streams[0].model_copy()
+        stream2 = acq.data_streams[0].model_copy()
+
+        stream1.active_devices = ["Device1", "Device2"]
+        stream2.active_devices = ["Device3", "Device4"]
+        stream1.notes = "Note 1"
+        stream2.notes = "Note 2"
+
+        combined_stream = stream1 + stream2
+
+        self.assertIn("Note 1", combined_stream.notes)
+        self.assertIn("Note 2", combined_stream.notes)
+
+    def test_datastream_add_with_duplicate_devices(self):
+        """Test that overlapping active devices are logged as warning when combining"""
+        acq = ephys_acquisition.model_copy()
+        stream1 = acq.data_streams[0]
+        stream2 = acq.data_streams[0].model_copy()
+
+        combined_stream = stream1 + stream2
+
+        self.assertIsNotNone(combined_stream)
+
+    def test_datastream_add_combines_connections(self):
+        """Test that connections are properly combined"""
+        acq = ephys_acquisition.model_copy()
+
+        stream1 = acq.data_streams[0].model_copy()
+        stream2 = acq.data_streams[0].model_copy()
+
+        stream1.active_devices = ["Device1", "Device2"]
+        stream2.active_devices = ["Device3", "Device4"]
+        stream1.connections = [Connection(source_device="Device1", target_device="Device2")]
+        stream2.connections = [Connection(source_device="Device3", target_device="Device4")]
+
+        combined_stream = stream1 + stream2
+
+        self.assertEqual(len(combined_stream.connections), 2)
+
+    def test_merge_data_stream_lists_single_streams(self):
+        """Test merging lists with single streams"""
+        acq = ephys_acquisition.model_copy()
+        stream1 = acq.data_streams[0].model_copy()
+        stream2 = acq.data_streams[1].model_copy()
+
+        stream1.active_devices = ["Device1"]
+        stream2.active_devices = ["Device2"]
+
+        merged = Acquisition._merge_data_streams([stream1] + [stream2])
+
+        self.assertEqual(len(merged), 2)
+
+    def test_merge_data_stream_lists_overlapping_streams(self):
+        """Test merging streams with overlapping start/end times"""
+        acq = ephys_acquisition.model_copy()
+
+        stream1 = acq.data_streams[0].model_copy()
+        stream2 = acq.data_streams[0].model_copy()
+
+        stream1.active_devices = ["Device1", "Device2"]
+        stream2.active_devices = ["Device3", "Device4"]
+
+        merged = Acquisition._merge_data_streams([stream1] + [stream2])
+
+        self.assertEqual(len(merged), 1)
+        self.assertEqual(len(merged[0].active_devices), 4)
+
+    def test_merge_data_stream_lists_non_overlapping_streams(self):
+        """Test merging streams with different start/end times"""
+        acq = ephys_acquisition.model_copy()
+
+        stream1 = acq.data_streams[1].model_copy()
+        stream2 = acq.data_streams[0].model_copy()
+
+        stream1.active_devices = ["Device1"]
+        stream2.active_devices = ["Device2"]
+
+        merged = Acquisition._merge_data_streams([stream1] + [stream2])
+
+        self.assertEqual(len(merged), 2)
+
+    def test_merge_data_stream_lists_multiple_overlapping_groups(self):
+        """Test merging multiple streams with multiple overlapping groups"""
+        acq = ephys_acquisition.model_copy()
+
+        stream1 = acq.data_streams[0].model_copy()
+        stream2 = acq.data_streams[0].model_copy()
+        stream3 = acq.data_streams[0].model_copy()
+
+        stream1.active_devices = ["DeviceA", "DeviceB"]
+        stream2.active_devices = ["DeviceC"]
+        stream3.active_devices = ["DeviceD", "DeviceE"]
+
+        start1 = datetime(year=2023, month=4, day=25, hour=2, minute=0, second=0, tzinfo=timezone.utc)
+        end1 = datetime(year=2023, month=4, day=25, hour=2, minute=30, second=0, tzinfo=timezone.utc)
+        stream1.stream_start_time = start1
+        stream1.stream_end_time = end1
+
+        start2 = datetime(year=2023, month=4, day=25, hour=2, minute=0, second=30, tzinfo=timezone.utc)
+        end2 = datetime(year=2023, month=4, day=25, hour=2, minute=29, second=30, tzinfo=timezone.utc)
+        stream2.stream_start_time = start2
+        stream2.stream_end_time = end2
+
+        start3 = datetime(year=2023, month=4, day=25, hour=3, minute=0, second=0, tzinfo=timezone.utc)
+        end3 = datetime(year=2023, month=4, day=25, hour=3, minute=30, second=0, tzinfo=timezone.utc)
+        stream3.stream_start_time = start3
+        stream3.stream_end_time = end3
+
+        merged = Acquisition._merge_data_streams([stream1] + [stream2, stream3])
+
+        for m in merged:
+            print(m.stream_start_time, m.stream_end_time, m.active_devices)
+
+        self.assertEqual(len(merged), 2)
+        self.assertEqual(len(merged[0].active_devices), 3)
+
+    def test_datastream_add_with_exaspim_example(self):
+        """Test combining DataStreams using ExaSPIM example"""
+        acq = exaspim_acquisition.model_copy()
+        stream1 = acq.data_streams[0].model_copy()
+
+        stream2 = acq.data_streams[0].model_copy()
+        stream2.active_devices = ["Device99"]
+
+        combined_stream = stream1 + stream2
+
+        self.assertIsNotNone(combined_stream)
+        self.assertIn(Modality.SPIM, combined_stream.modalities)
+        self.assertEqual(len(combined_stream.active_devices), 3)
+
 
 if __name__ == "__main__":
     unittest.main()

tests/test_composability_merge.py

@@ -143,7 +143,7 @@ def test_merge_acquisition(self):
         self.assertEqual(len(merged_acq.experimenters), 1)
         self.assertEqual(len(merged_acq.maintenance), 2)
         self.assertEqual(len(merged_acq.calibrations), 2)
-        self.assertEqual(len(merged_acq.data_streams), 2)
+        self.assertEqual(len(merged_acq.data_streams), 1)
         self.assertEqual(merged_acq.acquisition_start_time, t)
         self.assertEqual(merged_acq.acquisition_end_time, t)
         self.assertEqual(merged_acq.acquisition_type, "ExaSPIM")