lundulator = 2.27e-2
lquad = 2*lundulator
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
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
print u'\u03B2_y =', betay
#print u'<\u03B2_x> =', (2*ldrift1+ldrift2+4*lquad)/1.7
#print u'<\u03B2_y> =', (2*ldrift1+ldrift2+4*lquad)/1.25
print u'\u03B1_x =', alphax
