Add Toroid device support and update WireMetadata for charge toroids

slac_devices/reader.py

@@ -9,6 +9,7 @@
 from slac_devices.pmt import PMT, PMTCollection
 from slac_devices.bpm import BPM, BPMCollection
 from slac_devices.tcav import TCAV, TCAVCollection
+from slac_devices.toroid import Toroid, ToroidCollection
 from slac_devices.area import Area
 from slac_devices.beampath import Beampath
 
@@ -21,6 +22,7 @@
     "lblms",
     "pmts",
     "tcavs",
+    "toroids",
 }
 
 _CONSTRUCTOR_MAP = {
@@ -31,8 +33,10 @@
     "lblms": LBLM,
     "pmts": PMT,
     "tcavs": TCAV,
+    "toroids": Toroid,
 }
 
+
 def create_device(name):
     data = slac_db.db_to_yaml.get_device(name)
     if data is None:
@@ -46,6 +50,7 @@ def create_device(name):
         print(field_error)
         return None
 
+
 def create_magnet(
     area: str = None, name: str = None
 ) -> Union[None, Magnet, MagnetCollection]:
@@ -167,6 +172,21 @@ def create_pmt(area: str = None, name: str = None) -> Union[None, PMT]:
         return PMTCollection(**device_data)
 
 
+def create_toroid(area: str = None, name: str = None) -> Union[None, Toroid]:
+    device_data = slac_db.get_device(area=area, device_type="toroids", name=name)
+    if not device_data:
+        return None
+    if name:
+        try:
+            device_data.update({"name": name})
+            return Toroid(**device_data)
+        except ValidationError as field_error:
+            print(field_error)
+            return None
+    else:
+        return ToroidCollection(**device_data)
+
+
 def create_area(
     area: str = None, device_types: Optional[set] = None
 ) -> Union[None, Area]:

slac_devices/toroid.py

@@ -0,0 +1,57 @@
+from typing import Dict
+
+from pydantic import BaseModel, SerializeAsAny, field_validator
+
+from slac_devices.device import Device, PVSet, ControlInformation, Metadata
+from slac_timing import Buffer
+from epics import PV
+
+
+class ToroidPVSet(PVSet):
+    tmit: PV
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+
+
+class ToroidControlInformation(ControlInformation):
+    PVs: SerializeAsAny[ToroidPVSet]
+
+    def __init__(self, *args, **kwargs):
+        super(ToroidControlInformation, self).__init__(*args, **kwargs)
+
+
+class ToroidMetadata(Metadata):
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+
+
+class Toroid(Device):
+    controls_information: SerializeAsAny[ToroidControlInformation]
+    metadata: SerializeAsAny[ToroidMetadata]
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+
+    @property
+    def tmit(self):
+        """Get current TMIT value."""
+        return self.controls_information.PVs.tmit.get()
+
+    def tmit_buffer(self, buffer: Buffer, **kwargs):
+        """Retrieve per-pulse TMIT data from timing buffer."""
+        return buffer.get(f"{self.controls_information.control_name}:TMIT", **kwargs)
+
+
+class ToroidCollection(BaseModel):
+    toroids: Dict[str, SerializeAsAny[Toroid]]
+
+    @field_validator("toroids", mode="before")
+    def validate_toroids(cls, v) -> Dict[str, Toroid]:
+        for name, toroid in v.items():
+            if isinstance(toroid, Toroid):
+                continue
+            toroid = dict(toroid)
+            toroid.update({"name": name})
+            v.update({name: toroid})
+        return v

slac_devices/wire.py

@@ -107,6 +107,7 @@ class WireMetadata(Metadata):
     default_detector: str
     tmitloss: Optional[TMITLossBPMs] = None
     jitter_bpms: Optional[List[str]] = None
+    charge_toroids: Optional[List[str]] = None
     type: str
     wire_type: str