Orbit correction: RF response for orbit correction and handling of weights

.github/workflows/tests.yml

@@ -32,7 +32,7 @@ jobs:
         pip install scipy
         pip install pydantic
         pip install accelerator-toolbox
-        pip install accelerator-commissioning==1.0.0
+        pip install accelerator-commissioning==1.1.0
         pip install matplotlib
         pip install h5py
         pip install PyYaml

examples/BESSY2_example/bessy2-orbit.ipynb

@@ -31,7 +31,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": null,
    "id": "aff4d20f",
    "metadata": {},
    "outputs": [
@@ -45,7 +45,8 @@
    ],
    "source": [
     "# Load the configuration\n",
-    "# Remember to change the prefix for the live mode to the one matching your virtual accelerator before loading.\n",
+    "# Remember to change the prefix for the live mode to the one matching your virtual\n",
+    "# accelerator before loading.\n",
     "sr = Accelerator.load(\"BESSY2Orbit.yaml\")"
    ]
   },
@@ -62,7 +63,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": null,
    "id": "9caca764",
    "metadata": {},
    "outputs": [],
@@ -91,8 +92,8 @@
     "        \"matrix\": orm_data[\"matrix\"],\n",
     "        \"input_names\": orm_data[\"input_names\"],\n",
     "        \"output_names\": orm_data[\"output_names\"],\n",
-    "        \"inputs_plane\": orm_data[\"inputs_plane\"],\n",
-    "        \"outputs_plane\": orm_data[\"outputs_plane\"],\n",
+    "        \"input_planes\": orm_data[\"input_planes\"],\n",
+    "        \"output_planes\": orm_data[\"output_planes\"],\n",
     "    }\n",
     "    json.dump(ORM_data, open(\"orm.json\", \"w\"))"
    ]
@@ -669,7 +670,7 @@
     }
    ],
    "source": [
-    "plot(orbit.get()[:, 1])"
+    "plt.plot(orbit.get()[:, 1])"
    ]
   },
   {
@@ -826,8 +827,8 @@
    "outputs": [],
    "source": [
     "# Correct the orbit\n",
-    "# If you are using the ORM measured on design to correct on live you need to set the gain\n",
-    "# since the unit for the BPMs are not the same for both modes yet.\n",
+    "# If you are using the ORM measured on design to correct on live you need to set the\n",
+    "# gain since the unit for the BPMs are not the same for both modes yet.\n",
     "\n",
     "sr.live.orbit.correct(gain=1e-9)\n",
     "# sr.design.orbit.correct()\n",
@@ -852,7 +853,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Remember: if you change the example from live to design you need to change the unit for the orbit when plotting since not the same yet.\n",
+    "# Remember: if you change the example from live to design you need to change the unit\n",
+    "# for the orbit when plotting since not the same yet.\n",
     "\n",
     "fig = plt.figure()\n",
     "ax1 = fig.add_subplot(311)\n",
@@ -890,7 +892,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "bessy2-orbit",
+   "display_name": "pyaml-tango",
    "language": "python",
    "name": "python3"
   },
@@ -904,7 +906,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.14"
+   "version": "3.14.2"
   }
  },
  "nbformat": 4,

examples/BESSY2_example/bessy2-orbit.py

@@ -1,7 +1,8 @@
 """Orbit correction
 
-This example shows how to run orbit correction using the simulator or the virtual accelerator.
-If you want to test the virtual accelerator you need to start the container before running the script.
+This example shows how to run orbit correction using the simulator or the virtual
+accelerator. If you want to test the virtual accelerator you need to start the
+container before running the script.
 
 """
 
@@ -16,13 +17,16 @@
 from pyaml.tuning_tools.orbit_response_matrix import OrbitResponseMatrix
 
 # ----- Load the configuration -----
-# Remember to change the prefix for the live mode to the one matching your virtual accelerator before loading.
+# Remember to change the prefix for the live mode to the one matching your virtual
+# accelerator before loading.
 
 sr = Accelerator.load("BESSY2Orbit.yaml")
 
 # ----- Measure the orbit response matrix -----
-# If there is no existing orbit response matrix you need to measure it. This can be done on either the design or live mode.
-# It is also possible to measure the ORM using the design mode and use it to correct the live mode.
+# If there is no existing orbit response matrix you need to measure it.
+# This can be done on either the design or live mode.
+# It is also possible to measure the ORM using the design mode and use it to correct
+# the live mode.
 
 # Choose which backend to use.
 SR = sr.design
