def _extract_bunch_plot(report, frame_index, run_dir):
filename = _particle_file_list(run_dir)[frame_index]
with h5py.File(str(filename), 'r') as f:
x = f['particles'][:, _COORD6.index(report['x'])].tolist()
y = f['particles'][:, _COORD6.index(report['y'])].tolist()
data = simulation_db.read_json(
run_dir.join(template_common.INPUT_BASE_NAME))
if 'bunchAnimation' not in data.models:
# In case the simulation was run before the bunchAnimation was added
return {
'error': 'Report not generated',
}
tlen = f['tlen'][()]
s_n = f['s_n'][()]
rep = 0 if s_n == 0 else int(round(tlen / s_n))
model = data.models.bunchAnimation
model.update(report)
return template_common.heatmap(
[x, y], model, {
'x_label':
label(report['x']),
'y_label':
label(report['y']),
'title':
'{}-{} at {:.1f}m, turn {}'.format(report['x'], report['y'],
tlen, rep),
})
def _extract_particle_plot(report, run_dir, page_index):
xfield = _map_field_name(report['x'])
yfield = _map_field_name(report['y'])
filename = _ion_files(run_dir)[page_index]
data = simulation_db.read_json(run_dir.join(template_common.INPUT_BASE_NAME))
settings = data.models.simulationSettings
time = settings.time / settings.step_number * settings.save_particle_interval * page_index
if time > settings.time:
time = settings.time
x_col = sdds_util.extract_sdds_column(filename, xfield, 0)
if x_col['err']:
return x_col['err']
x = x_col['values']
y_col = sdds_util.extract_sdds_column(filename, yfield, 0)
if y_col['err']:
return y_col['err']
y = y_col['values']
model = data.models.particleAnimation
model.update(report)
model['x'] = xfield
model['y'] = yfield
return template_common.heatmap([x, y], model, {
'x_label': _field_label(xfield, x_col['column_def']),
'y_label': _field_label(yfield, y_col['column_def']),
'title': 'Ions at time {:.2f} [s]'.format(time),
})
def sim_frame_particleAnimation(frame_args):
page_index = frame_args.frameIndex
xfield = _map_field_name(frame_args.x)
yfield = _map_field_name(frame_args.y)
filename = _ion_files(frame_args.run_dir)[page_index]
data = frame_args.sim_in
settings = data.models.simulationSettings
time = settings.time / settings.step_number * settings.save_particle_interval * page_index
if time > settings.time:
time = settings.time
x_col = sdds_util.extract_sdds_column(filename, xfield, 0)
if x_col['err']:
return x_col['err']
x = x_col['values']
y_col = sdds_util.extract_sdds_column(filename, yfield, 0)
if y_col['err']:
return y_col['err']
y = y_col['values']
model = data.models.particleAnimation
model.update(frame_args)
model['x'] = xfield
model['y'] = yfield
return template_common.heatmap(
[x, y], model, {
'x_label': _field_label(xfield, x_col['column_def']),
'y_label': _field_label(yfield, y_col['column_def']),
'title': 'Ions at time {:.2f} [s]'.format(time),
})
def sim_frame_beamAnimation(frame_args):
data = simulation_db.read_json(
run_dir.join(template_common.INPUT_BASE_NAME))
model = data.models.beamAnimation
model.update(frame_args)
beam_info = hellweg_dump_reader.beam_info(_dump_file(run_dir),
frame_args.frameIndex)
x, y = frame_args.reportType.split('-')
values = [
hellweg_dump_reader.get_points(beam_info, x),
hellweg_dump_reader.get_points(beam_info, y),
]
model['x'] = x
model['y'] = y
# see issue #872
if not np.any(values):
values = [[], []]
return template_common.heatmap(
values, model, {
'x_label':
hellweg_dump_reader.get_label(x),
'y_label':
hellweg_dump_reader.get_label(y),
'title':
_report_title(frame_args.reportType, 'BeamReportType', beam_info),
'z_label':
'Number of Particles',
'summaryData':
_summary_text(run_dir),
})
def save_sequential_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 = _SIM_DATA.lib_file_name_with_model_field(
'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'])]
res = template_common.heatmap(
[x, y], report, {
'title': '',
'x_label': label(report.x, _SCHEMA.enum.PhaseSpaceCoordinate8),
'y_label': label(report.y, _SCHEMA.enum.PhaseSpaceCoordinate8),
'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,
}
template_common.write_sequential_result(res, run_dir=run_dir)
def _extract_heatmap_data(report, col_names, rows, title):
x_info = _ANIMATION_FIELD_INFO[report.x]
y_info = _ANIMATION_FIELD_INFO[report.y]
x = np.array(column_data(report.x, col_names, rows)) * x_info[1]
y = np.array(column_data(report.y, col_names, rows)) * y_info[1]
return template_common.heatmap([x, y], report, PKDict(
x_label=x_info[0],
y_label=y_info[0],
title=title,
z_label='Number of Particles',
))
def _extract_report_data(xFilename, frame_args, page_count=0):
page_index = frame_args.frameIndex
xfield = frame_args.x if 'x' in frame_args else frame_args[_X_FIELD]
# x, column_names, x_def, err
x_col = sdds_util.extract_sdds_column(xFilename, xfield, page_index)
if x_col['err']:
return x_col['err']
x = x_col['values']
if not _is_histogram_file(xFilename, x_col['column_names']):
# parameter plot
plots = []
filename = PKDict(
y1=xFilename,
#TODO(pjm): y2Filename, y3Filename are not currently used. Would require rescaling x value across files.
y2=xFilename,
y3=xFilename,
)
for f in ('y1', 'y2', 'y3'):
if re.search(r'^none$', frame_args[f],
re.IGNORECASE) or frame_args[f] == ' ':
continue
yfield = frame_args[f]
y_col = sdds_util.extract_sdds_column(filename[f], yfield,
page_index)
if y_col['err']:
return y_col['err']
y = y_col['values']
plots.append(
PKDict(
field=yfield,
points=y,
label=_field_label(yfield, y_col['column_def'][1]),
))
title = ''
if page_count > 1:
title = 'Plot {} of {}'.format(page_index + 1, page_count)
return template_common.parameter_plot(
x, plots, frame_args,
PKDict(
title=title,
y_label='',
x_label=_field_label(xfield, x_col['column_def'][1]),
))
yfield = frame_args['y1'] if 'y1' in frame_args else frame_args['y']
y_col = sdds_util.extract_sdds_column(xFilename, yfield, page_index)
if y_col['err']:
return y_col['err']
return template_common.heatmap(
[x, y_col['values']], frame_args,
PKDict(
x_label=_field_label(xfield, x_col['column_def'][1]),
y_label=_field_label(yfield, y_col['column_def'][1]),
title=_plot_title(xfield, yfield, page_index, page_count),
))
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 _extract_heatmap_data(report, col_names, rows, title):
x_info = _ANIMATION_FIELD_INFO[report['x']]
y_info = _ANIMATION_FIELD_INFO[report['y']]
x = np.array(column_data(report['x'], col_names, rows)) * x_info[1]
y = np.array(column_data(report['y'], col_names, rows)) * y_info[1]
return template_common.heatmap(
[x, y], report, {
'x_label': x_info[0],
'y_label': y_info[0],
'title': title,
'z_label': 'Number of Particles',
})
def _extract_report_bunchReport(data, run_dir):
parts = simulation_db.read_json(run_dir.join(BUNCH_PARTICLES_FILE))
m = data.models[data.report]
res = template_common.heatmap([
parts[m.x],
parts[m.y],
], m, PKDict(
x_label=_field_label(m.x),
y_label=_field_label(m.y),
))
bunch = data.models.bunch
res.summaryData = parts.summaryData
return res
def extract_report_data(xFilename, data, page_index, page_count=0):
xfield = data['x'] if 'x' in data else data[_X_FIELD]
# x, column_names, x_def, err
x_col = sdds_util.extract_sdds_column(xFilename, xfield, page_index)
if x_col['err']:
return x_col['err']
x = x_col['values']
if not _is_histogram_file(xFilename, x_col['column_names']):
# parameter plot
plots = []
filename = {
'y1': xFilename,
#TODO(pjm): y2Filename, y3Filename are not currently used. Would require rescaling x value across files.
'y2': xFilename,
'y3': xFilename,
}
for f in ('y1', 'y2', 'y3'):
if re.search(r'^none$', data[f], re.IGNORECASE) or data[f] == ' ':
continue
yfield = data[f]
y_col = sdds_util.extract_sdds_column(filename[f], yfield,
page_index)
if y_col['err']:
return y_col['err']
y = y_col['values']
plots.append({
'field': yfield,
'points': y,
'label': _field_label(yfield, y_col['column_def'][1]),
})
title = ''
if page_count > 1:
title = 'Plot {} of {}'.format(page_index + 1, page_count)
return template_common.parameter_plot(
x, plots, data, {
'title': title,
'y_label': '',
'x_label': _field_label(xfield, x_col['column_def'][1]),
})
yfield = data['y1'] if 'y1' in data else data['y']
y_col = sdds_util.extract_sdds_column(xFilename, yfield, page_index)
if y_col['err']:
return y_col['err']
return template_common.heatmap(
[x, y_col['values']], data, {
'x_label': _field_label(xfield, x_col['column_def'][1]),
'y_label': _field_label(yfield, y_col['column_def'][1]),
'title': _plot_title(xfield, yfield, page_index, page_count),
})
def _fit_animation(frame_args):
idx = int(frame_args.columnNumber)
header, v = _read_file(frame_args.run_dir, _OUTPUT_FILE.predictOutputFile)
_, y = _read_file(frame_args.run_dir, _OUTPUT_FILE.testOutputFile)
frame_args.histogramBins = 30
return template_common.heatmap(
[v[:, idx], y[:, idx]],
frame_args,
PKDict(
x_label='',
y_label='',
title=header[idx],
hideColorBar=True,
),
)
def sim_frame_bunchAnimation(frame_args):
a = frame_args.sim_in.models.bunchAnimation
a.update(frame_args)
res = PKDict()
with h5py.File(str(a.run_dir.join(_OPAL_H5_FILE)), 'r') as f:
for field in ('x', 'y'):
res[field] = PKDict(
name=a[field],
points=np.array(f['/Step#{}/{}'.format(a.frameIndex,
a[field])]),
)
_units_from_hdf5(f, res[field])
return template_common.heatmap([res.x.points, res.y.points], a,
PKDict(
x_label=res.x.label,
y_label=res.y.label,
))
def _bunch_plot(report, run_dir, idx, filename=_OPAL_H5_FILE):
res = PKDict()
title = 'Step {}'.format(idx)
with h5py.File(str(run_dir.join(filename)), 'r') as f:
for field in ('x', 'y'):
res[field] = PKDict(
name=report[field],
points=np.array(f['/Step#{}/{}'.format(idx, report[field])]),
label=report[field],
)
_units_from_hdf5(f, res[field])
if 'SPOS' in f['/Step#{}'.format(idx)].attrs:
title += ', SPOS {0:.5f}m'.format(f['/Step#{}'.format(idx)].attrs['SPOS'][0])
return template_common.heatmap([res.x.points, res.y.points], report, PKDict(
x_label=res.x.label,
y_label=res.y.label,
title=title,
))
def extract_beam_report(report, run_dir, frame):
data = simulation_db.read_json(run_dir.join(template_common.INPUT_BASE_NAME))
model = data.models.beamAnimation
model.update(report)
beam_info = hellweg_dump_reader.beam_info(_dump_file(run_dir), frame)
x, y = report.reportType.split('-')
values = [
hellweg_dump_reader.get_points(beam_info, x),
hellweg_dump_reader.get_points(beam_info, y),
]
model['x'] = x
model['y'] = y
return template_common.heatmap(values, model, {
'x_label': hellweg_dump_reader.get_label(x),
'y_label': hellweg_dump_reader.get_label(y),
'title': _report_title(report.reportType, 'BeamReportType', beam_info),
'z_label': 'Number of Particles',
'summaryData': _summary_text(run_dir),
})
def _confusion_matrix_to_heatmap_report(frame_args, filename, title):
r = pkjson.load_any(frame_args.run_dir.join(filename))
a = None
for y, _ in enumerate(r.matrix):
for x, v in enumerate(r.matrix[y]):
t = np.repeat([[x, y]], v, axis=0)
a = t if a is None else np.vstack([t, a])
return template_common.heatmap(
a,
PKDict(histogramBins=len(r.matrix)),
plot_fields=PKDict(
labels=list(
_get_classification_output_col_encoding(frame_args).values(),
),
title=title.format(**r),
x_label='Predicted',
y_label='True',
),
)
def _fit_animation(frame_args):
idx = int(frame_args.columnNumber)
frame_args.histogramBins = 30
info = frame_args.sim_in.models.columnInfo
header = []
for i in range(len(info.inputOutput)):
if info.inputOutput[i] == 'output':
header.append(info.header[i])
return template_common.heatmap(
[
_read_file(frame_args.run_dir, _OUTPUT_FILE.predictFile)[:, idx],
_read_file(frame_args.run_dir, _OUTPUT_FILE.testFile)[:, idx],
],
frame_args,
PKDict(
x_label='',
y_label='',
title=header[idx],
hideColorBar=True,
),
)
def _extract_impact_density_2d(run_dir, data):
# use a simple heatmap instead due to a normalization problem in rswarp
all_particles = np.load('all-particles.npy')
grid = data.models.simulationGrid
plate_spacing = _meters(grid.plate_spacing)
channel_width = _meters(grid.channel_width)
m = PKDict(
histogramBins=200,
plotRangeType='fixed',
horizontalSize=plate_spacing * 1.02,
horizontalOffset=plate_spacing / 2,
verticalSize=channel_width,
verticalOffset=0,
)
return template_common.heatmap(
[all_particles[1].tolist(), all_particles[0].tolist()], m,
PKDict(
x_label='z [m]',
y_label='x [m]',
title='Impact Density',
aspectRatio=4.0 / 7,
))
def _extract_report_bunchReport(data, run_dir):
# read from file? store on model?
parts = _ptc_particles(
_get_initial_twiss_params(data),
data.models.simulation.numberOfParticles
)
labels = []
res = []
r_model = data.models[data.report]
for dim in ('x', 'y'):
c = _SCHEMA.constants.phaseSpaceParticleMap[r_model[dim]]
res.append(parts[c[0]][c[1]])
labels.append(r_model[dim])
return template_common.heatmap(
res,
PKDict(histogramBins=100),
PKDict(
x_label=labels[0],
y_label=labels[1],
)
)
def sim_frame_bunchAnimation(frame_args):
n = _particle_file_list(frame_args.run_dir)[frame_args.frameIndex]
with h5py.File(str(n), 'r') as f:
x = f['particles'][:, _COORD6.index(frame_args.x)].tolist()
y = f['particles'][:, _COORD6.index(frame_args.y)].tolist()
if 'bunchAnimation' not in frame_args.sim_in.models:
# In case the simulation was run before the bunchAnimation was added
return PKDict(error='report not generated')
tlen = f['tlen'][()]
s_n = f['s_n'][()]
rep = 0 if s_n == 0 else int(round(tlen / s_n))
model = frame_args.sim_in.models.bunchAnimation
model.update(frame_args)
return template_common.heatmap(
[x, y],
model,
PKDict(
x_label=label(frame_args.x),
y_label=label(frame_args.y),
title='{}-{} at {:.1f}m, turn {}'.format(frame_args.x, frame_args.y, tlen, rep),
),
)
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,
),
),
)
