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/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/elegant.js b/sirepo/package_data/static/js/elegant.js index 2b622b7e94..4219a58428 100644 --- a/sirepo/package_data/static/js/elegant.js +++ b/sirepo/package_data/static/js/elegant.js @@ -485,6 +485,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) { @@ -562,9 +593,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/package_data/static/js/omega.js b/sirepo/package_data/static/js/omega.js index 781bb14c58..786c7778d9 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 = []; @@ -183,7 +186,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: { @@ -196,20 +199,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; diff --git a/sirepo/package_data/static/js/sirepo-components.js b/sirepo/package_data/static/js/sirepo-components.js index 55b3593925..c2047d6a87 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: ` +
@@ -3484,6 +3485,7 @@ SIREPO.app.directive('completeRegistration', function() {
+
`, }; }); @@ -3493,6 +3495,7 @@ SIREPO.app.directive('emailLogin', function(requestSender, errorService) { restrict: 'A', scope: {}, template: ` +
@@ -3521,6 +3524,7 @@ SIREPO.app.directive('emailLogin', function(requestSender, errorService) {

We just emailed a confirmation link to {{ data.sentEmail }}. Click the link and you'll be signed in. You may close this window.

+
`, controller: function($scope) { function handleResponse(data) { @@ -3570,6 +3574,7 @@ SIREPO.app.directive('emailLoginConfirm', function() { return { restrict: 'A', template: ` +

Please click the button below to complete the login process.

@@ -3579,6 +3584,7 @@ SIREPO.app.directive('emailLoginConfirm', function() {
+ `, }; }); @@ -5246,23 +5252,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)); + } + + function itemName(sim) { + const n = sim.folder === '/' + ? `/${sim.name}` + : `${sim.folder}/${sim.name}`; + return sim.invalidMsg + ? `${n} <${sim.invalidMsg}>` + : n; + } - $scope.openSimulation = function() { + $scope.openSimulation = () => { if ($scope.model && $scope.model[$scope.field]) { requestSender.openSimulation( $scope.code, @@ -5271,14 +5292,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); } }, { 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 6669a9eb01..a6c278f3e7 100644 --- a/sirepo/package_data/template/omega/parameters.py.jinja +++ b/sirepo/package_data/template/omega/parameters.py.jinja @@ -1,284 +1,36 @@ - +from pmd_beamphysics import pmd_init from pykern import pkio -from pykern.pkcollections import PKDict -from rsbeams.rsdata import switchyard -from sirepo.template.lattice import LatticeUtil +from rslume import Elegant, OPAL, Genesis2 +import h5py 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 - - -_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): - 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' - 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): - 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}" - ) - 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 _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): - 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}') - 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) - 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) - else: - raise AssertionError(f"unhandled sim_type={s.sim_type}") - 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( - sim_type="{{sim.sim_type}}", - sim_id="{{sim.sim_id}}", +# patch numpy +numpy.float = float + + +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="/" ) + p.write(f) + return sim + + +{% for sim in simCall %} +{{ 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/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/genesis.py b/sirepo/template/genesis.py index ca4cf47d3f..960231773e 100644 --- a/sirepo/template/genesis.py +++ b/sirepo/template/genesis.py @@ -264,7 +264,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): @@ -421,15 +421,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}") @@ -497,14 +497,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( diff --git a/sirepo/template/omega.py b/sirepo/template/omega.py index f299de13e3..e3a4047366 100644 --- a/sirepo/template/omega.py +++ b/sirepo/template/omega.py @@ -10,66 +10,27 @@ 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 -import pmd_beamphysics.interfaces.elegant -import pmd_beamphysics.interfaces.genesis -import pmd_beamphysics.interfaces.opal import re +import pmd_beamphysics import sirepo.sim_data import sirepo.template + _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", @@ -107,12 +68,24 @@ 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( 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): @@ -157,60 +130,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): @@ -227,7 +157,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: @@ -246,6 +176,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) @@ -254,10 +191,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: @@ -282,6 +221,7 @@ def stateful_compute_get_opal_sim_list(**kwargs): def write_parameters(data, run_dir, is_parallel): + _prepare_subsims(data, run_dir) pkio.write_text( run_dir.join(template_common.PARAMETERS_PYTHON_FILE), _generate_parameters_file(data), @@ -324,9 +264,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) @@ -341,7 +280,24 @@ def _generate_parameters_file(data): break if not sim_list: raise AssertionError("No simulations selected") - v.simList = sim_list + return sim_list + + +def _generate_parameters_file(data): + res, v = template_common.generate_parameters_file(data) + v.simCall = [] + for idx, s in enumerate(_coupled_sims_list(data)): + 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) @@ -363,7 +319,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) @@ -385,7 +341,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 @@ -403,7 +359,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] @@ -418,9 +374,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)) @@ -433,41 +386,107 @@ def _plot_field_dist(sim_type, frame_args): def _plot_phase(sim_type, frame_args): - _phase_plot_args(sim_type, frame_args) + def _col(column_name): + p = pmd_beamphysics.ParticleGroup(h5=str(frame_args.run_dir.join("openpmd.h5"))) + c = PKDict( + x=p.x, + y=p.y, + z=p.z if sim_type == "opal" else p.t, + px=p.px, + py=p.py, + pz=p.pz, + )[column_name] + return numpy.array(c) + + def _x_label(x_name): + if x_name == "z" and sim_type != "opal": + return "t" + return x_name - if sim_type == "opal": - 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 + _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=[x, y], + model=frame_args, + plot_fields=PKDict( + x_label=_x_label(frame_args.x), + y_label=frame_args.y, + title=_PLOT_TITLE[frame_args.x + "-" + frame_args.y], ), ) - 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, - ) - 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], - ) - ) raise AssertionError("unhandled sim_type for sim_frame(): {}".format(sim_type)) -def _sim_dir(run_dir, sim_count): +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) + 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): + data = _read_sim("genesis", genesis_id) + data.computeModel = "animation" + if "report" in data: + del data["report"] + sirepo.simulation_db.prepare_simulation(data, run_dir) + 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): + 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) + + +def _prepare_subsims(data, run_dir): + sim_list = _coupled_sims_list(data) + for idx in range(len(sim_list)): + s = sim_list[idx] + d = _sim_in_dir(run_dir, idx + 1) + pkio.unchecked_remove(d) + pkio.mkdir_parent(d) + if s.sim_type == "opal": + _prepare_opal(d, s.sim_id) + elif s.sim_type == "elegant": + _prepare_elegant(d, s.sim_id) + elif s.sim_type == "genesis": + _prepare_genesis(d, s.sim_id) + else: + raise AssertionError(f"unhandled sim_type={s.sim_type}") + + +def _read_sim(sim_type, sim_id): + return sirepo.sim_data.get_class(sim_type).sim_db_read_sim(sim_id) + + +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_DIR) + + def _sim_info(dm, idx): s = dm.simWorkflow.coupledSims if len(s) > idx: diff --git a/tests/template/omega_data/genesis.txt b/tests/template/omega_data/genesis.txt new file mode 100644 index 0000000000..3a5364da2d --- /dev/null +++ b/tests/template/omega_data/genesis.txt @@ -0,0 +1,33 @@ +#!/usr/bin/env python + +from pmd_beamphysics import pmd_init +from pykern import pkio +from rslume import Elegant, OPAL, Genesis2 +import h5py +import numpy +import os + +# patch numpy +numpy.float = float + + +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="/" ) + p.write(f) + return sim + + +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 new file mode 100644 index 0000000000..0269e6e4d7 --- /dev/null +++ b/tests/template/omega_data/opal_elegant.txt @@ -0,0 +1,44 @@ +#!/usr/bin/env python + +from pmd_beamphysics import pmd_init +from pykern import pkio +from rslume import Elegant, OPAL, Genesis2 +import h5py +import numpy +import os + +# patch numpy +numpy.float = float + + +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="/" ) + p.write(f) + return sim + + +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, + ), +) + 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)