Correction of serialized magnets array and new tune correction test with serialized magnets

Author: gupichon-soleil

pyaml/arrays/serialized_magnet_array.py

@@ -5,7 +5,7 @@
 from ..magnet.serialized_magnet import SerializedMagnets
 from .element_array import ElementArray
 
-# TODO handle aggregator for CFM
+# TODO handle aggregator for serialized magnets
 
 
 class RWMagnetStrengths(ReadWriteFloatArray):
@@ -16,19 +16,13 @@ def __init__(self, name: str, magnets: list[SerializedMagnets]):
 
     # Gets the values
     def get(self) -> np.array:
-        r = np.zeros(self.__nb)
-        idx = 0
-        for m in self.__magnets:
-            r[idx : idx + m.get_nb_magnets()] = m.strengths.get()
-            idx += m.get_nb_magnets()
-        return r
+        return np.array([m.strengths.get() for m in self.__magnets])
 
     # Sets the values
     def set(self, value: np.array):
-        nvalue = np.ones(self.__nb) * value if isinstance(value, float) else value
-        for idx, m in enumerate(self.__magnets):
-            m.strengths.set(nvalue[idx])
-            idx += m.get_nb_magnets()
+        nvalue = np.ones(len(self.__magnets)) * value if isinstance(value, float) else value
+        for value, m in zip(nvalue, self.__magnets, strict=True):
+            m.strengths.set(value)
 
     # Sets the values and waits that the read values reach their setpoint
     def set_and_wait(self, value: np.array):
@@ -50,18 +44,13 @@ def __init__(self, name: str, magnets: list[SerializedMagnets]):
 
     # Gets the values
     def get(self) -> np.array:
-        r = np.zeros(self.__nb)
-        idx = 0
-        for m in self.__magnets:
-            r[idx : idx + m.get_nb_magnets()] = m.hardwares.get()
-            idx += m.get_nb_magnets()
-        return r
+        return np.array([m.hardwares.get() for m in self.__magnets])
 
     # Sets the values
     def set(self, value: np.array):
-        nvalue = np.ones(self.__nb) * value if isinstance(value, float) else value
-        for idx, m in enumerate(self.__magnets):
-            m.hardwares.set(nvalue[idx])
+        nvalue = np.ones(len(self.__magnets)) * value if isinstance(value, float) else value
+        for value, m in zip(nvalue, self.__magnets, strict=True):
+            m.hardwares.set(value)
 
     # Sets the values and waits that the read values reach their setpoint
     def set_and_wait(self, value: np.array):
@@ -103,11 +92,7 @@ def __init__(
         self.__rwhardwares = RWMagnetHardwares(arrayName, magnets)
 
         if use_aggregator:
-            raise (
-                PyAMLException(
-                    "Aggregator not implemented for CombinedFunctionMagnetArray"
-                )
-            )
+            raise (PyAMLException("Aggregator not implemented for CombinedFunctionMagnetArray"))
 
     @property
     def strengths(self) -> RWMagnetStrengths:

pyaml/common/element_holder.py

@@ -204,7 +204,7 @@ def fill_serialized_magnet_array(self, arrayName: str, elementNames: list[str]):
             self.__SERIALIZED_MAGNETS_ARRAYS,
         )
 
-    def get_serialized_magnet(self, name: str) -> Magnet:
+    def get_serialized_magnet(self, name: str) -> SerializedMagnets:
         return self.__get("SerializedMagnets", name, self.__SERIALIZED_MAGNETS)
 
     def add_serialized_magnet(self, m: Magnet):

tests/config/sr_serialized_magnets.yaml

@@ -12,6 +12,11 @@ arrays:
     name: series
     elements:
       - mySeriesOfMagnets
+  - type: pyaml.arrays.serialized_magnet
+    name: QForTune
+    elements:
+      - QD2?
+      - QF1?
 devices:
   - type: pyaml.magnet.serialized_magnet
     name: mySeriesOfMagnets
@@ -34,3 +39,211 @@ devices:
         type: tango.pyaml.attribute
         attribute: srmag/ps-corr-sh1/c01-a-ch1/current
         unit: A
