rdt optics module with some simple utilities, and two tests

pyhdtoolkit/optics/rdt.py

@@ -0,0 +1,59 @@
+"""
+.. _optics-rdt:
+
+Resonance Driving Terms Utilities
+-----------------
+
+Module implementing utilities for the handling of resonance driving terms.
+"""
+
+from typing import Union, Tuple
+
+
+def rdt_to_order_and_type(rdt: Union[int, str]) -> str:
+    """
+    Decompose the input RDT into its four various components
+    and return the type of RDT (normal or skew) and its order.
+
+    Args:
+        rdt (Union[int, str]): the RDT to decompose.
+
+    Returns:
+        A string with the type and (magnet) order of
+        the provided RDT.
+    """
+    j, k, l, m = map(int, str(rdt))  # noqa: E741
+    rdt_type = "normal" if (l + m) % 2 == 0 else "skew"
+    orders = dict(
+        (
+            (1, "dipole"),
+            (2, "quadrupole"),
+            (3, "sextupole"),
+            (4, "octupole"),
+            (5, "decapole"),
+            (6, "dodecapole"),
+            (7, "tetradecapole"),
+            (8, "hexadecapole"),
+        )
+    )
+    return f"{rdt_type}_{orders[j + k + l + m]}"
+
+
+def determine_rdt_line(rdt: Union[int, str], plane: str) -> Tuple[int, int, int]:
+    """
+    Find the given line to look for in the spectral analysis of
+    the given plane that corresponds to the given RDT.
+
+    Args:
+        rdt (Union[int, str]): the RDT to look for.
+        plane (str): the plane to look for the RDT in.
+
+    Returns:
+        A tuple of three integers representing the line
+        to look for in the spectral analysis. For instance,
+        f1001 corresponds to the line (0, 1, 0) in the X plane
+        which means the line located at 1 * Qy = Qy.
+    """
+    j, k, l, m = map(int, str(rdt))  # noqa: E741
+    lines = dict(X=(1 - j + k, m - l, 0), Y=(k - j, 1 - l + m, 0))
+    return lines[plane]

tests/test_optics.py

@@ -6,6 +6,7 @@
 from pyhdtoolkit.models.beam import BeamParameters
 from pyhdtoolkit.optics import ripken, twiss
 from pyhdtoolkit.optics.beam import Beam, compute_beam_parameters
+from pyhdtoolkit.optics.rdt import determine_rdt_line, rdt_to_order_and_type
 
 CURRENT_DIR = pathlib.Path(__file__).parent
 INPUTS_DIR = CURRENT_DIR / "inputs"
@@ -156,6 +157,32 @@ def test_add_beam_size_to_df(_non_matched_lhc_madx):
     assert "SIZE_Y" in df.columns
 
 
+def test_rdt_order_and_type():
+    assert rdt_to_order_and_type(1001) == "skew_quadrupole"
+    assert rdt_to_order_and_type(1010) == "skew_quadrupole"
+    assert rdt_to_order_and_type(1020) == "normal_sextupole"
+    assert rdt_to_order_and_type(2002) == "normal_octupole"
+    assert rdt_to_order_and_type(1003) == "skew_octupole"
+    assert rdt_to_order_and_type(1004) ==  "normal_decapole"
+    assert rdt_to_order_and_type(3003) == "skew_dodecapole"
+    assert rdt_to_order_and_type(4004) == "normal_hexadecapole"
+
+    with pytest.raises(KeyError):
+        rdt_to_order_and_type(8888)
+        rdt_to_order_and_type(1090)
+
+
+def test_rdt_spectrum_line():
+    # Some simple coupling RDTs lines
+    assert determine_rdt_line(1001, "X") == (0, 1, 0)
+    assert determine_rdt_line(1010, "Y") == (-1, 0, 0)
+    assert determine_rdt_line("0220", "X") == (3, -2, 0)
+
+    with pytest.raises(KeyError):
+        determine_rdt_line("0220", "Z")
+
+
+
 # ----- Fixtures ----- #