"""
.. _cpymadtools-ptc:
PTC Routines
------------
Module with functions to manipulate ``MAD-X`` ``PTC`` functionality through a
`~cpymad.madx.Madx` object.
"""
from logging import getLogger
from pathlib import Path
from typing import Dict, Optional, Sequence, Tuple, Union
import pandas as pd
import tfs
from cpymad.madx import Madx
from cpymadtools.utils import get_table_tfs
logger = getLogger(__name__)
[docs]def get_amplitude_detuning(
madx: Madx, order: int = 2, file: Union[Path, str] = None, fringe: bool = False, **kwargs
) -> tfs.TfsDataFrame:
"""
.. versionadded:: 0.7.0
Calculates amplitude detuning coefficients via ``PTC_NORMAL``, with sensible defaults set for
other relevant ``PTC`` commands used in the process. The result table is returned as a
`~tfs.frame.TfsDataFrame`, the headers of which are the contents of the internal ``SUMM`` table.
This is a heavily refactored version of an initial implementation by :user:`Joschua Dilly <joschd>`.
.. important::
The ``PTC_CREATE_LAYOUT`` command is issued with ``model=3`` (``SixTrack`` model), ``method=4``
(integration order), ``nst=3`` (number of integration steps, aka body slices for elements) and
``exact=True`` (use exact Hamiltonian, not an approximated one).
The ``PTC_NORMAL`` command is explicitely given ``icase=6`` to enforce 6D calculations (see the
`MAD-X manual <http://madx.web.cern.ch/madx/releases/last-rel/madxuguide.pdf>`_ for details),
``no=5`` (map order for derivative evaluation of Twiss parameters), ``closedorbit=True`` (triggers
closed orbit calculation) and ``normal=True`` (activate calculation of the Normal Form).
Args:
madx (cpymad.madx.Madx): an instanciated `~cpymad.madx.Madx` object.
order (int): maximum derivative order coefficient (only 0, 1 or 2 implemented in ``PTC``).
Defaults to 2.
file (Union[Path, str]): path to output file. Defaults to `None`.
fringe (bool): boolean flag to include fringe field effects in the calculation. Defaults to
``False``.
**kwargs: any keyword argument is transmitted to the ``PTC_NORMAL`` command. See above which
arguments are already set for ``PTC_NORMAL`` to avoid trying to override them.
Returns:
A `~tfs.frame.TfsDataframe` with the calculated coefficients.
Example:
.. code-block:: python
>>> ampdet_coeffs = get_amplitude_detuning(madx, order=2, closedorbit=True)
"""
if order >= 3:
logger.error(f"Maximum amplitude detuning order in PTC is 2, but {order:d} was requested")
raise NotImplementedError("PTC amplitude detuning is not implemented for order > 2")
logger.debug("Creating PTC universe")
madx.ptc_create_universe()
logger.debug("Creating PTC layout")
madx.ptc_create_layout(model=3, method=4, nst=3, exact=True)
logger.debug("Incorporating MAD-X alignment errors")
madx.ptc_align() # use madx alignment errors
madx.ptc_setswitch(fringe=fringe)
logger.debug("Selecting tune orders")
madx.select_ptc_normal(q1="0", q2="0")
for ii in range(1, order + 1): # These are d^iQ/ddp^i
madx.select_ptc_normal(dq1=f"{ii:d}", dq2=f"{ii:d}")
# ANH = anharmonicities (ex, ey, deltap), works only with parameters as full strings
# could be done nicer with permutations ...
logger.debug("Selecting anharmonicities")
if order >= 1:
# madx.select_ptc_normal('anhx=0, 0, 1') # dQx/ddp
# madx.select_ptc_normal('anhy=0, 0, 1') # dQy/ddp
madx.select_ptc_normal("anhx=1, 0, 0") # dQx/dex
madx.select_ptc_normal("anhx=0, 1, 0") # dQx/dey
madx.select_ptc_normal("anhy=1, 0, 0") # dQy/dex
madx.select_ptc_normal("anhy=0, 1, 0") # dQy/dey
if order >= 2:
# madx.select_ptc_normal('anhx=0, 0, 2') # d^2Qx/ddp^2
# madx.select_ptc_normal('anhy=0, 0, 2') # d^2Qy/ddp^2
madx.select_ptc_normal("anhx=2, 0, 0") # d^2Qx/dex^2
madx.select_ptc_normal("anhx=1, 1, 0") # d^2Qx/dexdey
madx.select_ptc_normal("anhx=0, 2, 0") # d^2Qx/dey^2
madx.select_ptc_normal("anhy=2, 0, 0") # d^2Qy/dex^2
madx.select_ptc_normal("anhy=1, 1, 0") # d^2Qy/dexdey
madx.select_ptc_normal("anhy=0, 2, 0") # d^2Qy/dey^2
logger.debug("Executing PTC Normal")
madx.ptc_normal(icase=6, no=5, closed_orbit=True, normal=True, **kwargs)
madx.ptc_end()
dframe = get_table_tfs(madx, table_name="normal_results")
dframe.index = range(len(dframe.NAME)) # table has a weird index
if file:
logger.debug(f"Exporting results to disk at '{Path(file).absolute()}'")
tfs.write(file, dframe)
return dframe
[docs]def get_rdts(
madx: Madx, order: int = 4, file: Union[Path, str] = None, fringe: bool = False, **kwargs
) -> tfs.TfsDataFrame:
"""
.. versionadded:: 0.7.0
Calculate the resonance driving terms up to *order* via ``PTC_TWISS``, with sensible defaults
set for other relevant ``PTC`` commands. The result table is returned as a `~tfs.frame.TfsDataFrame`,
the headers of which are the contents of the internal ``SUMM`` table. This is a heavily refactored
version of an initial implementation by :user:`Joschua Dilly <joschd>`.
.. important::
The ``PTC_CREATE_LAYOUT`` command is issued with ``model=3`` (``SixTrack`` model), ``method=4``
(integration order), ``nst=3`` (number of integration steps, aka body slices for elements) and
``exact=True`` (use exact Hamiltonian, not an approximated one).
The ``PTC_NORMAL`` command is explicitely given ``icase=6`` to enforce 6D calculations (see the
`MAD-X manual <http://madx.web.cern.ch/madx/releases/last-rel/madxuguide.pdf>`_ for details),
``trackrdts=True`` (for this function to fullfill its purpose) and ``normal=True`` to trigger
saving the normal form analysis results in a table called ``NONLIN`` which will then be available
through the provided `~cpymad.madx.Madx` instance.
Args:
madx (cpymad.madx.Madx): an instanciated `~cpymad.madx.Madx` object.
order (int): map order for derivative evaluation of Twiss parameters. Defaults to 4.
file (Union[Path, str]): path to output file. Default to `None`.
fringe (bool): boolean flag to include fringe field effects in the calculation. Defaults to `False`.
**kwargs: any keyword argument is transmitted to the ``PTC_TWISS`` command. See above which arguments
are already set for ``PTC_TWISS`` to avoid trying to override them.
Returns:
A `TfsDataframe` with the calculated RDTs.
Example:
.. code-block:: python
>>> rdts_df = get_rdts(madx, order=3, fringe=True)
"""
logger.debug("Creating PTC universe")
madx.ptc_create_universe()
logger.debug("Creating PTC layout")
madx.ptc_create_layout(model=3, method=4, nst=3, exact=True)
logger.debug("Incorporating MAD-X alignment errors")
madx.ptc_align() # use madx alignment errors
madx.ptc_setswitch(fringe=fringe)
logger.debug("Executing PTC Twiss")
madx.ptc_twiss(icase=6, no=order, normal=True, trackrdts=True, **kwargs)
madx.ptc_end()
dframe = get_table_tfs(madx, table_name="twissrdt", headers_table="ptc_twiss_summary")
if file:
logger.debug(f"Exporting results to disk at '{Path(file).absolute()}'")
tfs.write(file, dframe)
return dframe
[docs]def ptc_twiss(
madx: Madx, order: int = 4, file: Union[Path, str] = None, fringe: bool = False, table: str = "ptc_twiss", **kwargs
) -> tfs.TfsDataFrame:
"""
.. versionadded:: 0.12.0
Calculates the ``TWISS`` parameters according to the :cite:t:`Ripken:optics:1989` formalism
via ``PTC_TWISS``, with sensible defaults set for other relevant ``PTC`` commands. The result
table is returned as a `~tfs.frame.TfsDataFrame`, the headers of which are the contents of the
internal ``SUMM`` table.
This is very similar to the `~.ptc.get_rdts` function as both use ``PTC_TWISS`` internally,
however this function does not track RDTs which makes the calculations significantly faster.
.. important::
The ``PTC_CREATE_LAYOUT`` command is issued with ``model=3`` (``SixTrack`` model), ``method=4``
(integration order), ``nst=3`` (number of integration steps, aka body slices for elements) and
``exact=True`` (use exact Hamiltonian, not an approximated one).
The ``PTC_TWISS`` command is explicitely given ``icase=6`` to enforce 6D calculations (see the
`MAD-X manual <http://madx.web.cern.ch/madx/releases/last-rel/madxuguide.pdf>`_ for details),
``normal=True`` to trigger saving the normal form analysis results in a table called ``NONLIN``
which will then be available through the provided `~cpymad.madx.Madx` instance.
Args:
madx (cpymad.madx.Madx): an instanciated `~cpymad.madx.Madx` object.
order (int): map order for derivative evaluation of ``TWISS`` parameters. Defaults to 4.
file (Union[Path, str]): path to output file. Default to `None`.
fringe (bool): boolean flag to include fringe field effects in the calculation. Defaults to `False`.
table (str): the name of the internal table in which to save the results. Defaults to **ptc_twiss**.
**kwargs: Any keyword argument is transmitted to the ``PTC_TWISS`` command. See above which arguments
are already set for `PTC_TWISS` to avoid trying to override them.
Returns:
A `TfsDataframe` with the calculated ``TWISS`` parameters.
Example:
.. code-block:: python
>>> twiss_ptc_df = ptc_twiss(madx, order=3)
"""
logger.debug(f"Creating PTC universe")
madx.ptc_create_universe()
logger.debug("Creating PTC layout")
madx.ptc_create_layout(model=3, method=4, nst=3, exact=True)
logger.debug("Incorporating MAD-X alignment errors")
madx.ptc_align() # use madx alignment errors
madx.ptc_setswitch(fringe=fringe)
logger.debug("Executing PTC Twiss")
madx.ptc_twiss(icase=6, no=order, normal=True, table=table, **kwargs)
madx.ptc_end()
dframe = get_table_tfs(madx, table_name=table, headers_table="ptc_twiss_summary")
if file:
logger.debug(f"Exporting results to disk at '{Path(file).absolute()}'")
tfs.write(file, dframe)
return dframe
[docs]def ptc_track_particle(
madx: Madx,
initial_coordinates: Tuple[float, float, float, float, float, float],
nturns: int,
sequence: Optional[str] = None,
observation_points: Sequence[str] = None,
onetable: bool = False,
fringe: bool = False,
**kwargs,
) -> Dict[str, pd.DataFrame]:
"""
.. versionadded:: 0.12.0
Tracks a single particle for *nturns* through ``PTC_TRACK``, based on its initial coordinates. The
use of this function is similar to that of `~.track.track_single_particle`.
.. important::
The ``PTC_CREATE_LAYOUT`` command is issued with ``model=3`` (``SixTrack`` model), ``method=4``
(integration order), ``nst=3`` (number of integration steps, aka body slices for elements) and
``exact=True`` (use exact Hamiltonian, not an approximated one).
The ``PTC_TRACK`` command is explicitely given ``ELEMENT_BY_ELEMENT=True`` to force element by
element tracking mode.
.. warning::
If the *sequence* argument is given a string value, the ``USE`` command will be ran on the provided
sequence name. This means the caveats of ``USE`` apply, for instance the erasing of previously
defined errors, orbits corrections etc. In this case a warning will be logged but the function will
proceed. If `None` is given (by default) then the sequence already in use will be the one tracking
is performed on.
Args:
madx (cpymad.madx.Madx): an instantiated cpymad.madx.Madx object.
initial_coordinates (Tuple[float, float, float, float, float, float]): a tuple with the ``X, PX,
Y, PY, T, PT`` starting coordinates of the particle to track. Defaults to all 0 if `None` given.
nturns (int): the number of turns to track for.
sequence (Optional[str]): the sequence to use for tracking. If no value is provided, it is assumed
that a sequence is already defined and in use, and this one will be picked up by ``MAD-X``.
Beware of the dangers of giving a sequence that will be used by ``MAD-X``, see the warning below
for more information.
observation_points (Sequence[str]): sequence of all element names at which to ``OBSERVE`` during the
tracking.
onetable (bool): flag to combine all observation points data into a single table. Defaults to `False`.
fringe (bool): boolean flag to include fringe field effects in the calculation. Defaults to `False`.
**kwargs: Any keyword argument is transmitted to the ``PTC_TRACK`` command such as the ``CLOSED_ORBIT``
flag to activate closed orbit calculation before tracking. Refer to the
`MAD-X manual <http://madx.web.cern.ch/madx/releases/last-rel/madxuguide.pdf>`_ for options.
Returns:
A `dict` with a copy of the track table's dataframe for each defined observation point,
with as columns the coordinates ``x, px, y, py, t, pt, s and e`` (energy). The keys of the dictionary
are simply named ``observation_point_1``, ``observation_point_2`` etc. The first observation point
always corresponds to the start of machine, the others correspond to the ones manually defined,
in the order they are defined in.
If the user has set ``onetable`` to `True`, only one entry is in the dictionary under the key
``trackone`` and it has the combined table as a pandas DataFrame for value.
Example:
.. code-block:: python
>>> tracks_dict = ptc_track_particle(
... madx, nturns=1023, initial_coordinates=(2e-4, 0, 1e-4, 0, 0, 0)
... )
"""
logger.debug("Performing single particle PTC (thick) tracking")
start = initial_coordinates if initial_coordinates else [0, 0, 0, 0, 0, 0]
observation_points = observation_points if observation_points else []
if isinstance(sequence, str):
logger.warning(f"Sequence '{sequence}' was provided and will be USEd, beware that this will erase errors etc.")
logger.debug(f"Using sequence '{sequence}' for tracking")
madx.use(sequence=sequence)
logger.debug(f"Creating PTC universe")
madx.ptc_create_universe()
logger.debug("Creating PTC layout")
madx.ptc_create_layout(model=3, method=4, nst=3, exact=True)
logger.debug("Incorporating MAD-X alignment errors")
madx.ptc_align() # use madx alignment errors
madx.ptc_setswitch(fringe=fringe)
logger.debug(f"Tracking coordinates with initial X, PX, Y, PY, T, PT of '{initial_coordinates}'")
madx.command.ptc_start(X=start[0], PX=start[1], Y=start[2], PY=start[3], T=start[4], PT=start[5])
for element in observation_points:
logger.debug(f"Setting observation point for tracking with OBSERVE at element '{element}'")
madx.command.ptc_observe(place=element)
madx.command.ptc_track(turns=nturns, element_by_element=True, onetable=onetable, **kwargs)
madx.ptc_end()
if onetable: # user asked for ONETABLE, there will only be one table 'trackone' given back by MAD-X
logger.debug("Because of option ONETABLE only one table 'TRACKONE' exists to be returned.")
return {"trackone": madx.table.trackone.dframe().copy()}
return {
f"observation_point_{point:d}": madx.table[f"track.obs{point:04d}.p0001"].dframe().copy()
for point in range(1, len(observation_points) + 2) # len(observation_points) + 1 for start of
# machine + 1 because MAD-X starts indexing these at 1
}