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 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: 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 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 """
beam = self.beam()
self.particles = ParticleGroup(data=beam.data())
vprint(f'Created particles in .particles: \n {self.particles}',
self.verbose > 0, 1, False)
def read_particles_h5(h5):
"""
Reads particles from h5
"""
dat = {}
for g in h5:
dat[g] = ParticleGroup(h5=h5[g])
return dat
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 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 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 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 load_archive(self, h5, configure=True):
"""
Loads input and output from archived h5 file.
See: Impact.archive
"""
if isinstance(h5, str):
fname = os.path.expandvars(h5)
g = h5py.File(fname, 'r')
glist = archive.find_impact_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 {fname}: {glist}')
self.vprint(f'Reading {message} archive file {h5}')
else:
g = h5
self.input = archive.read_input_h5(g['input'], verbose=self.verbose)
self.output, self._units = archive.read_output_h5(g['output'],
verbose=self.verbose)
if 'initial_particles' in g:
self.initial_particles = ParticleGroup(h5=g['initial_particles'])
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
if configure:
self.configure()
# Re-link groups
# TODO: cleaner logic
for _, cg in self.group.items():
cg.link(self.ele)
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 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 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))
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 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 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 load_initial_particles(self, h5):
"""Loads a openPMD-beamphysics particle h5 handle or file"""
P = ParticleGroup(h5=h5)
self.initial_particles = P
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
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)
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]