Improve cpymadtools.ptc module functionality (#46)

Docstrings now detail inner workings and parameters used for functions in cpymadtools.ptc.
Both functions now have a fringe argument and transmit kwargs to either ptc_normal or ptc_twiss depending on which is used.

Adds a cpymadtools.utils module with convenient functionality to export an internal MAD-X table

pyhdtoolkit/cpymadtools/__init__.py

@@ -31,3 +31,4 @@
 from .track import track_single_particle
 from .tune import make_footprint_table
 from .twiss import get_ips_twiss, get_ir_twiss, get_twiss_tfs
+from .utils import get_table_tfs

pyhdtoolkit/cpymadtools/ptc.py

@@ -15,21 +15,42 @@
 from cpymad.madx import Madx
 from loguru import logger
 
+from pyhdtoolkit.cpymadtools.utils import get_table_tfs
+
 # ----- Utilities ----- #
 
 
-def get_amplitude_detuning(madx: Madx, order: int = 2, file: Union[Path, str] = None) -> tfs.TfsDataFrame:
+def get_amplitude_detuning(
+    madx: Madx, order: int = 2, file: Union[Path, str] = None, fringe: bool = False, **kwargs
+) -> tfs.TfsDataFrame:
     """
-    INITIAL IMPLEMENTATION CREDITS GO TO JOSCHUA DILLY (@JoschD).
-    Calculate amplitude detuning via PTC_NORMAL.
+    INITIAL IMPLEMENTATION CREDITS GO TO JOSCHUA DILLY (@JoschD), but this has been heavily refactored.
+    Calculate amplitude detuning via `PTC_NORMAL`, with sensible defaults set for other relevant `PTC`
+    commands. The result table is returned as a `TfsDataFrame`, the headers of which are the contents of the
+    internal `SUMM` table.
+
+    The `PTC_CREATE_LAYOUT` command is issued with `model=3` (`SixTrack` model), `method=4` (integration
+    order), `nst=3` (number of integratioin 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 for details), `no=5` (map order for derivative evaluation of Twiss parameters) and
+    `closedorbit=True` (triggers closed orbit calculation) and `normal=True` (activate calculation of the
+    Normal Form).
 
     Args:
         madx (cpymad.madx.Madx): an instanciated cpymad Madx object.
-        order (int): maximum derivative order (only 0, 1 or 2 implemented in PTC). Defaults to `2`.
+        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. Default `None`
+        fringe (bool): boolean flag to include fringe field effects in the calculation. Defaults to `False`.
+
+    Keyword Args:
+        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 TfsDataframe with results.
+        A `TfsDataframe` with results.
     """
     if order >= 3:
         logger.error(f"Maximum amplitude detuning order in PTC is 2, but {order:d} was requested")
@@ -38,20 +59,12 @@ def get_amplitude_detuning(madx: Madx, order: int = 2, file: Union[Path, str] =
     logger.info("Entering PTC to calculate amplitude detuning")
     madx.ptc_create_universe()
 
-    # layout I got with mask (jdilly)
-    # model=3 (Sixtrack code model: Delta-Matrix-Kick-Matrix)
-    # method=4 (integration order), nst=3 (integration steps), exact=True (exact Hamiltonian)
     logger.trace("Creating PTC layout")
     madx.ptc_create_layout(model=3, method=4, nst=3, exact=True)
 
-    # alternative layout: model=3, method=6, nst=3
-    # resplit=True (adaptive splitting of magnets), thin=0.0005 (splitting of quads),
-    # xbend=0.0005 (splitting of dipoles)
-    # madx.ptc_create_layout(model=3, method=6, nst=3, resplit=True, thin=0.0005, xbend=0.0005)
-
     logger.trace("Incorporating MAD-X alignment errors")
     madx.ptc_align()  # use madx alignment errors
-    # madx.ptc_setswitch(fringe=True)  # include fringe effects
+    madx.ptc_setswitch(fringe=fringe)
 
     logger.trace("Selecting tune orders")
     madx.select_ptc_normal(q1="0", q2="0")
@@ -79,15 +92,11 @@ def get_amplitude_detuning(madx: Madx, order: int = 2, file: Union[Path, str] =
         madx.select_ptc_normal("anhy=1, 1, 0")  # d^2Qy/dexdey
         madx.select_ptc_normal("anhy=0, 2, 0")  # d^2Qy/dey^2
 
-    # icase = phase-space dimensionality, no = order of map
     logger.debug("Executing PTC Normal")
-    madx.ptc_normal(closed_orbit=True, normal=True, icase=5, no=5)
+    madx.ptc_normal(icase=6, no=5, closed_orbit=True, normal=True, **kwargs)
     madx.ptc_end()
 
-    logger.debug("Extracting results to TfsDataFrame")
-    dframe = tfs.TfsDataFrame(madx.table.normal_results.dframe())
-    dframe.columns = dframe.columns.str.upper()
-    dframe.NAME = dframe.NAME.str.upper()
+    dframe = get_table_tfs(madx, table_name="normal_results")
     dframe.index = range(len(dframe.NAME))  # table has a weird index
 
     if file:
@@ -97,38 +106,52 @@ def get_amplitude_detuning(madx: Madx, order: int = 2, file: Union[Path, str] =
     return dframe
 
 
-def get_rdts(madx: Madx, order: int = 4, file: Union[Path, str] = None) -> tfs.TfsDataFrame:
+def get_rdts(
+    madx: Madx, order: int = 4, file: Union[Path, str] = None, fringe: bool = False, **kwargs
+) -> tfs.TfsDataFrame:
     """
-    INITIAL IMPLEMENTATION CREDITS GO TO JOSCHUA DILLY (@JoschD).
-    Calculate the RDTs via PTC_TWISS.
+    INITIAL IMPLEMENTATION CREDITS GO TO JOSCHUA DILLY (@JoschD), but this has been heavily refactored.
+    Calculate the `RDTs` via `PTC_TWISS`, with sensible defaults set for other relevant `PTC` commands.
+    The result table is returned as a `TfsDataFrame`, the headers of which are the contents of the
+    internal `SUMM` table.
+
+    The `PTC_CREATE_LAYOUT` command is issued with `model=3` (`SixTrack` model), `method=4` (integration
+    order), `nst=3` (number of integratioin 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 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 `Madx` instance.
 
     Args:
         madx (cpymad.madx.Madx): an instanciated cpymad Madx object.
-        order (int): maximum derivative order (only 0, 1 or 2 implemented in PTC). Defaults to `2`.
+        order (int): map order for derivative evaluation of Twiss parameters. Defaults to `4`.
         file (Union[Path, str]): path to output file. Default `None`
+        fringe (bool): boolean flag to include fringe field effects in the calculation. Defaults to `False`.
+
+    Keyword Args:
+        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 results.
+        A `TfsDataframe` with results.
     """
     logger.info(f"Entering PTC to calculate RDTs up to order {order}")
     madx.ptc_create_universe()
 
     logger.trace("Creating PTC layout")
     madx.ptc_create_layout(model=3, method=4, nst=3, exact=True)
-    # madx.ptc_create_layout(model=3, method=6, nst=1)  # from Michi
 
     logger.trace("Incorporating MAD-X alignment errors")
     madx.ptc_align()  # use madx alignment errors
-    # madx.ptc_setswitch(fringe=True)  # include fringe effects
+    madx.ptc_setswitch(fringe=fringe)
 
     logger.debug("Executing PTC Twiss")
-    madx.ptc_twiss(icase=6, no=order, normal=True, trackrdts=True)
+    madx.ptc_twiss(icase=6, no=order, normal=True, trackrdts=True, **kwargs)
     madx.ptc_end()
 
-    logger.debug("Extracting results to TfsDataFrame")
-    dframe = tfs.TfsDataFrame(madx.table.twissrdt.dframe())
-    dframe.columns = dframe.columns.str.upper()
-    dframe.NAME = dframe.NAME.str.upper()
+    dframe = get_table_tfs(madx, table_name="twissrdt")
 
     if file:
         logger.debug(f"Exporting results to disk at '{Path(file).absolute()}'")

pyhdtoolkit/cpymadtools/track.py

@@ -62,9 +62,11 @@ def track_single_particle(
 
     logger.debug(f"Tracking coordinates with initial X, PX, Y, PY, T, PT of '{initial_coordinates}'")
     madx.command.track(**kwargs)
+
     for element in observation_points:
         logger.trace(f"Setting observation point for tracking with OBSERVE at element '{element}'")
         madx.command.observe(place=element)
+
     madx.command.start(
         X=start[0], PX=start[1], Y=start[2], PY=start[3], T=start[4], PT=start[5],
     )

pyhdtoolkit/cpymadtools/twiss.py

@@ -50,6 +50,7 @@ def get_pattern_twiss(
     """
     logger.trace("Clearing 'TWISS' flag")
     madx.select(flag="twiss", clear=True)
+
     for pattern in patterns:
         logger.trace(f"Adding pattern {pattern} to 'TWISS' flag")
         madx.select(flag="twiss", pattern=pattern, column=columns)

pyhdtoolkit/cpymadtools/utils.py

@@ -0,0 +1,37 @@
+"""
+Module cpymadtools.utils
+------------------------
+
+Created on 2021.07.22
+:author: Felix Soubelet (felix.soubelet@cern.ch)
+
+A module with utility functions to do mundane operatiions with `cpymad.madx.Madx` objects.
+"""
+import tfs
+
+from cpymad.madx import Madx
+from loguru import logger
+
+
+def get_table_tfs(madx: Madx, table_name: str) -> tfs.TfsDataFrame:
+    """
+    Turns an internal table from the `MAD-X` process into a `TfsDataFrame` object.
+
+    Args:
+        madx (cpymad.madx.Madx): an instanciated cpymad Madx object.
+        table_name (str): the name of the internal table.
+
+    Returns:
+        A `TfsDataFrame` object with the table data, and the `SUMM` table as headers.
+    """
+    logger.debug(f"Extracting table {table_name} into a TfsDataFrame")
+    dframe = tfs.TfsDataFrame(madx.table[table_name].dframe())
+    dframe.columns = dframe.columns.str.upper()
+
+    if "NAME" in dframe.columns:
+        logger.trace("Uppercasing 'NAME' column contents")
+        dframe.NAME = dframe.NAME.str.upper()
+
+    logger.trace("Turning SUMM table into headers")
+    dframe.headers = {var.upper(): madx.table.summ[var][0] for var in madx.table.summ}
+    return dframe

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "pyhdtoolkit"
-version = "0.10.0"
+version = "0.11.0"
 description = "An all-in-one toolkit package to easy my Python work in my PhD."
 authors = ["Felix Soubelet <felix.soubelet@cern.ch>"]
 license = "MIT"