From e0433936bc31e0c60b8ac9384665319913aca93b Mon Sep 17 00:00:00 2001 From: ThorstenHellert Date: Sat, 10 Oct 2026 09:14:44 -0700 Subject: [PATCH 1/2] Write corrector strength to KickAngle for CorrectorPass elements CorrectorPass integrates KickAngle and ignores PolynomA/PolynomB, so a strength set on an at.Corrector element round-tripped through the accessor but never moved the orbit. The simulator accessors now address the dipole component of such an element through KickAngle, keeping a polynom the element carries in step with it. Correctors defined without polynoms are supported, and a kick already present in the lattice is read back. Tests compare the closed orbit against a KickAngle reference. --- pyaml/lattice/abstract_impl.py | 67 ++++++++++++++++++-- tests/lattice/test_magnet_accessors.py | 85 ++++++++++++++++++++++++++ 2 files changed, 146 insertions(+), 6 deletions(-) diff --git a/pyaml/lattice/abstract_impl.py b/pyaml/lattice/abstract_impl.py index a734a587..764509d0 100644 --- a/pyaml/lattice/abstract_impl.py +++ b/pyaml/lattice/abstract_impl.py @@ -36,6 +36,61 @@ def _effective_lengths(elements: list[at.Element]) -> list[float]: return lengths +class _KickAngleCoefficients: + """ + Present the ``KickAngle`` of a ``CorrectorPass`` element as its dipole polynom coefficient. + + ``CorrectorPass`` integrates ``KickAngle`` only and ignores ``PolynomA``/``PolynomB``, so the dipole component of + such an element lives in ``KickAngle``. This view exposes it 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. A polynom carried by the element is kept in step with the kick; coefficients above the + dipole are passed through to it. + + Parameters + ---------- + element : at.Element + Element tracked with ``CorrectorPass``. + 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._plane, self._kick_sign = (0, -1.0) if att_name == "PolynomB" else (1, 1.0) + + def __getitem__(self, index: int) -> float: + if index != 0: + return getattr(self._element, self._att_name)[index] + if self._length == 0.0: + return 0.0 + return self._kick_sign * self._element.KickAngle[self._plane] / self._length + + def __setitem__(self, index: int, value: float): + polynom = getattr(self._element, self._att_name, None) + if index != 0: + polynom[index] = value + return + self._element.KickAngle[self._plane] = self._kick_sign * value * self._length + if polynom is not None: + polynom[0] = value + + +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 tracked with ``CorrectorPass``, whose dipole kick is held by + ``KickAngle`` (see :class:`_KickAngleCoefficients`). + """ + if element.PassMethod == "CorrectorPass": + 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 +130,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 +232,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 +266,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 +294,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 +593,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 +684,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..b8548f33 100644 --- a/tests/lattice/test_magnet_accessors.py +++ b/tests/lattice/test_magnet_accessors.py @@ -98,3 +98,88 @@ 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) From 718e93c9d859e1d41561e9a3fca6ff401455f726 Mon Sep 17 00:00:00 2001 From: ThorstenHellert Date: Sat, 10 Oct 2026 11:36:30 -0700 Subject: [PATCH 2/2] Count the KickAngle of multipole elements in the dipole strength The multipole pass methods add KickAngle to the polynom. An element that carries one read back only its polynom part, and a write added to the kick instead of replacing it. The accessors now read the sum and, on a write, store the value in the polynom and reset the KickAngle component. A KickAngle on a pass method that ignores it is not read. --- pyaml/lattice/abstract_impl.py | 56 ++++++++++++++++++-------- tests/lattice/test_magnet_accessors.py | 34 ++++++++++++++++ 2 files changed, 74 insertions(+), 16 deletions(-) diff --git a/pyaml/lattice/abstract_impl.py b/pyaml/lattice/abstract_impl.py index 764509d0..a7764ce9 100644 --- a/pyaml/lattice/abstract_impl.py +++ b/pyaml/lattice/abstract_impl.py @@ -38,18 +38,22 @@ def _effective_lengths(elements: list[at.Element]) -> list[float]: class _KickAngleCoefficients: """ - Present the ``KickAngle`` of a ``CorrectorPass`` element as its dipole polynom coefficient. + Present the dipole strength of an element that carries a ``KickAngle`` as a single polynom coefficient. - ``CorrectorPass`` integrates ``KickAngle`` only and ignores ``PolynomA``/``PolynomB``, so the dipole component of - such an element lives in ``KickAngle``. This view exposes it 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. A polynom carried by the element is kept in step with the kick; coefficients above the - dipole are passed through to it. + 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 tracked with ``CorrectorPass``. + Element carrying a ``KickAngle`` that its pass method applies. att_name : str Polynom the view stands for, ``"PolynomA"`` or ``"PolynomB"``. length : float @@ -60,33 +64,53 @@ 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 __getitem__(self, index: int) -> float: - if index != 0: - return getattr(self._element, self._att_name)[index] + 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 - return - self._element.KickAngle[self._plane] = self._kick_sign * value * self._length - if polynom is not None: + 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 tracked with ``CorrectorPass``, whose dipole kick is held by - ``KickAngle`` (see :class:`_KickAngleCoefficients`). + 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": + if element.PassMethod == "CorrectorPass" or (hasattr(element, "KickAngle") and _applies_kick_angle(element)): return _KickAngleCoefficients(element, att_name, length) return getattr(element, att_name) diff --git a/tests/lattice/test_magnet_accessors.py b/tests/lattice/test_magnet_accessors.py index b8548f33..ca3b52d1 100644 --- a/tests/lattice/test_magnet_accessors.py +++ b/tests/lattice/test_magnet_accessors.py @@ -183,3 +183,37 @@ def test_multipole_corrector_agrees_with_corrector_pass(): 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