diff --git a/slac_devices/reader.py b/slac_devices/reader.py index 4bc1b53..3f9386d 100644 --- a/slac_devices/reader.py +++ b/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]: diff --git a/slac_devices/toroid.py b/slac_devices/toroid.py new file mode 100644 index 0000000..cf6a253 --- /dev/null +++ b/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 diff --git a/slac_devices/wire.py b/slac_devices/wire.py index 2a2c5a1..61bdcec 100644 --- a/slac_devices/wire.py +++ b/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