Support for transfer line

pyaml/lattice/abstract_impl.py

@@ -708,13 +708,14 @@ class BPMScalarAggregator(ScalarAggregator):
         Return the value unit.
     """
 
-    def __init__(self, ring: at.Lattice):
+    def __init__(self, ring: at.Lattice, twiss_in: dict):
         """
         Initialize the BPMScalarAggregator.
         """
         self._lattice = ring
         self._refpts = []
         self._matrices = []
+        self._twiss_in = twiss_in
 
     def add_elem(self, elem: at.Element):
         """
@@ -764,9 +765,15 @@ def unit(self) -> str:
 
     def _transform(self) -> np.array:
         """Transform the configured value for the lattice interface."""
-        _, orbit = at.find_orbit(self._lattice, refpts=self._refpts)
+        if self._twiss_in is None:
+            _, orbit = at.find_orbit(self._lattice, refpts=self._refpts)
+            pts = orbit[:, [0, 2]]  # Extract x,y
+        else:
+            _, _, elemdata = at.get_optics(self._lattice, refpts=self._refpts, twiss_in=self._twiss_in)
+            pts = [e.closed_orbit[[0, 2]] for e in elemdata]
+
         ones = np.ones(len(self._refpts))
-        pts = orbit[:, [0, 2]]  # Extract x,y
+        # Transfom BPM coordinates (offset and tilt)
         # Batch matrices multiplication (homogeneous coordinates)
         return np.matmul(self._matrices, np.column_stack([pts, ones])[:, :, None]).squeeze(-1)
 
@@ -856,19 +863,25 @@ class RBpmArray(abstract.ReadFloatArray):
         Return the value unit.
     """
 
-    def __init__(self, element: at.Element, lattice: at.Lattice):
+    def __init__(self, element: at.Element, lattice: at.Lattice, twiss_in: dict):
         """
         Initialize the RBpmArray.
         """
         self._element = element
         self._lattice = lattice
+        self._twiss_in = twiss_in
 
     # Gets the value
     def get(self) -> np.array:
         """Return the current value."""
         index = self._lattice.index(self._element)
-        _, orbit = at.find_orbit(self._lattice, refpts=index)
-        pts = orbit[0, [0, 2]]
+        if self._twiss_in is None:
+            _, orbit = at.find_orbit(self._lattice, refpts=index)
+            pts = orbit[0, [0, 2]]
+        else:
+            _, _, elemdata = at.get_optics(self._lattice, refpts=index, twiss_in=self._twiss_in)
+            pts = elemdata[0].closed_orbit[[0, 2]]
+        # Transfom BPM coordinates (offset and tilt)
         return np.dot(self._element._transform, np.hstack([pts, 1]))  # Use homogeneous coordinates
 
     # Gets the unit of the value

pyaml/lattice/simulator.py

@@ -48,6 +48,7 @@
 from ..tuning_tools.tuning_tool import TuningTool
 from ..validation import DynamicValidation, register_schema
 from .lattice_elements_linker import LatticeElementsLinker
+from .twiss_in import TwissIn
 
 if TYPE_CHECKING:
     from ..configuration.unbound_element import UnboundElement
@@ -79,6 +80,8 @@ class Simulator(ElementHolder, DynamicValidation):
     mat_key : str, optional
         Variable name of the lattice when loading a MATLAB ``.mat``
         lattice file.
+    twiss_in : TwissIn, optional
+        Describe input beam parameters of a tranfer line, when not specified the lattice is considered as a ring
     linker : LatticeElementsLinker, optional
         Custom linker used to associate PyAML elements with PyAT
         lattice elements. If omitted, elements are matched by name.
@@ -123,6 +126,7 @@ def __init__(
         name: str,
         lattice: str,
         mat_key: str | None = None,
+        twiss_in: TwissIn | None = None,
         linker: LatticeElementsLinker | None = None,
         description: str | None = None,
     ):
@@ -135,6 +139,11 @@ def __init__(
         self._lattice = lattice
         self._mat_key = mat_key
         self.description = description
+        self.twiss_in = twiss_in
+        if twiss_in is None:
+            self._at_twiss_in = None
+        else:
+            self._at_twiss_in = self.twiss_in._to_at()
 
         path: Path = ROOT.get() / self._lattice
 
@@ -226,9 +235,9 @@ def create_bpm_aggregators(self, bpms: list[BPM]) -> list[ScalarAggregator]:
         list[ScalarAggregator]
             Aggregators for combined, horizontal, and vertical BPM positions.
         """
