def _generate_transverse_dist(models):
dist_type = models.beam.transversalDistribution
if dist_type == 'twiss4d':
dist = models.twissDistribution
return 'TWISS4D {} {} {} {} {} {}'.format(dist.horizontalAlpha,
dist.horizontalBeta,
dist.horizontalEmittance,
dist.verticalAlpha,
dist.verticalBeta,
dist.verticalEmittance)
if dist_type == 'sph2d':
dist = models.sphericalDistribution
if dist.curvature == 'flat':
dist.curvatureFactor = 0
return 'SPH2D {} {} {}'.format(dist.radialLimit, dist.curvatureFactor,
dist.thermalEmittance)
if dist_type == 'ell2d':
dist = models.ellipticalDistribution
return 'ELL2D {} {} {} {}'.format(dist.aX, dist.bY, dist.rotationAngle,
dist.rmsDeviationFactor)
beam = models.beam
if dist_type == 'file2d':
return 'FILE2D {}'.format(
template_common.lib_file_name('beam', 'transversalFile2d',
beam.transversalFile2d))
if dist_type == 'file4d':
return 'FILE4D {}'.format(
template_common.lib_file_name('beam', 'transversalFile4d',
beam.transversalFile4d))
raise RuntimeError('unknown transverse distribution: {}'.format(dist_type))
def _generate_longitude_dist(models):
dist_type = models.beam.longitudinalDistribution
if dist_type == 'norm2d':
dist = models.energyPhaseDistribution
if dist.distributionType == 'uniform':
return 'NORM2D {} {} {} {}'.format(dist.meanEnergy,
dist.energySpread,
dist.meanPhase,
dist.phaseLength)
if dist.distributionType == 'gaussian':
return 'NORM2D {} {} {} {} {} {}'.format(
dist.meanEnergy, dist.energySpread, dist.energyDeviation,
dist.meanPhase, dist.phaseLength, dist.phaseDeviation)
raise RuntimeError('unknown longitudinal distribution type: {}'.format(
models.longitudinalDistribution.distributionType))
if dist_type == 'file1d':
return 'FILE1D {} {}'.format(
template_common.lib_file_name('beam', 'longitudinalFile1d',
models.beam.longitudinalFile1d),
_generate_energy_phase_distribution(
models.energyPhaseDistribution),
)
if dist_type == 'file2d':
return 'FILE2D {}'.format(
template_common.lib_file_name('beam', 'transversalFile2d',
beam.transversalFile2d))
raise RuntimeError('unknown longitudinal distribution: {}'.format(
models.beam.longitudinalDistribution))
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 template_common.lib_file_name('ring', 'lattice', ring['lattice'])
if ring['latticeSource'] == 'elegant':
elegant_twiss_file = template_common.lib_file_name('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 = template.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
def lib_files(data, source_lib):
res = []
ring = data['models']['ring']
lattice_source = ring['latticeSource']
if lattice_source == 'madx':
res.append(template_common.lib_file_name('ring', 'lattice', ring['lattice']))
elif lattice_source == 'elegant':
res.append(template_common.lib_file_name('ring', 'elegantTwiss', ring['elegantTwiss']))
res = template_common.filename_to_path(res, source_lib)
if lattice_source == 'elegant-sirepo' and 'elegantSirepo' in ring:
f = _elegant_dir().join(ring['elegantSirepo'], _ELEGANT_TWISS_PATH)
if f.exists():
res.append(f)
return res
def _iterator_lattice_elements(state,
model,
element_schema=None,
field_name=None):
# only interested in elements, not commands
if '_type' in model:
return
if element_schema:
state['field_index'] += 1
if field_name in ['name', 'type', '_id'] or re.search(
'(X|Y|File)$', field_name):
return
value = model[field_name]
default_value = element_schema[2]
if value is not None and default_value is not None:
if str(value) != str(default_value):
el_type = element_schema[1]
if model['type'] == 'SCRIPT' and field_name == 'command':
for f in ('commandFile', 'commandInputFile'):
if f in model and model[f]:
fn = template_common.lib_file_name(
model['type'], f, model[f])
value = re.sub(r'\b' + re.escape(model[f]) + r'\b',
fn, value)
if model['commandFile']:
value = './' + value
if el_type == 'RPNValue':
value = _format_rpn_value(value)
if el_type.startswith('InputFile'):
value = template_common.lib_file_name(
model['type'], field_name, value)
if el_type == 'InputFileXY':
value += '={}+{}'.format(model[field_name + 'X'],
model[field_name + 'Y'])
elif el_type == 'OutputFile':
value = state['filename_map']['{}{}{}'.format(
model['_id'], _FILE_ID_SEP, state['field_index'])]
if not _is_numeric(el_type, value):
value = '"{}"'.format(value)
state['lattice'] += '{}={},'.format(field_name, value)
else:
state['field_index'] = 0
if state['lattice']:
state['lattice'] = state['lattice'][:-1]
state['lattice'] += "\n"
state['lattice'] += '"{}": {},'.format(model['name'].upper(),
model['type'])
state['beamline_map'][model['_id']] = model['name']
def python_source_for_model(data, model):
ring = data['models']['ring']
elegant_twiss_file = None
if ring['latticeSource'] == 'elegant':
elegant_twiss_file = template_common.lib_file_name(
'ring', 'elegantTwiss', ring['elegantTwiss'])
elif ring['latticeSource'] == 'elegant-sirepo':
elegant_twiss_file = ELEGANT_TWISS_FILENAME
convert_twiss_to_tfs = ''
if elegant_twiss_file:
convert_twiss_to_tfs = '''
from pykern import pkconfig
pkconfig.append_load_path('sirepo')
from sirepo.template import sdds_util
sdds_util.twiss_to_madx('{}', '{}')
'''.format(elegant_twiss_file, JSPEC_TWISS_FILENAME)
return '''
{}
with open('{}', 'w') as f:
f.write(jspec_file)
{}
import os
os.system('jspec {}')
'''.format(_generate_parameters_file(data), JSPEC_INPUT_FILENAME,
convert_twiss_to_tfs, JSPEC_INPUT_FILENAME)
def _generate_bunch_simulation(data, v):
for f in _SCHEMA['model']['bunch']:
info = _SCHEMA['model']['bunch'][f]
if info[1] == 'RPNValue':
field = 'bunch_{}'.format(f)
v[field] = _format_rpn_value(v[field], is_command=True)
longitudinal_method = int(data['models']['bunch']['longitudinalMethod'])
# sigma s, sigma dp, dp s coupling
if longitudinal_method == 1:
v['bunch_emit_z'] = 0
v['bunch_beta_z'] = 0
v['bunch_alpha_z'] = 0
# sigma s, sigma dp, alpha z
elif longitudinal_method == 2:
v['bunch_emit_z'] = 0
v['bunch_beta_z'] = 0
v['bunch_dp_s_coupling'] = 0
# emit z, beta z, alpha z
elif longitudinal_method == 3:
v['bunch_sigma_dp'] = 0
v['bunch_sigma_s'] = 0
v['bunch_dp_s_coupling'] = 0
if data['models']['bunchSource']['inputSource'] == 'sdds_beam':
v['bunch_beta_x'] = 5
v['bunch_beta_y'] = 5
v['bunch_alpha_x'] = 0
v['bunch_alpha_x'] = 0
if v['bunchFile_sourceFile'] and v['bunchFile_sourceFile'] != 'None':
v['bunchInputFile'] = template_common.lib_file_name(
'bunchFile', 'sourceFile', v['bunchFile_sourceFile'])
v['bunchFileType'] = _sdds_beam_type_from_file(v['bunchInputFile'])
v['bunchOutputFile'] = BUNCH_OUTPUT_FILE
return template_common.render_jinja(SIM_TYPE, v, 'bunch.py')
def _generate_bunch_simulation(data, v):
for f in _SCHEMA['model']['bunch']:
info = _SCHEMA['model']['bunch'][f]
if info[1] == 'RPNValue':
field = 'bunch_{}'.format(f)
v[field] = _format_rpn_value(v[field], is_command=True)
longitudinal_method = int(data['models']['bunch']['longitudinalMethod'])
# sigma s, sigma dp, dp s coupling
if longitudinal_method == 1:
v['bunch_emit_z'] = 0
v['bunch_beta_z'] = 0
v['bunch_alpha_z'] = 0
# sigma s, sigma dp, alpha z
elif longitudinal_method == 2:
v['bunch_emit_z'] = 0
v['bunch_beta_z'] = 0
v['bunch_dp_s_coupling'] = 0
# emit z, beta z, alpha z
elif longitudinal_method == 3:
v['bunch_sigma_dp'] = 0
v['bunch_sigma_s'] = 0
v['bunch_dp_s_coupling'] = 0
if data['models']['bunchSource']['inputSource'] == 'sdds_beam':
v['bunch_beta_x'] = 5
v['bunch_beta_y'] = 5
v['bunch_alpha_x'] = 0
v['bunch_alpha_x'] = 0
if v['bunchFile_sourceFile'] and v['bunchFile_sourceFile'] != 'None':
v['bunchInputFile'] = template_common.lib_file_name('bunchFile', 'sourceFile', v['bunchFile_sourceFile'])
v['bunchFileType'] = _sdds_beam_type_from_file(v['bunchInputFile'])
v['bunchOutputFile'] = _report_output_filename('bunchReport')
return template_common.render_jinja(SIM_TYPE, v, 'bunch.py')
def _generate_beam(models):
# BEAM SPH2D 0.564 -15 5 NORM2D 0.30 0.0000001 90 180
beam_def = models.beam.beamDefinition
if beam_def == 'transverse_longitude':
return 'BEAM {} {}'.format(_generate_transverse_dist(models), _generate_longitude_dist(models))
if beam_def == 'cst_pit':
return 'BEAM CST_PIT {} {}'.format(
template_common.lib_file_name('beam', 'cstFile', models.beam.cstFile),
'COMPRESS' if models.beam.cstCompress else '',
)
if beam_def == 'cst_pid':
return 'BEAM CST_PID {} {}'.format(
template_common.lib_file_name('beam', 'cstFile', models.beam.cstFile),
_generate_energy_phase_distribution(models.energyPhaseDistribution),
)
raise RuntimeError('invalid beam def: {}'.format(beam_def))
def _lib_file_datetime(filename):
lib_filename = template_common.lib_file_name('analysisData', 'file',
filename)
path = simulation_db.simulation_lib_dir(SIM_TYPE).join(lib_filename)
if path.exists():
return {filename: path.mtime()}
pkdlog('error, missing lib file: {}', path)
return 0
def lib_files(data, source_lib):
res = []
for el in data.models.elements:
if el.type == 'TOSCA' and el.magnetFile:
res.append(
template_common.lib_file_name('TOSCA', 'magnetFile',
el.magnetFile))
return template_common.filename_to_path(res, source_lib)
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 template_common.lib_file_name('ring', 'lattice',
ring['lattice'])
if ring['latticeSource'] == 'elegant':
elegant_twiss_file = template_common.lib_file_name(
'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 = template.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
def _simulation_files(data):
res = []
bunch = data.models.bunch
if bunch.distribution == 'file':
res.append(
template_common.lib_file_name('bunch', 'particleFile',
bunch.particleFile))
return res
def save_report_data(data, run_dir):
if 'bunchReport' in data['report']:
import synergia.bunch
with h5py.File(str(run_dir.join(OUTPUT_FILE['twissReport'])), 'r') as f:
twiss0 = dict(map(
lambda k: (k, format_float(f[k][0])),
('alpha_x', 'alpha_y', 'beta_x', 'beta_y'),
))
report = data.models[data['report']]
bunch = data.models.bunch
if bunch.distribution == 'file':
bunch_file = template_common.lib_file_name('bunch', 'particleFile', bunch.particleFile)
else:
bunch_file = OUTPUT_FILE['bunchReport']
if not run_dir.join(bunch_file).exists():
return
with h5py.File(str(run_dir.join(bunch_file)), 'r') as f:
x = f['particles'][:, getattr(synergia.bunch.Bunch, report['x'])]
y = f['particles'][:, getattr(synergia.bunch.Bunch, report['y'])]
hist, edges = np.histogramdd([x, y], template_common.histogram_bins(report['histogramBins']))
res = {
'title': '',
'x_range': [float(edges[0][0]), float(edges[0][-1]), len(hist)],
'y_range': [float(edges[1][0]), float(edges[1][-1]), len(hist[0])],
'x_label': label(report['x'], _SCHEMA['enum']['PhaseSpaceCoordinate8']),
'y_label': label(report['y'], _SCHEMA['enum']['PhaseSpaceCoordinate8']),
'z_matrix': hist.T.tolist(),
'summaryData': {
'bunchTwiss': twiss0,
},
}
else:
report_name = data['report']
x = None
plots = []
report = data['models'][report_name]
with h5py.File(str(run_dir.join(OUTPUT_FILE[report_name])), 'r') as f:
x = f['s'][:].tolist()
for yfield in ('y1', 'y2', 'y3'):
if report[yfield] == 'none':
continue
name = report[yfield]
plots.append({
'name': name,
'label': label(report[yfield], _SCHEMA['enum']['TwissParameter']),
'points': f[name][:].tolist(),
})
res = {
'title': '',
'x_range': [min(x), max(x)],
'y_range': template_common.compute_plot_color_and_range(plots),
'x_label': 's [m]',
'y_label': '',
'x_points': x,
'plots': plots,
}
simulation_db.write_result(res, run_dir=run_dir)
def _generate_element_tosca(el):
el = copy.deepcopy(el)
el['MOD'] = _MAGNET_TYPE_TO_MOD[el.meshType][el.magnetType]
if '{{ fileCount }}' in el['MOD']:
el['MOD'] = el['MOD'].replace('{{ fileCount }}', str(el['fileCount']))
el['hasFields'] = True
if '-sf' in el.magnetType or ('-f' in el.magnetType
and el['fileCount'] == 1):
filename = el['magnetFile']
if _is_zip_file(filename):
filename = el['fileNames'][0]
assert filename, 'missing file name from zip for single file magnet'
el['fileNames'] = [
template_common.lib_file_name('TOSCA', 'magnetFile', filename)
]
else:
if '-f' in el.magnetType:
file_count = el.fileCount
elif el.magnetType == '3d-mf-2v':
file_count = math.floor((el.IZ + 1) / 2)
elif el.magnetType == '3d-mf-1v':
file_count = el.IZ
else:
assert false, 'unhandled magnetType: {}'.format(el.magnetType)
el['fileNames'] = el['fileNames'][:int(file_count)]
for f in _MODEL_UNITS['TOSCA']:
el[f] = _element_value(el, f)
template = '''
'TOSCA'
0 0
{{ BNORM }} {{ XN }} {{ YN }} {{ ZN }}
{{ name -}}
{%- if flipX == '1' %} FLIP {% endif -%}
{%- if headerLineCount > 0 %} HEADER_{{ headerLineCount }} {% endif -%}
{%- if zeroBXY == '1' %} ZroBXY {% endif -%}
{%- if normalizeHelix == '1' %} RHIC_helix {% endif %}
{{ IX }} {{ IY }} {{ IZ }} {{ MOD -}}
{%- if hasFields %} {{ field1 }} {{ field2 }} {{ field3 }} {{ field4 }}{% endif %}
{%- for fileName in fileNames %}
{{ fileName -}}
{% endfor %}
{{ ID }} {{ A }} {{ B }} {{ C }} {{ Ap }} {{ Bp }} {{ Cp }} {{ App }} {{ Bpp }} {{ Cpp }}
{{ IORDRE }}
{{ XPAS }}
{% if meshType == 'cartesian' -%}
{{ KPOS }} {{ XCE }} {{ YCE }} {{ ALE }}
{% else -%}
{{ KPOS }}
{{ RE }} {{ TE }} {{ RS }} {{ TS }}
{% endif %}
'''
return 'line.add(core.FAKE_ELEM("""{}"""))\n'.format(
jinja2.Template(template).render(el))
def _generate_solenoid(models):
solenoid = models.solenoid
if solenoid.sourceDefinition == 'none':
return ''
if solenoid.sourceDefinition == 'values':
#TODO(pjm): latest version also has solenoid.fringeRegion
return 'SOLENOID {} {} {}'.format(
solenoid.fieldStrength, solenoid.length, solenoid.z0)
if solenoid.sourceDefinition == 'file':
return 'SOLENOID {}'.format(
template_common.lib_file_name('solenoid', 'solenoidFile', solenoid.solenoidFile))
raise RuntimeError('unknown solenoidDefinition: {}'.format(solenoid.sourceDefinition))
def _simulation_files(data):
res = []
solenoid = data.models.solenoid
if solenoid.sourceDefinition == 'file' and solenoid.solenoidFile:
res.append(
template_common.lib_file_name('solenoid', 'solenoidFile',
solenoid.solenoidFile))
beam = data.models.beam
if beam.beamDefinition == 'cst_pit' or beam.beamDefinition == 'cst_pid':
res.append(
template_common.lib_file_name('beam', 'cstFile', beam.cstFile))
if beam.beamDefinition == 'transverse_longitude':
if beam.transversalDistribution == 'file2d':
res.append(
template_common.lib_file_name('beam', 'transversalFile2d',
beam.transversalFile2d))
elif beam.transversalDistribution == 'file4d':
res.append(
template_common.lib_file_name('beam', 'transversalFile4d',
beam.transversalFile4d))
if beam.longitudinalDistribution == 'file1d':
res.append(
template_common.lib_file_name('beam', 'longitudinalFile1d',
beam.longitudinalFile1d))
if beam.longitudinalDistribution == 'file2d':
res.append(
template_common.lib_file_name('beam', 'longitudinalFile2d',
beam.longitudinalFile2d))
return res
def _iterator_lattice_elements(state, model, element_schema=None, field_name=None):
# only interested in elements, not commands
if '_type' in model:
return
if element_schema:
state['field_index'] += 1
if field_name in ['name', 'type', '_id'] or re.search('(X|Y|File)$', field_name):
return
value = model[field_name]
default_value = element_schema[2]
if value is not None and default_value is not None:
if str(value) != str(default_value):
if model['type'] == 'SCRIPT' and field_name == 'command':
for f in ('commandFile', 'commandInputFile'):
if f in model and model[f]:
fn = template_common.lib_file_name(model['type'], f, model[f])
value = re.sub(r'\b' + re.escape(model[f]) + r'\b', fn, value)
if model['commandFile']:
value = './' + value
if element_schema[1] == 'RPNValue':
value = _format_rpn_value(value)
if element_schema[1].startswith('InputFile'):
value = template_common.lib_file_name(model['type'], field_name, value)
if element_schema[1] == 'InputFileXY':
value += '={}+{}'.format(model[field_name + 'X'], model[field_name + 'Y'])
elif element_schema[1] == 'OutputFile':
value = state['filename_map']['{}{}{}'.format(model['_id'], _FILE_ID_SEP, state['field_index'])]
#TODO(pjm): don't quote numeric constants
state['lattice'] += '{}="{}",'.format(field_name, value)
else:
state['field_index'] = 0
if state['lattice']:
state['lattice'] = state['lattice'][:-1]
state['lattice'] += "\n"
state['lattice'] += '"{}": {},'.format(model['name'].upper(), model['type'])
state['beamline_map'][model['_id']] = model['name']
def _generate_parameters_file(data):
report = data['report'] if 'report' in data else None
template_common.validate_models(data, simulation_db.get_schema(SIM_TYPE))
v = template_common.flatten_data(data['models'], {})
v['beamEvolutionOutputFilename'] = _BEAM_EVOLUTION_OUTPUT_FILENAME
v['runSimulation'] = report is None or report == 'animation'
v['runRateCalculation'] = report is None or report == 'rateCalculationReport'
if data['models']['ring']['latticeSource'] == 'madx':
v['latticeFilename'] = template_common.lib_file_name('ring', 'lattice', v['ring_lattice'])
else:
v['latticeFilename'] = JSPEC_TWISS_FILENAME
if v['ionBeam_beam_type'] == 'continuous':
v['ionBeam_rms_bunch_length'] = 0
v['simulationSettings_ibs'] = 'on' if v['simulationSettings_ibs'] == '1' else 'off'
v['simulationSettings_e_cool'] = 'on' if v['simulationSettings_e_cool'] == '1' else 'off'
return template_common.render_jinja(SIM_TYPE, v)
def _iterator_commands(state, model, element_schema=None, field_name=None):
# only interested in commands, not elements
if '_type' not in model:
return
if element_schema:
state['field_index'] += 1
value = model[field_name]
default_value = element_schema[2]
if value is not None and default_value is not None:
if str(value) != str(default_value):
el_type = element_schema[1]
if el_type.endswith('StringArray'):
state['commands'] += ' {}[0] = {},'.format(
field_name, value) + "\n"
else:
#TODO(pjm): combine with lattice file input formatting below
if el_type == 'RPNValue':
value = _format_rpn_value(value, is_command=True)
elif el_type == 'OutputFile':
value = state['filename_map']['{}{}{}'.format(
model['_id'], _FILE_ID_SEP, state['field_index'])]
elif el_type.startswith('InputFile'):
value = template_common.lib_file_name(
'command_{}'.format(model['_type']), field_name,
value)
elif el_type == 'BeamInputFile':
value = 'bunchFile-sourceFile.{}'.format(value)
elif el_type == 'LatticeBeamlineList':
value = state['beamline_map'][int(value)]
elif el_type == 'ElegantLatticeList':
if value and value == 'Lattice':
value = 'elegant.lte'
else:
value = value + '.filename.lte'
if not _is_numeric(el_type, str(value)):
value = '"{}"'.format(value)
state['commands'] += ' {} = {},'.format(
field_name, value) + "\n"
else:
state['field_index'] = 0
if state['commands']:
state['commands'] += '&end' + "\n"
state['commands'] += "\n" + '&{}'.format(model['_type']) + "\n"
if model['_type'] == 'run_setup':
state['commands'] += ' semaphore_file = {},'.format(
_ELEGANT_SEMAPHORE_FILE) + "\n"
def _generate_parameters_file(data):
_validate_data(data, _SCHEMA)
v = template_common.flatten_data(data['models'], {})
beamline_map = _build_beamline_map(data)
v['lattice'] = _generate_lattice(data, beamline_map, v)
v['bunchFileName'] = OUTPUT_FILE['bunchReport']
v['diagnosticFilename'] = OUTPUT_FILE['beamEvolutionAnimation']
v['twissFileName'] = OUTPUT_FILE['twissReport']
if data.models.bunch.distribution == 'file':
v['bunchFile'] = template_common.lib_file_name('bunch', 'particleFile', data.models.bunch.particleFile)
v['bunch'] = template_common.render_jinja(SIM_TYPE, v, 'bunch.py')
res = template_common.render_jinja(SIM_TYPE, v, 'base.py')
report = data['report'] if 'report' in data else ''
if 'bunchReport' in report or 'twissReport' in report:
res += template_common.render_jinja(SIM_TYPE, v, 'twiss.py')
if 'bunchReport' in report:
res += template_common.render_jinja(SIM_TYPE, v, 'bunch-report.py')
else:
res += template_common.render_jinja(SIM_TYPE, v, 'parameters.py')
return res
def write_parameters(data,
run_dir,
is_parallel,
python_file=template_common.PARAMETERS_PYTHON_FILE):
pkio.write_text(
run_dir.join(python_file),
_generate_parameters_file(data),
)
# unzip the required magnet files
for el in data.models.elements:
if el.type != 'TOSCA':
continue
filename = str(
run_dir.join(
template_common.lib_file_name('TOSCA', 'magnetFile',
el.magnetFile)))
if _is_zip_file(filename):
with zipfile.ZipFile(filename, 'r') as z:
for info in z.infolist():
if info.filename in el.fileNames:
z.extract(info, str(run_dir))
def _tosca_info(data):
# determine the list of available files (from zip if necessary)
# compute the tosca length from datafile
tosca = data['tosca']
#TODO(pjm): keep a cache on the tosca model?
datafile = simulation_db.simulation_lib_dir(SIM_TYPE).join(
template_common.lib_file_name('TOSCA', 'magnetFile',
tosca['magnetFile']))
if not datafile.exists():
return {
'error': 'missing or invalid file: {}'.format(tosca['magnetFile']),
}
error = None
length = None
if _is_zip_file(datafile):
with zipfile.ZipFile(str(datafile), 'r') as z:
filenames = []
if 'fileNames' not in tosca or not tosca['fileNames']:
tosca['fileNames'] = []
for info in z.infolist():
filenames.append(info.filename)
if not length and info.filename in tosca['fileNames']:
length, error = _tosca_length(tosca,
z.read(info).splitlines())
if length:
error = None
else:
filenames = [tosca['magnetFile']]
with pkio.open_text(str(datafile)) as f:
length, error = _tosca_length(tosca, f)
if error:
return {'error': error}
return {
'toscaInfo': {
'toscaLength': length,
'fileList': sorted(filenames) if filenames else None,
'magnetFile': tosca['magnetFile'],
},
}
def _iterator_commands(state, model, element_schema=None, field_name=None):
# only interested in commands, not elements
if '_type' not in model:
return
if element_schema:
state['field_index'] += 1
value = model[field_name]
default_value = element_schema[2]
if value is not None and default_value is not None:
if str(value) != str(default_value):
if element_schema[1] == 'RPNValue':
state['commands'] += ' {} = "{}",'.format(field_name, _format_rpn_value(value, is_command=True)) + "\n"
elif element_schema[1].endswith('StringArray'):
state['commands'] += ' {}[0] = {},'.format(field_name, value) + "\n"
else:
#TODO(pjm): combine with lattice file input formatting below
if element_schema[1] == 'OutputFile':
value = state['filename_map']['{}{}{}'.format(model['_id'], _FILE_ID_SEP, state['field_index'])]
elif element_schema[1].startswith('InputFile'):
value = template_common.lib_file_name('command_{}'.format(model['_type']), field_name, value)
elif element_schema[1] == 'BeamInputFile':
value = 'bunchFile-sourceFile.{}'.format(value)
elif element_schema[1] == 'LatticeBeamlineList':
value = state['beamline_map'][int(value)]
elif element_schema[1] == 'ElegantLatticeList':
if value and value == 'Lattice':
value = 'elegant.lte'
else:
value = value + '.filename.lte'
state['commands'] += ' {} = "{}",'.format(field_name, value) + "\n"
else:
state['field_index'] = 0
if state['commands']:
state['commands'] += '&end' + "\n"
state['commands'] += "\n" + '&{}'.format(model['_type']) + "\n"
if model['_type'] == 'run_setup':
state['commands'] += ' semaphore_file = {},'.format(_ELEGANT_SEMAPHORE_FILE) + "\n"
def _run_bunch_report(data):
import synergia.bunch
with h5py.File(template.OUTPUT_FILE['twissReport'], 'r') as f:
twiss0 = dict(
map(
lambda k: (k, template.format_float(f[k][0])),
('alpha_x', 'alpha_y', 'beta_x', 'beta_y'),
))
report = data.models[data['report']]
bunch = data.models.bunch
if bunch.distribution == 'file':
bunch_file = template_common.lib_file_name('bunch', 'particleFile',
bunch.particleFile)
else:
bunch_file = template.OUTPUT_FILE['bunchReport']
with h5py.File(bunch_file, 'r') as f:
x = f['particles'][:, getattr(synergia.bunch.Bunch, report['x'])]
y = f['particles'][:, getattr(synergia.bunch.Bunch, report['y'])]
hist, edges = np.histogramdd([x, y],
template_common.histogram_bins(
report['histogramBins']))
return {
'title': '{}-{}'.format(report['x'], report['y']),
'x_range': [float(edges[0][0]),
float(edges[0][-1]),
len(hist)],
'y_range': [float(edges[1][0]),
float(edges[1][-1]),
len(hist[0])],
'x_label': template.label(report['x']),
'y_label': template.label(report['y']),
'z_matrix': hist.T.tolist(),
'summaryData': {
'bunchTwiss': twiss0,
},
}
def _iterator_input_files(state, model, element_schema=None, field_name=None):
if element_schema:
if model[field_name] and element_schema[1].startswith('InputFile'):
state.append(
template_common.lib_file_name(_model_name_for_data(model),
field_name, model[field_name]))
def _wiggler_file(value):
return template_common.lib_file_name('wiggler', 'trajFile', value)
def _simulation_files(data):
res = []
bunch = data.models.bunch
if bunch.distribution == 'file':
res.append(template_common.lib_file_name('bunch', 'particleFile', bunch.particleFile))
return res
def _wiggler_file(value):
return template_common.lib_file_name('wiggler', 'trajFile', value)
def _iterator_input_files(state, model, element_schema=None, field_name=None):
if element_schema:
if model[field_name] and element_schema[1].startswith('InputFile'):
state.append(template_common.lib_file_name(_model_name_for_data(model), field_name, model[field_name]))
def _analysis_data_file(data):
return template_common.lib_file_name('analysisData', 'file', data.models.analysisData.file)
