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'])
self.output = self.particles
else:
vprint('No particles found.', self.verbose>0,1,False)
def raw_data_to_particle_groups(touts, screens, verbose=False, ref_ccs=False):
"""
Coverts a list of touts to a list of ParticleGroup objects
"""
if (verbose):
print(' Converting tout and screen data to ParticleGroup(s)')
if (ref_ccs):
pg_touts = [
ParticleGroup(data=raw_data_to_particle_data(datum))
for datum in touts
]
pg_screens = [
ParticleGroup(data=raw_data_to_particle_data(datum))
for datum in screens
]
new_touts = [
transform_to_centroid_coordinates(tout) for tout in pg_touts
]
return new_touts + pg_screens
else:
return [
ParticleGroup(data=raw_data_to_particle_data(datum))
for datum in touts + screens
]
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 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 run(self):
""" Runs the generator.beam function stores the partice in
an openPMD-beamphysics ParticleGroup in self.particles """
if self.input is not None:
beam = self.beam()
self.particles = ParticleGroup(data=beam.data())
self.output = self.particles
vprint(f'Created particles in .particles: \n {self.particles}',
self.verbose > 0, 1, False)
else:
print('No input data specified.')
return self.output
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 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 read_particles_h5(h5):
"""
Reads particles from h5
"""
dat = {}
for g in h5:
dat[g] = ParticleGroup(h5=h5[g])
return dat
def load_output(self):
pfile = self.output_file
data = parsers.parse_astra_phase_file(pfile)
# Clock time is used when at cathode
data['t'] = data['t_clock']
P = ParticleGroup(data=data)
self.output['particles'] = P
def raw_data_to_particle_groups(touts, screens):
"""
Coverts a list of touts to a list of ParticleGroup objects
"""
return [
ParticleGroup(data=raw_data_to_particle_data(datum))
for datum in touts + screens
]
def run(self):
""" Runs the generator.beam function stores the partice in
an openPMD-beamphysics ParticleGroup in self.particles """
if self.input is not None:
beam = self.beam()
if self.params['start']['type'] == "cathode":
status = ParticleStatus.CATHODE
else:
status = ParticleStatus.ALIVE
self.particles = ParticleGroup(data=beam.data(status=status))
self.output = self.particles
vprint(f'Created particles in .particles: \n {self.particles}', self.verbose > 0, 1, False)
else:
print('No input data specified.')
return self.output
def read_particles_h5(h5):
"""
Reads particles from h5
See: write_particles_h5
"""
# This should be a list of '0', '1', etc.
# Cast to int, sort, reform to get the list order correct.
ilist = sorted([int(x) for x in list(h5)])
glist = [str(i) for i in ilist]
return [ParticleGroup(h5=h5[g]) for g in glist]
def load_particles(self):
# Standard output
self.vprint('Loading particles')
self.output['particles'] = load_many_fort(self.path,
FORT_PARTICLE_TYPES,
verbose=self.verbose)
# Additional particle files:
for e in self.input['lattice']:
if e['type'] == 'write_beam':
name = e['name']
fname = e['filename']
full_fname = os.path.join(self.path, fname)
if os.path.exists(full_fname):
self.particles[name] = parse_impact_particles(full_fname)
self.vprint(f'Loaded write beam particles {name} {fname}')
# Convert all to ParticleGroup
# Interpolate stats to get the time.
time_f = interp1d(self.output['stats']['mean_z'],
self.output['stats']['t'],
assume_sorted=True,
fill_value='extrapolate')
for name, pdata in self.particles.items():
# Initial particles have special z = beta_ref*c. See: impact_particles_to_particle_data
if name == 'initial_particles' and self.header['Flagimg']:
cathode_kinetic_energy_ref = self.header['Bkenergy']
else:
cathode_kinetic_energy_ref = None
time = time_f(pdata['z'].mean())
pg_data = impact_particles_to_particle_data(
pdata,
mc2=self.mc2,
species=self.species,
time=time,
macrocharge=self.macrocharge,
cathode_kinetic_energy_ref=cathode_kinetic_energy_ref,
verbose=self.verbose)
self.particles[name] = ParticleGroup(data=pg_data)
self.vprint(f'Converted {name} to ParticleGroup')
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 initial_beam_to_particle_group(
gdffile,
verbose=0,
extra_screen_keys=['q', 'nmacro', 'ID', 'm'],
missing_data=None):
screen = read_particle_gdf_file(gdffile,
verbose=verbose,
extra_screen_keys=extra_screen_keys)
if (missing_data is not None):
for mdatum in missing_data:
if (mdatum not in screen.keys()
and len(missing_data[mdatum]) == len(screen['x'])):
screen[mdatum] = missing_data[mdatum]
return ParticleGroup(data=raw_data_to_particle_data(screen))
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
def load_initial_particles(self, h5):
"""Loads a openPMD-beamphysics particle h5 handle or file"""
P = ParticleGroup(h5=h5)
self.initial_particles = P
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
def read_lucretia(filename,
ele_name='',
t_ref=0,
kill_dead_particles=True,
verbose=False):
"""
Lucretia's format is described in:
https://www.slac.stanford.edu/accel/ilc/codes/Lucretia/web/beam.html
One Lucretia ".mat" file can include beams at multiple lattice elements.
To find the beam at one element, one has to follow down this order of "fields":
bstore >> ele_name >> Bunch >> x,
in which x is a 6-to-Np array with:
Lucretia x = x in m
Lucretia px = px/p in radian
Lucretia y = y in m
Lucretia py = py/p in radian
Lucretia z = (t - t_ref)*c in m
Lucretia p = p in GeV/c
Note that p is the total, not reference, momentum.
t_ref is zero by default.
To access valid element names in a Lucretia beam file, do:
dat = sio.loadmat('filename.mat')
print(dat['bstore'].dtype)
"""
mdat = sio.loadmat(filename)
coords = mdat['bstore'][ele_name][0, 0]['Bunch'][0, 0]['x'][0, 0]
charges = mdat['bstore'][ele_name][0, 0]['Bunch'][0, 0]['Q'][0, 0]
Np = coords.shape[1]
x = coords[0]
px_luc = coords[1] # normalized by total momentum
y = coords[2]
py_luc = coords[3] # normalized by total momentum
z_luc = coords[4]
ptot = coords[5] # total momentum in GeV/c
px = px_luc * ptot * 1E9 # in eV/c
py = py_luc * ptot * 1E9
pz = np.sqrt((ptot * 1E9)**2 - px**2 - py**2)
t = z_luc / 299792458 + t_ref
status = np.ones(Np)
ix = np.where(ptot == 0)
status[ix] = 0
n_dead = len(ix[0])
if verbose:
print(Np, 'particles detected,', n_dead, 'found dead!')
data = {
'x': x,
'px': px,
'y': y,
'py': py,
'z': np.zeros(Np),
'pz': pz,
't': t,
'status': status,
'weight': charges,
'species': 'electron'
}
P = ParticleGroup(data=data)
if (kill_dead_particles):
if verbose:
print('Excluding dead particles (if any)...')
P = P.where(P.p > 0)
return P
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.'
def run(self, inputs, verbose=False):
tag = f'vb24@{self.id}:'
#----------------------------------------------------------------------------
# Get laser distribution, cathode quantities, and gun current
#----------------------------------------------------------------------------
r_params = {'sigma_xy': dunits(str(inputs[f'{tag}laser:sigma_xy'])),
'alpha': dunits(str(inputs[f'{tag}laser:alpha_xy']))}
count = self.pvdefs[f'{tag}laser:r']['count']
laser_wavelength = inputs[f'{tag}laser:wavelength']
laser_power = inputs[f'{tag}laser:power']
laser_sigma_xy = inputs[f'{tag}laser:sigma_xy']
laser_alpha_xy = inputs[f'{tag}laser:alpha_xy']
laser_avg_x = inputs[f'{tag}laser:mean_x']
laser_avg_y = inputs[f'{tag}laser:mean_y']
r_dist = SuperGaussianRad(verbose=False, **r_params)
rs = (r_dist.get_r_pts(count)).to(self.pvdefs[f'{tag}laser:r']['unit'])
Pr = (dunits(str(laser_power))*r_dist.rho(rs)).to(self.pvdefs[f'{tag}laser:Pr']['unit'])
cathode_QE = inputs[f'{tag}cathode:QE']
cathode_MTE = inputs[f'{tag}cathode:MTE']
hc = 1*units.h*units.c
photon_flux = (laser_power/(hc/laser_wavelength) ).to_base_units()
gun_current = (photon_flux*cathode_QE*(1*units.e)).to(self.pvdefs[f'{tag}gun:current']['unit'])
#----------------------------------------------------------------------------
#----------------------------------------------------------------------------
# Create Distgen input and run generator
#----------------------------------------------------------------------------
distgen_input = yaml.dump(
{'n_particle':inputs[f'{tag}gpt:n_particle'].magnitude,
'random_type':'hammersley',
'total_charge': {'value': 0.0, 'units': 'pC'},
'start': {
'type':'cathode',
'MTE': {'value': cathode_MTE.magnitude, 'units': str(cathode_MTE.units)}},
'r_dist': {
'type':'rsg',
'sigma_xy':{'value': laser_sigma_xy.magnitude, 'units': str(laser_sigma_xy.units)},
'alpha':{'value': laser_alpha_xy.magnitude, 'units': str(laser_alpha_xy.units)},},
'transforms':{
't1':{'type':'set_avg x', 'avg_x': {'value': laser_avg_x.magnitude, 'units': str(laser_avg_x.units)}},
't2':{'type':'set_avg y', 'avg_y': {'value': laser_avg_y.magnitude, 'units': str(laser_avg_y.units)}}
}})
gen = Generator(distgen_input, verbose=True)
beam = gen.beam()
#----------------------------------------------------------------------------
#----------------------------------------------------------------------------
# Configure GPT and run
#----------------------------------------------------------------------------
G = GPT(input_file=os.path.join(os.getcwd(),'templates/gpt.in'),
initial_particles = ParticleGroup(data=beam.data()),
use_tempdir=True,
workdir=os.path.join(os.getcwd(),'tmp'),
timeout = 5,
verbose=True)
settings = {'gun_voltage': inputs[f'{tag}gun:voltage'].magnitude,
'sol01_current': inputs[f'{tag}sol1:current'].magnitude,
'sol02_current': inputs[f'{tag}sol2:current'].magnitude,
'npts': inputs[f'{tag}gpt:n_screen'].magnitude+1}
result = G.set_variables(settings)
G.run()
#----------------------------------------------------------------------------
#----------------------------------------------------------------------------
# Load all relevant data into output structure
#----------------------------------------------------------------------------
# laser distribution
output = {f'{tag}laser:r':rs.magnitude, f'{tag}laser:Pr':Pr.magnitude, f'{tag}gun:current':gun_current.magnitude}
# GPT statistical data
stats = {'max':['r'], 'mean':['x', 'y', 'z', 'kinetic_energy'], 'sigma':['x','y']}
for stat, variables in stats.items():
output = {**output, **{f'{tag}beam:{stat}_{var}': self.gpt_stat_to_pv(G, f'{stat}_{var}', 'screen').magnitude for var in variables} }
scr_numbers = [1]
for scr_number in scr_numbers:
z = inputs[f'{tag}scr{scr_number}:mean_z'].magnitude
for var in ['x' ,'y']:
output[f'{tag}scr{scr_number}:mean_{var}'] = np.interp(z, output[f'{tag}beam:mean_z'], output[f'{tag}beam:mean_{var}'])
output[f'{tag}scr{scr_number}:sigma_{var}'] = np.interp(z, output[f'{tag}beam:mean_z'], output[f'{tag}beam:sigma_{var}'])
# transmission
output[f'{tag}beam:transmission'] = [100*len(screen['x'])/inputs[f'{tag}gpt:n_particle'].magnitude for screen in G.screen]
min_clearance = np.min( (inputs[f'{tag}beampipe:radius']-self.gpt_stat_to_pv(G, f'{stat}_{var}', 'screen') ) ).to('mm')
output[f'{tag}beam:radiation'] = output[f'{tag}gun:current']*np.max(output[f'{tag}beam:mean_kinetic_energy'])/min_clearance.magnitude
#----------------------------------------------------------------------------
return output
def read_lucretia(filename,
ele_name='BEGINNING',
t_ref=0,
exclude_dead_particles=True,
verbose=False):
"""
Load one beam in a Lucretia beam file into a ParticleGroup
Parameters:
----------
filename : str
Lucretia '.mat' file name.
ele_name : str
name of the element at which the beam is located.
An invalid name results in an error.
If the beam file has one element, this only one beam is read.
Default: 'BEGINNING'
t_ref : float, optional
reference time of the beam in seconds. Default: 0.
exclude_dead_particles : bool, optional
if True, excludes dead particles. Default: True.
----------
Returns:
----------
A ParticleGroup object
----------
Lucretia's format is described in:
https://www.slac.stanford.edu/accel/ilc/codes/Lucretia/web/beam.html
One Lucretia ".mat" file can include beams at multiple lattice elements.
To find the beam at one element, one has to follow down this order of "fields":
bstore >> ele_name >> Bunch >> x,
in which x is a 6-to-Np array with:
Lucretia x = x in m
Lucretia px = px/p in radian
Lucretia y = y in m
Lucretia py = py/p in radian
Lucretia z = (t - t_ref)*c in m
Lucretia p = p in GeV/c
Note that p is the total, not reference, momentum.
To access valid element names in a Lucretia beam file,
use the helper function list_element_names(filename).
dat = sio.loadmat('filename.mat')
print(dat['bstore'].dtype)
"""
ele_list = list_element_names(filename)
if verbose:
print(len(ele_list), 'elements found in the file!')
# Check if the element exists
if (ele_name not in ele_list):
raise ValueError('The provided element name ' + str(ele_name) +
' does not exist in the file!')
elif (len(ele_list) == 1):
ele_name = ele_list[0]
mdat = sio.loadmat(filename)
coords = mdat['bstore'][ele_name][0, 0]['Bunch'][0, 0]['x'][0, 0]
charges = mdat['bstore'][ele_name][0, 0]['Bunch'][0, 0]['Q'][0, 0]
Np = coords.shape[1]
x = coords[0]
px_luc = coords[1] # normalized by total momentum
y = coords[2]
py_luc = coords[3] # normalized by total momentum
z_luc = coords[4]
ptot = coords[5] # total momentum in GeV/c
px = px_luc * ptot * 1E9 # in eV/c
py = py_luc * ptot * 1E9
pz = np.sqrt((ptot * 1E9)**2 - px**2 - py**2)
t = z_luc / 299792458 + t_ref
status = np.ones(Np)
ix = np.where(ptot == 0)
status[ix] = 0
n_dead = len(ix[0])
if verbose:
print(Np, 'particles detected,', n_dead, 'found dead!')
data = {
'x': x,
'px': px,
'y': y,
'py': py,
'z': np.zeros(Np),
'pz': pz,
't': t,
'status': status,
'weight': charges,
'species': 'electron'
}
P = ParticleGroup(data=data)
if (exclude_dead_particles):
if verbose:
print('Excluding dead particles (if any)...')
P = P.where(P.p > 0)
return P
def touts_to_particlegroups(touts):
"""
Coverts a list of touts to a list of ParticleGroup objects
"""
return [ParticleGroup(data=tout_to_particle_data(tout)) for tout in touts]
def bmad_to_particle_group(bmad_beam,
p0c=None,
charges=None,
t_ref=0,
verbose=False):
"""
Converts a bmad beam to a particle group.
Assumes electrons.
Parameters:
----------
bmad_beam : float, array
a 2D array of size (6, number_of_particles)
p0c: float,
reference momentum in eV/c
t_ref : float, optional
reference time of the beam in seconds. Default: 0.
charges : float, array
an 1D array of size (number_of_particles)
----------
Returns:
----------
A ParticleGroup object
----------
Bmad's ASCII format is described in:
https://www.classe.cornell.edu/bmad/manual.html
Bmad normally uses s-based coordinates, with momenta:
bmad px = px/p0
bmad py = py/p0
bmad pz = p/p0 - 1
and longitudinal coordinate
bmad z = -beta*c(t - t_ref)
"""
if (p0c <= 0):
raise ValueError(' invalid p0c value given!! ')
if (np.any(charges <= 0)):
raise ValueError(' invalid charges value(s) given!! ')
Np = bmad_beam.shape[1] # Number of macro particles
delta = bmad_beam[5]
p = p0c * (1 + delta) # in eV/c
E = np.sqrt(p**2 + 510998.950**2) # in eV, assuming electrons
beta = p / E
x = bmad_beam[0] # in m
px = bmad_beam[1] * p0c # in eV/c
y = bmad_beam[2] # in m
py = bmad_beam[3] * p0c # in eV/c
t = bmad_beam[4] / (-1. * beta * 299792458) + t_ref
pz = np.sqrt(p**2 - px**2 - py**2)
data = {
'x': x,
'px': px,
'y': y,
'py': py,
'z': np.zeros(Np),
'pz': pz,
't': t,
'status': np.ones(Np),
'weight': charges,
'species': 'electron'
}
return ParticleGroup(data=data)
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.'
)
