orbit object tuning tool for orbit correction

Author: kparasch

examples/ESRF_ORM_example/correct_orbit.py

@@ -4,29 +4,42 @@
 
 import matplotlib.pyplot as plt
 import numpy as np
-from pyaml.external.pySC.pySC.apps import orbit_correction
-from pyaml.external.pySC.pySC.tuning.response_matrix import ResponseMatrix
 
 from pyaml.accelerator import Accelerator
-from pyaml.external.pySC import pySC
-from pyaml.external.pySC_interface import pySCInterface
-
-logging.getLogger("pyaml.external.pySC").setLevel(logging.WARNING)
+from pyaml.tuning_tools.orbit import ConfigModel as Orbit_ConfigModel
+from pyaml.tuning_tools.orbit import Orbit
 
 parent_folder = Path(__file__).parent
 pyaml_folder = parent_folder.parent.parent
 config_path = pyaml_folder.joinpath("tests/config/EBSOrbit.yaml").resolve()
 sr = Accelerator.load(config_path)
 ebs = sr.design
 
+orbit = Orbit(
+    element_holder=ebs,
+    cfg=Orbit_ConfigModel(
+        bpm_array_name="BPM",
+        hcorr_array_name="HCorr",
+        vcorr_array_name="VCorr",
+        singular_values=162,
+        response_matrix_file=str(
+            pyaml_folder.joinpath("examples/ESRF_ORM_example/ideal_orm.json").resolve()
+        ),
+    ),
+)
+
+## get reference
+ref_h, ref_v = orbit.element_holder.get_bpms("BPM").positions.get().T
+reference = np.concat((ref_h, ref_v))
+########################################################
+
+
+## generate some orbit
 std_kick = 1e-6
 hcorr = ebs.get_magnets("HCorr")
 vcorr = ebs.get_magnets("VCorr")
 bpms = ebs.get_bpms("BPM")
 
-x0, y0 = bpms.positions.get().T
-reference = np.concat([x0, y0])
-
 # mangle orbit
 hcorr.strengths.set(
     hcorr.strengths.get() + std_kick * np.random.normal(size=len(hcorr))
@@ -41,36 +54,15 @@
     "R.m.s. orbit before correction "
     f"H: {1e6 * std_bc[0]: .1f} µm, V: {1e6 * std_bc[1]: .1f} µm."
 )
+########################################################
 
-response_matrix = ResponseMatrix.from_json(parent_folder / Path("ideal_orm.json"))
-interface = pySCInterface(element_holder=ebs)
-
-# response_matrix.disable_inputs(["SH3E-C03-V"])
-
-for _ in range(5):
-    trims = orbit_correction(
-        interface=interface,
-        response_matrix=response_matrix,
-        method="svd_values",
-        parameter=100,
-        zerosum=True,
-        apply=True,
-        plane="H",
-        reference=reference,
-    )
-    trims = orbit_correction(
-        interface=interface,
-        response_matrix=response_matrix,
-        method="svd_values",
-        parameter=100,
-        zerosum=False,
-        apply=True,
-        plane="V",
-        reference=reference,
-    )
-
-# print("SH3E-C03-V" in trims.keys())
+## Correct the orbit
+orbit.correct(reference=reference)
+# orbit.correct(plane="H")
+# orbit.correct(plane="V")
+########################################################
 
+## inspect orbit correction
 positions_ac = bpms.positions.get()
 std_ac = np.std(positions_ac, axis=0)
 print(

pyaml/tuning_tools/orbit.py

@@ -0,0 +1,90 @@
+import logging
+from pathlib import Path
+from typing import List, Literal, Optional
+
+from pydantic import BaseModel, ConfigDict
+
+from ..common.element_holder import ElementHolder
+from ..external.pySC import pySC
+from ..external.pySC.pySC import ResponseMatrix
+from ..external.pySC.pySC.apps import orbit_correction
+from ..external.pySC_interface import pySCInterface
+
+logger = logging.getLogger(__name__)
+logging.getLogger("pyaml.external.pySC").setLevel(logging.WARNING)
+
+PYAMLCLASS = "Orbit"
+
+
+class ConfigModel(BaseModel):
+    model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid")
+
+    bpm_array_name: str
+    hcorr_array_name: str
+    vcorr_array_name: str
+    singular_values: int
+    response_matrix_file: str
+
+
+class Orbit(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.hcorr_array_name = cfg.hcorr_array_name
+        self.vcorr_array_name = cfg.vcorr_array_name
+        self.singular_values = cfg.singular_values
+        self.response_matrix = ResponseMatrix.from_json(Path(cfg.response_matrix_file))
+
+    def correct(
+        self,
+        reference=None,
+        gain: float = 1.0,
+        plane: Optional[Literal["H", "V"]] = None,
+    ):
+        interface = pySCInterface(
+            element_holder=self.element_holder,
+            bpm_array_name=self.bpm_array_name,
+            hcorr_array_name=self.hcorr_array_name,
+            vcorr_array_name=self.vcorr_array_name,
+        )
+
+        if plane is None or plane == "H":
+            trims_h = orbit_correction(
+                interface=interface,
+                response_matrix=self.response_matrix,
+                method="svd_values",
+                parameter=self.singular_values,
+                zerosum=True,
+                apply=False,
+                plane="H",
+                reference=reference,
+            )
+
+        if plane is None or plane == "V":
+            trims_v = orbit_correction(
+                interface=interface,
+                response_matrix=self.response_matrix,
+                method="svd_values",
+                parameter=self.singular_values,
+                zerosum=False,
+                apply=False,
+                plane="V",
+                reference=reference,
+            )
+
+        # this is now the only place where interface.set_many and get_many are used.
+        # if we can remove it from here then the pySCInterface will be even simpler.
+        correctors = self.response_matrix.input_names
+        data_to_send = interface.get_many(correctors)
+        if plane is None or plane == "H":
+            for name in trims_h.keys():
+                data_to_send[name] += trims_h[name] * gain
+        if plane is None or plane == "V":
+            for name in trims_v.keys():
+                data_to_send[name] += trims_v[name] * gain
+
+        interface.set_many(data_to_send)
+
+        return