-        agg = BPMScalarAggregator(self.get_lattice())
-        aggh = BPMHScalarAggregator(self.get_lattice())
-        aggv = BPMVScalarAggregator(self.get_lattice())
+        agg = BPMScalarAggregator(self.get_lattice(), self._at_twiss_in)
+        aggh = BPMHScalarAggregator(self.get_lattice(), self._at_twiss_in)
+        aggv = BPMVScalarAggregator(self.get_lattice(), self._at_twiss_in)
         for b in bpms:
             e = self.get_at_elems(b)[0]
             agg.add_elem(e)
@@ -293,7 +302,9 @@ def _fill_bpm(self, bpm: BPM) -> None:
             raise PyAMLException(f"BPM {bpm.get_name()} offset must be a 2-element array.")
         update_bpm_transform_matrix(bpm_elt)
         self.diagnostic.bpm.add(
-            bpm.attach(self, RBpmArray(bpm_elt, self.ring), RWBpmOffsetArray(bpm_elt), RWBpmTiltScalar(bpm_elt))
+            bpm.attach(
+                self, RBpmArray(bpm_elt, self.ring, self._at_twiss_in), RWBpmOffsetArray(bpm_elt), RWBpmTiltScalar(bpm_elt)
+            )
         )
 
     def _fill_rf_plant(self, rf_plant: RFPlant) -> None:

pyaml/lattice/twiss_in.py

