Add beamline-specific zebra mappings to zebra device (#976)

* Add beamline-specific zebra mappings to zebra device

* Add more structure to the zebra mappings

* Use -1 as sentinel value

src/dodal/beamlines/i03.py

@@ -40,8 +40,13 @@
 from dodal.devices.webcam import Webcam
 from dodal.devices.xbpm_feedback import XBPMFeedback
 from dodal.devices.xspress3.xspress3 import Xspress3
-from dodal.devices.zebra import Zebra
-from dodal.devices.zebra_controlled_shutter import ZebraShutter
+from dodal.devices.zebra.zebra import Zebra
+from dodal.devices.zebra.zebra_constants_mapping import (
+    ZebraMapping,
+    ZebraSources,
+    ZebraTTLOutputs,
+)
+from dodal.devices.zebra.zebra_controlled_shutter import ZebraShutter
 from dodal.devices.zocalo import ZocaloResults
 from dodal.log import set_beamline as set_log_beamline
 from dodal.utils import BeamlinePrefix, get_beamline_name, skip_device
@@ -58,6 +63,12 @@
 
 set_path_provider(PandASubpathProvider())
 
+I03_ZEBRA_MAPPING = ZebraMapping(
+    outputs=ZebraTTLOutputs(TTL_DETECTOR=1, TTL_SHUTTER=2, TTL_XSPRESS3=3, TTL_PANDA=4),
+    sources=ZebraSources(),
+    AND_GATE_FOR_AUTO_SHUTTER=2,
+)
+
 
 def aperture_scatterguard(
     wait_for_connection: bool = True,
@@ -368,6 +379,7 @@ def zebra(wait_for_connection: bool = True, fake_with_ophyd_sim: bool = False) -
         "-EA-ZEBRA-01:",
         wait_for_connection,
         fake_with_ophyd_sim,
+        mapping=I03_ZEBRA_MAPPING,
     )
 
 

src/dodal/beamlines/i04.py

@@ -28,8 +28,13 @@
 from dodal.devices.thawer import Thawer
 from dodal.devices.undulator import Undulator
 from dodal.devices.xbpm_feedback import XBPMFeedback
-from dodal.devices.zebra import Zebra
-from dodal.devices.zebra_controlled_shutter import ZebraShutter
+from dodal.devices.zebra.zebra import Zebra
+from dodal.devices.zebra.zebra_constants_mapping import (
+    ZebraMapping,
+    ZebraSources,
+    ZebraTTLOutputs,
+)
+from dodal.devices.zebra.zebra_controlled_shutter import ZebraShutter
 from dodal.log import set_beamline as set_log_beamline
 from dodal.utils import BeamlinePrefix, get_beamline_name, skip_device
 
@@ -47,6 +52,11 @@
 set_log_beamline(BL)
 set_utils_beamline(BL)
 
+I04_ZEBRA_MAPPING = ZebraMapping(
+    outputs=(ZebraTTLOutputs(TTL_DETECTOR=1, TTL_FAST_SHUTTER=2, TTL_XSPRESS3=3)),
+    sources=ZebraSources(),
+)
+
 
 def smargon(
     wait_for_connection: bool = True, fake_with_ophyd_sim: bool = False
@@ -341,6 +351,7 @@ def zebra(wait_for_connection: bool = True, fake_with_ophyd_sim: bool = False) -
         "-EA-ZEBRA-01:",
         wait_for_connection,
         fake_with_ophyd_sim,
+        mapping=I04_ZEBRA_MAPPING,
     )
 
 

src/dodal/beamlines/i24.py

@@ -16,7 +16,12 @@
 from dodal.devices.i24.vgonio import VerticalGoniometer
 from dodal.devices.oav.oav_detector import OAV
 from dodal.devices.oav.oav_parameters import OAVConfig
-from dodal.devices.zebra import Zebra
+from dodal.devices.zebra.zebra import Zebra
+from dodal.devices.zebra.zebra_constants_mapping import (
+    ZebraMapping,
+    ZebraSources,
+    ZebraTTLOutputs,
+)
 from dodal.log import set_beamline as set_log_beamline
 from dodal.utils import get_beamline_name, skip_device
 
@@ -29,6 +34,11 @@
 set_log_beamline(BL)
 set_utils_beamline(BL)
 
+I24_ZEBRA_MAPPING = ZebraMapping(
+    outputs=ZebraTTLOutputs(TTL_EIGER=1, TTL_PILATUS=2, TTL_FAST_SHUTTER=4),
+    sources=ZebraSources(),
+)
+
 
 def attenuator(
     wait_for_connection: bool = True, fake_with_ophyd_sim: bool = False
@@ -191,6 +201,7 @@ def zebra(wait_for_connection: bool = True, fake_with_ophyd_sim: bool = False) -
         "-EA-ZEBRA-01:",
         wait_for_connection,
         fake_with_ophyd_sim,
+        mapping=I24_ZEBRA_MAPPING,
     )
 
 

src/dodal/devices/zebra.py → src/dodal/devices/zebra/zebra.py

@@ -14,37 +14,7 @@
 )
 from ophyd_async.epics.core import epics_signal_r, epics_signal_rw
 
