def filter_beam_sigma(beam_matrix, max_sigma, w=None):
"""Filter the beam particles keeping only those within a certain number of
sigmas.
Parameters
----------
beam_matrix : array
M x N matrix containing all M components of the N particles.
max_sigma : array
Array of size M with the maximum number of sigmas of each component.
For each compoonent x, particles with x < x_avg - x_std * max_sigma or
x > x_avg + x_std * max_sigma will be discarded. If max_sigma is None
for some component, no filtering is performed.
w : array
Statistical weight of the particles.
Returns
-------
A M x N' matrix containing the particles within the given range. N' is the
number of particles after filtering.
"""
elements_to_keep = np.ones_like(beam_matrix[0])
for i, arr in enumerate(beam_matrix):
if max_sigma[i] is not None:
arr_mean = np.average(arr, weights=w)
arr_std = weighted_std(arr, weights=w)
min_range = arr_mean - max_sigma[i] * arr_std
max_range = arr_mean + max_sigma[i] * arr_std
elements_to_keep = np.where((arr < min_range) | (arr > max_range),
0, elements_to_keep)
elements_to_keep = np.array(elements_to_keep, dtype=bool)
return beam_matrix[:, elements_to_keep]
def rms_relative_correlated_energy_spread(z, px, py, pz, w=None):
"""Calculate the correlated energy spread of the provided particle
distribution
Parameters
----------
z : array
Contains the longitudinal position of the particles in units of meters
px : array
Contains the transverse momentum in the x direction of the
beam particles in non-dimmensional units (beta*gamma)
py : array
Contains the transverse momentum in the x direction of the
beam particles in non-dimmensional units (beta*gamma)
pz : array
Contains the longitudonal momentum of the beam particles in
non-dimmensional units (beta*gamma)
w : array or single value
Statistical weight of the particles.
Returns
-------
A float with the energy spread value in non-dimmensional units,
i.e. [1/(m_e c**2)]
"""
K = longitudinal_energy_chirp(z, px, py, pz, w)
mean_z = np.average(z, weights=w)
dz = z - mean_z
corr_ene = K * dz
corr_ene_sp = weighted_std(corr_ene, w)
return corr_ene_sp
def rms_energy_spread(px, py, pz, w=None):
"""Calculate the absotule RMS energy spread of the provided particle
distribution
Parameters
----------
px : array
Contains the transverse momentum in the x direction of the
beam particles in non-dimmensional units (beta*gamma)
py : array
Contains the transverse momentum in the x direction of the
beam particles in non-dimmensional units (beta*gamma)
pz : array
Contains the longitudonal momentum of the beam particles in
non-dimmensional units (beta*gamma)
w : array or single value
Statistical weight of the particles.
Returns
-------
A float with the energy spread value in non-dimmensional units,
i.e. [1/(m_e c**2)]
"""
part_ene = np.sqrt(1 + np.square(px) + np.square(py) + np.square(pz))
ene_std = weighted_std(part_ene, weights=w)
return ene_std
def rms_size(x, w=None):
"""Calculate the RMS bunch size of the provided particle
distribution
Parameters
----------
x : array
Contains the transverse position of the particles in units of meters
w : array or single value
Statistical weight of the particles.
Returns
-------
A float with the RMS length value in meters.
"""
s_x = weighted_std(x, weights=w)
return s_x
def rms_length(z, w=None):
"""Calculate the RMS bunch length of the provided particle
distribution
Parameters
----------
z : array
Contains the longitudinal position of the particles in units of meters
w : array or single value
Statistical weight of the particles.
Returns
-------
A float with the RMS length value in meters.
"""
s_z = weighted_std(z, weights=w)
return s_z
def rms_relative_uncorrelated_energy_spread(z, px, py, pz, w=None):
"""Calculate the uncorrelated energy spread of the provided particle
distribution
Parameters
----------
z : array
Contains the longitudinal position of the particles in units of meters
px : array
Contains the transverse momentum in the x direction of the
beam particles in non-dimmensional units (beta*gamma)
py : array
Contains the transverse momentum in the x direction of the
beam particles in non-dimmensional units (beta*gamma)
pz : array
Contains the longitudonal momentum of the beam particles in
non-dimmensional units (beta*gamma)
w : array or single value
Statistical weight of the particles.
Returns
-------
A float with the energy spread value in non-dimmensional units,
i.e. [1/(m_e c**2)]
"""
if len(z) > 1:
ene = np.sqrt(1 + np.square(px) + np.square(py) + np.square(pz))
mean_ene = np.average(ene, weights=w)
mean_z = np.average(z, weights=w)
dE = ene - mean_ene
dz = z - mean_z
p = np.polyfit(dz, dE, 1)
K = p[0]
unc_ene = ene - K * dz
unc_ene_sp = weighted_std(unc_ene, w) / mean_ene
else:
unc_ene_sp = 0
return unc_ene_sp