def track_on_axis(element,
t,
p,
xacc=6.5,
GBacc=12,
dtmin=1e-15,
dtmax=1e-8,
n_screen=1,
workdir=None):
if (workdir is None):
tempdir = tempfile.TemporaryDirectory(dir=workdir)
gpt_file = os.path.join(tempdir.name, f'{element.name}.gpt.in')
workdir = tempdir.name
else:
gpt_file = os.path.join(workdir, f'{element.name}.gpt.in')
element.write_element_to_gpt_file(ztrack1_template(gpt_file))
G = GPT(gpt_file,
ccs_beg=element.ccs_beg,
workdir=workdir,
use_tempdir=False)
G.set_variables({
'xacc': xacc,
'GBacc': GBacc,
'dtmin': dtmin,
'dtmax': dtmax
})
z_beg = np.linalg.norm(element._ccs_beg_origin - element._p_beg)
G.track1_in_ccs(z_beg=z_beg,
z_end=z_beg + element.length,
ccs=element.ccs_beg,
pz0=p,
t0=t,
weight=1,
status=1,
species='electron',
xacc=xacc,
GBacc=GBacc,
n_screen=n_screen)
return G
def run_gpt_with_particlegroup(settings=None,
gpt_input_file=None,
input_particle_group=None,
workdir=None,
use_tempdir=True,
gpt_bin='$GPT_BIN',
timeout=2500,
auto_phase=False,
verbose=False,
gpt_verbose=False,
asci2gdf_bin='$ASCI2GDF_BIN'):
"""
Run gpt with particles from ParticleGroup.
settings: dict with keys that are in gpt input file.
"""
# Call simpler evaluation if there is no input_particle_group:
if (input_particle_group is None):
return run_gpt(settings=settings,
gpt_input_file=gpt_input_file,
workdir=workdir,
use_tempdir=use_tempdir,
gpt_bin=gpt_bin,
timeout=timeout,
verbose=verbose)
if (verbose):
print('Run GPT with ParticleGroup:')
unit_registry = UnitRegistry()
# Make gpt and generator objects
G = GPT(gpt_bin=gpt_bin,
input_file=gpt_input_file,
initial_particles=input_particle_group,
workdir=workdir,
use_tempdir=use_tempdir)
G.timeout = timeout
G.verbose = verbose
# Set inputs
if settings:
for k, v in settings.items():
G.set_variable(k, v)
if ('final_charge' in settings):
raise ValueError(
'final_charge is deprecated, please specify value and units instead.'
)
# Run
if (auto_phase):
if (verbose):
print('\nAuto Phasing >------\n')
t1 = time.time()
# Create the distribution used for phasing
if (verbose):
print('****> Creating initial distribution for phasing...')
phasing_beam = get_distgen_beam_for_phasing_from_particlegroup(
input_particle_group, n_particle=10, verbose=verbose)
phasing_particle_file = os.path.join(G.path,
'gpt_particles.phasing.gdf')
write_gpt(phasing_beam,
phasing_particle_file,
verbose=verbose,
asci2gdf_bin=asci2gdf_bin)
if (verbose):
print('<**** Created initial distribution for phasing.\n')
G.write_input_file() # Write the unphased input file
phased_file_name, phased_settings = gpt_phasing(
G.input_file,
path_to_gpt_bin=G.gpt_bin[:-3],
path_to_phasing_dist=phasing_particle_file,
verbose=verbose)
G.set_variables(phased_settings)
t2 = time.time()
if (verbose):
print(f'Time Ellapsed: {t2-t1} sec.')
print('------< Auto Phasing\n')
# If here, either phasing successful, or no phasing requested
G.run(gpt_verbose=gpt_verbose)
if ('final_charge:value' in settings and 'final_charge:units' in settings
and len(G.screen) > 0):
final_charge = settings[
'final_charge:value'] * unit_registry.parse_expression(
settings['final_charge:units'])
final_charge = final_charge.to('coulomb').magnitude
clip_to_charge(G.screen[-1], final_charge, make_copy=False)
if ('final_radius:value' in settings and 'final_radius:units' in settings
and len(G.screen) > 0):
final_radius = settings[
'final_radius:value'] * unit_registry.parse_expression(
settings['final_radius:units'])
final_radius = final_radius.to('meter').magnitude
take_range(G.screen[-1], 'r', 0, final_radius)
if (input_particle_group['sigma_t'] == 0.0):
# Initial distribution is a tout
if (G.output['n_tout'] > 0):
G.output['particles'].insert(0, input_particle_group)
G.output['n_tout'] = G.output['n_tout'] + 1
else:
# Initial distribution is a screen
if (G.output['n_screen'] > 0):
G.output['particles'].insert(G.output['n_tout'],
input_particle_group)
G.output['n_screen'] = G.output['n_screen'] + 1
return G
def multirun_gpt_with_particlegroup(settings=None,
gpt_input_file=None,
input_particle_group=None,
workdir=None,
use_tempdir=True,
gpt_bin='$GPT_BIN',
timeout=2500,
auto_phase=False,
verbose=False,
gpt_verbose=False,
asci2gdf_bin='$ASCI2GDF_BIN'):
"""
Run gpt with particles from ParticleGroup.
settings: dict with keys that are in gpt input file.
"""
unit_registry = UnitRegistry()
# Call simpler evaluation if there is no input_particle_group:
if (input_particle_group is None):
raise ValueError('Must supply input_particle_group')
if (verbose):
print('Run GPT with ParticleGroup:')
if ('clipping_charge' in settings):
raise ValueError(
'clipping_charge is deprecated, please specify value and units instead.'
)
if ('final_charge' in settings):
raise ValueError(
'final_charge is deprecated, please specify value and units instead.'
)
if ('t_restart' not in settings):
raise ValueError('t_restart must be supplied')
t_restart = settings['t_restart']
if ('restart_file' not in settings):
# Make gpt and generator objects
G = GPT(gpt_bin=gpt_bin,
input_file=gpt_input_file,
initial_particles=input_particle_group,
workdir=workdir,
use_tempdir=use_tempdir)
G.timeout = timeout
G.verbose = verbose
# Set inputs
if settings:
for k, v in settings.items():
G.set_variable(k, v)
else:
raise ValueError('Must supply settings')
G.set_variable('multi_run', 0)
if (auto_phase):
if (verbose):
print('\nAuto Phasing >------\n')
t1 = time.time()
# Create the distribution used for phasing
if (verbose):
print('****> Creating initial distribution for phasing...')
phasing_beam = get_distgen_beam_for_phasing_from_particlegroup(
input_particle_group, n_particle=10, verbose=verbose)
phasing_particle_file = os.path.join(G.path,
'gpt_particles.phasing.gdf')
write_gpt(phasing_beam,
phasing_particle_file,
verbose=verbose,
asci2gdf_bin=asci2gdf_bin)
if (verbose):
print('<**** Created initial distribution for phasing.\n')
G.write_input_file() # Write the unphased input file
phased_file_name, phased_settings = gpt_phasing(
G.input_file,
path_to_gpt_bin=G.gpt_bin[:-3],
path_to_phasing_dist=phasing_particle_file,
verbose=verbose)
G.set_variables(phased_settings)
t2 = time.time()
if (verbose):
print(f'Time Ellapsed: {t2-t1} sec.')
print('------< Auto Phasing\n')
G.set_variable('multi_run', 1)
G.set_variable('last_run', 2)
G.set_variable('t_start', 0.0)
G.set_variable('t_restart', t_restart)
# If here, either phasing successful, or no phasing requested
G.run(gpt_verbose=gpt_verbose)
else:
G = GPT()
G.load_archive(settings['restart_file'])
if settings:
for k, v in settings.items():
G.set_variable(k, v)
# Remove touts and screens that are after t_restart
t_restart_with_fudge = t_restart + 1.0e-18 # slightly larger that t_restart to avoid floating point comparison problem
G.output['n_tout'] = np.count_nonzero(
G.stat('mean_t', 'tout') <= t_restart_with_fudge)
G.output['n_screen'] = np.count_nonzero(
G.stat('mean_t', 'screen') <= t_restart_with_fudge)
for p in reversed(G.particles):
if (p['mean_t'] > t_restart_with_fudge):
G.particles.remove(p)
G_all = G # rename it, and then overwrite G
if (verbose):
print(f'Looking for tout at t = {t_restart}')
restart_particles = get_screen_data(G,
tout_t=t_restart,
use_extension=False,
verbose=verbose)[0]
if ('clipping_charge:value' in settings
and 'clipping_charge:units' in settings):
clipping_charge = settings[
'clipping_charge:value'] * unit_registry.parse_expression(
settings['clipping_charge:units'])
clipping_charge = clipping_charge.to('coulomb').magnitude
clip_to_charge(restart_particles, clipping_charge, make_copy=False)
G = GPT(gpt_bin=gpt_bin,
input_file=gpt_input_file,
initial_particles=restart_particles,
workdir=workdir,
use_tempdir=use_tempdir)
G.timeout = timeout
G.verbose = verbose
for k, v in G_all.input["variables"].items():
G.set_variable(k, v)
G.set_variable('multi_run', 2)
G.set_variable('last_run', 2)
G.set_variable('t_start', t_restart)
G.run(gpt_verbose=gpt_verbose)
G_all.output['particles'][G_all.output['n_tout']:G_all.
output['n_tout']] = G.tout
G_all.output['particles'] = G_all.output['particles'] + G.screen
G_all.output['n_tout'] = G_all.output['n_tout'] + G.output['n_tout']
G_all.output['n_screen'] = G_all.output['n_screen'] + G.output['n_screen']
if ('final_charge:value' in settings and 'final_charge:units' in settings
and len(G_all.screen) > 0):
final_charge = settings[
'final_charge:value'] * unit_registry.parse_expression(
settings['final_charge:units'])
final_charge = final_charge.to('coulomb').magnitude
clip_to_charge(G_all.screen[-1], final_charge, make_copy=False)
if (input_particle_group['sigma_t'] == 0.0):
# Initial distribution is a tout
if (G_all.output['n_tout'] > 0):
# Don't include the cathode if there are no other screens. Screws up optimizations of "final" screen when there is an error
G_all.output['particles'].insert(0, input_particle_group)
G_all.output['n_tout'] = G_all.output['n_tout'] + 1
else:
# Initial distribution is a screen
if (G_all.output['n_screen'] > 0):
# Don't include the cathode if there are no other screens. Screws up optimizations of "final" screen when there is an error
G_all.output['particles'].insert(G_all.output['n_tout'],
input_particle_group)
G_all.output['n_screen'] = G_all.output['n_screen'] + 1
return G_all
def run_gpt_with_distgen(settings=None,
gpt_input_file=None,
distgen_input_file=None,
workdir=None,
use_tempdir=True,
gpt_bin='$GPT_BIN',
timeout=2500,
auto_phase=False,
verbose=False,
gpt_verbose=False,
asci2gdf_bin='$ASCI2GDF_BIN'):
"""
Run gpt with particles generated by distgen.
settings: dict with keys that can appear in an gpt or distgen Generator input file.
Example usage:
G = run_gpt_with_distgen({'lspch':False},
gpt_input_file='$LCLS_LATTICE/gpt/models/gunb_eic/gpt.in',
distgen_input_file='$LCLS_LATTICE/distgen/models/gunb_gaussian/gunb_gaussian.json',
verbose=True,
timeout=None
)
"""
# Call simpler evaluation if there is no generator:
if not distgen_input_file:
return run_gpt(settings=settings,
gpt_input_file=gpt_input_file,
workdir=workdir,
use_tempdir=use_tempdir,
gpt_bin=gpt_bin,
timeout=timeout,
verbose=verbose)
if (verbose):
print('Run GPT with Distgen:')
# Make gpt and generator objects
G = GPT(gpt_bin=gpt_bin,
input_file=gpt_input_file,
workdir=workdir,
use_tempdir=use_tempdir)
G.timeout = timeout
G.verbose = verbose
# Distgen generator
gen = Generator(verbose=verbose)
f = full_path(distgen_input_file)
distgen_params = yaml.safe_load(open(f))
# Set inputs
if settings:
G, distgen_params = set_gpt_and_distgen(G,
distgen_params,
settings,
verbose=verbose)
# Link particle files
particle_file = os.path.join(G.path, G.get_dist_file())
if (verbose):
print('Linking particle files, distgen output will point to -> "' +
os.path.basename(particle_file) + '" in working directory.')
G.set_dist_file(particle_file)
if ('output' in distgen_params and verbose):
print('Replacing Distgen output params')
distgen_params['output'] = {'type': 'gpt', 'file': particle_file}
if (verbose):
print('\nDistgen >------\n')
# Configure distgen
gen.parse_input(distgen_params)
# Run
beam = gen.beam()
write_gpt(beam, particle_file, verbose=verbose, asci2gdf_bin=asci2gdf_bin)
if (verbose):
print('------< Distgen\n')
if (auto_phase):
if (verbose):
print('\nAuto Phasing >------\n')
t1 = time.time()
# Create the distribution used for phasing
if (verbose):
print('****> Creating intiial distribution for phasing...')
phasing_beam = get_distgen_beam_for_phasing(beam,
n_particle=10,
verbose=verbose)
phasing_particle_file = os.path.join(G.path,
'gpt_particles.phasing.gdf')
write_gpt(phasing_beam,
phasing_particle_file,
verbose=verbose,
asci2gdf_bin=asci2gdf_bin)
if (verbose):
print('<**** Created intiial distribution for phasing.\n')
G.write_input_file() # Write the unphased input file
phased_file_name, phased_settings = gpt_phasing(
G.input_file,
path_to_gpt_bin=G.gpt_bin[:-3],
path_to_phasing_dist=phasing_particle_file,
verbose=verbose)
G.set_variables(phased_settings)
t2 = time.time()
if (verbose):
print(f'Time Ellapsed: {t2-t1} sec.')
print('------< Auto Phasing\n')
G.run(gpt_verbose=gpt_verbose)
return G
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 track_ref(self,
t0=0,
p0=1e-15,
xacc=6.5,
GBacc=5.5,
dtmin=1e-14,
dtmax=1e-8,
Ntout=100,
workdir=None):
dz_ccs_beg = np.linalg.norm(self.p_beg - self._ccs_beg_origin)
dz_fringe = 0
if (np.abs(self._b1) > 0):
dz_fringe = 10.0 / self._b1
else:
dz_fringe = 0
settings = {
'xacc': xacc,
'GBacc': GBacc,
'dtmin': dtmin,
'dtmax': dtmax,
'Ntout': Ntout,
'ZSTART':
-2 * np.sign(dz_ccs_beg - dz_fringe) * dz_ccs_beg - dz_fringe
}
particle = single_particle(z=dz_ccs_beg - dz_fringe,
pz=p0,
t=0,
weight=1,
status=1,
species=self.species)
if (workdir is None):
tempdir = tempfile.TemporaryDirectory(dir=workdir)
gpt_file = os.path.join(tempdir.name,
f'track_to_{self.name}.gpt.in')
workdir = tempdir.name
else:
gpt_file = os.path.join(workdir, f'{self.name}.gpt.in')
self.write_element_to_gpt_file(basic_template(gpt_file))
G = GPT(gpt_file,
initial_particles=particle,
ccs_beg=self.ccs_beg,
workdir=workdir,
use_tempdir=False)
G.set_variables(settings)
G.track1_to_z(z_end=dz_fringe,
ds=self.length + 2 * dz_fringe,
ccs_beg=self.ccs_beg,
ccs_end=self.ccs_end,
z0=dz_ccs_beg - dz_fringe,
t0=t0,
pz0=p0,
species=self.species,
s_screen=self.s_end + dz_fringe)
#os.remove(gpt_file)
return G