fix: modify instrument merge (#1619)

* fix: raise warnings when users include duplicate components, except for HarpDevice with is_clock_generator True

* refactor: fix logic and move combine code into its own function

* chore: lint

Author: dbirman

src/aind_data_schema/core/instrument.py

@@ -1,6 +1,7 @@
 """Core Instrument model"""
 
 from datetime import date
+import logging
 from typing import List, Literal, Optional
 
 from aind_data_schema_models.modalities import Modality
@@ -267,6 +268,46 @@ def validate_modality_device_dependencies(self):
 
         return self
 
+    def _is_harp_clock_generator(self, component) -> bool:
+        """Check if a component is a HarpDevice and a clock generator"""
+        return (
+            isinstance(component, HarpDevice)
+            and hasattr(component, "is_clock_generator")
+            and component.is_clock_generator is True
+        )
+
+    def _combine_components(self, components: List[Device], other_components: List[Device]) -> List[Device]:
+        """Combine components from two instruments, handling duplicates appropriately."""
+        seen_names = set()
+        combined_components = []
+
+        for component in components:
+            combined_components.append(component)
+            seen_names.add(component.name)
+
+        for other_component in other_components:
+            if other_component.name in seen_names:
+                matching_component = next((c for c in components if c.name == other_component.name), None)
+                if (
+                    matching_component
+                    and self._is_harp_clock_generator(matching_component)
+                    and self._is_harp_clock_generator(other_component)
+                ):
+                    logging.info(
+                        f"{other_component.name} is a HarpDevice clock generator, "
+                        f"only one instance will be kept in the combined instrument."
+                    )
+                else:
+                    logging.error(
+                        f"Instruments should not have duplicated components,"
+                        f" this will raise an error in future versions: {other_component.name}"
+                    )
+            else:
+                combined_components.append(other_component)
+                seen_names.add(other_component.name)
+
+        return combined_components
+
     def __add__(self, other: "Instrument") -> "Instrument":
         """Combine two Instrument objects"""
 
@@ -306,7 +347,7 @@ def __add__(self, other: "Instrument") -> "Instrument":
         combined_connections = self.connections + other.connections
 
         # Combine components
-        combined_components = self.components + other.components
+        combined_components = self._combine_components(self.components, other.components)
 
         # Combine notes
         combined_notes = merge_notes(self.notes, other.notes)

tests/test_instrument.py

@@ -3,11 +3,13 @@
 import json
 import unittest
 from datetime import date
+from unittest.mock import patch
 
 from aind_data_schema_models.coordinates import AnatomicalRelative
 from aind_data_schema_models.modalities import Modality
 from aind_data_schema_models.organizations import Organization
 from aind_data_schema_models.units import FrequencyUnit, PowerUnit, SizeUnit
+from aind_data_schema_models.harp_types import HarpDeviceType
 from pydantic import ValidationError
 
 from aind_data_schema.components.coordinates import CoordinateSystemLibrary
@@ -25,6 +27,7 @@
     EphysAssembly,
     EphysProbe,
     FiberPatchCord,
+    HarpDevice,
     Laser,
     LaserAssembly,
     Lens,
@@ -655,8 +658,8 @@ def test_instrument_addition(self):
         # Check that modalities are combined and sorted (should be the same since we're adding identical instruments)
         self.assertEqual(len(combined.modalities), len(set(inst1.modalities + inst2.modalities)))
 
-        # Check that components are combined
-        self.assertEqual(len(combined.components), len(inst1.components) + len(inst2.components))
+        # Check that components are deduplicated (same names from both instruments result in keeping only one)
+        self.assertEqual(len(combined.components), len(inst1.components))
 
         # Check that connections are combined
         self.assertEqual(len(combined.connections), len(inst1.connections) + len(inst2.connections))
@@ -710,6 +713,128 @@ def test_instrument_addition(self):
         combined = inst1 + inst2
         self.assertEqual(combined.notes, "Only note")
 
+    def test_duplicate_non_harp_device_components(self):
+        """Test that duplicate non-HarpDevice components log an error when combining instruments"""
+
+        inst1 = Instrument(
+            instrument_id="test_inst",
+            modification_date=date(2020, 10, 10),
+            modalities=[Modality.ECEPHYS],
+            coordinate_system=CoordinateSystemLibrary.BREGMA_ARI,
+            components=[Computer(name="Computer1")],
+        )
+        inst2 = Instrument(
+            instrument_id="test_inst",
+            modification_date=date(2020, 10, 10),
+            modalities=[Modality.ECEPHYS],
+            coordinate_system=CoordinateSystemLibrary.BREGMA_ARI,
+            components=[Computer(name="Computer1")],
+        )
+
+        with patch("aind_data_schema.core.instrument.logging") as mock_logging:
+            combined = inst1 + inst2
+            mock_logging.error.assert_called_once()
+            error_call_args = mock_logging.error.call_args[0][0]
+            self.assertIn("Computer1", error_call_args)
+            self.assertIn("duplicated", error_call_args)
+
+        self.assertEqual(len(combined.components), 1)
+
+    def test_duplicate_harp_clock_generator_devices(self):
+        """Test that duplicate HarpDevice clock generators are allowed when combining instruments"""
+
+        harp_clock_gen = HarpDevice(
+            name="Harp Clock Generator",
+            harp_device_type=HarpDeviceType.CLOCKSYNCHRONIZER,
+            core_version="2.1",
+            channels=[],
+            is_clock_generator=True,
+        )
+
+        inst1 = Instrument(
+            instrument_id="test_inst",
+            modification_date=date(2020, 10, 10),
+            modalities=[Modality.ECEPHYS],
+            coordinate_system=CoordinateSystemLibrary.BREGMA_ARI,
+            components=[harp_clock_gen],
+        )
+        inst2 = Instrument(
+            instrument_id="test_inst",
+            modification_date=date(2020, 10, 10),
+            modalities=[Modality.ECEPHYS],
+            coordinate_system=CoordinateSystemLibrary.BREGMA_ARI,
+            components=[harp_clock_gen.model_copy(deep=True)],
+        )
+
+        with patch("aind_data_schema.core.instrument.logging") as mock_logging:
+            combined = inst1 + inst2
+            mock_logging.info.assert_called_once()
+            info_call_args = mock_logging.info.call_args[0][0]
+            self.assertIn("Harp Clock Generator", info_call_args)
+
+        self.assertEqual(len(combined.components), 1)
+
+    def test_duplicate_non_harp_device_with_clock_generator_attribute(self):
+        """Test that duplicate non-HarpDevice components with is_clock_generator log error"""
+
+        harp_clock_gen = HarpDevice(
+            name="CustomClockGenerator",
+            harp_device_type=HarpDeviceType.BEHAVIOR,
+            is_clock_generator=True,
+            channels=[],
+        )
+
+        harp_non_clock_gen = HarpDevice(
+            name="CustomClockGenerator",
+            harp_device_type=HarpDeviceType.BEHAVIOR,
+            is_clock_generator=False,
+            channels=[],
+        )
+
+        inst1 = Instrument(
+            instrument_id="test_inst",
+            modification_date=date(2020, 10, 10),
+            modalities=[Modality.BEHAVIOR],
+            coordinate_system=CoordinateSystemLibrary.BREGMA_ARI,
+            components=[harp_clock_gen, LickSpoutAssembly(
+                name="Lick spout assembly A",
+                lick_spouts=[
+                    LickSpout(
+                        name="Left spout",
+                        spout_diameter=1.2,
+                        solenoid_valve=Device(name="Solenoid Left"),
+                        lick_sensor=Device(name="Lick-o-meter Left"),
+                    ),
+                ],
+            )],
+        )
+        inst2 = Instrument(
+            instrument_id="test_inst",
+            modification_date=date(2020, 10, 10),
+            modalities=[Modality.BEHAVIOR],
+            coordinate_system=CoordinateSystemLibrary.BREGMA_ARI,
+            components=[harp_non_clock_gen, LickSpoutAssembly(
+                name="Lick spout assembly B",
+                lick_spouts=[
+                    LickSpout(
+                        name="Left spout",
+                        spout_diameter=1.2,
+                        solenoid_valve=Device(name="Solenoid Left"),
+                        lick_sensor=Device(name="Lick-o-meter Left"),
+                    ),
+                ],
+            )],
+        )
+
+        with patch("aind_data_schema.core.instrument.logging") as mock_logging:
+            combined = inst1 + inst2
+            mock_logging.error.assert_called_once()
+            error_call_args = mock_logging.error.call_args[0][0]
+            self.assertIn("CustomClockGenerator", error_call_args)
+            self.assertIn("duplicated", error_call_args)
+
+        self.assertEqual(len(combined.components), 3)
+
 
 class ConnectionTest(unittest.TestCase):
     """Test the Connection schema"""