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()
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
print u'\u03B1_y =', alphay
print 'X radius =', np.sqrt(1.2e-6*betax/1202.5)
print 'Y radius =', np.sqrt(1.2e-6*betay/1202.5)
print 'transfer map =\n', transfermap
betax, betay, s = my_lattice.compute_beta_func()
plt.plot(s, betax, c='r', label=r'$\beta_x$')
plt.plot(s, betay, c='g', label=r'$\beta_y$')
plt.ylabel(r'$\beta~\textrm{[m/rad]}$')
plt.xlabel(r'$s~\textrm{[m]}$')
plt.ylim([0., 1.5*max(betax)])
plt.xlim([0., s[-1]])
plt.legend()
plt.tight_layout()
plt.show()
my_lattice.export_genesis_lattice(0.0227, 1202.5)
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
print u'\u03B1_y =', alphay
print 'X radius =', np.sqrt(1.2e-6*betax/1202.5)
print 'Y radius =', np.sqrt(1.2e-6*betay/1202.5)
print 'transfer map =\n', transfermap
betax, betay, s = my_lattice.compute_beta_func()
#plt.plot(s, betax, c='r', label=r'$\beta_x$')
#plt.plot(s, betay, c='g', label=r'$\beta_y$')
#plt.ylabel(r'$\beta~\textrm{[m/rad]}$')
#plt.xlabel(r'$s~\textrm{[m]}$')
#plt.ylim([0., 1.5*max(betax)])
#plt.xlim([0., s[-1]])
#plt.legend()
#plt.tight_layout()
#plt.show()
my_lattice.export_genesis_lattice(l_undulator, gamma0)
run_genesis = rbExecuteGenesis()
run_genesis.execute_genesis('ttf.in')
