diff --git a/pyaml/bpm/bpm.py b/pyaml/bpm/bpm.py index a3da7175..22f3514a 100644 --- a/pyaml/bpm/bpm.py +++ b/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 diff --git a/pyaml/lattice/abstract_impl.py b/pyaml/lattice/abstract_impl.py index 69354e91..abe26360 100644 --- a/pyaml/lattice/abstract_impl.py +++ b/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,8 +300,7 @@ 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() @@ -307,18 +308,24 @@ def readback(self) -> np.array: 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] # ------------------------------------------------------------------------------ @@ -326,17 +333,27 @@ def get(self) -> np.array: 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): diff --git a/pyaml/lattice/simulator.py b/pyaml/lattice/simulator.py index 91a8e878..fd824db1 100644 --- a/pyaml/lattice/simulator.py +++ b/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) diff --git a/tests/arrays/test_arrays.py b/tests/arrays/test_arrays.py index eb72fdc8..001d6b14 100644 --- a/tests/arrays/test_arrays.py +++ b/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"): diff --git a/tests/bpm/test_bpm.py b/tests/bpm/test_bpm.py index 914d179f..de0c8577 100644 --- a/tests/bpm/test_bpm.py +++ b/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(