Skip to content

Commit cd27bef

Browse files
jgray-19jgray-19
authored
Convert MAD-NG track to TBT data (#21)
Adds a module provides functions to read and write ``MAD-NG`` turn-by-turn measurement files. These files are in the **TFS** format. --------- Co-authored-by: jgray-19 <jgray-19@users.noreply.github.com>
1 parent 3163c5e commit cd27bef

File tree

10 files changed

+377
-7
lines changed

10 files changed

+377
-7
lines changed

‎doc/readers/index.rst‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,3 +31,7 @@
3131
.. automodule:: turn_by_turn.trackone
3232
:members:
3333
:noindex:
34+
35+
.. automodule:: turn_by_turn.madng
36+
:members:
37+
:noindex:

‎pyproject.toml‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,7 @@ authors = [
2424
]
2525
license = "MIT"
2626
dynamic = ["version"]
27-
requires-python = ">=3.9"
27+
requires-python = ">=3.10"
2828

2929
classifiers = [
3030
"Development Status :: 5 - Production/Stable",
@@ -33,7 +33,6 @@ classifiers = [
3333
"Natural Language :: English",
3434
"Operating System :: OS Independent",
3535
"Programming Language :: Python :: 3 :: Only",
36-
"Programming Language :: Python :: 3.9",
3736
"Programming Language :: Python :: 3.10",
3837
"Programming Language :: Python :: 3.11",
3938
"Programming Language :: Python :: 3.12",
@@ -49,6 +48,7 @@ dependencies = [
4948
"pandas >= 2.1",
5049
"sdds >= 0.4",
5150
"h5py >= 2.9",
51+
"tfs-pandas >= 4.0.0", # for madng (could be an optional dependency)
5252
]
5353

5454
[project.optional-dependencies]

‎tests/inputs/madng/fodo_track.mad‎

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,42 @@
1+
MADX:open_env()
2+
circum = 60
3+
lcell = 20
4+
f =\ lcell/sin(pi/4)/4
5+
k =\ 1/f
6+
qf = multipole 'QF' { knl := {0,k} }
7+
qd = multipole 'QD' { knl := {0, -k} }
8+
bpm = monitor 'BPM' { }
9+
10+
seq = sequence 'SEQ' { refer = centre, l = circum,
11+
bpm 'BPM1' { at = 0 },
12+
qf 'QF' { at = 0 },
13+
qd 'QD' { at = 0.5 *lcell },
14+
qf 'QF' { at = 1.0 *lcell },
15+
bpm 'BPM3' { at = 1.5 *lcell },
16+
qd 'QD' { at = 1.5 *lcell },
17+
qf 'QF' { at = 2.0 *lcell },
18+
qd 'QD' { at = 2.5 *lcell },
19+
bpm 'BPM2' { at = 3 *lcell },
20+
}
21+
MADX:close_env()
22+
23+
local seq in MADX
24+
25+
local beam, track in MAD
26+
seq.beam = beam
27+
28+
local observed in MAD.element.flags
29+
seq:deselect(observed)
30+
seq:select(observed, {pattern="BPM"})
31+
local part1 = {x = 1e-3, px = 0, y =-1e-3, py = 0, t = 0, pt = 0}
32+
local part2 = {x =-1e-3, px = 0, y = 1e-3, py = 0, t = 0, pt = 0}
33+
34+
local mtbl = track {
35+
sequence = seq,
36+
X0 = {part1, part2},
37+
nturn = 3,
38+
}
39+
mtbl:write("fodo_track.tfs", {"name", "x", "y", "turn", "id"})
40+
41+
mtbl:remove(#mtbl) -- remove the final row to produce an incomplete file
42+
mtbl:write("fodo_track_error.tfs", {"name", "x", "y", "turn", "id"})

‎tests/inputs/madng/fodo_track.tfs‎

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,35 @@
1+
@ name %03s "SEQ"
2+
@ type %05s "track"
3+
@ title %03s "SEQ"
4+
@ origin %18s "MAD 1.0.0 Linux 64"
5+
@ date %08s "07/01/25"
6+
@ time %08s "15:38:30"
7+
@ refcol %04s "name"
8+
@ direction %le 1
9+
@ observe %le 1
10+
@ implicit %b false
11+
@ misalign %b false
12+
@ radiate %b false
13+
@ energy %le 1
14+
@ deltap %le 0
15+
@ lost %le 0
16+
* name x y turn id
17+
$ %s %le %le %le %le
18+
"BPM1" 0.001 -0.001 1 1
19+
"BPM1" -0.001 0.001 1 2
20+
"BPM3" -0.0009999999503 0.00100000029 1 1
21+
"BPM3" 0.0009999999503 -0.00100000029 1 2
22+
"BPM2" 0.002414213831 0.0004142133507 1 1
23+
"BPM2" -0.002414213831 -0.0004142133507 1 2
24+
"BPM1" 0.002414213831 0.0004142133507 2 1
25+
"BPM1" -0.002414213831 -0.0004142133507 2 2
26+
"BPM3" -0.0004142138307 -0.002414213351 2 1
27+
"BPM3" 0.0004142138307 0.002414213351 2 2
28+
"BPM2" -0.0009999991309 0.001000000149 2 1
29+
"BPM2" 0.0009999991309 -0.001000000149 2 2
30+
"BPM1" -0.0009999991309 0.001000000149 3 1
31+
"BPM1" 0.0009999991309 -0.001000000149 3 2
32+
"BPM3" 0.0009999998012 -0.001000001159 3 1
33+
"BPM3" -0.0009999998012 0.001000001159 3 2
34+
"BPM2" -0.002414214191 -0.0004142129907 3 1
35+
"BPM2" 0.002414214191 0.0004142129907 3 2

‎tests/inputs/madng/fodo_track_error.tfs‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,34 @@
1+
@ name %03s "SEQ"
2+
@ type %05s "track"
3+
@ title %03s "SEQ"
4+
@ origin %18s "MAD 1.0.0 Linux 64"
5+
@ date %08s "07/01/25"
6+
@ time %08s "15:38:30"
7+
@ refcol %04s "name"
8+
@ direction %le 1
9+
@ observe %le 1
10+
@ implicit %b false
11+
@ misalign %b false
12+
@ radiate %b false
13+
@ energy %le 1
14+
@ deltap %le 0
15+
@ lost %le 0
16+
* name x y turn id
17+
$ %s %le %le %le %le
18+
"BPM1" 0.001 -0.001 1 1
19+
"BPM1" -0.001 0.001 1 2
20+
"BPM3" -0.0009999999503 0.00100000029 1 1
21+
"BPM3" 0.0009999999503 -0.00100000029 1 2
22+
"BPM2" 0.002414213831 0.0004142133507 1 1
23+
"BPM2" -0.002414213831 -0.0004142133507 1 2
24+
"BPM1" 0.002414213831 0.0004142133507 2 1
25+
"BPM1" -0.002414213831 -0.0004142133507 2 2
26+
"BPM3" -0.0004142138307 -0.002414213351 2 1
27+
"BPM3" 0.0004142138307 0.002414213351 2 2
28+
"BPM2" -0.0009999991309 0.001000000149 2 1
29+
"BPM2" 0.0009999991309 -0.001000000149 2 2
30+
"BPM1" -0.0009999991309 0.001000000149 3 1
31+
"BPM1" 0.0009999991309 -0.001000000149 3 2
32+
"BPM3" 0.0009999998012 -0.001000001159 3 1
33+
"BPM3" -0.0009999998012 0.001000001159 3 2
34+
"BPM2" -0.002414214191 -0.0004142129907 3 1

‎tests/test_madng.py‎

Lines changed: 86 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,86 @@
1+
2+
from datetime import datetime
3+
4+
import numpy as np
5+
import pandas as pd
6+
import pytest
7+
8+
from tests.test_lhc_and_general import INPUTS_DIR, compare_tbt
9+
from turn_by_turn import madng, read_tbt, write_tbt
10+
from turn_by_turn.structures import TbtData, TransverseData
11+
12+
13+
def test_read_ng(_ng_file):
14+
original = _original_simulation_data()
15+
16+
# Check directly from the module
17+
new = madng.read_tbt(_ng_file)
18+
compare_tbt(original, new, no_binary=True)
19+
20+
# Check from the main function
21+
new = read_tbt(_ng_file, datatype="madng")
22+
compare_tbt(original, new, no_binary=True)
23+
24+
def test_write_ng(_ng_file, tmp_path):
25+
original_tbt = _original_simulation_data()
26+
27+
# Write the data
28+
from_tbt = tmp_path / "from_tbt.tfs"
29+
madng.write_tbt(from_tbt, original_tbt)
30+
31+
# Read the written data
32+
new_tbt = madng.read_tbt(from_tbt)
33+
compare_tbt(original_tbt, new_tbt, no_binary=True)
34+
35+
# Check from the main function
36+
original_tbt = read_tbt(_ng_file, datatype="madng")
37+
write_tbt(from_tbt, original_tbt, datatype="madng")
38+
39+
new_tbt = read_tbt(from_tbt, datatype="madng")
40+
compare_tbt(original_tbt, new_tbt, no_binary=True)
41+
assert original_tbt.date == new_tbt.date
42+
43+
def test_error_ng(_error_file):
44+
with pytest.raises(ValueError):
45+
read_tbt(_error_file, datatype="madng")
46+
47+
# ---- Helpers ---- #
48+
def _original_simulation_data() -> TbtData:
49+
# Create a TbTData object with the original data
50+
names = np.array(["BPM1", "BPM3", "BPM2"])
51+
bpm1_p1_x = np.array([ 1e-3, 0.002414213831,-0.0009999991309])
52+
bpm1_p1_y = np.array([-1e-3, 0.0004142133507, 0.001000000149])
53+
bpm1_p2_x = np.array([-1e-3,-0.002414213831, 0.0009999991309])
54+
bpm1_p2_y = np.array([ 1e-3,-0.0004142133507,-0.001000000149])
55+
56+
bpm2_p1_x = np.array([-0.0009999999503,-0.0004142138307, 0.0009999998012])
57+
bpm2_p1_y = np.array([ 0.00100000029,-0.002414213351,-0.001000001159])
58+
bpm2_p2_x = np.array([ 0.0009999999503, 0.0004142138307,-0.0009999998012])
59+
bpm2_p2_y = np.array([-0.00100000029, 0.002414213351, 0.001000001159])
60+
61+
bpm3_p1_x = np.array([ 0.002414213831,-0.0009999991309,-0.002414214191])
62+
bpm3_p1_y = np.array([ 0.0004142133507, 0.001000000149,-0.0004142129907])
63+
bpm3_p2_x = np.array([-0.002414213831, 0.0009999991309, 0.002414214191])
64+
bpm3_p2_y = np.array([-0.0004142133507,-0.001000000149, 0.0004142129907])
65+
66+
matrix = [
67+
TransverseData( # first particle
68+
X=pd.DataFrame(index=names, data=[bpm1_p1_x, bpm2_p1_x, bpm3_p1_x]),
69+
Y=pd.DataFrame(index=names, data=[bpm1_p1_y, bpm2_p1_y, bpm3_p1_y]),
70+
),
71+
TransverseData( # second particle
72+
X=pd.DataFrame(index=names, data=[bpm1_p2_x, bpm2_p2_x, bpm3_p2_x]),
73+
Y=pd.DataFrame(index=names, data=[bpm1_p2_y, bpm2_p2_y, bpm3_p2_y]),
74+
),
75+
]
76+
return TbtData(matrices=matrix, bunch_ids=[1, 2], nturns=3)
77+
78+
79+
# ---- Fixtures ---- #
80+
@pytest.fixture
81+
def _ng_file(tmp_path):
82+
return INPUTS_DIR / "madng" / "fodo_track.tfs"
83+
84+
@pytest.fixture
85+
def _error_file(tmp_path):
86+
return INPUTS_DIR / "madng" / "fodo_track_error.tfs"

‎turn_by_turn/__init__.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55
__title__ = "turn_by_turn"
66
__description__ = "Read and write turn-by-turn measurement files from different particle accelerator formats."
77
__url__ = "https://github.com/pylhc/turn_by_turn"
8-
__version__ = "0.7.2"
8+
__version__ = "0.8.0"
99
__author__ = "pylhc"
1010
__author_email__ = "pylhc@github.com"
1111
__license__ = "MIT"

‎turn_by_turn/io.py‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2,15 +2,15 @@
22
IO
33
--
44
5-
This module contains high-level I/O functions to read and write tur-by-turn data objects in different
5+
This module contains high-level I/O functions to read and write turn-by-turn data objects in different
66
formats. While data can be loaded from the formats of different machines / codes, each format getting its
77
own reader module, writing functionality is at the moment always done in the ``LHC``'s **SDDS** format.
88
"""
99
import logging
1010
from pathlib import Path
1111
from typing import Union, Any
1212

13-
from turn_by_turn import ascii, doros, esrf, iota, lhc, ptc, sps, trackone
13+
from turn_by_turn import ascii, doros, esrf, iota, lhc, ptc, sps, trackone, madng
1414
from turn_by_turn.ascii import write_ascii
1515
from turn_by_turn.errors import DataTypeError
1616
from turn_by_turn.structures import TbtData
@@ -29,8 +29,9 @@
2929
ptc=ptc,
3030
trackone=trackone,
3131
ascii=ascii,
32+
madng=madng,
3233
)
33-
WRITERS = ("lhc", "sps", "doros", "doros_positions", "doros_oscillations", "ascii") # implemented writers
34+
WRITERS = ("lhc", "sps", "doros", "doros_positions", "doros_oscillations", "ascii", "madng") # implemented writers
3435

3536
write_lhc_ascii = write_ascii # Backwards compatibility <0.4
3637

0 commit comments

Comments
 (0)