From a4793a7b28d9985a36a233b350bc5041c176eefe Mon Sep 17 00:00:00 2001 From: Felix Soubelet <19598248+fsoubelet@users.noreply.github.com> Date: Mon, 28 Jun 2021 17:11:18 +0000 Subject: [PATCH] Tracking improvements (#43) --- pyhdtoolkit/__init__.py | 2 +- pyhdtoolkit/cpymadtools/special.py | 1 + pyhdtoolkit/cpymadtools/track.py | 42 ++++++++++++++++++++++++------ pyproject.toml | 2 +- tests/test_cpymadtools.py | 35 ++++++++++++++++++++----- 5 files changed, 66 insertions(+), 16 deletions(-) diff --git a/pyhdtoolkit/__init__.py b/pyhdtoolkit/__init__.py index 532c3d79..a96c3180 100644 --- a/pyhdtoolkit/__init__.py +++ b/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" diff --git a/pyhdtoolkit/cpymadtools/special.py b/pyhdtoolkit/cpymadtools/special.py index f91b94bd..19afeedf 100644 --- a/pyhdtoolkit/cpymadtools/special.py +++ b/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 ----- # diff --git a/pyhdtoolkit/cpymadtools/track.py b/pyhdtoolkit/cpymadtools/track.py index 6f8c71bf..fc00fa4e 100644 --- a/pyhdtoolkit/cpymadtools/track.py +++ b/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 + } diff --git a/pyproject.toml b/pyproject.toml index ad6914d8..59ca5341 100644 --- a/pyproject.toml +++ b/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 "] license = "MIT" diff --git a/tests/test_cpymadtools.py b/tests/test_cpymadtools.py index 90af35c7..eb54a108 100644 --- a/tests/test_cpymadtools.py +++ b/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")] )