Address the dipole kick through KickAngle where the pass method uses it

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)
 

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