Xsuite Conversion to Tbt Objects

doc/readers/index.rst

@@ -33,5 +33,9 @@
     :noindex:
 
 .. automodule:: turn_by_turn.madng
+    :members:
+    :noindex:
+
+.. automodule:: turn_by_turn.xtrack
     :members:
     :noindex:

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]

tests/conftest.py

@@ -0,0 +1,41 @@
+import numpy as np
+import pandas as pd
+import pytest
+from turn_by_turn.structures import TbtData, TransverseData
+
+@pytest.fixture(scope="session")
+def example_tbt():
+    """
+    Returns a TbtData object with the original simulation data.
+    This fixture is session-scoped and will only be created once per test run.
+    """
+    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)

tests/test_madng.py

@@ -1,28 +1,22 @@
-
-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_tbt: TbtData):
     # Check directly from the module
     new = madng.read_tbt(_ng_file)
-    compare_tbt(original, new, no_binary=True)
+    compare_tbt(example_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_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_tbt: TbtData):
+    original_tbt = example_tbt
 
     # Write the data
     from_tbt = tmp_path / "from_tbt.tfs"
@@ -40,47 +34,15 @@ def test_write_ng(_ng_file, tmp_path):
     compare_tbt(original_tbt, new_tbt, no_binary=True)
     assert original_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):
     return INPUTS_DIR / "madng" / "fodo_track.tfs"
 
 @pytest.fixture
-def _error_file(tmp_path):
+def _error_file(tmp_path: Path):
     return INPUTS_DIR / "madng" / "fodo_track_error.tfs"

tests/test_xtrack.py

@@ -0,0 +1,62 @@
+import numpy as np
+import xtrack as xt
+import pytest
+import sys
+
+from tests.test_lhc_and_general import compare_tbt
+from turn_by_turn.structures import TbtData
+from turn_by_turn import xtrack
+from turn_by_turn.io import convert_to_tbt
+
+@pytest.mark.skipif(sys.platform == "win32", reason="xtrack kernel compilation not supported on Windows CI")
+def test_convert_xsuite(example_line: xt.Line, example_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.convert_to_tbt(example_line)
+    compare_tbt(example_tbt, tbt_data, no_binary=True)
+
+    # Now convert using the generic function
+    tbt_data = convert_to_tbt(example_line, datatype="xtrack")
+    compare_tbt(example_tbt, tbt_data, no_binary=True)
+
+# --- Fixtures ---- #
+@pytest.fixture(scope="module")
+def example_line():
+    lcell = 20.
+    f = lcell/np.sin(np.pi/4.)/4.
+    k = 1/f
+    nturns = 3
+    qf = xt.Multipole(knl=[0., k], ksl=[0.,0.])
+    qd = xt.Multipole(knl=[0.,-k], ksl=[0.,0.])
+    drift = xt.Drift(length=10.)
+    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'
+            ]
+        )
+    line.particle_ref = xt.Particles(p0c=1e9, q0=1., mass0=xt.ELECTRON_MASS_EV)
+    return line

turn_by_turn/__init__.py

@@ -1,5 +1,5 @@
 """Exposes TbtData, read_tbt and write_tbt directly in the package's namespace."""
-from .io import read_tbt, write_tbt
+from .io import read_tbt, write_tbt, convert_to_tbt, load_tbt_data
 from .structures import TbtData, TransverseData
 
 __title__ = "turn_by_turn"
@@ -13,6 +13,8 @@
 # aliases
 write = write_tbt
 read = read_tbt
+convert = convert_to_tbt
+load = load_tbt_data
 
 # Importing * is a bad practice and you should be punished for using it
 __all__ = []