+  - type: pyaml.magnet.serialized_magnet
+    name: QF1A
+    function: B1
+    elements:
+      - QF1A-C05
+      - QF1A-C06
+      - QF1A-C07
+      - QF1A-C08
+      - QF1A-C09
+      - QF1A-C10
+      - QF1A-C11
+      - QF1A-C12
+      - QF1A-C13
+      - QF1A-C14
+      - QF1A-C15
+      - QF1A-C16
+      - QF1A-C17
+      - QF1A-C18
+      - QF1A-C19
+      - QF1A-C20
+      - QF1A-C21
+      - QF1A-C22
+      - QF1A-C23
+      - QF1A-C24
+      - QF1A-C25
+      - QF1A-C26
+      - QF1A-C27
+      - QF1A-C28
+      - QF1A-C29
+      - QF1A-C30
+      - QF1A-C31
+      - QF1A-C32
+      - QF1A-C01
+      - QF1A-C02
+      - QF1A-C03
+    model:
+      type: pyaml.magnet.linear_serialized_model
+      calibration_factors: 1
+      calibration_offsets: 0.0
+      unit: m-1
+      curves:
+        type: pyaml.configuration.csvcurve
+        file: sr/magnet_models/QF1_strength.csv
+      powerconverter:
+        type: tango.pyaml.attribute
+        attribute: srmag/ps-corr-sh1/c01-a-ch1/current
+        unit: A
+  - type: pyaml.magnet.serialized_magnet
+    name: QF1E
+    function: B1
+    elements:
+      - QF1E-C04
+      - QF1E-C05
+      - QF1E-C06
+      - QF1E-C07
+      - QF1E-C08
+      - QF1E-C09
+      - QF1E-C10
+      - QF1E-C11
+      - QF1E-C12
+      - QF1E-C13
+      - QF1E-C14
+      - QF1E-C15
+      - QF1E-C16
+      - QF1E-C17
+      - QF1E-C18
+      - QF1E-C19
+      - QF1E-C20
+      - QF1E-C21
+      - QF1E-C22
+      - QF1E-C23
+      - QF1E-C24
+      - QF1E-C25
+      - QF1E-C26
+      - QF1E-C27
+      - QF1E-C28
+      - QF1E-C29
+      - QF1E-C30
+      - QF1E-C31
+      - QF1E-C32
+      - QF1E-C01
+      - QF1E-C02
+    model:
+      type: pyaml.magnet.linear_serialized_model
+      calibration_factors: 1
+      calibration_offsets: 0.0
+      unit: m-1
+      curves:
+        type: pyaml.configuration.csvcurve
+        file: sr/magnet_models/QF1_strength.csv
+      powerconverter:
+        type: tango.pyaml.attribute
+        attribute: srmag/ps-corr-sh1/c01-a-ch1/current
+        unit: A
+  - type: pyaml.magnet.serialized_magnet
+    name: QD2A
+    function: B1
+    elements:
+      - QD2A-C05
+      - QD2A-C06
+      - QD2A-C07
+      - QD2A-C08
+      - QD2A-C09
+      - QD2A-C10
+      - QD2A-C11
+      - QD2A-C12
+      - QD2A-C13
+      - QD2A-C14
+      - QD2A-C15
+      - QD2A-C16
+      - QD2A-C17
+      - QD2A-C18
+      - QD2A-C19
+      - QD2A-C20
+      - QD2A-C21
+      - QD2A-C22
+      - QD2A-C23
+      - QD2A-C24
+      - QD2A-C25
+      - QD2A-C26
+      - QD2A-C27
+      - QD2A-C28
+      - QD2A-C29
+      - QD2A-C30
+      - QD2A-C31
+      - QD2A-C32
+      - QD2A-C01
+      - QD2A-C02
+      - QD2A-C03
+    model:
+      type: pyaml.magnet.linear_serialized_model
+      calibration_factors: 1
+      calibration_offsets: 0.0
+      unit: m-1
+      curves:
+        type: pyaml.configuration.csvcurve
+        file: sr/magnet_models/QD2_strength.csv
+      powerconverter:
+        type: tango.pyaml.attribute
+        attribute: srmag/ps-corr-sh1/c01-a-ch1/current
+        unit: A
+  - type: pyaml.magnet.serialized_magnet
+    name: QD2E
+    function: B1
+    elements:
+      - QD2E-C04
+      - QD2E-C05
+      - QD2E-C06
+      - QD2E-C07
+      - QD2E-C08
+      - QD2E-C09
+      - QD2E-C10
+      - QD2E-C11
+      - QD2E-C12
+      - QD2E-C13
+      - QD2E-C14
+      - QD2E-C15
+      - QD2E-C16
+      - QD2E-C17
+      - QD2E-C18
+      - QD2E-C19
+      - QD2E-C20
+      - QD2E-C21
+      - QD2E-C22
+      - QD2E-C23
+      - QD2E-C24
+      - QD2E-C25
+      - QD2E-C26
+      - QD2E-C27
+      - QD2E-C28
+      - QD2E-C29
+      - QD2E-C30
+      - QD2E-C31
+      - QD2E-C32
+      - QD2E-C01
+      - QD2E-C02
+    model:
+      type: pyaml.magnet.linear_serialized_model
+      calibration_factors: 1
+      calibration_offsets: 0.0
+      unit: m-1
+      curves:
+        type: pyaml.configuration.csvcurve
+        file: sr/magnet_models/QD2_strength.csv
+      powerconverter:
+        type: tango.pyaml.attribute
+        attribute: srmag/ps-corr-sh1/c01-a-ch1/current
+        unit: A
+  - type: pyaml.diagnostics.tune_monitor
+    name: BETATRON_TUNE
+    tune_h:
+      type: tango.pyaml.attribute_read_only
+      attribute: srdiag/beam-tune/main/Qh
+      unit: mm
+    tune_v:
+      type: tango.pyaml.attribute_read_only
+      attribute: srdiag/beam-tune/main/Qv
+      unit: mm
+  - type: pyaml.tuning_tools.tune
+    name: DEFAULT_TUNE_CORRECTION
+    quad_array_name: QForTune
+    betatron_tune_name: BETATRON_TUNE
+    response_matrix: file:tune_response.json
+  - type: pyaml.tuning_tools.tune_response_matrix
+    name: DEFAULT_TUNE_RESPONSE_MATRIX
+    quad_array_name: QForTune
+    betatron_tune_name: BETATRON_TUNE
+    quad_delta: 1e-4

