Implemented BPM offset/tilt

pyaml/bpm/bpm.py

@@ -54,7 +54,6 @@ def __init__(self, cfg: ConfigModel):
 
         super().__init__(cfg.name)
 
-        self.__model = cfg.model if hasattr(cfg, "model") else None
         self._cfg = cfg
         self._positions = None
         self._offset = None
@@ -146,7 +145,6 @@ def attach(
         # Attach positions, offset and tilt attributes and returns a new
         # reference
         obj = self.__class__(self._cfg)
-        obj.__model = self.__model
         obj._positions = positions
         obj._offset = offset
         obj._tilt = tilt

pyaml/lattice/abstract_impl.py

@@ -285,11 +285,13 @@ class BPMScalarAggregator(ScalarAggregator):
     """
 
     def __init__(self, ring: at.Lattice):
-        self.lattice = ring
-        self.refpts = []
+        self._lattice = ring
+        self._refpts = []
+        self._matrices = []
 
     def add_elem(self, elem: at.Element):
-        self.refpts.append(self.lattice.index(elem))
+        self._refpts.append(self._lattice.index(elem))
+        self._matrices.append(elem._transform)
 
     def set(self, value: NDArray[np.float64]):
         pass
@@ -298,45 +300,60 @@ def set_and_wait(self, value: NDArray[np.float64]):
         pass
 
     def get(self) -> np.array:
-        _, orbit = at.find_orbit(self.lattice, refpts=self.refpts)
-        return orbit[:, [0, 2]].flatten()
+        return self._transform().flatten()
 
     def readback(self) -> np.array:
         return self.get()
 
     def unit(self) -> str:
         return "m"
 
+    def _transform(self) -> np.array:
+        _, orbit = at.find_orbit(self._lattice, refpts=self._refpts)
+        ones = np.ones(len(self._refpts))
+        pts = orbit[:, [0, 2]]  # Extract x,y
+        # Batch matrices multiplication (homogeneous coordinates)
+        return np.matmul(self._matrices, np.column_stack([pts, ones])[:, :, None]).squeeze(-1)
+
 
 # ------------------------------------------------------------------------------
 
 
 class BPMHScalarAggregator(BPMScalarAggregator):
     """
-    Vertical BPM simulator aggregator
+    Horizontal BPM simulator aggregator
     """
 
     def get(self) -> np.array:
-        _, orbit = at.find_orbit(self.lattice, refpts=self.refpts)
-        return orbit[:, 0]
+        return self._transform()[:, 0]
 
 
 # ------------------------------------------------------------------------------
 
 
 class BPMVScalarAggregator(BPMScalarAggregator):
     """
-    Horizontal BPM simulator aggregator
+    Vertical BPM simulator aggregator
     """
 
     def get(self) -> np.array:
-        _, orbit = at.find_orbit(self.lattice, refpts=self.refpts)
-        return orbit[:, 2]
+        return self._transform()[:, 1]
 
 
 # ------------------------------------------------------------------------------
 
 
+def update_bpm_transform_matrix(element: at.Element):
+    # BPM transformation matrix (homogeneous coordinates)
+    tx = element.Offset[0]
+    ty = element.Offset[1]
+    cos_theta = np.cos(element.Tilt)
+    sin_theta = np.sin(element.Tilt)
+    if not hasattr(element, "_transform"):
+        element._transform = np.empty((2, 3))
+    element._transform[:, :] = [[cos_theta, -sin_theta, tx], [sin_theta, cos_theta, ty]]
+
+
 class RBpmArray(abstract.ReadFloatArray):
     """
     Class providing read access to a BPM position (array) of a simulator.
@@ -346,14 +363,15 @@ class RBpmArray(abstract.ReadFloatArray):
     """
 
     def __init__(self, element: at.Element, lattice: at.Lattice):
-        self.__element = element
-        self.__lattice = lattice
+        self._element = element
+        self._lattice = lattice
 
     # Gets the value
     def get(self) -> np.array:
-        index = self.__lattice.index(self.__element)
-        _, orbit = at.find_orbit(self.__lattice, refpts=index)
-        return orbit[0, [0, 2]]
+        index = self._lattice.index(self._element)
+        _, orbit = at.find_orbit(self._lattice, refpts=index)
+        pts = orbit[0, [0, 2]]
+        return np.dot(self._element._transform, np.hstack([pts, 1]))  # Use homogeneous coordinates
 
     # Gets the unit of the value
     def unit(self) -> str:
@@ -370,25 +388,16 @@ class RWBpmOffsetArray(abstract.ReadWriteFloatArray):
     """
 
     def __init__(self, element: at.Element):
-        self.__element = element
-        try:
-            self.__offset = element.__getattribute__("Offset")
-        except AttributeError:
-            self.__offset = None
+        self._element = element
 
     # Gets the value
     def get(self) -> np.array:
-        if self.__offset is None:
-            raise PyAMLException("Element does not have an Offset attribute.")
-        return self.__offset
+        return self._element.Offset
 
     # Sets the value
     def set(self, value: np.array):
-        if self.__offset is None:
-            raise PyAMLException("Element does not have an Offset attribute.")
-        if len(value) != 2:
-            raise PyAMLException("BPM offset must be a 2-element array.")
-        self.__offset = value
+        self._element.Offset = value
+        update_bpm_transform_matrix(self._element)
 
     # Sets the value and wait that the read value reach the setpoint
     def set_and_wait(self, value: np.array):
@@ -409,25 +418,19 @@ class RWBpmTiltScalar(abstract.ReadWriteFloatScalar):
     """
 
     def __init__(self, element: at.Element):
-        self.__element = element
-        try:
-            self.__tilt = element.__getattribute__("Rotation")[0]
-        except AttributeError:
-            self.__tilt = None
+        self._element = element
 
     # Gets the value
     def get(self) -> float:
-        if self.__tilt is None:
-            raise ValueError("Element does not have a Tilt attribute.")
-        return self.__tilt
+        return self._element.Tilt
 
     # Sets the value
     def set(
         self,
         value: float,
     ):
-        self.__tilt = value
-        self.__element.__setattr__("Rotation", [value, None, None])
+        self._element.Tilt = value
+        update_bpm_transform_matrix(self._element)
 
     # Sets the value and wait that the read value reach the setpoint
     def set_and_wait(self, value: float):

pyaml/lattice/simulator.py

@@ -29,6 +29,7 @@
     RWSerializedStrength,
     RWStrengthArray,
     RWStrengthScalar,
+    update_bpm_transform_matrix,
 )
 from ..magnet.cfm_magnet import CombinedFunctionMagnet
 from ..magnet.magnet import Magnet
@@ -193,9 +194,20 @@ def fill_device(self, elements: list[Element]):
 
             elif isinstance(e, BPM):
                 # This assumes unique BPM names in the pyAT lattice
-                tilt = RWBpmTiltScalar(self.get_at_elems(e)[0])
-                offsets = RWBpmOffsetArray(self.get_at_elems(e)[0])
-                positions = RBpmArray(self.get_at_elems(e)[0], self.ring)
+
+                # Add Tilt and Offset fields if not present
+                bpm_elt = self.get_at_elems(e)[0]
+                if not hasattr(bpm_elt, "Tilt"):
+                    bpm_elt.Tilt = 0.0  # No tilt
+                if not hasattr(bpm_elt, "Offset"):
+                    bpm_elt.Offset = [0.0, 0.0]  # No offset
+                if len(bpm_elt.Offset) != 2:
+                    raise PyAMLException(f"BPM {e.get_name()} offset must be a 2-element array.")
+                update_bpm_transform_matrix(bpm_elt)
+
+                tilt = RWBpmTiltScalar(bpm_elt)
+                offsets = RWBpmOffsetArray(bpm_elt)
+                positions = RBpmArray(bpm_elt, self.ring)
                 e = e.attach(self, positions, offsets, tilt)
                 self.add_bpm(e)
 

tests/arrays/test_arrays.py

@@ -143,6 +143,17 @@ def test_arrays(install_test_package):
     assert np.all(np.isclose(pos[:, 0], pos_h, rtol=1e-15, atol=1e-15))
     assert np.all(np.isclose(pos[:, 1], pos_v, rtol=1e-15, atol=1e-15))
 
+    # Test BPM transformation matrices
+    sr.design.get_bpm("BPM_C04-01").offset.set([0.1, 0.2])
+    sr.design.get_bpm("BPM_C04-02").offset.set([0.3, 0.4])
+    pos = sr.design.get_bpms("BPMS").positions.get()
+    assert np.abs(pos[0][0] - 7.22262850488348e-05 - 0.1) < 1e-10
+    assert np.abs(pos[0][1] - 3.4291613955705856e-05 - 0.2) < 1e-10
+    assert np.abs(pos[1][0] + 1.1696152238807462e-04 - 0.3) < 1e-10
+    assert np.abs(pos[1][1] - 7.4265634524358045e-06 - 0.4) < 1e-10
+    sr.design.get_bpm("BPM_C04-01").offset.set([0.0, 0.0])
+    sr.design.get_bpm("BPM_C04-02").offset.set([0.0, 0.0])
+
     # No aggregator
     bpms = []
     for b in sr.design.get_bpms("BPMS"):

tests/bpm/test_bpm.py

@@ -7,10 +7,24 @@
 def test_simulator_bpm_tilt():
     sr: Accelerator = Accelerator.load("tests/config/bpms.yaml", ignore_external=True)
     sr.design.get_lattice().disable_6d()
+    sr.design.get_magnet("SH1A-C01-H").strength.set(10e-6)  # Add orbit
+    sr.design.get_magnet("SH1A-C01-V").strength.set(10e-6)  # Add orbit
     bpm = sr.design.get_bpm("BPM_C01-01")
+    assert np.allclose(bpm.positions.get(), np.array([5.90809968e-05, 2.24832853e-05]))
     assert bpm.tilt.get() == 0
-    bpm.tilt.set(0.01)
-    assert bpm.tilt.get() == 0.01
+    alpha = np.pi / 3
+    bpm.tilt.set(alpha)
+    assert bpm.tilt.get() == alpha
+
+    new_x = 5.908792e-05 * np.cos(alpha) - 2.24832853e-05 * np.sin(alpha)
+    new_y = 5.908792e-05 * np.sin(alpha) + 2.24832853e-05 * np.cos(alpha)
+    assert np.allclose(bpm.positions.get(), np.array([new_x, new_y]))
+
+    alpha = np.pi / 2
+    bpm.tilt.set(alpha)
+    new_x = 5.908792e-05 * np.cos(alpha) - 2.24832853e-05 * np.sin(alpha)
+    new_y = 5.908792e-05 * np.sin(alpha) + 2.24832853e-05 * np.cos(alpha)
+    assert np.allclose(bpm.positions.get(), np.array([new_x, new_y]))
 
 
 def test_simulator_bpm_offset():
@@ -23,7 +37,7 @@ def test_simulator_bpm_offset():
     bpm.offset.set(np.array([0.1, 0.2]))
     assert bpm.offset.get()[0] == 0.1
     assert bpm.offset.get()[1] == 0.2
-    assert np.allclose(bpm.positions.get(), np.array([0.0, 0.0]))
+    assert np.allclose(bpm.positions.get(), np.array([0.1, 0.2]))
 
 
 @pytest.mark.parametrize(