def save_bunch_to_file(bunch,
code_name,
folder_path,
file_name,
reposition=False,
avg_pos=[None, None, None],
n_part=None):
bunch_data = bunch.get_bunch_matrix()
ds.save_beam(code_name,
bunch_data,
folder_path,
file_name,
reposition=reposition,
avg_pos=avg_pos,
n_part=n_part)
def generate_from_file_modifying_twiss(code_name,
file_path,
read_kwargs={},
alphax_target=None,
betax_target=None,
alphay_target=None,
betay_target=None,
save_to_file=False,
save_to_code='astra',
save_to_path=None,
file_name=None,
save_kwargs={},
perform_checks=False):
"""
Creates a transversely Gaussian particle bunch with the specified Twiss
parameters.
Parameters
----------
code_name: str
Name of the tracking or PIC code of the data to read. Possible values
are 'csrtrack', 'astra', 'openpmd', 'osiris' and 'hipace'
file_path: str
Path of the file containing the data
read_kwargs: dict
Dictionary containing optional parameters for the read_beam function.
save_to_file: bool
Whether to save the generated distribution to a file.
save_to_code: string
(optional) Name of the target code that will use the saved file.
Possible values are 'csrtrack', 'astra' and 'fbpic'. Required if
save_to_file=True.
save_to_path: string
(optional) Path to the folder where to save the data. Required if
save_to_file=True.
file_name: string
(optional) Name of the file where to store the beam data. Required if
save_to_file=True.
save_kwargs: dict
Dictionary containing optional parameters for the save_beam function.
perform_checks: bool
Whether to compute and print the parameters of the generated bunch.
Returns
-------
A tuple with 7 arrays containing the 6D components and charge of the
modified distribution.
"""
# Read distribution
x, y, z, px, py, pz, q = dr.read_beam(code_name, file_path, **read_kwargs)
# Modify Twiss parameters
x, y, z, px, py, pz, q = bo.modify_twiss_parameters_all_beam(
[x, y, z, px, py, pz, q],
alphax_target=alphax_target,
betax_target=betax_target,
alphay_target=alphay_target,
betay_target=betay_target)
# Save to file
if save_to_file:
print('Saving to file... ', end='')
ds.save_beam(save_to_code, [x, y, z, px, py, pz, q], save_to_path,
file_name, **save_kwargs)
print('Done.')
# Perform checks
if perform_checks:
_check_beam_parameters(x, y, z, px, py, pz, q)
return x, y, z, px, py, pz, q
def generate_gaussian_bunch_from_twiss(a_x,
a_y,
b_x,
b_y,
en_x,
en_y,
ene,
ene_sp,
s_t,
q_tot,
n_part,
x_c=0,
y_c=0,
z_c=0,
lon_profile='gauss',
min_len_scale_noise=None,
sigma_trunc_lon=None,
smooth_sigma=None,
smooth_trunc=None,
save_to_file=False,
save_to_code='astra',
save_to_path=None,
file_name=None,
perform_checks=False):
"""
Creates a transversely Gaussian particle bunch with the specified Twiss
parameters.
Parameters
----------
a_x : float
Alpha parameter in the x-plane.
a_y : float
Alpha parameter in the y-plane.
b_x : float
Beta parameter in the x-plane in units of m.
b_y : float
Beta parameter in the y-plane in units of m.
en_x : float
Normalized trace-space emittance in the x-plane in units of m*rad.
en_y : float
Normalized trace-space emittance in the y-plane in units of m*rad.
ene: float
Mean bunch energy in non-dimmensional units (beta*gamma).
ene_sp: float
Relative energy spread in %.
s_t: float
Bunch duration in seconds. If lon_profile='gauss', this corresponds to
the RMS duration. If lon_profile='flattop' or
lon_profile='flattop_smoothed', this instead the whole flat-top lenght.
q_tot: float
Total bunch charge in C.
n_part: int
Total number of particles in the bunch.
x_c: float
Central bunch position in the x-plane in units of m.
y_c: float
Central bunch position in the y-plane in units of m.
z_c: float
Central bunch position in the z-plane in units of m.
lon_profile: string
Longitudonal profile of the bunch. Possible values are 'gauss' and
'flattop'.
min_len_scale_noise: float
(optional) If specified, a different algorithm to generate a less noisy
longitudinal profile is used. This algorithm creates a profile that is
smooth for a longitudinal binning of the bunch with
bin lengths >= min_len_scale_noise
sigma_trunc_lon: float
(optional) If specified, it truncates the longitudinal distribution of
the bunch between [z_c-sigma_trunc_lon*s_z, z_c+sigma_trunc_lon*s_z].
Only used when lon_profile = 'gauss' and required if
min_len_scale_noise is specified.
smooth_sigma: float
The sigma of the Gaussian longitudinal smoothing applied to the
flat-top profile when lon_profile='flattop_smoothed'. Units are in
seconds.
smooth_trunc: float
Number of sigmas after which to truncate the Gaussian smoothing when
lon_profile='flattop_smoothed'
save_to_file: bool
Whether to save the generated distribution to a file.
save_to_code: string
(optional) Name of the target code that will use the saved file.
Possible values are 'csrtrack', 'astra' and 'fbpic'. Required if
save_to_file=True.
save_to_path: string
(optional) Path to the folder where to save the data. Required if
save_to_file=True.
file_name: string
(optional) Name of the file where to store the beam data. Required if
save_to_file=True.
perform_checks: bool
Whether to compute and print the parameters of the generated bunch.
Returns
-------
The 6D components and charge of the bunch in 7 arrays.
"""
print('Generating particle distribution... ', end='')
# Calculate necessary values
n_part = int(n_part)
ene_sp = ene_sp / 100
ene_sp_abs = ene_sp * ene
s_z = s_t * ct.c
em_x = en_x / ene
em_y = en_y / ene
g_x = (1 + a_x**2) / b_x
g_y = (1 + a_y**2) / b_y
s_x = np.sqrt(em_x * b_x)
s_y = np.sqrt(em_y * b_y)
s_xp = np.sqrt(em_x * g_x)
s_yp = np.sqrt(em_y * g_y)
p_x = -a_x * em_x / (s_x * s_xp)
p_y = -a_y * em_y / (s_y * s_yp)
# Create longitudinal distributions
if lon_profile == 'gauss':
z = _create_gaussian_longitudinal_profile(z_c, s_z, n_part,
sigma_trunc_lon,
min_len_scale_noise)
elif lon_profile == 'flattop':
z = _create_flattop_longitudinal_profile(z_c, s_z, n_part,
min_len_scale_noise)
elif lon_profile == 'flattop_smoothed':
z = _create_flattop_longitudinal_profile_with_smoothing(
z_c, s_z, n_part, min_len_scale_noise, smooth_sigma, smooth_trunc)
# Define again n_part in case it changed when crealing long. profile
n_part = len(z)
pz = np.random.normal(ene, ene_sp_abs, n_part)
# Create normalized gaussian distributions
u_x = np.random.standard_normal(n_part)
v_x = np.random.standard_normal(n_part)
u_y = np.random.standard_normal(n_part)
v_y = np.random.standard_normal(n_part)
# Calculate transverse particle distributions
x = s_x * u_x
xp = s_xp * (p_x * u_x + np.sqrt(1 - np.square(p_x)) * v_x)
y = s_y * u_y
yp = s_yp * (p_y * u_y + np.sqrt(1 - np.square(p_y)) * v_y)
# Change from slope to momentum
px = xp * pz
py = yp * pz
# Charge
q = np.ones(n_part) * (q_tot / n_part)
print('Done.')
# Save to file
if save_to_file:
print('Saving to file... ', end='')
ds.save_beam(save_to_code, [x, y, z, px, py, pz, q], save_to_path,
file_name)
print('Done.')
if perform_checks:
_check_beam_parameters(x, y, z, px, py, pz, q)
return x, y, z, px, py, pz, q
def convert_beam(orig_code,
final_code,
orig_path,
final_path,
final_file_name,
reposition=False,
avg_pos=[None, None, None],
avg_mom=[None, None, None],
n_part=None,
read_kwargs={},
save_kwargs={}):
"""Converts particle data from one code to another.
Parameters
----------
orig_code : str
Name of the tracking or PIC code of the original data. Possible values
are 'csrtrack', 'astra' and 'openpmd'
final_code : str
Name of the tracking or PIC code in which to convert the data. Possible
values are 'csrtrack', 'astra', 'fbpic' and 'openpmd'
orig_path : str
Path of the file containing the original data
final_path : str
Path to the folder in which to save the converted data
final_file_name : str
Name of the file to save, without extension
reposition : bool
Optional. Whether to reposition de particle distribution in space
and/or momentum centered in the coordinates specified in avg_pos and
avg_mom
avg_pos : list
Optional, only used it reposition=True. Contains the new average
positions of the beam after repositioning. Should be specified as
[x_avg, y_avg, z_avg] in meters. Setting a component as None prevents
repositioning in that coordinate.
avg_mom : list
Optional, only used it reposition=True. Contains the new
average momentum of the beam after repositioning. Should be specified
as [px_avg, py_avg, pz_avg] in non-dimmesional units (beta*gamma).
Setting a component as None prevents repositioning in that coordinate.
n_part : int
Optional. Number of particles to save. Must be lower than the original
number of particles. Particles to save are chosen randomly.
Other Parameters
----------------
**kwargs
This method takes additional keyword parameters that might be needed
for some data readers. Currenlty, the only parameter is 'species_name',
for reading data from PIC codes.
"""
x, y, z, px, py, pz, q = read_beam(orig_code, orig_path, **read_kwargs)
beam_data = [x, y, z, px, py, pz, q]
save_beam(final_code, beam_data, final_path, final_file_name, reposition,
avg_pos, avg_mom, n_part, **save_kwargs)