diff --git a/doc/conf.py b/doc/conf.py index 91aa21e..e36821c 100644 --- a/doc/conf.py +++ b/doc/conf.py @@ -1,4 +1,3 @@ -# -*- coding: utf-8 -*- # # TFS-Pandas documentation build configuration file, created by # sphinx-quickstart on Tue Feb 6 12:10:18 2018. @@ -90,7 +89,7 @@ def about_package(init_posixpath: pathlib.Path) -> dict: # Override link in 'Edit on Github' rst_prolog = f""" -:github_url: {ABOUT_TBT['__url__']} +:github_url: {ABOUT_TBT["__url__"]} """ # The version info for the project you're documenting, acts as replacement for @@ -120,6 +119,9 @@ def about_package(init_posixpath: pathlib.Path) -> dict: # If true, `todo` and `todoList` produce output, else they produce nothing. todo_include_todos = True +# Activate nitpicky mode for sphinx to warn about missing references +# nitpicky = True + # -- Options for HTML output ---------------------------------------------- # The theme to use for HTML and HTML Help pages. See the documentation for @@ -130,7 +132,7 @@ def about_package(init_posixpath: pathlib.Path) -> dict: html_logo = "_static/img/omc_logo.svg" html_static_path = ["_static"] html_context = { -# "css_files": ["_static/css/custom.css"], + # "css_files": ["_static/css/custom.css"], "display_github": True, # the following are only needed if :github_url: is not set "github_user": author, @@ -141,17 +143,18 @@ def about_package(init_posixpath: pathlib.Path) -> dict: "css/custom.css", ] -smartquotes_action = "qe" # renders only quotes and ellipses (...) but not dashes (option: D) +# renders only quotes and ellipses (...) but not dashes (option: D) +smartquotes_action = "qe" # Theme options are theme-specific and customize the look and feel of a theme # further. For a list of options available for each theme, see the # documentation. # html_theme_options = { - 'collapse_navigation': False, - 'display_version': True, - 'logo_only': True, - 'navigation_depth': 1, + "collapse_navigation": False, + "version_selector": True, # sphinx-rtd-theme>=3.0, formerly 'display_version' + "logo_only": True, + "navigation_depth": 2, } # Add any paths that contain custom static files (such as style sheets) here, @@ -163,11 +166,11 @@ def about_package(init_posixpath: pathlib.Path) -> dict: # pages. Single values can also be put in this dictionary using the # -A command-line option of sphinx-build. html_context = { - 'display_github': True, + "display_github": True, # the following are only needed if :github_url: is not set - 'github_user': author, - 'github_repo': project, - 'github_version': 'master/doc/', + "github_user": author, + "github_repo": project, + "github_version": "master/doc/", } # Custom sidebar templates, must be a dictionary that maps document names # to template names. @@ -207,7 +210,13 @@ def about_package(init_posixpath: pathlib.Path) -> dict: # (source start file, target name, title, # author, documentclass [howto, manual, or own class]). latex_documents = [ - (master_doc, "turn_by_turn.tex", "turn_by_turn Documentation", "pyLHC/OMC-TEAM", "manual"), + ( + master_doc, + "turn_by_turn.tex", + "turn_by_turn Documentation", + "pyLHC/OMC-TEAM", + "manual", + ), ] # -- Options for manual page output --------------------------------------- @@ -232,3 +241,19 @@ def about_package(init_posixpath: pathlib.Path) -> dict: "Miscellaneous", ), ] + +# -- Instersphinx Configuration ---------------------------------------------- + +# Example configuration for intersphinx: refer to the Python standard library. +# use in refs e.g: +# :ref:`comparison manual ` +intersphinx_mapping = { + "python": ("https://docs.python.org/3/", None), + "numpy": ("https://numpy.org/doc/stable/", None), + "pandas": ("https://pandas.pydata.org/pandas-docs/stable/", None), + "matplotlib": ("https://matplotlib.org/stable/", None), + "scipy": ("https://docs.scipy.org/doc/scipy/", None), + "cpymad": ("https://hibtc.github.io/cpymad/", None), + "tfs": ("https://pylhc.github.io/tfs/", None), + "sdds": ("https://pylhc.github.io/sdds/", None), +} diff --git a/doc/index.rst b/doc/index.rst index 7228b52..d4d34fc 100644 --- a/doc/index.rst +++ b/doc/index.rst @@ -5,8 +5,41 @@ Welcome to turn_by_turn' documentation! It provides a custom dataclass ``TbtData`` to do so, with attributes corresponding to the relevant measurements information. +How to Use turn_by_turn +----------------------- + +There are two main ways to create a ``TbtData`` object: + +1. **Reading from file (disk):** + Use ``read_tbt`` to load measurement data from a file on disk. This is the standard entry point for working with measurement files in supported formats. + +2. **In-memory conversion:** + Use ``convert_to_tbt`` to convert data that is already loaded in memory (such as a pandas DataFrame, tfs DataFrame, or xtrack.Line) into a ``TbtData`` object. This is useful for workflows where you generate or manipulate data in Python before standardizing it. + +Both methods produce a ``TbtData`` object, which can then be used for further analysis or written out to supported formats. + +Supported Modules and Limitations +--------------------------------- + +Different modules support different file formats and workflows (disk reading vs. in-memory conversion). For a detailed table of which modules support which features, and any important limitations, see the documentation for the :mod:`turn_by_turn.io` module. + +- Only ``madng`` and ``xtrack`` support in-memory conversion. +- Most modules are for disk reading only. +- Some modules (e.g., ``esrf``) are experimental or have limited support. +- For writing, see the next section. + +Writing Data +------------ + +To write a ``TbtData`` object to disk, use the ``write_tbt`` function. This function supports writing in the LHC SDDS format by default, as well as other supported formats depending on the ``datatype`` argument. The output format is determined by the ``datatype`` you specify, but for most workflows, SDDS is the standard output. + +Example:: + + from turn_by_turn.io import write_tbt + write_tbt("output.sdds", tbt_data) + Package Reference -================= +----------------- .. toctree:: :caption: Modules @@ -24,9 +57,8 @@ Package Reference Indices and tables -================== +------------------ * :ref:`genindex` * :ref:`modindex` * :ref:`search` - diff --git a/doc/modules/index.rst b/doc/modules/index.rst index 930ad99..41181a0 100644 --- a/doc/modules/index.rst +++ b/doc/modules/index.rst @@ -16,4 +16,3 @@ .. automodule:: turn_by_turn.utils :members: :noindex: - diff --git a/doc/readers/index.rst b/doc/readers/index.rst index 36d0a95..1e657b8 100644 --- a/doc/readers/index.rst +++ b/doc/readers/index.rst @@ -33,5 +33,9 @@ :noindex: .. automodule:: turn_by_turn.madng + :members: + :noindex: + +.. automodule:: turn_by_turn.xtrack_line :members: :noindex: \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index 391385f..9b99867 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -29,7 +29,6 @@ requires-python = ">=3.10" classifiers = [ "Development Status :: 5 - Production/Stable", "Intended Audience :: Science/Research", - "License :: OSI Approved :: MIT License", "Natural Language :: English", "Operating System :: OS Independent", "Programming Language :: Python :: 3 :: Only", @@ -48,14 +47,26 @@ dependencies = [ "pandas >= 2.1", "sdds >= 0.4", "h5py >= 2.9", - "tfs-pandas >= 4.0.0", # for madng (could be an optional dependency) ] [project.optional-dependencies] +madng = [ + "tfs-pandas >= 4.0.0", # for reading MAD-NG files (Could do everything in memory with just pandas) +] + +xtrack = [ + "xtrack >= 0.84.7", # for xtrack + "setuptools >= 65", # for xtrack + "xpart >= 0.23.0", # for xtrack +] + test = [ "pytest>=7.0", "pytest-cov>=2.9", + "turn_by_turn[madng]", + "turn_by_turn[xtrack]", ] + doc = [ "sphinx >= 7.0", "sphinx_rtd_theme >= 2.0", @@ -64,6 +75,8 @@ doc = [ all = [ "turn_by_turn[test]", "turn_by_turn[doc]", + "turn_by_turn[madng]", + "turn_by_turn[xtrack]", ] [project.urls] diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..8093167 --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,46 @@ +import numpy as np +import pandas as pd +import pytest +from turn_by_turn.structures import TbtData, TransverseData + +@pytest.fixture(scope="session") +def example_fake_tbt(): + """ + Returns a TbtData object using simulation data taken from MAD-NG. + This data is also used for the tests in xtrack, so change the numbers + at your own risk. + + It is possible to run the MAD-NG in the inputs folder to regenerate the data. + Also, xtrack produces the same data, so you can use the xtrack test fixture + `example_line`. + """ + names = np.array(["BPM1", "BPM3", "BPM2"]) + # First BPM + bpm1_p1_x = np.array([ 1e-3, 0.002414213831,-0.0009999991309]) + bpm1_p1_y = np.array([-1e-3, 0.0004142133507, 0.001000000149]) + bpm1_p2_x = np.array([-1e-3,-0.002414213831, 0.0009999991309]) + bpm1_p2_y = np.array([ 1e-3,-0.0004142133507,-0.001000000149]) + + # Second BPM + bpm3_p1_x = np.array([ 0.002414213831,-0.0009999991309,-0.002414214191]) + bpm3_p1_y = np.array([ 0.0004142133507, 0.001000000149,-0.0004142129907]) + bpm3_p2_x = np.array([-0.002414213831, 0.0009999991309, 0.002414214191]) + bpm3_p2_y = np.array([-0.0004142133507,-0.001000000149, 0.0004142129907]) + + # Third BPM + bpm2_p1_x = np.array([-0.0009999999503,-0.0004142138307, 0.0009999998012]) + bpm2_p1_y = np.array([ 0.00100000029,-0.002414213351,-0.001000001159]) + bpm2_p2_x = np.array([ 0.0009999999503, 0.0004142138307,-0.0009999998012]) + bpm2_p2_y = np.array([-0.00100000029, 0.002414213351, 0.001000001159]) + + matrix = [ + TransverseData( # first particle + X=pd.DataFrame(index=names, data=[bpm1_p1_x, bpm2_p1_x, bpm3_p1_x]), + Y=pd.DataFrame(index=names, data=[bpm1_p1_y, bpm2_p1_y, bpm3_p1_y]), + ), + TransverseData( # second particle + X=pd.DataFrame(index=names, data=[bpm1_p2_x, bpm2_p2_x, bpm3_p2_x]), + Y=pd.DataFrame(index=names, data=[bpm1_p2_y, bpm2_p2_y, bpm3_p2_y]), + ), + ] + return TbtData(matrices=matrix, bunch_ids=[0, 1], nturns=3) diff --git a/tests/test_madng.py b/tests/test_madng.py index 8ef0390..40d5f21 100644 --- a/tests/test_madng.py +++ b/tests/test_madng.py @@ -1,86 +1,46 @@ - -from datetime import datetime - -import numpy as np -import pandas as pd import pytest +from pathlib import Path from tests.test_lhc_and_general import INPUTS_DIR, compare_tbt from turn_by_turn import madng, read_tbt, write_tbt -from turn_by_turn.structures import TbtData, TransverseData +from turn_by_turn.structures import TbtData -def test_read_ng(_ng_file): - original = _original_simulation_data() - +def test_read_ng(_ng_file: Path, example_fake_tbt: TbtData): # Check directly from the module new = madng.read_tbt(_ng_file) - compare_tbt(original, new, no_binary=True) + compare_tbt(example_fake_tbt, new, no_binary=True) # Check from the main function new = read_tbt(_ng_file, datatype="madng") - compare_tbt(original, new, no_binary=True) + compare_tbt(example_fake_tbt, new, no_binary=True) -def test_write_ng(_ng_file, tmp_path): - original_tbt = _original_simulation_data() - +def test_write_ng(_ng_file: Path, tmp_path: Path, example_fake_tbt: TbtData): # Write the data from_tbt = tmp_path / "from_tbt.tfs" - madng.write_tbt(from_tbt, original_tbt) + madng.write_tbt(from_tbt, example_fake_tbt) # Read the written data new_tbt = madng.read_tbt(from_tbt) - compare_tbt(original_tbt, new_tbt, no_binary=True) + compare_tbt(example_fake_tbt, new_tbt, no_binary=True) # Check from the main function - original_tbt = read_tbt(_ng_file, datatype="madng") - write_tbt(from_tbt, original_tbt, datatype="madng") + written_tbt = read_tbt(_ng_file, datatype="madng") + write_tbt(from_tbt, written_tbt, datatype="madng") new_tbt = read_tbt(from_tbt, datatype="madng") - compare_tbt(original_tbt, new_tbt, no_binary=True) - assert original_tbt.date == new_tbt.date + compare_tbt(written_tbt, new_tbt, no_binary=True) + assert written_tbt.date == new_tbt.date -def test_error_ng(_error_file): +def test_error_ng(_error_file: Path): with pytest.raises(ValueError): read_tbt(_error_file, datatype="madng") -# ---- Helpers ---- # -def _original_simulation_data() -> TbtData: - # Create a TbTData object with the original data - names = np.array(["BPM1", "BPM3", "BPM2"]) - bpm1_p1_x = np.array([ 1e-3, 0.002414213831,-0.0009999991309]) - bpm1_p1_y = np.array([-1e-3, 0.0004142133507, 0.001000000149]) - bpm1_p2_x = np.array([-1e-3,-0.002414213831, 0.0009999991309]) - bpm1_p2_y = np.array([ 1e-3,-0.0004142133507,-0.001000000149]) - - bpm2_p1_x = np.array([-0.0009999999503,-0.0004142138307, 0.0009999998012]) - bpm2_p1_y = np.array([ 0.00100000029,-0.002414213351,-0.001000001159]) - bpm2_p2_x = np.array([ 0.0009999999503, 0.0004142138307,-0.0009999998012]) - bpm2_p2_y = np.array([-0.00100000029, 0.002414213351, 0.001000001159]) - - bpm3_p1_x = np.array([ 0.002414213831,-0.0009999991309,-0.002414214191]) - bpm3_p1_y = np.array([ 0.0004142133507, 0.001000000149,-0.0004142129907]) - bpm3_p2_x = np.array([-0.002414213831, 0.0009999991309, 0.002414214191]) - bpm3_p2_y = np.array([-0.0004142133507,-0.001000000149, 0.0004142129907]) - - matrix = [ - TransverseData( # first particle - X=pd.DataFrame(index=names, data=[bpm1_p1_x, bpm2_p1_x, bpm3_p1_x]), - Y=pd.DataFrame(index=names, data=[bpm1_p1_y, bpm2_p1_y, bpm3_p1_y]), - ), - TransverseData( # second particle - X=pd.DataFrame(index=names, data=[bpm1_p2_x, bpm2_p2_x, bpm3_p2_x]), - Y=pd.DataFrame(index=names, data=[bpm1_p2_y, bpm2_p2_y, bpm3_p2_y]), - ), - ] - return TbtData(matrices=matrix, bunch_ids=[1, 2], nturns=3) - - # ---- Fixtures ---- # @pytest.fixture -def _ng_file(tmp_path): +def _ng_file(tmp_path: Path) -> Path: return INPUTS_DIR / "madng" / "fodo_track.tfs" @pytest.fixture -def _error_file(tmp_path): +def _error_file(tmp_path: Path) -> Path: return INPUTS_DIR / "madng" / "fodo_track_error.tfs" \ No newline at end of file diff --git a/tests/test_xtrack.py b/tests/test_xtrack.py new file mode 100644 index 0000000..a9b56b7 --- /dev/null +++ b/tests/test_xtrack.py @@ -0,0 +1,73 @@ +import sys + +import numpy as np +import pytest +import xtrack as xt + +from tests.test_lhc_and_general import compare_tbt +from turn_by_turn import xtrack_line +from turn_by_turn.io import convert_to_tbt +from turn_by_turn.structures import TbtData + + +@pytest.mark.skipif(sys.platform == "win32", reason="xtrack not supported on Windows") +def test_convert_xsuite(example_line: xt.Line, example_fake_tbt: TbtData): + # Build the particles + particles = example_line.build_particles(x=[1e-3, -1e-3], y=[-1e-3, 1e-3]) + + # Track the particles through the line + example_line.track(particles, num_turns=3) + + # Convert to TbtData using xtrack + tbt_data = xtrack_line.convert_to_tbt(example_line) + compare_tbt(example_fake_tbt, tbt_data, no_binary=True) + + # Now convert using the generic function + tbt_data = convert_to_tbt(example_line, datatype="xtrack") + compare_tbt(example_fake_tbt, tbt_data, no_binary=True) + + +# --- Fixtures ---- # +@pytest.fixture(scope="module") +def example_line(): + """ + Creates a simple xtrack Line with three BPMs and two quadrupoles. + This replicates the MAD-NG example used in the tests exactly, changes + will likely break the tests. + See tests/inputs/madng/fodo_track.mad for the original MAD-NG file. + """ + lcell = 20.0 + f = lcell / np.sin(np.pi / 4.0) / 4.0 + k = 1 / f + nturns = 3 + qf = xt.Multipole(knl=[0.0, k], ksl=[0.0, 0.0]) + qd = xt.Multipole(knl=[0.0, -k], ksl=[0.0, 0.0]) + drift = xt.Drift(length=10.0) + # fmt: off + line = xt.Line( + elements=[ + xt.ParticlesMonitor(start_at_turn=0, stop_at_turn=nturns, num_particles=2), + qf, drift, + qd, drift, + qf, drift, + xt.ParticlesMonitor(start_at_turn=0, stop_at_turn=nturns, num_particles=2), + qd, drift, + qf, drift, + qd, drift, + xt.ParticlesMonitor(start_at_turn=0, stop_at_turn=nturns, num_particles=2), + ], + element_names=[ + 'BPM1', + 'qf_0', 'drift_0', + 'qd_0', 'drift_1', + 'qf_1', 'drift_2', + 'BPM3', # Deliberately not in order to test the conversion + 'qd_1', 'drift_3', + 'qf_2', 'drift_4', + 'qd_2', 'drift_5', + 'BPM2' + ] + ) + # fmt: on + line.particle_ref = xt.Particles(p0c=1e9, q0=1.0, mass0=xt.ELECTRON_MASS_EV) + return line diff --git a/turn_by_turn/__init__.py b/turn_by_turn/__init__.py index 54835f5..999d476 100644 --- a/turn_by_turn/__init__.py +++ b/turn_by_turn/__init__.py @@ -1,11 +1,12 @@ """Exposes TbtData, read_tbt and write_tbt directly in the package's namespace.""" -from .io import read_tbt, write_tbt -from .structures import TbtData, TransverseData + +from .io import convert_to_tbt, read_tbt, write_tbt +from .structures import TbtData, TransverseData # noqa: F401 __title__ = "turn_by_turn" __description__ = "Read and write turn-by-turn measurement files from different particle accelerator formats." __url__ = "https://github.com/pylhc/turn_by_turn" -__version__ = "0.8.0" +__version__ = "0.9.0" __author__ = "pylhc" __author_email__ = "pylhc@github.com" __license__ = "MIT" @@ -13,6 +14,7 @@ # aliases write = write_tbt read = read_tbt +convert = convert_to_tbt # Importing * is a bad practice and you should be punished for using it __all__ = [] diff --git a/turn_by_turn/ascii.py b/turn_by_turn/ascii.py index 4bfe75c..4e5fa29 100644 --- a/turn_by_turn/ascii.py +++ b/turn_by_turn/ascii.py @@ -10,11 +10,13 @@ - BPM index/longitunial location - Value Turn 1, Turn 2, etc. """ +from __future__ import annotations + import logging from datetime import datetime from pathlib import Path import re -from typing import TextIO, Union, Tuple, List, Optional +from typing import TextIO import numpy as np import pandas as pd @@ -23,7 +25,7 @@ from turn_by_turn.constants import FORMAT_STRING, PLANE_TO_NUM, PLANES, NUM_TO_PLANE from turn_by_turn.structures import TbtData, TransverseData -LOGGER = logging.getLogger() +LOGGER = logging.getLogger(__name__) # ASCII IDs ASCII_COMMENT: str = "#" @@ -32,12 +34,12 @@ ACQ_DATE_FORMAT: str = "%Y-%m-%d at %H:%M:%S" -def is_ascii_file(file_path: Union[str, Path]) -> bool: +def is_ascii_file(file_path: str | Path) -> bool: """ Returns ``True`` only if the file looks like a readable TbT ASCII file, else ``False``. Args: - file_path (Union[str, Path]): path to the turn-by-turn measurement file. + file_path (str | Path): path to the turn-by-turn measurement file. Returns: A boolean. @@ -57,12 +59,12 @@ def is_ascii_file(file_path: Union[str, Path]) -> bool: # ----- Writer ----- # -def write_tbt(output_path: Union[str, Path], tbt_data: TbtData) -> None: +def write_tbt(output_path: str | Path, tbt_data: TbtData) -> None: """ Write a ``TbtData`` object's data to file, in the TbT ASCII format. Args: - output_path (Union[str, Path]): path to the disk location where to write the data. + output_path (str | Path): path to the disk location where to write the data. tbt_data (TbtData): the ``TbtData`` object to write to disk. """ output_path = Path(output_path) @@ -99,13 +101,13 @@ def _write_tbt_data(tbt_data: TbtData, bunch_id: int, output_file: TextIO) -> No # ----- Reader ----- # -def read_tbt(file_path: Union[str, Path], bunch_id: int = None) -> TbtData: +def read_tbt(file_path: str | Path, bunch_id: int = None) -> TbtData: """ Reads turn-by-turn data from an ASCII turn-by-turn format file, and return the date as well as parsed matrices for construction of a ``TbtData`` object. Args: - file_path (Union[str, Path]): path to the turn-by-turn measurement file. + file_path (str | Path): path to the turn-by-turn measurement file. bunch_id (int, optional): the bunch id associated with this file. Defaults to `None`, but is then attempted to parsed from the filename. If not found, `0` is used. @@ -155,7 +157,7 @@ def read_tbt(file_path: Union[str, Path], bunch_id: int = None) -> TbtData: # ----- Helpers ----- # -def _parse_samples(line: str) -> Tuple[str, str, np.ndarray]: +def _parse_samples(line: str) -> tuple[str, str, np.ndarray]: """Parse a line into its different elements.""" parts = line.split() plane_num = parts[0] diff --git a/turn_by_turn/doros.py b/turn_by_turn/doros.py index b819c30..9fcf9a3 100644 --- a/turn_by_turn/doros.py +++ b/turn_by_turn/doros.py @@ -45,7 +45,7 @@ from turn_by_turn.structures import TbtData, TransverseData from turn_by_turn.utils import all_elements_equal -LOGGER = logging.getLogger() +LOGGER = logging.getLogger(__name__) DEFAULT_BUNCH_ID: int = 0 # bunch ID not saved in the DOROS file diff --git a/turn_by_turn/esrf.py b/turn_by_turn/esrf.py index ae59b11..17df28b 100644 --- a/turn_by_turn/esrf.py +++ b/turn_by_turn/esrf.py @@ -18,7 +18,7 @@ from turn_by_turn.utils import numpy_to_tbt BPM_NAMES_FILE: str = "bpm_names.json" -LOGGER = logging.getLogger() +LOGGER = logging.getLogger(__name__) def read_tbt(file_path: Union[str, Path]) -> TbtData: diff --git a/turn_by_turn/io.py b/turn_by_turn/io.py index 4e8a318..d317274 100644 --- a/turn_by_turn/io.py +++ b/turn_by_turn/io.py @@ -2,41 +2,140 @@ IO -- -This module contains high-level I/O functions to read and write turn-by-turn data objects in different -formats. While data can be loaded from the formats of different machines / codes, each format getting its -own reader module, writing functionality is at the moment always done in the ``LHC``'s **SDDS** format. +This module contains high-level I/O functions to read and write turn-by-turn data objects to and from different formats. + +Reading Data +============ +Since version ``0.9.0`` of the package, data can be loaded either from file or from in-memory structures exclusive to certain codes (for some tracking simulation in *MAD-NG* or *xtrack*). +Two different APIs are provided for these use cases. + +1. **To read from file**, use the ``read_tbt`` function (exported as ``read`` at the package's level). The file format is detected or specified by the ``datatype`` parameter. +2. **To load in-memory data**, use the ``convert_to_tbt`` function (exported as ``convert`` at the package's level). This is valid for tracking simulation results from e.g. *xtrack* or sent back by *MAD-NG*. + +In both cases, the returned value is a structured ``TbtData`` object. + +Writing Data +============ +The single entry point for writing to disk is the ``write_tbt`` function (exported as ``write`` at the package's level). This writes a ``TbtData`` object to disk, typically in the LHC SDDS format (by default). The output file extension and format are determined by the ``datatype`` argument. + +The following cases arise: +- If ``datatype`` is set to ``lhc``, ``sps`` or ``ascii``, the output will be in SDDS format and the file extension will be set to ``.sdds`` if not already present. +- If ``datatype`` is set to ``madng``, the output will be in a TFS file (extension ``.tfs`` is recommended). +- Other supported datatypes (see ``WRITERS``) will use their respective formats and conventions if implemented. + +The ``datatype`` parameter controls both the output format and any additional options passed to the underlying writer. +Should the ``noise`` parameter be used, random noise will be added to the data before writing. A ``seed`` can be provided for reproducibility. + +Example:: + + from turn_by_turn import write + write("output.sdds", tbt_data) # writes in SDDS format by default + write("output.tfs", tbt_data, datatype="madng") # writes a TFS file in MAD-NG's tracking results format + write("output.sdds", tbt_data, noise=0.01, seed=42) # reproducibly adds noise before writing + +While data can be loaded from the formats of different machines/codes (each through its own reader module), writing functionality is at the moment always done in the ``LHC``'s **SDDS** format by default, unless another supported format is specified. The interface is designed to be future-proof and easy to extend for new formats. + + +Supported Modules and Limitations +================================= + +The following table summarizes which modules support disk reading and in-memory conversion, and any important limitations: + ++----------------+---------------------+-----------------------+----------------------------------------------------------+ +| Module | Disk Reading | In-Memory Conversion | Notes / Limitations | ++================+=====================+=======================+==========================================================+ +| lhc | Yes (SDDS, ASCII) | No | Reads LHC SDDS and legacy ASCII files. | ++----------------+---------------------+-----------------------+----------------------------------------------------------+ +| sps | Yes (SDDS, ASCII) | No | Reads SPS SDDS and legacy ASCII files. | ++----------------+---------------------+-----------------------+----------------------------------------------------------+ +| doros | Yes (HDF5) | No | Reads DOROS HDF5 files. | ++----------------+---------------------+-----------------------+----------------------------------------------------------+ +| madng | Yes (TFS) | Yes | In-memory: only via pandas/tfs DataFrame. | ++----------------+---------------------+-----------------------+----------------------------------------------------------+ +| xtrack | No | Yes | Only in-memory via xtrack.Line. | ++----------------+---------------------+-----------------------+----------------------------------------------------------+ +| ptc | Yes (trackone) | No | Reads MAD-X PTC trackone files. | ++----------------+---------------------+-----------------------+----------------------------------------------------------+ +| esrf | Yes (Matlab .mat) | No | Experimental/untested. | ++----------------+---------------------+-----------------------+----------------------------------------------------------+ +| iota | Yes (HDF5) | No | Reads IOTA HDF5 files. | ++----------------+---------------------+-----------------------+----------------------------------------------------------+ +| ascii | Yes (legacy ASCII) | No | For legacy ASCII files only. | ++----------------+---------------------+-----------------------+----------------------------------------------------------+ +| trackone | Yes (MAD-X) | No | Reads MAD-X trackone files. | ++----------------+---------------------+-----------------------+----------------------------------------------------------+ + +- Only ``madng`` and ``xtrack`` support in-memory conversion. +- Most modules are for disk reading only. +- Some modules (e.g., ``esrf``) are experimental or have limited support. + +API +=== """ + +from __future__ import annotations + import logging from pathlib import Path -from typing import Union, Any - -from turn_by_turn import ascii, doros, esrf, iota, lhc, ptc, sps, trackone, madng +from typing import TYPE_CHECKING, Any + +from turn_by_turn import ( + ascii, # noqa: A004 + doros, + esrf, + iota, + lhc, + madng, + ptc, + sps, + trackone, + xtrack_line, +) from turn_by_turn.ascii import write_ascii from turn_by_turn.errors import DataTypeError -from turn_by_turn.structures import TbtData from turn_by_turn.utils import add_noise_to_tbt -LOGGER = logging.getLogger() - -TBT_MODULES = dict( - lhc=lhc, - doros=doros, - doros_positions=doros, - doros_oscillations=doros, - sps=sps, - iota=iota, - esrf=esrf, - ptc=ptc, - trackone=trackone, - ascii=ascii, - madng=madng, +LOGGER = logging.getLogger(__name__) + +if TYPE_CHECKING: + from pandas import DataFrame + from xtrack import Line + + from turn_by_turn.structures import TbtData + +TBT_MODULES = { + "lhc": lhc, + "doros": doros, + "doros_positions": doros, + "doros_oscillations": doros, + "sps": sps, + "iota": iota, + "esrf": esrf, + "ptc": ptc, + "trackone": trackone, + "ascii": ascii, + "madng": madng, + "xtrack": xtrack_line, +} + +# Modules supporting in-memory conversion to TbtData (not file readers) +TBT_CONVERTERS = ("madng", "xtrack") + +# implemented writers +WRITERS = ( + "lhc", + "sps", + "doros", + "doros_positions", + "doros_oscillations", + "ascii", + "madng", ) -WRITERS = ("lhc", "sps", "doros", "doros_positions", "doros_oscillations", "ascii", "madng") # implemented writers write_lhc_ascii = write_ascii # Backwards compatibility <0.4 -def read_tbt(file_path: Union[str, Path], datatype: str = "lhc") -> TbtData: +def read_tbt(file_path: str | Path, datatype: str = "lhc") -> TbtData: """ Calls the appropriate loader for the provided matrices type and returns a ``TbtData`` object of the loaded matrices. @@ -62,7 +161,32 @@ def read_tbt(file_path: Union[str, Path], datatype: str = "lhc") -> TbtData: return module.read_tbt(file_path, **additional_args(datatype)) -def write_tbt(output_path: Union[str, Path], tbt_data: TbtData, noise: float = None, seed: int = None, datatype: str = "lhc") -> None: +# Note: I don't specify tfs.TfsDataFrame as this inherits from pandas.DataFrame +def convert_to_tbt(file_data: DataFrame | Line, datatype: str = "xtrack") -> TbtData: + """ + Convert a pandas or tfs DataFrame (MAD-NG) or a Line (XTrack) to a TbtData object. + Args: + file_data (Union[DataFrame, xt.Line]): The data to convert. + datatype (str): The type of the data, either 'xtrack' or 'madng'. Defaults to 'xtrack'. + Returns: + TbtData: The converted TbtData object. + """ + if datatype.lower() not in TBT_CONVERTERS: + raise DataTypeError( + f"Only {','.join(TBT_CONVERTERS)} converters are implemented for now." + ) + + module = TBT_MODULES[datatype.lower()] + return module.convert_to_tbt(file_data) # No additional arguments as no doros. + + +def write_tbt( + output_path: str | Path, + tbt_data: TbtData, + noise: float = None, + seed: int = None, + datatype: str = "lhc", +) -> None: """ Write a ``TbtData`` object's data to file, in the ``LHC``'s **SDDS** format. @@ -71,26 +195,29 @@ def write_tbt(output_path: Union[str, Path], tbt_data: TbtData, noise: float = N tbt_data (TbtData): the ``TbtData`` object to write to disk. noise (float): optional noise to add to the data. seed (int): A given seed to initialise the RNG if one chooses to add noise. This is useful - to ensure the exact same RNG state across operations. Defaults to `None`, which means + to ensure the exact same RNG state across operations. Defaults to ``None``, which means any new RNG operation in noise addition will pull fresh entropy from the OS. datatype (str): type of matrices in the file, determines the reader to use. Case-insensitive, defaults to ``lhc``. """ output_path = Path(output_path) if datatype.lower() not in WRITERS: - raise DataTypeError(f"Only {','.join(WRITERS)} writers are implemented for now.") + raise DataTypeError( + f"Only {','.join(WRITERS)} writers are implemented for now." + ) - if datatype.lower() in ("lhc", "sps", "ascii") and output_path.suffix != ".sdds": - # I would like to remove this, but I'm afraid of compatibility issues with omc3 (jdilly, 2024) + if datatype.lower() in ("lhc", "sps", "ascii") and output_path.suffix != ".sdds": + # I would like to remove this, but I'm afraid of compatibility issues with omc3 (jdilly, 2024) output_path = output_path.with_name(f"{output_path.name}.sdds") + # If the datatype is not in the list of writers, we raise an error. Therefore the datatype + # must be in the TBT_MODULES dictionary -> No need for a try-except block here. try: module = TBT_MODULES[datatype.lower()] - except KeyError as error: - LOGGER.exception( - f"Unsupported datatype '{datatype}' was provided, should be one of {list(TBT_MODULES.keys())}" + except KeyError: + raise DataTypeError( + f"Invalid datatype: {datatype}. Ensure it is one of {', '.join(TBT_MODULES)}." ) - raise DataTypeError(datatype) from error else: if noise is not None: tbt_data = add_noise_to_tbt(tbt_data, noise=noise, seed=seed) @@ -98,16 +225,15 @@ def write_tbt(output_path: Union[str, Path], tbt_data: TbtData, noise: float = N def additional_args(datatype: str) -> dict[str, Any]: - """ Additional parameters to be added to the reader/writer function. - + """Additional parameters to be added to the reader/writer function. + Args: datatype (str): Type of the data. """ if datatype.lower() == "doros_oscillations": - return dict(data_type=doros.DataKeys.OSCILLATIONS) + return {"data_type": doros.DataKeys.OSCILLATIONS} if datatype.lower() == "doros_positions": - return dict(data_type=doros.DataKeys.POSITIONS) - - return dict() + return {"data_type": doros.DataKeys.POSITIONS} + return {} diff --git a/turn_by_turn/iota.py b/turn_by_turn/iota.py index 2e3bf4b..7c12c26 100644 --- a/turn_by_turn/iota.py +++ b/turn_by_turn/iota.py @@ -15,7 +15,7 @@ from turn_by_turn.errors import HDF5VersionError from turn_by_turn.structures import TbtData, TransverseData -LOGGER = logging.getLogger() +LOGGER = logging.getLogger(__name__) VERSIONS = (1, 2) PLANES_CONV: Dict[int, Dict[str, str]] = { diff --git a/turn_by_turn/lhc.py b/turn_by_turn/lhc.py index b5fe0f2..1707057 100644 --- a/turn_by_turn/lhc.py +++ b/turn_by_turn/lhc.py @@ -19,7 +19,7 @@ from turn_by_turn.structures import TbtData, TransverseData from turn_by_turn.utils import matrices_to_array -LOGGER = logging.getLogger() +LOGGER = logging.getLogger(__name__) # IDs N_BUNCHES: str = "nbOfCapBunches" diff --git a/turn_by_turn/madng.py b/turn_by_turn/madng.py index b4e1c41..de3f450 100644 --- a/turn_by_turn/madng.py +++ b/turn_by_turn/madng.py @@ -2,23 +2,35 @@ MAD-NG ------ -This module provides functions to read and write ``MAD-NG`` turn-by-turn measurement files. These files -are in the **TFS** format. +This module provides functions to read and write turn-by-turn measurement data +produced by the ``MAD-NG`` code. MAD-NG stores its tracking data in the **TFS** +(Table File System) file format. + +Data is loaded into the standardized ``TbtData`` structure used by ``turn_by_turn``, +allowing easy post-processing and conversion between formats. + +Dependencies: + - Requires the ``tfs-pandas >= 4.0.0`` package for compatibility with MAD-NG + features. Earlier versions does not support MAD-NG TFS files. """ from __future__ import annotations -from datetime import datetime import logging -from pathlib import Path +from datetime import datetime +from typing import TYPE_CHECKING import pandas as pd -import tfs + +if TYPE_CHECKING: + from pathlib import Path # Only used for type hinting + + import tfs from turn_by_turn.structures import TbtData, TransverseData -LOGGER = logging.getLogger() +LOGGER = logging.getLogger(__name__) # Define the column names in the TFS file NAME = "name" @@ -36,21 +48,69 @@ def read_tbt(file_path: str | Path) -> TbtData: """ - Reads turn-by-turn data from the ``MAD-NG`` **TFS** format file. + Read turn-by-turn data from a MAD-NG TFS file. + + Loads the TFS file using ``tfs-pandas`` and converts its contents into a + ``TbtData`` object for use with the ``turn_by_turn`` toolkit. Args: - file_path (str | Path): path to the turn-by-turn measurement file. + file_path (str | Path): Path to the MAD-NG TFS measurement file. Returns: - A ``TbTData`` object with the loaded data. + TbtData: The loaded turn-by-turn data. + + Raises: + ImportError: If the ``tfs-pandas`` package is not installed. """ + try: + import tfs + except ImportError as e: + raise ImportError( + "The 'tfs' package is required to read MAD-NG TFS files. " + "Install it with: pip install 'turn_by_turn[madng]'" + ) from e + LOGGER.debug("Starting to read TBT data from dataframe") df = tfs.read(file_path) + return convert_to_tbt(df) + + +def convert_to_tbt(df: pd.DataFrame | tfs.TfsDataFrame) -> TbtData: + """ + Convert a TFS or pandas DataFrame to a ``TbtData`` object. + + This function parses the required turn-by-turn columns, reconstructs the + particle-by-particle tracking data, and returns a ``TbtData`` instance + that can be written or converted to other formats. + + Args: + df (pd.DataFrame | TfsDataFrame): + DataFrame containing MAD-NG turn-by-turn tracking data. + + Returns: + TbtData: The extracted and structured turn-by-turn data. + + Raises: + TypeError: If the input is not a recognized DataFrame type. + ValueError: If the data structure is inconsistent (e.g., lost particles). + """ # Get the date and time from the headers (return None if not found) - date_str = df.headers.get(DATE) - time_str = df.headers.get(TIME) - + try: + import tfs + + is_tfs_df = isinstance(df, tfs.TfsDataFrame) + except ImportError: + LOGGER.debug("The 'tfs' package is not installed. Assuming a pandas DataFrame.") + is_tfs_df = False + + if is_tfs_df: + date_str = df.headers.get(DATE) + time_str = df.headers.get(TIME) + else: + date_str = df.attrs.get(DATE) + time_str = df.attrs.get(TIME) + # Combine the date and time into a datetime object date = None if date_str and time_str: @@ -70,17 +130,19 @@ def read_tbt(file_path: str | Path) -> TbtData: # Check if the number of observed points is consistent for all particles/turns if len(df[NAME]) / nturns / npart != num_observables: raise ValueError( - "The number of BPMs (or observed points) is not consistent for all particles/turns. Simulation may have lost particles." + "The number of observed points is not consistent for all particles/turns. " + "Simulation may have lost particles." ) matrices = [] - bunch_ids = range(1, npart + 1) # Particle IDs start from 1 (not 0) - for particle_id in bunch_ids: + # Particle IDs start from 1, but we use 0-based indexing in Python + particle_ids = range(npart) + for particle_id in particle_ids: LOGGER.info(f"Processing particle ID: {particle_id}") # Filter the dataframe for the current particle - df_particle = df.loc[particle_id] + df_particle = df.loc[particle_id + 1] # Create a dictionary of the TransverseData fields tracking_data_dict = { @@ -99,17 +161,32 @@ def read_tbt(file_path: str | Path) -> TbtData: matrices.append(TransverseData(**tracking_data_dict)) LOGGER.debug("Finished reading TBT data") - return TbtData(matrices=matrices, bunch_ids=list(bunch_ids), nturns=nturns, date=date) + return TbtData( + matrices=matrices, bunch_ids=list(particle_ids), nturns=nturns, date=date + ) def write_tbt(output_path: str | Path, tbt_data: TbtData) -> None: """ - Writes turn-by-turn data to a TFS file for MAD-NG. + Write turn-by-turn data to a MAD-NG TFS file. + + Takes a ``TbtData`` object and writes its contents to disk in the standard + TFS format used by MAD-NG, including relevant headers (date, time, origin). Args: - tbt_data (TbtData): Turn-by-turn data to write. - file_path (str | Path): Target file path. + output_path (str | Path): Destination file path for the TFS file. + tbt_data (TbtData): The turn-by-turn data to write. + + Raises: + ImportError: If the ``tfs-pandas`` package is not installed. """ + try: + import tfs + except ImportError as e: + raise ImportError( + "The 'tfs' package is required to write MAD-NG TFS files. Install it with: pip install 'turn_by_turn[madng]'" + ) from e + planes = [plane.lower() for plane in TransverseData.fieldnames()] # x, y plane_dfs = {plane: [] for plane in planes} @@ -134,7 +211,7 @@ def write_tbt(output_path: str | Path, tbt_data: TbtData) -> None: ) # Add the particle ID column - particle_df[PARTICLE_ID] = particle_id + particle_df[PARTICLE_ID] = particle_id + 1 # Convert the turn column to integer and increment by 1 (MAD-NG uses 1-based indexing) particle_df[TURN] = particle_df[TURN].astype(int) + 1 diff --git a/turn_by_turn/ptc.py b/turn_by_turn/ptc.py index 999437a..162f827 100644 --- a/turn_by_turn/ptc.py +++ b/turn_by_turn/ptc.py @@ -22,7 +22,7 @@ from turn_by_turn.errors import PTCFormatError from turn_by_turn.structures import TbtData, TransverseData -LOGGER = logging.getLogger() +LOGGER = logging.getLogger(__name__) Segment = namedtuple("Segment", ["number", "turns", "particles", "element", "name"]) # IDs --- diff --git a/turn_by_turn/sps.py b/turn_by_turn/sps.py index 2a97fe2..534ff8c 100644 --- a/turn_by_turn/sps.py +++ b/turn_by_turn/sps.py @@ -18,7 +18,7 @@ from turn_by_turn.ascii import is_ascii_file, read_tbt as read_ascii from turn_by_turn.structures import TbtData, TransverseData -LOGGER = logging.getLogger() +LOGGER = logging.getLogger(__name__) # IDs N_TURNS: str = "nbOfTurns" diff --git a/turn_by_turn/trackone.py b/turn_by_turn/trackone.py index 323f8b9..10f3de1 100644 --- a/turn_by_turn/trackone.py +++ b/turn_by_turn/trackone.py @@ -16,7 +16,7 @@ from turn_by_turn.structures import TbtData, TrackingData, TransverseData from turn_by_turn.utils import numpy_to_tbt -LOGGER = logging.getLogger() +LOGGER = logging.getLogger(__name__) def read_tbt(file_path: Union[str, Path], is_tracking_data: bool = False) -> TbtData: diff --git a/turn_by_turn/xtrack_line.py b/turn_by_turn/xtrack_line.py new file mode 100644 index 0000000..99f360b --- /dev/null +++ b/turn_by_turn/xtrack_line.py @@ -0,0 +1,168 @@ +""" +XTRACK_LINE +----------- + +This module provides functions to convert tracking results from an ``xtrack.Line`` +into the standardized ``TbtData`` format used by ``turn_by_turn``. + +Prerequisites for using ``convert_to_tbt``: + + 1. The input ``Line`` must contain one or more ``ParticlesMonitor`` elements + positioned at each location where turn-by-turn data is required (e.g., all BPMs). + + A valid monitor setup involves: + + - Placing a ``xt.ParticlesMonitor`` instance in the line's element sequence + at all the places you would like to observe. + - Configuring each monitor with identical settings: + + * ``start_at_turn`` (first turn to record, usually 0) + * ``stop_at_turn`` (The total number of turns to record, e.g., 100) + * ``num_particles`` (number of tracked particles) + + If any monitor is configured with different parameters, ``convert_to_tbt`` + will either find no data or raise an inconsistency error. + + Also, if you specify more turns than were actually tracked, the resulting + TBT data will include all turns up to the monitor's configured limit. + This may result in extra rows filled with zeros for turns where no real + data was recorded, which might not be desirable for your analysis. + + 2. Before conversion, you must: + + - Build particles with the desired initial coordinates + (using ``line.build_particles(...)``). + - Track those particles through the line for the intended number of turns + (using ``line.track(..., num_turns=num_turns)``). + +Once these conditions are met, pass the tracked ``Line`` to ``convert_to_tbt`` to +extract the data from each particle monitor into a ``TbtData`` object. +""" + +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING + +import numpy as np +import pandas as pd + +from turn_by_turn.structures import TbtData, TransverseData + +if TYPE_CHECKING: + from pathlib import Path + + import xtrack as xt + +LOGGER = logging.getLogger(__name__) + + +def convert_to_tbt(xline: xt.Line) -> TbtData: + """ + Convert tracking results from an ``xtrack`` Line into a ``TbtData`` object. + + This function extracts all ``ParticlesMonitor`` elements found in the Line, + verifies they contain consistent turn-by-turn data, and assembles the results + into the standard ``TbtData`` format. One ``TransverseData`` matrix is created + per tracked particle. + + Args: + xline (Line): An ``xtrack.Line`` containing at least one ``ParticlesMonitor``. + + Returns: + TbtData: The extracted turn-by-turn data for all particles and monitors. + + Raises: + ImportError: If the ``xtrack`` library is not installed. + TypeError: If the input is not a valid ``xtrack.Line``. + ValueError: If no monitors are found or data is inconsistent. + """ + try: + import xtrack as xt + except ImportError as e: + raise ImportError( + "The 'xtrack' package is required to convert xtrack Line objects. Install it with: pip install 'turn_by_turn[xtrack]'" + ) from e + + if not isinstance(xline, xt.Line): + raise TypeError(f"Expected an xtrack Line object, got {type(xline)} instead.") + + # Collect monitor names and monitor objects in order from the line + monitor_pairs = [ + (name, elem) + for name, elem in zip(xline.element_names, xline.elements) + if isinstance(elem, xt.ParticlesMonitor) + ] + # Check that we have at least one monitor + if not monitor_pairs: + raise ValueError( + "No ParticlesMonitor found in the Line. Please add a ParticlesMonitor to the Line." + ) + monitor_names, monitors = zip(*monitor_pairs) + + # First check that no particles were lost during tracking. There will be trailing + # zeros in the data if particles were lost. This might be difficult to detect. + assert all( + mon.data.particle_id[-1] == mon.data.particle_id.max() for mon in monitors + ), ( + "Some particles were lost during tracking, which is not supported by this function. " + "Ensure that all particles are tracked through the entire line without loss." + ) + + # Check that all monitors have the same number of turns + nturns_set = {mon.data.at_turn.max() + 1 for mon in monitors} + if len(nturns_set) != 1: + raise ValueError( + "Monitors have different number of turns, have you set the monitors with different 'start_at_turn' or 'stop_at_turn' parameters?" + ) + nturns = nturns_set.pop() + + # Check that all monitors have the same number of particles + npart_set = {len(set(mon.data.particle_id)) for mon in monitors} + if len(npart_set) != 1: + raise ValueError( + "Monitors have different number of particles, maybe some lost particles?" + ) + npart = npart_set.pop() + + # Precompute masks for each monitor and particle_id + monitor_pid_masks = [ + mon.data.particle_id[:, None] == np.arange(npart)[None, :] for mon in monitors + ] + + matrices = [] + # Loop over each particle ID (pid) + for pid in range(npart): + # For each plane (e.g., 'X', 'Y'), build a DataFrame: rows=BPMs, cols=turns + tracking_data_dict = {} + for plane in TransverseData.fieldnames(): + # fmt: off + stacked = np.vstack([ + getattr(mon.data, plane.lower())[monitor_pid_masks[i][:, pid]] + for i, mon in enumerate(monitors) + ]) + # fmt: on + tracking_data_dict[plane] = pd.DataFrame( + stacked, + index=monitor_names, + ) + # Create a TransverseData object for this particle and add to the list + matrices.append(TransverseData(**tracking_data_dict)) + + # Return the TbtData object containing all particles' data + return TbtData( + matrices=matrices, bunch_ids=list(range(npart)), nturns=nturns, date=None + ) + + +# Added this function to match the interface, but it is not implemented. +def read_tbt(path: str | Path) -> None: + """ + Not implemented. + + Reading TBT data directly from files is not supported for xtrack. + Use ``convert_to_tbt`` to convert an in-memory ``xtrack.Line`` instead. + """ + raise NotImplementedError( + "Reading TBT data from xtrack Line files is not implemented." + )