def extract_parameter_report(data, run_dir, filename=_TWISS_OUTPUT_FILE):
t = madx_parser.parse_tfs_file(run_dir.join(filename))
plots = []
m = data.models[data.report]
for f in ('y1', 'y2', 'y3'):
if m[f] == 'None':
continue
plots.append(
PKDict(field=m[f],
points=to_floats(t[m[f]]),
label=_field_label(m[f])), )
x = m.get('x') or 's'
res = template_common.parameter_plot(
to_floats(t[x]), plots, m, PKDict(
y_label='',
x_label=_field_label(x),
))
if 'betx' in t and 'bety' in t and 'alfx' in t and 'alfy' in t:
res.initialTwissParameters = PKDict(
betx=t.betx[0],
bety=t.bety[0],
alfx=t.alfx[0],
alfy=t.alfx[0],
)
return res
def _extract_twiss_report(data):
report = data['models'][data['report']]
report['x'] = _X_FIELD
values = madx_parser.parse_tfs_file(
_elegant_to_madx(data['models']['ring']))
# special case if dy and/or dpy are missing, default to 0s
for opt_col in ('dy', 'dpy'):
if opt_col not in values and 'dx' in values:
values[opt_col] = ['0'] * len(values['dx'])
x = _float_list(values[report['x']])
y_range = None
plots = []
for f in ('y1', 'y2', 'y3'):
if report[f] == 'none':
continue
plots.append({
'points':
_float_list(values[report[f]]),
'label':
'{} [{}]'.format(report[f], _FIELD_UNITS[report[f]])
if report[f] in _FIELD_UNITS else report[f],
})
return {
'title': '',
'x_range': [min(x), max(x)],
'y_label': '',
'x_label': '{} [{}]'.format(report['x'], 'm'),
'x_points': x,
'plots': plots,
'y_range': template_common.compute_plot_color_and_range(plots),
}
def _extract_twiss_report(data):
report = data['models'][data['report']]
report['x'] = _X_FIELD
values = madx_parser.parse_tfs_file(
_elegant_to_madx(data['models']['ring']))
x = _float_list(values[report['x']])
y_range = None
plots = []
for f in ('y1', 'y2', 'y3'):
if report[f] == 'none':
continue
plots.append({
'points':
_float_list(values[report[f]]),
'label':
'{} [{}]'.format(report[f], _FIELD_UNITS[report[f]])
if report[f] in _FIELD_UNITS else report[f],
})
return {
'title': '',
'x_range': [min(x), max(x)],
'y_label': '',
'x_label': '{} [{}]'.format(report['x'], 'm'),
'x_points': x,
'plots': plots,
'y_range': template_common.compute_plot_color_and_range(plots),
}
def _extract_report_ptcAnimation(data, run_dir, filename):
m = data.models[data.report]
t = madx_parser.parse_tfs_file(run_dir.join(filename))
return template_common.heatmap(
[_to_floats(t[m.x]), _to_floats(t[m.y])],
m,
PKDict(
x_label=_FIELD_LABEL[m.x],
y_label=_FIELD_LABEL[m.y],
),
)
def test_extract_tfs_pages():
from pykern.pkunit import pkeq
from sirepo.template import madx_parser
path = pkunit.data_dir().join('ptc_track.file.tfs')
header = madx_parser.parse_tfs_file(path, header_only=True)
pkeq(header, ['number', 'turn', 'x', 'px', 'y', 'py', 't', 'pt', 's', 'e'])
info = madx_parser.parse_tfs_page_info(path)
pkeq(
info[3],
dict(
name='BPMY1',
turn='1',
s='2.72',
),
)
res = madx_parser.parse_tfs_file(path, want_page=3)
pkeq(len(res.s), 5)
pkeq(res.s[0], info[3].s)
res = madx_parser.parse_tfs_file(path)
pkeq(len(res.s), 75)
def _extract_report_twissReport(data, run_dir, filename=_TWISS_OUTPUT_FILE):
t = madx_parser.parse_tfs_file(run_dir.join(filename))
plots = []
m = data.models[data.report]
for f in ('y1', 'y2', 'y3'):
if m[f] == 'none':
continue
plots.append(
PKDict(field=m[f], points=_to_floats(t[m[f]]), label=_FIELD_LABEL[m[f]])
)
x = m.get('x') or 's'
return template_common.parameter_plot(
_to_floats(t[x]),
plots,
m,
PKDict(
y_label='',
x_label=_FIELD_LABEL[x],
)
)
def _file_info(filename, run_dir, file_id):
path = str(run_dir.join(filename))
plottable = []
tfs = madx_parser.parse_tfs_file(path)
for f in tfs:
if f in _ALPHA_COLUMNS:
continue
v = to_floats(tfs[f])
if np.any(v):
plottable.append(f)
count = 1
if 'turn' in tfs:
info = madx_parser.parse_tfs_page_info(path)
count = len(info)
return PKDict(
modelKey='elementAnimation{}'.format(file_id),
filename=filename,
isHistogram=not _is_parameter_report_file(filename),
plottableColumns=plottable,
pageCount=count,
)
def test_parse_tfs_file():
from pykern.pkunit import pkeq
from sirepo.template import madx_parser
path = pkunit.data_dir().join('twiss.file.tfs')
res = madx_parser.parse_tfs_file(path)
pkeq(
res.betx,
['11.639', '14.73204923'],
)
pkeq(
res.name,
['"RING$START"', '"D3"'],
)
pkeq(
res.sig44,
['1.048547761e-07', '1.048547761e-07'],
)
pkeq(
res.n1,
['0', '0'],
)
def _extract_report_elementAnimation(data, run_dir, filename):
if _is_parameter_report_file(filename):
return extract_parameter_report(data, run_dir, filename)
m = data.models[data.report]
t = madx_parser.parse_tfs_file(run_dir.join(filename),
want_page=m.frameIndex)
info = madx_parser.parse_tfs_page_info(
run_dir.join(filename))[m.frameIndex]
return template_common.heatmap(
[to_floats(t[m.x]), to_floats(t[m.y1])],
m,
PKDict(
x_label=_field_label(m.x),
y_label=_field_label(m.y1),
title='{}-{} at {}m, {} turn {}'.format(
m.x,
m.y1,
info.s,
info.name,
info.turn,
),
),
)
