Tracking improvements (#43)

pyhdtoolkit/__init__.py

@@ -13,7 +13,7 @@
 __title__ = "pyhdtoolkit"
 __description__ = "An all-in-one toolkit package to easy my Python work in my PhD."
 __url__ = "https://github.com/fsoubelet/PyhDToolkit"
-__version__ = "0.9.2"
+__version__ = "0.10.0"
 __author__ = "Felix Soubelet"
 __author_email__ = "felix.soubelet@cern.ch"
 __license__ = "MIT"

pyhdtoolkit/cpymadtools/special.py

@@ -448,6 +448,7 @@ def match_no_coupling_through_ripkens(
     madx.command.lmdif(calls=500, tolerance=1e-21)
     madx.command.endmatch()
 
+
 # ----- Helpers ----- #
 
 

pyhdtoolkit/cpymadtools/track.py

@@ -7,7 +7,7 @@
 
 A module with functions to manipulate MAD-X TRACK functionality through a cpymad.madx.Madx object.
 """
-from typing import Tuple
+from typing import Dict, Sequence, Tuple
 
 import pandas as pd
 
@@ -22,33 +22,59 @@ def track_single_particle(
     initial_coordinates: Tuple[float, float, float, float, float, float],
     nturns: int,
     sequence: str = None,
-) -> pd.DataFrame:
+    observation_points: Sequence[str] = None,
+    **kwargs,
+) -> Dict[str, pd.DataFrame]:
     """
     Tracks a single particle for nturns, based on its initial coordinates.
 
     Args:
         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 the particle to track.
+            T, PT starting coordinates the particle to track. Defaults to all 0 if none given.
         nturns (int): the number of turns to track for.
         sequence (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.
+        observation_points (Sequence[str]): sequence of all element names at which to OBSERVE during the
+            tracking.
+
+    Keyword Args:
+        Any keyword argument to be given to the TRACK command like it would be given directly into MAD-X,
+        for instance ONETABLE etc. Refer to the MAD-X manual for options.
 
     Returns:
-        A copy of the track table's dataframe, with as columns the coordinates x, px, y, py, t, pt,
-        s and e (energy).
+        A dictionary 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 numbered '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 provided the TRACKONE option, only one entry is in the dictionary under the key
+        'trackone' and it has the combine table as a pandas DataFrame for value.
     """
-    start = initial_coordinates
+    onetable = kwargs.get("onetable", False) if "onetable" in kwargs else kwargs.get("ONETABLE", False)
+    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.debug(f"Using sequence '{sequence}' for tracking")
         madx.use(sequence=sequence)
 
     logger.debug(f"Tracking coordinates with initial X, PX, Y, PY, T, PT of '{initial_coordinates}'")
-    madx.command.track()
+    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],
     )
     madx.command.run(turns=nturns)
     madx.command.endtrack()
-    return madx.table["track.obs0001.p0001"].dframe().copy()
+    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
+    }

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "pyhdtoolkit"
-version = "0.9.2"
+version = "0.10.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"

tests/test_cpymadtools.py

@@ -829,10 +829,11 @@ def test_makethin_lhc(self, _matched_lhc_madx):
         madx = _matched_lhc_madx
         make_lhc_thin(madx, sequence="lhcb1", slicefactor=4)
 
-        tracks = track_single_particle(
+        tracks_dict = track_single_particle(
             madx, initial_coordinates=(1e-4, 0, 1e-4, 0, 0, 0), nturns=10, sequence="lhcb1"
         )
-        assert isinstance(tracks, DataFrame)
+        assert isinstance(tracks_dict, dict)
+        tracks = tracks_dict["observation_point_1"]
         assert len(tracks) == 11  # nturns + 1 because $start coordinates also given by MAD-X
         assert all(
             [coordinate in tracks.columns for coordinate in ("x", "px", "y", "py", "t", "pt", "s", "e")]
@@ -880,14 +881,36 @@ def test_vary_independent_ir_quads_raises_on_wrong_quads(self, _non_matched_lhc_
 
 
 class TestTrack:
-    def test_single_particle_tracking(self, _matched_base_lattice):
+    @pytest.mark.parametrize("obs_points", [[], ["qf", "mb", "msf"]])
+    def test_single_particle_tracking(self, _matched_base_lattice, obs_points):
+        madx = _matched_base_lattice
+        tracks_dict = track_single_particle(
+            madx,
+            sequence="CAS3",
+            nturns=100,
+            initial_coordinates=(1e-4, 0, 2e-4, 0, 0, 0),
+            observation_points=obs_points,
+        )
+
+        assert isinstance(tracks_dict, dict)
+        assert len(tracks_dict.keys()) == len(obs_points) + 1
+        for tracks in tracks_dict.values():
+            assert isinstance(tracks, DataFrame)
+            assert all(
+                [coordinate in tracks.columns for coordinate in ("x", "px", "y", "py", "t", "pt", "s", "e")]
+            )
+
+    def test_single_particle_tracking_with_onepass(self, _matched_base_lattice):
         madx = _matched_base_lattice
-        tracks = track_single_particle(
-            madx, initial_coordinates=(1e-4, 0, 2e-4, 0, 0, 0), nturns=100, sequence="CAS3"
+        tracks_dict = track_single_particle(
+            madx, sequence="CAS3", nturns=100, initial_coordinates=(2e-4, 0, 1e-4, 0, 0, 0), ONETABLE=True,
         )
 
+        assert isinstance(tracks_dict, dict)
+        assert len(tracks_dict.keys()) == 1  # should be only one because of ONETABLE option
+        assert "trackone" in tracks_dict.keys()
+        tracks = tracks_dict["trackone"]
         assert isinstance(tracks, DataFrame)
-        assert len(tracks) == 101  # nturns + 1 because $start coordinates also given by MAD-X
         assert all(
             [coordinate in tracks.columns for coordinate in ("x", "px", "y", "py", "t", "pt", "s", "e")]
         )