diff --git a/pyaml/lattice/abstract_impl.py b/pyaml/lattice/abstract_impl.py index 91c07975..615d7ea0 100644 --- a/pyaml/lattice/abstract_impl.py +++ b/pyaml/lattice/abstract_impl.py @@ -12,6 +12,7 @@ from ..common import abstract from ..common.abstract_aggregator import ScalarAggregator +from ..common.exception import PyAMLException from ..magnet.model import MagnetModel from .polynom_info import PolynomInfo @@ -20,6 +21,21 @@ # ------------------------------------------------------------------------------ +def _effective_lengths(elements: list[at.Element]) -> list[float]: + """ + Return the lengths used to convert between integrated strengths and polynom coefficients. + + AT convention is followed: a zero-length (thin) element stores the *integrated* strength in + ``PolynomA``/``PolynomB``, so it behaves as a unit-length element for the conversion. When every element of the + group is thin, each one is given a unit length, so that the integrated strength is shared equally between the + slices. Otherwise the real lengths are used and any thin slice of the group is left untouched. + """ + lengths = [float(e.Length) for e in elements] + if all(L == 0.0 for L in lengths): + return [1.0] * len(lengths) + return lengths + + class RWHardwareScalar(abstract.ReadWriteFloatScalar): """ Provide read/write access to a simulated magnet in hardware units. @@ -62,19 +78,23 @@ def __init__(self, elements: list[at.Element], poly: PolynomInfo, model: MagnetM self._poly = [e.__getattribute__(poly.attName) for e in elements] self._sign = poly.sign self._polyIdx = poly.index - self._length: float = 0.0 - for e in elements: - self._length += e.Length + self._is_thin = all(e.Length == 0 for e in elements) + self._lengths = _effective_lengths(elements) + self._length: float = sum(self._lengths) def get_length(self) -> float: - """Return the total length of the lattice elements.""" + """Return the total effective length of the lattice elements (1.0 per element when they are all thin).""" return self._length + def is_thin(self) -> bool: + """Return True when every lattice element of this magnet has zero length.""" + return self._is_thin + def get(self) -> float: """Return the current value.""" s = 0 - for idx, e in enumerate(self._elements): - s += self._poly[idx][self._polyIdx] * self._sign * e.Length + for idx, _ in enumerate(self._elements): + s += self._poly[idx][self._polyIdx] * self._sign * self._lengths[idx] return self._model.compute_hardware_values([s])[0] def set(self, value: float): @@ -160,9 +180,8 @@ def __init__(self, elements: list[at.Element], poly: PolynomInfo, model: MagnetM self._poly = [e.__getattribute__(poly.attName) for e in elements] self._sign = poly.sign self._polyIdx = poly.index - self._length = 0 - for e in elements: - self._length += e.Length + self._lengths = _effective_lengths(elements) + self._length: float = sum(self._lengths) def get_element_length(self) -> float: """Return the total length of the represented element.""" @@ -195,8 +214,8 @@ def get(self, polynom: str = None, polyidx: int = None) -> float: poly = [e.__getattribute__(polynom) for e in self._elements] s = 0 - for idx, e in enumerate(self._elements): - s += poly[idx][pIdx] * self._sign * e.Length + for idx, _ in enumerate(self._elements): + s += poly[idx][pIdx] * self._sign * self._lengths[idx] return s # Sets the value @@ -398,6 +417,14 @@ def __init__( self.__elements_strength = elements_strength self.__elements_hardware = elements_hardware self.__element_index = element_index + # The strength is shared between the serialized magnets in proportion to their (effective) length: thick + # magnets share by length, thin magnets share equally. Mixing both has no meaningful share rule. + thin_flags = [e.is_thin() for e in elements_hardware] + if any(thin_flags) and not all(thin_flags): + raise PyAMLException( + "Serialized magnets must be either all thin (zero length) or all thick; " + f"got {sum(thin_flags)} thin out of {len(thin_flags)} magnets" + ) self.__total_length = 0 for e in self.__elements_hardware: self.__total_length += e.get_length() @@ -509,6 +536,7 @@ def __init__(self, elements: list[at.Element], poly: list[PolynomInfo], model: M self.__polyIdx = [] self.__sign = [] self.__model = model + self.__length = _effective_lengths(elements[:1])[0] for p in poly: self.__poly.append(elements[0].__getattribute__(p.attName)) self.__polyIdx.append(p.index) @@ -520,7 +548,7 @@ def get(self) -> np.array: nbStrength = len(self.__poly) s = np.zeros(nbStrength) for i in range(nbStrength): - s[i] = self.__poly[i][self.__polyIdx[i]] * self.__sign[i] * self.__elements[0].Length + s[i] = self.__poly[i][self.__polyIdx[i]] * self.__sign[i] * self.__length return self.__model.compute_hardware_values(s) # Sets the value @@ -536,7 +564,7 @@ def set(self, value: np.array): nbStrength = len(self.__poly) s = self.__model.compute_strengths(value) for i in range(nbStrength): - self.__poly[i][self.__polyIdx[i]] = s[i] / (self.__elements[0].Length * self.__sign[i]) + self.__poly[i][self.__polyIdx[i]] = s[i] / (self.__length * self.__sign[i]) # Sets the value and wait that the read value reach the setpoint def set_and_wait(self, value: np.array): @@ -599,6 +627,7 @@ def __init__(self, elements: list[at.Element], poly: list[PolynomInfo], model: M self.__polyIdx = [] self.__sign = [] self.__model = model + self.__length = _effective_lengths(elements[:1])[0] for p in poly: self.__poly.append(elements[0].__getattribute__(p.attName)) self.__polyIdx.append(p.index) @@ -610,7 +639,7 @@ def get(self) -> np.array: nbStrength = len(self.__poly) s = np.zeros(nbStrength) for i in range(nbStrength): - s[i] = self.__poly[i][self.__polyIdx[i]] * self.__sign[i] * self.__elements[0].Length + s[i] = self.__poly[i][self.__polyIdx[i]] * self.__sign[i] * self.__length return s # Sets the value @@ -626,7 +655,7 @@ def set(self, value: np.array): nbStrength = len(self.__poly) s = np.zeros(nbStrength) for i in range(nbStrength): - self.__poly[i][self.__polyIdx[i]] = value[i] / (self.__elements[0].Length * self.__sign[i]) + self.__poly[i][self.__polyIdx[i]] = value[i] / (self.__length * self.__sign[i]) # Sets the value and wait that the read value reach the setpoint def set_and_wait(self, value: np.array): diff --git a/tests/lattice/test_magnet_accessors.py b/tests/lattice/test_magnet_accessors.py new file mode 100644 index 00000000..e585b37e --- /dev/null +++ b/tests/lattice/test_magnet_accessors.py @@ -0,0 +1,100 @@ +import at +import numpy as np +import pytest + +from pyaml import PyAMLException +from pyaml.lattice.abstract_impl import ( + RWHardwareArray, + RWHardwareScalar, + RWSerializedStrength, + RWStrengthArray, + RWStrengthScalar, +) +from pyaml.magnet.hcorrector import HCorrector +from pyaml.magnet.identity_cfm_model import IdentityCFMagnetModel +from pyaml.magnet.identity_model import IdentityMagnetModel +from pyaml.magnet.vcorrector import VCorrector + + +def _thin_corrector(name="COR"): + return at.Corrector(name, 0.0, [0.0, 0.0], PolynomA=[0.0], PolynomB=[0.0]) + + +@pytest.mark.parametrize("accessor_type", [RWStrengthScalar, RWHardwareScalar]) +def test_zero_length_scalar_magnet_stores_integrated_strength(accessor_type): + element = _thin_corrector() + model = IdentityMagnetModel(physics="COR", unit="rad") + accessor = accessor_type([element], HCorrector.polynom, model) + + accessor.set(1.0e-6) + + # AT convention: a thin element stores the integrated strength in its polynom + assert element.PolynomB[0] == pytest.approx(-1.0e-6) + assert accessor.get() == pytest.approx(1.0e-6) + + +@pytest.mark.parametrize("accessor_type", [RWStrengthScalar, RWHardwareScalar]) +def test_zero_length_split_magnet_shares_integrated_strength_equally(accessor_type): + elements = [_thin_corrector("COR_A"), _thin_corrector("COR_B")] + model = IdentityMagnetModel(physics="COR", unit="rad") + accessor = accessor_type(elements, HCorrector.polynom, model) + + accessor.set(1.0e-6) + + assert elements[0].PolynomB[0] == pytest.approx(-0.5e-6) + assert elements[1].PolynomB[0] == pytest.approx(-0.5e-6) + assert accessor.get() == pytest.approx(1.0e-6) + + +@pytest.mark.parametrize("accessor_type", [RWStrengthScalar, RWHardwareScalar]) +def test_thick_scalar_magnet_is_unchanged(accessor_type): + element = at.Corrector("COR", 0.2, [0.0, 0.0], PolynomA=[0.0], PolynomB=[0.0]) + model = IdentityMagnetModel(physics="COR", unit="rad") + accessor = accessor_type([element], HCorrector.polynom, model) + + accessor.set(1.0e-6) + + assert element.PolynomB[0] == pytest.approx(-5.0e-6) + assert accessor.get() == pytest.approx(1.0e-6) + + +@pytest.mark.parametrize("accessor_type", [RWStrengthArray, RWHardwareArray]) +def test_zero_length_combined_function_magnet_stores_integrated_strengths(accessor_type): + element = _thin_corrector() + model = IdentityCFMagnetModel( + multipoles=["B0", "A0"], + physics=["HCOR", "VCOR"], + units=["rad", "rad"], + ) + accessor = accessor_type([element], [HCorrector.polynom, VCorrector.polynom], model) + + accessor.set(np.array([1.0e-6, -2.0e-6])) + + assert element.PolynomB[0] == pytest.approx(-1.0e-6) + assert element.PolynomA[0] == pytest.approx(-2.0e-6) + np.testing.assert_allclose(accessor.get(), [1.0e-6, -2.0e-6]) + + +def test_zero_length_serialized_magnets_share_strength_equally(): + elements = [_thin_corrector("COR_A"), _thin_corrector("COR_B")] + model = IdentityMagnetModel(physics="COR", unit="rad") + strengths = [RWStrengthScalar([e], HCorrector.polynom, model) for e in elements] + currents = [RWHardwareScalar([e], HCorrector.polynom, model) for e in elements] + serialized = [RWSerializedStrength(strengths, currents, i) for i in range(len(elements))] + + serialized[0].set(1.0e-6) + + assert serialized[0].get() == pytest.approx(0.5e-6) + assert serialized[1].get() == pytest.approx(0.5e-6) + assert elements[0].PolynomB[0] == pytest.approx(-0.5e-6) + assert elements[1].PolynomB[0] == pytest.approx(-0.5e-6) + + +def test_serialized_magnets_mixing_thin_and_thick_is_rejected(): + elements = [_thin_corrector("COR_A"), at.Corrector("COR_B", 0.2, [0.0, 0.0], PolynomA=[0.0], PolynomB=[0.0])] + model = IdentityMagnetModel(physics="COR", unit="rad") + strengths = [RWStrengthScalar([e], HCorrector.polynom, model) for e in elements] + currents = [RWHardwareScalar([e], HCorrector.polynom, model) for e in elements] + + with pytest.raises(PyAMLException, match="all thin .* or all thick"): + RWSerializedStrength(strengths, currents, 0)