From 9d7c10940191154bd881248661eac1be454f70fc Mon Sep 17 00:00:00 2001 From: git-user Date: Thu, 8 Feb 2024 21:27:46 +0000 Subject: [PATCH 01/23] Fix #6653: opal ckp --- .../template/omega/parameters.py.jinja | 130 +++++++++--------- sirepo/template/omega.py | 103 +++++++++++++- 2 files changed, 166 insertions(+), 67 deletions(-) diff --git a/sirepo/package_data/template/omega/parameters.py.jinja b/sirepo/package_data/template/omega/parameters.py.jinja index c45dd4fe31..a11b35b11d 100644 --- a/sirepo/package_data/template/omega/parameters.py.jinja +++ b/sirepo/package_data/template/omega/parameters.py.jinja @@ -59,9 +59,9 @@ def genesis_to_pmd(sim): return v -def prep_run_dir(run_dir, data): - sirepo.simulation_db.prepare_simulation(data, run_dir) - return pkio.save_chdir(run_dir) +#def prep_run_dir(run_dir, data): +# sirepo.simulation_db.prepare_simulation(data, run_dir) +# return pkio.save_chdir(run_dir) def read_sim(sim_type, sim_id): @@ -77,7 +77,7 @@ def run_elegant(run_dir, elegant_id, prev_sim=None): c.copy( c.uri(c.LIB_DIR, b), c.uri( - c.LIB_DIR, + c.LIB_DIR, s.lib_file_name_with_model_field("command_run_setup", "expand_for", filename), ), ) @@ -188,65 +188,67 @@ def run_genesis(run_dir, genesis_id, prev_sim=None): def run_opal(run_dir, opal_id, prev_sim=None): - def _save_lib_file(tmp, filename): - s = sirepo.sim_data.get_class("opal") - s.lib_file_write( - s.lib_file_name_with_model_field("command_distribution", "fname", filename), - tmp, - ) - tmp.remove() - - def _update_sim(data, filename, tmp): - d = LatticeUtil.find_first_command(data, "distribution") - d.type = "FROMFILE" - d.fname = filename - n = 0 - zsum = 0 - psum = 0 - # calculate z offset - with pkio.open_text(tmp) as f: - for line in f: - if not n: - n = int(line) - continue - r = [v for v in re.split(r"\s+", line.strip())] - zsum += float(r[4]) - psum += float(r[5]) - d.offsetz = -zsum / n - - b = LatticeUtil.find_first_command(data, "beam") - b.npart = n - b.gamma = 0 - b.energy = 0 - #TODO(pjm): handle other particle types - mass_and_charge = PKDict( - ELECTRON=0.51099895000e-03, - PROTON=0.93827208816, - ) - b.pc = psum / n * mass_and_charge[b.particle] - _save_lib_file(tmp, filename) - write_sim(data) - - data = read_sim('opal', opal_id) - if prev_sim: - filename = file_name_from_sim_name('opal', f'{_OMEGA_SIM_NAME}-{run_dir.basename}') - t = run_dir.join("opal-command_distribution") - if prev_sim.sim_type == 'genesis': - import pmd_beamphysics.interfaces.opal - - pmd_beamphysics.interfaces.opal.write_opal( - genesis_to_pmd(prev_sim), - str(t), - ) - else: - assert prev_sim.outfile_path - sw = switchyard.Switchyard() - sw.read(f'{prev_sim.outfile_path}', prev_sim.sim_type) - sw.write(str(t), "opal") - _update_sim(data, filename, t) - data.computeModel = 'animation' - LatticeUtil.find_first_command(data, "option").psdumpfreq = 0 - with prep_run_dir(run_dir, data): +# def _save_lib_file(tmp, filename): +# s = sirepo.sim_data.get_class("opal") +# s.lib_file_write( +# s.lib_file_name_with_model_field("command_distribution", "fname", filename), +# tmp, +# ) +# tmp.remove() +# +# def _update_sim(data, filename, tmp): +# d = LatticeUtil.find_first_command(data, "distribution") +# d.type = "FROMFILE" +# d.fname = filename +# n = 0 +# zsum = 0 +# psum = 0 +# # calculate z offset +# with pkio.open_text(tmp) as f: +# for line in f: +# if not n: +# n = int(line) +# continue +# r = [v for v in re.split(r"\s+", line.strip())] +# zsum += float(r[4]) +# psum += float(r[5]) +# d.offsetz = -zsum / n +# +# b = LatticeUtil.find_first_command(data, "beam") +# b.npart = n +# b.gamma = 0 +# b.energy = 0 +# #TODO(pjm): handle other particle types +# mass_and_charge = PKDict( +# ELECTRON=0.51099895000e-03, +# PROTON=0.93827208816, +# ) +# b.pc = psum / n * mass_and_charge[b.particle] +# _save_lib_file(tmp, filename) +# write_sim(data) +# +# data = read_sim('opal', opal_id) +# if prev_sim: +# filename = file_name_from_sim_name('opal', f'{_OMEGA_SIM_NAME}-{run_dir.basename}') +# t = run_dir.join("opal-command_distribution") +# if prev_sim.sim_type == 'genesis': +# import pmd_beamphysics.interfaces.opal +# +# pmd_beamphysics.interfaces.opal.write_opal( +# genesis_to_pmd(prev_sim), +# str(t), +# ) +# else: +# assert prev_sim.outfile_path +# sw = switchyard.Switchyard() +# sw.read(f'{prev_sim.outfile_path}', prev_sim.sim_type) +# sw.write(str(t), "opal") +# _update_sim(data, filename, t) +# data.computeModel = 'animation' +# LatticeUtil.find_first_command(data, "option").psdumpfreq = 0 +# with prep_run_dir(run_dir, data): + with pkio.save_chdir(run_dir): + print("run dir", run_dir) sirepo.pkcli.opal.run_opal(with_mpi=True) return f'{run_dir}/{sirepo.template.opal._OPAL_H5_FILE}' @@ -255,8 +257,6 @@ def run_sims(sim_list): prev = None for idx in range(len(sim_list)): run_dir = pkio.py_path(f'run{idx + 1}') - pkio.unchecked_remove(run_dir) - pkio.mkdir_parent(run_dir) s = sim_list[idx] if s.sim_type == 'opal': s.outfile_path = run_opal(run_dir, s.sim_id, prev) diff --git a/sirepo/template/omega.py b/sirepo/template/omega.py index 00ead9493b..0dd51827ef 100644 --- a/sirepo/template/omega.py +++ b/sirepo/template/omega.py @@ -9,6 +9,7 @@ from pykern.pkdebug import pkdp, pkdc, pkdlog from sirepo import simulation_db from sirepo.template import template_common +from sirepo.template.lattice import LatticeUtil import h5py import numpy import pmd_beamphysics @@ -281,12 +282,105 @@ def stateful_compute_get_opal_sim_list(**kwargs): def write_parameters(data, run_dir, is_parallel): + prev = None + sim_list = _coupled_sims_list(data) + pkdp("\n\n\n sim_list={}", sim_list) + for idx in range(len(sim_list)): + s = sim_list[idx] + if s.sim_type == "opal": + d = run_dir.join(f'run{idx + 1}') + pkio.unchecked_remove(d) + pkio.mkdir_parent(d) + _prepare_opal(d, s.sim_id, prev) pkio.write_text( run_dir.join(template_common.PARAMETERS_PYTHON_FILE), _generate_parameters_file(data), ) +def read_sim(sim_type, sim_id): + return sirepo.sim_data.get_class(sim_type).sim_db_read_sim(sim_id) + + +def write_sim(data): + sirepo.sim_data.get_class(data.simulationType).sim_db_save_sim(data) + + +def file_name_from_sim_name(sim_type, sim_name): + res = re.sub(r'[^0-9a-zA-Z]', '_', sim_name) + res = re.sub(r'^\_+|\_+$', '', res) + res = re.sub(r'\_+', '_', res) + ext = "sdds" if sim_type == "elegant" else "dat" + return f"{res}.{ext}" + + +def _prepare_opal(run_dir, opal_id, prev_sim=None): + def _save_lib_file(tmp, filename): + s = sirepo.sim_data.get_class("opal") + s.lib_file_write( + s.lib_file_name_with_model_field("command_distribution", "fname", filename), + tmp, + ) + tmp.remove() + + def _update_sim(data, filename, tmp): + d = LatticeUtil.find_first_command(data, "distribution") + d.type = "FROMFILE" + d.fname = filename + n = 0 + zsum = 0 + psum = 0 + # calculate z offset + with pkio.open_text(tmp) as f: + for line in f: + if not n: + n = int(line) + continue + r = [v for v in re.split(r"\s+", line.strip())] + zsum += float(r[4]) + psum += float(r[5]) + d.offsetz = -zsum / n + + b = LatticeUtil.find_first_command(data, "beam") + b.npart = n + b.gamma = 0 + b.energy = 0 + #TODO(pjm): handle other particle types + mass_and_charge = PKDict( + ELECTRON=0.51099895000e-03, + PROTON=0.93827208816, + ) + b.pc = psum / n * mass_and_charge[b.particle] + _save_lib_file(tmp, filename) + write_sim(data) + + data = read_sim('opal', opal_id) + if prev_sim: + filename = file_name_from_sim_name('opal', f'{_OMEGA_SIM_NAME}-{run_dir.basename}') + t = run_dir.join("opal-command_distribution") + if prev_sim.sim_type == 'genesis': + import pmd_beamphysics.interfaces.opal + + pmd_beamphysics.interfaces.opal.write_opal( + genesis_to_pmd(prev_sim), + str(t), + ) + else: + assert prev_sim.outfile_path + sw = switchyard.Switchyard() + sw.read(f'{prev_sim.outfile_path}', prev_sim.sim_type) + sw.write(str(t), "opal") + _update_sim(data, filename, t) + data.computeModel = 'animation' + LatticeUtil.find_first_command(data, "option").psdumpfreq = 0 + sirepo.simulation_db.prepare_simulation(data, run_dir) + # pkio.save_chdir(run_dir) + # with prep_run_dir(run_dir, data): + # sirepo.pkcli.opal.run_opal(with_mpi=True) + # return f'{run_dir}/{sirepo.template.opal._OPAL_H5_FILE}' + + + def _extract_elegant_beam_plot(frame_args): # this is tricky because the data could come from 2 different elegant output files files = PKDict() @@ -323,9 +417,8 @@ def _extract_elegant_beam_plot(frame_args): return res -def _generate_parameters_file(data): +def _coupled_sims_list(data): dm = data.models - res, v = template_common.generate_parameters_file(data) sim_list = [] for idx in range(len(dm.simWorkflow.coupledSims)): s = _sim_info(dm, idx) @@ -340,6 +433,12 @@ def _generate_parameters_file(data): break if not sim_list: raise AssertionError("No simulations selected") + return sim_list + + +def _generate_parameters_file(data): + res, v = template_common.generate_parameters_file(data) + sim_list = _coupled_sims_list(data) v.simList = sim_list return res + template_common.render_jinja(SIM_TYPE, v) From 6020ae8b0c437e6a943a905ef7fb39dc491b4ee9 Mon Sep 17 00:00:00 2001 From: git-user Date: Thu, 8 Feb 2024 21:43:44 +0000 Subject: [PATCH 02/23] Fix #6653: ckp elegant --- .../template/omega/parameters.py.jinja | 111 +----------------- sirepo/template/omega.py | 66 ++++++++++- 2 files changed, 64 insertions(+), 113 deletions(-) diff --git a/sirepo/package_data/template/omega/parameters.py.jinja b/sirepo/package_data/template/omega/parameters.py.jinja index a11b35b11d..c708d14ecc 100644 --- a/sirepo/package_data/template/omega/parameters.py.jinja +++ b/sirepo/package_data/template/omega/parameters.py.jinja @@ -69,56 +69,7 @@ def read_sim(sim_type, sim_id): def run_elegant(run_dir, elegant_id, prev_sim=None): - def _save_lib_files(tmp, filename): - s = sirepo.sim_data.get_class("elegant") - b = s.lib_file_name_with_model_field("bunchFile", "sourceFile", filename) - s.lib_file_write(b, tmp) - c = s.sim_db_client() - c.copy( - c.uri(c.LIB_DIR, b), - c.uri( - c.LIB_DIR, - s.lib_file_name_with_model_field("command_run_setup", "expand_for", filename), - ), - ) - tmp.remove() - - def _update_sim(data, filename, prev_sim): - if data.models.bunchSource.inputSource == 'bunched_beam': - convert_bunched_beam_to_sdds_beam(data, filename) - assert data.models.bunchSource.inputSource == 'sdds_beam' - cmd = LatticeUtil.find_first_command(data, "sdds_beam") - cmd.input = filename - cmd.center_arrival_time = '1' - cmd.center_transversely = '1' - cmd.reverse_t_sign = '1' if prev_sim.sim_type == 'genesis' else '0' - LatticeUtil.find_first_command(data, "run_setup").expand_for = filename - write_sim(data) - - data = read_sim('elegant', elegant_id) - if prev_sim: - assert prev_sim.outfile_path - filename = file_name_from_sim_name('elegant', f'{_OMEGA_SIM_NAME}-{run_dir.basename}') - _update_sim(data, filename, prev_sim) - t = run_dir.join("omega-elegant-bunch") - if prev_sim.sim_type == 'elegant': - pkio.py_path(prev_sim.outfile_path).copy(t) - elif prev_sim.sim_type == 'genesis': - import pmd_beamphysics.interfaces.elegant - - pmd_beamphysics.interfaces.elegant.write_elegant( - genesis_to_pmd(prev_sim), - str(t), - ) - else: - sw = switchyard.Switchyard() - sw.read(f'{prev_sim.outfile_path}', prev_sim.sim_type) - sw.write(str(t), 'elegant') - _save_lib_files(t, filename) - data.computeModel = 'animation' - if 'report' in data: - del data['report'] - with prep_run_dir(run_dir, data): + with pkio.save_chdir(run_dir): sirepo.pkcli.elegant.run_elegant() return f'{run_dir}/run_setup.output.sdds' @@ -188,67 +139,7 @@ def run_genesis(run_dir, genesis_id, prev_sim=None): def run_opal(run_dir, opal_id, prev_sim=None): -# def _save_lib_file(tmp, filename): -# s = sirepo.sim_data.get_class("opal") -# s.lib_file_write( -# s.lib_file_name_with_model_field("command_distribution", "fname", filename), -# tmp, -# ) -# tmp.remove() -# -# def _update_sim(data, filename, tmp): -# d = LatticeUtil.find_first_command(data, "distribution") -# d.type = "FROMFILE" -# d.fname = filename -# n = 0 -# zsum = 0 -# psum = 0 -# # calculate z offset -# with pkio.open_text(tmp) as f: -# for line in f: -# if not n: -# n = int(line) -# continue -# r = [v for v in re.split(r"\s+", line.strip())] -# zsum += float(r[4]) -# psum += float(r[5]) -# d.offsetz = -zsum / n -# -# b = LatticeUtil.find_first_command(data, "beam") -# b.npart = n -# b.gamma = 0 -# b.energy = 0 -# #TODO(pjm): handle other particle types -# mass_and_charge = PKDict( -# ELECTRON=0.51099895000e-03, -# PROTON=0.93827208816, -# ) -# b.pc = psum / n * mass_and_charge[b.particle] -# _save_lib_file(tmp, filename) -# write_sim(data) -# -# data = read_sim('opal', opal_id) -# if prev_sim: -# filename = file_name_from_sim_name('opal', f'{_OMEGA_SIM_NAME}-{run_dir.basename}') -# t = run_dir.join("opal-command_distribution") -# if prev_sim.sim_type == 'genesis': -# import pmd_beamphysics.interfaces.opal -# -# pmd_beamphysics.interfaces.opal.write_opal( -# genesis_to_pmd(prev_sim), -# str(t), -# ) -# else: -# assert prev_sim.outfile_path -# sw = switchyard.Switchyard() -# sw.read(f'{prev_sim.outfile_path}', prev_sim.sim_type) -# sw.write(str(t), "opal") -# _update_sim(data, filename, t) -# data.computeModel = 'animation' -# LatticeUtil.find_first_command(data, "option").psdumpfreq = 0 -# with prep_run_dir(run_dir, data): with pkio.save_chdir(run_dir): - print("run dir", run_dir) sirepo.pkcli.opal.run_opal(with_mpi=True) return f'{run_dir}/{sirepo.template.opal._OPAL_H5_FILE}' diff --git a/sirepo/template/omega.py b/sirepo/template/omega.py index 0dd51827ef..2a1852d8db 100644 --- a/sirepo/template/omega.py +++ b/sirepo/template/omega.py @@ -287,11 +287,13 @@ def write_parameters(data, run_dir, is_parallel): pkdp("\n\n\n sim_list={}", sim_list) for idx in range(len(sim_list)): s = sim_list[idx] + d = run_dir.join(f'run{idx + 1}') + pkio.unchecked_remove(d) + pkio.mkdir_parent(d) if s.sim_type == "opal": - d = run_dir.join(f'run{idx + 1}') - pkio.unchecked_remove(d) - pkio.mkdir_parent(d) _prepare_opal(d, s.sim_id, prev) + elif s.sim_type == "elegant": + _prepare_elegant(d, s.sim_id, prev) pkio.write_text( run_dir.join(template_common.PARAMETERS_PYTHON_FILE), _generate_parameters_file(data), @@ -381,6 +383,64 @@ def _update_sim(data, filename, tmp): +def _prepare_elegant(run_dir, elegant_id, prev_sim=None): + def _save_lib_files(tmp, filename): + s = sirepo.sim_data.get_class("elegant") + b = s.lib_file_name_with_model_field("bunchFile", "sourceFile", filename) + s.lib_file_write(b, tmp) + c = s.sim_db_client() + c.copy( + c.uri(c.LIB_DIR, b), + c.uri( + c.LIB_DIR, + s.lib_file_name_with_model_field("command_run_setup", "expand_for", filename), + ), + ) + tmp.remove() + + def _update_sim(data, filename, prev_sim): + if data.models.bunchSource.inputSource == 'bunched_beam': + convert_bunched_beam_to_sdds_beam(data, filename) + assert data.models.bunchSource.inputSource == 'sdds_beam' + cmd = LatticeUtil.find_first_command(data, "sdds_beam") + cmd.input = filename + cmd.center_arrival_time = '1' + cmd.center_transversely = '1' + cmd.reverse_t_sign = '1' if prev_sim.sim_type == 'genesis' else '0' + LatticeUtil.find_first_command(data, "run_setup").expand_for = filename + write_sim(data) + + data = read_sim('elegant', elegant_id) + if prev_sim: + assert prev_sim.outfile_path + filename = file_name_from_sim_name('elegant', f'{_OMEGA_SIM_NAME}-{run_dir.basename}') + _update_sim(data, filename, prev_sim) + t = run_dir.join("omega-elegant-bunch") + if prev_sim.sim_type == 'elegant': + pkio.py_path(prev_sim.outfile_path).copy(t) + elif prev_sim.sim_type == 'genesis': + import pmd_beamphysics.interfaces.elegant + + pmd_beamphysics.interfaces.elegant.write_elegant( + genesis_to_pmd(prev_sim), + str(t), + ) + else: + sw = switchyard.Switchyard() + sw.read(f'{prev_sim.outfile_path}', prev_sim.sim_type) + sw.write(str(t), 'elegant') + _save_lib_files(t, filename) + data.computeModel = 'animation' + if 'report' in data: + del data['report'] + sirepo.simulation_db.prepare_simulation(data, run_dir) + # pkio.save_chdir(run_dir) + # with prep_run_dir(run_dir, data): + # sirepo.pkcli.opal.run_opal(with_mpi=True) + # return f'{run_dir}/{sirepo.template.opal._OPAL_H5_FILE}' + + + def _extract_elegant_beam_plot(frame_args): # this is tricky because the data could come from 2 different elegant output files files = PKDict() From 6602b8e325b5008e930a3717968c77317084f6ec Mon Sep 17 00:00:00 2001 From: git-user Date: Thu, 8 Feb 2024 22:03:25 +0000 Subject: [PATCH 03/23] Fix #6653: cleanup ckp --- .../template/omega/parameters.py.jinja | 51 +----------- sirepo/template/omega.py | 82 +++++++++++++------ 2 files changed, 59 insertions(+), 74 deletions(-) diff --git a/sirepo/package_data/template/omega/parameters.py.jinja b/sirepo/package_data/template/omega/parameters.py.jinja index c708d14ecc..a1b328e047 100644 --- a/sirepo/package_data/template/omega/parameters.py.jinja +++ b/sirepo/package_data/template/omega/parameters.py.jinja @@ -13,61 +13,17 @@ import sirepo.sim_data import sirepo.simulation_db import sirepo.template.opal - +# TODO (gurhar1133): figure out move this to template _OMEGA_SIM_NAME = '{{ simulation_name }}' _GENESIS_PARTICLE_COLUMN_COUNT = 6 assert os.environ['SIREPO_SIMULATION_DB_LOGGED_IN_USER'], 'missing user id env var' -def convert_bunched_beam_to_sdds_beam(data, filename): - s = sirepo.sim_data.get_class('elegant') - cmd = LatticeUtil.find_first_command(data, "bunched_beam") - for k in list(cmd.keys()): - if k != '_id': - del cmd[k] - s.update_model_defaults(cmd, 'command_sdds_beam') - cmd._type = 'sdds_beam' - data.models.bunchSource.inputSource = 'sdds_beam' - - -def file_name_from_sim_name(sim_type, sim_name): - res = re.sub(r'[^0-9a-zA-Z]', '_', sim_name) - res = re.sub(r'^\_+|\_+$', '', res) - res = re.sub(r'\_+', '_', res) - ext = "sdds" if sim_type == "elegant" else "dat" - return f"{res}.{ext}" - - -def genesis_to_pmd(sim): - import pmd_beamphysics.interfaces.genesis - import pmd_beamphysics.particles - - g = read_sim("genesis", sim.sim_id) - d = numpy.fromfile(sim.outfile_path, dtype=numpy.float64) - d = d.reshape(int(len(d) / _GENESIS_PARTICLE_COLUMN_COUNT / g.models.electronBeam.npart), _GENESIS_PARTICLE_COLUMN_COUNT, g.models.electronBeam.npart) - # phase must be > 0 to avoid wrapping - d[:,1,:] -= numpy.min(d[:,1,:]) - v = pmd_beamphysics.particles.ParticleGroup( - data=pmd_beamphysics.interfaces.genesis.genesis2_dpa_to_data( - d, xlamds=g.models.radiation.xlamds, current=numpy.array([g.models.electronBeam.curpeak]), - # compute required wavelengths to hold the whole beam - zsep=numpy.max(d[:,1,:]) / (2 * numpy.pi), - ) - ) - # center psi - v.t -= numpy.mean(v.t) - return v - - #def prep_run_dir(run_dir, data): # sirepo.simulation_db.prepare_simulation(data, run_dir) # return pkio.save_chdir(run_dir) -def read_sim(sim_type, sim_id): - return sirepo.sim_data.get_class(sim_type).sim_db_read_sim(sim_id) - - def run_elegant(run_dir, elegant_id, prev_sim=None): with pkio.save_chdir(run_dir): sirepo.pkcli.elegant.run_elegant() @@ -160,11 +116,6 @@ def run_sims(sim_list): prev = s - -def write_sim(data): - sirepo.sim_data.get_class(data.simulationType).sim_db_save_sim(data) - - run_sims([ {% for sim in simList %} PKDict( diff --git a/sirepo/template/omega.py b/sirepo/template/omega.py index 2a1852d8db..a4391355c3 100644 --- a/sirepo/template/omega.py +++ b/sirepo/template/omega.py @@ -19,7 +19,7 @@ import re import sirepo.sim_data import sirepo.template - +# TODO (gurhar1133): need to alphabetical sort functions _SIM_DATA, SIM_TYPE, SCHEMA = sirepo.sim_data.template_globals() _PHASE_PLOT_COUNT = 4 _PHASE_PLOTS = PKDict( @@ -300,22 +300,6 @@ def write_parameters(data, run_dir, is_parallel): ) -def read_sim(sim_type, sim_id): - return sirepo.sim_data.get_class(sim_type).sim_db_read_sim(sim_id) - - -def write_sim(data): - sirepo.sim_data.get_class(data.simulationType).sim_db_save_sim(data) - - -def file_name_from_sim_name(sim_type, sim_name): - res = re.sub(r'[^0-9a-zA-Z]', '_', sim_name) - res = re.sub(r'^\_+|\_+$', '', res) - res = re.sub(r'\_+', '_', res) - ext = "sdds" if sim_type == "elegant" else "dat" - return f"{res}.{ext}" - - def _prepare_opal(run_dir, opal_id, prev_sim=None): def _save_lib_file(tmp, filename): s = sirepo.sim_data.get_class("opal") @@ -356,15 +340,16 @@ def _update_sim(data, filename, tmp): _save_lib_file(tmp, filename) write_sim(data) - data = read_sim('opal', opal_id) + data = _read_sim('opal', opal_id) if prev_sim: - filename = file_name_from_sim_name('opal', f'{_OMEGA_SIM_NAME}-{run_dir.basename}') + # TODO (gurhar1133): figure out _OMEGA_SIM_NAME + filename = _file_name_from_sim_name('opal', f'{_OMEGA_SIM_NAME}-{run_dir.basename}') t = run_dir.join("opal-command_distribution") if prev_sim.sim_type == 'genesis': import pmd_beamphysics.interfaces.opal pmd_beamphysics.interfaces.opal.write_opal( - genesis_to_pmd(prev_sim), + _genesis_to_pmd(prev_sim), str(t), ) else: @@ -400,7 +385,7 @@ def _save_lib_files(tmp, filename): def _update_sim(data, filename, prev_sim): if data.models.bunchSource.inputSource == 'bunched_beam': - convert_bunched_beam_to_sdds_beam(data, filename) + _convert_bunched_beam_to_sdds_beam(data, filename) assert data.models.bunchSource.inputSource == 'sdds_beam' cmd = LatticeUtil.find_first_command(data, "sdds_beam") cmd.input = filename @@ -410,10 +395,11 @@ def _update_sim(data, filename, prev_sim): LatticeUtil.find_first_command(data, "run_setup").expand_for = filename write_sim(data) - data = read_sim('elegant', elegant_id) + data = _read_sim('elegant', elegant_id) if prev_sim: assert prev_sim.outfile_path - filename = file_name_from_sim_name('elegant', f'{_OMEGA_SIM_NAME}-{run_dir.basename}') + # TODO (gurhar1133): figure out _OMEGA_SIM_NAME + filename = _file_name_from_sim_name('elegant', f'{_OMEGA_SIM_NAME}-{run_dir.basename}') _update_sim(data, filename, prev_sim) t = run_dir.join("omega-elegant-bunch") if prev_sim.sim_type == 'elegant': @@ -422,7 +408,7 @@ def _update_sim(data, filename, prev_sim): import pmd_beamphysics.interfaces.elegant pmd_beamphysics.interfaces.elegant.write_elegant( - genesis_to_pmd(prev_sim), + _genesis_to_pmd(prev_sim), str(t), ) else: @@ -496,6 +482,46 @@ def _coupled_sims_list(data): return sim_list +def _convert_bunched_beam_to_sdds_beam(data, filename): + s = sirepo.sim_data.get_class('elegant') + cmd = LatticeUtil.find_first_command(data, "bunched_beam") + for k in list(cmd.keys()): + if k != '_id': + del cmd[k] + s.update_model_defaults(cmd, 'command_sdds_beam') + cmd._type = 'sdds_beam' + data.models.bunchSource.inputSource = 'sdds_beam' + + +def _file_name_from_sim_name(sim_type, sim_name): + res = re.sub(r'[^0-9a-zA-Z]', '_', sim_name) + res = re.sub(r'^\_+|\_+$', '', res) + res = re.sub(r'\_+', '_', res) + ext = "sdds" if sim_type == "elegant" else "dat" + return f"{res}.{ext}" + + +def _genesis_to_pmd(sim): + import pmd_beamphysics.interfaces.genesis + import pmd_beamphysics.particles + + g = _read_sim("genesis", sim.sim_id) + d = numpy.fromfile(sim.outfile_path, dtype=numpy.float64) + d = d.reshape(int(len(d) / _GENESIS_PARTICLE_COLUMN_COUNT / g.models.electronBeam.npart), _GENESIS_PARTICLE_COLUMN_COUNT, g.models.electronBeam.npart) + # phase must be > 0 to avoid wrapping + d[:,1,:] -= numpy.min(d[:,1,:]) + v = pmd_beamphysics.particles.ParticleGroup( + data=pmd_beamphysics.interfaces.genesis.genesis2_dpa_to_data( + d, xlamds=g.models.radiation.xlamds, current=numpy.array([g.models.electronBeam.curpeak]), + # compute required wavelengths to hold the whole beam + zsep=numpy.max(d[:,1,:]) / (2 * numpy.pi), + ) + ) + # center psi + v.t -= numpy.mean(v.t) + return v + + def _generate_parameters_file(data): res, v = template_common.generate_parameters_file(data) sim_list = _coupled_sims_list(data) @@ -648,3 +674,11 @@ def _sim_list(sim_type): def _template_for_sim_type(sim_type): return sirepo.template.import_module(sim_type) + + +def _read_sim(sim_type, sim_id): + return sirepo.sim_data.get_class(sim_type).sim_db_read_sim(sim_id) + + +def _write_sim(data): + sirepo.sim_data.get_class(data.simulationType).sim_db_save_sim(data) From bde26a8c44a34b736b864feb4bf2f69144ead01c Mon Sep 17 00:00:00 2001 From: git-user Date: Thu, 8 Feb 2024 23:38:00 +0000 Subject: [PATCH 04/23] Fix #6653: ckp --- .../template/omega/parameters.py.jinja | 113 +----- sirepo/template/omega.py | 354 +++++++++++------- 2 files changed, 223 insertions(+), 244 deletions(-) diff --git a/sirepo/package_data/template/omega/parameters.py.jinja b/sirepo/package_data/template/omega/parameters.py.jinja index a1b328e047..22597c9476 100644 --- a/sirepo/package_data/template/omega/parameters.py.jinja +++ b/sirepo/package_data/template/omega/parameters.py.jinja @@ -1,119 +1,32 @@ from pykern import pkio from pykern.pkcollections import PKDict -from rsbeams.rsdata import switchyard -from sirepo.template.lattice import LatticeUtil -import numpy import os -import re import sirepo.pkcli.elegant import sirepo.pkcli.genesis import sirepo.pkcli.opal -import sirepo.sim_data -import sirepo.simulation_db -import sirepo.template.opal -# TODO (gurhar1133): figure out move this to template -_OMEGA_SIM_NAME = '{{ simulation_name }}' -_GENESIS_PARTICLE_COLUMN_COUNT = 6 assert os.environ['SIREPO_SIMULATION_DB_LOGGED_IN_USER'], 'missing user id env var' -#def prep_run_dir(run_dir, data): -# sirepo.simulation_db.prepare_simulation(data, run_dir) -# return pkio.save_chdir(run_dir) - - -def run_elegant(run_dir, elegant_id, prev_sim=None): - with pkio.save_chdir(run_dir): - sirepo.pkcli.elegant.run_elegant() - return f'{run_dir}/run_setup.output.sdds' - - -def run_genesis(run_dir, genesis_id, prev_sim=None): - def _save_lib_file(tmp, filename): - s = sirepo.sim_data.get_class("genesis") - s.lib_file_write( - s.lib_file_name_with_model_field("io", "partfile", filename), - tmp, - ) - tmp.remove() - - def _update_sim(data, filename, tmp): - io = data.models.io - io.partfile = filename - io.ippart = 0 - io.ipradi = 0 - d = numpy.fromfile(tmp, dtype=numpy.float64) - n = len(d) / _GENESIS_PARTICLE_COLUMN_COUNT - factor = 4 * data.models.particleLoading.nbins - data.models.electronBeam.npart = int(n / factor) * factor - # clip particle count to match npart - d = d.reshape(_GENESIS_PARTICLE_COLUMN_COUNT, int(n)) - d = d[:, :data.models.electronBeam.npart] - with open(tmp, 'wb') as f: - d.tofile(f) - _save_lib_file(tmp, filename) - write_sim(data) - - data = read_sim('genesis', genesis_id) - if prev_sim: - assert prev_sim.outfile_path - filename = file_name_from_sim_name( - "genesis", f"{_OMEGA_SIM_NAME}-{run_dir.basename}" +def run_sims(sim_list): + for idx in range(len(sim_list)): + _run( + pkio.py_path(f'run{idx + 1}'), + sim_list[idx].sim_type, ) - t = run_dir.join("omega-genesis-partfile") - if prev_sim.sim_type == "genesis": - # center longitudinal - with open(prev_sim.outfile_path, 'rb') as f: - d = numpy.fromfile(f, dtype=numpy.float64) - d = d.reshape((_GENESIS_PARTICLE_COLUMN_COUNT, int(len(d) / _GENESIS_PARTICLE_COLUMN_COUNT))) - d[1,:] -= numpy.mean(d[1,:]) - with open(t, 'wb') as f: - d.tofile(f) - else: - if prev_sim.sim_type == "elegant": - p = switchyard.read_elegant(f"{prev_sim.outfile_path}") - else: - assert prev_sim.sim_type == "opal" - p = switchyard.read_opal(f"{prev_sim.outfile_path}") - particle_data = numpy.zeros([_GENESIS_PARTICLE_COLUMN_COUNT, len(p.x)]) - for i, col in enumerate(['pt', 'ct', 'x', 'y', 'ux', 'uy']): - particle_data[i, :] = getattr(p, col) - #TODO(pjm): [0] should be gamma not momentum - particle_data[1] *= 2 * numpy.pi / data.models.radiation.xlamds - particle_data[1] -= numpy.mean(particle_data[1]) - with open(t, "wb") as f: - particle_data.tofile(f) - _update_sim(data, filename, t) - data.computeModel = 'animation' - if 'report' in data: - del data['report'] - with prep_run_dir(run_dir, data): - sirepo.pkcli.genesis.run_genesis(run_dir) - return f"{run_dir}/genesis.out.dpa" -def run_opal(run_dir, opal_id, prev_sim=None): +def _run(run_dir, sim_type): with pkio.save_chdir(run_dir): - sirepo.pkcli.opal.run_opal(with_mpi=True) - return f'{run_dir}/{sirepo.template.opal._OPAL_H5_FILE}' - - -def run_sims(sim_list): - prev = None - for idx in range(len(sim_list)): - run_dir = pkio.py_path(f'run{idx + 1}') - s = sim_list[idx] - if s.sim_type == 'opal': - s.outfile_path = run_opal(run_dir, s.sim_id, prev) - elif s.sim_type == 'elegant': - s.outfile_path = run_elegant(run_dir, s.sim_id, prev) - elif s.sim_type == 'genesis': - s.outfile_path = run_genesis(run_dir, s.sim_id, prev) + if sim_type == 'opal': + sirepo.pkcli.opal.run_opal(with_mpi=True) + elif sim_type == 'elegant': + sirepo.pkcli.elegant.run_elegant() + elif sim_type == 'genesis': + sirepo.pkcli.genesis.run_genesis(run_dir) else: - raise AssertionError(f"unhandled sim_type={s.sim_type}") - prev = s + raise AssertionError(f"unhandled sim_type={sim_type}") run_sims([ diff --git a/sirepo/template/omega.py b/sirepo/template/omega.py index a4391355c3..fcc2acc6f0 100644 --- a/sirepo/template/omega.py +++ b/sirepo/template/omega.py @@ -7,6 +7,7 @@ from pykern import pkjson from pykern.pkcollections import PKDict from pykern.pkdebug import pkdp, pkdc, pkdlog +from rsbeams.rsdata import switchyard from sirepo import simulation_db from sirepo.template import template_common from sirepo.template.lattice import LatticeUtil @@ -19,7 +20,7 @@ import re import sirepo.sim_data import sirepo.template -# TODO (gurhar1133): need to alphabetical sort functions + _SIM_DATA, SIM_TYPE, SCHEMA = sirepo.sim_data.template_globals() _PHASE_PLOT_COUNT = 4 _PHASE_PLOTS = PKDict( @@ -107,6 +108,7 @@ Cx="run_setup.centroid.sdds", Cy="run_setup.centroid.sdds", ) +_GENESIS_PARTICLE_COLUMN_COUNT = 6 _RELATED_SIMS_FOLDER = "/Omega" _SUCCESS_OUTPUT_FILE = PKDict( elegant="run_setup.output.sdds", @@ -282,151 +284,14 @@ def stateful_compute_get_opal_sim_list(**kwargs): def write_parameters(data, run_dir, is_parallel): - prev = None - sim_list = _coupled_sims_list(data) - pkdp("\n\n\n sim_list={}", sim_list) - for idx in range(len(sim_list)): - s = sim_list[idx] - d = run_dir.join(f'run{idx + 1}') - pkio.unchecked_remove(d) - pkio.mkdir_parent(d) - if s.sim_type == "opal": - _prepare_opal(d, s.sim_id, prev) - elif s.sim_type == "elegant": - _prepare_elegant(d, s.sim_id, prev) + pkdp("\n\n\n data.simulation_name={}", data.models.simulation.name) + _prepare_subsims(data, run_dir, data.models.simulation.name) pkio.write_text( run_dir.join(template_common.PARAMETERS_PYTHON_FILE), _generate_parameters_file(data), ) -def _prepare_opal(run_dir, opal_id, prev_sim=None): - def _save_lib_file(tmp, filename): - s = sirepo.sim_data.get_class("opal") - s.lib_file_write( - s.lib_file_name_with_model_field("command_distribution", "fname", filename), - tmp, - ) - tmp.remove() - - def _update_sim(data, filename, tmp): - d = LatticeUtil.find_first_command(data, "distribution") - d.type = "FROMFILE" - d.fname = filename - n = 0 - zsum = 0 - psum = 0 - # calculate z offset - with pkio.open_text(tmp) as f: - for line in f: - if not n: - n = int(line) - continue - r = [v for v in re.split(r"\s+", line.strip())] - zsum += float(r[4]) - psum += float(r[5]) - d.offsetz = -zsum / n - - b = LatticeUtil.find_first_command(data, "beam") - b.npart = n - b.gamma = 0 - b.energy = 0 - #TODO(pjm): handle other particle types - mass_and_charge = PKDict( - ELECTRON=0.51099895000e-03, - PROTON=0.93827208816, - ) - b.pc = psum / n * mass_and_charge[b.particle] - _save_lib_file(tmp, filename) - write_sim(data) - - data = _read_sim('opal', opal_id) - if prev_sim: - # TODO (gurhar1133): figure out _OMEGA_SIM_NAME - filename = _file_name_from_sim_name('opal', f'{_OMEGA_SIM_NAME}-{run_dir.basename}') - t = run_dir.join("opal-command_distribution") - if prev_sim.sim_type == 'genesis': - import pmd_beamphysics.interfaces.opal - - pmd_beamphysics.interfaces.opal.write_opal( - _genesis_to_pmd(prev_sim), - str(t), - ) - else: - assert prev_sim.outfile_path - sw = switchyard.Switchyard() - sw.read(f'{prev_sim.outfile_path}', prev_sim.sim_type) - sw.write(str(t), "opal") - _update_sim(data, filename, t) - data.computeModel = 'animation' - LatticeUtil.find_first_command(data, "option").psdumpfreq = 0 - sirepo.simulation_db.prepare_simulation(data, run_dir) - # pkio.save_chdir(run_dir) - # with prep_run_dir(run_dir, data): - # sirepo.pkcli.opal.run_opal(with_mpi=True) - # return f'{run_dir}/{sirepo.template.opal._OPAL_H5_FILE}' - - - -def _prepare_elegant(run_dir, elegant_id, prev_sim=None): - def _save_lib_files(tmp, filename): - s = sirepo.sim_data.get_class("elegant") - b = s.lib_file_name_with_model_field("bunchFile", "sourceFile", filename) - s.lib_file_write(b, tmp) - c = s.sim_db_client() - c.copy( - c.uri(c.LIB_DIR, b), - c.uri( - c.LIB_DIR, - s.lib_file_name_with_model_field("command_run_setup", "expand_for", filename), - ), - ) - tmp.remove() - - def _update_sim(data, filename, prev_sim): - if data.models.bunchSource.inputSource == 'bunched_beam': - _convert_bunched_beam_to_sdds_beam(data, filename) - assert data.models.bunchSource.inputSource == 'sdds_beam' - cmd = LatticeUtil.find_first_command(data, "sdds_beam") - cmd.input = filename - cmd.center_arrival_time = '1' - cmd.center_transversely = '1' - cmd.reverse_t_sign = '1' if prev_sim.sim_type == 'genesis' else '0' - LatticeUtil.find_first_command(data, "run_setup").expand_for = filename - write_sim(data) - - data = _read_sim('elegant', elegant_id) - if prev_sim: - assert prev_sim.outfile_path - # TODO (gurhar1133): figure out _OMEGA_SIM_NAME - filename = _file_name_from_sim_name('elegant', f'{_OMEGA_SIM_NAME}-{run_dir.basename}') - _update_sim(data, filename, prev_sim) - t = run_dir.join("omega-elegant-bunch") - if prev_sim.sim_type == 'elegant': - pkio.py_path(prev_sim.outfile_path).copy(t) - elif prev_sim.sim_type == 'genesis': - import pmd_beamphysics.interfaces.elegant - - pmd_beamphysics.interfaces.elegant.write_elegant( - _genesis_to_pmd(prev_sim), - str(t), - ) - else: - sw = switchyard.Switchyard() - sw.read(f'{prev_sim.outfile_path}', prev_sim.sim_type) - sw.write(str(t), 'elegant') - _save_lib_files(t, filename) - data.computeModel = 'animation' - if 'report' in data: - del data['report'] - sirepo.simulation_db.prepare_simulation(data, run_dir) - # pkio.save_chdir(run_dir) - # with prep_run_dir(run_dir, data): - # sirepo.pkcli.opal.run_opal(with_mpi=True) - # return f'{run_dir}/{sirepo.template.opal._OPAL_H5_FILE}' - - - def _extract_elegant_beam_plot(frame_args): # this is tricky because the data could come from 2 different elegant output files files = PKDict() @@ -648,6 +513,211 @@ def _plot_phase(sim_type, frame_args): raise AssertionError("unhandled sim_type for sim_frame(): {}".format(sim_type)) +def _prepare_elegant(run_dir, elegant_id, omega_sim_name, prev_sim=None): + def _save_lib_files(tmp, filename): + s = sirepo.sim_data.get_class("elegant") + b = s.lib_file_name_with_model_field("bunchFile", "sourceFile", filename) + s.lib_file_write(b, tmp) + c = s.sim_db_client() + c.copy( + c.uri(c.LIB_DIR, b), + c.uri( + c.LIB_DIR, + s.lib_file_name_with_model_field("command_run_setup", "expand_for", filename), + ), + ) + tmp.remove() + + def _update_sim(data, filename, prev_sim): + if data.models.bunchSource.inputSource == 'bunched_beam': + _convert_bunched_beam_to_sdds_beam(data, filename) + assert data.models.bunchSource.inputSource == 'sdds_beam' + cmd = LatticeUtil.find_first_command(data, "sdds_beam") + cmd.input = filename + cmd.center_arrival_time = '1' + cmd.center_transversely = '1' + cmd.reverse_t_sign = '1' if prev_sim.sim_type == 'genesis' else '0' + LatticeUtil.find_first_command(data, "run_setup").expand_for = filename + _write_sim(data) + + data = _read_sim('elegant', elegant_id) + if prev_sim: + assert prev_sim.outfile_path + filename = _file_name_from_sim_name('elegant', f'{omega_sim_name}-{run_dir.basename}') + _update_sim(data, filename, prev_sim) + t = run_dir.join("omega-elegant-bunch") + if prev_sim.sim_type == 'elegant': + pkio.py_path(prev_sim.outfile_path).copy(t) + elif prev_sim.sim_type == 'genesis': + import pmd_beamphysics.interfaces.elegant + + pmd_beamphysics.interfaces.elegant.write_elegant( + _genesis_to_pmd(prev_sim), + str(t), + ) + else: + # TODO (gurhar1133): this read/write from the output of a previous sim stuff needs to be done in parameters.py + sw = switchyard.Switchyard() + sw.read(f'{prev_sim.outfile_path}', prev_sim.sim_type) + sw.write(str(t), 'elegant') + _save_lib_files(t, filename) + data.computeModel = 'animation' + if 'report' in data: + del data['report'] + sirepo.simulation_db.prepare_simulation(data, run_dir) + return f'{run_dir}/run_setup.output.sdds' + + +def _prepare_genesis(run_dir, genesis_id, omega_sim_name, prev_sim=None): + def _save_lib_file(tmp, filename): + s = sirepo.sim_data.get_class("genesis") + s.lib_file_write( + s.lib_file_name_with_model_field("io", "partfile", filename), + tmp, + ) + tmp.remove() + + def _update_sim(data, filename, tmp): + io = data.models.io + io.partfile = filename + io.ippart = 0 + io.ipradi = 0 + d = numpy.fromfile(tmp, dtype=numpy.float64) + n = len(d) / _GENESIS_PARTICLE_COLUMN_COUNT + factor = 4 * data.models.particleLoading.nbins + data.models.electronBeam.npart = int(n / factor) * factor + # clip particle count to match npart + d = d.reshape(_GENESIS_PARTICLE_COLUMN_COUNT, int(n)) + d = d[:, :data.models.electronBeam.npart] + with open(tmp, 'wb') as f: + d.tofile(f) + _save_lib_file(tmp, filename) + _write_sim(data) + + data = _read_sim('genesis', genesis_id) + if prev_sim: + assert prev_sim.outfile_path + filename = file_name_from_sim_name( + "genesis", f"{omega_sim_name}-{run_dir.basename}" + ) + t = run_dir.join("omega-genesis-partfile") + if prev_sim.sim_type == "genesis": + # center longitudinal + with open(prev_sim.outfile_path, 'rb') as f: + d = numpy.fromfile(f, dtype=numpy.float64) + d = d.reshape((_GENESIS_PARTICLE_COLUMN_COUNT, int(len(d) / _GENESIS_PARTICLE_COLUMN_COUNT))) + d[1,:] -= numpy.mean(d[1,:]) + with open(t, 'wb') as f: + d.tofile(f) + else: + if prev_sim.sim_type == "elegant": + p = switchyard.read_elegant(f"{prev_sim.outfile_path}") + else: + assert prev_sim.sim_type == "opal" + p = switchyard.read_opal(f"{prev_sim.outfile_path}") + particle_data = numpy.zeros([_GENESIS_PARTICLE_COLUMN_COUNT, len(p.x)]) + for i, col in enumerate(['pt', 'ct', 'x', 'y', 'ux', 'uy']): + particle_data[i, :] = getattr(p, col) + #TODO(pjm): [0] should be gamma not momentum + particle_data[1] *= 2 * numpy.pi / data.models.radiation.xlamds + particle_data[1] -= numpy.mean(particle_data[1]) + with open(t, "wb") as f: + particle_data.tofile(f) + _update_sim(data, filename, t) + data.computeModel = 'animation' + if 'report' in data: + del data['report'] + sirepo.simulation_db.prepare_simulation(data, run_dir) + return f"{run_dir}/genesis.out.dpa" + + +def _prepare_opal(run_dir, opal_id, omega_sim_name, prev_sim=None): + def _save_lib_file(tmp, filename): + s = sirepo.sim_data.get_class("opal") + s.lib_file_write( + s.lib_file_name_with_model_field("command_distribution", "fname", filename), + tmp, + ) + tmp.remove() + + def _update_sim(data, filename, tmp): + d = LatticeUtil.find_first_command(data, "distribution") + d.type = "FROMFILE" + d.fname = filename + n = 0 + zsum = 0 + psum = 0 + # calculate z offset + with pkio.open_text(tmp) as f: + for line in f: + if not n: + n = int(line) + continue + r = [v for v in re.split(r"\s+", line.strip())] + zsum += float(r[4]) + psum += float(r[5]) + d.offsetz = -zsum / n + + b = LatticeUtil.find_first_command(data, "beam") + b.npart = n + b.gamma = 0 + b.energy = 0 + #TODO(pjm): handle other particle types + mass_and_charge = PKDict( + ELECTRON=0.51099895000e-03, + PROTON=0.93827208816, + ) + b.pc = psum / n * mass_and_charge[b.particle] + _save_lib_file(tmp, filename) + _write_sim(data) + + data = _read_sim('opal', opal_id) + if prev_sim: + filename = _file_name_from_sim_name('opal', f'{omega_sim_name}-{run_dir.basename}') + t = run_dir.join("opal-command_distribution") + if prev_sim.sim_type == 'genesis': + import pmd_beamphysics.interfaces.opal + + pmd_beamphysics.interfaces.opal.write_opal( + _genesis_to_pmd(prev_sim), + str(t), + ) + else: + assert prev_sim.outfile_path + sw = switchyard.Switchyard() + sw.read(f'{prev_sim.outfile_path}', prev_sim.sim_type) + sw.write(str(t), "opal") + _update_sim(data, filename, t) + data.computeModel = 'animation' + LatticeUtil.find_first_command(data, "option").psdumpfreq = 0 + sirepo.simulation_db.prepare_simulation(data, run_dir) + return f'{run_dir}/{sirepo.template.opal._OPAL_H5_FILE}' + + +def _prepare_subsims(data, run_dir, omega_sim_name): + prev = None + sim_list = _coupled_sims_list(data) + pkdp("\n\n\n sim_list={}", sim_list) + for idx in range(len(sim_list)): + s = sim_list[idx] + d = run_dir.join(f'run{idx + 1}') + pkio.unchecked_remove(d) + pkio.mkdir_parent(d) + if s.sim_type == "opal": + s.outfile_path = _prepare_opal(d, s.sim_id, omega_sim_name, prev_sim=prev) + elif s.sim_type == "elegant": + s.outfile_path = _prepare_elegant(d, s.sim_id, omega_sim_name, prev_sim=prev) + elif s.sim_type == "genesis": + s.outfile_path = _prepare_genesis(d, s.sim_id, omega_sim_name, prev_sim=prev) + else: + raise AssertionError(f"unhandled sim_type={s.sim_type}") + prev = s + + +def _read_sim(sim_type, sim_id): + return sirepo.sim_data.get_class(sim_type).sim_db_read_sim(sim_id) + + def _sim_dir(run_dir, sim_count): return run_dir.join(f"run{sim_count}") @@ -676,9 +746,5 @@ def _template_for_sim_type(sim_type): return sirepo.template.import_module(sim_type) -def _read_sim(sim_type, sim_id): - return sirepo.sim_data.get_class(sim_type).sim_db_read_sim(sim_id) - - def _write_sim(data): sirepo.sim_data.get_class(data.simulationType).sim_db_save_sim(data) From 4b805b6cb6b2588210addcfbe254541c54d5bcc7 Mon Sep 17 00:00:00 2001 From: git-user Date: Fri, 9 Feb 2024 21:51:14 +0000 Subject: [PATCH 05/23] Fix #6653: ckp --- sirepo/template/omega.py | 34 +++++++++++++++++----------------- 1 file changed, 17 insertions(+), 17 deletions(-) diff --git a/sirepo/template/omega.py b/sirepo/template/omega.py index fcc2acc6f0..fa008aef22 100644 --- a/sirepo/template/omega.py +++ b/sirepo/template/omega.py @@ -544,23 +544,23 @@ def _update_sim(data, filename, prev_sim): if prev_sim: assert prev_sim.outfile_path filename = _file_name_from_sim_name('elegant', f'{omega_sim_name}-{run_dir.basename}') - _update_sim(data, filename, prev_sim) - t = run_dir.join("omega-elegant-bunch") - if prev_sim.sim_type == 'elegant': - pkio.py_path(prev_sim.outfile_path).copy(t) - elif prev_sim.sim_type == 'genesis': - import pmd_beamphysics.interfaces.elegant - - pmd_beamphysics.interfaces.elegant.write_elegant( - _genesis_to_pmd(prev_sim), - str(t), - ) - else: - # TODO (gurhar1133): this read/write from the output of a previous sim stuff needs to be done in parameters.py - sw = switchyard.Switchyard() - sw.read(f'{prev_sim.outfile_path}', prev_sim.sim_type) - sw.write(str(t), 'elegant') - _save_lib_files(t, filename) + # _update_sim(data, filename, prev_sim) + # t = run_dir.join("omega-elegant-bunch") + # if prev_sim.sim_type == 'elegant': + # pkio.py_path(prev_sim.outfile_path).copy(t) + # elif prev_sim.sim_type == 'genesis': + # import pmd_beamphysics.interfaces.elegant + + # pmd_beamphysics.interfaces.elegant.write_elegant( + # _genesis_to_pmd(prev_sim), + # str(t), + # ) + # else: + # # TODO (gurhar1133): this read/write from the output of a previous sim stuff needs to be done in parameters.py + # sw = switchyard.Switchyard() + # sw.read(f'{prev_sim.outfile_path}', prev_sim.sim_type) + # sw.write(str(t), 'elegant') + # _save_lib_files(t, filename) data.computeModel = 'animation' if 'report' in data: del data['report'] From 8f96c61189e901edc7db53ee5729d793dbec0ff2 Mon Sep 17 00:00:00 2001 From: git-user Date: Fri, 9 Feb 2024 22:48:06 +0000 Subject: [PATCH 06/23] Fix #6653: just gen sim dirs no prev output --- sirepo/template/omega.py | 245 ++++----------------------------------- 1 file changed, 21 insertions(+), 224 deletions(-) diff --git a/sirepo/template/omega.py b/sirepo/template/omega.py index fa008aef22..b72ca291ed 100644 --- a/sirepo/template/omega.py +++ b/sirepo/template/omega.py @@ -284,8 +284,7 @@ def stateful_compute_get_opal_sim_list(**kwargs): def write_parameters(data, run_dir, is_parallel): - pkdp("\n\n\n data.simulation_name={}", data.models.simulation.name) - _prepare_subsims(data, run_dir, data.models.simulation.name) + _prepare_subsims(data, run_dir) pkio.write_text( run_dir.join(template_common.PARAMETERS_PYTHON_FILE), _generate_parameters_file(data), @@ -347,46 +346,6 @@ def _coupled_sims_list(data): return sim_list -def _convert_bunched_beam_to_sdds_beam(data, filename): - s = sirepo.sim_data.get_class('elegant') - cmd = LatticeUtil.find_first_command(data, "bunched_beam") - for k in list(cmd.keys()): - if k != '_id': - del cmd[k] - s.update_model_defaults(cmd, 'command_sdds_beam') - cmd._type = 'sdds_beam' - data.models.bunchSource.inputSource = 'sdds_beam' - - -def _file_name_from_sim_name(sim_type, sim_name): - res = re.sub(r'[^0-9a-zA-Z]', '_', sim_name) - res = re.sub(r'^\_+|\_+$', '', res) - res = re.sub(r'\_+', '_', res) - ext = "sdds" if sim_type == "elegant" else "dat" - return f"{res}.{ext}" - - -def _genesis_to_pmd(sim): - import pmd_beamphysics.interfaces.genesis - import pmd_beamphysics.particles - - g = _read_sim("genesis", sim.sim_id) - d = numpy.fromfile(sim.outfile_path, dtype=numpy.float64) - d = d.reshape(int(len(d) / _GENESIS_PARTICLE_COLUMN_COUNT / g.models.electronBeam.npart), _GENESIS_PARTICLE_COLUMN_COUNT, g.models.electronBeam.npart) - # phase must be > 0 to avoid wrapping - d[:,1,:] -= numpy.min(d[:,1,:]) - v = pmd_beamphysics.particles.ParticleGroup( - data=pmd_beamphysics.interfaces.genesis.genesis2_dpa_to_data( - d, xlamds=g.models.radiation.xlamds, current=numpy.array([g.models.electronBeam.curpeak]), - # compute required wavelengths to hold the whole beam - zsep=numpy.max(d[:,1,:]) / (2 * numpy.pi), - ) - ) - # center psi - v.t -= numpy.mean(v.t) - return v - - def _generate_parameters_file(data): res, v = template_common.generate_parameters_file(data) sim_list = _coupled_sims_list(data) @@ -513,205 +472,47 @@ def _plot_phase(sim_type, frame_args): raise AssertionError("unhandled sim_type for sim_frame(): {}".format(sim_type)) -def _prepare_elegant(run_dir, elegant_id, omega_sim_name, prev_sim=None): - def _save_lib_files(tmp, filename): - s = sirepo.sim_data.get_class("elegant") - b = s.lib_file_name_with_model_field("bunchFile", "sourceFile", filename) - s.lib_file_write(b, tmp) - c = s.sim_db_client() - c.copy( - c.uri(c.LIB_DIR, b), - c.uri( - c.LIB_DIR, - s.lib_file_name_with_model_field("command_run_setup", "expand_for", filename), - ), - ) - tmp.remove() - - def _update_sim(data, filename, prev_sim): - if data.models.bunchSource.inputSource == 'bunched_beam': - _convert_bunched_beam_to_sdds_beam(data, filename) - assert data.models.bunchSource.inputSource == 'sdds_beam' - cmd = LatticeUtil.find_first_command(data, "sdds_beam") - cmd.input = filename - cmd.center_arrival_time = '1' - cmd.center_transversely = '1' - cmd.reverse_t_sign = '1' if prev_sim.sim_type == 'genesis' else '0' - LatticeUtil.find_first_command(data, "run_setup").expand_for = filename - _write_sim(data) - - data = _read_sim('elegant', elegant_id) - if prev_sim: - assert prev_sim.outfile_path - filename = _file_name_from_sim_name('elegant', f'{omega_sim_name}-{run_dir.basename}') - # _update_sim(data, filename, prev_sim) - # t = run_dir.join("omega-elegant-bunch") - # if prev_sim.sim_type == 'elegant': - # pkio.py_path(prev_sim.outfile_path).copy(t) - # elif prev_sim.sim_type == 'genesis': - # import pmd_beamphysics.interfaces.elegant - - # pmd_beamphysics.interfaces.elegant.write_elegant( - # _genesis_to_pmd(prev_sim), - # str(t), - # ) - # else: - # # TODO (gurhar1133): this read/write from the output of a previous sim stuff needs to be done in parameters.py - # sw = switchyard.Switchyard() - # sw.read(f'{prev_sim.outfile_path}', prev_sim.sim_type) - # sw.write(str(t), 'elegant') - # _save_lib_files(t, filename) - data.computeModel = 'animation' - if 'report' in data: - del data['report'] +def _prepare_elegant(run_dir, elegant_id): + data = _read_sim("elegant", elegant_id) + data.computeModel = "animation" + if "report" in data: + del data["report"] sirepo.simulation_db.prepare_simulation(data, run_dir) - return f'{run_dir}/run_setup.output.sdds' + return f"{run_dir}/run_setup.output.sdds" -def _prepare_genesis(run_dir, genesis_id, omega_sim_name, prev_sim=None): - def _save_lib_file(tmp, filename): - s = sirepo.sim_data.get_class("genesis") - s.lib_file_write( - s.lib_file_name_with_model_field("io", "partfile", filename), - tmp, - ) - tmp.remove() - - def _update_sim(data, filename, tmp): - io = data.models.io - io.partfile = filename - io.ippart = 0 - io.ipradi = 0 - d = numpy.fromfile(tmp, dtype=numpy.float64) - n = len(d) / _GENESIS_PARTICLE_COLUMN_COUNT - factor = 4 * data.models.particleLoading.nbins - data.models.electronBeam.npart = int(n / factor) * factor - # clip particle count to match npart - d = d.reshape(_GENESIS_PARTICLE_COLUMN_COUNT, int(n)) - d = d[:, :data.models.electronBeam.npart] - with open(tmp, 'wb') as f: - d.tofile(f) - _save_lib_file(tmp, filename) - _write_sim(data) - - data = _read_sim('genesis', genesis_id) - if prev_sim: - assert prev_sim.outfile_path - filename = file_name_from_sim_name( - "genesis", f"{omega_sim_name}-{run_dir.basename}" - ) - t = run_dir.join("omega-genesis-partfile") - if prev_sim.sim_type == "genesis": - # center longitudinal - with open(prev_sim.outfile_path, 'rb') as f: - d = numpy.fromfile(f, dtype=numpy.float64) - d = d.reshape((_GENESIS_PARTICLE_COLUMN_COUNT, int(len(d) / _GENESIS_PARTICLE_COLUMN_COUNT))) - d[1,:] -= numpy.mean(d[1,:]) - with open(t, 'wb') as f: - d.tofile(f) - else: - if prev_sim.sim_type == "elegant": - p = switchyard.read_elegant(f"{prev_sim.outfile_path}") - else: - assert prev_sim.sim_type == "opal" - p = switchyard.read_opal(f"{prev_sim.outfile_path}") - particle_data = numpy.zeros([_GENESIS_PARTICLE_COLUMN_COUNT, len(p.x)]) - for i, col in enumerate(['pt', 'ct', 'x', 'y', 'ux', 'uy']): - particle_data[i, :] = getattr(p, col) - #TODO(pjm): [0] should be gamma not momentum - particle_data[1] *= 2 * numpy.pi / data.models.radiation.xlamds - particle_data[1] -= numpy.mean(particle_data[1]) - with open(t, "wb") as f: - particle_data.tofile(f) - _update_sim(data, filename, t) - data.computeModel = 'animation' - if 'report' in data: - del data['report'] +def _prepare_genesis(run_dir, genesis_id): + data = _read_sim("genesis", genesis_id) + data.computeModel = "animation" + if "report" in data: + del data["report"] sirepo.simulation_db.prepare_simulation(data, run_dir) return f"{run_dir}/genesis.out.dpa" -def _prepare_opal(run_dir, opal_id, omega_sim_name, prev_sim=None): - def _save_lib_file(tmp, filename): - s = sirepo.sim_data.get_class("opal") - s.lib_file_write( - s.lib_file_name_with_model_field("command_distribution", "fname", filename), - tmp, - ) - tmp.remove() - - def _update_sim(data, filename, tmp): - d = LatticeUtil.find_first_command(data, "distribution") - d.type = "FROMFILE" - d.fname = filename - n = 0 - zsum = 0 - psum = 0 - # calculate z offset - with pkio.open_text(tmp) as f: - for line in f: - if not n: - n = int(line) - continue - r = [v for v in re.split(r"\s+", line.strip())] - zsum += float(r[4]) - psum += float(r[5]) - d.offsetz = -zsum / n - - b = LatticeUtil.find_first_command(data, "beam") - b.npart = n - b.gamma = 0 - b.energy = 0 - #TODO(pjm): handle other particle types - mass_and_charge = PKDict( - ELECTRON=0.51099895000e-03, - PROTON=0.93827208816, - ) - b.pc = psum / n * mass_and_charge[b.particle] - _save_lib_file(tmp, filename) - _write_sim(data) - - data = _read_sim('opal', opal_id) - if prev_sim: - filename = _file_name_from_sim_name('opal', f'{omega_sim_name}-{run_dir.basename}') - t = run_dir.join("opal-command_distribution") - if prev_sim.sim_type == 'genesis': - import pmd_beamphysics.interfaces.opal - - pmd_beamphysics.interfaces.opal.write_opal( - _genesis_to_pmd(prev_sim), - str(t), - ) - else: - assert prev_sim.outfile_path - sw = switchyard.Switchyard() - sw.read(f'{prev_sim.outfile_path}', prev_sim.sim_type) - sw.write(str(t), "opal") - _update_sim(data, filename, t) - data.computeModel = 'animation' +def _prepare_opal(run_dir, opal_id): + data = _read_sim("opal", opal_id) + data.computeModel = "animation" LatticeUtil.find_first_command(data, "option").psdumpfreq = 0 sirepo.simulation_db.prepare_simulation(data, run_dir) - return f'{run_dir}/{sirepo.template.opal._OPAL_H5_FILE}' + return f"{run_dir}/{sirepo.template.opal._OPAL_H5_FILE}" -def _prepare_subsims(data, run_dir, omega_sim_name): - prev = None +def _prepare_subsims(data, run_dir): sim_list = _coupled_sims_list(data) - pkdp("\n\n\n sim_list={}", sim_list) for idx in range(len(sim_list)): s = sim_list[idx] - d = run_dir.join(f'run{idx + 1}') + d = run_dir.join(f"run{idx + 1}") pkio.unchecked_remove(d) pkio.mkdir_parent(d) if s.sim_type == "opal": - s.outfile_path = _prepare_opal(d, s.sim_id, omega_sim_name, prev_sim=prev) + _prepare_opal(d, s.sim_id) elif s.sim_type == "elegant": - s.outfile_path = _prepare_elegant(d, s.sim_id, omega_sim_name, prev_sim=prev) + _prepare_elegant(d, s.sim_id) elif s.sim_type == "genesis": - s.outfile_path = _prepare_genesis(d, s.sim_id, omega_sim_name, prev_sim=prev) + _prepare_genesis(d, s.sim_id) else: raise AssertionError(f"unhandled sim_type={s.sim_type}") - prev = s def _read_sim(sim_type, sim_id): @@ -744,7 +545,3 @@ def _sim_list(sim_type): def _template_for_sim_type(sim_type): return sirepo.template.import_module(sim_type) - - -def _write_sim(data): - sirepo.sim_data.get_class(data.simulationType).sim_db_save_sim(data) From b87684cc8a7288a01d4109bc57bddeada1f0598e Mon Sep 17 00:00:00 2001 From: git-user Date: Fri, 9 Feb 2024 22:49:57 +0000 Subject: [PATCH 07/23] Fix #6653: cleanup --- sirepo/template/omega.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/sirepo/template/omega.py b/sirepo/template/omega.py index b72ca291ed..2d6a07c53c 100644 --- a/sirepo/template/omega.py +++ b/sirepo/template/omega.py @@ -7,7 +7,6 @@ from pykern import pkjson from pykern.pkcollections import PKDict from pykern.pkdebug import pkdp, pkdc, pkdlog -from rsbeams.rsdata import switchyard from sirepo import simulation_db from sirepo.template import template_common from sirepo.template.lattice import LatticeUtil @@ -108,7 +107,6 @@ Cx="run_setup.centroid.sdds", Cy="run_setup.centroid.sdds", ) -_GENESIS_PARTICLE_COLUMN_COUNT = 6 _RELATED_SIMS_FOLDER = "/Omega" _SUCCESS_OUTPUT_FILE = PKDict( elegant="run_setup.output.sdds", From d7664eca55c9b322fadd784cd5f51228143e35e0 Mon Sep 17 00:00:00 2001 From: moellep Date: Mon, 15 Apr 2024 18:02:57 +0000 Subject: [PATCH 08/23] include folder name in sim list --- .../static/js/sirepo-components.js | 40 +++++++++++++------ 1 file changed, 27 insertions(+), 13 deletions(-) diff --git a/sirepo/package_data/static/js/sirepo-components.js b/sirepo/package_data/static/js/sirepo-components.js index c7bbf22ed7..a3ca8a7549 100644 --- a/sirepo/package_data/static/js/sirepo-components.js +++ b/sirepo/package_data/static/js/sirepo-components.js @@ -4901,23 +4901,38 @@ SIREPO.app.directive('simList', function(appState, requestSender) { route: '@', }, template: ` - +
- +
`, controller: function($scope) { - $scope.simList = null; - // special processing of the item's name if necessary - $scope.itemName = function(item) { - return item.invalidMsg ? `${item.name} <${item.invalidMsg}>` : item.name; - }; + function buildList(simList) { + $scope.items = []; + for (const s of simList) { + $scope.items.push({ + simulationId: s.simulationId, + isInvalid: s.invalidMsg ? true : false, + name: itemName(s), + }); + } + $scope.items.sort((a, b) => a.name.localeCompare(b.name)); + } - $scope.openSimulation = function() { + function itemName(sim) { + const n = sim.folder === '/' + ? `/${sim.name}` + : `${sim.folder}/${sim.name}`; + return sim.invalidMsg + ? `${n} <${sim.invalidMsg}>` + : n; + } + + $scope.openSimulation = () => { if ($scope.model && $scope.model[$scope.field]) { requestSender.openSimulation( $scope.code, @@ -4926,14 +4941,13 @@ SIREPO.app.directive('simList', function(appState, requestSender) { ); } }; - appState.whenModelsLoaded($scope, function() { + + appState.whenModelsLoaded($scope, () => { requestSender.sendStatefulCompute( appState, - function(data) { + (data) => { if (appState.isLoaded() && data.simList) { - $scope.simList = data.simList.sort(function(a, b) { - return a.name.localeCompare(b.name); - }); + buildList(data.simList); } }, { From f51cca2d425d80d476285b7786d05e72b7abbdb7 Mon Sep 17 00:00:00 2001 From: moellep Date: Mon, 15 Apr 2024 18:04:33 +0000 Subject: [PATCH 09/23] for #6653 use lume classes for running omega sims --- sirepo/lib.py | 12 +- sirepo/package_data/static/js/omega.js | 3 + .../genesis/lib/io-maginfile.tesla.lat | 220 +++++++++++++++++- .../genesis/lib/io-maginfile.ttf-ii.lat | 57 +++-- .../template/omega/parameters.py.jinja | 66 +++--- sirepo/template/genesis.py | 10 +- sirepo/template/omega.py | 132 +++++------ 7 files changed, 357 insertions(+), 143 deletions(-) diff --git a/sirepo/lib.py b/sirepo/lib.py index 63711cc407..7c450ca9cd 100644 --- a/sirepo/lib.py +++ b/sirepo/lib.py @@ -21,11 +21,12 @@ class LibAdapterBase: """Common functionality between code specific LibAdapter implementations.""" - def __init__(self, ignore_files=None): + def __init__(self, ignore_files=None, update_filenames=False): m = inspect.getmodule(self) self._sim_data, _, self._schema = sirepo.sim_data.template_globals(m.SIM_TYPE) self._code_var = m.code_var self._ignore_files = ignore_files if ignore_files else [] + self._update_filenames = update_filenames def _convert(self, data): def _model(model, name): @@ -77,7 +78,9 @@ def _write_input_files(self, data, source_path, dest_dir): for f in set( LatticeUtil(data, self._schema) .iterate_models( - lattice.InputFileIterator(self._sim_data, update_filenames=False), + lattice.InputFileIterator( + self._sim_data, update_filenames=self._update_filenames + ), ) .result, ): @@ -112,11 +115,12 @@ class Importer: ignore_files (list): files ignored during verification and symlink routines [None] """ - def __init__(self, sim_type, ignore_files=None): + def __init__(self, sim_type, ignore_files=None, update_filenames=False): import sirepo.template self.__adapter = sirepo.template.import_module(sim_type).LibAdapter( - ignore_files or [] + ignore_files or [], + update_filenames, ) def parse_file(self, path): diff --git a/sirepo/package_data/static/js/omega.js b/sirepo/package_data/static/js/omega.js index 781bb14c58..c35535d369 100644 --- a/sirepo/package_data/static/js/omega.js +++ b/sirepo/package_data/static/js/omega.js @@ -164,6 +164,9 @@ SIREPO.app.directive('beamAndPhasePlots', function(appState, omegaService) { }; $scope.$on('modelChanged', (e, name) => { + if (! $scope.reports) { + return; + } for (const sim of $scope.reports) { if (name === sim[1][0].modelKey) { const updated = []; diff --git a/sirepo/package_data/template/genesis/lib/io-maginfile.tesla.lat b/sirepo/package_data/template/genesis/lib/io-maginfile.tesla.lat index ef7e7605bb..efd3353cb4 100644 --- a/sirepo/package_data/template/genesis/lib/io-maginfile.tesla.lat +++ b/sirepo/package_data/template/genesis/lib/io-maginfile.tesla.lat @@ -1,13 +1,209 @@ -! LOOP=50 -AW 2.8280E+00 4.5000E-02 100 -AW 0.0000E+00 4.5000E-02 12 -! ENDLOOP -QF 0.0000E+00 4.5000E-02 104 -! LOOP=25 -QF 1.5000E+01 4.5000E-02 4 -QF 0.0000E+00 4.5000E-02 108 -QF -1.5000E+01 4.5000E-02 4 -QF 0.0000E+00 4.5000E-02 108 -!ENDLOOP -QF 1.5000E+01 4.5000E-02 4 +? VERSION = 1 +? UNITLENGTH = 0.045 # meters +#------------ +# QF +QF 15.0 4.0 104.0 +QF 0 108.0 0.0 +QF -15.0 4.0 0.0 +QF 0 108.0 0.0 +QF 15.0 4.0 0.0 +QF 0 108.0 0.0 +QF -15.0 4.0 0.0 +QF 0 108.0 0.0 +QF 15.0 4.0 0.0 +QF 0 108.0 0.0 +QF -15.0 4.0 0.0 +QF 0 108.0 0.0 +QF 15.0 4.0 0.0 +QF 0 108.0 0.0 +QF -15.0 4.0 0.0 +QF 0 108.0 0.0 +QF 15.0 4.0 0.0 +QF 0 108.0 0.0 +QF -15.0 4.0 0.0 +QF 0 108.0 0.0 +QF 15.0 4.0 0.0 +QF 0 108.0 0.0 +QF -15.0 4.0 0.0 +QF 0 108.0 0.0 +QF 15.0 4.0 0.0 +QF 0 108.0 0.0 +QF -15.0 4.0 0.0 +QF 0 108.0 0.0 +QF 15.0 4.0 0.0 +QF 0 108.0 0.0 +QF -15.0 4.0 0.0 +QF 0 108.0 0.0 +QF 15.0 4.0 0.0 +QF 0 108.0 0.0 +QF -15.0 4.0 0.0 +QF 0 108.0 0.0 +QF 15.0 4.0 0.0 +QF 0 108.0 0.0 +QF -15.0 4.0 0.0 +QF 0 108.0 0.0 +QF 15.0 4.0 0.0 +QF 0 108.0 0.0 +QF -15.0 4.0 0.0 +QF 0 108.0 0.0 +QF 15.0 4.0 0.0 +QF 0 108.0 0.0 +QF -15.0 4.0 0.0 +QF 0 108.0 0.0 +QF 15.0 4.0 0.0 +QF 0 108.0 0.0 +QF -15.0 4.0 0.0 +QF 0 108.0 0.0 +QF 15.0 4.0 0.0 +QF 0 108.0 0.0 +QF -15.0 4.0 0.0 +QF 0 108.0 0.0 +QF 15.0 4.0 0.0 +QF 0 108.0 0.0 +QF -15.0 4.0 0.0 +QF 0 108.0 0.0 +QF 15.0 4.0 0.0 +QF 0 108.0 0.0 +QF -15.0 4.0 0.0 +QF 0 108.0 0.0 +QF 15.0 4.0 0.0 +QF 0 108.0 0.0 +QF -15.0 4.0 0.0 +QF 0 108.0 0.0 +QF 15.0 4.0 0.0 +QF 0 108.0 0.0 +QF -15.0 4.0 0.0 +QF 0 108.0 0.0 +QF 15.0 4.0 0.0 +QF 0 108.0 0.0 +QF -15.0 4.0 0.0 +QF 0 108.0 0.0 +QF 15.0 4.0 0.0 +QF 0 108.0 0.0 +QF -15.0 4.0 0.0 +QF 0 108.0 0.0 +QF 15.0 4.0 0.0 +QF 0 108.0 0.0 +QF -15.0 4.0 0.0 +QF 0 108.0 0.0 +QF 15.0 4.0 0.0 +QF 0 108.0 0.0 +QF -15.0 4.0 0.0 +QF 0 108.0 0.0 +QF 15.0 4.0 0.0 +QF 0 108.0 0.0 +QF -15.0 4.0 0.0 +QF 0 108.0 0.0 +QF 15.0 4.0 0.0 +QF 0 108.0 0.0 +QF -15.0 4.0 0.0 +QF 0 108.0 0.0 +QF 15.0 4.0 0.0 +QF 0 108.0 0.0 +QF -15.0 4.0 0.0 +QF 0 108.0 0.0 +QF 15.0 4.0 0.0 + +#------------ +# AW +AW 2.828 100.0 0.0 +AW 0 12.0 0.0 +AW 2.828 100.0 0.0 +AW 0 12.0 0.0 +AW 2.828 100.0 0.0 +AW 0 12.0 0.0 +AW 2.828 100.0 0.0 +AW 0 12.0 0.0 +AW 2.828 100.0 0.0 +AW 0 12.0 0.0 +AW 2.828 100.0 0.0 +AW 0 12.0 0.0 +AW 2.828 100.0 0.0 +AW 0 12.0 0.0 +AW 2.828 100.0 0.0 +AW 0 12.0 0.0 +AW 2.828 100.0 0.0 +AW 0 12.0 0.0 +AW 2.828 100.0 0.0 +AW 0 12.0 0.0 +AW 2.828 100.0 0.0 +AW 0 12.0 0.0 +AW 2.828 100.0 0.0 +AW 0 12.0 0.0 +AW 2.828 100.0 0.0 +AW 0 12.0 0.0 +AW 2.828 100.0 0.0 +AW 0 12.0 0.0 +AW 2.828 100.0 0.0 +AW 0 12.0 0.0 +AW 2.828 100.0 0.0 +AW 0 12.0 0.0 +AW 2.828 100.0 0.0 +AW 0 12.0 0.0 +AW 2.828 100.0 0.0 +AW 0 12.0 0.0 +AW 2.828 100.0 0.0 +AW 0 12.0 0.0 +AW 2.828 100.0 0.0 +AW 0 12.0 0.0 +AW 2.828 100.0 0.0 +AW 0 12.0 0.0 +AW 2.828 100.0 0.0 +AW 0 12.0 0.0 +AW 2.828 100.0 0.0 +AW 0 12.0 0.0 +AW 2.828 100.0 0.0 +AW 0 12.0 0.0 +AW 2.828 100.0 0.0 +AW 0 12.0 0.0 +AW 2.828 100.0 0.0 +AW 0 12.0 0.0 +AW 2.828 100.0 0.0 +AW 0 12.0 0.0 +AW 2.828 100.0 0.0 +AW 0 12.0 0.0 +AW 2.828 100.0 0.0 +AW 0 12.0 0.0 +AW 2.828 100.0 0.0 +AW 0 12.0 0.0 +AW 2.828 100.0 0.0 +AW 0 12.0 0.0 +AW 2.828 100.0 0.0 +AW 0 12.0 0.0 +AW 2.828 100.0 0.0 +AW 0 12.0 0.0 +AW 2.828 100.0 0.0 +AW 0 12.0 0.0 +AW 2.828 100.0 0.0 +AW 0 12.0 0.0 +AW 2.828 100.0 0.0 +AW 0 12.0 0.0 +AW 2.828 100.0 0.0 +AW 0 12.0 0.0 +AW 2.828 100.0 0.0 +AW 0 12.0 0.0 +AW 2.828 100.0 0.0 +AW 0 12.0 0.0 +AW 2.828 100.0 0.0 +AW 0 12.0 0.0 +AW 2.828 100.0 0.0 +AW 0 12.0 0.0 +AW 2.828 100.0 0.0 +AW 0 12.0 0.0 +AW 2.828 100.0 0.0 +AW 0 12.0 0.0 +AW 2.828 100.0 0.0 +AW 0 12.0 0.0 +AW 2.828 100.0 0.0 +AW 0 12.0 0.0 +AW 2.828 100.0 0.0 +AW 0 12.0 0.0 +AW 2.828 100.0 0.0 +AW 0 12.0 0.0 +AW 2.828 100.0 0.0 +AW 0 12.0 0.0 +AW 2.828 100.0 0.0 +AW 0 12.0 0.0 +AW 2.828 100.0 0.0 +AW 0 120.0 0.0 diff --git a/sirepo/package_data/template/genesis/lib/io-maginfile.ttf-ii.lat b/sirepo/package_data/template/genesis/lib/io-maginfile.ttf-ii.lat index 6a51d616ef..795b3dc367 100644 --- a/sirepo/package_data/template/genesis/lib/io-maginfile.ttf-ii.lat +++ b/sirepo/package_data/template/genesis/lib/io-maginfile.ttf-ii.lat @@ -1,14 +1,43 @@ -!LOOP=6 -AW 0.8960E+00 2.7300E-02 163 -AW 0.0000E+00 2.7300E-02 26 -!ENDLOOP -!LOOP=6 -QF 0.0000E+00 2.7300E-02 163 -QF 0.0000E+00 2.7300E-02 2 -QF 3.7000E+01 2.7300E-02 3 -QF 0.0000E+00 2.7300E-02 2 -QF 0.0000E+00 2.7300E-02 12 -QF 0.0000E+00 2.7300E-02 2 -QF -3.7000E+01 2.7300E-02 3 -QF 0.0000E+00 2.7300E-02 2 -!ENDLOOP +? VERSION = 1 +? UNITLENGTH = 0.0273 # meters + +#------------ +# AW +AW 0.896 163.0 0.0 +AW 0 26.0 0.0 +AW 0.896 163.0 0.0 +AW 0 26.0 0.0 +AW 0.896 163.0 0.0 +AW 0 26.0 0.0 +AW 0.896 163.0 0.0 +AW 0 26.0 0.0 +AW 0.896 163.0 0.0 +AW 0 26.0 0.0 +AW 0.896 163.0 0.0 +AW 0 24.0 0.0 + +#------------ +# QF +QF 37.0 3.0 165.0 +QF 0 16.0 0.0 +QF -37.0 3.0 0.0 +QF 0 167.0 0.0 +QF 37.0 3.0 0.0 +QF 0 16.0 0.0 +QF -37.0 3.0 0.0 +QF 0 167.0 0.0 +QF 37.0 3.0 0.0 +QF 0 16.0 0.0 +QF -37.0 3.0 0.0 +QF 0 167.0 0.0 +QF 37.0 3.0 0.0 +QF 0 16.0 0.0 +QF -37.0 3.0 0.0 +QF 0 167.0 0.0 +QF 37.0 3.0 0.0 +QF 0 16.0 0.0 +QF -37.0 3.0 0.0 +QF 0 167.0 0.0 +QF 37.0 3.0 0.0 +QF 0 16.0 0.0 +QF -37.0 3.0 0.0 diff --git a/sirepo/package_data/template/omega/parameters.py.jinja b/sirepo/package_data/template/omega/parameters.py.jinja index 22597c9476..64361f5b55 100644 --- a/sirepo/package_data/template/omega/parameters.py.jinja +++ b/sirepo/package_data/template/omega/parameters.py.jinja @@ -1,39 +1,33 @@ - from pykern import pkio -from pykern.pkcollections import PKDict +from rslume import Elegant, OPAL, Genesis2 +import h5py import os -import sirepo.pkcli.elegant -import sirepo.pkcli.genesis -import sirepo.pkcli.opal - -assert os.environ['SIREPO_SIMULATION_DB_LOGGED_IN_USER'], 'missing user id env var' - - -def run_sims(sim_list): - for idx in range(len(sim_list)): - _run( - pkio.py_path(f'run{idx + 1}'), - sim_list[idx].sim_type, - ) - - -def _run(run_dir, sim_type): - with pkio.save_chdir(run_dir): - if sim_type == 'opal': - sirepo.pkcli.opal.run_opal(with_mpi=True) - elif sim_type == 'elegant': - sirepo.pkcli.elegant.run_elegant() - elif sim_type == 'genesis': - sirepo.pkcli.genesis.run_genesis(run_dir) - else: - raise AssertionError(f"unhandled sim_type={sim_type}") - - -run_sims([ -{% for sim in simList %} - PKDict( - sim_type="{{sim.sim_type}}", - sim_id="{{sim.sim_id}}", - ), + +# patch numpy +import numpy +numpy.float = float + + +def code_kwargs(src_dir, input_file): + workdir = f"{src_dir}/out" + pkio.mkdir_parent(workdir) + return dict( + input_file=f"{src_dir}/{input_file}", + workdir=workdir, + verbose=True, + use_temp_dir=False, + ) + + +def run_and_save_particles(sim): + sim.run() + p = sim.final_particles() + if p: + with h5py.File(os.path.join(sim.workdir, "{{ particleOutfile }}"), "w") as f: + p.write(f) + + +{% for sim in simCall %} +{{ sim }} + {% endfor %} -]) diff --git a/sirepo/template/genesis.py b/sirepo/template/genesis.py index a39c713101..a05b7a3bae 100644 --- a/sirepo/template/genesis.py +++ b/sirepo/template/genesis.py @@ -425,15 +425,15 @@ def _parse_maginfile(filepath): u = 1 with pkio.open_text(filepath) as f: for line in f: + m = re.match(r"\?\s+(\w+)\s*=\s*([\S]+)", line) + if m: + if m.group(1) == "UNITLENGTH": + u = float(m.group(2)) + continue row = line.split() if row: if row[0] == _MAGIN_PLOT_FIELD: p.append(row[1]) - if row[0] == "?" and "UNITLENGTH" in row[1]: - if row[2] == "=": - u = row[3] - else: - u = row[2] if p: return PKDict(unit_length=u, points=p) raise AssertionError(f"No AW fields present in maginfile={filepath.basename}") diff --git a/sirepo/template/omega.py b/sirepo/template/omega.py index 000a43257a..433e962d53 100644 --- a/sirepo/template/omega.py +++ b/sirepo/template/omega.py @@ -13,10 +13,6 @@ from sirepo.template.lattice import LatticeUtil import h5py import numpy -import pmd_beamphysics -import pmd_beamphysics.interfaces.elegant -import pmd_beamphysics.interfaces.genesis -import pmd_beamphysics.interfaces.opal import re import sirepo.sim_data import sirepo.template @@ -108,12 +104,23 @@ Cx="run_setup.centroid.sdds", Cy="run_setup.centroid.sdds", ) +_PARTICLE_OUTFILE = "openpmd.h5" _RELATED_SIMS_FOLDER = "/Omega" _SUCCESS_OUTPUT_FILE = PKDict( elegant="run_setup.output.sdds", opal="opal.h5", genesis="genesis.out.dpa", ) +_SIM_TYPE_TO_CODE = PKDict( + elegant="Elegant", + opal="OPAL", + genesis="Genesis2", +) +_SIM_TYPE_TO_INPUT_FILE = PKDict( + elegant="elegant.ele", + opal="opal.in", + genesis="genesis.in", +) def background_percent_complete(report, run_dir, is_running): @@ -158,60 +165,17 @@ def copy_related_sims(data, qcall=None): def get_data_file(run_dir, model, frame, options): - def _particle_file_and_sim_info(): - i = int(re.search(r"Animation(\d+)\-", model).groups(1)[0]) - s = _sim_info( - simulation_db.read_json( - run_dir.join(template_common.INPUT_BASE_NAME) - ).models, - i - 1, - ) - return (pkio.py_path(f"run{i}").join(_SUCCESS_OUTPUT_FILE[s.sim_type]), s) - - def _particle_group(sim_type, particle_file): - if sim_type == "elegant": - return pmd_beamphysics.ParticleGroup( - data=pmd_beamphysics.interfaces.elegant.elegant_to_data(particle_file), - ) - elif sim_type == "opal": - step = _template_for_sim_type(sim_type).read_frame_count(run_dir) - with h5py.File(particle_file, "r") as f: - return pmd_beamphysics.ParticleGroup( - data=pmd_beamphysics.interfaces.opal.opal_to_data( - f[f"/Step#{step}"] - ), - ) - elif sim_type == "genesis": - dm = simulation_db.read_json( - particle_file.dirpath().join(template_common.INPUT_BASE_NAME + ".json") - ).models - v = numpy.fromfile( - particle_file.dirpath().join(particle_file.purebasename + ".dpa"), - dtype=numpy.float64, - ) - v = v.reshape( - int(len(v) / 6 / dm.electronBeam.npart), - 6, - dm.electronBeam.npart, - ) - return pmd_beamphysics.ParticleGroup( - data=pmd_beamphysics.interfaces.genesis.genesis2_dpa_to_data( - v, - xlamds=dm.radiation.xlamds, - current=numpy.array([dm.electronBeam.curpeak]), - ) - ) - else: - raise AssertionError(f"unsupported sim_type={sim_type}") - - f, s = _particle_file_and_sim_info() + i = int(re.search(r"Animation(\d+)\-", model).groups(1)[0]) + out = _sim_out_dir(run_dir, i) + s = _sim_info( + simulation_db.read_json(run_dir.join(template_common.INPUT_BASE_NAME)).models, + i - 1, + ) if options.suffix is None: - return f - if options.suffix != "openpmd": - raise AssertionError(f"unknown data type={options.suffix} requested") - n = f"{s.sim_type}_openpmd.h5" - _particle_group(s.sim_type, f).write(n) - return n + return out.join(_SUCCESS_OUTPUT_FILE[s.sim_type]) + if options.suffix == "openpmd": + return out.join(_PARTICLE_OUTFILE) + raise AssertionError(f"unknown data type={options.suffix} requested") def post_execution_processing(success_exit, run_dir, **kwargs): @@ -228,7 +192,7 @@ def post_execution_processing(success_exit, run_dir, **kwargs): s = _sim_info(dm, idx) if not s.sim_type or not s.sid: continue - sim_dir = _sim_dir(run_dir, idx + 1) + sim_dir = _sim_out_dir(run_dir, idx + 1) sim_template = _template_for_sim_type(s.sim_type) res = f"{s.sim_type.upper()} failed\n" if success_exit: @@ -255,10 +219,12 @@ def sim_frame(frame_args): sim_type = frame_args.sim_in.models.simWorkflow.coupledSims[ int(frame_args.simCount) - 1 ].simulationType - frame_args.run_dir = _sim_dir(frame_args.run_dir, frame_args.simCount) frame_args.sim_in = simulation_db.read_json( - frame_args.run_dir.join(template_common.INPUT_BASE_NAME) + _sim_in_dir(frame_args.run_dir, frame_args.simCount).join( + template_common.INPUT_BASE_NAME + ) ) + frame_args.run_dir = _sim_out_dir(frame_args.run_dir, frame_args.simCount) if "Phase" in frame_args.frameReport: return _plot_phase(sim_type, frame_args) if "Beam" in frame_args.frameReport: @@ -346,9 +312,24 @@ def _coupled_sims_list(data): def _generate_parameters_file(data): + def _sim_call(idx, sim): + r = ( + f"sim{idx + 1} = {_SIM_TYPE_TO_CODE[sim.sim_type]}(" + + f'**code_kwargs("run{idx + 1 }", "{_SIM_TYPE_TO_INPUT_FILE[sim.sim_type]}")' + ) + if sim.sim_type in ("opal", "elegant"): + r += ", update_filenames=True" + r += ")\n" + if idx > 0: + r += f"sim{idx + 1}.set_particle_input(sim{idx}.final_particles())\n" + r += f"run_and_save_particles(sim{idx + 1})" + return r + res, v = template_common.generate_parameters_file(data) - sim_list = _coupled_sims_list(data) - v.simList = sim_list + v.simCall = [] + for idx, s in enumerate(_coupled_sims_list(data)): + v.simCall.append(_sim_call(idx, s)) + v.particleOutfile = _PARTICLE_OUTFILE return res + template_common.render_jinja(SIM_TYPE, v) @@ -370,7 +351,7 @@ def _report_info(sim_count, model_name, report_count): res = [] idx = 0 - sim_dir = _sim_dir(run_dir, idx + 1) + sim_dir = _sim_out_dir(run_dir, idx + 1) while sim_dir.exists() and _has_file(sim_dir): r = [] res.append(r) @@ -392,7 +373,7 @@ def _report_info(sim_count, model_name, report_count): ] ) idx += 1 - sim_dir = _sim_dir(run_dir, idx + 1) + sim_dir = _sim_out_dir(run_dir, idx + 1) return res @@ -425,9 +406,6 @@ def _plot_beam(sim_type, frame_args): if sim_type == "elegant": return _extract_elegant_beam_plot(frame_args) if sim_type == "genesis": - frame_args.sim_in = simulation_db.read_json( - frame_args.run_dir.join(template_common.INPUT_BASE_NAME) - ) return _template_for_sim_type(sim_type).sim_frame_parameterAnimation(frame_args) raise AssertionError("unhandled sim_type for sim_frame(): {}".format(sim_type)) @@ -477,7 +455,11 @@ def _prepare_elegant(run_dir, elegant_id): if "report" in data: del data["report"] sirepo.simulation_db.prepare_simulation(data, run_dir) - return f"{run_dir}/run_setup.output.sdds" + t = pkio.read_text(run_dir.join(template_common.PARAMETERS_PYTHON_FILE)) + m = re.search(r'"""(.*?)""".*?"""(.*?)"""', t, re.MULTILINE | re.DOTALL) + assert m + pkio.write_text(run_dir.join("elegant.lte"), m.group(1)) + pkio.write_text(run_dir.join("elegant.ele"), m.group(2)) def _prepare_genesis(run_dir, genesis_id): @@ -486,7 +468,10 @@ def _prepare_genesis(run_dir, genesis_id): if "report" in data: del data["report"] sirepo.simulation_db.prepare_simulation(data, run_dir) - return f"{run_dir}/genesis.out.dpa" + t = pkio.read_text(run_dir.join(template_common.PARAMETERS_PYTHON_FILE)) + m = re.search(r'"""(.*?)"""', t, re.MULTILINE | re.DOTALL) + assert m + pkio.write_text(run_dir.join("genesis.in"), m.group(1)) def _prepare_opal(run_dir, opal_id): @@ -494,14 +479,13 @@ def _prepare_opal(run_dir, opal_id): data.computeModel = "animation" LatticeUtil.find_first_command(data, "option").psdumpfreq = 0 sirepo.simulation_db.prepare_simulation(data, run_dir) - return f"{run_dir}/{sirepo.template.opal._OPAL_H5_FILE}" def _prepare_subsims(data, run_dir): sim_list = _coupled_sims_list(data) for idx in range(len(sim_list)): s = sim_list[idx] - d = run_dir.join(f"run{idx + 1}") + d = _sim_in_dir(run_dir, idx + 1) pkio.unchecked_remove(d) pkio.mkdir_parent(d) if s.sim_type == "opal": @@ -518,10 +502,14 @@ def _read_sim(sim_type, sim_id): return sirepo.sim_data.get_class(sim_type).sim_db_read_sim(sim_id) -def _sim_dir(run_dir, sim_count): +def _sim_in_dir(run_dir, sim_count): return run_dir.join(f"run{sim_count}") +def _sim_out_dir(run_dir, sim_count): + return _sim_in_dir(run_dir, sim_count).join("out") + + def _sim_info(dm, idx): s = dm.simWorkflow.coupledSims if len(s) > idx: From f78e8155dfe9ddc35c84061485603c9449599d6c Mon Sep 17 00:00:00 2001 From: moellep Date: Mon, 15 Apr 2024 23:54:57 +0000 Subject: [PATCH 10/23] moved sim list load to controller to avoid multiple calls --- sirepo/package_data/static/js/omega.js | 29 ++++++++++++-------------- 1 file changed, 13 insertions(+), 16 deletions(-) diff --git a/sirepo/package_data/static/js/omega.js b/sirepo/package_data/static/js/omega.js index c35535d369..c160c70c7f 100644 --- a/sirepo/package_data/static/js/omega.js +++ b/sirepo/package_data/static/js/omega.js @@ -37,7 +37,7 @@ SIREPO.app.factory('omegaService', function(appState) { return self; }); -SIREPO.app.controller('SourceController', function (appState, frameCache, omegaService, persistentSimulation, $scope) { +SIREPO.app.controller('SourceController', function (appState, frameCache, omegaService, persistentSimulation, requestSender, $scope) { const self = this; let errorMessage; self.omegaService = omegaService; @@ -45,6 +45,16 @@ SIREPO.app.controller('SourceController', function (appState, frameCache, omegaS self.reports = null; self.simAnalysisModel = 'animation'; + function requestSimListByType(simType) { + requestSender.sendRequest( + 'listSimulations', + () => {}, + { + simulationType: simType, + } + ); + } + self.simHandleStatus = data => { errorMessage = data.error; frameCache.setFrameCount(data.frameCount || 0); @@ -73,6 +83,7 @@ SIREPO.app.controller('SourceController', function (appState, frameCache, omegaS } }; + SIREPO.APP_SCHEMA.relatedSimTypes.forEach(simType => requestSimListByType(simType)); self.simState = persistentSimulation.initSimulationState(self); //TODO(pjm): this should be default behavior in simStatusPanel self.simState.errorMessage = () => errorMessage; @@ -186,7 +197,7 @@ SIREPO.app.directive('beamAndPhasePlots', function(appState, omegaService) { }; }); -SIREPO.app.directive('dynamicSimList', function(appState, requestSender) { +SIREPO.app.directive('dynamicSimList', function(appState) { return { restrict: 'A', scope: { @@ -199,20 +210,6 @@ SIREPO.app.directive('dynamicSimList', function(appState, requestSender) { `, controller: function($scope) { - const requestSimListByType = (simType) => { - requestSender.sendRequest( - 'listSimulations', - () => {}, - { - simulationType: simType, - } - ); - }; - if (SIREPO.APP_SCHEMA.relatedSimTypes) { - SIREPO.APP_SCHEMA.relatedSimTypes.forEach(simType => { - requestSimListByType(simType); - }); - } $scope.selectedCode = () => { if ($scope.model) { $scope.code = $scope.model.simulationType; From dd63d1099fe8716d97a251d883bdc120828932e7 Mon Sep 17 00:00:00 2001 From: moellep Date: Mon, 15 Apr 2024 23:56:04 +0000 Subject: [PATCH 11/23] for the dpa genesis plot, assume only 1 frame - remove dependence on electronBeam.npart when running from omega --- sirepo/template/genesis.py | 23 +++++++++++++++-------- 1 file changed, 15 insertions(+), 8 deletions(-) diff --git a/sirepo/template/genesis.py b/sirepo/template/genesis.py index a05b7a3bae..3d72e99e9c 100644 --- a/sirepo/template/genesis.py +++ b/sirepo/template/genesis.py @@ -265,7 +265,7 @@ def sim_frame_finalFieldAnimation(frame_args): def sim_frame_finalParticleAnimation(frame_args): - return _particle_plot(frame_args, _FINAL_PARTICLE_OUTPUT_FILENAME) + return _particle_plot(frame_args, _FINAL_PARTICLE_OUTPUT_FILENAME, one_frame=True) def sim_frame_particleAnimation(frame_args): @@ -501,14 +501,21 @@ def _parse_namelist(data, text): return data -def _particle_plot(frame_args, filename): - n = frame_args.sim_in.models.electronBeam.npart +def _particle_plot(frame_args, filename, one_frame=False): d = numpy.fromfile(str(frame_args.run_dir.join(filename)), dtype=numpy.float64) - b = d.reshape( - int(len(d) / len(SCHEMA.enum.ParticleColumn) / n), - len(SCHEMA.enum.ParticleColumn), - n, - ) + if one_frame: + b = d.reshape( + 1, + len(SCHEMA.enum.ParticleColumn), + int(len(d) / len(SCHEMA.enum.ParticleColumn)), + ) + else: + n = frame_args.sim_in.models.electronBeam.npart + b = d.reshape( + int(len(d) / len(SCHEMA.enum.ParticleColumn) / n), + len(SCHEMA.enum.ParticleColumn), + n, + ) x = _get_col(frame_args.x) y = _get_col(frame_args.y) return template_common.heatmap( From 8c2646fbd4358aec71bcf655aaa1dd0d03ad3cd6 Mon Sep 17 00:00:00 2001 From: gurhar1133 Date: Thu, 18 Apr 2024 22:49:04 +0000 Subject: [PATCH 12/23] Fix #7014: plotting ckp --- sirepo/template/omega.py | 102 +++++++++++++++++++++++++++++++-------- 1 file changed, 81 insertions(+), 21 deletions(-) diff --git a/sirepo/template/omega.py b/sirepo/template/omega.py index 433e962d53..e15273e686 100644 --- a/sirepo/template/omega.py +++ b/sirepo/template/omega.py @@ -420,32 +420,92 @@ def _plot_field_dist(sim_type, frame_args): def _plot_phase(sim_type, frame_args): _phase_plot_args(sim_type, frame_args) + out_data = template_common.read_dict_from_h5(frame_args.run_dir.join("openpmd.h5")) if sim_type == "opal": - r = _template_for_sim_type(sim_type).bunch_plot( - frame_args, - frame_args.run_dir, - frame_args.frameIndex, + col_map = PKDict( + x=out_data.position.x, + y=out_data.position.y, + z=out_data.position.z, + px=out_data.momentum.x, + py=out_data.momentum.y, + pz=out_data.momentum.z, ) - return r.pkupdate( - title=_PLOT_TITLE[sim_type][frame_args.x + "-" + frame_args.y], - y_label=_PLOT_Y_LABEL[sim_type].get( - frame_args.x + "-" + frame_args.y, r.y_label - ), + title_map = PKDict( + xy="Cross-Section", + xpx="Horizontal", + ypy="Vertical", + zpz="Longitudinal", ) - if sim_type == "elegant": - return _template_for_sim_type(sim_type).extract_report_data( - str(frame_args.run_dir.join(_SUCCESS_OUTPUT_FILE[sim_type])), - frame_args, + elif sim_type == "elegant": + col_map = PKDict( + x=out_data.position.x, + y=out_data.position.y, + t=out_data.time, + xp=out_data.momentum.x, + yp=out_data.momentum.y, + p=out_data.momentum.z, ) - if sim_type == "genesis": - frame_args.frameIndex = 0 - return ( - _template_for_sim_type(sim_type) - .sim_frame_finalParticleAnimation(frame_args) - .pkupdate( - title=_PLOT_TITLE[sim_type][frame_args.x + "-" + frame_args.y], - ) + title_map = PKDict( + xy="Cross-Section", + xxp="Horizontal", + yyp="Vertical", + tp="Longitudinal", + ) + else: + col_map = PKDict( + x=out_data.position.x, + y=out_data.position.y, + psi=out_data.time, + pxmc=out_data.momentum.x, + pymc=out_data.momentum.y, + gamma=out_data.momentum.z, + ) + title_map = PKDict( + xy="Cross-Section", + xpxmc="Horizontal", + ypymc="Vertical", + psigamma="Longitudinal", + ) + if sim_type in ("elegant", "opal", "genesis"): + y_col = numpy.array(col_map[frame_args.y]) + y_col[numpy.isnan(y_col)] = .0 + v = [col_map[frame_args.x], y_col] + return template_common.heatmap( + values=v, + model=frame_args, + plot_fields=PKDict( + x_label=frame_args.x, + y_label=frame_args.y, + title=title_map[frame_args.x + frame_args.y], + ), ) + # r = _template_for_sim_type(sim_type).bunch_plot( + # frame_args, + # frame_args.run_dir, + # frame_args.frameIndex, + # ) + # return r.pkupdate( + # title=_PLOT_TITLE[sim_type][frame_args.x + "-" + frame_args.y], + # y_label=_PLOT_Y_LABEL[sim_type].get( + # frame_args.x + "-" + frame_args.y, r.y_label + # ), + # ) + # if sim_type == "elegant": + # pkdp("\n\n\nelegant cols={}", frame_args.x + " " + frame_args.y) + # return _template_for_sim_type(sim_type).extract_report_data( + # str(frame_args.run_dir.join(_SUCCESS_OUTPUT_FILE[sim_type])), + # frame_args, + # ) + # if sim_type == "genesis": + # pkdp("\n\n\ngenesis cols={}", frame_args.x + " " + frame_args.y) + # frame_args.frameIndex = 0 + # return ( + # _template_for_sim_type(sim_type) + # .sim_frame_finalParticleAnimation(frame_args) + # .pkupdate( + # title=_PLOT_TITLE[sim_type][frame_args.x + "-" + frame_args.y], + # ) + # ) raise AssertionError("unhandled sim_type for sim_frame(): {}".format(sim_type)) From 6c3023097edf862e8dcb317b3a0906a880b7e527 Mon Sep 17 00:00:00 2001 From: gurhar1133 Date: Fri, 19 Apr 2024 16:56:50 +0000 Subject: [PATCH 13/23] Fix #7014: ckp --- sirepo/template/omega.py | 133 ++++++++++++++++----------------------- 1 file changed, 54 insertions(+), 79 deletions(-) diff --git a/sirepo/template/omega.py b/sirepo/template/omega.py index e15273e686..ddab7dc471 100644 --- a/sirepo/template/omega.py +++ b/sirepo/template/omega.py @@ -418,94 +418,69 @@ def _plot_field_dist(sim_type, frame_args): def _plot_phase(sim_type, frame_args): - _phase_plot_args(sim_type, frame_args) - out_data = template_common.read_dict_from_h5(frame_args.run_dir.join("openpmd.h5")) - if sim_type == "opal": - col_map = PKDict( - x=out_data.position.x, - y=out_data.position.y, - z=out_data.position.z, - px=out_data.momentum.x, - py=out_data.momentum.y, - pz=out_data.momentum.z, - ) - title_map = PKDict( - xy="Cross-Section", - xpx="Horizontal", - ypy="Vertical", - zpz="Longitudinal", - ) - elif sim_type == "elegant": - col_map = PKDict( - x=out_data.position.x, - y=out_data.position.y, - t=out_data.time, - xp=out_data.momentum.x, - yp=out_data.momentum.y, - p=out_data.momentum.z, - ) - title_map = PKDict( - xy="Cross-Section", - xxp="Horizontal", - yyp="Vertical", - tp="Longitudinal", - ) - else: - col_map = PKDict( - x=out_data.position.x, - y=out_data.position.y, - psi=out_data.time, - pxmc=out_data.momentum.x, - pymc=out_data.momentum.y, - gamma=out_data.momentum.z, - ) - title_map = PKDict( - xy="Cross-Section", - xpxmc="Horizontal", - ypymc="Vertical", - psigamma="Longitudinal", - ) + def _clean(column): + column = numpy.array(column) + column[numpy.isnan(column)] = 0.0 + return column + + def _col(column_name, data): + if column_name in ("x", "y"): + return data.position[column_name] + return PKDict( + elegant=PKDict( + t=data.time, + xp=data.momentum.x, + yp=data.momentum.y, + p=data.momentum.z, + ), + opal=PKDict( + z=data.position.z, + px=data.momentum.x, + py=data.momentum.y, + pz=data.momentum.z, + ), + genesis=PKDict( + psi=data.time, + pxmc=data.momentum.x, + pymc=data.momentum.y, + gamma=data.momentum.z, + ), + )[sim_type][column_name] + + def _title(x, y): + if x + y == "xy": + return "Cross-Section" + return PKDict( + elegant=PKDict( + xxp="Horizontal", + yyp="Vertical", + tp="Longitudinal", + ), + opal=PKDict( + xpx="Horizontal", + ypy="Vertical", + zpz="Longitudinal", + ), + genesis=PKDict( + xpxmc="Horizontal", + ypymc="Vertical", + psigamma="Longitudinal", + ), + )[sim_type][x + y] + + d = template_common.read_dict_from_h5(frame_args.run_dir.join("openpmd.h5")) + _phase_plot_args(sim_type, frame_args) if sim_type in ("elegant", "opal", "genesis"): - y_col = numpy.array(col_map[frame_args.y]) - y_col[numpy.isnan(y_col)] = .0 - v = [col_map[frame_args.x], y_col] return template_common.heatmap( - values=v, + values=[_clean(_col(frame_args.x, d)), _clean(_col(frame_args.y, d))], model=frame_args, plot_fields=PKDict( x_label=frame_args.x, y_label=frame_args.y, - title=title_map[frame_args.x + frame_args.y], + title=_title(frame_args.x, frame_args.y), ), ) - # r = _template_for_sim_type(sim_type).bunch_plot( - # frame_args, - # frame_args.run_dir, - # frame_args.frameIndex, - # ) - # return r.pkupdate( - # title=_PLOT_TITLE[sim_type][frame_args.x + "-" + frame_args.y], - # y_label=_PLOT_Y_LABEL[sim_type].get( - # frame_args.x + "-" + frame_args.y, r.y_label - # ), - # ) - # if sim_type == "elegant": - # pkdp("\n\n\nelegant cols={}", frame_args.x + " " + frame_args.y) - # return _template_for_sim_type(sim_type).extract_report_data( - # str(frame_args.run_dir.join(_SUCCESS_OUTPUT_FILE[sim_type])), - # frame_args, - # ) - # if sim_type == "genesis": - # pkdp("\n\n\ngenesis cols={}", frame_args.x + " " + frame_args.y) - # frame_args.frameIndex = 0 - # return ( - # _template_for_sim_type(sim_type) - # .sim_frame_finalParticleAnimation(frame_args) - # .pkupdate( - # title=_PLOT_TITLE[sim_type][frame_args.x + "-" + frame_args.y], - # ) - # ) raise AssertionError("unhandled sim_type for sim_frame(): {}".format(sim_type)) From c0ed930d1d9f8ae503beae652b301e8659e1f3ef Mon Sep 17 00:00:00 2001 From: gurhar1133 Date: Fri, 19 Apr 2024 17:57:59 +0000 Subject: [PATCH 14/23] Fix #7014: ckp --- sirepo/template/omega.py | 126 +++++++++++++++++++++++---------------- 1 file changed, 75 insertions(+), 51 deletions(-) diff --git a/sirepo/template/omega.py b/sirepo/template/omega.py index ddab7dc471..df61633d4f 100644 --- a/sirepo/template/omega.py +++ b/sirepo/template/omega.py @@ -17,6 +17,29 @@ import sirepo.sim_data import sirepo.template + +_AXES_LABELS = PKDict( + Cross=PKDict( + x_label="x", + y_label="y", + title="Cross-Section", + ), + Horizontal=PKDict( + x_label="x", + y_label="px", + title="Horizontal", + ), + Vertical=PKDict( + x_label="y", + y_label="py", + title="Vertical", + ), + Longitudinal=PKDict( + x_label="z", + y_label="pz", + title="Longitudinal", + ), +) _SIM_DATA, SIM_TYPE, SCHEMA = sirepo.sim_data.template_globals() _PHASE_PLOT_COUNT = 4 _PHASE_PLOTS = PKDict( @@ -418,67 +441,68 @@ def _plot_field_dist(sim_type, frame_args): def _plot_phase(sim_type, frame_args): - - def _clean(column): - column = numpy.array(column) - column[numpy.isnan(column)] = 0.0 - return column - def _col(column_name, data): if column_name in ("x", "y"): - return data.position[column_name] - return PKDict( - elegant=PKDict( - t=data.time, - xp=data.momentum.x, - yp=data.momentum.y, - p=data.momentum.z, - ), - opal=PKDict( - z=data.position.z, - px=data.momentum.x, - py=data.momentum.y, - pz=data.momentum.z, - ), - genesis=PKDict( - psi=data.time, - pxmc=data.momentum.x, - pymc=data.momentum.y, - gamma=data.momentum.z, - ), - )[sim_type][column_name] - - def _title(x, y): + c = data.position[column_name] + else: + c = PKDict( + elegant=PKDict( + t=data.time, + xp=data.momentum.x, + yp=data.momentum.y, + p=data.momentum.z, + ), + opal=PKDict( + z=data.position.z, + px=data.momentum.x, + py=data.momentum.y, + pz=data.momentum.z, + ), + genesis=PKDict( + psi=data.time, + pxmc=data.momentum.x, + pymc=data.momentum.y, + gamma=data.momentum.z, + ), + )[sim_type][column_name] + c = numpy.array(c) + c[numpy.isnan(c)] = 0.0 + return c + + def _title_and_label(x, y): if x + y == "xy": - return "Cross-Section" - return PKDict( - elegant=PKDict( - xxp="Horizontal", - yyp="Vertical", - tp="Longitudinal", - ), - opal=PKDict( - xpx="Horizontal", - ypy="Vertical", - zpz="Longitudinal", - ), - genesis=PKDict( - xpxmc="Horizontal", - ypymc="Vertical", - psigamma="Longitudinal", - ), - )[sim_type][x + y] + return _AXES_LABELS["Cross"] + return _AXES_LABELS[ + PKDict( + elegant=PKDict( + xxp="Horizontal", + yyp="Vertical", + tp="Longitudinal", + ), + opal=PKDict( + xpx="Horizontal", + ypy="Vertical", + zpz="Longitudinal", + ), + genesis=PKDict( + xpxmc="Horizontal", + ypymc="Vertical", + psigamma="Longitudinal", + ), + )[sim_type][x + y] + ] d = template_common.read_dict_from_h5(frame_args.run_dir.join("openpmd.h5")) _phase_plot_args(sim_type, frame_args) if sim_type in ("elegant", "opal", "genesis"): + t = _title_and_label(frame_args.x, frame_args.y) return template_common.heatmap( - values=[_clean(_col(frame_args.x, d)), _clean(_col(frame_args.y, d))], + values=[_col(frame_args.x, d), _col(frame_args.y, d)], model=frame_args, plot_fields=PKDict( - x_label=frame_args.x, - y_label=frame_args.y, - title=_title(frame_args.x, frame_args.y), + x_label=t.x_label, + y_label=t.y_label, + title=t.title, ), ) raise AssertionError("unhandled sim_type for sim_frame(): {}".format(sim_type)) From b0d06aa131e25100218d61fc731826a805b5a1d0 Mon Sep 17 00:00:00 2001 From: gurhar1133 Date: Fri, 19 Apr 2024 20:11:15 +0000 Subject: [PATCH 15/23] Fix #7014: cleanup --- sirepo/template/omega.py | 143 ++++++--------------------------------- 1 file changed, 22 insertions(+), 121 deletions(-) diff --git a/sirepo/template/omega.py b/sirepo/template/omega.py index df61633d4f..a180297b24 100644 --- a/sirepo/template/omega.py +++ b/sirepo/template/omega.py @@ -18,78 +18,18 @@ import sirepo.template -_AXES_LABELS = PKDict( - Cross=PKDict( - x_label="x", - y_label="y", - title="Cross-Section", - ), - Horizontal=PKDict( - x_label="x", - y_label="px", - title="Horizontal", - ), - Vertical=PKDict( - x_label="y", - y_label="py", - title="Vertical", - ), - Longitudinal=PKDict( - x_label="z", - y_label="pz", - title="Longitudinal", - ), -) _SIM_DATA, SIM_TYPE, SCHEMA = sirepo.sim_data.template_globals() _PHASE_PLOT_COUNT = 4 -_PHASE_PLOTS = PKDict( - genesis=[ - ["x", "pxmc"], - ["y", "pymc"], - ["x", "y"], - ["psi", "gamma"], - ], - opal=[ - ["x", "px"], - ["y", "py"], - ["x", "y"], - ["z", "pz"], - ], - elegant=[ - ["x", "xp"], - ["y", "yp"], - ["x", "y"], - ["t", "p"], - ], -) +_PHASE_PLOTS = [["x", "px"], ["y", "py"], ["x", "y"], ["z", "pz"]] _PLOT_TITLE = PKDict( - opal=PKDict( - { - "x-px": "Horizontal", - "y-py": "Vertical", - "x-y": "Cross-section", - "z-pz": "Longitudinal", - }, - ), - genesis=PKDict( - { - "x-pxmc": "Horizontal", - "y-pymc": "Vertical", - "x-y": "Cross-section", - "psi-gamma": "PSI/Gamma", - }, - ), -) -_PLOT_Y_LABEL = PKDict( - opal=PKDict( - { - # TODO(pjm): should format px and βx with subscripts - "x-px": "px (βx γ)", - "y-py": "py (βy γ)", - "z-pz": "pz (β γ)", - } - ) + { + "x-px": "Horizontal", + "y-py": "Vertical", + "x-y": "Cross-section", + "z-pz": "Longitudinal", + } ) + _BEAM_PARAMETERS = PKDict( genesis=PKDict( rmsx="xrms", @@ -414,7 +354,7 @@ def _is_genesis(run_dir, index): def _phase_plot_args(sim_type, frame_args): - xy = _PHASE_PLOTS[sim_type][int(frame_args.reportCount) - 1] + xy = _PHASE_PLOTS[int(frame_args.reportCount) - 1] del frame_args["y1"] frame_args.x = xy[0] frame_args.y = xy[1] @@ -441,68 +381,29 @@ def _plot_field_dist(sim_type, frame_args): def _plot_phase(sim_type, frame_args): - def _col(column_name, data): - if column_name in ("x", "y"): - c = data.position[column_name] - else: - c = PKDict( - elegant=PKDict( - t=data.time, - xp=data.momentum.x, - yp=data.momentum.y, - p=data.momentum.z, - ), - opal=PKDict( - z=data.position.z, - px=data.momentum.x, - py=data.momentum.y, - pz=data.momentum.z, - ), - genesis=PKDict( - psi=data.time, - pxmc=data.momentum.x, - pymc=data.momentum.y, - gamma=data.momentum.z, - ), - )[sim_type][column_name] + def _col(data, column_name): + c = PKDict( + x=data.position.x, + y=data.position.y, + z=data.position.z if sim_type == "opal" else data.time, + px=data.momentum.x, + py=data.momentum.y, + pz=data.momentum.z, + )[column_name] c = numpy.array(c) c[numpy.isnan(c)] = 0.0 return c - def _title_and_label(x, y): - if x + y == "xy": - return _AXES_LABELS["Cross"] - return _AXES_LABELS[ - PKDict( - elegant=PKDict( - xxp="Horizontal", - yyp="Vertical", - tp="Longitudinal", - ), - opal=PKDict( - xpx="Horizontal", - ypy="Vertical", - zpz="Longitudinal", - ), - genesis=PKDict( - xpxmc="Horizontal", - ypymc="Vertical", - psigamma="Longitudinal", - ), - )[sim_type][x + y] - ] - d = template_common.read_dict_from_h5(frame_args.run_dir.join("openpmd.h5")) _phase_plot_args(sim_type, frame_args) if sim_type in ("elegant", "opal", "genesis"): - t = _title_and_label(frame_args.x, frame_args.y) return template_common.heatmap( - values=[_col(frame_args.x, d), _col(frame_args.y, d)], + values=[_col(d, frame_args.x), _col(d, frame_args.y)], model=frame_args, plot_fields=PKDict( - x_label=t.x_label, - y_label=t.y_label, - title=t.title, + x_label=frame_args.x, + y_label=frame_args.y, + title=_PLOT_TITLE[frame_args.x + "-" + frame_args.y], ), ) raise AssertionError("unhandled sim_type for sim_frame(): {}".format(sim_type)) From 3813f92d95cd81c1a5c98e687d4e5a5a45301664 Mon Sep 17 00:00:00 2001 From: gurhar1133 Date: Mon, 22 Apr 2024 21:54:04 +0000 Subject: [PATCH 16/23] Fix #7021: ckp --- .../template/omega/parameters.py.jinja | 2 ++ sirepo/template/omega.py | 24 ++++++++++++++----- 2 files changed, 20 insertions(+), 6 deletions(-) diff --git a/sirepo/package_data/template/omega/parameters.py.jinja b/sirepo/package_data/template/omega/parameters.py.jinja index 64361f5b55..5cfce14840 100644 --- a/sirepo/package_data/template/omega/parameters.py.jinja +++ b/sirepo/package_data/template/omega/parameters.py.jinja @@ -1,3 +1,4 @@ +from pmd_beamphysics import pmd_init from pykern import pkio from rslume import Elegant, OPAL, Genesis2 import h5py @@ -24,6 +25,7 @@ def run_and_save_particles(sim): p = sim.final_particles() if p: with h5py.File(os.path.join(sim.workdir, "{{ particleOutfile }}"), "w") as f: + pmd_init(f, basePath='/', particlesPath='/' ) p.write(f) diff --git a/sirepo/template/omega.py b/sirepo/template/omega.py index a180297b24..6b2a1d3d32 100644 --- a/sirepo/template/omega.py +++ b/sirepo/template/omega.py @@ -382,13 +382,25 @@ def _plot_field_dist(sim_type, frame_args): def _plot_phase(sim_type, frame_args): def _col(data, column_name): + from pmd_beamphysics import ParticleGroup + + # pkdp("\n\n\n h5path={}", frame_args.run_dir.join("openpmd.h5")) + p = ParticleGroup(h5=str(frame_args.run_dir.join("openpmd.h5"))) + # pkdp("\n\n\n x={}, y={}, z={}\n\n\n", p.x, p.y, p.z) + if sim_type != "opal": + pkdp("\n\n\n\n BEFORE p.z={}", p.z) + t0 = p.avg('t') + p.drift_to_t(t0) + # p.drift_to_t() + pkdp("\n\n\n\n AFTER p.z={}", p.z) c = PKDict( - x=data.position.x, - y=data.position.y, - z=data.position.z if sim_type == "opal" else data.time, - px=data.momentum.x, - py=data.momentum.y, - pz=data.momentum.z, + x=p.x, + y=p.y, + # z=data.position.z if sim_type == "opal" else data.time, + z=p.z, + px=p.px, + py=p.py, + pz=p.pz, )[column_name] c = numpy.array(c) c[numpy.isnan(c)] = 0.0 From af15a6b84771f3db54b2a41bcde57df07fd25cdf Mon Sep 17 00:00:00 2001 From: gurhar1133 Date: Mon, 22 Apr 2024 22:52:49 +0000 Subject: [PATCH 17/23] Fix #7021: note --- sirepo/template/omega.py | 1 + 1 file changed, 1 insertion(+) diff --git a/sirepo/template/omega.py b/sirepo/template/omega.py index 6b2a1d3d32..f36aa631e7 100644 --- a/sirepo/template/omega.py +++ b/sirepo/template/omega.py @@ -389,6 +389,7 @@ def _col(data, column_name): # pkdp("\n\n\n x={}, y={}, z={}\n\n\n", p.x, p.y, p.z) if sim_type != "opal": pkdp("\n\n\n\n BEFORE p.z={}", p.z) + # TODO (gurhar1133): not working with genesis t0 = p.avg('t') p.drift_to_t(t0) # p.drift_to_t() From 96bb19258569d005a9d87bf38e8ef8c54f52d1fb Mon Sep 17 00:00:00 2001 From: gurhar1133 Date: Tue, 23 Apr 2024 17:36:17 +0000 Subject: [PATCH 18/23] Fix #7021: ckp --- sirepo/template/omega.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/sirepo/template/omega.py b/sirepo/template/omega.py index f36aa631e7..32fcf85d09 100644 --- a/sirepo/template/omega.py +++ b/sirepo/template/omega.py @@ -387,13 +387,17 @@ def _col(data, column_name): # pkdp("\n\n\n h5path={}", frame_args.run_dir.join("openpmd.h5")) p = ParticleGroup(h5=str(frame_args.run_dir.join("openpmd.h5"))) # pkdp("\n\n\n x={}, y={}, z={}\n\n\n", p.x, p.y, p.z) - if sim_type != "opal": + if sim_type == "elegant": pkdp("\n\n\n\n BEFORE p.z={}", p.z) # TODO (gurhar1133): not working with genesis t0 = p.avg('t') p.drift_to_t(t0) # p.drift_to_t() pkdp("\n\n\n\n AFTER p.z={}", p.z) + if sim_type == "genesis": + pkdp("\n\n\n genesis weight = {}", p.weight) + # z0 = p.avg('z') + p.drift_to_t(0.1) c = PKDict( x=p.x, y=p.y, From 1c8f94d5447cd508a9ca4a5fa214dbb902d6a835 Mon Sep 17 00:00:00 2001 From: gurhar1133 Date: Tue, 23 Apr 2024 19:38:25 +0000 Subject: [PATCH 19/23] Fix #7021: ckp --- sirepo/template/omega.py | 32 +++++++++++--------------------- 1 file changed, 11 insertions(+), 21 deletions(-) diff --git a/sirepo/template/omega.py b/sirepo/template/omega.py index 32fcf85d09..eb561ca64e 100644 --- a/sirepo/template/omega.py +++ b/sirepo/template/omega.py @@ -3,7 +3,6 @@ :copyright: Copyright (c) 2023 RadiaSoft LLC. All Rights Reserved. :license: http://www.apache.org/licenses/LICENSE-2.0.html """ - from pykern import pkio from pykern import pkjson from pykern.pkcollections import PKDict @@ -14,6 +13,7 @@ import h5py import numpy import re +import pmd_beamphysics import sirepo.sim_data import sirepo.template @@ -381,27 +381,18 @@ def _plot_field_dist(sim_type, frame_args): def _plot_phase(sim_type, frame_args): - def _col(data, column_name): - from pmd_beamphysics import ParticleGroup - - # pkdp("\n\n\n h5path={}", frame_args.run_dir.join("openpmd.h5")) - p = ParticleGroup(h5=str(frame_args.run_dir.join("openpmd.h5"))) - # pkdp("\n\n\n x={}, y={}, z={}\n\n\n", p.x, p.y, p.z) - if sim_type == "elegant": - pkdp("\n\n\n\n BEFORE p.z={}", p.z) - # TODO (gurhar1133): not working with genesis - t0 = p.avg('t') - p.drift_to_t(t0) - # p.drift_to_t() - pkdp("\n\n\n\n AFTER p.z={}", p.z) - if sim_type == "genesis": - pkdp("\n\n\n genesis weight = {}", p.weight) - # z0 = p.avg('z') - p.drift_to_t(0.1) + def _col(column_name): + p = pmd_beamphysics.ParticleGroup(h5=str(frame_args.run_dir.join("openpmd.h5"))) + if sim_type != "opal": + try: + t0 = p.avg('t') + p.drift_to_t(t0) + except Exception as e: + # t0 = 0 + pass c = PKDict( x=p.x, y=p.y, - # z=data.position.z if sim_type == "opal" else data.time, z=p.z, px=p.px, py=p.py, @@ -411,11 +402,10 @@ def _col(data, column_name): c[numpy.isnan(c)] = 0.0 return c - d = template_common.read_dict_from_h5(frame_args.run_dir.join("openpmd.h5")) _phase_plot_args(sim_type, frame_args) if sim_type in ("elegant", "opal", "genesis"): return template_common.heatmap( - values=[_col(d, frame_args.x), _col(d, frame_args.y)], + values=[_col(frame_args.x), _col(frame_args.y)], model=frame_args, plot_fields=PKDict( x_label=frame_args.x, From 67679e546a9243aa04782ea984e83af9c179a670 Mon Sep 17 00:00:00 2001 From: gurhar1133 Date: Tue, 23 Apr 2024 21:31:06 +0000 Subject: [PATCH 20/23] Fix #7021: just using particleGroup --- sirepo/template/omega.py | 17 ++++++++--------- 1 file changed, 8 insertions(+), 9 deletions(-) diff --git a/sirepo/template/omega.py b/sirepo/template/omega.py index eb561ca64e..6f9eb7f4fa 100644 --- a/sirepo/template/omega.py +++ b/sirepo/template/omega.py @@ -3,6 +3,7 @@ :copyright: Copyright (c) 2023 RadiaSoft LLC. All Rights Reserved. :license: http://www.apache.org/licenses/LICENSE-2.0.html """ + from pykern import pkio from pykern import pkjson from pykern.pkcollections import PKDict @@ -383,17 +384,10 @@ def _plot_field_dist(sim_type, frame_args): def _plot_phase(sim_type, frame_args): def _col(column_name): p = pmd_beamphysics.ParticleGroup(h5=str(frame_args.run_dir.join("openpmd.h5"))) - if sim_type != "opal": - try: - t0 = p.avg('t') - p.drift_to_t(t0) - except Exception as e: - # t0 = 0 - pass c = PKDict( x=p.x, y=p.y, - z=p.z, + z=p.z if sim_type == "opal" else p.t, px=p.px, py=p.py, pz=p.pz, @@ -402,13 +396,18 @@ def _col(column_name): c[numpy.isnan(c)] = 0.0 return c + def _x_label(x_name): + if x_name == "z" and sim_type != "opal": + return "t" + return x_name + _phase_plot_args(sim_type, frame_args) if sim_type in ("elegant", "opal", "genesis"): return template_common.heatmap( values=[_col(frame_args.x), _col(frame_args.y)], model=frame_args, plot_fields=PKDict( - x_label=frame_args.x, + x_label=_x_label(frame_args.x), y_label=frame_args.y, title=_PLOT_TITLE[frame_args.x + "-" + frame_args.y], ), From d862e0154e28fddcb2cdb11a49c79089f7079d1f Mon Sep 17 00:00:00 2001 From: Paul Moeller Date: Mon, 29 Apr 2024 20:23:48 +0000 Subject: [PATCH 21/23] for #6653 update omega_test, move relate sim check to server - also pip install rslume prior to running --- sirepo/package_data/static/js/omega.js | 13 +------ sirepo/pkcli/omega.py | 6 ++++ sirepo/template/omega.py | 7 ++++ tests/template/omega_data/genesis.txt | 36 +++++++++++++++++++ tests/template/omega_data/opal_elegant.txt | 40 ++++++++++++++++++++++ tests/template/omega_test.py | 14 ++++---- 6 files changed, 98 insertions(+), 18 deletions(-) create mode 100644 tests/template/omega_data/genesis.txt create mode 100644 tests/template/omega_data/opal_elegant.txt diff --git a/sirepo/package_data/static/js/omega.js b/sirepo/package_data/static/js/omega.js index c160c70c7f..786c7778d9 100644 --- a/sirepo/package_data/static/js/omega.js +++ b/sirepo/package_data/static/js/omega.js @@ -37,7 +37,7 @@ SIREPO.app.factory('omegaService', function(appState) { return self; }); -SIREPO.app.controller('SourceController', function (appState, frameCache, omegaService, persistentSimulation, requestSender, $scope) { +SIREPO.app.controller('SourceController', function (appState, frameCache, omegaService, persistentSimulation, $scope) { const self = this; let errorMessage; self.omegaService = omegaService; @@ -45,16 +45,6 @@ SIREPO.app.controller('SourceController', function (appState, frameCache, omegaS self.reports = null; self.simAnalysisModel = 'animation'; - function requestSimListByType(simType) { - requestSender.sendRequest( - 'listSimulations', - () => {}, - { - simulationType: simType, - } - ); - } - self.simHandleStatus = data => { errorMessage = data.error; frameCache.setFrameCount(data.frameCount || 0); @@ -83,7 +73,6 @@ SIREPO.app.controller('SourceController', function (appState, frameCache, omegaS } }; - SIREPO.APP_SCHEMA.relatedSimTypes.forEach(simType => requestSimListByType(simType)); self.simState = persistentSimulation.initSimulationState(self); //TODO(pjm): this should be default behavior in simStatusPanel self.simState.errorMessage = () => errorMessage; diff --git a/sirepo/pkcli/omega.py b/sirepo/pkcli/omega.py index 291dbcb418..8d03e856dc 100644 --- a/sirepo/pkcli/omega.py +++ b/sirepo/pkcli/omega.py @@ -6,7 +6,13 @@ """ from pykern.pkdebug import pkdp, pkdc, pkdlog from sirepo.template import template_common +import os def run_background(cfg_dir): + # TODO(pjm): work-around until rslume is bundled with sirepo + try: + import rslume.elegant + except ModuleNotFoundError as e: + os.system("pip install git+https://github.com/radiasoft/rslume") template_common.exec_parameters() diff --git a/sirepo/template/omega.py b/sirepo/template/omega.py index 6f9eb7f4fa..12e0212ff1 100644 --- a/sirepo/template/omega.py +++ b/sirepo/template/omega.py @@ -175,6 +175,13 @@ def post_execution_processing(success_exit, run_dir, **kwargs): return "An unknown error occurred" +def prepare_for_client(data, qcall, **kwargs): + # ensure the related codes are present + for s in SCHEMA.relatedSimTypes: + simulation_db.simulation_dir(s, qcall=qcall) + return data + + def python_source_for_model(data, model, qcall, **kwargs): return _generate_parameters_file(data) diff --git a/tests/template/omega_data/genesis.txt b/tests/template/omega_data/genesis.txt new file mode 100644 index 0000000000..e4334fcccc --- /dev/null +++ b/tests/template/omega_data/genesis.txt @@ -0,0 +1,36 @@ +#!/usr/bin/env python + +from pmd_beamphysics import pmd_init +from pykern import pkio +from rslume import Elegant, OPAL, Genesis2 +import h5py +import os + +# patch numpy +import numpy +numpy.float = float + + +def code_kwargs(src_dir, input_file): + workdir = f"{src_dir}/out" + pkio.mkdir_parent(workdir) + return dict( + input_file=f"{src_dir}/{input_file}", + workdir=workdir, + verbose=True, + use_temp_dir=False, + ) + + +def run_and_save_particles(sim): + sim.run() + p = sim.final_particles() + if p: + with h5py.File(os.path.join(sim.workdir, "openpmd.h5"), "w") as f: + pmd_init(f, basePath='/', particlesPath='/' ) + p.write(f) + + +sim1 = Genesis2(**code_kwargs("run1", "genesis.in")) +run_and_save_particles(sim1) + diff --git a/tests/template/omega_data/opal_elegant.txt b/tests/template/omega_data/opal_elegant.txt new file mode 100644 index 0000000000..f36b855184 --- /dev/null +++ b/tests/template/omega_data/opal_elegant.txt @@ -0,0 +1,40 @@ +#!/usr/bin/env python + +from pmd_beamphysics import pmd_init +from pykern import pkio +from rslume import Elegant, OPAL, Genesis2 +import h5py +import os + +# patch numpy +import numpy +numpy.float = float + + +def code_kwargs(src_dir, input_file): + workdir = f"{src_dir}/out" + pkio.mkdir_parent(workdir) + return dict( + input_file=f"{src_dir}/{input_file}", + workdir=workdir, + verbose=True, + use_temp_dir=False, + ) + + +def run_and_save_particles(sim): + sim.run() + p = sim.final_particles() + if p: + with h5py.File(os.path.join(sim.workdir, "openpmd.h5"), "w") as f: + pmd_init(f, basePath='/', particlesPath='/' ) + p.write(f) + + +sim1 = OPAL(**code_kwargs("run1", "opal.in"), update_filenames=True) +run_and_save_particles(sim1) + +sim2 = Elegant(**code_kwargs("run2", "elegant.ele"), update_filenames=True) +sim2.set_particle_input(sim1.final_particles()) +run_and_save_particles(sim2) + diff --git a/tests/template/omega_test.py b/tests/template/omega_test.py index dbdabba4da..5b0fc274f6 100644 --- a/tests/template/omega_test.py +++ b/tests/template/omega_test.py @@ -18,10 +18,11 @@ def test_sims(fc): from pykern import pkunit + import sirepo.template.omega def _case(name, sims): fc.sr_get_root(sim_type=_SIM_TYPE) - r = fc.sr_post( + sim = fc.sr_post( "newSimulation", PKDict( simulationType=fc.sr_sim_type, @@ -29,12 +30,13 @@ def _case(name, sims): name=name, ), ) - r.models.simWorkflow = _workflow(fc, sims) - r = fc.sr_run_sim( - data=fc.sr_post("saveSimulationData", r), - model="animation", + sim.models.simWorkflow = _workflow(fc, sims) + pkunit.file_eq( + pkunit.data_dir().join(f"{name}.txt"), + actual=sirepo.template.omega.python_source_for_model( + sim, model="animation", qcall=None + ), ) - pkunit.pkok("completed", r.state) def _coupled_sim(stype, sname): r = fc.sr_sim_data(sim_type=stype, sim_name=sname) From 13d9b004581efadcb2b4b9090b02bd11bf67f6e6 Mon Sep 17 00:00:00 2001 From: Paul Moeller Date: Mon, 29 Apr 2024 23:28:42 +0000 Subject: [PATCH 22/23] for #6653 improve omega code generation and dark-mode improvements --- sirepo/package_data/static/css/omega.css | 5 +++ .../static/js/sirepo-components.js | 6 +++ .../template/omega/parameters.py.jinja | 29 ++++++------- sirepo/template/omega.py | 27 ++++++------ tests/template/omega_data/genesis.txt | 27 ++++++------ tests/template/omega_data/opal_elegant.txt | 42 ++++++++++--------- 6 files changed, 73 insertions(+), 63 deletions(-) diff --git a/sirepo/package_data/static/css/omega.css b/sirepo/package_data/static/css/omega.css index 61749b6b34..dc15b4f6bf 100644 --- a/sirepo/package_data/static/css/omega.css +++ b/sirepo/package_data/static/css/omega.css @@ -215,4 +215,9 @@ .table-hover > tbody > tr:hover { background-color: var(--sr-item-bg-dark-mode); } + + .sr-login-panel, .help-block { + color: var(--sr-panel-text-dark-mode); + } + } diff --git a/sirepo/package_data/static/js/sirepo-components.js b/sirepo/package_data/static/js/sirepo-components.js index 96a2f98a14..6e030f52e4 100644 --- a/sirepo/package_data/static/js/sirepo-components.js +++ b/sirepo/package_data/static/js/sirepo-components.js @@ -3465,6 +3465,7 @@ SIREPO.app.directive('completeRegistration', function() { return { restrict: 'A', template: ` + `, }; }); @@ -3493,6 +3495,7 @@ SIREPO.app.directive('emailLogin', function(requestSender, errorService) { restrict: 'A', scope: {}, template: ` + `, controller: function($scope) { function handleResponse(data) { @@ -3570,6 +3574,7 @@ SIREPO.app.directive('emailLoginConfirm', function() { return { restrict: 'A', template: ` + + `, }; }); diff --git a/sirepo/package_data/template/omega/parameters.py.jinja b/sirepo/package_data/template/omega/parameters.py.jinja index 5cfce14840..a6c278f3e7 100644 --- a/sirepo/package_data/template/omega/parameters.py.jinja +++ b/sirepo/package_data/template/omega/parameters.py.jinja @@ -2,34 +2,35 @@ from pmd_beamphysics import pmd_init from pykern import pkio from rslume import Elegant, OPAL, Genesis2 import h5py +import numpy import os # patch numpy -import numpy numpy.float = float -def code_kwargs(src_dir, input_file): - workdir = f"{src_dir}/out" - pkio.mkdir_parent(workdir) - return dict( - input_file=f"{src_dir}/{input_file}", - workdir=workdir, - verbose=True, - use_temp_dir=False, - ) - - def run_and_save_particles(sim): + pkio.mkdir_parent(sim.workdir) sim.run() p = sim.final_particles() if p: with h5py.File(os.path.join(sim.workdir, "{{ particleOutfile }}"), "w") as f: - pmd_init(f, basePath='/', particlesPath='/' ) + pmd_init(f, basePath="/", particlesPath="/" ) p.write(f) + return sim {% for sim in simCall %} -{{ sim }} +{{ sim.name }} = run_and_save_particles( + {{ sim.code }}( + input_file="{{ sim.name }}/{{ sim.input_file }}", + use_temp_dir=False, + workdir="{{ sim.name}}/{{ sim.out }}", + initial_particles={{ sim.initial_particles }}, + {% if sim.update_filenames %} + update_filenames=True, + {% endif %} + ), +) {% endfor %} diff --git a/sirepo/template/omega.py b/sirepo/template/omega.py index 12e0212ff1..36ff863b95 100644 --- a/sirepo/template/omega.py +++ b/sirepo/template/omega.py @@ -68,6 +68,7 @@ Cx="run_setup.centroid.sdds", Cy="run_setup.centroid.sdds", ) +_OUT_DIR = "out" _PARTICLE_OUTFILE = "openpmd.h5" _RELATED_SIMS_FOLDER = "/Omega" _SUCCESS_OUTPUT_FILE = PKDict( @@ -283,23 +284,19 @@ def _coupled_sims_list(data): def _generate_parameters_file(data): - def _sim_call(idx, sim): - r = ( - f"sim{idx + 1} = {_SIM_TYPE_TO_CODE[sim.sim_type]}(" - + f'**code_kwargs("run{idx + 1 }", "{_SIM_TYPE_TO_INPUT_FILE[sim.sim_type]}")' - ) - if sim.sim_type in ("opal", "elegant"): - r += ", update_filenames=True" - r += ")\n" - if idx > 0: - r += f"sim{idx + 1}.set_particle_input(sim{idx}.final_particles())\n" - r += f"run_and_save_particles(sim{idx + 1})" - return r - res, v = template_common.generate_parameters_file(data) v.simCall = [] for idx, s in enumerate(_coupled_sims_list(data)): - v.simCall.append(_sim_call(idx, s)) + v.simCall.append( + PKDict( + name=f"run{idx + 1}", + input_file=_SIM_TYPE_TO_INPUT_FILE[s.sim_type], + code=_SIM_TYPE_TO_CODE[s.sim_type], + out=_OUT_DIR, + initial_particles=f"run{idx}.final_particles()" if idx else "None", + update_filenames=True if s.sim_type in ("opal", "elegant") else False, + ) + ) v.particleOutfile = _PARTICLE_OUTFILE return res + template_common.render_jinja(SIM_TYPE, v) @@ -480,7 +477,7 @@ def _sim_in_dir(run_dir, sim_count): def _sim_out_dir(run_dir, sim_count): - return _sim_in_dir(run_dir, sim_count).join("out") + return _sim_in_dir(run_dir, sim_count).join(_OUT_DIR) def _sim_info(dm, idx): diff --git a/tests/template/omega_data/genesis.txt b/tests/template/omega_data/genesis.txt index e4334fcccc..3a5364da2d 100644 --- a/tests/template/omega_data/genesis.txt +++ b/tests/template/omega_data/genesis.txt @@ -4,33 +4,30 @@ from pmd_beamphysics import pmd_init from pykern import pkio from rslume import Elegant, OPAL, Genesis2 import h5py +import numpy import os # patch numpy -import numpy numpy.float = float -def code_kwargs(src_dir, input_file): - workdir = f"{src_dir}/out" - pkio.mkdir_parent(workdir) - return dict( - input_file=f"{src_dir}/{input_file}", - workdir=workdir, - verbose=True, - use_temp_dir=False, - ) - - def run_and_save_particles(sim): + pkio.mkdir_parent(sim.workdir) sim.run() p = sim.final_particles() if p: with h5py.File(os.path.join(sim.workdir, "openpmd.h5"), "w") as f: - pmd_init(f, basePath='/', particlesPath='/' ) + pmd_init(f, basePath="/", particlesPath="/" ) p.write(f) + return sim -sim1 = Genesis2(**code_kwargs("run1", "genesis.in")) -run_and_save_particles(sim1) +run1 = run_and_save_particles( + Genesis2( + input_file="run1/genesis.in", + use_temp_dir=False, + workdir="run1/out", + initial_particles=None, + ), +) diff --git a/tests/template/omega_data/opal_elegant.txt b/tests/template/omega_data/opal_elegant.txt index f36b855184..0269e6e4d7 100644 --- a/tests/template/omega_data/opal_elegant.txt +++ b/tests/template/omega_data/opal_elegant.txt @@ -4,37 +4,41 @@ from pmd_beamphysics import pmd_init from pykern import pkio from rslume import Elegant, OPAL, Genesis2 import h5py +import numpy import os # patch numpy -import numpy numpy.float = float -def code_kwargs(src_dir, input_file): - workdir = f"{src_dir}/out" - pkio.mkdir_parent(workdir) - return dict( - input_file=f"{src_dir}/{input_file}", - workdir=workdir, - verbose=True, - use_temp_dir=False, - ) - - def run_and_save_particles(sim): + pkio.mkdir_parent(sim.workdir) sim.run() p = sim.final_particles() if p: with h5py.File(os.path.join(sim.workdir, "openpmd.h5"), "w") as f: - pmd_init(f, basePath='/', particlesPath='/' ) + pmd_init(f, basePath="/", particlesPath="/" ) p.write(f) + return sim -sim1 = OPAL(**code_kwargs("run1", "opal.in"), update_filenames=True) -run_and_save_particles(sim1) - -sim2 = Elegant(**code_kwargs("run2", "elegant.ele"), update_filenames=True) -sim2.set_particle_input(sim1.final_particles()) -run_and_save_particles(sim2) +run1 = run_and_save_particles( + OPAL( + input_file="run1/opal.in", + use_temp_dir=False, + workdir="run1/out", + initial_particles=None, + update_filenames=True, + ), +) + +run2 = run_and_save_particles( + Elegant( + input_file="run2/elegant.ele", + use_temp_dir=False, + workdir="run2/out", + initial_particles=run1.final_particles(), + update_filenames=True, + ), +) From 0b376f211476c4c25cbed5c834cfe601d7f8dcd0 Mon Sep 17 00:00:00 2001 From: Paul Moeller Date: Wed, 1 May 2024 22:37:56 +0000 Subject: [PATCH 23/23] for #6653 remove NaN particles. - also added better defaults for common elegant plots --- sirepo/package_data/static/js/elegant.js | 33 +++++++++++++++++++++++- sirepo/template/omega.py | 15 ++++++++--- 2 files changed, 43 insertions(+), 5 deletions(-) diff --git a/sirepo/package_data/static/js/elegant.js b/sirepo/package_data/static/js/elegant.js index d1031081ff..89a6d32c00 100644 --- a/sirepo/package_data/static/js/elegant.js +++ b/sirepo/package_data/static/js/elegant.js @@ -484,6 +484,37 @@ SIREPO.app.controller('VisualizationController', function(appState, elegantServi return columns[1]; } + const defaultColumns = { + twiss_output: { + "y1": "betax", + "y2": "betay", + "y3": "etax", + }, + "run_setup.sigma": { + "y1": "Sx", + "y2": "Sy", + "y3": "Ss", + }, + "run_setup.centroid": { + "y1": "Cx", + "y2": "Cy", + }, + }; + + function setDefaultColumns(model, plottableColumns) { + model.x = plottableColumns[0]; + for (const f in defaultColumns) { + if (model.xFile.includes(f)) { + for (const y in defaultColumns[f]) { + if (plottableColumns.includes(defaultColumns[f][y])) { + model[y] = defaultColumns[f][y]; + } + } + model.includeLattice = "1"; + } + } + } + self.simHandleStatus = function (data) { self.simulationAlerts = data.alert || ''; if (data.frameCount) { @@ -561,9 +592,9 @@ SIREPO.app.controller('VisualizationController', function(appState, elegantServi m = appState.models[modelKey] = { xFile: info.filename, y1File: info.filename, - x: info.plottableColumns[0], xFileId: info.id, }; + setDefaultColumns(m, info.plottableColumns); // Only display the first outputFile if (i > 0 && ! panelState.isHidden(modelKey)) { panelState.toggleHidden(modelKey); diff --git a/sirepo/template/omega.py b/sirepo/template/omega.py index 36ff863b95..e3a4047366 100644 --- a/sirepo/template/omega.py +++ b/sirepo/template/omega.py @@ -396,9 +396,7 @@ def _col(column_name): py=p.py, pz=p.pz, )[column_name] - c = numpy.array(c) - c[numpy.isnan(c)] = 0.0 - return c + return numpy.array(c) def _x_label(x_name): if x_name == "z" and sim_type != "opal": @@ -407,8 +405,17 @@ def _x_label(x_name): _phase_plot_args(sim_type, frame_args) if sim_type in ("elegant", "opal", "genesis"): + # remove any NaN particles + x = _col(frame_args.x) + y = _col(frame_args.y) + if numpy.isnan(x).any(): + x = x[~numpy.isnan(x)] + y = y[~numpy.isnan(x)] + if numpy.isnan(y).any(): + x = x[~numpy.isnan(y)] + y = y[~numpy.isnan(y)] return template_common.heatmap( - values=[_col(frame_args.x), _col(frame_args.y)], + values=[x, y], model=frame_args, plot_fields=PKDict( x_label=_x_label(frame_args.x),