Skip to content

Commit

Permalink
Convert MAD-NG track to TBT data (#21)
Browse files Browse the repository at this point in the history 
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 <[email protected]>
  • Loading branch information
@jgray-19
jgray-19 and jgray-19 authored Jan 8, 2025
1 parent 3163c5e commit cd27bef
Show file tree
Hide file tree
Showing 10 changed files with 377 additions and 7 deletions.
4 changes: 4 additions & 0 deletions doc/readers/index.rst
View file
Original file line number Diff line number Diff line change
Expand Up @@ -31,3 +31,7 @@
.. automodule:: turn_by_turn.trackone
:members:
:noindex:

.. automodule:: turn_by_turn.madng
:members:
:noindex:
4 changes: 2 additions & 2 deletions pyproject.toml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,7 @@ authors = [
]
license = "MIT"
dynamic = ["version"]
requires-python = ">=3.9"
requires-python = ">=3.10"

classifiers = [
"Development Status :: 5 - Production/Stable",
Expand All @@ -33,7 +33,6 @@ classifiers = [
"Natural Language :: English",
"Operating System :: OS Independent",
"Programming Language :: Python :: 3 :: Only",
"Programming Language :: Python :: 3.9",
"Programming Language :: Python :: 3.10",
"Programming Language :: Python :: 3.11",
"Programming Language :: Python :: 3.12",
Expand All @@ -49,6 +48,7 @@ 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]
Expand Down
42 changes: 42 additions & 0 deletions tests/inputs/madng/fodo_track.mad
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,42 @@
MADX:open_env()
circum = 60
lcell = 20
f =\ lcell/sin(pi/4)/4
k =\ 1/f
qf = multipole 'QF' { knl := {0,k} }
qd = multipole 'QD' { knl := {0, -k} }
bpm = monitor 'BPM' { }

seq = sequence 'SEQ' { refer = centre, l = circum,
bpm 'BPM1' { at = 0 },
qf 'QF' { at = 0 },
qd 'QD' { at = 0.5 *lcell },
qf 'QF' { at = 1.0 *lcell },
bpm 'BPM3' { at = 1.5 *lcell },
qd 'QD' { at = 1.5 *lcell },
qf 'QF' { at = 2.0 *lcell },
qd 'QD' { at = 2.5 *lcell },
bpm 'BPM2' { at = 3 *lcell },
}
MADX:close_env()

local seq in MADX

local beam, track in MAD
seq.beam = beam

local observed in MAD.element.flags
seq:deselect(observed)
seq:select(observed, {pattern="BPM"})
local part1 = {x = 1e-3, px = 0, y =-1e-3, py = 0, t = 0, pt = 0}
local part2 = {x =-1e-3, px = 0, y = 1e-3, py = 0, t = 0, pt = 0}

local mtbl = track {
sequence = seq,
X0 = {part1, part2},
nturn = 3,
}
mtbl:write("fodo_track.tfs", {"name", "x", "y", "turn", "id"})

mtbl:remove(#mtbl) -- remove the final row to produce an incomplete file
mtbl:write("fodo_track_error.tfs", {"name", "x", "y", "turn", "id"})
35 changes: 35 additions & 0 deletions tests/inputs/madng/fodo_track.tfs
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,35 @@
@ name %03s "SEQ"
@ type %05s "track"
@ title %03s "SEQ"
@ origin %18s "MAD 1.0.0 Linux 64"
@ date %08s "07/01/25"
@ time %08s "15:38:30"
@ refcol %04s "name"
@ direction %le 1
@ observe %le 1
@ implicit %b false
@ misalign %b false
@ radiate %b false
@ energy %le 1
@ deltap %le 0
@ lost %le 0
* name x y turn id
$ %s %le %le %le %le
"BPM1" 0.001 -0.001 1 1
"BPM1" -0.001 0.001 1 2
"BPM3" -0.0009999999503 0.00100000029 1 1
"BPM3" 0.0009999999503 -0.00100000029 1 2
"BPM2" 0.002414213831 0.0004142133507 1 1
"BPM2" -0.002414213831 -0.0004142133507 1 2
"BPM1" 0.002414213831 0.0004142133507 2 1
"BPM1" -0.002414213831 -0.0004142133507 2 2
"BPM3" -0.0004142138307 -0.002414213351 2 1
"BPM3" 0.0004142138307 0.002414213351 2 2
"BPM2" -0.0009999991309 0.001000000149 2 1
"BPM2" 0.0009999991309 -0.001000000149 2 2
"BPM1" -0.0009999991309 0.001000000149 3 1
"BPM1" 0.0009999991309 -0.001000000149 3 2
"BPM3" 0.0009999998012 -0.001000001159 3 1
"BPM3" -0.0009999998012 0.001000001159 3 2
"BPM2" -0.002414214191 -0.0004142129907 3 1
"BPM2" 0.002414214191 0.0004142129907 3 2
34 changes: 34 additions & 0 deletions tests/inputs/madng/fodo_track_error.tfs
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,34 @@
@ name %03s "SEQ"
@ type %05s "track"
@ title %03s "SEQ"
@ origin %18s "MAD 1.0.0 Linux 64"
@ date %08s "07/01/25"
@ time %08s "15:38:30"
@ refcol %04s "name"
@ direction %le 1
@ observe %le 1
@ implicit %b false
@ misalign %b false
@ radiate %b false
@ energy %le 1
@ deltap %le 0
@ lost %le 0
* name x y turn id
$ %s %le %le %le %le
"BPM1" 0.001 -0.001 1 1
"BPM1" -0.001 0.001 1 2
"BPM3" -0.0009999999503 0.00100000029 1 1
"BPM3" 0.0009999999503 -0.00100000029 1 2
"BPM2" 0.002414213831 0.0004142133507 1 1
"BPM2" -0.002414213831 -0.0004142133507 1 2
"BPM1" 0.002414213831 0.0004142133507 2 1
"BPM1" -0.002414213831 -0.0004142133507 2 2
"BPM3" -0.0004142138307 -0.002414213351 2 1
"BPM3" 0.0004142138307 0.002414213351 2 2
"BPM2" -0.0009999991309 0.001000000149 2 1
"BPM2" 0.0009999991309 -0.001000000149 2 2
"BPM1" -0.0009999991309 0.001000000149 3 1
"BPM1" 0.0009999991309 -0.001000000149 3 2
"BPM3" 0.0009999998012 -0.001000001159 3 1
"BPM3" -0.0009999998012 0.001000001159 3 2
"BPM2" -0.002414214191 -0.0004142129907 3 1
86 changes: 86 additions & 0 deletions tests/test_madng.py
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,86 @@

from datetime import datetime

import numpy as np
import pandas as pd
import pytest

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


def test_read_ng(_ng_file):
original = _original_simulation_data()

# Check directly from the module
new = madng.read_tbt(_ng_file)
compare_tbt(original, new, no_binary=True)

# Check from the main function
new = read_tbt(_ng_file, datatype="madng")
compare_tbt(original, new, no_binary=True)

def test_write_ng(_ng_file, tmp_path):
original_tbt = _original_simulation_data()

# Write the data
from_tbt = tmp_path / "from_tbt.tfs"
madng.write_tbt(from_tbt, original_tbt)

# Read the written data
new_tbt = madng.read_tbt(from_tbt)
compare_tbt(original_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")

new_tbt = read_tbt(from_tbt, datatype="madng")
compare_tbt(original_tbt, new_tbt, no_binary=True)
assert original_tbt.date == new_tbt.date

def test_error_ng(_error_file):
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):
return INPUTS_DIR / "madng" / "fodo_track.tfs"

@pytest.fixture
def _error_file(tmp_path):
return INPUTS_DIR / "madng" / "fodo_track_error.tfs"
2 changes: 1 addition & 1 deletion turn_by_turn/__init__.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@
__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.7.2"
__version__ = "0.8.0"
__author__ = "pylhc"
__author_email__ = "[email protected]"
__license__ = "MIT"
Expand Down
7 changes: 4 additions & 3 deletions turn_by_turn/io.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,15 +2,15 @@
IO
--
This module contains high-level I/O functions to read and write tur-by-turn data objects in different
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.
"""
import logging
from pathlib import Path
from typing import Union, Any

from turn_by_turn import ascii, doros, esrf, iota, lhc, ptc, sps, trackone
from turn_by_turn import ascii, doros, esrf, iota, lhc, ptc, sps, trackone, madng
from turn_by_turn.ascii import write_ascii
from turn_by_turn.errors import DataTypeError
from turn_by_turn.structures import TbtData
Expand All @@ -29,8 +29,9 @@
ptc=ptc,
trackone=trackone,
ascii=ascii,
madng=madng,
)
WRITERS = ("lhc", "sps", "doros", "doros_positions", "doros_oscillations", "ascii") # implemented writers
WRITERS = ("lhc", "sps", "doros", "doros_positions", "doros_oscillations", "ascii", "madng") # implemented writers

write_lhc_ascii = write_ascii # Backwards compatibility <0.4

Expand Down
Loading

0 comments on commit cd27bef

Please sign in to comment.