@@ -39,13 +43,14 @@
         "matrix": orm_data["matrix"],
         "input_names": orm_data["input_names"],
         "output_names": orm_data["output_names"],
-        "inputs_plane": orm_data["inputs_plane"],
-        "outputs_plane": orm_data["outputs_plane"],
+        "input_planes": orm_data["input_planes"],
+        "output_planes": orm_data["output_planes"],
     }
     json.dump(ORM_data, open("orm.json", "w"))
 
 # ----- Load the response matrix -----
-# The example does the correction for the live mode but it can also be done on the design mode.
+# The example does the correction for the live mode but it can also be done on the
+# design mode.
 
 # Load the ORM for the live mode
 sr.live.orbit.load_response_matrix("orm.json")
@@ -66,8 +71,8 @@
 orbit_initial = orbit.get()
 
 # Correct the orbit
-# If you are using the ORM measured on design to correct on live you need to set the gain
-# since the unit for the BPMs are not the same for both modes yet.
+# If you are using the ORM measured on design to correct on live you need to set the
+# gain since the unit for the BPMs are not the same for both modes yet.
 
 sr.live.orbit.correct(gain=1e-9)
 # sr.design.orbit.correct()
@@ -77,8 +82,10 @@
 orbit_after = orbit.get()
 
 # ----- Plot the results -----
-# Remember: if you change the example from live to design you need to change the unit for the orbit when plotting since not the same yet.
-# If you are running in VS code you might need to switch the matplotlib backend for the plot to show.
+# Remember: if you change the example from live to design you need to change the unit
+# for the orbit when plotting since not the same yet.
+# If you are running in VS code you might need to switch the matplotlib backend for
+# the plot to show.
 
 fig = plt.figure()
 ax1 = fig.add_subplot(311)

examples/ESRF_ORM_example/correct_orbit.py

@@ -27,6 +27,7 @@
 vcorr = ebs.get_magnets("VCorr")
 bpms = ebs.get_bpms("BPM")
 
+np.random.seed(1)
 # mangle orbit
 hcorr.strengths.set(
     hcorr.strengths.get() + std_kick * np.random.normal(size=len(hcorr))
@@ -35,6 +36,9 @@
     vcorr.strengths.get() + std_kick * np.random.normal(size=len(vcorr))
 )
 
+h0 = hcorr.strengths.get()
+v0 = vcorr.strengths.get()
+
 positions_bc = bpms.positions.get()
 std_bc = np.std(positions_bc, axis=0)
 print(
@@ -43,6 +47,7 @@
 )
 ########################################################
 
+ebs.orbit.set_virtual_weight(1000)
 ## Correct the orbit
 ebs.orbit.correct(reference=reference)
 # ebs.orbit.correct(plane="H")
@@ -76,4 +81,6 @@
 ax3.set_xlabel("Steerer number")
 fig.tight_layout()
 
+print(f"Sum of h. corr. trims: {np.sum(hcorr.strengths.get() - h0)}")
+print(f"Sum of v. corr. trims: {np.sum(vcorr.strengths.get() - v0)}")
 plt.show()

examples/ESRF_ORM_example/correct_orbit_disp.py

