Merge pull request #129 from slaclab/abstract_obj_classes

create measurement class and io function

lcls_tools/common/devices/ict.py

@@ -0,0 +1,34 @@
+from pydantic import field_validator, SerializeAsAny
+
+from lcls_tools.common.devices.device import Device, PVSet, ControlInformation, Metadata
+from epics import PV
+
+
+class ICTPVSet(PVSet):
+    """
+    We list the potential PVs below and only
+    use the ones that are set to be PV-typed after
+    initialisation.
+    """
+
+    charge_nC: PV
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+
+    @field_validator("*", mode="before")
+    def validate_pv_fields(cls, v: str):
+        """Convert each PV string from YAML into a PV object"""
+        return PV(v)
+
+
+class ICTControlInformation(ControlInformation):
+    PVs: SerializeAsAny[ICTPVSet]
+
+
+class ICT(Device):
+    controls_information: SerializeAsAny[ICTControlInformation]
+    metadata: SerializeAsAny[Metadata]
+
+    def get_charge(self) -> float:
+        return self.controls_information.PVs.charge_nC.get(as_numpy=True)

lcls_tools/common/io.py

@@ -0,0 +1,28 @@
+from typing import Optional
+
+from lcls_tools.common.measurements.measurement import Measurement
+
+
+class HDF5IO:
+    def __init__(self):
+        pass
+
+    def write(
+            self,
+            measurement_data: dict,
+            measurement_obj: Measurement,
+            filename: Optional[str] = None
+    ):
+        """
+        Write data to h5file
+        """
+        raise NotImplementedError
+
+    def read(
+            self,
+            filename: Optional[str] = None
+    ):
+        """
+        Read data from h5file
+        """
+        raise NotImplementedError

lcls_tools/common/measurements/beam_charge.py

@@ -0,0 +1,51 @@
+import numpy as np
+from pydantic import PositiveFloat
+
+from lcls_tools.common.devices.ict import ICT
+from lcls_tools.common.measurements.measurement import Measurement
+import time
+
+from lcls_tools.common.measurements.utils import calculate_statistics
+
+
+class BeamChargeMeasurement(Measurement):
+    name = "beam_charge"
+    ict_monitor: ICT
+    wait_time: PositiveFloat = 1.0
+
+    def measure(self, n_shots: int = 1) -> dict:
+        """
+        Measure the bunch charge using an ICT monitor.
+
+        Parameters:
+        - n_shots (int, optional): The number of measurements to perform. Defaults to 1.
+
+        Returns:
+        dict: A dictionary containing the measured bunch charge values and additional
+        statistics if multiple shots are taken.
+
+        If n_shots is 1, the function returns a dictionary with the key "bunch_charge_nC"
+        and the corresponding single measurement value.
+
+        If n_shots is greater than 1, the function performs multiple measurements with
+        the specified wait time and returns a dictionary with the key "bunch_charge_nC"
+        containing a list of measured values. Additionally, statistical information
+        (mean, standard deviation, etc.) is included in the dictionary.
+
+
+        """
+        if n_shots == 1:
+            return {"bunch_charge_nC": self.ict_monitor.get_charge()}
+        elif n_shots > 1:
+            bunch_charges = []
+            for i in range(n_shots):
+                bunch_charges += [self.ict_monitor.get_charge()]
+                time.sleep(self.wait_time)
+
+            # add statistics to results
+            results = {"bunch_charge_nC": bunch_charges}
+            results = results | calculate_statistics(
+                np.array(bunch_charges), "bunch_charge_nC"
+            )
+
+            return results

lcls_tools/common/measurements/measurement.py

@@ -0,0 +1,15 @@
+from abc import ABC, abstractmethod
+from typing import Optional
+
+from pydantic import BaseModel, DirectoryPath
+
+
+class Measurement(BaseModel, ABC):
+    name: str
+    save_data: bool = True
+    save_location: Optional[DirectoryPath] = None
+
+    @abstractmethod
+    def measure(self, **kwargs) -> dict:
+        """ Implements a measurement and returns a dictionary with the results"""
+        raise NotImplementedError

lcls_tools/common/measurements/utils.py

@@ -0,0 +1,11 @@
+import numpy as np
+from numpy import ndarray
+
+
+def calculate_statistics(data: ndarray, name):
+    return {
+        f"{name}_mean": np.mean(data),
+        f"{name}_std": np.std(data),
+        f"{name}_q05": np.quantile(data, 0.05),
+        f"{name}_q95": np.quantile(data, 0.95),
+    }