def __init__(self):
## Beam properties
self.beamdata = self.getDefaultBeamdata()
self.twiss = self.getDefaultTwiss()
self.multipart = gaussianTwiss3D(self.beamdata[5], self.twiss)
self.envelope = envelopeFromMultipart(self.multipart)
# empty lattice
self.lattice = Lattice("Facility", self.beamdata, self.twiss, self.multipart)
def setMultipart(self,multipart): # Also updates the envelope
self.lattice.setMultipart(multipart)
self.multipart = self.lattice.getMultipart()
self.envelope = envelopeFromMultipart(self.multipart)
def plotEverything(multipartin,twiss,multipartout, envlist, multipartafterall):#,envx,envy):
xin = [multipartin[i][0][0] for i in xrange(len(multipartin))]
xpin = [multipartin[i][0][1] for i in xrange(len(multipartin))]
yin = [multipartin[i][0][2] for i in xrange(len(multipartin))]
ypin = [multipartin[i][0][3] for i in xrange(len(multipartin))]
zin = [multipartin[i][0][4] for i in xrange(len(multipartin))]
zpin = [multipartin[i][0][5] for i in xrange(len(multipartin))]
alpha_x = twiss[0]
beta_x = twiss[1]
epsilon_rms_x = twiss[2]
alpha_y = twiss[3]
beta_y = twiss[4]
epsilon_rms_y = twiss[5]
alpha_z = twiss[6]
beta_z = twiss[7]
epsilon_rms_z = twiss[8]
envx = [envlist[i][0][0] for i in xrange(len(envlist))]
envy = [envlist[i][0][3] for i in xrange(len(envlist))]
envz = [envlist[i][0][6] for i in xrange(len(envlist))]
s = [envlist[i][1] for i in xrange(len(envlist))]
## begin ellipse, see appendix C in Wille
# initial
gamma_x = (1+alpha_x**2)/beta_x
#sigma_xp = np.sqrt(epsilon_rms_x*gamma_x)# b, from fig 3.23
#sigma_x = np.sqrt(epsilon_rms_x*beta_x) # a
sigma_xp = np.sqrt(envlist[0][0][2]) # sqrt since env has sigma**2. 3*sigma will cover 99.7%.
sigma_x = np.sqrt(envx[0])
three_sigma_x = 3*sigma_x
three_sigma_xp = 3*sigma_xp
angle_x = angleFrom_i_and_ip(xin, xpin) # in radians
ellipse_x = Ellipse((0,0), 2*three_sigma_x, 2*three_sigma_xp, angle_x*180/constants.pi) # 2* is because matplotlibs height is from bottom to top and not center to top...
ellipse_x.set_facecolor('none')
ellipse_x.set_edgecolor((0,0,0))
gamma_y = (1+alpha_y**2)/beta_y
#sigma_yp = np.sqrt(epsilon_rms_y*gamma_y)# b, from fig 3.23
#sigma_y = np.sqrt(epsilon_rms_y*beta_y) # a
sigma_yp = np.sqrt(envlist[0][0][5]) # sqrt since env has sigma**2. 3*sigma will cover 99.7%.
sigma_y = np.sqrt(envy[0])
three_sigma_y = 3*sigma_y
three_sigma_yp = 3*sigma_yp
angle_y = angleFrom_i_and_ip(yin, ypin)
ellipse_y = Ellipse((0,0), 2*three_sigma_y, 2*three_sigma_yp, angle_y*180/constants.pi) # 2* is because matplotlibs height is from bottom to top and not center to top...
ellipse_y.set_facecolor('none')
ellipse_y.set_edgecolor((0,0,0))
ellipse_xy = Ellipse((0,0), 2*three_sigma_x, 2*three_sigma_y, 0)
ellipse_xy.set_facecolor('none')
ellipse_xy.set_edgecolor((0,0,0))
gamma_z = (1+alpha_z**2)/beta_z
sigma_zp = np.sqrt(envlist[0][0][8]) # sqrt since env has sigma**2. 3*sigma will cover 99.7%.
sigma_z = np.sqrt(envz[0])
## end initial ellipses
xo = [multipartout[i][0][0] for i in xrange(len(multipartout))]
xpo = [multipartout[i][0][1] for i in xrange(len(multipartout))]
yo = [multipartout[i][0][2] for i in xrange(len(multipartout))]
ypo = [multipartout[i][0][3] for i in xrange(len(multipartout))]
zo = [multipartout[i][0][4] for i in xrange(len(multipartout))]
zpo = [multipartout[i][0][5] for i in xrange(len(multipartout))]
xallo = [multipartafterall[i][0][0] for i in xrange(len(multipartafterall))]
xpallo = [multipartafterall[i][0][1] for i in xrange(len(multipartafterall))]
#print "xallo: \n" + str(xallo)
#print "xpallo: \n" + str(xpallo)
#print "xo: \n" + str(xo)
#print "xpo: \n" + str(xpo)
# Cull inf values
#if len(xo) > 1:
# xo[xo > 1e6] = 0
# xpo[xpo > 1e6] = 0
# xallo[xo > 1e6] = 0
# xpallo[xpo > 1e6] = 0
# after ellipses
envelopeAfter = envelopeFromMultipart(multipartout)
#angle_x_after = angleFrom_i_and_ip(xo, xpo)
sigma_xTimessigma_xp_after = envelopeAfter[1]
sigma_xp_after = np.sqrt(envelopeAfter[2]) # sqrt since env has sigma**2. 3*sigma will cover 99.7%.
sigma_x_after = np.sqrt(envelopeAfter[0])
three_sigma_x_after = 3*sigma_x_after
three_sigma_xp_after = 3*sigma_xp_after
#ellipse_x_after = Ellipse((0,0), 2*three_sigma_x_after, 2*three_sigma_xp_after, angle_x_after*180/constants.pi) # 2* is because matplotlibs height is from bottom to top and not center to top....
#ellipse_x_after.set_facecolor('none')
#ellipse_x_after.set_edgecolor((0,0,0))
# y
#angle_y_after = angleFrom_i_and_ip(yo, ypo)
sigma_yTimessigma_yp_after = envelopeAfter[4]
sigma_yp_after = np.sqrt(envelopeAfter[5]) # sqrt since env has sigma**2. 3*sigma will cover 99.7%.
sigma_y_after = np.sqrt(envelopeAfter[3])
three_sigma_y_after = 3*sigma_y_after
three_sigma_yp_after = 3*sigma_yp_after
#ellipse_y_after = Ellipse((0,0), 2*three_sigma_y_after, 2*three_sigma_yp_after, angle_y_after*180/constants.pi) # 2* is because matplotlibs height is from bottom to top and not center to top....
#ellipse_y_after.set_facecolor('none')
#ellipse_y_after.set_edgecolor((0,0,0))
# z
#angle_z_after = angleFrom_i_and_ip(zo, zpo)
sigma_zTimessigma_zp_after = envelopeAfter[7]
sigma_zp_after = np.sqrt(envelopeAfter[8]) # sqrt since env has sigma**2. 3*sigma will cover 99.7%.
sigma_z_after = np.sqrt(envelopeAfter[6])
three_sigma_z_after = 3*sigma_z_after
three_sigma_zp_after = 3*sigma_zp_after
#ellipse_z_after = Ellipse((0,0), 2*three_sigma_z_after, 2*three_sigma_zp_after, angle_z_after*180/constants.pi) # 2* is because matplotlibs height is from bottom to top and not center to top....
#ellipse_z_after.set_facecolor('none')
#ellipse_z_after.set_edgecolor((0,0,0))
# xy
ellipse_xy_after = Ellipse((0,0), 2*three_sigma_x_after, 2*three_sigma_y_after, 0)
ellipse_xy_after.set_facecolor('none')
ellipse_xy_after.set_edgecolor((0,0,0))
# end after ellipses
plt.figure(0)
ax1 = plt.subplot2grid((4,4), (0,0))
#ax1.add_artist(ellipse_x)
ax1.plot(xin,xpin,'ro', zorder=1)
if not sigma_x == 0:
ax1.set_xlim(-5*sigma_x, 5*sigma_x)
if not sigma_xp == 0:
ax1.set_ylim(-5*sigma_xp, 5*sigma_xp)
plt.title('Initial values in x')
plt.xlabel('x [m]')
plt.ylabel('x\' []')
ax2 = plt.subplot2grid((4,4), (0,1))
#ax2.add_artist(ellipse_y)
ax2.plot(yin,ypin,'bo', zorder=1)
if not sigma_y == 0:
ax2.set_xlim(-5*sigma_y, 5*sigma_y)
if not sigma_yp == 0:
ax2.set_ylim(-5*sigma_yp, 5*sigma_yp)
plt.title('Initial values in y')
plt.xlabel('y [m]')
plt.ylabel('y\' []')
ax3 = plt.subplot2grid((4,4), (0, 2))
#ax8.add_artist(ellipse_y_after)
ax3.plot(zin,zpin,'yo', zorder=1)
if not sigma_z == 0:
ax3.set_xlim(-5*sigma_z, 5*sigma_z)
if not sigma_zp == 0:
ax3.set_ylim(-5*sigma_zp, 5*sigma_zp)
#ax3.set_xlim(-4e-3, 4e-3)
#ax3.set_ylim(-4e-3, 4e-3)
plt.title('Values before lattice in z')
plt.xlabel('z [m]')
plt.ylabel('$\delta$ []')
ax4 = plt.subplot2grid((4,4), (0,3))
ax4.add_artist(ellipse_xy)
ax4.plot(xin,yin,'go', zorder=1)
if not sigma_x == 0:
ax4.set_xlim(-5*sigma_x, 5*sigma_x)
if not sigma_y == 0:
ax4.set_ylim(-5*sigma_y, 5*sigma_y)
plt.title('Initial values x and y')
plt.xlabel('x [m]')
plt.ylabel('y [m]')
ax5 = plt.subplot2grid((4,4), (1, 0), colspan=3)
ax5.plot(s,envx,'ro')
plt.title('Envelope in sigma_x^2 by s')
plt.xlabel('s [m]')
plt.ylabel('Envelope in sigma_x^2')
ax6 = plt.subplot2grid((4,4), (2, 0), colspan=3)
ax6.plot(s,envy,'bo')
plt.title('Envelope in sigma_y^2 by s')
plt.xlabel('s [m]')
plt.ylabel('Envelope in sigma_y^2')
ax7 = plt.subplot2grid((4,4), (2, 3))
#ax7.add_artist(ellipse_x_after)
ax7.plot(xallo,xpallo,'r.', zorder=1)
if not sigma_x_after == 0:
ax7.set_xlim(-5*sigma_x_after, 5*sigma_x_after)
if not sigma_xp_after == 0:
ax7.set_ylim(-5*sigma_xp_after, 5*sigma_xp_after)
#ax7.set_xlim(-0.004, 0.004)
#ax7.set_ylim(-0.004, 0.004)
plt.title('Values after all elems in lattice in x')
plt.xlabel('x [m]')
plt.ylabel('x\' []')
ax8 = plt.subplot2grid((4,4), (3, 0))
#ax8.add_artist(ellipse_x_after)
ax8.plot(xo,xpo,'ro', zorder=1)
if not sigma_x_after == 0:
ax8.set_xlim(-5*sigma_x_after, 5*sigma_x_after)
if not sigma_xp_after == 0:
ax8.set_ylim(-5*sigma_xp_after, 5*sigma_xp_after)
#ax8.set_xlim(-0.004, 0.004)
#ax8.set_ylim(-0.004, 0.004)
plt.title('Values after lattice in x')
plt.xlabel('x [m]')
plt.ylabel('x\' []')
ax9 = plt.subplot2grid((4,4), (3, 1))
#ax9.add_artist(ellipse_y_after)
ax9.plot(yo,ypo,'bo', zorder=1)
if not sigma_y_after == 0:
ax9.set_xlim(-5*sigma_y_after, 5*sigma_y_after)
if not sigma_yp_after == 0:
ax9.set_ylim(-5*sigma_yp_after, 5*sigma_yp_after)
plt.title('Values after lattice in y')
plt.xlabel('y [m]')
plt.ylabel('y\' []')
ax10 = plt.subplot2grid((4,4), (3, 2))
#ax10.add_artist(ellipse_y_after)
ax10.plot(zo,zpo,'yo', zorder=1)
if not sigma_z_after == 0:
ax10.set_xlim(-5*sigma_z_after, 5*sigma_z_after)
if not sigma_zp_after == 0:
ax10.set_ylim(-5*sigma_zp_after, 5*sigma_zp_after)
plt.title('Values after lattice in z')
plt.xlabel('z [m]')
plt.ylabel('$\delta$ []')
ax11 = plt.subplot2grid((4,4), (3, 3))
ax11.add_artist(ellipse_xy_after)
ax11.plot(xo,yo,'go', zorder=1)
if not sigma_x_after == 0:
ax11.set_xlim(-5*sigma_x_after, 5*sigma_x_after)
if not sigma_y_after == 0:
ax11.set_ylim(-5*sigma_y_after, 5*sigma_y_after)
plt.title('Values after lattice in y and x')
plt.xlabel('x [m]')
plt.ylabel('y [m]')
plt.suptitle("Plots")
left = 0.125 # the left side of the subplots of the figure
right = 0.9 # the right side of the subplots of the figure
bottom = 0.1 # the bottom of the subplots of the figure
top = 0.9 # the top of the subplots of the figure
wspace = 0.2 # the amount of width reserved for blank space between subplots
hspace = 0.4 # the amount of height reserved for white space between subplots
plt.subplots_adjust(left, bottom, right, top, wspace, hspace) # Increases the space between subplots
plt.show()
#multipartfromold = gaussianTwiss3D(nbrOfParticles, twissfromold)
#print "multipartfromoldcopy: \n" + str(multipartfromoldcopy)
spaceChargeOnInComp = 0
E = 2e9*constants.e # 2GeV to joule from ref F.
freq = 704.42e6 # (Hz) from ref. F
rf_lambda = constants.c/freq # beam data needed
m = constants.m_p
beta = betaFromE(m, E)
q = constants.e
beamdata = [beta, rf_lambda, m, q, E]
#envelopeInComp = np.array([1, 0, 0, 1, 0, 0, 1, 0, 0])
envelopeInComp = envelopeFromMultipart(multipartfromold)
#print "envelopeInComp: " + str(envelopeInComp)
nbrOfSplits = 1
## the lattice will be a FODSO cell (Focusing Quad, Drift, Defocusing Quad, Sextupole, Drift)
compLattice = Lattice('compLattice', beamdata, twissfromold, multipartfromold)
cavityName = "cavity"
cavityLength = 2.0
cavityOscillations = 2
cavityAmplitudeA = 0
cavityAmplitudeB = 30 # 30 MeV / m
cavityE_0 = cavityAmplitudeB
cavitySigma = 1
cavityP = 3