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 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
