def load_particle_array_from_npz(filename, print_params=False):
"""
Load beam file in npz format and return ParticleArray
:param filename:
:return:
"""
p_array = ParticleArray()
with np.load(filename) as data:
for key in data.keys():
p_array.__dict__[key] = data[key]
return p_array
def csrtrackBeam2particleArray(filename, orient="H"):
#H z x y pz px py -> x y z px py pz
#V z y x pz py px -> x y -z px py -pz
PD = np.loadtxt(filename)
#PD = load(infile)
n = np.shape(PD)[0] - 1 #length(PD(:,1))- 1
#print(n)
PD1 = np.zeros((n, 6))
Q = np.sum(PD[1:, 6])*1e9
t0 = PD[0, 0]
if orient=='H':
PD1[:, 1-1] = PD[1:, 2-1]
PD1[:, 2-1] = PD[1:, 3-1]
PD1[:, 3-1] = PD[1:, 1-1]
PD1[:, 4-1] = PD[1:, 5-1]
PD1[:, 5-1] = PD[1:, 6-1]
PD1[:, 6-1] = PD[1:, 4-1]
else:
PD1[:, 1-1] = -PD[1:, 3-1]
PD1[:, 2-1] = PD[1:, 2-1]
PD1[:, 3-1] = PD[1:, 1-1]
PD1[:, 4-1] = -PD[1:, 6-1]
PD1[:, 5-1] = PD[1:, 5-1]
PD1[:, 6-1] = PD[1:, 4-1]
#print("CSR", PD1[0, :])
for i in range(6):
PD1[1:n, i] = PD1[1:n, i] + PD1[0, i]
p_ref = np.sqrt(PD1[0, 3]**2 + PD1[0, 4]**2 + PD1[0, 5]**2)
px = PD1[:, 3]/ p_ref
py = PD1[:, 4] / p_ref
Eref = np.sqrt(m_e_eV ** 2 + p_ref ** 2)
pe = (np.sqrt(m_e_eV**2 + (PD1[:, 3]**2 + PD1[:, 4]**2 + PD1[:, 5]**2)) - Eref) / p_ref
p_array = ParticleArray(n)
p_array.rparticles[0] = PD1[:, 0] - 0*PD1[0, 0]
p_array.rparticles[2] = PD1[:, 1] - 0*PD1[0, 1]
p_array.rparticles[4] = -(PD1[:, 2] - PD1[0, 2])
p_array.rparticles[1] = px[:]
p_array.rparticles[3] = py[:]
p_array.rparticles[5] = pe[:]
p_array.q_array[:] = PD[1:, 6]
p_array.s = PD1[0, 2]
p_array.E = Eref*1e-9
return p_array
def slicing_p_array(p_array, SIZE):
"""
Divide ParticleArray into chuncks
:param p_array: ParticleArray
:param SIZE: number of threads
:return: list of ParticleArray
"""
if SIZE == 0:
raise ValueError("SIZE must be not zero")
if SIZE > p_array.n:
raise ValueError("SIZE must be less or equal number of paricles")
p_array_list = []
nchunk = int(p_array.n / SIZE)
#print(p_array.n, nchunk, SIZE)
nparticles_list = np.ones(SIZE, dtype=int)*nchunk
nparticles_last = p_array.n - nchunk * (SIZE - 1)
if nparticles_last > nchunk:
nchunk += 1
nextra = nparticles_last - nchunk
nparticles_list[:nextra] += 1
nparticles_list[-1] = nparticles_last - nextra
#nparticles = nchunk
ncontrol = 0
for i, nparticles in enumerate(nparticles_list):
p_array_part = ParticleArray(n=nparticles)
p_array_part.rparticles[:, :] = p_array.rparticles[:, ncontrol:ncontrol+nparticles]
p_array_part.q_array[:] = p_array.q_array[ncontrol:ncontrol+nparticles]
ncontrol += nparticles
p_array_part.s = p_array.s
p_array_part.E = p_array.E
p_array_list.append(p_array_part)
return p_array_list
def test_lost_particle_recorder():
n = 20
p_array = np.arange(n)
lpr = ParticleArray.LostParticleRecorder(n)
for pos, inds in enumerate([[1, 3], [0, 1], [4, 10], [0, 1],
[1, 0, 10, 7, 8]]):
lpr.add(inds, pos)
p_array = np.delete(p_array, inds)
lpr.lost_particles.sort()
assert all([
lp_idx == truth_idx for lp_idx, truth_idx in zip(
lpr.lost_particles, [0, 1, 2, 3, 4, 5, 6, 7, 8, 14, 15, 16, 18])
])
assert all([
lp_idx == truth_idx
for lp_idx, truth_idx in zip(lpr.lp_to_pos_hist, [(0, 2), (1, 2), (
2, 2), (3, 2), (4, 5)])
])
def apply(self, prcl_series):
"""
:param prcl_series: can be list of Particles [Particle_1, Particle_2, ... ] or ParticleArray
:return: None
"""
if prcl_series.__class__ == ParticleArray:
self.map(prcl_series.rparticles, energy=prcl_series.E)
prcl_series.E += self.delta_e
prcl_series.s += self.length
elif prcl_series.__class__ == Particle:
p = prcl_series
p.x, p.px, p.y, p.py, p.tau, p.p = self.map(
array([[p.x], [p.px], [p.y], [p.py], [p.tau], [p.p]]), p.E)[:,
0]
p.s += self.length
p.E += self.delta_e
elif prcl_series.__class__ == list and prcl_series[
0].__class__ == Particle:
# If the energy is not the same (p.E) for all Particles in the list of Particles
# in that case cycle is applied. For particles with the same energy p.E
list_e = array([p.E for p in prcl_series])
if False in (list_e[:] == list_e[0]):
for p in prcl_series:
self.map(array([[p.x], [p.px], [p.y], [p.py], [p.tau],
[p.p]]),
energy=p.E)
p.E += self.delta_e
p.s += self.length
else:
pa = ParticleArray()
pa.list2array(prcl_series)
pa.E = prcl_series[0].E
self.map(pa.rparticles, energy=pa.E)
pa.E += self.delta_e
pa.s += self.length
pa.array2ex_list(prcl_series)
else:
print(prcl_series)
exit("Unknown type of Particle_series. class TransferMap.apply()")
def apply(self, prcl_series):
"""
:param prcl_series: can be list of Particles [Particle_1, Particle_2, ... ] or ParticleArray
:return: None
"""
if prcl_series.__class__ == ParticleArray:
self.map(prcl_series.rparticles, energy=prcl_series.E)
prcl_series.E += self.delta_e
prcl_series.s += self.length
elif prcl_series.__class__ == Particle:
p = prcl_series
p.x, p.px, p.y, p.py, p.tau, p.p = self.map(array([[p.x], [p.px], [p.y], [p.py], [p.tau], [p.p]]), p.E)[:,0]
p.s += self.length
p.E += self.delta_e
elif prcl_series.__class__ == list and prcl_series[0].__class__ == Particle:
# If the energy is not the same (p.E) for all Particles in the list of Particles
# in that case cycle is applied. For particles with the same energy p.E
list_e = array([p.E for p in prcl_series])
if False in (list_e[:] == list_e[0]):
for p in prcl_series:
self.map(array([[p.x], [p.px], [p.y], [p.py], [p.tau], [p.p]]), energy=p.E)
p.E += self.delta_e
p.s += self.length
else:
pa = ParticleArray()
pa.list2array(prcl_series)
pa.E = prcl_series[0].E
self.map(pa.rparticles, energy=pa.E)
pa.E += self.delta_e
pa.s += self.length
pa.array2ex_list(prcl_series)
else:
print(prcl_series)
exit("Unknown type of Particle_series. class TransferMap.apply()")
def csrtrackBeam2particleArray(filename, orient="H"):
#H z x y pz px py -> x y z px py pz
#V z y x pz py px -> x y -z px py -pz
PD = np.loadtxt(filename)
#PD = load(infile)
n = np.shape(PD)[0] - 1 #length(PD(:,1))- 1
#print(n)
PD1 = np.zeros((n, 6))
Q = np.sum(PD[1:, 6]) * 1e9
t0 = PD[0, 0]
if orient == 'H':
PD1[:, 1 - 1] = PD[1:, 2 - 1]
PD1[:, 2 - 1] = PD[1:, 3 - 1]
PD1[:, 3 - 1] = PD[1:, 1 - 1]
PD1[:, 4 - 1] = PD[1:, 5 - 1]
PD1[:, 5 - 1] = PD[1:, 6 - 1]
PD1[:, 6 - 1] = PD[1:, 4 - 1]
else:
PD1[:, 1 - 1] = -PD[1:, 3 - 1]
PD1[:, 2 - 1] = PD[1:, 2 - 1]
PD1[:, 3 - 1] = PD[1:, 1 - 1]
PD1[:, 4 - 1] = -PD[1:, 6 - 1]
PD1[:, 5 - 1] = PD[1:, 5 - 1]
PD1[:, 6 - 1] = PD[1:, 4 - 1]
#print("CSR", PD1[0, :])
for i in range(6):
PD1[1:n, i] = PD1[1:n, i] + PD1[0, i]
p_ref = np.sqrt(PD1[0, 3]**2 + PD1[0, 4]**2 + PD1[0, 5]**2)
px = PD1[:, 3] / p_ref
py = PD1[:, 4] / p_ref
Eref = np.sqrt(m_e_eV**2 + p_ref**2)
pe = (np.sqrt(m_e_eV**2 +
(PD1[:, 3]**2 + PD1[:, 4]**2 + PD1[:, 5]**2)) - Eref) / p_ref
p_array = ParticleArray(n)
p_array.rparticles[0] = PD1[:, 0] - 0 * PD1[0, 0]
p_array.rparticles[2] = PD1[:, 1] - 0 * PD1[0, 1]
p_array.rparticles[4] = -(PD1[:, 2] - PD1[0, 2])
p_array.rparticles[1] = px[:]
p_array.rparticles[3] = py[:]
p_array.rparticles[5] = pe[:]
p_array.q_array[:] = PD[1:, 6]
p_array.s = PD1[0, 2]
p_array.E = Eref * 1e-9
return p_array
def csrtrackBeam2particleArray(filename, orient="H"):
"""
Function to read CSRtrack beam files ".fmt1"
H: z x y pz px py -> x y z px py pz
V: z y x pz py px -> x y -z px py -pz
:param filename: filename
:param orient: str, "H" or "V" horizontal or vertical orientation
:return: ParticleArray
"""
pd = np.loadtxt(filename)
n = np.shape(pd)[0] - 1
pd1 = np.zeros((n, 6))
if orient == 'H':
pd1[:, 0] = pd[1:, 1]
pd1[:, 1] = pd[1:, 2]
pd1[:, 2] = pd[1:, 0]
pd1[:, 3] = pd[1:, 4]
pd1[:, 4] = pd[1:, 5]
pd1[:, 5] = pd[1:, 3]
else:
pd1[:, 0] = -pd[1:, 2]
pd1[:, 1] = pd[1:, 1]
pd1[:, 2] = pd[1:, 0]
pd1[:, 3] = -pd[1:, 5]
pd1[:, 4] = pd[1:, 4]
pd1[:, 5] = pd[1:, 3]
for i in range(6):
pd1[1:n, i] = pd1[1:n, i] + pd1[0, i]
p_ref = np.sqrt(pd1[0, 3]**2 + pd1[0, 4]**2 + pd1[0, 5]**2)
px = pd1[:, 3] / p_ref
py = pd1[:, 4] / p_ref
Eref = np.sqrt(m_e_eV**2 + p_ref**2)
pe = (np.sqrt(m_e_eV**2 +
(pd1[:, 3]**2 + pd1[:, 4]**2 + pd1[:, 5]**2)) - Eref) / p_ref
p_array = ParticleArray(n)
p_array.rparticles[0] = pd1[:, 0]
p_array.rparticles[2] = pd1[:, 1]
p_array.rparticles[4] = -(pd1[:, 2] - pd1[0, 2])
p_array.rparticles[1] = px[:]
p_array.rparticles[3] = py[:]
p_array.rparticles[5] = pe[:]
p_array.q_array[:] = pd[1:, 6]
p_array.s = pd1[0, 2]
p_array.E = Eref * 1e-9
return p_array
def wake_t_beam_to_parray(wake_t_beam, gamma_ref=None, z_ref=None):
"""
Converts a Wake-T particle beam to an Ocelot ParticleArray.
Parameters:
-----------
wake_t_beam : ParticleBunch (Wake-T class)
The original particle distribution from Wake-T.
gamma_ref : float (Optional)
Reference energy of the particle beam used for tracking in Ocelot. If
not specified, the reference energy will be taken as the average
energy of the input distribution.
z_ref : float (Optional)
Reference longitudinal position of the particle beam used for tracking
in Ocelot. If not specified, the reference value will be taken as the
average longitudinal position of the input distribution.
Returns:
--------
An Ocelot ParticleArray.
"""
# Extract particle coordinates.
x = wake_t_beam.x # [m]
y = wake_t_beam.y # [m]
z = wake_t_beam.xi + wake_t_beam.prop_distance # [m]
px = wake_t_beam.px # [m_e * c]
py = wake_t_beam.py # [m_e * c]
pz = wake_t_beam.pz # [m_e * c]
q = wake_t_beam.q # [C]
# Calculate gamma.
gamma = np.sqrt(1 + px**2 + py**2 + pz**2)
# Determine reference energy and longitudinal position.
if gamma_ref is None:
gamma_ref = np.average(gamma, weights=q)
if z_ref is None:
z_ref = wake_t_beam.prop_distance
# Calculate momentum deviation (dp) and kinetic momentum (p_kin).
b_ref = np.sqrt(1 - gamma_ref**(-2))
dp = (gamma-gamma_ref)/(gamma_ref*b_ref)
p_kin = np.sqrt(gamma_ref**2 - 1)
# Create particle array
p_array = ParticleArray(len(q))
p_array.rparticles[0] = x
p_array.rparticles[2] = y
p_array.rparticles[4] = -(z - z_ref)
p_array.rparticles[1] = px/p_kin
p_array.rparticles[3] = py/p_kin
p_array.rparticles[5] = dp
p_array.q_array[:] = q
p_array.s = z_ref
p_array.E = gamma_ref * m_e_GeV
return p_array
"""
if len(lst) == 0:
raise ValueError("Length of list must be not zero")
if nthreads == 0:
raise ValueError("Number of threads must be not zero")
if nthreads > len(lst):
raise ValueError(
"Number of threads ({}) must be less or equal list length ({})".
format(nthreads, len(lst)))
n = int(len(lst) / nthreads)
return [lst[i:i + n] for i in range(0, len(lst), n)]
if __name__ == "__main__":
p_array = ParticleArray(n=10)
p_array_list = slicing_p_array(p_array, SIZE=1)
print("test1: len(list) = ", len(p_array_list), "n particles in each: ",
[pa.n for pa in p_array_list])
p_array_list = slicing_p_array(p_array, SIZE=2)
print("test2: len(list) = ", len(p_array_list), "n particles in each: ",
[pa.n for pa in p_array_list])
p_array_list = slicing_p_array(p_array, SIZE=3)
print("test3: len(list) = ", len(p_array_list), "n particles in each: ",
[pa.n for pa in p_array_list])
p_array_list = slicing_p_array(p_array, SIZE=10)
print("test4: len(list) = ", len(p_array_list), "n particles in each: ",
[pa.n for pa in p_array_list])