class Astra: """ Astra simulation object. Essential methods: .__init__(...) .configure() .run() Input deck is held in .input Output data is parsed into .output .load_particles() will load particle data into .output['particles'][...] The Astra binary file can be set on init. If it doesn't exist, configure will check the $ASTRA_BIN environmental variable. """ def __init__(self, input_file=None, initial_particles=None, astra_bin='$ASTRA_BIN', use_tempdir=True, workdir=None, verbose=False): # Save init self.original_input_file = input_file self.initial_particles = initial_particles self.use_tempdir = use_tempdir self.workdir = workdir if workdir: assert os.path.exists( workdir), 'workdir does not exist: ' + workdir self.verbose = verbose self.astra_bin = astra_bin # These will be set self.log = [] self.output = {'stats': {}, 'particles': {}, 'run_info': {}} self.timeout = None self.error = False # Run control self.finished = False self.configured = False self.using_tempdir = False # Call configure if input_file: self.load_input(input_file) self.configure() else: self.vprint('Warning: Input file does not exist. Not configured.') self.original_input_file = 'astra.in' def clean_output(self): run_number = parsers.astra_run_extension(self.input['newrun']['run']) outfiles = parsers.find_astra_output_files(self.input_file, run_number) for f in outfiles: os.remove(f) def clean_particles(self): run_number = parsers.astra_run_extension(self.input['newrun']['run']) phase_files = parsers.find_phase_files(self.input_file, run_number) files = [x[0] for x in phase_files] # This is sorted by approximate z for f in files: os.remove(f) # Convenience routines @property def particles(self): return self.output['particles'] def stat(self, key): return self.output['stats'][key] def particle_stat(self, key, alive_only=True): """ Compute a statistic from the particles. Alive particles have status == 1. By default, statistics will only be computed on these. n_dead will override the alive_only flag, and return the number of particles with status < -6 (Astra convention) """ if key == 'n_dead': return np.array( [len(np.where(P.status < -6)[0]) for P in self.particles]) if key == 'n_alive': return np.array( [len(np.where(P.status > -6)[0]) for P in self.particles]) pstats = [] for P in self.particles: if alive_only and P.n_dead > 0: P = P.where(P.status == 1) pstats.append(P[key]) return np.array(pstats) def configure(self): self.configure_astra(workdir=self.workdir) def configure_astra(self, input_filePath=None, workdir=None): if input_filePath: self.load_input(input_filePath) # Check that binary exists self.astra_bin = tools.full_path(self.astra_bin) assert os.path.exists( self.astra_bin ), 'ERROR: Astra binary does not exist:' + self.astra_bin # Set paths if self.use_tempdir: # Need to attach this to the object. Otherwise it will go out of scope. self.tempdir = tempfile.TemporaryDirectory(dir=workdir) self.path = self.tempdir.name else: # Work in place self.path = self.original_path self.input_file = os.path.join(self.path, self.original_input_file) self.configured = True def load_initial_particles(self, h5): """Loads a openPMD-beamphysics particle h5 handle or file""" P = ParticleGroup(h5=h5) self.initial_particles = P def load_input(self, input_filePath, absolute_paths=True): f = tools.full_path(input_filePath) self.original_path, self.original_input_file = os.path.split( f) # Get original path, filename self.input = parsers.parse_astra_input_file(f) if absolute_paths: parsers.fix_input_paths(self.input, root=self.original_path) def load_output(self, include_particles=True): """ Loads Astra output files into .output .output is a dict with dicts: .stats .run_info .other and if include_particles, .particles = list of ParticleGroup objects """ run_number = parsers.astra_run_extension(self.input['newrun']['run']) outfiles = parsers.find_astra_output_files(self.input_file, run_number) #assert len(outfiles)>0, 'No output files found' stats = self.output['stats'] = {} for f in outfiles: type = parsers.astra_output_type(f) d = parsers.parse_astra_output_file(f) if type in ['Cemit', 'Xemit', 'Yemit', 'Zemit']: stats.update(d) elif type in ['LandF']: self.output['other'] = d else: raise ValueError(f'Unknown output type: {type}') # Check that the lengths of all arrays are the same nlist = {len(stats[k]) for k in stats} assert len( nlist ) == 1, f'Stat keys do not all have the same length: { [len(stats[k]) for k in stats]}' if include_particles: self.load_particles() def load_particles(self, end_only=False): # Clear existing particles self.output['particles'] = [] # Sort files by approximate z run_number = parsers.astra_run_extension(self.input['newrun']['run']) phase_files = parsers.find_phase_files(self.input_file, run_number) files = [x[0] for x in phase_files] # This is sorted by approximate z zapprox = [x[1] for x in phase_files] if end_only: files = files[-1:] if self.verbose: print('loading ' + str(len(files)) + ' particle files') print(zapprox) for f in files: pdat = parsers.parse_astra_phase_file(f) P = ParticleGroup(data=pdat) self.output['particles'].append(P) def run(self): if not self.configured: print('not configured to run') return self.run_astra(verbose=self.verbose, timeout=self.timeout) def get_run_script(self, write_to_path=True): """ Assembles the run script. Optionally writes a file 'run' with this line to path. """ #_, infile = os.path.split(self.input_file) runscript = [self.astra_bin, self.input_file] if write_to_path: with open(os.path.join(self.path, 'run'), 'w') as f: f.write(' '.join(runscript)) return runscript def run_astra(self, verbose=False, parse_output=True, timeout=None): """ Runs Astra Does not change directory - this has problems with multithreading. """ run_info = self.output['run_info'] = {} t1 = time() run_info['start_time'] = t1 if self.initial_particles: fname = self.write_initial_particles() self.input['newrun']['distribution'] = fname # Write input file from internal dict self.write_input_file() runscript = self.get_run_script() run_info['run_script'] = ' '.join(runscript) try: if timeout: res = tools.execute2(runscript, timeout=timeout) log = res['log'] self.error = res['error'] run_info['why_error'] = res['why_error'] # Log file must have this to have finished properly if log.find('finished simulation') == -1: run_info['error'] = True run_info.update({ 'error': True, 'why_error': "Couldn't find finished simulation" }) else: # Interactive output, for Jupyter log = [] for path in tools.execute(runscript): self.vprint(path, end="") log.append(path) self.log = log if parse_output: self.load_output() except Exception as ex: print('Run Aborted', ex) self.error = True run_info['why_error'] = str(ex) finally: run_info['run_time'] = time() - t1 run_info['run_error'] = self.error self.finished = True self.vprint(run_info) def fingerprint(self): """ Data fingerprint using the input. """ return tools.fingerprint(self.input) def vprint(self, *args, **kwargs): # Verbose print if self.verbose: print(*args, **kwargs) def units(self, key): if key in parsers.OutputUnits: return parsers.OutputUnits[key] else: return 'unknown unit' def write_input_file(self): parsers.write_namelists(self.input, self.input_file) def write_initial_particles(self, fname=None): if not fname: fname = os.path.join(self.path, 'astra.particles') self.initial_particles.write_astra(fname) self.vprint(f'Initial particles written to {fname}') return fname def load_archive(self, h5=None): """ Loads input and output from archived h5 file. See: Astra.archive """ if isinstance(h5, str): g = h5py.File(h5, 'r') glist = archive.find_astra_archives(g) n = len(glist) if n == 0: # legacy: try top level message = 'legacy' elif n == 1: gname = glist[0] message = f'group {gname} from' g = g[gname] else: raise ValueError( f'Multiple archives found in file {h5}: {glist}') self.vprint(f'Reading {message} archive file {h5}') else: g = h5 self.input = archive.read_input_h5(g['input']) self.output = archive.read_output_h5(g['output']) if 'initial_particles' in g: self.initial_particles = ParticleGroup(h5=g['initial_particles']) self.vprint( 'Loaded from archive. Note: Must reconfigure to run again.') self.configured = False def archive(self, h5=None): """ Archive all data to an h5 handle or filename. If no file is given, a file based on the fingerprint will be created. """ if not h5: h5 = 'astra_' + self.fingerprint() + '.h5' if isinstance(h5, str): g = h5py.File(h5, 'w') self.vprint(f'Archiving to file {h5}') else: # store directly in the given h5 handle g = h5 # Write basic attributes archive.astra_init(g) # Initial particles if self.initial_particles: self.initial_particles.write(g, name='initial_particles') # All input archive.write_input_h5(g, self.input) # All output archive.write_output_h5(g, self.output) return h5
class Astra: """ Astra simulation object. Essential methods: .__init__(...) .configure() .run() Input deck is held in .input Output data is parsed into .output .load_particles() will load particle data into .output['particles'][...] The Astra binary file can be set on init. If it doesn't exist, configure will check the $ASTRA_BIN environmental variable. """ def __init__(self, input_file=None, initial_particles=None, astra_bin='$ASTRA_BIN', use_tempdir=True, workdir=None, group=None, verbose=False): # Save init self.original_input_file = input_file self.initial_particles = initial_particles self.use_tempdir = use_tempdir self.workdir = workdir if workdir: assert os.path.exists( workdir), 'workdir does not exist: ' + workdir self.verbose = verbose self.astra_bin = astra_bin # These will be set self.log = [] self.output = {'stats': {}, 'particles': {}, 'run_info': {}} self.timeout = None self.error = False self.group = {} # Control Groups self.fieldmap = {} # Fieldmaps # Run control self.finished = False self.configured = False self.using_tempdir = False # Call configure if input_file: self.load_input(input_file) self.configure() # Add groups, if any. if group: for k, v in group.items(): self.add_group(k, **v) else: self.vprint('Warning: Input file does not exist. Not configured.') self.original_input_file = 'astra.in' def add_group(self, name, **kwargs): """ Add a control group. See control.py """ assert name not in self.input, f'{name} not allowed to be overwritten by group.' if name in self.group: self.vprint(f'Warning: group {name} already exists, overwriting.') g = ControlGroup(**kwargs) g.link(self.input) self.group[name] = g return self.group[name] def clean_output(self): run_number = parsers.astra_run_extension(self.input['newrun']['run']) outfiles = parsers.find_astra_output_files(self.input_file, run_number) for f in outfiles: os.remove(f) def clean_particles(self): run_number = parsers.astra_run_extension(self.input['newrun']['run']) phase_files = parsers.find_phase_files(self.input_file, run_number) files = [x[0] for x in phase_files] # This is sorted by approximate z for f in files: os.remove(f) # Convenience routines @property def particles(self): return self.output['particles'] def stat(self, key): return self.output['stats'][key] def particle_stat(self, key, alive_only=True): """ Compute a statistic from the particles. Alive particles have status == 1. By default, statistics will only be computed on these. n_dead will override the alive_only flag, and return the number of particles with status < -6 (Astra convention) """ if key == 'n_dead': return np.array( [len(np.where(P.status < -6)[0]) for P in self.particles]) if key == 'n_alive': return np.array( [len(np.where(P.status > -6)[0]) for P in self.particles]) pstats = [] for P in self.particles: if alive_only and P.n_dead > 0: P = P.where(P.status == 1) pstats.append(P[key]) return np.array(pstats) def configure(self): self.configure_astra(workdir=self.workdir) def configure_astra(self, input_filePath=None, workdir=None): if input_filePath: self.load_input(input_filePath) # Check that binary exists self.astra_bin = tools.full_path(self.astra_bin) assert os.path.exists( self.astra_bin ), 'ERROR: Astra binary does not exist:' + self.astra_bin # Set paths if self.use_tempdir: # Need to attach this to the object. Otherwise it will go out of scope. self.tempdir = tempfile.TemporaryDirectory(dir=workdir) self.path = self.tempdir.name else: # Work in place self.path = self.original_path self.input_file = os.path.join(self.path, self.original_input_file) self.configured = True def load_fieldmaps(self): """ Loads fieldmaps into Astra.fieldmap as a dict """ # Do not consider files if fieldmaps have been loaded. if self.fieldmap: strip_path = False else: strip_path = True self.fieldmap = load_fieldmaps(self, fieldmap_dict=self.fieldmap, search_paths=[self.path], verbose=self.verbose, strip_path=strip_path) def load_initial_particles(self, h5): """Loads a openPMD-beamphysics particle h5 handle or file""" P = ParticleGroup(h5=h5) self.initial_particles = P def load_input(self, input_filePath, absolute_paths=True): f = tools.full_path(input_filePath) self.original_path, self.original_input_file = os.path.split( f) # Get original path, filename self.input = parsers.parse_astra_input_file(f) if absolute_paths: parsers.fix_input_paths(self.input, root=self.original_path) def load_output(self, include_particles=True): """ Loads Astra output files into .output .output is a dict with dicts: .stats .run_info .other and if include_particles, .particles = list of ParticleGroup objects """ run_number = parsers.astra_run_extension(self.input['newrun']['run']) outfiles = parsers.find_astra_output_files(self.input_file, run_number) #assert len(outfiles)>0, 'No output files found' stats = self.output['stats'] = {} for f in outfiles: type = parsers.astra_output_type(f) d = parsers.parse_astra_output_file(f) if type in ['Cemit', 'Xemit', 'Yemit', 'Zemit']: stats.update(d) elif type in ['LandF']: self.output['other'] = d else: raise ValueError(f'Unknown output type: {type}') # Check that the lengths of all arrays are the same nlist = {len(stats[k]) for k in stats} assert len( nlist ) == 1, f'Stat keys do not all have the same length: { [len(stats[k]) for k in stats]}' if include_particles: self.load_particles() def load_particles(self, end_only=False): # Clear existing particles self.output['particles'] = [] # Sort files by approximate z run_number = parsers.astra_run_extension(self.input['newrun']['run']) phase_files = parsers.find_phase_files(self.input_file, run_number) files = [x[0] for x in phase_files] # This is sorted by approximate z zapprox = [x[1] for x in phase_files] if end_only: files = files[-1:] if self.verbose: print('loading ' + str(len(files)) + ' particle files') print(zapprox) for f in files: pdat = parsers.parse_astra_phase_file(f) P = ParticleGroup(data=pdat) self.output['particles'].append(P) def run(self): if not self.configured: print('not configured to run') return self.run_astra(verbose=self.verbose, timeout=self.timeout) def get_run_script(self, write_to_path=True): """ Assembles the run script. Optionally writes a file 'run' with this line to path. This expect to run with .path as the cwd. """ _, infile = os.path.split( self.input_file ) # Expect to run locally. Astra has problems with long paths. runscript = [self.astra_bin, infile] if write_to_path: with open(os.path.join(self.path, 'run'), 'w') as f: f.write(' '.join(runscript)) return runscript def run_astra(self, verbose=False, parse_output=True, timeout=None): """ Runs Astra Changes directory, so does not work with threads. """ run_info = self.output['run_info'] = {} t1 = time() run_info['start_time'] = t1 if self.initial_particles: fname = self.write_initial_particles() self.input['newrun']['distribution'] = fname # Write all input self.write_input() runscript = self.get_run_script() run_info['run_script'] = ' '.join(runscript) try: if timeout: res = tools.execute2(runscript, timeout=timeout, cwd=self.path) log = res['log'] self.error = res['error'] run_info['why_error'] = res['why_error'] # Log file must have this to have finished properly if log.find('finished simulation') == -1: run_info['error'] = True run_info.update({ 'error': True, 'why_error': "Couldn't find finished simulation" }) else: # Interactive output, for Jupyter log = [] for path in tools.execute(runscript, cwd=self.path): self.vprint(path, end="") log.append(path) self.log = log if parse_output: self.load_output() except Exception as ex: print('Run Aborted', ex) self.error = True run_info['why_error'] = str(ex) finally: run_info['run_time'] = time() - t1 run_info['run_error'] = self.error self.finished = True self.vprint(run_info) def fingerprint(self): """ Data fingerprint using the input. """ return tools.fingerprint(self.input) def vprint(self, *args, **kwargs): # Verbose print if self.verbose: print(*args, **kwargs) def units(self, key): if key in parsers.OutputUnits: return parsers.OutputUnits[key] else: return 'unknown unit' def load_archive(self, h5=None): """ Loads input and output from archived h5 file. See: Astra.archive """ if isinstance(h5, str): h5 = os.path.expandvars(h5) g = h5py.File(h5, 'r') glist = archive.find_astra_archives(g) n = len(glist) if n == 0: # legacy: try top level message = 'legacy' elif n == 1: gname = glist[0] message = f'group {gname} from' g = g[gname] else: raise ValueError( f'Multiple archives found in file {h5}: {glist}') self.vprint(f'Reading {message} archive file {h5}') else: g = h5 self.input = archive.read_input_h5(g['input']) self.output = archive.read_output_h5(g['output']) if 'initial_particles' in g: self.initial_particles = ParticleGroup(h5=g['initial_particles']) if 'fieldmap' in g: self.fieldmap = archive.read_fieldmap_h5(g['fieldmap']) if 'control_groups' in g: self.group = archive.read_control_groups_h5(g['control_groups'], verbose=self.verbose) self.vprint( 'Loaded from archive. Note: Must reconfigure to run again.') self.configured = False # Re-link groups # TODO: cleaner logic for _, cg in self.group.items(): cg.link(self.input) def archive(self, h5=None): """ Archive all data to an h5 handle or filename. If no file is given, a file based on the fingerprint will be created. """ if not h5: h5 = 'astra_' + self.fingerprint() + '.h5' if isinstance(h5, str): h5 = os.path.expandvars(h5) g = h5py.File(h5, 'w') self.vprint(f'Archiving to file {h5}') else: # store directly in the given h5 handle g = h5 # Write basic attributes archive.astra_init(g) # Initial particles if self.initial_particles: self.initial_particles.write(g, name='initial_particles') # Fieldmaps if self.fieldmap: archive.write_fieldmap_h5(g, self.fieldmap, name='fieldmap') # All input archive.write_input_h5(g, self.input) # All output archive.write_output_h5(g, self.output) # Control groups if self.group: archive.write_control_groups_h5(g, self.group, name='control_groups') return h5 @classmethod def from_archive(cls, archive_h5): """ Class method to return an GPT object loaded from an archive file """ c = cls() c.load_archive(archive_h5) return c @classmethod def from_yaml(cls, yaml_file): """ Returns an Astra object instantiated from a YAML config file Will load intial_particles from an h5 file. """ # Try file if os.path.exists(os.path.expandvars(yaml_file)): config = yaml.safe_load(open(yaml_file)) # The input file might be relative to the yaml file if 'input_file' in config: f = os.path.expandvars(config['input_file']) if not os.path.isabs(f): # Get the yaml file root root, _ = os.path.split(tools.full_path(yaml_file)) config['input_file'] = os.path.join(root, f) else: #Try raw string config = yaml.safe_load(yaml_file) # Form ParticleGroup from file if 'initial_particles' in config: f = config['initial_particles'] if not os.path.isabs(f): root, _ = os.path.split(tools.full_path(yaml_file)) f = os.path.join(root, f) config['initial_particles'] = ParticleGroup(f) return cls(**config) def write_fieldmaps(self): """ Writes any loaded fieldmaps to path """ if self.fieldmap: write_fieldmaps(self.fieldmap, self.path) self.vprint( f'{len(self.fieldmap)} fieldmaps written to {self.path}') def write_input(self): """ Writes all input. If fieldmaps have been loaded, these will also be written. """ self.write_fieldmaps() self.write_input_file() def write_input_file(self): if self.use_tempdir: make_symlinks = True else: make_symlinks = False writers.write_namelists(self.input, self.input_file, make_symlinks=make_symlinks, verbose=self.verbose) def write_initial_particles(self, fname=None): if not fname: fname = os.path.join(self.path, 'astra.particles') self.initial_particles.write_astra(fname) self.vprint(f'Initial particles written to {fname}') return fname def plot(self, y=['sigma_x', 'sigma_y'], x='mean_z', xlim=None, y2=[], nice=True, include_layout=True, include_labels=False, include_particles=True, include_legend=True, **kwargs): """ Plots stat output multiple keys. If a list of ykeys2 is given, these will be put on the right hand axis. This can also be given as a single key. Logical switches, all default to True: nice: a nice SI prefix and scaling will be used to make the numbers reasonably sized. include_legend: The plot will include the legend include_layout: the layout plot (fieldmaps) will be displayed at the bottom include_labels: the layout will include element labels. If there is no output to plot, the fieldmaps will be plotted with .plot_fieldmaps """ # Just plot fieldmaps if there are no if not self.output['stats']: return plot_fieldmaps(self, xlim=xlim, fieldmap_dict=self.fieldmap, **kwargs) plot_stats_with_layout(self, ykeys=y, ykeys2=y2, xkey=x, xlim=xlim, nice=nice, include_layout=include_layout, include_labels=include_labels, include_particles=include_particles, include_legend=include_legend, **kwargs) def plot_fieldmaps(self, **kwargs): return plot_fieldmaps(self, **kwargs) def copy(self): """ Returns a deep copy of this object. If a tempdir is being used, will clear this and deconfigure. """ A2 = deepcopy(self) # Clear this if A2.use_tempdir: A2.path = None A2.configured = False return A2 def __getitem__(self, key): """ Convenience syntax to get a header or element attribute. Special syntax: end_X will return the final item in a stat array X Example: 'end_norm_emit_x' particles:N will return a ParticleGroup N from the .particles list Example: 'particles:-1' returns the readback of the final particles particles:N:Y ParticleGroup N's property Y Example: 'particles:-1:sigma_x' returns sigma_x from the end of the particles list. See: __setitem__ """ # Object attributes if hasattr(self, key): return getattr(self, key) # Send back top level input (namelist) or group object. # Do not add these to __setitem__. The user shouldn't be allowed to change them as a whole, # because it will break all the links. if key in self.group: return self.group[key] if key in self.input: return self.input[key] if key.startswith('end_'): key2 = key[len('end_'):] assert key2 in self.output[ 'stats'], f'{key} does not have valid output stat: {key2}' return self.output['stats'][key2][-1] if key.startswith('particles:'): key2 = key[len('particles:'):] x = key2.split(':') if len(x) == 1: return self.particles[int(x[0])] else: return self.particles[int(x[0])][x[1]] # key isn't an ele or group, should have property s x = key.split(':') assert len( x ) == 2, f'{x} was not found in group or input dict, so should have : ' name, attrib = x[0], x[1] # Look in input and group if name in self.input: return self.input[name][attrib] elif name in self.group: return self.group[name][attrib] def __setitem__(self, key, item): """ Convenience syntax to set namelist or group attribute. attribute_string should be 'header:key' or 'ele_name:key' Examples of attribute_string: 'header:Np', 'SOL1:solenoid_field_scale' Settable attributes can also be given: ['stop'] = 1.2345 will set Impact.stop = 1.2345 """ # Set attributes if hasattr(self, key): setattr(self, key, item) return # Must be in input or group name, attrib = key.split(':') if name in self.input: self.input[name][attrib] = item elif name in self.group: self.group[name][attrib] = item else: raise ValueError( f'{name} does not exist in eles or groups of the Impact object.' )
class GPT: """ GPT simulation object. Essential methods: .__init__(...) .configure() .run() Input deck is held in .input Output data is parsed into .output .load_screens() will load particle data into .screen[...] The GPT binary file can be set on init. If it doesn't exist, configure will check the $GPT_BIN environmental variable. """ def __init__(self, input_file=None, initial_particles=None, gpt_bin='$GPT_BIN', use_tempdir=True, workdir=None, timeout=None, verbose=False, ccs_beg='wcs', ref_ccs=False, kill_msgs=DEFAULT_KILL_MSGS): # Save init self.original_input_file = input_file self.initial_particles = initial_particles self.use_tempdir = use_tempdir self.workdir = workdir if workdir: assert os.path.exists( workdir), 'workdir does not exist: ' + workdir self.workdir = os.path.abspath(workdir) self.verbose = verbose self.gpt_bin = gpt_bin # These will be set self.log = [] self.output = {} #self.screen = [] # list of screens self.timeout = timeout self.error = False # Run control self.finished = False self.configured = False self.using_tempdir = False self.ccs_beg = ccs_beg self.ref_ccs = ref_ccs self.kill_msgs = kill_msgs # Call configure if input_file: self.load_input(input_file) self.configure() else: self.vprint('Warning: Input file does not exist. Not configured.') def configure(self): """ Convenience wrapper for configure_gpt """ self.configure_gpt(workdir=self.workdir) def configure_gpt(self, input_filePath=None, workdir=None): """ Configure the GPT object """ if input_filePath: self.load_input(input_filePath) # Check that binary exists self.gpt_bin = tools.full_path(self.gpt_bin) assert os.path.exists( self.gpt_bin), 'ERROR: GPT binary does not exist:' + self.gpt_bin # Set paths if self.use_tempdir: # Need to attach this to the object. Otherwise it will go out of scope. self.tempdir = tempfile.TemporaryDirectory(dir=workdir) self.path = self.tempdir.name elif (workdir): # Use the top level of the provided workdir self.path = workdir else: # Work in location of the template file self.path = self.original_path self.input_file = os.path.join(self.path, self.original_input_file) parsers.set_support_files(self.input['lines'], self.original_path) self.vprint('GPT.configure_gpt:') self.vprint( f' Original input file "{self.original_input_file}" in "{self.original_path}"' ) self.vprint(f' Configured to run in "{self.path}"') self.configured = True def load_input(self, input_filePath, absolute_paths=True): """ Load the GPT template file """ f = tools.full_path(input_filePath) self.original_path, self.original_input_file = os.path.split( f) # Get original path, filename self.input = parsers.parse_gpt_input_file(f) def get_dist_file(self): """ Find the distribution input file name in the GPT file """ for line in self.input['lines']: if ('setfile' in line): return parse_gpt_string(line)[1] def set_dist_file(self, dist_file): """ Set the input distirbution file name in a GPT file """ dist_file_set = False for ii, line in enumerate(self.input['lines']): if ('setfile' in line): gpt_strs = parse_gpt_string(line) assert len( gpt_strs ) == 2, "Couldn't find distribution input file strs." assert gpt_strs[ 0] == 'beam', "Could not find beam defintion in setfile str." self.input['lines'][ii] = f'setfile("beam", "{dist_file}");' dist_file_set = True if (not dist_file_set): self.input['lines'].append(f'setfile("beam", "{dist_file}");') def set_variable(self, variable, value): """ Set variable in the GPT input file to a new value """ if (variable in self.input["variables"]): self.input['variables'][variable] = value return True else: return False def set_variables(self, variables): """ Set a list of variables (variable.keys) to new values (variables.values()) in the GPT Input file """ return { var: self.set_variable(var, variables[var]) for var in variables.keys() } def load_output(self, file='gpt.out.gdf'): """ loads the GPT raw data and puts it into particle groups """ self.vprint(f' Loading GPT data from {self.get_gpt_output_file()}') touts, screens = parsers.read_gdf_file(file, self.verbose) # Raw GPT data self.output['particles'] = raw_data_to_particle_groups( touts, screens, verbose=self.verbose, ref_ccs=self.ref_ccs) self.output['n_tout'] = len(touts) self.output['n_screen'] = len(screens) @property def n_tout(self): """ number of tout particle groups """ return self.output['n_tout'] @property def n_screen(self): """ number of screen particle groups""" return self.output['n_screen'] @property def tout(self): """ Returns output particle groups for touts """ if ('particles' in self.output): return self.output['particles'][:self.output['n_tout']] @property def tout_ccs(self): """ Returns output particle groups for touts transformed into centroid coordinate system """ if ('particles' in self.output): return [ transform_to_centroid_coordinates(tout) for tout in self.tout ] @property def s_ccs(self): s = [ np.sqrt(self.tout[0]['mean_x']**2 + self.tout[0]['mean_y']**2 + self.tout[0]['mean_z']**2) ] current_tout = self.tout[0] current_p = current_tout['mean_p'] for next_tout in self.tout[1:]: next_p = next_tout['mean_p'] if (np.abs(current_p - next_p) / current_p < 1e-5): # Beam drifting beta = current_tout['mean_beta'] dt = next_tout['mean_t'] - current_tout['mean_t'] ds = dt * beta * c else: # Assume straight line acceleration dx = next_tout['mean_x'] - current_tout['mean_x'] dy = next_tout['mean_y'] - current_tout['mean_y'] dz = next_tout['mean_z'] - current_tout['mean_z'] ds = np.sqrt(dx**2 + dy**2 + dz**2) s.append(s[-1] + ds) current_tout = next_tout return np.array(s) def tout_stat(self, key=None): """ Returns array of stats for key from tout particle groups """ return self.stat(key, data_type='tout') def tout_ccs_stat(self, key=None): """ Returns array of stats for key from tout particle groups """ return self.stat(key, data_type='tout_ccs') @property def screen(self): """ Returns output particle groups for screens """ if ('particles' in self.output): return self.output['particles'][self.output['n_tout']:] def screen_stat(self, key): """ Returns array of stats for key from screen particle groups """ return self.stat(key, data_type='screen') @property def particles(self): """ Returns output particle groups for touts + screens """ if ('particles' in self.output): return self.output['particles'] def trajectory(self, pid, data_type='tout'): """ Returns a 3d particle trajectory for particle with id = pid """ if (data_type == 'tout'): particle_groups = self.tout elif (data_type == 'screen'): particle_groups = self.screen else: raise ValueError( f'GPT.trajectory got an unsupported data type = {data_type}.') #for pg in particle_groups: # print(pg, pid in pg['id']) pgs_with_pid = [pg for pg in particle_groups if (pid in pg['id'])] if (len(pgs_with_pid) == 0): return None variables = ['x', 'y', 'z', 'px', 'py', 'pz', 't'] trajectory = { var: np.zeros((len(pgs_with_pid), )) for var in variables } for ii, pg in enumerate(pgs_with_pid): for var in variables: trajectory[var][ii] = pg[var][pg['id'] == pid] return trajectory def run(self, gpt_verbose=False): """ performs a basic GPT simulation configured in the current GPT object """ if not self.configured: self.configure() #pass self.run_gpt(verbose=self.verbose, timeout=self.timeout, gpt_verbose=gpt_verbose) def get_run_script(self, write_to_path=True): """ Assembles the run script. Optionally writes a file 'run' with this line to path. """ _, infile = os.path.split(self.input_file) tokens = infile.split('.') if (len(tokens) > 1): outfile = '.'.join(tokens[:-1]) + '.out.gdf' else: outfile = tokens[0] + '.out.gdf' runscript = [ self.gpt_bin, '-j1', '-v', '-o', self.get_gpt_output_file(), self.input_file ] if write_to_path: with open(os.path.join(self.path, 'run'), 'w') as f: f.write(' '.join(runscript)) return runscript def get_gpt_output_file(self): """ get the name of the GPT output file """ path, infile = os.path.split(self.input_file) tokens = infile.split('.') if (len(tokens) > 1): outfile = '.'.join(tokens[:-1]) + '.out.gdf' else: outfile = tokens[0] + '.out.gdf' return os.path.join(path, outfile) def run_gpt(self, verbose=False, parse_output=True, timeout=None, gpt_verbose=False): """ RUN GPT and read in results """ self.vprint('GPT.run_gpt:') run_info = {} t1 = time() run_info['start_time'] = t1 if self.initial_particles: fname = self.write_initial_particles() #print(fname) # Link input file to new particle file self.set_dist_file(fname) #init_dir = os.getcwd() self.vprint(f' Running GPT...') # Write input file from internal dict self.write_input_file() runscript = self.get_run_script() self.vprint(f' Running with timeout = {self.timeout} sec.') run_time, exception, log = tools.execute(runscript, kill_msgs=self.kill_msgs, timeout=timeout, verbose=gpt_verbose) if (exception is not None): self.error = True run_info["error"] = True run_info['why_error'] = exception.strip() self.log = log if parse_output: self.load_output(file=self.get_gpt_output_file()) run_info['run_time'] = time() - t1 run_info['run_error'] = self.error self.vprint( f' Run finished, total time ellapsed: {run_info["run_time"]:G} (sec)' ) # Add run_info self.output.update(run_info) self.finished = True def fingerprint(self): """ Data fingerprint using the input. """ return tools.fingerprint(self.input) def vprint(self, *args, **kwargs): # Verbose print if self.verbose: print(*args, **kwargs) def plot(self, y=['sigma_x', 'sigma_y'], x='mean_z', xlim=None, y2=[], nice=True, include_layout=False, include_labels=False, include_particles=True, include_legend=True, **kwargs): """ Convenience plotting function for making nice plots. """ plot_stats_with_layout(self, ykeys=y, ykeys2=y2, xkey=x, xlim=xlim, nice=nice, include_layout=include_layout, include_labels=include_labels, include_legend=include_legend, **kwargs) def stat(self, key, data_type='all'): """ Calculates any statistic that the ParticleGroup class can calculate, on all particle groups, or just touts, or screens """ if (data_type == 'all'): particle_groups = self.output['particles'] elif (data_type == 'tout'): particle_groups = self.tout elif (data_type == 'tout_ccs'): particle_groups = self.tout_ccs elif (data_type == 'screen'): particle_groups = self.screen else: raise ValueError(f'Unsupported GPT data type: {data_type}') return particle_stats(particle_groups, key) def units(self, key): """ Calculates any statistic that the ParticleGroup class can calculate, on all particle groups. """ """Returns a str decribing the physical units of a stat key.""" return pg_units(key) def write_input_file(self): """ Write the updated GPT input file """ self.vprint(f' Writing gpt input file to "{self.input_file}"') parsers.write_gpt_input_file(self.input, self.input_file, self.ccs_beg) def write_initial_particles(self, fname=None): """ Write the initial particle data to file for use with GPT """ if not fname: fname = os.path.join(self.path, 'gpt.particles.gdf') self.initial_particles.write_gpt(fname, asci2gdf_bin='$ASCI2GDF_BIN', verbose=False) self.vprint(f' Initial particles written to "{fname}"') return fname def load_initial_particles(self, h5): """Loads a openPMD-beamphysics particle h5 handle or file""" P = ParticleGroup(h5=h5) self.initial_particles = P def load_archive(self, h5=None): """ Loads input and output from archived h5 file. See: GPT.archive """ if isinstance(h5, str): h5 = os.path.expandvars(h5) g = h5py.File(h5, 'r') glist = gpt.archive.find_gpt_archives(g) n = len(glist) if n == 0: # legacy: try top level message = 'legacy' elif n == 1: gname = glist[0] message = f'group {gname} from' g = g[gname] else: raise ValueError( f'Multiple archives found in file {h5}: {glist}') self.vprint(f'Reading {message} archive file {h5}') else: g = h5 self.input = gpt.archive.read_input_h5(g['input']) if 'initial_particles' in g: self.initial_particles = ParticleGroup(g['initial_particles']) self.output = gpt.archive.read_output_h5(g['output']) self.vprint( 'Loaded from archive. Note: Must reconfigure to run again.') self.configured = False def archive(self, h5=None): """ Archive all data to an h5 handle or filename. If no file is given, a file based on the fingerprint will be created. """ if not h5: h5 = 'gpt_' + self.fingerprint() + '.h5' if isinstance(h5, str): h5 = os.path.expandvars(h5) g = h5py.File(h5, 'w') self.vprint(f'Archiving to file {h5}') else: # store directly in the given h5 handle g = h5 # Write basic attributes gpt.archive.gpt_init(g) # All input gpt.archive.write_input_h5(g, self.input, name='input') if self.initial_particles: self.initial_particles.write(g, name='initial_particles') # All output gpt.archive.write_output_h5(g, self.output, name='output') return h5 @classmethod def from_archive(cls, archive_h5): """ Class method to return an GPT object loaded from an archive file """ c = cls() c.load_archive(archive_h5) return c def __str__(self): outstr = '\nGPT object:' if (self.configured): outstr = outstr + "\n Original input file: " + self.original_input_file outstr = outstr + "\n Template location: " + self.original_path if (self.workdir): outstr = outstr + "\n Top level work dir: " + self.workdir if (self.use_tempdir): outstr = outstr + f"\n Use temp directory: {self.use_tempdir}" #else: # outstr = outstr+f"\n Work directory: {self.path}" # Run control outstr = outstr + "\n\nRun Control" outstr = outstr + f"\n Run configured: {self.configured}" if (self.configured): outstr = outstr + f"\n Work location: {self.path}" outstr = outstr + f"\n Timeout: {self.timeout} (sec)" # Results outstr = outstr + "\n\nResults" outstr = outstr + f"\n Finished: {self.finished}" outstr = outstr + f"\n Error occured: {self.error}" if (self.error): outstr = outstr + f'\n Cause: {self.output["why_error"]}' errline = self.get_syntax_error_line(self.output["why_error"]) if (errline): outstr = outstr + f'\n Suspected input file line: "{errline}"' rtime = self.output['run_time'] outstr = outstr + f'\n Run time: {rtime} (sec)' #outstr = outstr+f"\n #outstr = outstr+f'\n Log: {self.log}\n' return outstr def get_syntax_error_line(self, error_msg): s = error_msg.strip().replace('\n', '') if (s.endswith('Error: syntax error')): error_line_index = int(s[s.find("(") + 1:s.find(")")]) return self.input['lines'][error_line_index] else: return None def track(self, particles, s=None, output='tout'): return track(self, particles, s=s, output=output) def track1(self, x0=0, px0=0, y0=0, py0=0, z0=0, pz0=1e-15, t0=0, s=None, species='electron', output='tout'): return track1(self, x0=x0, px0=px0, y0=y0, py0=py0, z0=z0, pz0=pz0, t0=t0, species=species, s=s, output=output) def track1_to_z(self, z_end, ds=0, ccs_beg='wcs', ccs_end='wcs', x0=0, px0=0, y0=0, py0=0, z0=0, pz0=1e-15, t0=0, weight=1, status=1, species='electron', s_screen=0): return track1_to_z(self, z_end=z_end, ds=ds, ccs_beg=ccs_beg, ccs_end=ccs_end, x0=x0, px0=px0, y0=y0, py0=py0, z0=z0, pz0=pz0, t0=t0, weight=weight, status=status, species=species, s_screen=s_screen) def track1_in_ccs(self, z_beg=0, z_end=0, ccs='wcs', x0=0, px0=0, y0=0, py0=0, pz0=1e-15, t0=0, weight=1, status=1, species='electron', xacc=6.5, GBacc=6.5, workdir=None, use_tempdir=True, n_screen=1, s_beg=0): return track1_in_ccs(self, z_beg=z_beg, z_end=z_end, ccs=ccs, x0=x0, px0=px0, y0=y0, py0=py0, pz0=pz0, t0=t0, weight=weight, status=status, species=species, xacc=xacc, GBacc=GBacc, workdir=workdir, use_tempdir=use_tempdir, n_screen=n_screen, s_beg=s_beg) def get_zminmax_line(self, z_beg, z_end, ccs='wcs'): return get_zminmax_line(self, z_beg, z_end, ccs=ccs) def copy(self): """ Returns a deep copy of this object. If a tempdir is being used, will clear this and deconfigure. """ G2 = deepcopy(self) # Clear this if G2.use_tempdir: G2.path = None G2.configured = False return G2
class Generator: """ This class defines the main run engine object for distgen and is responsible for 1. Parsing the input data dictionary passed from a Reader object 2. Check the input for internal consistency 3. Collect the parameters for distributions requested in the params dictionary 4. Form a the Beam object and populated the particle phase space coordinates """ def __init__(self, input=None, verbose=0): """ The class initialization takes in a verbose level for controlling text output to the user """ self.verbose = verbose self.input = input # This will be set with .beam() self.rands = None # This will be set with .run() self.particles = None if input: self.parse_input(input) self.configure() def parse_input(self, input): """ Parse the input structure passed from a Reader object. The structure is then converted to an easier form for use in populating the Beam object. YAML or JSON is accepted if params is a filename (str) Relative paths for input 'file' keys will be expanded. """ if isinstance(input, str): if os.path.exists(os.path.expandvars(input)): # File filename = full_path(input) with open(filename) as fid: input = yaml.safe_load(fid) # Fill any 'file' keys expand_input_filepaths(input, root=os.path.split(filename)[0], ignore_keys=['output']) else: #Try raw string input = yaml.safe_load(input) assert isinstance( input, dict ), f'ERROR: parsing unsuccessful, could not read {input}' expand_input_filepaths(input) self.input = input def configure(self): """ Configures the generator for creating a 6d particle distribution: 1. Copies the input dictionary read in from a file or passed directly 2. Converts physical quantities to PINT quantities in the params dictionary 3. Runs consistency checks on the resulting params dict """ self.params = copy.deepcopy(self.input) # Copy the input dictionary if ('start' not in self.params): self.params['start'] = {'type': 'free'} convert_params( self.params ) # Conversion of the input dictionary using tools.convert_params self.check_input_consistency( self.params) # Check that the result is logically sound def check_input_consistency(self, params): ''' Perform consistency checks on the user input data''' # Make sure all required top level params are present required_params = ['n_particle', 'random_type', 'total_charge'] for rp in required_params: assert rp in params, 'Required generator parameter ' + rp + ' not found.' # Check that only allowed params present at top level allowed_params = required_params + ['output', 'transforms', 'start'] for p in params: #assert p in allowed_params or '_dist'==p[-5:], 'Unexpected distgen input parameter: ' + p[-5:] assert p in allowed_params or p.endswith( '_dist'), 'Unexpected distgen input parameter: ' + p assert params[ 'n_particle'] > 0, 'User must speficy n_particle must > 0.' # Check consistency of transverse coordinate definitions if (("r_dist" in params) or ("x_dist" in params) or ("xy_dist" in params)): assert ("r_dist" in params) ^ ("x_dist" in params) ^ ( "xy_dist" in params ), "User must specify only one transverse distribution." if (("r_dist" in params) or ("y_dist" in params) or ("xy_dist" in params)): assert ("r_dist" in params) ^ ("y_dist" in params) ^ ( "xy_dist" in params), "User must specify r dist OR y dist NOT BOTH." if (params['start']['type'] == "cathode"): vprint("Ignoring user specified z distribution for cathode start.", self.verbose > 0 and "z_dist" in params, 0, True) vprint( "Ignoring user specified px distribution for cathode start.", self.verbose > 0 and "px_dist" in params, 0, True) vprint( "Ignoring user specified py distribution for cathode start.", self.verbose > 0 and "py_dist" in params, 0, True) vprint( "Ignoring user specified pz distribution for cathode start.", self.verbose > 0 and "pz_dist" in params, 0, True) assert "MTE" in params[ 'start'], "User must specify the MTE for cathode start." # Handle momentum distribution for cathode MTE = self.params['start']["MTE"] sigma_pxyz = (np.sqrt((MTE / MC2).to_reduced_units()) * unit_registry("GB")).to("eV/c") self.params["px_dist"] = {"type": "g", "sigma_px": sigma_pxyz} self.params["py_dist"] = {"type": "g", "sigma_py": sigma_pxyz} self.params["pz_dist"] = {"type": "g", "sigma_pz": sigma_pxyz} elif (params['start']['type'] == 'time'): vprint("Ignoring user specified t distribution for time start.", self.verbose > 0 and "t_dist" in params, 0, True) if ('t_dist' in params): warnings.warn( 'Ignoring user specified t distribution for time start.') self.params.pop('t_dist') if ('output' in self.params): out_params = self.params["output"] for op in out_params: assert op in ['file', 'type' ], f'Unexpected output parameter specified: {op}' else: self.params['output'] = {"type": None} def __getitem__(self, varstr): return get_nested_dict(self.input, varstr, sep=':', prefix='distgen') def __setitem__(self, varstr, val): return set_nested_dict(self.input, varstr, val, sep=':', prefix='distgen') def get_dist_params(self): """ Loops through the input params dict and collects all distribution definitions """ dist_vars = [ p.replace('_dist', '') for p in self.params if (p.endswith('_dist')) ] dist_params = { p.replace('_dist', ''): self.params[p] for p in self.params if (p.endswith('_dist')) } if ('r' in dist_vars and 'theta' not in dist_vars): vprint("Assuming cylindrical symmetry...", self.verbose > 0, 1, True) dist_params['theta'] = { 'type': 'ut', 'min_theta': 0 * unit_registry('rad'), 'max_theta': 2 * pi } if (self.params['start']['type'] == 'time' and 't_dist' in self.params): raise ValueError('Error: t_dist should not be set for time start') return dist_params def get_rands(self, variables): """ Gets random numbers [0,1] for the coordinatess in variables using either the Hammersley sequence or rand """ specials = ['xy'] self.rands = {var: None for var in variables if var not in specials} if ('xy' in variables): self.rands['x'] = None self.rands['y'] = None elif ('r' in variables and 'theta' not in variables): self.rands['theta'] = None n_coordinate = len(self.rands.keys()) n_particle = int(self.params['n_particle']) shape = (n_coordinate, n_particle) if (n_coordinate > 0): rns = random_generator(shape, sequence=self.params['random_type']) for ii, key in enumerate(self.rands.keys()): if (len(rns.shape) > 1): self.rands[key] = rns[ii, :] * unit_registry('dimensionless') else: self.rands[key] = rns[:] * unit_registry('dimensionless') var_list = list(self.rands.keys()) for ii, vii in enumerate(var_list[:-1]): viip1 = var_list[ii + 1] assert ( not np.array_equal(self.rands[vii].magnitude, self.rands[viip1].magnitude) ) or n_particle == 1, f'Error: coordinate probalitiies for {vii} and {viip1} are the same!' # These lines can be used to check for unwanted correlations #v0 = self.rands[vii].magnitude-self.rands[vii].magnitude.mean() #v1 = self.rands[viip1].magnitude-self.rands[viip1].magnitude.mean() #print( np.mean(v0*v1) ) def beam(self): """ Creates a 6d particle distribution and returns it in a distgen.beam class """ watch = StopWatch() watch.start() self.configure() verbose = self.verbose #outputfile = [] beam_params = { 'total_charge': self.params['total_charge'], 'n_particle': self.params['n_particle'] } if ('transforms' in self.params): transforms = self.params['transforms'] else: transforms = None #dist_params = {p.replace('_dist',''):self.params[p] for p in self.params if(p.endswith('_dist')) } #self.get_rands() vprint(f'Distribution format: {self.params["output"]["type"]}', self.verbose > 0, 0, True) N = int(self.params['n_particle']) bdist = Beam(**beam_params) if ("file" in self.params['output']): outfile = self.params['output']["file"] else: outfile = "None" vprint( f'Warning: no output file specified, defaulting to "{outfile}".', verbose > 0, 1, True) vprint(f'Output file: {outfile}', verbose > 0, 0, True) vprint('\nCreating beam distribution....', verbose > 0, 0, True) vprint(f"Beam starting from: {self.params['start']['type']}", verbose > 0, 1, True) vprint(f'Total charge: {bdist.q:G~P}.', verbose > 0, 1, True) vprint(f'Number of macroparticles: {N}.', verbose > 0, 1, True) units = { 'x': 'm', 'y': 'm', 'z': 'm', 'px': 'eV/c', 'py': 'eV/c', 'pz': 'eV/c', 't': 's' } # Initialize coordinates to zero for var, unit in units.items(): bdist[var] = np.full(N, 0.0) * unit_registry(units[var]) bdist["w"] = np.full((N, ), 1 / N) * unit_registry("dimensionless") avgs = {var: 0 * unit_registry(units[var]) for var in units} stds = {var: 0 * unit_registry(units[var]) for var in units} dist_params = self.get_dist_params( ) # Get the relevant dist params, setting defaults as needed, and samples random number generator self.get_rands(list(dist_params.keys())) # Do radial dist first if requested if ('r' in dist_params and 'theta' in dist_params): vprint('r distribution: ', verbose > 0, 1, False) # Get r distribution rdist = get_dist('r', dist_params['r'], verbose=verbose) if (rdist.rms() > 0): r = rdist.cdfinv( self.rands['r']) # Sample to get beam coordinates # Sample to get beam coordinates vprint('theta distribution: ', verbose > 0, 1, False) theta_dist = get_dist('theta', dist_params['theta'], verbose=verbose) theta = theta_dist.cdfinv(self.rands['theta']) rrms = rdist.rms() avgr = rdist.avg() avgCos = 0 avgSin = 0 avgCos2 = 0.5 avgSin2 = 0.5 bdist['x'] = r * np.cos(theta) bdist['y'] = r * np.sin(theta) avgs['x'] = avgr * avgCos avgs['y'] = avgr * avgSin stds['x'] = rrms * np.sqrt(avgCos2) stds['y'] = rrms * np.sqrt(avgSin2) # remove r, theta from list of distributions to sample del dist_params['r'] del dist_params['theta'] # Do 2D distributions if ("xy" in dist_params): vprint('xy distribution: ', verbose > 0, 1, False) dist = get_dist('xy', dist_params['xy'], verbose=verbose) bdist['x'], bdist['y'] = dist.cdfinv(self.rands['x'], self.rands['y']) dist_params.pop('xy') avgs['x'] = bdist.avg('x') avgs['y'] = bdist.avg('y') stds['x'] = bdist.std('x') stds['y'] = bdist.std('y') # Do all other specified single coordinate dists for x in dist_params.keys(): vprint(x + " distribution: ", verbose > 0, 1, False) dist = get_dist(x, dist_params[x], verbose=verbose) # Get distribution if (dist.std() > 0): # Only reach here if the distribution has > 0 size bdist[x] = dist.cdfinv( self.rands[x]) # Sample to get beam coordinates # Fix up the avg and std so they are exactly what user asked for if ("avg_" + x in dist_params[x]): avgs[x] = dist_params[x]["avg_" + x] else: avgs[x] = dist.avg() stds[x] = dist.std() #if("sigma_"+x in dist_params[x]): # stds[x] = dist_params[x]["sigma_"+x] #else: #stds[x] = dist.std() #print(x, stds[x]) # Shift and scale coordinates to undo sampling error for x in avgs: vprint(f'Shifting avg_{x} = {bdist.avg(x):G~P} -> {avgs[x]:G~P}', verbose > 0 and bdist[x].mean() != avgs[x], 1, True) vprint(f'Scaling sigma_{x} = {bdist.std(x):G~P} -> {stds[x]:G~P}', verbose > 0 and bdist[x].std() != stds[x], 1, True) #bdist = transform(bdist, {'type':f'set_avg_and_std {x}', 'avg_'+x:avgs[x],'sigma_'+x:stds[x], 'verbose':0}) bdist = set_avg_and_std( bdist, **{ 'variables': x, 'avg_' + x: avgs[x], 'sigma_' + x: stds[x], 'verbose': 0 }) # Handle any start type specific settings if (self.params['start']['type'] == "cathode"): bdist['pz'] = np.abs(bdist['pz']) # Only take forward hemisphere vprint('Cathode start: fixing pz momenta to forward hemisphere', verbose > 0, 1, True) vprint( f'avg_pz -> {bdist.avg("pz"):G~P}, sigma_pz -> {bdist.std("pz"):G~P}', verbose > 0, 2, True) elif (self.params['start']['type'] == 'time'): if ('tstart' in self.params['start']): tstart = self.params['start']['tstart'] else: vprint( "Time start: no start time specified, defaulting to 0 sec.", verbose > 0, 1, True) tstart = 0 * unit_registry('sec') vprint( f'Time start: fixing all particle time values to start time: {tstart:G~P}.', verbose > 0, 1, True) bdist = set_avg( bdist, **{ 'variables': 't', 'avg_t': 0.0 * unit_registry('sec'), 'verbose': verbose > 0 }) elif (self.params['start']['type'] == 'free'): pass else: raise ValueError( f'Beam start type "{self.params["start"]["type"]}" is not supported!' ) # Apply any user desired coordinate transformations if (transforms): # Check if the user supplied the transform order, otherwise just go through the dictionary if ('order' in transforms): order = transforms['order'] if (not isinstance(order, list)): raise ValueError( 'Transform "order" key must be associated a list of transform IDs' ) del transforms['order'] else: order = transforms.keys() for name in order: T = transforms[name] T['verbose'] = verbose > 0 vprint( f'Applying user supplied transform: "{name}" = {T["type"]}...', verbose > 0, 1, True) bdist = transform(bdist, T) watch.stop() vprint(f'...done. Time Ellapsed: {watch.print()}.\n', verbose > 0, 0, True) return bdist def run(self): """ Runs the generator.beam function stores the partice in an openPMD-beamphysics ParticleGroup in self.particles """ beam = self.beam() self.particles = ParticleGroup(data=beam.data()) vprint(f'Created particles in .particles: \n {self.particles}', self.verbose > 0, 1, False) def fingerprint(self): """ Data fingerprint using the input. """ return fingerprint(self.input) def load_archive(self, h5=None): """ Loads input and output from archived h5 file. See: Generator.archive """ if isinstance(h5, str): g = h5py.File(h5, 'r') glist = archive.find_distgen_archives(g) n = len(glist) if n == 0: # legacy: try top level message = 'legacy' elif n == 1: gname = glist[0] message = f'group {gname} from' g = g[gname] else: raise ValueError( f'Multiple archives found in file {h5}: {glist}') vprint(f'Reading {message} archive file {h5}', self.verbose > 0, 1, False) else: g = h5 vprint(f'Reading Distgen archive file {h5}', self.verbose > 0, 1, False) self.input = archive.read_input_h5(g['input']) if 'particles' in g: self.particles = ParticleGroup(g['particles']) else: vprint(f'No particles found.', self.verbose > 0, 1, False) def archive(self, h5=None): """ Archive all data to an h5 handle or filename. If no file is given, a file based on the fingerprint will be created. """ if not h5: h5 = 'distgen_' + self.fingerprint() + '.h5' if isinstance(h5, str): g = h5py.File(h5, 'w') # Proper openPMD init pmd_init(g, basePath='/', particlesPath='particles/') g.attrs['software'] = np.string_('distgen') # makes a fixed string #TODO: add version: g.attrs('version') = np.string_(__version__) else: g = h5 # Init archive.distgen_init(g) # Input archive.write_input_h5(g, self.input, name='input') # Particles if self.particles: self.particles.write(g, name='particles') return h5 def __repr__(self): s = '<disgten.Generator with input: \n' return s + yaml.dump(self.input) + '\n>' def check_inputs(self, params): """ Checks the params sent to the generator only contain allowed inputs """ # Make sure user isn't passing the wrong parameters: allowed_params = self.optional_params + self.required_params + [ 'verbose' ] for param in params: assert param in allowed_params, 'Incorrect param given to ' + self.__class__.__name__ + '.__init__(**kwargs): ' + param + '\nAllowed params: ' + str( allowed_params) # Make sure all required parameters are specified for req in self.required_params: assert req in params, 'Required input parameter ' + req + ' to ' + self.__class__.__name__ + '.__init__(**kwargs) was not found.'