Example-02: Workflow (MADX)
[1]:
from pathlib import Path
from model.command.external import load_lattice
from model.command.external import rift_lattice
from model.command.external import text_lattice
from model.command.external import load_tfs
from model.command.external import convert
from model.command.external import add_rc
[2]:
# Given some initial MADX lattice file (FODO)
file = Path('initial.madx')
with file.open('r') as stream:
print(stream.read())
# Several regular elements are defined
# HEAD and TAIL should appear as the first and the last elements
# All elements should be defined on a single line with numerical parameters
# Lattice should be defined using lines
# Comma after element type is mandatory
# Comments appearing after definitions should also represent an element definition
DR: DRIFT, L=2.0;
BM: SBEND, L=1.0, ANGLE=0.17453292519943295;
QF: QUADRUPOLE, L=1.0, K1=+0.2;
QD: QUADRUPOLE, L=0.5, K1=-0.2;
M: MONITOR,;
HEAD: MARKER,; ! TEST: DRIFT,
TAIL: MARKER,; ! TEST: DRIFT,
FODO: LINE=(HEAD, M, QD, DR, BM, DR, QF, DR, BM, DR, QD, TAIL) ;
[3]:
# If element and beamline definitions comply with the above requirements
# The lattice file can be loaded as a python dictionary
lattice = load_lattice(file)
for key, value in lattice.items():
print(key, value)
# For each element and beamline, a key-value pair in created
# Value is itself a dictionary containing all information about the original elements
# Each element parameter is casted from string to int, float or string
# Comment after element definition is saved into RC (it has a special use case, see below)
DR {'KIND': 'DRIFT', 'RC': '', 'L': 2.0}
BM {'KIND': 'SBEND', 'RC': '', 'L': 1.0, 'ANGLE': 0.17453292519943295}
QF {'KIND': 'QUADRUPOLE', 'RC': '', 'L': 1.0, 'K1': 0.2}
QD {'KIND': 'QUADRUPOLE', 'RC': '', 'L': 0.5, 'K1': -0.2}
M {'KIND': 'MONITOR', 'RC': ''}
HEAD {'KIND': 'MARKER', 'RC': 'TEST: DRIFT,'}
TAIL {'KIND': 'MARKER', 'RC': 'TEST: DRIFT,'}
FODO {'KIND': 'LINE', 'SEQUENCE': ['HEAD', 'M', 'QD', 'DR', 'BM', 'DR', 'QF', 'DR', 'BM', 'DR', 'QD', 'TAIL']}
[4]:
# Error lattice is defined by a set of linear transformations between selected locations
# Each locations can be a MONITOR (beam observation) or a VIRTUAL (error)
# Two special locatons (HEAD and TAIL) should present in the lattice
# Using the above dictionary representation, new observation locations can be inserted
# Locations are inserted at the middle of selected elements (selected by type or name)
# Selected elements are splitted in half and renamed, old names are binded to beamlines
# Original element definitions are added to created location RC
# Typicaly, monitor locations correspond to MONITOR elements, but new monitor elements can be also inserted
# Virtual locations can be inserted into quadrupole or other elements to represent errors
lattice = rift_lattice(lattice,
'MONITOR',
'MARKER',
['DRIFT'],
['SBEND', 'QUADRUPOLE'],
exclude_virtual=['QD'])
for key, value in lattice.items():
print(key, value)
M_DR {'KIND': 'MONITOR', 'RC': ['DR', {'KIND': 'DRIFT', 'L': 2.0}]}
H_DR {'KIND': 'DRIFT', 'L': 1.0}
DR {'KIND': 'LINE', 'SEQUENCE': ['H_DR', 'M_DR', 'H_DR']}
V_BM {'KIND': 'MARKER', 'RC': ['BM', {'KIND': 'SBEND', 'L': 1.0, 'ANGLE': 0.17453292519943295}]}
H_BM {'KIND': 'SBEND', 'L': 0.5, 'ANGLE': 0.08726646259971647}
BM {'KIND': 'LINE', 'SEQUENCE': ['H_BM', 'V_BM', 'H_BM']}
V_QF {'KIND': 'MARKER', 'RC': ['QF', {'KIND': 'QUADRUPOLE', 'L': 1.0, 'K1': 0.2}]}
H_QF {'KIND': 'QUADRUPOLE', 'L': 0.5, 'K1': 0.2}
QF {'KIND': 'LINE', 'SEQUENCE': ['H_QF', 'V_QF', 'H_QF']}
QD {'KIND': 'QUADRUPOLE', 'RC': '', 'L': 0.5, 'K1': -0.2}
M {'KIND': 'MONITOR', 'RC': ''}
HEAD {'KIND': 'MARKER', 'RC': 'TEST: DRIFT,'}
TAIL {'KIND': 'MARKER', 'RC': 'TEST: DRIFT,'}
FODO {'KIND': 'LINE', 'SEQUENCE': ['HEAD', 'M', 'QD', 'DR', 'BM', 'DR', 'QF', 'DR', 'BM', 'DR', 'QD', 'TAIL']}
[5]:
# Modified lattice can be converted to text
# Comments are added to locations while original comments are preserved
text = text_lattice('MADX', lattice, rc=True)
print(text)
M_DR: MONITOR,; ! DR: DRIFT, L=2.0,;
H_DR: DRIFT, L=1.0,;
DR: LINE=(H_DR, M_DR, H_DR);
V_BM: MARKER,; ! BM: SBEND, L=1.0, ANGLE=0.17453292519943295,;
H_BM: SBEND, L=0.5, ANGLE=0.08726646259971647,;
BM: LINE=(H_BM, V_BM, H_BM);
V_QF: MARKER,; ! QF: QUADRUPOLE, L=1.0, K1=0.2,;
H_QF: QUADRUPOLE, L=0.5, K1=0.2,;
QF: LINE=(H_QF, V_QF, H_QF);
QD: QUADRUPOLE, L=0.5, K1=-0.2,;
M: MONITOR,;
HEAD: MARKER,; ! TEST: DRIFT,
TAIL: MARKER,; ! TEST: DRIFT,
FODO: LINE=(HEAD, M, QD, DR, BM, DR, QF, DR, BM, DR, QD, TAIL);
[6]:
# Compute TWISS parameters using MADX
# TWISS command is appended to modified lattice
task = """
BEAM;
USE, PERIOD=FODO;
SET,FORMAT="20.20f","-20s";
TWISS;
WRITE,TABLE=TWISS,FILE="final.tfs";
RETURN;
""" ;
with Path('final.madx').open('w') as stream:
stream.write(text)
stream.write(task)
!madx final.madx > /dev/null
[7]:
# Load lattice can be also loaded from file
# Original comments will be parsed as elements (look at HEAD and TAIL)
# Empty RC will be nested in this case
file = Path('final.madx')
with file.open('w') as stream:
stream.write(text)
lattice = load_lattice(file, rc=True)
for key, value in lattice.items():
print(key, value)
M_DR {'KIND': 'MONITOR', 'RC': ['DR', {'KIND': 'DRIFT', 'RC': '', 'L': 2.0}]}
H_DR {'KIND': 'DRIFT', 'RC': ['', {'KIND': '', 'RC': ''}], 'L': 1.0}
DR {'KIND': 'LINE', 'SEQUENCE': ['H_DR', 'M_DR', 'H_DR']}
V_BM {'KIND': 'MARKER', 'RC': ['BM', {'KIND': 'SBEND', 'RC': '', 'L': 1.0, 'ANGLE': 0.17453292519943295}]}
H_BM {'KIND': 'SBEND', 'RC': ['', {'KIND': '', 'RC': ''}], 'L': 0.5, 'ANGLE': 0.08726646259971647}
BM {'KIND': 'LINE', 'SEQUENCE': ['H_BM', 'V_BM', 'H_BM']}
V_QF {'KIND': 'MARKER', 'RC': ['QF', {'KIND': 'QUADRUPOLE', 'RC': '', 'L': 1.0, 'K1': 0.2}]}
H_QF {'KIND': 'QUADRUPOLE', 'RC': ['', {'KIND': '', 'RC': ''}], 'L': 0.5, 'K1': 0.2}
QF {'KIND': 'LINE', 'SEQUENCE': ['H_QF', 'V_QF', 'H_QF']}
QD {'KIND': 'QUADRUPOLE', 'RC': ['', {'KIND': '', 'RC': ''}], 'L': 0.5, 'K1': -0.2}
M {'KIND': 'MONITOR', 'RC': ['', {'KIND': '', 'RC': ''}]}
HEAD {'KIND': 'MARKER', 'RC': ['TEST', {'KIND': 'DRIFT', 'RC': ''}]}
TAIL {'KIND': 'MARKER', 'RC': ['TEST', {'KIND': 'DRIFT', 'RC': ''}]}
FODO {'KIND': 'LINE', 'SEQUENCE': ['HEAD', 'M', 'QD', 'DR', 'BM', 'DR', 'QF', 'DR', 'BM', 'DR', 'QD', 'TAIL']}
[8]:
# TWISS results can be loaded into python dictionaries
data = Path('final.tfs')
parameters, columns = load_tfs(data)
[9]:
# Optics data can be converted into model table
# Note, all locations have different name
# If an element appear several times in a line, locations are renamed
table = convert(columns, 'TFS', ['MONITOR'], ['MARKER'], rc=True)
table
[9]:
{'HEAD': {'TYPE': 'VIRTUAL',
'S': 0.0,
'BX': 4.287017735718936,
'AX': -3.338e-16,
'FX': 0.0,
'BY': 19.818489282044894,
'AY': -2.975e-17,
'FY': 0.0,
'DQX': 1.5490410441348796,
'DPX': 5.551e-17,
'DQY': 0.0,
'DPY': -0.0,
'RC': None},
'M': {'TYPE': 'MONITOR',
'S': 0.0,
'BX': 4.287017735718936,
'AX': -3.338e-16,
'FX': 0.0,
'BY': 19.818489282044894,
'AY': -2.975e-17,
'FY': 0.0,
'DQX': 1.5490410441348796,
'DPX': 5.551e-17,
'DQY': 0.0,
'DPY': -0.0,
'RC': None},
'M_DR': {'TYPE': 'MONITOR',
'S': 1.5,
'BX': 5.980356480296974,
'AX': -0.8524040348462865,
'FX': 0.3068185833848285,
'BY': 15.315159900296642,
'AY': 1.6491678474974667,
'FY': 0.08453149247893033,
'DQX': 1.744126900814982,
'DPX': 0.1561982029636459,
'DQY': 0.0,
'DPY': 0.0,
'RC': None},
'V_BM': {'TYPE': 'VIRTUAL',
'S': 3.0,
'BX': 9.120629409698159,
'AX': -1.1465687400023221,
'FX': 0.5107301354647854,
'BY': 10.914137614299745,
'AY': 1.2848470098337954,
'FY': 0.20081415045191314,
'DQX': 1.992896582518225,
'DPX': 0.21385269164968151,
'DQY': 0.0,
'DPY': 0.0,
'RC': None},
'M_DR_1': {'TYPE': 'MONITOR',
'S': 4.5,
'BX': 12.70896979457465,
'AX': -1.3363966727488819,
'FX': 0.6500917125372264,
'BY': 7.606077841293864,
'AY': 0.9205261721701236,
'FY': 0.3661997931359615,
'DQX': 2.3837861006470566,
'DPX': 0.26987963222618605,
'DQY': 0.0,
'DPY': 0.0,
'RC': None},
'V_QF': {'TYPE': 'VIRTUAL',
'S': 6.0,
'BX': 16.392007194825762,
'AX': -8.0415e-16,
'FX': 0.7521811118347426,
'BY': 5.674619594695141,
'AY': -2.5871e-16,
'FY': 0.601131166393595,
'DQX': 2.7214181783177667,
'DPX': 5.551e-17,
'DQY': 0.0,
'DPY': 0.0,
'RC': None},
'M_DR_2': {'TYPE': 'MONITOR',
'S': 7.5,
'BX': 12.708969794574653,
'AX': 1.336396672748881,
'FX': 0.8542705111322586,
'BY': 7.606077841293868,
'AY': -0.9205261721701244,
'FY': 0.8360625396512283,
'DQX': 2.3837861006470566,
'DPX': -0.26987963222618594,
'DQY': 0.0,
'DPY': 0.0,
'RC': None},
'V_BM_1': {'TYPE': 'VIRTUAL',
'S': 9.0,
'BX': 9.120629409698164,
'AX': 1.146568740002322,
'FX': 0.9936320882046995,
'BY': 10.914137614299747,
'AY': -1.2848470098337954,
'FY': 1.0014481823352765,
'DQX': 1.9928965825182252,
'DPX': -0.21385269164968146,
'DQY': 0.0,
'DPY': 0.0,
'RC': None},
'M_DR_3': {'TYPE': 'MONITOR',
'S': 10.5,
'BX': 5.980356480296976,
'AX': 0.8524040348462865,
'FX': 1.1975436402846564,
'BY': 15.315159900296639,
'AY': -1.6491678474974665,
'FY': 1.1177308403082595,
'DQX': 1.7441269008149822,
'DPX': -0.1561982029636459,
'DQY': 0.0,
'DPY': 0.0,
'RC': None},
'TAIL': {'TYPE': 'VIRTUAL',
'S': 12.0,
'BX': 4.287017735718939,
'AX': 3.6486e-16,
'FX': 1.5043622236694847,
'BY': 19.81848928204488,
'AY': -5.5051e-16,
'FY': 1.2022623327871897,
'DQX': 1.54904104413488,
'DPX': -2.776e-17,
'DQY': 0.0,
'DPY': 0.0,
'RC': None}}
[10]:
# RC parameter from lattice data can be added to model table
# Configuration table can be saved using util.save
table = add_rc(table, lattice)
table
[10]:
{'HEAD': {'TYPE': 'VIRTUAL',
'S': 0.0,
'BX': 4.287017735718936,
'AX': -3.338e-16,
'FX': 0.0,
'BY': 19.818489282044894,
'AY': -2.975e-17,
'FY': 0.0,
'DQX': 1.5490410441348796,
'DPX': 5.551e-17,
'DQY': 0.0,
'DPY': -0.0,
'RC': None},
'M': {'TYPE': 'MONITOR',
'S': 0.0,
'BX': 4.287017735718936,
'AX': -3.338e-16,
'FX': 0.0,
'BY': 19.818489282044894,
'AY': -2.975e-17,
'FY': 0.0,
'DQX': 1.5490410441348796,
'DPX': 5.551e-17,
'DQY': 0.0,
'DPY': -0.0,
'RC': None},
'M_DR': {'TYPE': 'MONITOR',
'S': 1.5,
'BX': 5.980356480296974,
'AX': -0.8524040348462865,
'FX': 0.3068185833848285,
'BY': 15.315159900296642,
'AY': 1.6491678474974667,
'FY': 0.08453149247893033,
'DQX': 1.744126900814982,
'DPX': 0.1561982029636459,
'DQY': 0.0,
'DPY': 0.0,
'RC': ['DR', {'KIND': 'DRIFT', 'L': 2.0}]},
'V_BM': {'TYPE': 'VIRTUAL',
'S': 3.0,
'BX': 9.120629409698159,
'AX': -1.1465687400023221,
'FX': 0.5107301354647854,
'BY': 10.914137614299745,
'AY': 1.2848470098337954,
'FY': 0.20081415045191314,
'DQX': 1.992896582518225,
'DPX': 0.21385269164968151,
'DQY': 0.0,
'DPY': 0.0,
'RC': ['BM', {'KIND': 'SBEND', 'L': 1.0, 'ANGLE': 0.17453292519943295}]},
'M_DR_1': {'TYPE': 'MONITOR',
'S': 4.5,
'BX': 12.70896979457465,
'AX': -1.3363966727488819,
'FX': 0.6500917125372264,
'BY': 7.606077841293864,
'AY': 0.9205261721701236,
'FY': 0.3661997931359615,
'DQX': 2.3837861006470566,
'DPX': 0.26987963222618605,
'DQY': 0.0,
'DPY': 0.0,
'RC': ['DR', {'KIND': 'DRIFT', 'L': 2.0}]},
'V_QF': {'TYPE': 'VIRTUAL',
'S': 6.0,
'BX': 16.392007194825762,
'AX': -8.0415e-16,
'FX': 0.7521811118347426,
'BY': 5.674619594695141,
'AY': -2.5871e-16,
'FY': 0.601131166393595,
'DQX': 2.7214181783177667,
'DPX': 5.551e-17,
'DQY': 0.0,
'DPY': 0.0,
'RC': ['QF', {'KIND': 'QUADRUPOLE', 'L': 1.0, 'K1': 0.2}]},
'M_DR_2': {'TYPE': 'MONITOR',
'S': 7.5,
'BX': 12.708969794574653,
'AX': 1.336396672748881,
'FX': 0.8542705111322586,
'BY': 7.606077841293868,
'AY': -0.9205261721701244,
'FY': 0.8360625396512283,
'DQX': 2.3837861006470566,
'DPX': -0.26987963222618594,
'DQY': 0.0,
'DPY': 0.0,
'RC': ['DR', {'KIND': 'DRIFT', 'L': 2.0}]},
'V_BM_1': {'TYPE': 'VIRTUAL',
'S': 9.0,
'BX': 9.120629409698164,
'AX': 1.146568740002322,
'FX': 0.9936320882046995,
'BY': 10.914137614299747,
'AY': -1.2848470098337954,
'FY': 1.0014481823352765,
'DQX': 1.9928965825182252,
'DPX': -0.21385269164968146,
'DQY': 0.0,
'DPY': 0.0,
'RC': ['BM', {'KIND': 'SBEND', 'L': 1.0, 'ANGLE': 0.17453292519943295}]},
'M_DR_3': {'TYPE': 'MONITOR',
'S': 10.5,
'BX': 5.980356480296976,
'AX': 0.8524040348462865,
'FX': 1.1975436402846564,
'BY': 15.315159900296639,
'AY': -1.6491678474974665,
'FY': 1.1177308403082595,
'DQX': 1.7441269008149822,
'DPX': -0.1561982029636459,
'DQY': 0.0,
'DPY': 0.0,
'RC': ['DR', {'KIND': 'DRIFT', 'L': 2.0}]},
'TAIL': {'TYPE': 'VIRTUAL',
'S': 12.0,
'BX': 4.287017735718939,
'AX': 3.6486e-16,
'FX': 1.5043622236694847,
'BY': 19.81848928204488,
'AY': -5.5051e-16,
'FY': 1.2022623327871897,
'DQX': 1.54904104413488,
'DPX': -2.776e-17,
'DQY': 0.0,
'DPY': 0.0,
'RC': None}}