/
jspec.py
152 lines (133 loc) · 5.26 KB
/
jspec.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
# -*- coding: utf-8 -*-
"""Wrapper to run JSPEC from the command line.
:copyright: Copyright (c) 2017 RadiaSoft LLC. All Rights Reserved.
:license: http://www.apache.org/licenses/LICENSE-2.0.html
"""
from __future__ import absolute_import, division, print_function
from pykern import pkio
from pykern import pksubprocess
from pykern.pkdebug import pkdp, pkdc, pkdlog
from sirepo import simulation_db
from sirepo.template import sdds_util, template_common
import os.path
import re
import sirepo.template.jspec as template
import sirepo.sim_data
_SIM_DATA = sirepo.sim_data.get_class('jspec')
def run(cfg_dir):
with pkio.save_chdir(cfg_dir):
data = simulation_db.read_json(template_common.INPUT_BASE_NAME)
if data['report'] == 'twissReport':
simulation_db.write_result(_extract_twiss_report(data))
elif data['report'] == 'rateCalculationReport':
text = _run_jspec(data)
res = {
#TODO(pjm): x_range is needed for sirepo-plotting.js, need a better valid-data check
'x_range': [],
'rate': [],
}
for line in text.split("\n"):
m = re.match(r'^(.*? rate.*?)\:\s+(\S+)\s+(\S+)\s+(\S+)', line)
if m:
row = [m.group(1), [m.group(2), m.group(3), m.group(4)]]
row[0] = re.sub('\(', '[', row[0]);
row[0] = re.sub('\)', ']', row[0]);
res['rate'].append(row)
simulation_db.write_result(res)
else:
assert False, 'unknown report: {}'.format(data['report'])
def run_background(cfg_dir):
with pkio.save_chdir(cfg_dir):
_run_jspec(simulation_db.read_json(template_common.INPUT_BASE_NAME))
simulation_db.write_result({})
def _elegant_to_madx(ring):
# if the lattice source is an elegant twiss file, convert it to MAD-X format
if ring['latticeSource'] == 'madx':
return _SIM_DATA.lib_file_name_with_model_field('ring', 'lattice', ring['lattice'])
if ring['latticeSource'] == 'elegant':
elegant_twiss_file = _SIM_DATA.lib_file_name_with_model_field('ring', 'elegantTwiss', ring['elegantTwiss'])
else: # elegant-sirepo
if 'elegantSirepo' not in ring or not ring['elegantSirepo']:
raise RuntimeError('elegant simulation not selected')
elegant_twiss_file = _SIM_DATA.JSPEC_ELEGANT_TWISS_FILENAME
if not os.path.exists(elegant_twiss_file):
raise RuntimeError('elegant twiss output unavailable. Run elegant simulation.')
sdds_util.twiss_to_madx(elegant_twiss_file, template.JSPEC_TWISS_FILENAME)
return template.JSPEC_TWISS_FILENAME
_X_FIELD = 's'
_FIELD_UNITS = {
'betx': 'm',
#'alfx': '',
'mux': 'rad/2π',
'dx': 'm',
#'dpx': '',
'bety': 'm',
#'alfy': '',
'muy': 'rad/2π',
'dx': 'm',
#'dpx': '',
}
def _extract_twiss_report(data):
report = data['models'][data['report']]
report['x'] = _X_FIELD
values = _parse_madx(_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 _float_from_str(v):
# handle misformatted floats, ex. -9.29135e-00E-25
v = re.sub(r'(e[+\-]\d+)(e[+\-]\d+)', r'\1', v, flags=re.IGNORECASE)
return float(v)
def _float_list(ar):
return map(lambda x: _float_from_str(x), ar)
def _run_jspec(data):
_elegant_to_madx(data['models']['ring'])
exec(pkio.read_text(template_common.PARAMETERS_PYTHON_FILE), locals(), locals())
jspec_filename = template.JSPEC_INPUT_FILENAME
pkio.write_text(jspec_filename, jspec_file)
pksubprocess.check_call_with_signals(['jspec', jspec_filename], msg=pkdlog, output=template.JSPEC_LOG_FILE)
return pkio.read_text(template.JSPEC_LOG_FILE)
def _parse_madx(tfs_file):
text = pkio.read_text(tfs_file)
mode = 'header'
col_names = []
rows = []
for line in text.split("\n"):
if mode == 'header':
# header row starts with *
if re.search('^\*\s', line):
col_names = re.split('\s+', line)
col_names = col_names[1:]
mode = 'data'
elif mode == 'data':
# data rows after header, start with blank
if re.search('^\s+\S', line):
data = re.split('\s+', line)
rows.append(data[1:])
res = dict(map(lambda x: (x.lower(), []), col_names))
for i in range(len(col_names)):
name = col_names[i].lower()
if name:
for row in rows:
res[name].append(row[i])
# special case if dy and/or dpy are missing, default to 0s
for opt_col in ('dy', 'dpy'):
if opt_col not in res:
res[opt_col] = ['0'] * len(res['dx'])
return res