Feature: shifted quadrupole support

.github/workflows/pySC_tests.yml

@@ -0,0 +1,29 @@
+name: Tests
+
+on:
+  push:
+    branches: [main, develop]
+  pull_request:
+    branches: [main, develop]
+
+jobs:
+  tests:
+    runs-on: ubuntu-latest
+
+    strategy:
+      matrix:
+        python-version: ["3.12", "3.13", "3.14"]
+
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - name: Install package with test dependencies
+        run: pip install ".[test]"
+
+      - name: Run tests
+        run: pytest --cov

.github/workflows/pyaml_tests.yml

@@ -0,0 +1,40 @@
+name: pyAML tests
+
+on:
+  push:
+    branches: [main, develop]
+  pull_request:
+    branches: [main, develop]
+
+jobs:
+  tests:
+    runs-on: ubuntu-latest
+
+    strategy:
+      matrix:
+        python-version: ["3.12", "3.14"]
+
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Clone second repository
+        uses: actions/checkout@v4
+        with:
+          repository: python-accelerator-middle-layer/pyaml
+          path: pyaml
+          ref: main
+
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - name: Install package with test dependencies
+        run: |
+          pip install "./pyaml[test]"
+          pip install ./pyaml/tests/dummy_cs/tango-pyaml
+          pip install ".[test]"
+
+      - name: Run tests
+        working-directory: pyaml
+        run: pytest -k "tuning"

pySC/configuration/magnets_conf.py

@@ -1,4 +1,5 @@
 from typing import Any
+import math
 from ..core.simulated_commissioning import SimulatedCommissioning
 from ..core.lattice import ATLattice
 from ..core.magnet import MAGNET_NAME_TYPE
