"""
.. _lhc-queries:
**Querying Utilities**
The functions below are settings query utilities for the ``LHC``.
"""
from typing import Dict, Sequence, Union
import tfs
from cpymad.madx import Madx
from loguru import logger
from pyhdtoolkit.cpymadtools import twiss
from pyhdtoolkit.cpymadtools.constants import (
LHC_KCD_KNOBS,
LHC_KCO_KNOBS,
LHC_KCOSX_KNOBS,
LHC_KCOX_KNOBS,
LHC_KCS_KNOBS,
LHC_KCSSX_KNOBS,
LHC_KCSX_KNOBS,
LHC_KCTX_KNOBS,
LHC_KO_KNOBS,
LHC_KQS_KNOBS,
LHC_KQSX_KNOBS,
LHC_KQTF_KNOBS,
LHC_KSF_KNOBS,
LHC_KSS_KNOBS,
)
from pyhdtoolkit.cpymadtools.lhc._setup import lhc_orbit_variables
from pyhdtoolkit.cpymadtools.utils import _get_k_strings
[docs]def get_magnets_powering(
madx: Madx, /, patterns: Sequence[str] = [r"^mb\.", r"^mq\.", r"^ms\."], brho: Union[str, float] = None, **kwargs
) -> tfs.TfsDataFrame:
r"""
.. versionadded:: 0.17.0
Gets the twiss table with additional defined columns for the given *patterns*.
.. note::
Here are below certain useful patterns for the ``LHC`` and their meaning:
* ``^mb\.`` :math:`\rightarrow` main bends.
* ``^mq\.`` :math:`\rightarrow` main quadrupoles.
* ``^ms\.`` :math:`\rightarrow` main sextupoles.
* ``^mb[rswx]`` :math:`\rightarrow` separation dipoles.
* ``^mq[mwxy]`` :math:`\rightarrow` insertion quads.
* ``^mqt.1[23]`` :math:`\rightarrow` short tuning quads (12 & 13).
* ``^mqtl`` :math:`\rightarrow` long tuning quads.
* ``^mcbx`` :math:`\rightarrow` crossing scheme magnets.
* ``^mcb[cy]`` :math:`\rightarrow` crossing scheme magnets.
To make no selection, one can give ``patterns=[""]`` and this will give back
the results for *all* elements. One can also give a specific magnet's exact
name to include it in the results.
.. note::
The ``TWISS`` flag will be fully cleared after running this function.
Args:
madx (cpymad.madx.Madx): an instanciated `~cpymad.madx.Madx` object. Positional only.
patterns (Sequence[str]): a list of regex patterns to define which elements
should be selected and included in the returned table. Defaults to selecting
the main bends, quads and sextupoles. See the note admonition above for
useful patterns to select specific ``LHC`` magnet families.
brho (Union[str, float]): optional, an explicit definition for the magnetic
rigidity in :math:`Tm^{-1}`. If not given, it will be assumed that
a ``brho`` quantity is defined in the ``MAD-X`` globals.
**kwargs: any keyword argument will be passed to `~.twiss.get_pattern_twiss` and
later on to the ``TWISS`` command executed in ``MAD-X``.
Returns:
A `~tfs.TfsDataFrame` of the ``TWISS`` table, with the relevant newly defined columns
and including the elements matching the regex *patterns* that were provided.
Example:
.. code-block:: python
sextupoles_powering = get_magnets_powering(madx, patterns=[r"^ms\."])
"""
logger.debug("Computing magnets field and powering limits proportions")
NEW_COLNAMES = ["name", "keyword", "ampere", "imax", "percent", "kn", "kmax", "integrated_field", "L"]
NEW_COLNAMES = list(set(NEW_COLNAMES + kwargs.pop("columns", []))) # in case user gives explicit columns
_list_field_currents(madx, brho=brho)
return twiss.get_pattern_twiss(madx, columns=NEW_COLNAMES, patterns=patterns, **kwargs)
[docs]def query_arc_correctors_powering(madx: Madx, /) -> Dict[str, float]:
"""
.. versionadded:: 0.15.0
Queries for the arc corrector strengths and returns their values as a percentage of
their max powering. This is a port of one of the **corr_value.madx** file's macros
Args:
madx (cpymad.madx.Madx): an instanciated `~cpymad.madx.Madx` object with an
active (HL)LHC sequence.
Returns:
A `dict` with the percentage for each corrector.
Example:
.. code-block:: python
arc_knobs = query_arc_correctors_powering(madx)
"""
logger.debug("Querying triplets correctors powering")
result: Dict[str, float] = {}
logger.debug("Querying arc tune trim quadrupole correctors (MQTs) powering")
k_mqt_max = 120 / madx.globals.brho # 120 T/m
result.update({knob: 100 * _knob_value(madx, knob) / k_mqt_max for knob in LHC_KQTF_KNOBS})
logger.debug("Querying arc short straight sections skew quadrupole correctors (MQSs) powering")
k_mqs_max = 120 / madx.globals.brho # 120 T/m
result.update({knob: 100 * _knob_value(madx, knob) / k_mqs_max for knob in LHC_KQS_KNOBS})
logger.debug("Querying arc sextupole correctors (MSs) powering")
k_ms_max = 1.280 * 2 / 0.017**2 / madx.globals.brho # 1.28 T @ 17 mm
result.update({knob: 100 * _knob_value(madx, knob) / k_ms_max for knob in LHC_KSF_KNOBS})
logger.debug("Querying arc skew sextupole correctors (MSSs) powering")
k_mss_max = 1.280 * 2 / 0.017**2 / madx.globals.brho # 1.28 T @ 17 mm
result.update({knob: 100 * _knob_value(madx, knob) / k_mss_max for knob in LHC_KSS_KNOBS})
logger.debug("Querying arc spool piece (skew) sextupole correctors (MCSs) powering")
k_mcs_max = 0.471 * 2 / 0.017**2 / madx.globals.brho # 0.471 T @ 17 mm
result.update({knob: 100 * _knob_value(madx, knob) / k_mcs_max for knob in LHC_KCS_KNOBS})
logger.debug("Querying arc spool piece (skew) octupole correctors (MCOs) powering")
k_mco_max = 0.040 * 6 / 0.017**3 / madx.globals.brho # 0.04 T @ 17 mm
result.update({knob: 100 * _knob_value(madx, knob) / k_mco_max for knob in LHC_KCO_KNOBS})
logger.debug("Querying arc spool piece (skew) decapole correctors (MCDs) powering")
k_mcd_max = 0.100 * 24 / 0.017**4 / madx.globals.brho # 0.1 T @ 17 mm
result.update({knob: 100 * _knob_value(madx, knob) / k_mcd_max for knob in LHC_KCD_KNOBS})
logger.debug("Querying arc short straight sections octupole correctors (MOs) powering")
k_mo_max = 0.29 * 6 / 0.017**3 / madx.globals.brho # 0.29 T @ 17 mm
result.update({knob: 100 * _knob_value(madx, knob) / k_mo_max for knob in LHC_KO_KNOBS})
return result
[docs]def query_triplet_correctors_powering(madx: Madx, /) -> Dict[str, float]:
"""
.. versionadded:: 0.15.0
Queries for the triplet corrector strengths and returns their values as a percentage of
their max powering. This is a port of one of the **corr_value.madx** file's macros.
Args:
madx (cpymad.madx.Madx): an instanciated `~cpymad.madx.Madx` object with an
active (HL)LHC sequence.
Returns:
A `dict` with the percentage for each corrector.
Example:
.. code-block:: python
triplet_knobs = query_triplet_correctors_powering(madx)
"""
logger.debug("Querying triplets correctors powering")
result: Dict[str, float] = {}
logger.debug("Querying triplet skew quadrupole correctors (MQSXs) powering")
k_mqsx_max = 1.360 / 0.017 / madx.globals.brho # 1.36 T @ 17mm
result.update({knob: 100 * _knob_value(madx, knob) / k_mqsx_max for knob in LHC_KQSX_KNOBS})
logger.debug("Querying triplet sextupole correctors (MCSXs) powering")
k_mcsx_max = 0.028 * 2 / 0.017**2 / madx.globals.brho # 0.028 T @ 17 mm
result.update({knob: 100 * _knob_value(madx, knob) / k_mcsx_max for knob in LHC_KCSX_KNOBS})
logger.debug("Querying triplet skew sextupole correctors (MCSSXs) powering")
k_mcssx_max = 0.11 * 2 / 0.017**2 / madx.globals.brho # 0.11 T @ 17 mm
result.update({knob: 100 * _knob_value(madx, knob) / k_mcssx_max for knob in LHC_KCSSX_KNOBS})
logger.debug("Querying triplet octupole correctors (MCOXs) powering")
k_mcox_max = 0.045 * 6 / 0.017**3 / madx.globals.brho # 0.045 T @ 17 mm
result.update({knob: 100 * _knob_value(madx, knob) / k_mcox_max for knob in LHC_KCOX_KNOBS})
logger.debug("Querying triplet skew octupole correctors (MCOSXs) powering")
k_mcosx_max = 0.048 * 6 / 0.017**3 / madx.globals.brho # 0.048 T @ 17 mm
result.update({knob: 100 * _knob_value(madx, knob) / k_mcosx_max for knob in LHC_KCOSX_KNOBS})
logger.debug("Querying triplet decapole correctors (MCTXs) powering")
k_mctx_max = 0.01 * 120 / 0.017**5 / madx.globals.brho # 0.010 T @ 17 mm
result.update({knob: 100 * _knob_value(madx, knob) / k_mctx_max for knob in LHC_KCTX_KNOBS})
return result
[docs]def get_current_orbit_setup(madx: Madx, /) -> Dict[str, float]:
"""
.. versionadded:: 0.8.0
Get the current values for the (HL)LHC orbit variables. Initial implementation credits go to
:user:`Joschua Dilly <joschd>`.
Args:
madx (cpymad.madx.Madx): an instanciated `~cpymad.madx.Madx` object. Positional only.
Returns:
A `dict` of all orbit variables set, and their values as set in the ``MAD-X`` globals.
Example:
.. code-block:: python
orbit_setup = get_current_orbit_setup(madx)
"""
logger.debug("Extracting orbit variables from global table")
variables, specials = lhc_orbit_variables()
return {orbit_variable: madx.globals[orbit_variable] for orbit_variable in variables + list(specials.keys())}
# ----- Helpers ----- #
def _list_field_currents(madx: Madx, /, brho: Union[str, float] = None) -> None:
"""
Creates additional columns for the ``TWISS`` table with the magnets' total fields
and currents, to help later on determine which proportion of their maximum powering
the current setting is using. This is an implementation of the old utility script
located at **/afs/cern.ch/eng/lhc/optics/V6.503/toolkit/list_fields_currents.madx**.
.. important::
Certain quantities are assumed to be defined in the ``MAD-X`` globals, such as
``brho``, or available in the magnets definition, such as ``calib``. For this
reason, this script most likely only works for the ``(HL)LHC`` sequences where
those are defined.
Args:
madx (cpymad.madx.Madx): an instanciated `~cpymad.madx.Madx` object. Positional only.
brho (Union[str, float]): optional, an explicit definition for the magnetic
rigidity in :math:`Tm^{-1}`. If not given, it will be assumed that
a ``brho`` quantity is defined in the ``MAD-X`` globals and this one will
be used.
"""
logger.debug("Creating additional TWISS table columns for magnets' fields and currents")
if brho is not None:
logger.trace(f"Setting 'brho' to explicitely defined '{brho}'")
madx.globals["brho"] = brho
# Define strength := table(twiss, k0l) + ... + table(twiss, k5sl) + table(twiss, hkick) + table(twiss, vkick);
madx.globals["strength"] = (
" + ".join(f"table(twiss, {a.lower()})" for a in _get_k_strings(stop=6))
+ " + table(twiss, hkick) + table(twiss, vkick)"
)
# All here are given as strings to make it deferred expressions in MAD-X
madx.globals["epsilon"] = 1e-20 # to avoid divisions by zero
madx.globals["length"] = "table(twiss, l) + table(twiss, lrad) + epsilon"
madx.globals["kmaxx"] = "table(twiss, kmax) + epsilon"
madx.globals["calibration"] = "table(twiss, calib) + epsilon"
madx.globals["kn"] = "abs(strength) / length"
# madx.globals["rho"] = "kn / (kn + epsilon) / (kn + epsilon)"
madx.globals["field"] = "kn * brho"
madx.globals["percent"] = "field * 100 / (kmaxx + epsilon)"
madx.globals["ampere"] = "field / calibration"
madx.globals["imax"] = "kmaxx / calibration"
madx.globals["integrated_field"] = "field * length"
def _knob_value(madx: Madx, /, knob: str) -> float:
"""
Queryies the current value of a given *knob* name in the ``MAD-X`` process, and defaults
to 0 (as ``MAD-X`` does) in case that knob has not been defined in the current process.
Args:
madx (cpymad.madx.Madx): an instanciated `~cpymad.madx.Madx` object. Positional only.
knob (str): the name the knob.
Returns:
The knob value if it was defined, otherwise 0.
Example:
.. code-block:: python
_knob_value(madx, knob="underfined_for_sure") # returns 0
"""
try:
return madx.globals[knob]
except KeyError: # cpymad gives a 'Variable not defined: var_name'
return 0