def transverse_trace_space_rms_emittance(x,
px,
py=None,
pz=None,
w=None,
disp_corrected=False,
corr_order=1):
"""Calculate the trasnverse trace-space RMS emittance of the
particle distribution in a given plane.
Parameters
----------
x : array
Contains the transverse position of the particles in one of the
transverse planes in units of meters
px : array
Contains the transverse momentum of the beam particles in the same
plane as x in non-dimmensional units (beta*gamma)
py : array
Contains the transverse momentum of the beam particles in the opposite
plane as as x in non-dimmensional units (beta*gamma). Necessary if
disp_corrected=True.
pz : array
Contains the longitudinal momentum of the beam particles in
non-dimmensional units (beta*gamma). Necessary if disp_corrected=True.
w : array or single value
Statistical weight of the particles.
disp_corrected : bool
Whether ot not to correct for dispersion contributions.
corr_order : int
Highest order up to which dispersion effects should be corrected.
Returns
-------
A float with the emmitance value in units of m * rad
"""
if len(x) > 1:
xp = px / pz
if disp_corrected:
# remove x-gamma correlation
gamma = np.sqrt(1 + np.square(px) + np.square(py) + np.square(pz))
gamma_avg = np.average(gamma, weights=w)
dgamma = (gamma - gamma_avg) / gamma_avg
x = remove_correlation(dgamma, x, w, corr_order)
# remove xp-gamma correlation
xp = remove_correlation(dgamma, xp, w, corr_order)
cov_x = np.cov(x, xp, aweights=np.abs(w))
em_x = np.sqrt(np.linalg.det(cov_x.astype(np.float32, copy=False)))
else:
em_x = 0
return em_x
def normalized_transverse_rms_slice_emittance(z,
x,
px,
py=None,
pz=None,
w=None,
disp_corrected=False,
corr_order=1,
n_slices=10,
len_slice=None):
"""Calculate the normalized transverse RMS slice emittance of the particle
distribution in a given plane.
Parameters
----------
z : array
Contains the longitudinal position of the particles in units of meters
x : array
Contains the transverse position of the particles in one of the
transverse planes in units of meters
px : array
Contains the transverse momentum of the beam particles in the same
plane as x in non-dimmensional units (beta*gamma)
py : array
Contains the transverse momentum of the beam particles in the opposite
plane as as x in non-dimmensional units (beta*gamma). Necessary if
disp_corrected=True.
pz : array
Contains the longitudinal momentum of the beam particles in
non-dimmensional units (beta*gamma). Necessary if disp_corrected=True.
w : array or single value
Statistical weight of the particles.
disp_corrected : bool
Whether ot not to correct for dispersion contributions.
corr_order : int
Highest order up to which dispersion effects should be corrected.
n_slices : array
Number of longitudinal slices in which to divite the particle
distribution. Not used if len_slice is specified.
len_slice : array
Length of the longitudinal slices. If not None, replaces n_slices.
Returns
-------
A tuple containing:
- An array with the emmitance value in each slice in units of m * rad.
- An array with the statistical weight of each slice.
- An array with the slice edges.
- A float with the weigthed average of the slice values.
"""
if disp_corrected:
# remove x-gamma correlation
gamma = np.sqrt(1 + np.square(px) + np.square(py) + np.square(pz))
gamma_avg = np.average(gamma, weights=w)
dgamma = (gamma - gamma_avg) / gamma_avg
x = remove_correlation(dgamma, x, w, corr_order)
slice_lims, n_slices = create_beam_slices(z, n_slices, len_slice)
slice_em = np.zeros(n_slices)
slice_weight = np.zeros(n_slices)
for i in np.arange(0, n_slices):
a = slice_lims[i]
b = slice_lims[i + 1]
slice_particle_filter = (z > a) & (z <= b)
if slice_particle_filter.any():
x_slice = x[slice_particle_filter]
px_slice = px[slice_particle_filter]
# if py is not None:
# py_slice = py[slice_particle_filter]
# else:
# py_slice=None
# if pz is not None:
# pz_slice = pz[slice_particle_filter]
# else:
# pz_slice=None
if hasattr(w, '__iter__'):
w_slice = w[slice_particle_filter]
else:
w_slice = w
slice_em[i] = normalized_transverse_rms_emittance(x_slice,
px_slice,
w=w_slice)
slice_weight[i] = np.sum(w_slice)
slice_avg = calculate_slice_average(slice_em, slice_weight)
return slice_em, slice_weight, slice_lims, slice_avg
def twiss_parameters(x,
px,
pz,
py=None,
w=None,
emitt='tr',
disp_corrected=False,
corr_order=1):
"""Calculate the alpha and beta functions of the beam in a certain
transverse plane
Parameters
----------
x : array
Contains the transverse position of the particles in one of the
transverse planes in units of meters
px : array
Contains the transverse momentum of the beam particles in the same
plane as x in non-dimmensional units (beta*gamma)
py : array
Contains the transverse momentum of the beam particles in the opposite
plane as as x in non-dimmensional units (beta*gamma). Necessary if
disp_corrected=True or emitt='ph'.
pz : array
Contains the longitudinal momentum of the beam particles in
non-dimmensional units (beta*gamma).
w : array or single value
Statistical weight of the particles.
emitt : str
Determines which emittance to use to calculate the Twiss parameters.
Possible values are 'tr' for trace-space emittance and 'ph' for
phase-space emittance
disp_corrected : bool
Whether ot not to correct for dispersion contributions.
corr_order : int
Highest order up to which dispersion effects should be corrected.
Returns
-------
A tuple with the value of the alpha, beta [m] and gamma [m^-1] functions
"""
if emitt == 'ph':
em_x = normalized_transverse_rms_emittance(x, px, py, pz, w,
disp_corrected, corr_order)
gamma = np.sqrt(1 + np.square(px) + np.square(py) + np.square(pz))
gamma_avg = np.average(gamma, weights=w)
x_avg = np.average(x, weights=w)
px_avg = np.average(px, weights=w)
# center x and x
x = x - x_avg
px = px - px_avg
if disp_corrected:
# remove x-gamma correlation
dgamma = (gamma - gamma_avg) / gamma_avg
x = remove_correlation(dgamma, x, w, corr_order)
b_x = np.average(x**2, weights=w) * gamma_avg / em_x
a_x = -np.average(x * px, weights=w) / em_x
elif emitt == 'tr':
em_x = transverse_trace_space_rms_emittance(x, px, py, pz, w,
disp_corrected, corr_order)
xp = px / pz
# center x and xp
x_avg = np.average(x, weights=w)
xp_avg = np.average(xp, weights=w)
x = x - x_avg
xp = xp - xp_avg
if disp_corrected:
# remove x-gamma correlation
gamma = np.sqrt(1 + np.square(px) + np.square(py) + np.square(pz))
gamma_avg = np.average(gamma, weights=w)
dgamma = (gamma - gamma_avg) / gamma_avg
x = remove_correlation(dgamma, x, w, corr_order)
# remove xp-gamma correlation
xp = remove_correlation(dgamma, xp, w, corr_order)
b_x = np.average(x**2, weights=w) / em_x
a_x = -np.average(x * xp, weights=w) / em_x
g_x = (1 + a_x**2) / b_x
return (a_x, b_x, g_x)