tests/test_serialized_magnets.py

@@ -48,6 +48,11 @@ def test_config_load(sr_file):
     assert check_no_diff(currents)
 
 
+def print_magnet_list(magnet_list: list) -> None:
+    for magnet in magnet_list:
+        print(f"- {magnet.FamName}")
+
+
 @pytest.mark.parametrize(
     "sr_file",
     [
@@ -57,6 +62,8 @@ def test_config_load(sr_file):
 def test_magnet_modification(sr_file):
     sr = Accelerator.load(sr_file, use_fast_loader=True, ignore_external=True)
 
+    print(sr.yellow_pages)
+
     sm: SerializedMagnets = sr.design.get_serialized_magnet("mySeriesOfMagnets")
     element_names = sm._SerializedMagnets__elements
 
@@ -98,3 +105,53 @@ def test_magnet_modification(sr_file):
         print(element_names[ii], sm.strengths.elements[ii].get())
 
     assert check_no_diff([sm.strengths.elements[ii].get() for ii in range(len(sm.strengths.elements))])
+
+
+@pytest.mark.parametrize(
+    "sr_file",
+    [
+        "tests/config/sr_serialized_magnets.yaml",
+    ],
+)
+def test_tune(sr_file):
+    sr = Accelerator.load(sr_file, use_fast_loader=True, ignore_external=True)
+    sr.design.get_lattice().disable_6d()
+
+    quadForTuneDesign = sr.design.get_serialized_magnets("QForTune")
+    tune_monitor = sr.design.get_betatron_tune_monitor("BETATRON_TUNE")
+    # Build tune response matrix
+    tunemat = np.zeros((len(quadForTuneDesign), 2))
+
+    # Magnet are not actually in series. Here a trick to set them to the same strengths
+    for m in quadForTuneDesign:
+        strength = m.strengths.get()
+        m.strengths.set(strength)
+    tune = tune_monitor.tune.get()
+    print(f"tune={tune}")
+
+    for idx, m in enumerate(quadForTuneDesign):
+        strength = m.strengths.get()
+        m.strengths.set(strength + 1e-4)
+        dq = tune_monitor.tune.get() - tune
+        tunemat[idx] = dq * 1e4
+        m.strengths.set(strength)
+
+    # Compute correction matrix
+    correctionmat = np.linalg.pinv(tunemat.T)
+    print(f"correctionmat.shape={correctionmat.shape}")
+    print(f"correctionmat={correctionmat}")
+
+    # Correct tune
+    strengths = quadForTuneDesign.strengths.get()
+    print(f"len(strengths)={len(strengths)}")
+    print(f"strengths={strengths}")
+    strengths += np.matmul(correctionmat, [0.1, 0.05])  # Ask for correction [dqx,dqy]
+    print(f"strengths={strengths}")
+    quadForTuneDesign.strengths.set(strengths)
+    newTune = tune_monitor.tune.get()
+    print(f"newTune={newTune}")
+    diffTune = newTune - tune
+
+    print(f"diffTune={diffTune}")
+    assert np.abs(diffTune[0] - 0.1) < 1e-3
+    assert np.abs(diffTune[1] - 0.05) < 1e-3

tests/test_tune.py

@@ -29,9 +29,13 @@ def test_tune():
 
     # Compute correction matrix
     correctionmat = np.linalg.pinv(tunemat.T)
+    print(correctionmat.shape)
+    print(correctionmat)
 
     # Correct tune
     strs = quadForTuneDesign.strengths.get()
+    print(len(strs))
+    print(strs)
     strs += np.matmul(correctionmat, [0.1, 0.05])  # Ask for correction [dqx,dqy]
     quadForTuneDesign.strengths.set(strs)
     newTune = tune_monitor.tune.get()