AC dipole hotfix and kicker plus matrix implementations (#51)

pyhdtoolkit/cpymadtools/__init__.py

@@ -20,7 +20,8 @@
     apply_lhc_coupling_knob,
     apply_lhc_rigidity_waist_shift_knob,
     deactivate_lhc_arc_sextupoles,
-    install_ac_dipole,
+    install_ac_dipole_as_kicker,
+    install_ac_dipole_as_matrix,
     make_lhc_beams,
     make_lhc_thin,
     make_sixtrack_output,

pyhdtoolkit/cpymadtools/special.py

@@ -197,28 +197,35 @@ def apply_lhc_coupling_knob(
     logger.trace(f"Set '{knob_name}' to {madx.globals[knob_name]}")
 
 
-def install_ac_dipole(
+def install_ac_dipole_as_kicker(
     madx: Madx,
     deltaqx: float,
     deltaqy: float,
     sigma_x: float,
     sigma_y: float,
+    beam: int = 1,
     geometric_emit: float = None,
     start_turn: int = 100,
     ramp_turns: int = 2000,
     top_turns: int = 6600,
 ) -> None:
     """
-    Installs an AC dipole for BEAM 1 ONLY.
-    This function assumes that you have already defined lhcb1, made a beam for it (BEAM command or
-    `make_lhc_beams` function) and matched to your desired working point.
+    Installs an AC dipole for (HL)LHC BEAM 1 OR 2 ONLY.
+    The AC Dipole does impact the orbit as well as the betatron functions when turned on. Unfortunately in
+    MAD-X, it cannot be modeled to do both at the same time. This routine introduces an AC Dipole as a
+    kicker element so that its effect can be seen on particle orbit in tracking. It DOES NOT affect TWISS
+    functions. See https://journals.aps.org/prab/abstract/10.1103/PhysRevSTAB.11.084002 (part III).
+
+    This function assumes that you have already defined lhcb1/lhcb2, made a beam for it (BEAM command or
+    `make_lhc_beams` function), matched to your desired working point and made a TWISS.
 
     Args:
         madx (cpymad.madx.Madx): an instanciated cpymad Madx object.
         deltaqx (float): the deltaQx (horizontal tune excitation) used by the AC dipole.
         deltaqy (float): the deltaQy (vertical tune excitation) used by the AC dipole.
         sigma_x (float): the horizontal amplitude to drive the beam to, in bunch sigma.
         sigma_y (float): the vertical amplitude to drive the beam to, in bunch sigma.
+        beam (int): the LHC beam to install the AC Dipole into, either 1 or 2. Defaults to 1.
         geometric_emit (float): the geometric emittance that was used when defining the beam. If not
             provided, it is assumed that 'geometric_emit' is a defined global in MAD-X, and the value will
             be directly queried from the internal tables.
@@ -228,6 +235,8 @@ def install_ac_dipole(
             as in the LHC.
         top_turns (int): the number of turns to drive the beam for. Defaults to 6600 as in the LHC.
     """
+    logger.warning("This AC Dipole is implemented as a kicker and will not affect TWISS functions!")
+    logger.info("This routine should be done after 'match', 'twiss' and 'makethin' for the appropriate beam.")
     if top_turns > 6600:
         logger.warning(
             f"Configuring the AC Dipole for {top_turns} of driving is fine for MAD-X but is "
@@ -251,22 +260,26 @@ def install_ac_dipole(
 
     logger.info(f"Installing AC Dipole to drive the tunes to Qx_D = {q1_dipole}  |  Qy_D = {q2_dipole}")
     madx.input(
-        f"MKACH.6L4.B1: hacdipole, l=0, freq:={q1_dipole}, lag=0, volt:=voltx, ramp1={ramp1}, "
+        f"MKACH.6L4.B{beam:d}: hacdipole, l=0, freq:={q1_dipole}, lag=0, volt:=voltx, ramp1={ramp1}, "
         f"ramp2={ramp2}, ramp3={ramp3}, ramp4={ramp4};"
     )
     madx.input(
-        f"MKACV.6L4.B1: vacdipole, l=0, freq:={q2_dipole}, lag=0, volt:=volty, ramp1={ramp1}, "
+        f"MKACV.6L4.B{beam:d}: vacdipole, l=0, freq:={q2_dipole}, lag=0, volt:=volty, ramp1={ramp1}, "
         f"ramp2={ramp2}, ramp3={ramp3}, ramp4={ramp4};"
     )
-    madx.command.seqedit(sequence="lhcb1")
+    madx.command.seqedit(sequence=f"lhcb{beam:d}")
     madx.command.flatten()
-    madx.command.install(element="MKACH.6L4.B1", at="0.0", from_="MKQA.6L4.B1")
-    madx.command.install(element="MKACV.6L4.B1", at="0.0", from_="MKQA.6L4.B1")
+    madx.command.install(  # same position as in model_creator macros to avoid negative drift
+        element=f"MKACH.6L4.B{beam:d}", at="1.583 / 2", from_=f"MKQA.6L4.B{beam:d}"
+    )
+    madx.command.install(  # same position as in model_creator macros to avoid negative drift
+        element=f"MKACV.6L4.B{beam:d}", at="1.583 / 2", from_=f"MKQA.6L4.B{beam:d}"
+    )
     madx.command.endedit()
 
     logger.trace("Querying BETX and BETY at AC Dipole location")
-    madx.input("betax_acd = table(twiss, MKQA.6L4.B1, betx);")
-    madx.input("betay_acd = table(twiss, MKQA.6L4.B1, bety);")
+    madx.input(f"betax_acd = table(twiss, MKQA.6L4.B{beam:d}, betx);")
+    madx.input(f"betay_acd = table(twiss, MKQA.6L4.B{beam:d}, bety);")
 
     betax_acd = madx.globals["betax_acd"]
     betay_acd = madx.globals["betay_acd"]
@@ -276,6 +289,67 @@ def install_ac_dipole(
     madx.globals["voltx"] = voltx
     madx.globals["volty"] = volty
 
+    logger.warning(
+        f"Sequence LHCB{beam:d} is now re-used for changes to take effect. Beware that this will reset it "
+        "to its default state, remove errors etc."
+    )
+    madx.use(sequence=f"lhcb{beam:d}")
+
+
+def install_ac_dipole_as_matrix(madx: Madx, deltaqx: float, deltaqy: float, beam: int = 1) -> None:
+    """
+    Installs an AC dipole as a matrix element for (HL)LHC BEAM 1 OR 2 ONLY.
+    The AC Dipole does impact the orbit as well as the betatron functions when turned on. Unfortunately in
+    MAD-X, it cannot be modeled to do both at the same time. This routine introduces an AC Dipole as a
+    matrix element so that its effect can be seen on the TWISS functions. It DOES NOT affect tracking.
+    See https://journals.aps.org/prab/abstract/10.1103/PhysRevSTAB.11.084002 (part III).
+
+    This function assumes that you have already defined lhcb1/lhcb2, made a beam for it (BEAM command or
+    `make_lhc_beams` function), matched to your desired working point and made a TWISS.
+
+    Args:
+        madx (cpymad.madx.Madx): an instanciated cpymad Madx object.
+        deltaqx (float): the deltaQx (horizontal tune excitation) used by the AC dipole.
+        deltaqy (float): the deltaQy (vertical tune excitation) used by the AC dipole.
+        beam (int): the LHC beam to install the AC Dipole into, either 1 or 2. Defaults to 1.
+    """
+    logger.warning("This AC Dipole is implemented as a matrix and will not affect particle tracking!")
+    logger.info("This routine should be done after 'match', 'twiss' and 'makethin' for the appropriate beam.")
+
+    logger.debug("Retrieving tunes from internal tables")
+    q1, q2 = madx.table.summ.q1[0], madx.table.summ.q2[0]
+    logger.trace(f"Retrieved values are q1 = {q1}, q2 = {q2}")
+    q1_dipole, q2_dipole = q1 + deltaqx, q2 + deltaqy
+
+    logger.trace("Querying BETX and BETY at AC Dipole location")
+    madx.input(f"betax_acd = table(twiss, MKQA.6L4.B{beam:d}, betx);")
+    madx.input(f"betay_acd = table(twiss, MKQA.6L4.B{beam:d}, bety);")
+    betax_acd = madx.globals["betax_acd"]
+    betay_acd = madx.globals["betay_acd"]
+
+    logger.trace("Calculating AC Dipole matrix terms")
+    hacmap21 = (
+        2 * (np.cos(2 * np.pi * q1_dipole) - np.cos(2 * np.pi * q1)) / (betax_acd * np.sin(2 * np.pi * q1))
+    )
+    vacmap43 = (
+        2 * (np.cos(2 * np.pi * q2_dipole) - np.cos(2 * np.pi * q2)) / (betay_acd * np.sin(2 * np.pi * q2))
+    )
+    madx.input(f"hacmap: matrix, l=0, rm21={hacmap21};")
+    madx.input(f"vacmap: matrix, l=0, rm43={vacmap43};")
+
+    logger.info(f"Installing AC Dipole matrix with driven tunes of Qx_D = {q1_dipole}  |  Qy_D = {q2_dipole}")
+    madx.command.seqedit(sequence=f"lhcb{beam:d}")
+    madx.command.flatten()
+    madx.command.install(element="hacmap", at="1.583 / 2", from_=f"MKQA.6L4.B{beam:d}")  # no negative drift
+    madx.command.install(element="vacmap", at="1.583 / 2", from_=f"MKQA.6L4.B{beam:d}")  # no negative drift
+    madx.command.endedit()
+
+    logger.warning(
+        f"Sequence LHCB{beam:d} is now re-used for changes to take effect. Beware that this will reset it "
+        "to its default state, remove errors etc."
+    )
+    madx.use(sequence=f"lhcb{beam:d}")
+
 
 def vary_independent_ir_quadrupoles(
     madx: Madx, quad_numbers: Sequence[int], ip: int, sides: Sequence[str] = ("r", "l"), beam: int = 1

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "pyhdtoolkit"
-version = "0.12.0"
+version = "0.13.0"
 description = "An all-in-one toolkit package to easy my Python work in my PhD."
 authors = ["Felix Soubelet <[email protected]>"]
 license = "MIT"

tests/test_cpymadtools.py

@@ -54,7 +54,8 @@
     apply_lhc_coupling_knob,
     apply_lhc_rigidity_waist_shift_knob,
     deactivate_lhc_arc_sextupoles,
-    install_ac_dipole,
+    install_ac_dipole_as_kicker,
+    install_ac_dipole_as_matrix,
     make_lhc_beams,
     make_lhc_thin,
     make_sixtrack_output,
@@ -816,9 +817,9 @@ def test_make_lhc_beams(self, energy, _bare_lhc_madx):
         assert madx.sequence.lhcb2.beam.energy == energy
 
     @pytest.mark.parametrize("top_turns", [1000, 6000, 10_000])
-    def test_install_ac_dipole(self, top_turns, _matched_lhc_madx):
+    def test_install_ac_dipole_as_kicker(self, top_turns, _matched_lhc_madx):
         madx = _matched_lhc_madx
-        install_ac_dipole(madx, deltaqx=-0.01, deltaqy=0.012, sigma_x=3, sigma_y=3, top_turns=top_turns)
+        install_ac_dipole_as_kicker(madx, deltaqx=-0.01, deltaqy=0.012, sigma_x=3, sigma_y=3, top_turns=top_turns)
         ramp3 = 2100 + top_turns
         ramp4 = ramp3 + 2000
 
@@ -831,6 +832,19 @@ def test_install_ac_dipole(self, top_turns, _matched_lhc_madx):
         assert math.isclose(madx.elements["MKACH.6L4.B1"].at, 9846.0765, rel_tol=1e-2)
         assert math.isclose(madx.elements["MKACH.6L4.B1"].freq, 62.3, rel_tol=1e-2)
 
+    def test_install_ac_dipole_matrix(self, _matched_lhc_madx):
+        madx = _matched_lhc_madx
+        twiss_before = madx.twiss().dframe().copy()
+        install_ac_dipole_as_matrix(madx, deltaqx=-0.01, deltaqy=0.012)
+        twiss_after = madx.twiss().dframe().copy()
+
+        for acd_name in ("hacmap", "vacmap"):
+            assert acd_name in madx.elements
+            assert math.isclose(madx.elements[acd_name].at, 9846.0765, rel_tol=1e-2)
+
+        with pytest.raises(AssertionError):
+            assert_frame_equal(twiss_before, twiss_after)  # they should be different!
+
     def test_makethin_lhc(self, _matched_lhc_madx):
         """
         Little trick: if we haven't sliced properly, tracking will fail so we can check all is ok by