Allow usage of zero-length element

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):

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)