print 'lquad =', lquad
kq = 1/(f*lquad)
print 'kq =', kq
ldrift1 = 10*lundulator
print 'l =', ldrift1
ldrift2 = 90*lundulator
print 'L =', ldrift2
n_cells = 1
pos = 0.
for x in range(0, n_cells):
my_lattice.add_quadrupole(kq, pos, lquad)
pos += lquad
#my_lattice.add_drift(pos, ldrift1)
my_lattice.add_undulator(0.849, pos, ldrift1)
pos += ldrift1
my_lattice.add_quadrupole(-kq, pos, lquad)
pos += lquad
my_lattice.add_undulator(0.849, pos, ldrift2)
pos += ldrift2
my_lattice.add_quadrupole(-kq, pos, lquad)
pos += lquad
#my_lattice.add_drift(pos, ldrift1)
my_lattice.add_undulator(0.849, pos, ldrift1)
pos += ldrift1
my_lattice.add_quadrupole(kq, pos, lquad)
pos += lquad
betax, alphax, betay, alphay, transfermap = my_lattice.compute_beamline()
emittance = 2.000000e-6 #2 mm-mrad
l_undulator = 2.73e-2
dBdx_quad = 37. #T/m
k_quadrupole = dBdx_quad*consts.e/(consts.m_e*gamma0*consts.c)
aw = 0.896
l_quad = 3*l_undulator
l_drift1 = 16*l_undulator
l_drift2 = 2*l_undulator
l_und_segment = 163*l_undulator
und_spacing = 26*l_undulator
n_cells = 6
pos = 0.
# add quads and undulators
for x in range(0, n_cells):
my_lattice.add_undulator(0.896,pos, l_und_segment)
pos += l_und_segment
my_lattice.add_drift(pos, l_drift2)
pos += l_drift2
my_lattice.add_quadrupole(k_quadrupole, pos, l_quad)
pos += l_quad
my_lattice.add_drift(pos, l_drift1)
pos += l_drift1
my_lattice.add_quadrupole(-k_quadrupole, pos, l_quad)
pos += l_quad
my_lattice.add_drift(pos, l_drift2)
pos += l_drift2
betax, alphax, betay, alphay, transfermap = my_lattice.compute_beamline()
print u'\u03B2_x =', betax