-# These constants refer to I03's Zebra. See https://github.com/DiamondLightSource/dodal/issues/772
-# Sources
-DISCONNECT = 0
-IN1_TTL = 1
-IN2_TTL = 4
-IN3_TTL = 7
-IN4_TTL = 10
-PC_ARM = 29
-PC_GATE = 30
-PC_PULSE = 31
-AND3 = 34
-AND4 = 35
-OR1 = 36
-PULSE1 = 52
-PULSE2 = 53
-SOFT_IN1 = 60
-SOFT_IN2 = 61
-SOFT_IN3 = 62
-
-# Instrument specific
-TTL_DETECTOR = 1
-TTL_SHUTTER = 2
-TTL_XSPRESS3 = 3
-TTL_PANDA = 4
-
-# The AND gate that controls the automatic shutter
-AUTO_SHUTTER_GATE = 2
-
-# The first two inputs of the auto shutter gate.
-AUTO_SHUTTER_INPUT_1 = 1
-AUTO_SHUTTER_INPUT_2 = 2
+from dodal.devices.zebra.zebra_constants_mapping import ZebraMapping
 
 
 class ArmSource(StrictEnum):
@@ -303,7 +273,8 @@ def __init__(self, prefix: str, name: str = "") -> None:
 class Zebra(StandardReadable):
     """The Zebra device."""
 
-    def __init__(self, name: str, prefix: str) -> None:
+    def __init__(self, mapping: ZebraMapping, name: str, prefix: str) -> None:
+        self.mapping = mapping
         self.pc = PositionCompare(prefix, name)
         self.output = ZebraOutputPanel(prefix, name)
         self.inputs = SoftInputs(prefix, name)

src/dodal/devices/zebra/zebra_constants_mapping.py