@@ -0,0 +1,85 @@
+import json
+import logging
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+import numpy as np
+
+from pyaml.accelerator import Accelerator
+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
+
+## get reference
+ref_h, ref_v = ebs.get_bpms("BPM").positions.get().T
+reference = np.concatenate((ref_h, ref_v))
+########################################################
+
+plant = sr.design.get_rf_plant("RF")
+frf = plant.frequency.get()
+plant.frequency.set(frf + 100)
+
+
+## generate some orbit
+std_kick = 0e-6
+hcorr = ebs.get_magnets("HCorr")
+vcorr = ebs.get_magnets("VCorr")
+bpms = ebs.get_bpms("BPM")
+
+# mangle orbit
+hcorr.strengths.set(
+    hcorr.strengths.get() + std_kick * np.random.normal(size=len(hcorr))
+)
+vcorr.strengths.set(
+    vcorr.strengths.get() + std_kick * np.random.normal(size=len(vcorr))
+)
+
+positions_bc = bpms.positions.get()
+std_bc = np.std(positions_bc, axis=0)
+print(
+    "R.m.s. orbit before correction "
+    f"H: {1e6 * std_bc[0]: .1f} µm, V: {1e6 * std_bc[1]: .1f} µm."
+)
+########################################################
+
+## Correct the orbit
+# ebs.orbit.set_rf_weight(1e7)
+ebs.orbit.correct(reference=reference, rf=True, gain_H=1, gain_V=1, gain_RF=1)
+ebs.orbit.correct(reference=reference, rf=True, gain_H=1, gain_V=1, gain_RF=1)
+# ebs.orbit.correct(plane="H")
+# ebs.orbit.correct(plane="V")
+########################################################
+
+## inspect orbit correction
+positions_ac = bpms.positions.get()
+std_ac = np.std(positions_ac, axis=0)
+print(
+    "R.m.s. orbit after correction H: "
+    f"{1e6 * std_ac[0]: .1f} µm, V: {1e6 * std_ac[1]: .1f} µm,"
+)
+
+fig = plt.figure()
+ax1 = fig.add_subplot(311)
+ax2 = fig.add_subplot(312)
+ax3 = fig.add_subplot(313)
+ax1.plot(positions_bc[:, 0] * 1e6, label="Orbit before correction")
+ax2.plot(positions_bc[:, 1] * 1e6, label="Orbit before correction")
+ax1.plot(positions_ac[:, 0] * 1e6, label="Orbit after correction")
+ax2.plot(positions_ac[:, 1] * 1e6, label="Orbit after correction")
+
+ax3.plot(hcorr.strengths.get())
+ax3.plot(vcorr.strengths.get())
+
+ax1.set_ylabel("Horizontal pos. [μm]")
+ax2.set_ylabel("Vertical pos. [μm]")
+ax2.set_xlabel("BPM number")
+ax3.set_ylabel("Strength (rad)")
+ax3.set_xlabel("Steerer number")
+fig.tight_layout()
+
+plt.show()

examples/ESRF_ORM_example/measure_ideal_ORM_and_disp.py

@@ -4,10 +4,6 @@
 import numpy as np
 
 from pyaml.accelerator import Accelerator
-from pyaml.tuning_tools.dispersion import ConfigModel as Disp_ConfigModel
-from pyaml.tuning_tools.dispersion import Dispersion
-from pyaml.tuning_tools.orbit_response_matrix import ConfigModel as ORM_ConfigModel
-from pyaml.tuning_tools.orbit_response_matrix import OrbitResponseMatrix
 
 parent_folder = Path(__file__).parent
 config_path = parent_folder.parent.parent.joinpath(
@@ -32,8 +28,8 @@
     "matrix": orm_data["matrix"],
     "input_names": orm_data["input_names"],
     "output_names": orm_data["output_names"],
-    "inputs_plane": orm_data["inputs_plane"],
-    "outputs_plane": orm_data["outputs_plane"],
+    "input_planes": orm_data["input_planes"],
+    "output_planes": orm_data["output_planes"],
     "rf_response": rf_response,
 }
 

examples/ESRF_ORM_example/scan_bpm_weight.py

@@ -0,0 +1,61 @@
+import json
+import logging
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+import numpy as np
+
+from pyaml.accelerator import Accelerator
+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
+
+## get reference
+ref_h, ref_v = ebs.get_bpms("BPM").positions.get().T
+reference = np.concatenate((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")
+
+np.random.seed(1)
+# mangle orbit
+h0 = hcorr.strengths.get() + std_kick * np.random.normal(size=len(hcorr))
+v0 = vcorr.strengths.get() + std_kick * np.random.normal(size=len(vcorr))
+
+my_bpm_name = "BPM_C04-05"
+my_bpm = ebs.get_bpm(my_bpm_name)
+
+weights = np.linspace(0.01, 10, 40)
+bpm_reading = np.zeros_like(weights)
+ref_bpm_reading = my_bpm.positions.get()[0]
+########################################################
+for i, w in enumerate(weights):
+    hcorr.strengths.set(h0)
+    vcorr.strengths.set(v0)
+    ebs.orbit.set_weight(my_bpm_name, w, plane="H")
+
+    ebs.orbit.correct(singular_values_H=40, singular_values_V=40)
+    ebs.orbit.correct(singular_values_H=40, singular_values_V=40)
+    ebs.orbit.correct(singular_values_H=40, singular_values_V=40)
+    bpm_reading[i] = my_bpm.positions.get()[0]
+
+fig = plt.figure()
+ax1 = fig.add_subplot(111)
+ax1.plot(weights, bpm_reading * 1e6, ".")
+
+
+ax1.set_ylabel("BPM reading [μm]")
+ax1.set_xlabel("BPM weight")
+fig.tight_layout()
+
+plt.show()