diff --git a/pySC/core/magnet.py b/pySC/core/magnet.py index 8e797c0..a9dbbe4 100644 --- a/pySC/core/magnet.py +++ b/pySC/core/magnet.py @@ -92,6 +92,11 @@ def state(self): link_value = link.value(setpoint) if link.is_integrated: link_value = link_value / self.length + # Mirror the order-1 B (dipole/corrector) sign flip + # applied in update(), so this diagnostic agrees with + # the stored PolynomB[0] value printed above. + if link.component == "B" and link.order == 1: + link_value = -link_value error = link.error print( f" - {link.control_name}: setpoint = {setpoint}, error = {repr(error)} -> {link_value}" @@ -109,7 +114,15 @@ def update(self): if link.is_integrated: assert self.length is not None, f'ERROR: magnet length not specified for integrated strength link: {repr(link)}' value = value / self.length - # if it is equal to zero then assume A and B are already integrated strengths :( + # AT sign convention for the dipole/corrector term: Δx' = -PolynomB[0]*L, + # so a positive horizontal kick requires negative PolynomB[0]. Negate the + # integrated B-component *only* for the order-1 term (PolynomB[0]) so that a + # positive corrector setpoint produces a positive physical kick, matching + # MATLAB's SCsetCMs2SetPoints (normBy = [-1, 1] * Length). Higher-order + # integrated B strengths (e.g. an integrated quadrupole B2L -> PolynomB[1]) + # must NOT be flipped -- that would invert focusing. + if link.component == "B" and link.order == 1: + value = -value if link.component == "A": self.A[link.order - 1] += value elif link.component == "B": diff --git a/pySC/utils/sc_tools.py b/pySC/utils/sc_tools.py index 63041db..13d5f43 100644 --- a/pySC/utils/sc_tools.py +++ b/pySC/utils/sc_tools.py @@ -102,23 +102,39 @@ # return matrix_inv -# def scale_circumference(RING, circ, mode='abs'): # TODO -# allowed_modes = ("abs", "rel") -# if mode not in allowed_modes: -# raise ValueError(f'Unsupported circumference scaling mode: ``{mode}``. Allowed are {allowed_modes}.') -# C = at_wrapper.findspos(RING)[-1] -# D = 0 -# for ind in range(len(RING)): -# if RING[ind].PassMethod == 'DriftPass': -# D += RING[ind].Length -# if mode == 'rel': -# Dscale = 1 - (1 - circ) * C / D -# else: # mode == 'abs' -# Dscale = 1 - (C - circ) / D -# for ind in range(len(RING)): -# if RING[ind].PassMethod == 'DriftPass': -# RING[ind].Length = RING[ind].Length * Dscale -# return RING +def scale_circumference(ring, circ, mode='rel'): + """Scale the circumference of a ring by adjusting drift space lengths. + + Parameters + ---------- + ring : at.Lattice + The AT lattice to modify (modified in place). + circ : float + Circumference value. Interpretation depends on *mode*. + mode : str + ``'rel'`` — *circ* is a multiplicative factor (e.g. 1 + 1e-6). + ``'abs'`` — *circ* is the desired circumference in metres. + + Returns + ------- + at.Lattice + The modified lattice (same object, returned for chaining). + """ + allowed_modes = ('abs', 'rel') + if mode not in allowed_modes: + raise ValueError(f'Unsupported circumference scaling mode: {mode!r}. Allowed are {allowed_modes}.') + C = sum(elem.Length for elem in ring) + D = sum(elem.Length for elem in ring if elem.PassMethod == 'DriftPass') + if D == 0: + raise ValueError('No drift elements found in ring — cannot scale circumference.') + if mode == 'rel': + Dscale = 1 - (1 - circ) * C / D + else: # mode == 'abs' + Dscale = 1 - (C - circ) / D + for elem in ring: + if elem.PassMethod == 'DriftPass': + elem.Length = elem.Length * Dscale + return ring def update_transformation(element, dx=None, dy=None, dz=None, roll=None, yaw=None, pitch=None): diff --git a/tests/core/test_magnet.py b/tests/core/test_magnet.py index 69bf5ea..84a70e8 100644 --- a/tests/core/test_magnet.py +++ b/tests/core/test_magnet.py @@ -85,7 +85,13 @@ def test_magnet_update_resets_to_offsets(): def test_magnet_update_integrated_strength(): - """Link with is_integrated=True divides link value by magnet length.""" + """Link with is_integrated=True divides link value by magnet length. + + For the order-1 B-component (PolynomB[0], the dipole/corrector term) the + AT sign convention applies: Δx' = -PolynomB[0]*L, so positive setpoint → + negative PolynomB[0] (positive physical kick). The value is negated after + division for this term only. + """ m, parent = _make_magnet_with_parent(max_order=1, length=2.0) ctrl = Control(name="c1", setpoint=6.0) @@ -97,8 +103,48 @@ def test_magnet_update_integrated_strength(): m._links = [link] m.update() - # 6.0 / 2.0 = 3.0 - assert m.B[0] == pytest.approx(3.0) + # 6.0 / 2.0 = 3.0, negated for the order-1 B (corrector/dipole) term: -3.0 + assert m.B[0] == pytest.approx(-3.0) + + +def test_magnet_update_integrated_higher_order_b_not_negated(): + """Integrated higher-order B strengths (e.g. B2L -> PolynomB[1], a quadrupole) + must NOT be sign-flipped -- only the order-1 dipole/corrector term is. + + Guards against an over-broad negation that would invert integrated + quadrupole/sextupole strengths (focusing -> defocusing). + """ + m, parent = _make_magnet_with_parent(max_order=2, length=2.0) + + ctrl = Control(name="c1", setpoint=6.0) + 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() + # 6.0 / 2.0 = 3.0, NOT negated for the order-2 B (quadrupole) term. + assert m.B[1] == pytest.approx(3.0) + + +def test_magnet_update_integrated_a_not_negated(): + """Integrated A-component (skew) strengths are never negated -- the sign + convention flip is specific to the normal dipole term PolynomB[0].""" + m, parent = _make_magnet_with_parent(max_order=1, length=2.0) + + ctrl = Control(name="c1", setpoint=6.0) + parent.controls["c1"] = ctrl + link = ControlMagnetLink( + link_name="lk1", magnet_name=0, control_name="c1", + component="A", order=1, is_integrated=True, + ) + m._links = [link] + + m.update() + # 6.0 / 2.0 = 3.0, A-component is not negated. + assert m.A[0] == pytest.approx(3.0) def test_magnet_update_no_length_raises(): diff --git a/tests/utils/test_sc_tools.py b/tests/utils/test_sc_tools.py index f29cc03..0546820 100644 --- a/tests/utils/test_sc_tools.py +++ b/tests/utils/test_sc_tools.py @@ -1,9 +1,9 @@ -"""Tests for pySC.utils.sc_tools: rotation matrices and coordinate transformations.""" +"""Tests for pySC.utils.sc_tools: rotation matrices, coordinate transformations, and circumference scaling.""" import numpy as np import pytest from types import SimpleNamespace -from pySC.utils.sc_tools import rotation, update_transformation +from pySC.utils.sc_tools import rotation, update_transformation, scale_circumference class TestRotation: @@ -151,3 +151,86 @@ def test_update_transformation_modifies_element_in_place(self): elem = _make_element(length=1.0) result = update_transformation(elem, dx=0.001, dy=0, dz=0, roll=0, yaw=0, pitch=0) assert result is elem + + +# --------------------------------------------------------------------------- +# scale_circumference +# --------------------------------------------------------------------------- + +def _make_ring(drift_lengths, magnet_lengths): + """Create a list of mock elements with PassMethod and Length attributes.""" + elements = [] + for L in drift_lengths: + elements.append(SimpleNamespace(PassMethod='DriftPass', Length=L)) + for L in magnet_lengths: + elements.append(SimpleNamespace(PassMethod='BndMPoleSymplectic4Pass', Length=L)) + return elements + + +class TestScaleCircumference: + + def test_relative_mode_identity(self): + """circ=1.0 in relative mode should not change any lengths.""" + ring = _make_ring([5.0, 3.0], [2.0]) + scale_circumference(ring, 1.0, mode='rel') + assert ring[0].Length == pytest.approx(5.0) + assert ring[1].Length == pytest.approx(3.0) + assert ring[2].Length == pytest.approx(2.0) + + def test_relative_mode_scales_drifts(self): + """Relative scaling changes total circumference by the requested factor.""" + ring = _make_ring([5.0, 3.0], [2.0]) + C_before = sum(e.Length for e in ring) # 10.0 + circ_scaling = 1 + 1e-4 # stretch by 0.01% + scale_circumference(ring, circ_scaling, mode='rel') + C_after = sum(e.Length for e in ring) + np.testing.assert_allclose(C_after, C_before * circ_scaling, rtol=1e-12) + + def test_relative_mode_preserves_magnet_lengths(self): + """Only drifts are modified; magnet lengths are untouched.""" + ring = _make_ring([5.0, 3.0], [2.0, 1.5]) + scale_circumference(ring, 1.001, mode='rel') + assert ring[2].Length == pytest.approx(2.0) + assert ring[3].Length == pytest.approx(1.5) + + def test_absolute_mode(self): + """Absolute mode sets total circumference to the given value.""" + ring = _make_ring([5.0, 3.0], [2.0]) + target_C = 10.5 + scale_circumference(ring, target_C, mode='abs') + C_after = sum(e.Length for e in ring) + np.testing.assert_allclose(C_after, target_C, rtol=1e-12) + + def test_absolute_mode_preserves_magnets(self): + """Absolute mode only changes drift lengths.""" + ring = _make_ring([5.0, 3.0], [2.0]) + scale_circumference(ring, 10.5, mode='abs') + assert ring[2].Length == pytest.approx(2.0) + + def test_drift_amplification_factor(self): + """Verify the Dscale formula: drifts absorb all the circumference change. + + For C=10, D=8, circ_scaling=1+1e-3: + Dscale = 1 - (1 - 1.001) * 10/8 = 1 + 0.001 * 10/8 = 1.00125 + """ + ring = _make_ring([5.0, 3.0], [2.0]) + scale_circumference(ring, 1.001, mode='rel') + expected_Dscale = 1 + 0.001 * 10.0 / 8.0 # 1.00125 + assert ring[0].Length == pytest.approx(5.0 * expected_Dscale) + assert ring[1].Length == pytest.approx(3.0 * expected_Dscale) + + def test_invalid_mode_raises(self): + ring = _make_ring([5.0], [2.0]) + with pytest.raises(ValueError, match='Unsupported'): + scale_circumference(ring, 1.0, mode='invalid') + + def test_no_drifts_raises(self): + ring = [SimpleNamespace(PassMethod='BndMPoleSymplectic4Pass', Length=2.0)] + with pytest.raises(ValueError, match='No drift elements'): + scale_circumference(ring, 1.001, mode='rel') + + def test_returns_same_ring(self): + """scale_circumference modifies in place and returns the same object.""" + ring = _make_ring([5.0], [2.0]) + result = scale_circumference(ring, 1.001, mode='rel') + assert result is ring