diff --git a/pyaml/lattice/abstract_impl.py b/pyaml/lattice/abstract_impl.py index a734a587..a7764ce9 100644 --- a/pyaml/lattice/abstract_impl.py +++ b/pyaml/lattice/abstract_impl.py @@ -36,6 +36,85 @@ def _effective_lengths(elements: list[at.Element]) -> list[float]: return lengths +class _KickAngleCoefficients: + """ + Present the dipole strength of an element that carries a ``KickAngle`` as a single polynom coefficient. + + AT holds the dipole kick of an element in two places. ``CorrectorPass`` integrates ``KickAngle`` only and ignores + ``PolynomA``/``PolynomB``; the multipole pass methods integrate the polynom *and* add ``KickAngle`` to it. This + view exposes the total at index 0 with the polynom convention (``PolynomB[0] * L = -KickAngle[0]``, + ``PolynomA[0] * L = KickAngle[1]``), so that the accessors address every element the same way and read the kick + the beam sees. + + Writing index 0 stores the value where the pass method reads it: in ``KickAngle`` for ``CorrectorPass`` (a polynom + carried by the element is kept in step), in the polynom otherwise, with the ``KickAngle`` component reset so that + it is not applied twice. Coefficients above the dipole are passed through to the polynom. + + Parameters + ---------- + element : at.Element + Element carrying a ``KickAngle`` that its pass method applies. + att_name : str + Polynom the view stands for, ``"PolynomA"`` or ``"PolynomB"``. + length : float + Length used to convert between the kick and the coefficient (see :func:`_effective_lengths`). + """ + + def __init__(self, element: at.Element, att_name: str, length: float): + self._element = element + self._att_name = att_name + self._length = length + self._kick_only = element.PassMethod == "CorrectorPass" + self._plane, self._kick_sign = (0, -1.0) if att_name == "PolynomB" else (1, 1.0) + + def _kick_coefficient(self) -> float: + if self._length == 0.0: + return 0.0 + return self._kick_sign * self._element.KickAngle[self._plane] / self._length + + def __getitem__(self, index: int) -> float: + if index != 0: + return getattr(self._element, self._att_name)[index] + if self._kick_only: + return self._kick_coefficient() + return getattr(self._element, self._att_name)[0] + self._kick_coefficient() + + def __setitem__(self, index: int, value: float): + polynom = getattr(self._element, self._att_name, None) + if index != 0: + polynom[index] = value + elif self._kick_only: + self._element.KickAngle[self._plane] = self._kick_sign * value * self._length + if polynom is not None: + polynom[0] = value + else: + polynom[0] = value + self._element.KickAngle[self._plane] = 0.0 + + +def _applies_kick_angle(element: at.Element) -> bool: + """ + Return True when the pass method of ``element`` applies its ``KickAngle``. + + That is ``CorrectorPass`` and the multipole integrators (``StrMPole*``, ``BndMPole*``, ``ThinMPolePass``, + ``ExactMultipole*``). Other pass methods ignore the attribute. + """ + pass_method = element.PassMethod + return pass_method == "CorrectorPass" or "MPole" in pass_method or "Multipole" in pass_method + + +def _coefficients(element: at.Element, att_name: str, length: float): + """ + Return the indexable coefficients that hold the strength of ``element`` for the polynom ``att_name``. + + This is the polynom itself, except for an element whose pass method applies a ``KickAngle``: its dipole strength + is then addressed through :class:`_KickAngleCoefficients`. + """ + if element.PassMethod == "CorrectorPass" or (hasattr(element, "KickAngle") and _applies_kick_angle(element)): + return _KickAngleCoefficients(element, att_name, length) + return getattr(element, att_name) + + class RWHardwareScalar(abstract.ReadWriteFloatScalar): """ Provide read/write access to a simulated magnet in hardware units. @@ -75,11 +154,11 @@ def __init__(self, elements: list[at.Element], poly: PolynomInfo, model: MagnetM """ self._model = model self._elements = elements - self._poly = [e.__getattribute__(poly.attName) for e in elements] self._sign = poly.sign self._polyIdx = poly.index self._is_thin = all(e.Length == 0 for e in elements) self._lengths = _effective_lengths(elements) + self._poly = [_coefficients(e, poly.attName, L) for e, L in zip(elements, self._lengths, strict=True)] self._length: float = sum(self._lengths) def get_length(self) -> float: @@ -177,10 +256,10 @@ def __init__(self, elements: list[at.Element], poly: PolynomInfo, model: MagnetM """ self._model = model self._elements = elements - self._poly = [e.__getattribute__(poly.attName) for e in elements] self._sign = poly.sign self._polyIdx = poly.index self._lengths = _effective_lengths(elements) + self._poly = [_coefficients(e, poly.attName, L) for e, L in zip(elements, self._lengths, strict=True)] self._length: float = sum(self._lengths) def get_element_length(self) -> float: @@ -211,7 +290,7 @@ def get(self, polynom: str = None, polyidx: int = None) -> float: else: # Override strength access pIdx = polyidx - poly = [e.__getattribute__(polynom) for e in self._elements] + poly = [_coefficients(e, polynom, L) for e, L in zip(self._elements, self._lengths, strict=True)] s = 0 for idx, _ in enumerate(self._elements): @@ -239,7 +318,7 @@ def set(self, value: float, polynom: str = None, polyidx: int = None): else: # Override strength access pIdx = polyidx - poly = [e.__getattribute__(polynom) for e in self._elements] + poly = [_coefficients(e, polynom, L) for e, L in zip(self._elements, self._lengths, strict=True)] for idx, _ in enumerate(self._elements): poly[idx][pIdx] = value / (self._length * self._sign) @@ -538,7 +617,7 @@ def __init__(self, elements: list[at.Element], poly: list[PolynomInfo], model: M self.__model = model self.__length = _effective_lengths(elements[:1])[0] for p in poly: - self.__poly.append(elements[0].__getattribute__(p.attName)) + self.__poly.append(_coefficients(elements[0], p.attName, self.__length)) self.__polyIdx.append(p.index) self.__sign.append(p.sign) @@ -629,7 +708,7 @@ def __init__(self, elements: list[at.Element], poly: list[PolynomInfo], model: M self.__model = model self.__length = _effective_lengths(elements[:1])[0] for p in poly: - self.__poly.append(elements[0].__getattribute__(p.attName)) + self.__poly.append(_coefficients(elements[0], p.attName, self.__length)) self.__polyIdx.append(p.index) self.__sign.append(p.sign) diff --git a/tests/lattice/test_magnet_accessors.py b/tests/lattice/test_magnet_accessors.py index e585b37e..ca3b52d1 100644 --- a/tests/lattice/test_magnet_accessors.py +++ b/tests/lattice/test_magnet_accessors.py @@ -98,3 +98,122 @@ def test_serialized_magnets_mixing_thin_and_thick_is_rejected(): with pytest.raises(PyAMLException, match="all thin .* or all thick"): RWSerializedStrength(strengths, currents, 0) + + +def _ring(corrector): + qf = at.Quadrupole("QF", 0.5, 1.2) + qd = at.Quadrupole("QD", 0.5, -1.2) + cell = [qf, at.Drift("D1", 1.0), at.Monitor("BPM"), qd, at.Drift("D2", 1.0)] + return at.Lattice([corrector] + cell * 8, energy=2e9, periodicity=1) + + +def _orbit(corrector): + """Closed orbit (x, y) at the monitors of a small ring holding ``corrector``.""" + orbit = at.find_orbit4(_ring(corrector), refpts=at.Monitor)[1] + return orbit[:, [0, 2]] + + +@pytest.mark.parametrize("length", [0.0, 0.2]) +@pytest.mark.parametrize( + "polynom, kick_angle", + [(HCorrector.polynom, [1.0e-4, 0.0]), (VCorrector.polynom, [0.0, 1.0e-4])], +) +def test_corrector_pass_strength_moves_the_orbit(length, polynom, kick_angle): + # CorrectorPass integrates KickAngle and ignores the polynoms: the strength has to reach KickAngle + reference = _orbit(at.Corrector("REF", length, kick_angle)) + assert np.abs(reference).max() > 1.0e-5 + + element = at.Corrector("COR", length, [0.0, 0.0], PolynomA=[0.0], PolynomB=[0.0]) + accessor = RWStrengthScalar([element], polynom, IdentityMagnetModel(physics="COR", unit="rad")) + accessor.set(1.0e-4) + + assert accessor.get() == pytest.approx(1.0e-4) + np.testing.assert_allclose(element.KickAngle, kick_angle) + np.testing.assert_allclose(_orbit(element), reference, rtol=0, atol=1.0e-12) + + +def test_corrector_pass_without_polynoms_is_supported(): + element = at.Corrector("COR", 0.0, [0.0, 0.0]) + accessor = RWStrengthScalar([element], HCorrector.polynom, IdentityMagnetModel(physics="COR", unit="rad")) + + accessor.set(1.0e-4) + + assert accessor.get() == pytest.approx(1.0e-4) + np.testing.assert_allclose(element.KickAngle, [1.0e-4, 0.0]) + + +def test_corrector_pass_reads_the_kick_of_the_lattice(): + element = at.Corrector("COR", 0.2, [2.0e-5, -3.0e-5]) + model = IdentityMagnetModel(physics="COR", unit="rad") + + assert RWStrengthScalar([element], HCorrector.polynom, model).get() == pytest.approx(2.0e-5) + assert RWStrengthScalar([element], VCorrector.polynom, model).get() == pytest.approx(-3.0e-5) + + +@pytest.mark.parametrize("accessor_type", [RWStrengthArray, RWHardwareArray]) +def test_corrector_pass_combined_function_strengths_move_the_orbit(accessor_type): + reference = _orbit(at.Corrector("REF", 0.0, [1.0e-4, -2.0e-4])) + element = at.Corrector("COR", 0.0, [0.0, 0.0], PolynomA=[0.0], PolynomB=[0.0]) + 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-4, -2.0e-4])) + + np.testing.assert_allclose(accessor.get(), [1.0e-4, -2.0e-4]) + np.testing.assert_allclose(_orbit(element), reference, rtol=0, atol=1.0e-12) + + +def test_split_corrector_pass_shares_the_kick_by_length(): + elements = [at.Corrector("COR_A", 0.1, [0.0, 0.0]), at.Corrector("COR_B", 0.3, [0.0, 0.0])] + accessor = RWStrengthScalar(elements, HCorrector.polynom, IdentityMagnetModel(physics="COR", unit="rad")) + + accessor.set(1.0e-4) + + assert elements[0].KickAngle[0] == pytest.approx(0.25e-4) + assert elements[1].KickAngle[0] == pytest.approx(0.75e-4) + assert accessor.get() == pytest.approx(1.0e-4) + + +def test_multipole_corrector_agrees_with_corrector_pass(): + # Same strength, same orbit, whichever pass method the lattice uses for its correctors + model = IdentityMagnetModel(physics="COR", unit="rad") + thin = at.ThinMultipole("COR", [0.0], [0.0]) + kicker = at.Corrector("COR", 0.0, [0.0, 0.0]) + for element in (thin, kicker): + RWStrengthScalar([element], HCorrector.polynom, model).set(1.0e-4) + + np.testing.assert_allclose(_orbit(thin), _orbit(kicker), rtol=0, atol=1.0e-12) + + +@pytest.mark.parametrize( + "element", + [ + at.ThinMultipole("COR", [0.0], [0.0], KickAngle=[1.0e-4, 0.0]), + at.Multipole("COR", 0.2, [0.0], [0.0], KickAngle=[1.0e-4, 0.0]), + at.Dipole("COR", 0.2, 0.0, KickAngle=[1.0e-4, 0.0]), + ], + ids=["thin multipole", "multipole", "dipole"], +) +def test_multipole_kick_angle_is_part_of_the_strength(element): + # The multipole pass methods add KickAngle to the polynom: the strength is the sum, and a write replaces it + element = element.deepcopy() + kicked = _orbit(element) + assert np.abs(kicked).max() > 1.0e-5 + accessor = RWStrengthScalar([element], HCorrector.polynom, IdentityMagnetModel(physics="COR", unit="rad")) + + assert accessor.get() == pytest.approx(1.0e-4) + + accessor.set(1.0e-4) + assert accessor.get() == pytest.approx(1.0e-4) + # a thick integrator does not treat the two attributes identically to the last digit + np.testing.assert_allclose(_orbit(element), kicked, rtol=1.0e-6, atol=0) + + accessor.set(0.0) + np.testing.assert_allclose(_orbit(element), 0.0, atol=1.0e-12) + + +def test_kick_angle_ignored_by_the_pass_method_is_not_read(): + element = at.Quadrupole("QUAD", 0.2, 1.0, PassMethod="QuadLinearPass", KickAngle=[1.0e-4, 0.0]) + accessor = RWStrengthScalar([element], HCorrector.polynom, IdentityMagnetModel(physics="COR", unit="rad")) + + assert accessor.get() == 0.0