@@ -0,0 +1,96 @@
+from collections import Counter
+
+from pydantic import BaseModel, Field, model_validator
+
+
+class ZebraMappingValidations(BaseModel):
+    """Raises an exception if field set to -1 is accessed, and validate against
+    multiple fields mapping to the same integer"""
+
+    def __getattribute__(self, name: str):
+        """To protect against mismatch between the Zebra configuration that a plan expects and the Zebra which has
+        been instantiated, raise exception if a field which has been set to -1 is accessed."""
+        value = object.__getattribute__(self, name)
+        if not name.startswith("__"):
+            if value == -1:
+                raise UnmappedZebraException(
+                    f"'{type(self).__name__}.{name}' was accessed but is set to -1. Please check the zebra mappings against the zebra's physical configuration"
+                )
+        return value
+
+    @model_validator(mode="after")
+    def ensure_no_duplicate_connections(self):
+        """Check that TTL outputs and sources are mapped to unique integers"""
+
+        integer_fields = {
+            key: value
+            for key, value in self.model_dump().items()
+            if isinstance(value, int) and value != -1
+        }
+        counted_vals = Counter(integer_fields.values())
+        integer_fields_with_duplicates = {
+            k: v for k, v in integer_fields.items() if counted_vals[v] > 1
+        }
+        if len(integer_fields_with_duplicates):
+            raise ValueError(
+                f"Each field in {type(self)} must be mapped to a unique integer. Duplicate fields: {integer_fields_with_duplicates}"
+            )
+        return self
+
+
+class ZebraTTLOutputs(ZebraMappingValidations):
+    """Maps hardware to the Zebra TTL output (1-4) that they're physically wired to, or
+    None if that hardware is not connected. A value of -1 means this hardware is not connected."""
+
+    TTL_EIGER: int = Field(default=-1, ge=-1, le=4)
+    TTL_PILATUS: int = Field(default=-1, ge=-1, le=4)
+    TTL_FAST_SHUTTER: int = Field(default=-1, ge=-1, le=4)
+    TTL_DETECTOR: int = Field(default=-1, ge=-1, le=4)
+    TTL_SHUTTER: int = Field(default=-1, ge=-1, le=4)
+    TTL_XSPRESS3: int = Field(default=-1, ge=-1, le=4)
+    TTL_PANDA: int = Field(default=-1, ge=-1, le=4)
+
+
+class ZebraSources(ZebraMappingValidations):
+    """Maps internal Zebra signal source to their integer PV value"""
+
+    DISCONNECT: int = Field(default=0, ge=0, le=63)
+    IN1_TTL: int = Field(default=1, ge=0, le=63)
+    IN2_TTL: int = Field(default=63, ge=0, le=63)
+    IN3_TTL: int = Field(default=7, ge=0, le=63)
+    IN4_TTL: int = Field(default=10, ge=0, le=63)
+    PC_ARM: int = Field(default=29, ge=0, le=63)
+    PC_GATE: int = Field(default=30, ge=0, le=63)
+    PC_PULSE: int = Field(default=31, ge=0, le=63)
+    AND3: int = Field(default=34, ge=0, le=63)
+    AND4: int = Field(default=35, ge=0, le=63)
+    OR1: int = Field(default=36, ge=0, le=63)
+    PULSE1: int = Field(default=52, ge=0, le=63)
+    PULSE2: int = Field(default=53, ge=0, le=63)
+    SOFT_IN1: int = Field(default=60, ge=0, le=63)
+    SOFT_IN2: int = Field(default=61, ge=0, le=63)
+    SOFT_IN3: int = Field(default=62, ge=0, le=63)
+
+
+class ZebraMapping(ZebraMappingValidations):
+    """Mappings to locate a Zebra device's Ophyd signals based on a specific
+    Zebra's hardware configuration and wiring.
+    """
+
+    # Zebra ophyd signal for connection can be accessed
+    # with, eg, zebra.output.out_pvs[zebra.mapping.outputs.TTL_DETECTOR]
+    outputs: ZebraTTLOutputs = ZebraTTLOutputs()
+
+    # Zebra ophyd signal sources can be mapped to a zebra output by doing, eg,
+    # bps.abs_set(zebra.output.out_pvs[zebra.mapping.outputs.TTL_DETECTOR],
+    # zebra.mapping.sources.AND3)
+    sources: ZebraSources = ZebraSources()
+
+    # Which of the Zebra's four AND gates is used to control the automatic shutter, if it's being used.
+    # After defining, the correct GateControl device can be accessed with, eg,
+    # zebra.logic_gates.and_gates[zebra.mapping.AND_GATE_FOR_AUTO_SHUTTER]. Set to -1 if not being used.
+    AND_GATE_FOR_AUTO_SHUTTER: int = Field(default=-1, ge=-1, le=4)
+
+
+class UnmappedZebraException(Exception):
+    pass

system_tests/test_zebra_system.py

@@ -1,11 +1,12 @@
 import pytest
 
-from dodal.devices.zebra import ArmDemand, Zebra
+from dodal.beamlines.i03 import I03_ZEBRA_MAPPING
+from dodal.devices.zebra.zebra import ArmDemand, Zebra
 
 
 @pytest.fixture()
 async def zebra():
-    zebra = Zebra(name="zebra", prefix="BL03S-EA-ZEBRA-01:")
+    zebra = Zebra(name="zebra", prefix="BL03S-EA-ZEBRA-01:", mapping=I03_ZEBRA_MAPPING)
     yield zebra
     await zebra.pc.arm.set(ArmDemand.DISARM)
 

tests/common/beamlines/test_beamline_utils.py

@@ -16,7 +16,6 @@
 from dodal.devices.focusing_mirror import FocusingMirror
 from dodal.devices.motors import XYZPositioner
 from dodal.devices.smargon import Smargon
-from dodal.devices.zebra import Zebra
 from dodal.log import LOGGER
 from dodal.utils import DeviceInitializationController, make_all_devices
 
@@ -40,7 +39,7 @@ def setup():
 
 
 def test_instantiate_function_makes_supplied_device():
