SF = Sextupole(l=0.0001, k2= 17673.786254063251*0, eid= "SF")
SD = Sextupole(l=0.0001, k2=-36169.817233025707*0, eid= "SD")
Q1 = Quadrupole(l=0.293, k1=2.62, eid= "Q1")
Q2 = Quadrupole(l=0.293, k1=-3.1, eid= "Q2")
Q3 = Quadrupole(l=0.327, k1=2.8, eid= "Q3")
Q4 = Quadrupole(l=0.291, k1=-3.7, eid= "Q4")
Q5 = Quadrupole(l=0.391, k1=4.0782, eid= "Q5")
Q6 = Quadrupole(l=0.291, k1=-3.534859, eid= "D6")
B1 = SBend(l = 0.23, angle = 0.23/19.626248, eid= "B1")
B2 = SBend(l = 1.227, angle = 1.227/4.906312, eid= "B2")
u = Undulator(lperiod=0.02, nperiods=50, Kx=2)
u.mag_field = lambda x, y, z: und_field(x, y, z, u.lperiod, u.Kx)
u.npoints = 500
cell_u = ( D1,SF, D2,Q1,D3, Q2,D2,SD,D4,B1,B2,D5,Q3,D5,B2,B1,D6,Q4,D7,Q5,D8,Q6,
Drift(l=1.0045), u, Drift(l=1.0045),
Q6,D8,Q5,D7,Q4,D6,B1,B2,D5,Q3,D5,B2,B1,D4,SD,D2,Q2,D3,Q1,D2,SF,D1)
cell = ( D1,SF, D2,Q1,D3, Q2,D2,SD,D4,B1,B2,D5,Q3,D5,B2,B1,D6,Q4,D7,Q5,D8,Q6,D9,Q6,D8,Q5,D7,Q4,D6,B1,B2,D5,Q3,D5,B2,B1,D4,SD,D2,Q2,D3,Q1,D2,SF,D1)
ring = 3*cell + cell_u + 2*cell
method = MethodTM()
method.params[Sextupole] = KickTM
#method.params[Undulator] = UndulatorTestTM
method.params[Undulator] = RungeKuttaTM
method.global_method = TransferMap
def prepare(self, lat):
"""
calculation of trajectory in rectangular coordinates
calculation of the z_csr_start
:param lat: Magnetic Lattice
:return: self.csr_traj: trajectory. traj[0,:] - longitudinal coordinate,
traj[1,:], traj[2,:], traj[3,:] - rectangular coordinates, \
traj[4,:], traj[5,:], traj[6,:] - tangential unit vectors
"""
# if pict_debug = True import matplotlib
if self.pict_debug:
self.plt = importlib.import_module("matplotlib.pyplot")
self.napply = 0
self.total_wake = 0
self.z_csr_start = sum([p.l for p in lat.sequence[:self.indx0]])
p = Particle()
beta = 1. if self.energy == None else np.sqrt(
1. - 1. / (self.energy / m_e_GeV)**2)
self.csr_traj = np.transpose([[0, p.x, p.y, p.s, p.px, p.py, 1.]])
if Undulator in [
elem.__class__
for elem in lat.sequence[self.indx0:self.indx1 + 1]
]:
self.rk_traj = True
for elem in lat.sequence[self.indx0:self.indx1 + 1]:
if elem.l == 0:
continue
delta_s = elem.l
step = self.traj_step
if elem.__class__ in [Bend, RBend, SBend] and not self.rk_traj:
if elem.angle != 0:
R = -elem.l / elem.angle
else:
R = 0
Rx = R * np.cos(elem.tilt)
Ry = R * np.sin(elem.tilt)
#B = energy*1e9*beta/(R*speed_of_light)
R_vect = [-Ry, Rx, 0]
self.csr_traj = arcline(self.csr_traj, delta_s, step, R_vect)
elif elem.__class__ in [
Bend, RBend, SBend, XYQuadrupole, Undulator
] and self.energy is not None and self.rk_traj:
"""
rk_track_in_field accepts initial conditions (initial coordinates) in the ParticleArray.rparticles format
from another hand the csr module use trajectory in another format (see csr.py)
"""
if elem.l < 1e-10:
continue
delta_z = delta_s
if elem.__class__ is XYQuadrupole:
hx = elem.k1 * elem.x_offs
hy = -elem.k1 * elem.y_offs
By = self.energy * 1e9 * beta * hx / speed_of_light
Bx = -self.energy * 1e9 * beta * hy / speed_of_light
mag_field = lambda x, y, z: (Bx, By, 0)
elif elem.__class__ == Undulator:
gamma = self.energy / m_e_GeV
ku = 2 * np.pi / elem.lperiod
delta_z = elem.lperiod * elem.nperiods
# delta_s += 0.5* delta_z / (gamma ) ** 2 * (1 + 0.5 * (elem.Kx ) ** 2)
delta_s = delta_z * (1 + 0.25 * (elem.Kx / gamma)**2)
By = elem.Kx * m_e_eV * 2. * pi / (elem.lperiod *
speed_of_light)
Bx = elem.Ky * m_e_eV * 2. * pi / (elem.lperiod *
speed_of_light)
mag_field = lambda x, y, z: (0, -By * np.cos(ku * z), 0)
# ending poles 1/4, -3/4, 1, -1, ... (or -1/4, 3/4, -1, 1)
if self.end_poles:
mag_field = lambda x, y, z: und_field(x,
y,
z,
elem.lperiod,
elem.Kx,
nperiods=elem.
nperiods)
else:
delta_z = delta_s * np.sin(
elem.angle
) / elem.angle if elem.angle != 0 else delta_s
hx = elem.angle / elem.l * np.cos(elem.tilt)
hy = elem.angle / elem.l * np.sin(elem.tilt)
By = self.energy * 1e9 * beta * hx / speed_of_light
Bx = -self.energy * 1e9 * beta * hy / speed_of_light
mag_field = lambda x, y, z: (Bx, By, 0)
sre0 = self.csr_traj[:, -1]
N = int(max(1, np.round(delta_s / step)))
SRE2 = np.zeros((7, N))
rparticle0 = np.array(
[[self.csr_traj[1, -1]],
[self.csr_traj[4, -1] / self.csr_traj[6, -1]],
[self.csr_traj[2, -1]],
[self.csr_traj[5, -1] / self.csr_traj[6, -1]], [0], [0]])
traj = rk_track_in_field(rparticle0,
s_stop=delta_z,
N=N + 1,
energy=self.energy,
mag_field=mag_field,
s_start=0)
x = traj[0::9].flatten()
y = traj[2::9].flatten()
xp = traj[1::9].flatten()
yp = traj[3::9].flatten()
z = traj[4::9].flatten()
xp2 = xp * xp
yp2 = yp * yp
zp = np.sqrt(1. / (1. + xp2 + yp2))
s = cumtrapz(np.sqrt(1. + xp2 + yp2), z, initial=0)
if elem.__class__ == Undulator: delta_s = s[-1]
dS = float(delta_s) / N
s_unif = np.arange(0, N + 1) * dS
z_s = np.interp(s_unif, s, z)
x = np.interp(z_s, z, x)
y = np.interp(z_s, z, y)
zp_s = np.interp(z_s, z, zp)
xp = np.interp(z_s, z, xp * zp)
yp = np.interp(z_s, z, yp * zp)
SRE2[0, :] = sre0[0] + s_unif[1:]
SRE2[1, :] = x[1:]
SRE2[2, :] = y[1:]
SRE2[3, :] = sre0[3] + z_s[1:]
SRE2[4, :] = xp[1:]
SRE2[5, :] = yp[1:]
SRE2[6, :] = zp_s[1:]
self.csr_traj = np.append(self.csr_traj, SRE2, axis=1)
else:
R_vect = [0, 0, 0.]
self.csr_traj = arcline(self.csr_traj, delta_s, step, R_vect)
# plot trajectory of the refernece particle
if self.pict_debug:
fig = self.plt.figure(figsize=(10, 8))
ax1 = self.plt.subplot(211)
#ax1.plot(self.csr_traj[0, :], self.csr_traj[0, :] -self.csr_traj[3, :], "r", label="X")
ax1.plot(self.csr_traj[0, :],
self.csr_traj[1, :] * 1000,
"r",
label="X")
ax1.plot(self.csr_traj[0, :],
self.csr_traj[2, :] * 1000,
"b",
label="Y")
self.plt.legend()
self.plt.ylabel("X/Y [mm]")
self.plt.setp(ax1.get_xticklabels(), visible=False)
ax3 = self.plt.subplot(212, sharex=ax1)
ax3.plot(self.csr_traj[0, :],
self.csr_traj[1 + 3, :] * 1000,
"r",
label=r"$X'$")
ax3.plot(self.csr_traj[0, :],
self.csr_traj[2 + 3, :] * 1000,
"b",
label=r"$Y'$")
self.plt.legend()
self.plt.ylabel(r"$X'/Y'$ [mrad]")
self.plt.xlabel("s [m]")
self.plt.show()
# data = np.array([np.array(self.s), np.array(self.total_wake)])
# np.savetxt("trajectory_cos.txt", self.csr_traj)
return self.csr_traj
SF = Sextupole(l=0.0001, k2=17673.786254063251 * 0.0, eid="SF")
SD = Sextupole(l=0.0001, k2=-36169.817233025707 * 0.0, eid="SD")
Q1 = Quadrupole(l=0.293, k1=2.62, eid="Q1")
Q2 = Quadrupole(l=0.293, k1=-3.1, eid="Q2")
Q3 = Quadrupole(l=0.327, k1=2.8, eid="Q3")
Q4 = Quadrupole(l=0.291, k1=-3.7, eid="Q4")
Q5 = Quadrupole(l=0.391, k1=4.0782, eid="Q5")
Q6 = Quadrupole(l=0.291, k1=-3.534859, eid="D6")
B1 = SBend(l=0.23, angle=0.23 / 19.626248, eid="B1")
B2 = SBend(l=1.227, angle=1.227 / 4.906312, eid="B2")
U = Undulator(lperiod=0.02, nperiods=50, Kx=2)
U.mag_field = lambda x, y, z: und_field(x, y, z, U.lperiod, U.Kx)
U.npoints = 500
"""pytest fixtures description"""
@pytest.fixture(scope='module')
def cell_u():
return (D1, SF, D2, Q1, D3, Q2, D2, SD, D4, B1, B2, D5, Q3,
D5, B2, B1, D6, Q4, D7, Q5, D8, Q6, Drift(l=1.0045), U,
Drift(l=1.0045), Q6, D8, Q5, D7, Q4, D6, B1, B2, D5, Q3, D5, B2,
B1, D4, SD, D2, Q2, D3, Q1, D2, SF, D1)
@pytest.fixture(scope='module')
def cell():
return (D1, SF, D2, Q1, D3, Q2, D2, SD, D4, B1, B2, D5, Q3, D5, B2, B1, D6,