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, ), ), )