-    device_types = [Zebra, XYZPositioner, Smargon]
+    device_types = [XYZPositioner, Smargon]
     for device in device_types:
         dev = beamline_utils.device_instantiation(
             device, device.__name__, "", False, True, None
@@ -50,7 +49,7 @@ def test_instantiate_function_makes_supplied_device():
 
 def test_instantiating_different_device_with_same_name():
     dev1 = beamline_utils.device_instantiation(  # noqa
-        Zebra, "device", "", False, True, None
+        XYZPositioner, "device", "", False, True, None
     )
     with pytest.raises(TypeError):
         dev2 = beamline_utils.device_instantiation(
@@ -76,11 +75,11 @@ def test_instantiate_v1_function_fake_makes_fake():
 
 def test_instantiate_v2_function_fake_makes_fake():
     RE()
-    fake_zeb: Zebra = beamline_utils.device_instantiation(
-        i03.Zebra, "zebra", "", True, True, None
+    fake_smargon: Smargon = beamline_utils.device_instantiation(
+        i03.Smargon, "smargon", "", True, True, None
     )
-    assert isinstance(fake_zeb, StandardReadable)
-    assert fake_zeb.pc.arm.armed.source.startswith("mock+ca")
+    assert isinstance(fake_smargon, StandardReadable)
+    assert fake_smargon.omega.user_setpoint.source.startswith("mock+ca")
 
 
 def test_clear_devices(RE):

tests/devices/unit_tests/test_shutter.py

@@ -2,7 +2,7 @@
 from ophyd_async.core import DeviceCollector
 from ophyd_async.testing import callback_on_mock_put, set_mock_value
 
-from dodal.devices.zebra_controlled_shutter import (
+from dodal.devices.zebra.zebra_controlled_shutter import (
     ZebraShutter,
     ZebraShutterControl,
     ZebraShutterState,

tests/devices/unit_tests/test_zebra.py

@@ -4,7 +4,7 @@
 from bluesky.run_engine import RunEngine
 from ophyd_async.testing import set_mock_value
 
-from dodal.devices.zebra import (
+from dodal.devices.zebra.zebra import (
     ArmDemand,
     ArmingDevice,
     ArmSource,

tests/devices/unit_tests/test_zebra_constants_mapping.py

@@ -0,0 +1,42 @@
+import pytest
+from ophyd_async.core import DeviceCollector
+
+from dodal.beamlines.i03 import I03_ZEBRA_MAPPING
+from dodal.devices.zebra.zebra import Zebra
+from dodal.devices.zebra.zebra_constants_mapping import (
+    UnmappedZebraException,
+    ZebraMapping,
+    ZebraTTLOutputs,
+)
+
+
+async def fake_zebra(zebra_mapping: ZebraMapping):
+    async with DeviceCollector(mock=True):
+        zebra = Zebra(mapping=zebra_mapping, name="", prefix="")
+    return zebra
+
+
+async def test_exception_when_accessing_mapping_set_to_minus_1():
+    mapping_no_output = ZebraMapping(outputs=ZebraTTLOutputs())
+    with pytest.raises(
+        UnmappedZebraException,
+        match="'ZebraTTLOutputs.TTL_EIGER' was accessed but is set to -1. Please check the zebra mappings against the zebra's physical configuration",
+    ):
+        zebra = await fake_zebra(mapping_no_output)
+        zebra.mapping.outputs.TTL_EIGER  # noqa: B018
+
+
+def test_exception_when_multiple_fields_set_to_same_integer():
+    expected_error_dict = {"TTL_DETECTOR": 1, "TTL_PANDA": 1}
+    with pytest.raises(
+        ValueError,
+        match=f"must be mapped to a unique integer. Duplicate fields: {expected_error_dict}",
+    ):
+        ZebraMapping(outputs=ZebraTTLOutputs(TTL_DETECTOR=1, TTL_PANDA=1))
+
+
+async def test_validly_mapped_zebra_is_happy():
+    zebra = await fake_zebra(zebra_mapping=I03_ZEBRA_MAPPING)
+    assert zebra.mapping.outputs.TTL_DETECTOR == 1
+    assert zebra.mapping.sources.DISCONNECT == 0
+    assert zebra.mapping.AND_GATE_FOR_AUTO_SHUTTER == 2