diff --git a/irnl_rdt_correction/__init__.py b/irnl_rdt_correction/__init__.py index 5d922b4..72ea7ae 100644 --- a/irnl_rdt_correction/__init__.py +++ b/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__ = "pylhc@github.com" __license__ = "MIT" diff --git a/irnl_rdt_correction/equation_system.py b/irnl_rdt_correction/equation_system.py index e9d3587..ec5e320 100644 --- a/irnl_rdt_correction/equation_system.py +++ b/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 diff --git a/irnl_rdt_correction/utilities.py b/irnl_rdt_correction/utilities.py index 7c8fe32..bd48d63 100644 --- a/irnl_rdt_correction/utilities.py +++ b/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( diff --git a/tests/helpers.py b/tests/helpers.py index da137da..2b7bf7c 100644 --- a/tests/helpers.py +++ b/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"] ) diff --git a/tests/test_standard_correction.py b/tests/test_standard_correction.py index 2cf3203..fdacaf1 100644 --- a/tests/test_standard_correction.py +++ b/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)