@@ -0,0 +1,52 @@
+"""
+Initial input beam parameters for a transfer line
+"""
+
+import numpy
+
+from ..validation import DynamicValidation, register_schema
+
+# Define the main class name for this module
+PYAMLCLASS = "TwissIn"
+
+
+@register_schema
+class TwissIn(DynamicValidation):
+    """
+    Class that describes input beam parameters of a transfer line
+
+    Attributes
+    ----------
+    alpha : list[float]
+        Initial alpha parameters (x,y)
+    beta : list[float] | None
+        Initial beta parameters (x,y), default [0,0]
+    position : list[float] | None
+        Initial position (6D), default [0,0,0,0,0,0]
+    dispersion : list[float] | None
+        Initial dispersion (4D), default [0,0,0,0]
+    """
+
+    def __init__(
+        self,
+        alpha: list[float],
+        beta: list[float] | None = None,
+        position: list[float] | None = None,
+        dispersion: list[float] | None = None,
+    ):
+        """
+        Create TwissIn object
+        """
+        super().__init__()
+        self.alpha = alpha
+        self.beta = beta if beta is not None else [0.0, 0.0]
+        self.position = position if position is not None else [0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
+        self.dispersion = dispersion if dispersion is not None else [0.0, 0.0, 0.0, 0.0]
+
+    def _to_at(self) -> dict:
+        return {
+            "beta": numpy.array(self.alpha, dtype=float),
+            "alpha": numpy.array(self.beta, dtype=float),
+            "closed_orbit": numpy.array(self.position, dtype=float),
+            "dispersion": numpy.array(self.dispersion, dtype=float),
+        }

tests/config/tl2/tl2.yml

@@ -0,0 +1,215 @@
+type: pyaml.accelerator
+facility: ESRF
+machine: tl2
+energy: 6e9
+
+simulators:
+  - class: pyaml.lattice.simulator.Simulator
+    lattice: tl2.mat
+    name: design
+    twiss_in:
+      class: pyaml.lattice.twiss_in.TwissIn
+      alpha: [-1.71,  1.47]
+      beta: [5.9 , 5.65]
+      dispersion: [-0.145, -0.092, 0, 0]
+
+data_folder: /data/store
+
+arrays:
+
+  - type: pyaml.arrays.magnet
+    name: HCORR
+    elements:
+      - CH*
+  - type: pyaml.arrays.magnet
+    name: VCORR
+    elements:
+      - CV*
+  - type: pyaml.arrays.bpm
+    name: BPMS
+    elements:
+      - BPM*
+
+devices:
+
+  - class: pyaml.tuning_tools.orbit.Orbit
+    bpm_array_name: BPMS
+    hcorr_array_name: HCORR
+    vcorr_array_name: VCORR
+    name: DEFAULT_ORBIT_CORRECTION
+    singular_values: 5
+    response_matrix: ideal_tl2_orm.json
+
+  - class: pyaml.tuning_tools.orbit_response_matrix.OrbitResponseMatrix
+    bpm_array_name: BPMS
+    hcorr_array_name: HCORR
+    vcorr_array_name: VCORR
+    corrector_delta: 1e-6
+    name: DEFAULT_ORBIT_RESPONSE_MATRIX
+
+  - class: pyaml.diagnostics.bpm.bpm.BPM
+    name: BPM_CV1
+    x_pos: tl2/d-bpm/cv1/X_Position
+    y_pos: tl2/d-bpm/cv1/Z_Position
+  - class: pyaml.diagnostics.bpm.bpm.BPM
+    name: BPM_CH3
+    x_pos: tl2/d-bpm/d2/X_Position
+    y_pos: tl2/d-bpm/d2/Z_Position
+  - class: pyaml.diagnostics.bpm.bpm.BPM
+    name: BPM_QD5
+    x_pos: tl2/d-bpm/q5/X_Position
+    y_pos: tl2/d-bpm/q5/Z_Position
+  - class: pyaml.diagnostics.bpm.bpm.BPM
+    name: BPM_QF6
+    x_pos: tl2/d-bpm/q6/X_Position
+    y_pos: tl2/d-bpm/q6/Z_Position
+  - class: pyaml.diagnostics.bpm.bpm.BPM
+    name: BPM_SX
+    x_pos: tl2/d-bpm/sx/X_Position
+    y_pos: tl2/d-bpm/sx/Z_Position
+  - class: pyaml.diagnostics.bpm.bpm.BPM
+    name: BPM_QD8
+    x_pos: tl2/d-bpm/q8/X_Position
+    y_pos: tl2/d-bpm/q8/Z_Position
+  - class: pyaml.diagnostics.bpm.bpm.BPM
+    name: BPM_QD12
+    x_pos: tl2/d-bpm/q12/X_Position
+    y_pos: tl2/d-bpm/q12/Z_Position
+  - class: pyaml.diagnostics.bpm.bpm.BPM
+    name: BPM_QF14
+    x_pos: tl2/d-bpm/q14/X_Position
+    y_pos: tl2/d-bpm/q14/Z_Position
+  - class: pyaml.diagnostics.bpm.bpm.BPM
+    name: BPM_S1PM
+    x_pos: tl2/d-bpm/s1pm/X_Position
+    y_pos: tl2/d-bpm/s1pm/Z_Position
+
+  - class: pyaml.magnet.hcorrector.HCorrector
+    name: CH1
+    model:
+      class: pyaml.magnet.linear_model.LinearMagnetModel
+      calibration_factor: 161e-6
+      unit: 1/m
+      hardware_unit: A
+      powerconverter: tl2/ps-c1/ch1/Current
+  - class: pyaml.magnet.hcorrector.HCorrector
+    name: CH2
+    model:
+      class: pyaml.magnet.linear_model.LinearMagnetModel
+      calibration_factor: 161e-6
+      unit: 1/m
+      hardware_unit: A
+      powerconverter: tl2/ps-c1/ch2/Current
+  - class: pyaml.magnet.hcorrector.HCorrector
+    name: CH3
+    model:
+      class: pyaml.magnet.linear_model.LinearMagnetModel
+      calibration_factor: 161e-6
+      unit: 1/m
+      hardware_unit: A
+      powerconverter: tl2/ps-c1/ch3/Current
+  - class: pyaml.magnet.hcorrector.HCorrector
+    name: CH4
+    model:
+      class: pyaml.magnet.linear_model.LinearMagnetModel
+      calibration_factor: 161e-6
+      unit: 1/m
+      hardware_unit: A
+      powerconverter: tl2/ps-c1/ch4/Current
+  - class: pyaml.magnet.hcorrector.HCorrector
+    name: CH5
+    model:
+      class: pyaml.magnet.linear_model.LinearMagnetModel
+      calibration_factor: 161e-6
+      unit: 1/m
+      hardware_unit: A
+      powerconverter: tl2/ps-c1/ch5/Current
+  - class: pyaml.magnet.hcorrector.HCorrector
+    name: CH6
+    model:
+      class: pyaml.magnet.linear_model.LinearMagnetModel
+      calibration_factor: 161e-6
+      unit: 1/m
+      hardware_unit: A
+      powerconverter: tl2/ps-c1/ch6/Current
+  - class: pyaml.magnet.hcorrector.HCorrector
+    name: CH7
+    model:
+      class: pyaml.magnet.linear_model.LinearMagnetModel
+      calibration_factor: 161e-6
+      unit: 1/m
+      hardware_unit: A
+      powerconverter: tl2/ps-c1/ch7/Current
+  - class: pyaml.magnet.hcorrector.HCorrector
+    name: CH8
+    model:
+      class: pyaml.magnet.linear_model.LinearMagnetModel
+      calibration_factor: 161e-6
+      unit: 1/m
+      hardware_unit: A
+      powerconverter: tl2/ps-c1/ch8/Current
+
+  - class: pyaml.magnet.vcorrector.VCorrector
+    name: CV0
+    model:
+      class: pyaml.magnet.linear_model.LinearMagnetModel
+      calibration_factor: 161e-6
+      unit: 1/m
+      hardware_unit: A
+      powerconverter: tl2/ps-c1/cv0/Current
+  - class: pyaml.magnet.vcorrector.VCorrector
+    name: CV1
+    model:
+      class: pyaml.magnet.linear_model.LinearMagnetModel
+      calibration_factor: 161e-6
+      unit: 1/m
+      hardware_unit: A
+      powerconverter: tl2/ps-c1/cv1/Current
+  - class: pyaml.magnet.vcorrector.VCorrector
+    name: CV3
+    model:
+      class: pyaml.magnet.linear_model.LinearMagnetModel
+      calibration_factor: 161e-6
+      unit: 1/m
+      hardware_unit: A
+      powerconverter: tl2/ps-c1/cv3/Current
+  - class: pyaml.magnet.vcorrector.VCorrector
+    name: CV4
+    model:
+      class: pyaml.magnet.linear_model.LinearMagnetModel
+      calibration_factor: 161e-6
+      unit: 1/m
+      hardware_unit: A
+      powerconverter: tl2/ps-c1/cv4/Current
+  - class: pyaml.magnet.vcorrector.VCorrector
+    name: CV5
+    model:
+      class: pyaml.magnet.linear_model.LinearMagnetModel
+      calibration_factor: 161e-6
+      unit: 1/m
+      hardware_unit: A
+      powerconverter: tl2/ps-c1/cv5/Current
+  - class: pyaml.magnet.vcorrector.VCorrector
+    name: CV6
+    model:
+      class: pyaml.magnet.linear_model.LinearMagnetModel
+      calibration_factor: 161e-6
+      unit: 1/m
+      hardware_unit: A
+      powerconverter: tl2/ps-c1/cv6/Current
+  - class: pyaml.magnet.vcorrector.VCorrector
+    name: CV7
+    model:
+      class: pyaml.magnet.linear_model.LinearMagnetModel
+      calibration_factor: 161e-6
+      unit: 1/m
+      hardware_unit: A
+      powerconverter: tl2/ps-c1/cv7/Current
+  - class: pyaml.magnet.vcorrector.VCorrector
+    name: CV8
+    model:
+      class: pyaml.magnet.linear_model.LinearMagnetModel
+      calibration_factor: 161e-6
+      unit: 1/m
+      hardware_unit: A
+      powerconverter: tl2/ps-c1/cv8/Current