def test_from_elegant_to_madx_and_back():
from pykern import pkio
from pykern.pkunit import pkeq
from sirepo.template import elegant, madx, madx_converter, madx_parser
with pkunit.save_chdir_work() as d:
for name in ('SPEAR3', 'Compact Storage Ring',
'Los Alamos Proton Storage Ring'):
data = _example_data(name)
mad = madx_parser.parse_file(
elegant.python_source_for_model(data, 'madx'))
madx._fixup_madx(mad)
outfile = name.lower().replace(' ', '-') + '.madx'
actual = madx.python_source_for_model(mad, None)
pkio.write_text(outfile, actual)
e = pkunit.data_dir().join(outfile)
expect = pkio.read_text(e)
pkeq(expect, actual, 'diff {} {}', e, d.join(outfile))
data = madx_parser.parse_file(actual)
lattice = madx_converter.from_madx(elegant.SIM_TYPE, data)
outfile = name.lower().replace(' ', '-') + '.lte'
actual = elegant.python_source_for_model(lattice, None)
pkio.write_text(outfile, actual)
e = pkunit.data_dir().join(outfile)
expect = pkio.read_text(e)
pkeq(expect, actual, 'diff {} {}', e, d.join(outfile))
def import_file(req, unit_test_mode=False, **kwargs):
from sirepo.template import opal_parser
text = pkcompat.from_bytes(req.file_stream.read())
if re.search(r'\.in$', req.filename, re.IGNORECASE):
data, input_files = opal_parser.parse_file(text, filename=req.filename)
missing_files = []
for infile in input_files:
if not _SIM_DATA.lib_file_exists(infile.lib_filename):
missing_files.append(infile)
if len(missing_files):
return PKDict(
error='Missing data files',
missingFiles=missing_files,
)
elif re.search(r'\.madx$', req.filename, re.IGNORECASE):
from sirepo.template import madx_converter, madx_parser
madx = madx_parser.parse_file(text)
data = madx_converter.from_madx(SIM_TYPE, madx)
data.models.simulation.name = re.sub(r'\.madx$',
'',
req.filename,
flags=re.IGNORECASE)
_fixup_madx(madx, data)
else:
raise IOError('invalid file extension, expecting .in or .madx')
return data
def import_file(req, test_data=None, **kwargs):
# input_data is passed by test cases only
input_data = test_data
text = pkcompat.from_bytes(req.file_stream.read())
if 'simulationId' in req.req_data:
input_data = simulation_db.read_simulation_json(
SIM_TYPE, sid=req.req_data.simulationId)
if re.search(r'\.ele$', req.filename, re.IGNORECASE):
data = elegant_command_importer.import_file(text)
elif re.search(r'\.lte$', req.filename, re.IGNORECASE):
data = elegant_lattice_importer.import_file(text, input_data)
if input_data:
_map_commands_to_lattice(data)
elif re.search(r'\.madx$', req.filename, re.IGNORECASE):
from sirepo.template import madx_converter, madx_parser
data = madx_converter.from_madx(
SIM_TYPE, madx_parser.parse_file(text, downcase_variables=True))
else:
raise IOError('invalid file extension, expecting .ele or .lte')
data.models.simulation.name = re.sub(r'\.(lte|ele|madx)$',
'',
req.filename,
flags=re.IGNORECASE)
if input_data and not test_data:
simulation_db.delete_simulation(
SIM_TYPE, input_data.models.simulation.simulationId)
return data
def parse_input_text(path, text=None, input_data=None):
def _map(data):
for cmd in data.models.commands:
if cmd._type == 'run_setup':
cmd.lattice = 'Lattice'
break
for cmd in data.models.commands:
if cmd._type == 'run_setup':
name = cmd.use_beamline.upper()
for bl in data.models.beamlines:
if bl.name.upper() == name:
cmd.use_beamline = bl.id
break
if text is None:
text = pkio.read_text(path)
e = path.ext.lower()
if e == '.ele':
return elegant_command_importer.import_file(text)
if e == '.lte':
data = elegant_lattice_importer.import_file(text, input_data)
if input_data:
_map(data)
return data
if e == '.madx':
from sirepo.template import madx_converter, madx_parser
return madx_converter.from_madx(
SIM_TYPE,
madx_parser.parse_file(text, downcase_variables=True),
)
raise IOError(
f'{path.basename}: invalid file format; expecting .madx, .ele, or .lte'
)
def test_parse_madx_file_downcase():
from sirepo.template import madx_parser
import re
parsed = madx_parser.parse_file('''
REAL energy = 1.6;
REAL gamma = (ENERGY + 0.0005109989) / 0.0005109989;
''', True)
assert re.search(r'energy', str(parsed))
assert not re.search(r'ENERGY', str(parsed))
def import_file(req, **kwargs):
text = pkcompat.from_bytes(req.file_stream.read())
assert re.search(r'\.madx$', req.filename, re.IGNORECASE), \
'invalid file extension, expecting .madx'
data = madx_parser.parse_file(text, downcase_variables=True)
# TODO(e-carlin): need to clean this up. copied from elegant
data.models.simulation.name = re.sub(r'\.madx$',
'',
req.filename,
flags=re.IGNORECASE)
return data
def test_parse_madx_file():
from pykern import pkio, pkjson
from pykern.pkunit import pkeq
from sirepo.template import madx, madx_parser
for name in ('particle_track', ):
actual = madx_parser.parse_file(
pkio.read_text(pkunit.data_dir().join(f'{name}.madx')))
madx._fixup_madx(actual)
del actual['version']
expect = pkjson.load_any(pkunit.data_dir().join(f'{name}.json'))
pkeq(expect, actual)
def test_parse_madx_file():
from pykern import pkio, pkjson
from pykern.pkunit import pkeq
from sirepo.template import madx, madx_parser
with pkunit.save_chdir_work():
for name in ('particle_track', 'alba'):
actual = madx_parser.parse_file(pkio.read_text(
pkunit.data_dir().join(f'{name}.madx')))
del actual['version']
outfile = f'{name}.json'
pkjson.dump_pretty(actual, outfile)
expect = pkjson.load_any(pkunit.data_dir().join(outfile))
pkeq(expect, actual)
def from_madx_text(self, text):
return self.from_madx(madx_parser.parse_file(text, self.downcase_variables))