Bugfix: Wrong sign starting values (#11)

* bugfix and test
* bump version

irnl_rdt_correction/__init__.py

@@ -13,7 +13,7 @@
 __title__ = "irnl-rdt-correction"
 __description__ = "Correction script to power the nonlinear correctors in the (HL-)LHC insertion regions based on RDTs."
 __url__ = "https://github.com/pylhc/irnl_rdt_correction"
-__version__ = "1.1.0"
+__version__ = "1.1.1"
 __author__ = "pylhc"
 __author_email__ = "[email protected]"
 __license__ = "MIT"

irnl_rdt_correction/equation_system.py

@@ -7,17 +7,16 @@
 
 """
 import logging
-from typing import Sequence, Tuple, Set, Dict, List
+from typing import Dict, List, Sequence, Set, Tuple
 
 import numpy as np
 from numpy.typing import ArrayLike
 from pandas import DataFrame, Series
 
-from irnl_rdt_correction.constants import BETA, SIDES, PLANES, DELTA, KEYWORD, MULTIPOLE
-from irnl_rdt_correction.rdt_handling import IRCorrector, RDT, RDTMap
-from irnl_rdt_correction.utilities import (
-    list2str, i_pow, is_even, is_odd, is_anti_mirror_symmetric, idx2str, Optics
-)
+from irnl_rdt_correction.constants import BETA, DELTA, KEYWORD, MULTIPOLE, PLANES, SIDES
+from irnl_rdt_correction.rdt_handling import RDT, IRCorrector, RDTMap
+from irnl_rdt_correction.utilities import (Optics, corrector_sign_beam_symmetry, i_pow, idx2str,
+                                           is_even, is_odd, list2str)
 
 LOG = logging.getLogger(__name__)
 
@@ -237,6 +236,8 @@ def init_corrector_and_optics_values(correctors: Sequence[IRCorrector], optics_s
 
     Returns:
         Dict[IRCorrector, Sequence[float]]: The saved values per corrector (per Optics).
+                                            These are the values directly taken from the optics, 
+                                            i.e. the sign might be different to the sign stored in the correctors objects.
                                             If the optics are updated anyway, this is an empty dict.
     """
     saved_values = {}
@@ -245,7 +246,7 @@ def init_corrector_and_optics_values(correctors: Sequence[IRCorrector], optics_s
         values = [optic.twiss.loc[corrector.name, corrector.strength_component] for optic in optics_seq]
 
         if not update_optics:
-            saved_values[corrector] = values
+            saved_values[corrector] = values  # saved without sign-change
 
         if ignore_settings:
             # set corrector value in optics to zero
@@ -256,7 +257,8 @@ def init_corrector_and_optics_values(correctors: Sequence[IRCorrector], optics_s
                 raise ValueError(f"Initial value for corrector {corrector.name} differs "
                                  f"between optics.")
             # use optics value as initial value (as equation system calculates corrector delta!!)
-            corrector.value = values[0]
+            sign = corrector_sign_beam_symmetry(optics_seq[0].beam, corrector.strength_component)
+            corrector.value = sign * values[0]
     return saved_values
 
 
@@ -425,10 +427,7 @@ def get_corrector_coefficient(rdt: RDT, corrector: IRCorrector, optics: Optics)
     # in Beam 2 and Beam 4. The correct sign in MAD-X is then assured by the
     # lattice setup, where these correctors have a minus sign in Beam 4.
     # See Chapter 3.1 Beam Directions in [DillyNonlinearIRCorrections2023]_
-    sign_beam = 1
-    if is_even(optics.beam) and is_anti_mirror_symmetric(corrector.strength_component):
-        sign_beam = -1
-
+    sign_beam = corrector_sign_beam_symmetry(optics.beam, corrector.strength_component)
     return sign_beam * sign_corrector * z * betax * betay
 
 
@@ -447,8 +446,7 @@ def get_side_sign(n: int, side: str) -> int:
         int: Either -1 or 1
     """
     if side == "R":
-        # return (-1)**n
-        return -1 if n % 2 else 1
+        return -1 if n % 2 else 1  # == (-1)**n, but faster and exact
     return 1
 
 
@@ -535,7 +533,7 @@ def optics_update(correctors: Sequence[IRCorrector], optics_seq: Sequence[Optics
     """
     for optics in optics_seq:
         for corrector in correctors:
-            sign = -1 if is_even(optics.beam) and is_anti_mirror_symmetric(corrector.strength_component) else 1
+            sign = corrector_sign_beam_symmetry(optics.beam, corrector.strength_component)
             optics.twiss.loc[corrector.name, corrector.strength_component] = sign * corrector.value
 
 

irnl_rdt_correction/utilities.py

@@ -126,6 +126,24 @@ def i_pow(n: int) -> complex:
     return 1j**(n % 4)   # more exact with modulo
 
 
+def corrector_sign_beam_symmetry(beam: int, column_name: str) -> int:
+    """This function returns -1 if there is if there is a sign change
+    change for the powering of the corrector and it's field component,
+    due to the anti-symmetric magnetic field.
+    See Chapter 3.1 Beam Directions in [DillyNonlinearIRCorrections2023]_
+
+    Args:
+        beam (int): beam used
+        column_name (str): column name of the field strength
+
+    Returns:
+        int: -1 if the corrector circuit has opposite sign to its field, 1 otherwise.
+    """
+    if is_even(beam) and is_anti_mirror_symmetric(column_name):
+        return -1
+    return 1
+
+
 def log_setup():
     """ Set up a basic logger. """
     logging.basicConfig(

tests/helpers.py

@@ -77,11 +77,11 @@ def generate_pseudo_model(n_ips: int, n_magnets: int, accel: str,
     return df
 
 
-def generate_errortable(index: Sequence, value: float = 0) -> pd.DataFrame:
+def generate_errortable(index: Sequence, value: float = 0, max_order: int = MAX_N) -> pd.DataFrame:
     """Return DataFrame from index and KN(S)L + D[XY] columns."""
     return pd.DataFrame(value,
                         index=index,
-                        columns=[f"K{n}{o}L" for n in range(MAX_N) for o in ("", "S")] + [f"D{plane}" for plane in "XY"]
+                        columns=[f"K{n}{o}L" for n in range(max_order) for o in ("", "S")] + [f"D{plane}" for plane in "XY"]
                         )
 
 

tests/test_standard_correction.py

@@ -1,15 +1,18 @@
+from dataclasses import dataclass
 from pathlib import Path
 
 import numpy as np
+import pandas as pd
 import pytest
 
-from irnl_rdt_correction.constants import KEYWORD
+from irnl_rdt_correction.constants import KEYWORD, DELTA, BETA
 from irnl_rdt_correction.main import irnl_rdt_correction
 from tests.helpers import (
     generate_pseudo_model, generate_errortable, get_some_magnet_names,
     read_lhc_model, get_ir_magnets_mask, get_corrector_magnets_mask,
     get_opposite_sign_beam4_kl_columns,
     NAME, PLACEHOLDER, CIRCUIT, STRENGTH, IP, EPS, VALUE, MAX_N,
+    X, Y
 )
 
 
@@ -31,7 +34,7 @@ def test_basic_correction(tmp_path: Path, order: int, orientation: str, accel: s
         if order == 6 and orientation == 'skew':
             pytest.skip("LHC has no skew dodecapole correctors")
 
-    orientation = "S" if orientation is "skew" else ""
+    orientation = "S" if orientation == "skew" else ""
 
     correct_ips = (1, 3)
     error_value = 2
@@ -169,3 +172,179 @@ def test_lhc_correction(tmp_path: Path, beam: int):
     # All circuits in madx script ---
     for circuit in df_corrections[CIRCUIT]:
         assert circuit in madx_corrections
+
+
+@pytest.mark.parametrize('accel', ('lhc', 'hllhc'))
+@pytest.mark.parametrize('feeddown', (0, 2))
+def test_beams_symmetries(tmp_path: Path, accel: str, feeddown: int):
+    """ Very similar to the basic correction, but here we check that we actually 
+    get the correct symmetries between the beams. 
+    It also checks that the pre-powering of the magnets is stored with the correct sign.
+    This test should fail with irnl_correction versions <= v1.0.1.
+    """
+    np.random.seed(20231003) 
+
+    correct_ips = (1,)
+    n_magnets = 1
+    n_ips = 1
+
+    is_lhc = accel == 'lhc'
+    orientations = ("", "S")
+    orders = list(range(2, 6))
+    order_letters = ["S", "O", "D", "T"]
+
+    if feeddown:
+        rdts = ['F0003', 'F0003*', 'F1002', 'F1002*']
+    else:
+        rdts = ['F0003', 'F0003*', 'F1002', 'F1002*',
+                'F1003', 'F3001', 'F4000', 'F0004',] + ([] if is_lhc else ['F0005', 'F0005*', 'F5001', 'F1005'])
+
+
+    corrector_order_map = {f"{order_name}{skew}": f"K{order}{skew}L" for order, order_name in zip(orders, order_letters) for skew in orientations}  # madx order
+
+    beam_results = {}
+
+    @dataclass
+    class BeamResult:
+        beam: int
+        twiss: pd.DataFrame
+        errors: pd.DataFrame
+        madx_corrections: str = None
+        df_corrections: pd.DataFrame = None
+
+    # Setup ----------------------------------------------------------------
+    # The setup is designed such that without feeddown, all corrections 
+    # yield the same values, which are n_magnets * 1.5 on the left and n_magents * 2 on the right.
+    # With feeddown only a3 and b3 are corrected:
+    # a3 = a3 + x*a4 + y*b4 + xy*b5 + 1/2(x**2 - y**2)*a5
+    # b3 = b3 + x*b4 - y*b4 - xy*a5 + 1/2(x**2 - y**2)*b5
+    # as x == y == 1
+    # a3 = 4 * n_magnets * 1.5
+    # b3 = 0
+
+    twiss_0 = generate_pseudo_model(accel=accel, n_ips=n_ips, n_magnets=n_magnets, betax=1, betay=1)
+    errors_0 = generate_errortable(index=get_some_magnet_names(n_ips=n_ips, n_magnets=n_magnets), max_order=0)
+
+    # here: 2 == sextupole
+    kl_columns = [f"K{order}{orientation}L" for order in orders for orientation in orientations]
+
+    twiss_0.loc[:, kl_columns]  = 0
+    errors_0.loc[:, kl_columns]  = 0
+
+    # orbit at corrector should not matter, as we don't use feed-down to correct here)
+    # twiss_0.loc[:, [X, Y]] = 0.5
+    # errors_0.loc[:, [f"{DELTA}{X}", f"{DELTA}{Y}"]] = 0.5
+
+    # Pre-Power corrector magnets, to test this does not cause problems (<=v1.0.1 does) 
+    # only power the field that will be modified by the correction, 
+    # e.g. normal quadrupole field for normal quadrupole correctors
+    if feeddown == 0:
+        for magnet_order, kl in corrector_order_map.items():
+            if magnet_order in ("D", "DS") and is_lhc:
+                continue
+            corrector_magnets = twiss_0.index.str.match(rf"MC{magnet_order}X") 
+            twiss_0.loc[corrector_magnets, kl] = 1
+
+    # Power other magnets and assign errors
+    non_corrector_magnets = twiss_0.index.str.match(r"M.\.")  # M[A-Z]. is created above
+    twiss_0.loc[non_corrector_magnets, kl_columns] = 1
+    twiss_0.loc[non_corrector_magnets, [X, Y]] = 0.5
+
+    non_corrector_magnets = errors_0.index.str.match(r"M.\.")  # M[A-Z]. is created above
+    errors_0.loc[non_corrector_magnets, kl_columns] = 1
+    errors_0.loc[non_corrector_magnets, [f"{DELTA}{X}", f"{DELTA}{Y}"]] = 0.5
+
+    # modify the left side, so they don't fully compensate for even (madx)-orders
+    left_hand_magnets = twiss_0.index.str.match(r".*L\d$")
+    twiss_0.loc[left_hand_magnets, f"{BETA}{Y}"] = 2 * twiss_0.loc[left_hand_magnets, f"{BETA}{Y}"]
+
+    left_hand_magnets = errors_0.index.str.match(r".*L\d$")
+    errors_0.loc[left_hand_magnets, kl_columns] = errors_0.loc[left_hand_magnets, kl_columns] / 2  
+
+    # Pre-calculate the integral based on this setup. 
+    integral_left = 1.5 * n_magnets
+    integral_right = 2.0 * n_magnets
+
+    # Correction -----------------------------------------------------------
+    for beam in (1, 2, 4):
+        # Create twiss per beam based on twiss_0 and the symmetries involved, 
+        # using the conventions for twiss and errors as implemented in MAD-X.
+        # ------- Reminder of MAD-X conventions ----------
+        #           Orbit:                        antisymmetic K:
+        # twiss:     X B1 =  X B2 =  -X B4        K B1 = -K B2 =  K B4 (i.e. K B1 = -K B4)
+        # errors:   DX B1 = DX B2 = -DX B4        K B1 =  K B2 = -K B4 
+        #
+        # BUT ALL CORRECTIONS SHOULD BE THE SAME, 
+        # as the the correction circuits are the same between all beams!
+
+        twiss = twiss_0.copy()
+        errors = errors_0.copy()
+
+        # Transform symmetries ----
+        negative_columns = get_opposite_sign_beam4_kl_columns(orders)
+        if beam in (2, 4):
+            twiss.loc[:, negative_columns] = -twiss.loc[:, negative_columns]
+
+        if beam == 4:
+            twiss.loc[:, X] = -twiss.loc[:, X]
+            errors.loc[:, f"{DELTA}{X}"] = -errors.loc[:, f"{DELTA}{X}"]
+            errors.loc[:, negative_columns] = -errors.loc[:, negative_columns]
+
+        beam_results[beam] = BeamResult(
+            beam=beam,
+            twiss=twiss,
+            errors=errors,
+        )
+
+        # Calculate corrections ---
+        beam_results[beam].madx_corrections, beam_results[beam].df_corrections = irnl_rdt_correction(
+            accel=accel,
+            twiss=[twiss],
+            errors=[errors],
+            beams=[beam],
+            rdts=rdts,
+            output=tmp_path / "correct",
+            feeddown=feeddown,
+            ips=correct_ips,
+            ignore_missing_columns=True,
+            iterations=1,
+            # ignore_corrector_settings = True,
+        )
+
+    # Testing --------------------------------------------------------------
+    # Check output data ---
+
+    # DEBUGGING: ---------------------------------
+    # for i in (1, 2, 4):
+    #     print(f"\n--- BEAM {i} ---")
+    #     print(beam_results[i].madx_corrections)
+    # --------------------------------------------
+
+    eps = 1e-14
+
+    # Comapre values between beams (this should always be true!!)
+    df_beam1 = beam_results[1].df_corrections
+    for beam in (2, 4):
+        df = beam_results[beam].df_corrections
+        assert np.allclose(df[VALUE], df_beam1[VALUE], atol=eps)
+
+    # Check values per beam (this depends a bit on the setup above)
+    for beam in (1, 2, 4):
+        df = beam_results[beam].df_corrections
+        assert len(df.index) == len(rdts)
+        left_correctors = df[CIRCUIT].str.match(r".*L\d$")
+        right_correctors = df[CIRCUIT].str.match(r".*R\d$")
+        assert left_correctors.sum() == len(rdts) * n_ips / 2
+        assert right_correctors.sum() == len(rdts) * n_ips / 2
+
+        if feeddown == 0:
+            assert np.allclose(df.loc[left_correctors, VALUE], -integral_left, atol=eps, rtol=0)
+            assert np.allclose(df.loc[right_correctors, VALUE], -integral_right, atol=eps, rtol=0)
+        else:
+            b3_correctors = df[CIRCUIT].str.match(r"KCSX.*")
+            a3_correctors = df[CIRCUIT].str.match(r"KCSSX.*")
+
+            assert np.allclose(df.loc[left_correctors & a3_correctors, VALUE], -integral_left * 4, atol=eps, rtol=0)
+            assert np.allclose(df.loc[right_correctors & a3_correctors, VALUE], -integral_right * 4, atol=eps, rtol=0)
+
+            assert np.allclose(df.loc[b3_correctors, VALUE], 0, atol=eps, rtol=0)