dispersion measurement

Author: kparasch

examples/ESRF_ORM_example/measure_dispersion.py

@@ -0,0 +1,36 @@
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+
+from pyaml.accelerator import Accelerator
+from pyaml.tuning_tools.dispersion import ConfigModel as Disp_ConfigModel
+from pyaml.tuning_tools.dispersion import Dispersion
+
+parent_folder = Path(__file__).parent
+config_path = parent_folder.parent.parent.joinpath(
+    "tests", "config", "EBSOrbit.yaml"
+).resolve()
+sr = Accelerator.load(config_path)
+element_holder = sr.design
+
+dispersion = Dispersion(
+    cfg=Disp_ConfigModel(
+        bpm_array_name="BPM",
+        rf_plant_name="RF",
+        frequency_delta=10,
+    ),
+    element_holder=element_holder,
+)
+
+dispersion.measure()
+# orm.save(parent_folder / Path("ideal_orm.json"))
+# orm.save(parent_folder / Path("ideal_orm.yaml"), with_type="yaml")
+# orm.save(parent_folder / Path("ideal_orm.npz"), with_type="npz")
+
+# ormdata = orm.get()
+dispersion_data = dispersion.get()
+
+fig = plt.figure()
+ax = fig.add_subplot(111)
+ax.plot(dispersion_data["frequency_response_x"])
+plt.show()

pyaml/external/pySC_interface.py

@@ -13,6 +13,7 @@ def __init__(
         bpm_array_name: str = "BPM",
         hcorr_array_name: str = "HCorr",
         vcorr_array_name: str = "VCorr",
+        rf_plant_name: str = "RF",
     ):
         self.element_holder = element_holder
 
@@ -38,6 +39,9 @@ def __init__(
             name: ii for ii, name in enumerate(self.vcorr_names)
         }
 
+        self.rf_plant_name = rf_plant_name
+        self.rf_plant = element_holder.get_rf_plant(self.rf_plant_name)
+
     def get_orbit(self) -> Tuple[np.array, np.array]:
         # we should wait here somehow according to polling rate
         positions = self.bpm_array.positions.get()
@@ -154,3 +158,10 @@ def set_many(self, names_values: Dict[str, float]) -> None:
             )  # ideally set_and_wait but not implemented
 
         return
+
+    def get_rf_main_frequency(self) -> float:
+        return self.rf_plant.frequency.get()
+
+    def set_rf_main_frequency(self, value: float) -> None:
+        self.rf_plant.frequency.set(value)
+        return

pyaml/tuning_tools/dispersion.py

@@ -0,0 +1,64 @@
+import logging
+from pathlib import Path
+from typing import List, Optional
+
+from pydantic import BaseModel, ConfigDict
+
+from ..common.element_holder import ElementHolder
+from ..common.exception import PyAMLException
+from ..external.pySC import pySC
+from ..external.pySC.pySC.apps import measure_dispersion
+from ..external.pySC.pySC.apps.codes import DispersionCode
+from ..external.pySC_interface import pySCInterface
+
+logger = logging.getLogger(__name__)
+
+PYAMLCLASS = "OrbitResponseMatrix"
+
+
+class ConfigModel(BaseModel):
+    model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid")
+
+    bpm_array_name: str
+    rf_plant_name: str
+    frequency_delta: float
+
+
+class Dispersion(object):
+    def __init__(self, element_holder: ElementHolder, cfg: ConfigModel):
+        self._cfg = cfg
+
+        self.element_holder = element_holder
+        self.bpm_array_name = cfg.bpm_array_name
+        self.rf_plant_name = cfg.rf_plant_name
+        self.frequency_delta = cfg.frequency_delta
+        self.latest_measurement = None
+
+    def measure(self):
+        interface = pySCInterface(
+            element_holder=self.element_holder,
+            bpm_array_name=self.bpm_array_name,
+            rf_plant_name=self.rf_plant_name,
+        )
+
+        generator = measure_dispersion(
+            interface=interface,
+            delta=self.frequency_delta,
+            skip_save=True,
+        )
+
+        pySC.disable_pySC_rich()
+        _, measurement = next(generator)
+        for _, _ in generator:
+            pass
+
+        dispersion_data = measurement.dispersion_data
+        # contains also pre-processed data
+
+        # dispersion_data.output_names = self.element_holder.get_bpms(
+        #     self.bpm_array_name
+        # ).names()
+        self.latest_measurement = dispersion_data.model_dump()
+
+    def get(self):
+        return self.latest_measurement