@@ -18,6 +19,22 @@ def generate_default_magnet_control(SC: SimulatedCommissioning, index: int, magn
     components_to_invert = dict.get(magnet_category_conf, 'invert', []).copy() # defaults to empty list if not declared
     # we need to copy because we remove elements later to check for undeclared components to invert
 
+    is_shifted = magnet_category_conf.get('shifted', False)
+    bending_angle = SC.lattice.get_bending_angle(index)
+    bending_length = SC.lattice.get_length(index)
+    magnet_length = bending_length
+    design_shift = 0.0
+    design_k1 = 0.0
+    if is_shifted:
+        element = SC.lattice.design[index]
+        component_value = getattr(element, 'PolynomB')[1]
+        if not SC.lattice.is_dipole(index):
+            raise ValueError(f"magnets/{magnet_category_name}: shifted=true requires an sbend with non-zero quad component at index {index} ({magnet_name}).")
+        radius = bending_length/bending_angle
+        magnet_length = 2*abs(radius)*abs(math.sin(bending_angle/2))
+        design_shift = bending_angle/(component_value*bending_length)
+        design_k1 = component_value
+
     if 'components' in magnet_category_conf:
         components = []
         cal_errors = []
@@ -26,11 +43,17 @@ def generate_default_magnet_control(SC: SimulatedCommissioning, index: int, magn
             components.append(component)
             cal_errors.append(cal_error)
 
-        magnet_length = SC.lattice.get_length(index)
         magnet_settings.add_individually_powered_magnet(
-            sim_index=index, controlled_components=components,
-            magnet_name=magnet_name, magnet_length=magnet_length,
-            to_design=to_design)
+            sim_index=index, 
+            controlled_components=components,
+            magnet_name=magnet_name, 
+            magnet_length=magnet_length,
+            is_shifted=is_shifted,
+            bending_length=bending_length,
+            design_shift=design_shift,
+            design_k1=design_k1,
+            to_design=to_design
+        )
 
         for component, cal_error in zip(components, cal_errors):
             control_name = f'{magnet_name}/{component}'
@@ -59,7 +82,7 @@ def generate_default_magnet_control(SC: SimulatedCommissioning, index: int, magn
                 offset = 0
                 setpoint = SC.lattice.get_magnet_component(index, component_type=component_type, order=order)
                 if component[-1] == 'L':
-                    length = SC.lattice.get_length(index)
+                    length = magnet_length
                     setpoint = setpoint * length
 
             if component in components_to_invert:
@@ -113,4 +136,4 @@ def configure_magnets(SC: SimulatedCommissioning):
     SC.magnet_settings.connect_links()
     SC.magnet_settings.sendall()
     SC.design_magnet_settings.connect_links()
-    SC.design_magnet_settings.sendall()
+    SC.design_magnet_settings.sendall()

pySC/core/magnet.py

@@ -30,6 +30,10 @@ class Magnet(BaseModel, extra="forbid"):
     offset_B: Optional[list[float]] = None
     to_design: bool = False
     length: Optional[float] = None
+    is_shifted: bool = False
+    bending_length: Optional[float] = None
+    design_shift: float = 0.0
+    design_k1: float = 0.0
     _links: list[ControlMagnetLink] = PrivateAttr(default=[])
     _parent = PrivateAttr(default=None)
 
@@ -112,8 +116,38 @@ def update(self):
                     f"Invalid component '{link.component}' for magnet '{self.name}'"
                 )
 
+        if self.is_shifted:
+            assert self.length is not None, f"ERROR: quadrupole length not specified for shifted magnet: {repr(self)}"
+            shift = self.design_shift
+            if not self.to_design:
+                dx, _ = self._parent._parent.support_system.get_total_offset(self.sim_index)
+                shift += dx
+            k1 = self.B[1]
+            self.B[0] += shift * k1 - self.design_shift * self.design_k1
+
         for ii in range(self.max_order + 1):
             self._parent._parent.lattice.set_magnet_component(
                 self.sim_index, self.A[ii], 'A', ii, use_design=self.to_design)
             self._parent._parent.lattice.set_magnet_component(
                 self.sim_index, self.B[ii], 'B', ii, use_design=self.to_design)
+
+        if self.is_shifted:
+
+            if self.to_design:
+                dx = dy = dz = 0.0
+                roll = yaw = pitch = 0.0
+            else:
+                dx, dy = self._parent._parent.support_system.get_total_offset(self.sim_index)
+                dz = self._parent._parent.support_system.data['L0'][self.sim_index].dz
+                roll, pitch, yaw = self._parent._parent.support_system.get_total_rotation(self.sim_index)
+
+            self._parent._parent.lattice.update_misalignment(
+                self.sim_index,
+                dx=dx,
+                dy=dy,
+                dz=dz,
+                roll=roll,
+                yaw=yaw,
+                pitch=pitch,
+                use_design=self.to_design
+            )

pySC/core/magnetsettings.py

@@ -94,6 +94,10 @@ def add_individually_powered_magnet(self,
                                         controlled_components: list[str],
                                         magnet_name: Optional[str] = None,
                                         magnet_length: Optional[float] = None,
+                                        is_shifted: bool = False,
+                                        bending_length: Optional[float] = None,
+                                        design_shift: float = 0.0,
+                                        design_k1: float = 0.0,
                                         to_design: bool = False) -> None:
         """
         Add a magnet with individually powered components.
@@ -112,7 +116,11 @@ def add_individually_powered_magnet(self,
                         sim_index=sim_index,
                         max_order=max_order,
                         to_design=to_design,
-                        length=magnet_length)
+                        length=magnet_length,
+                        is_shifted=is_shifted,
+                        bending_length=bending_length,
+                        design_shift=design_shift,
+                        design_k1=design_k1)
         magnet._parent = self  # Set the parent to the current settings instance
 
         # check non-zero components that are not controlled and put them in offset_A/B.

pySC/core/supports.py

@@ -329,6 +329,11 @@ def trigger_update(self, level: str, index):
             else:
                 self._parent.lattice.update_misalignment(index=eo.index, dx=dx, dy=dy, dz=dz,
                                       roll=roll, yaw=yaw, pitch=pitch)
+                magnet_name = self._parent.magnet_settings.index_mapping.get(eo.index)
+                if magnet_name is not None:
+                    magnet = self._parent.magnet_settings.magnets[magnet_name]
+                    if magnet.is_shifted:
+                        magnet.update()
 
     def update_all(self) -> None:
         for index in self.data['L0'].keys():

tests/core/test_magnet.py

@@ -1,7 +1,6 @@
 """Tests for pySC.core.magnet: Magnet, ControlMagnetLink."""
 import pytest
-from unittest.mock import MagicMock, PropertyMock
-
+from unittest.mock import MagicMock
 from pySC.core.magnet import Magnet, ControlMagnetLink
 from pySC.core.control import Control, LinearConv
 
@@ -118,3 +117,81 @@ def test_magnet_update_no_length_raises():
 
     with pytest.raises(AssertionError, match="magnet length not specified"):
         m.update()
+
+
+def test_shifted_magnet_update():
+    """Shifted quadrupole feeds down"""
+
+    design_arc_length = 0.25
+    design_shift = 0.01
+    initial_k1 = 5.0
+    initial_bending_angle = design_shift * initial_k1 * design_arc_length
+    initial_b0 = 0.02
+
+    m, parent = _make_magnet_with_parent(max_order=1, length=design_arc_length)
+    m.is_shifted = True
+    m.bending_length = design_arc_length
+    m.design_shift = design_shift
+    m.design_k1 = initial_k1
+    m.offset_B[0] = initial_b0
+
+    element = MagicMock()
+    element.Length = design_arc_length
+    element.BendingAngle = initial_bending_angle
+    element.EntranceAngle = initial_bending_angle / 2
+    element.ExitAngle = initial_bending_angle / 2
+
+    support_system = MagicMock()
+    support_system.get_total_offset.return_value = (0.0, 0.0)
+    support_system.get_total_rotation.return_value = (0.0, 0.0, 0.0)
+    support_system.data = {'L0': {0: MagicMock(dz=0.0)}}
+    parent._parent.support_system = support_system
+    parent._parent.lattice.ring = {0: element}
+    parent._parent.lattice.design = {0: element}
+
+    ctrl = Control(name="c1", setpoint=initial_k1 * design_arc_length)
+    parent.controls["c1"] = ctrl
+    link = ControlMagnetLink(link_name="lk1", magnet_name=0, control_name="c1", component="B", order=2, is_integrated=True)
+    m._links = [link]
+
+    m.update()
+
+    expected_k1 = initial_k1
+    expected_b0 = initial_b0
+    expected_angle = initial_bending_angle
+    assert m.B[1] == pytest.approx(expected_k1)
+    assert m.B[0] == pytest.approx(expected_b0)
+    assert element.Length == pytest.approx(design_arc_length)
+    assert element.BendingAngle == pytest.approx(expected_angle)
+    assert element.EntranceAngle == pytest.approx(expected_angle/2)
+    assert element.ExitAngle == pytest.approx(expected_angle/2)
+    assert m.length == pytest.approx(design_arc_length)
+    assert m.bending_length == pytest.approx(design_arc_length)
+
+    ctrl.setpoint = 5.2 * design_arc_length
+    m.update()
+
+    expected_k1 = 5.2
+    expected_b0 = initial_b0 + design_shift * (expected_k1 - initial_k1)
+    assert m.B[1] == pytest.approx(expected_k1)
+    assert m.B[0] == pytest.approx(expected_b0)
+    assert element.Length == pytest.approx(design_arc_length)
+    assert element.BendingAngle == pytest.approx(initial_bending_angle)
+    assert element.EntranceAngle == pytest.approx(initial_bending_angle/2)
+    assert element.ExitAngle == pytest.approx(initial_bending_angle/2)
+    assert m.length == pytest.approx(design_arc_length)
+    assert m.bending_length == pytest.approx(design_arc_length)
+
+    support_system.get_total_offset.return_value = (0.001, 0.0)
+    m.update()
+
+    total_shift = design_shift + 0.001
+    expected_b0 = initial_b0 + total_shift * expected_k1 - design_shift * initial_k1
+    assert m.B[1] == pytest.approx(expected_k1)
+    assert m.B[0] == pytest.approx(expected_b0)
+    assert element.Length == pytest.approx(design_arc_length)
+    assert element.BendingAngle == pytest.approx(initial_bending_angle)
+    assert element.EntranceAngle == pytest.approx(initial_bending_angle/2)
+    assert element.ExitAngle == pytest.approx(initial_bending_angle/2)
+    assert m.length == pytest.approx(design_arc_length)
+    assert m.bending_length == pytest.approx(design_arc_length)