Merge pull request #109 from python-accelerator-middle-layer/bugfix/orbit_correction

Bugfix/orbit correction

Author: JeanLucPons

pyaml/tuning_tools/orbit.py

@@ -98,21 +98,18 @@ def correct(
             trims = trims_v
             corr_array = self._vcorr
 
-        corrector_was_used = {key: False for key in trims.keys()}
         corrector_names = corr_array.names()
+        corrector_to_index = {name: idx for idx, name in enumerate(corrector_names)}
         data_to_send = corr_array.strengths.get()
-        for idx, name in enumerate(corrector_names):
-            data_to_send[idx] += trims[name] * gain
-            corrector_was_used[name] = True
-
-        # check that all corrector trims will be sent
-        for key in trims.keys():
-            if not corrector_was_used[key]:
+        for name in trims.keys():
+            idx = corrector_to_index.get(name, None)
+            if idx is None:
                 raise PyAMLException(
-                    f"Corrector {key} was not used in the orbit correction. "
-                    "There is an inconcistency between corrector arrays and "
+                    f"Corrector {name} not found in the magnet array for orbit corr. "
+                    "Possible inconcistency between corrector arrays and "
                     "response matrix."
                 )
+            data_to_send[idx] += trims[name] * gain
 
         corr_array.strengths.set(data_to_send)
         return

tests/test_tuning_orbit_correction.py

@@ -24,6 +24,11 @@ def test_tuning_orm():
     vcorr = element_holder.get_magnets("VCorr")
     bpms = element_holder.get_bpms("BPM")
 
+    x, y = bpms.positions.get().T  # get reference orbit
+    reference = np.concat((x, y))
+    # there should be nothing to correct, but still work
+    element_holder.orbit.correct(reference=reference)
+
     # mangle orbit
     hcorr.strengths.set(
         hcorr.strengths.get() + std_kick * np.random.normal(size=len(hcorr))