beam.py

from __future__ import division
import numpy as np
from sys import exit

class beam():
	"""
	The python class of G4Beamline beam array;
	"""
	def __init__(self):
		#
		np.seterr(divide='raise',invalid='raise')
		#
		self.X, self.Y = [np.zeros(2),np.array([[1,0],[0,1]])],[np.zeros(2),np.array([[1,0],[0,1]])]
		self.T, self.Plimits = np.zeros(2),np.zeros(2)
		self.Xset, self.Yset, self.Tset, self.Pset, self.Weightset = [0,]*5
		self.numberOfparticles = 1000
		self.numberSet = 0
		self.PDGset = 0
		self.dataSet = 0
		# Header of the G4BL beam file:
		self.beamFileHeader = "BLTrackFile This is a BLTrackfile generated by frankliuao' program, with the extension name .beam;\n"+ \
		                      "x y z Px Py Pz t PDGid EvNum TrkId Parent weight\n"+ \
		                      "mm mm mm MeV/c MeV/c MeV/c ns - - - - -"

	def fitBeam(self,rmsEmitX=None,rmsEmitY=None,numBins=100,logFileDir = './'):
		"""Fit the beam phase space distribution with a 2D Gaussian.
		   The unit for rmsEmit is m*rad
		   Usage:
		        beamExp.fitBeam(rmsEmit=2e-3)
		"""
		print "The default logFileDir is the current Python working directory, unless specified by fitBeam(logFileDir='PathToDirToPutTheLog') otherwise."

		#==== First make sure the beam is set, and then get the statistics and set the number of bins: ====
		#== Set the logfile name and where to put it.
		if self.dataSet:
			self.g4blBeamMat = np.copy(self.data)
			self.g4blBeamMat[:,0:3] /= 1000
			self.rmsEmitX = rmsEmitX if rmsEmitX is not None else None
			self.rmsEmitY = rmsEmitY if rmsEmitY is not None else None
		else:
			self.genError('---Error: The beam is not set yet, can not fit the beam ---')
		self.numBins = numBins
		self.logFileName = logFileDir+'fitLog.txt'
		#==== ! First make sure the beam is set, and then get the statistics and set the number of bins: ! ====

		def getTWISS():
			'''Get the TWISS from the variance matrix'''
			paramListX,paramListY = self.xStat,self.yStat
			self.varMatX = np.array([[paramListX[1]**2,paramListX[4]*paramListX[1]*paramListX[3]],
			                            [paramListX[4]*paramListX[1]*paramListX[3],paramListX[3]**2]])
			self.varMatY = np.array([[paramListY[1]**2,paramListY[4]*paramListY[1]*paramListY[3]],
			                   [paramListY[4]*paramListY[1]*paramListY[3],paramListY[3]**2]])
			self.varMatDetX = np.linalg.det(self.varMatX)
			self.varMatDetY = np.linalg.det(self.varMatY)
			TWISSmatX = self.varMatX/np.sqrt(self.varMatDetX)
			TWISSmatY = self.varMatY/np.sqrt(self.varMatDetY)
			self.betaX, self.alphaX, self.gammaX = TWISSmatX[0,0], -1*TWISSmatX[0,1], TWISSmatX[1,1]
			self.betaY, self.alphaY, self.gammaY= TWISSmatY[0,0], -1*TWISSmatY[0,1], TWISSmatY[1,1]

		def getStat():
			"""Get the statistics of the current beam being fitted
			"""
			self.sumWeight = np.sum(self.g4blBeamMat[:,11])
			self.xPrime, self.yPrime = self.g4blBeamMat[:,3]/self.g4blBeamMat[:,5], self.g4blBeamMat[:,4]/self.g4blBeamMat[:,5]
			self.xMean = np.sum(self.g4blBeamMat[:,0]*self.g4blBeamMat[:,11])/self.sumWeight
			self.yMean = np.sum(self.g4blBeamMat[:,1]*self.g4blBeamMat[:,11])/self.sumWeight
			self.xPrimeMean = np.sum(self.xPrime*self.g4blBeamMat[:,11])/self.sumWeight
			self.yPrimeMean = np.sum(self.yPrime*self.g4blBeamMat[:,11])/self.sumWeight
			self.xSigma = np.sqrt(np.sum((self.g4blBeamMat[:,0]-self.xMean)**2*self.g4blBeamMat[:,11])/self.sumWeight)
			self.ySigma = np.sqrt(np.sum((self.g4blBeamMat[:,1]-self.yMean)**2*self.g4blBeamMat[:,11])/self.sumWeight)
			self.xPrimeSigma = np.sqrt(np.sum((self.xPrime-self.xPrimeMean)**2*self.g4blBeamMat[:,11])/self.sumWeight)
			self.yPrimeSigma = np.sqrt(np.sum((self.yPrime-self.yPrimeMean)**2*self.g4blBeamMat[:,11])/self.sumWeight)
			self.xRho = np.sum((self.g4blBeamMat[:,0]-self.xMean)*(self.xPrime-self.xPrimeMean)*self.g4blBeamMat[:,11])/self.sumWeight/self.xSigma/self.xPrimeSigma
			self.yRho = np.sum((self.g4blBeamMat[:,1]-self.yMean)*(self.yPrime-self.yPrimeMean)*self.g4blBeamMat[:,11])/self.sumWeight/self.ySigma/self.yPrimeSigma
			#
			self.xStat = [self.xMean,self.xSigma,self.xPrimeMean,self.xPrimeSigma,self.xRho]
			self.yStat = [self.yMean,self.ySigma,self.yPrimeMean,self.yPrimeSigma,self.yRho]

		def getDist():
			'''Get the histogram of the beam in the phase space,
			   and the coordinates of the bin centers of the phase space axes.
			'''
			self.xxPrimeCount, xEdge, xPrimeEdge = np.histogram2d(self.g4blBeamMat[:,0],self.xPrime,
			                                                      bins=self.numBins,normed=True,weights=self.g4blBeamMat[:,11])
			self.yyPrimeCount, yEdge, yPrimeEdge = np.histogram2d(self.g4blBeamMat[:,1],self.yPrime,
			                                                      bins=self.numBins,normed=True,weights=self.g4blBeamMat[:,11])
			self.xCenter, self.xPrimeCenter = (xEdge[1:]+xEdge[0:-1])/2, (xPrimeEdge[1:]+xPrimeEdge[0:-1])/2
			self.yCenter, self.yPrimeCenter = (yEdge[1:]+yEdge[0:-1])/2, (yPrimeEdge[1:]+yPrimeEdge[0:-1])/2

		def GaussNewton():
			'''Gauss Newton method to fit the statVec
			'''
			# Clear the variables from last call of GaussNewton()
			try:
				del(self.DeltaParamListX,self.DeltaParamListY)
			except AttributeError:
				pass

			# For algorithm's convenience, convert the paramList to 5*1 array;
			# [[\mu_u],[\sigma_u],[\mu_{u'}],[\sigma_{u'}],[\rho_{u,u'}]]
			#
			self.paramListX = np.copy(self.xStat).reshape(5,1)
			self.paramListY = np.copy(self.yStat).reshape(5,1)
			def residual(coordinateNow,coordinatePrimeNow,zNow,dimension):
				'''Calculate the residual in each bin
				'''
				paramListLocal = np.copy(self.paramListX) if dimension.lower() == 'x' else np.copy(self.paramListY)
				functionValue = 1/(2*np.pi*paramListLocal[1,0]*paramListLocal[3,0]*np.sqrt(1-paramListLocal[4,0]**2))*\
				                      np.exp(-1*((coordinateNow-paramListLocal[0,0])**2/(paramListLocal[1,0]**2)-
				                      2*paramListLocal[4,0]*(coordinateNow-paramListLocal[0,0])*(coordinatePrimeNow-paramListLocal[2,0])/paramListLocal[1,0]/paramListLocal[3,0]+
				                      (coordinatePrimeNow-paramListLocal[2,0])**2/paramListLocal[3,0]**2)/2/(1-paramListLocal[4,0]**2))
				residualValue = zNow - functionValue
				return residualValue

			def JacobianList(coordinateNow,coordinatePrimeNow,dimension):
				'''Jacobian of the Gauss Newton method
				'''
				paramListLocal = np.copy(self.paramListX) if dimension.lower() == 'x' else np.copy(self.paramListY)
				functionValue = 1/(2*np.pi*paramListLocal[1,0]*paramListLocal[3,0]*np.sqrt(1-paramListLocal[4,0]**2))*\
				                      np.exp(-1*((coordinateNow-paramListLocal[0,0])**2/(paramListLocal[1,0]**2)-
				                      2*paramListLocal[4,0]*(coordinateNow-paramListLocal[0,0])*(coordinatePrimeNow-paramListLocal[2,0])/paramListLocal[1,0]/paramListLocal[3,0]+
				                      (coordinatePrimeNow-paramListLocal[2,0])**2/paramListLocal[3,0]**2)/2/(1-paramListLocal[4,0]**2))
				JacobianVec = [-1*functionValue*(-1/(1-paramListLocal[4,0]**2)*
				                (-1*(coordinateNow-paramListLocal[0,0])/paramListLocal[1,0]**2+paramListLocal[4,0]*(coordinatePrimeNow-paramListLocal[2,0])/paramListLocal[1,0]/paramListLocal[3,0])),
								# Checked;
				                -1*functionValue*(-1/paramListLocal[1,0]-1/(1-paramListLocal[4,0]**2)*
				                                  (-1*(coordinateNow-paramListLocal[0,0])**2/paramListLocal[1,0]**3+
				                                    paramListLocal[4,0]*(coordinateNow-paramListLocal[0,0])*(coordinatePrimeNow-paramListLocal[2,0])/paramListLocal[1,0]**2/paramListLocal[3,0])),
								# Checked;
				                -1*functionValue*(-1/(1-paramListLocal[4,0]**2)*
				                (-1*(coordinatePrimeNow-paramListLocal[2,0])/paramListLocal[3,0]**2+paramListLocal[4,0]*(coordinateNow-paramListLocal[0,0])/paramListLocal[1,0]/paramListLocal[3,0])),
								# Checked;
				                -1*functionValue*(-1/paramListLocal[3,0]-1/(1-paramListLocal[4,0]**2)*
				                                  (-1*(coordinatePrimeNow-paramListLocal[2,0])**2/paramListLocal[3,0]**3+
				                                    paramListLocal[4,0]*(coordinateNow-paramListLocal[0,0])*(coordinatePrimeNow-paramListLocal[2,0])/paramListLocal[1,0]/paramListLocal[3,0]**2)),
								# Checked;
				                -1*functionValue*(paramListLocal[4,0]/(1-paramListLocal[4,0]**2)-
				                             paramListLocal[4,0]/(1-paramListLocal[4,0]**2)**2*((coordinateNow-paramListLocal[0,0])**2/(paramListLocal[1,0]**2)-
				                      2*paramListLocal[4,0]*(coordinateNow-paramListLocal[0,0])*(coordinatePrimeNow-paramListLocal[2,0])/paramListLocal[1,0]/paramListLocal[3,0]+
				                      (coordinatePrimeNow-paramListLocal[2,0])**2/paramListLocal[3,0]**2)+
				                               1/(1-paramListLocal[4,0]**2)*(coordinateNow-paramListLocal[0,0])*(coordinatePrimeNow-paramListLocal[2,0])/paramListLocal[1,0]/paramListLocal[3,0])
				                # Checked;
				                ]
				return JacobianVec

			def updateParamList(dimension):
				'''Update the paramListU to better fit the beam.
				'''
				if dimension.lower() == 'x':
					coordinate = self.xCenter
					coordinatePrime = self.xPrimeCenter
					histogram = self.xxPrimeCount
					paramList = self.paramListX
				else:
					coordinate = self.yCenter
					coordinatePrime = self.yPrimeCenter
					histogram = self.yyPrimeCount
					paramList = self.paramListY

				# Calculate the Jacobian Matrix, the delta vector, update the paramList
				iterations_GaussNewton = 1

				while True:
					# Residual r matrix where r_{i,j} = ln(z(u_i,uPrime_j))-ln(f(u_i,uPrime_j,paramList))
					rMat = np.zeros((self.numBins,self.numBins))
					for ii in xrange(self.numBins):
						for jj in xrange(self.numBins):
							coordinateNow, coordinatePrimeNow, zNow = coordinate[ii],coordinatePrime[jj],histogram[ii,jj]
							rMat[ii,jj] = residual(coordinateNow,coordinatePrimeNow,zNow,'x') if dimension.lower()=='x' else residual(coordinateNow,coordinatePrimeNow,zNow,'y')
					# Vectorize the residual matrix:
					rvec = rMat.reshape(self.numBins**2,1)
					# If the residuals are not decreasing, subtract the last deltaParamList and break;
					try:
						if (np.std(rvec)-rvecStd)/rvecStd <= -1e-4:
							rvecStd = np.std(rvec)
						else:
							paramList += DeltaParamList
							break
					except UnboundLocalError:
						rvecStd = np.std(rvec)
					# Jacobian matrix, the ij-th element is (\partial r_i)/(\partial param_j)
					JacobianMat = [0,]*self.numBins**2
					kk = 0
					for ii in xrange(self.numBins):
						for jj in xrange(self.numBins):
							JacobianMat[kk] = JacobianList(coordinate[ii],coordinatePrime[jj],'x') if dimension.lower()=='x' \
								else JacobianList(coordinate[ii],coordinatePrime[jj],'y')
							kk += 1
					JacobianMat = np.array(JacobianMat)
					JacobianMatTran = JacobianMat.transpose()

					# The change of paramList in this iteration:
					try:
						DeltaParamList = np.dot(np.dot(np.linalg.inv(np.dot(JacobianMatTran,JacobianMat)),JacobianMatTran),rvec)
					except Exception as e:
						print e
						print 'Error in GaussNewton -- Could not calculate the paramList change.'
						return
						exit()
					# Get the relative change amount:
					paramError = np.std(DeltaParamList/paramList)

					# If the change makes the parameters out of the boundaries, break:
					if paramError >= 4e-2 and iterations_GaussNewton!=100:
						if paramList[1,0]-DeltaParamList[1,0]>0 and paramList[3,0]-DeltaParamList[3,0]>0 and \
							abs(paramList[4,0]-DeltaParamList[4,0])<1:
							paramList -= DeltaParamList
							iterations_GaussNewton += 1
						else:
							break
					else:
						break

			updateParamList('x')
			updateParamList('y')

			# Finally, need to update the new parameter list to the x(y)StatVec for later use:
			self.xStat,self.yStat = list(self.paramListX.reshape(5)), list(self.paramListY.reshape(5))

		def reduceEmit():
			'''Reduce the emittance of the phase space.
			'''
			self.logFileHandle = open(self.logFileName,'w')
			self.logFileCont = '# Iteration Number (col 1); MuX (col 2); BetaX (col 3); AlphaX (col 4); EmitX (col 5); ' + \
		                       "MuY (col 6); BetaY (col 7); AlphaY (col 8); EmitY (col 9) \n" + '#Next line is the number of pions in ' + \
		                       'the acceptance.\n'
			# Get the first guess of the Gaussian parameters:
			getStat()
			self.xPrimeInit, self.yPrimeInit = np.copy(self.xPrime), np.copy(self.yPrime)
			# Get the probability density function in the phase space and get the TWISS matrix
			getDist()
			getTWISS()
			# Use the Gauss Newton method to fit the parameters, update the statVec
			GaussNewton()
			# Start to cut the beam to reduce the RMS emittance.
			iteration = 1
			# Get the TWISS and emittance again using the new statVec
			getTWISS()
			# If self.rmsEmitX or self.rmsEmitY, or both were set by the user, keep using the values desired by the user;
			# else, use the rms emit from the first Gauss-Newton fitting iteration.
			self.rmsEmitX = np.sqrt(self.varMatDetX) if self.rmsEmitX is None else self.rmsEmitX
			self.rmsEmitY = np.sqrt(self.varMatDetY) if self.rmsEmitY is None else self.rmsEmitY
			# Records how much more to reduce the emit;
			self.moreX, self.moreY = self.varMatDetX/self.rmsEmitX**2, self.varMatDetY/self.rmsEmitY**2
			#
			while self.moreX>1 or self.moreY>1 :
				if self.moreX>1 and self.moreY<=1:
					particleEmit = self.gammaX*self.g4blBeamMat[:,0]**2+2*self.alphaX*self.g4blBeamMat[:,0]*self.xPrime+self.betaX*self.xPrime**2
					self.more = self.moreX if self.moreX<=20 else 20
				elif self.moreX<=1 and self.moreY>1:
					particleEmit = self.gammaY*self.g4blBeamMat[:,1]**2+2*self.alphaY*self.g4blBeamMat[:,1]*self.yPrime+self.betaY*self.yPrime**2
					self.more = self.moreY if self.moreY<=20 else 20
				else:
					# ParticleEmit for both of the directions:
					particleEmitX = self.gammaX*self.g4blBeamMat[:,0]**2+2*self.alphaX*self.g4blBeamMat[:,0]*self.xPrime+self.betaX*self.xPrime**2
					particleEmitY = self.gammaY*self.g4blBeamMat[:,1]**2+2*self.alphaY*self.g4blBeamMat[:,1]*self.yPrime+self.betaY*self.yPrime**2
					# Average Emit:
					# Get the difference of emittance for each particle and RMS emittance:
					# x and y:
					deltaEmitX, deltaEmitY = (particleEmitX)/self.rmsEmitX, (particleEmitY)/self.rmsEmitY
					# mean:
					particleEmit = np.mean([deltaEmitX,deltaEmitY],axis=0)
					#
					self.more = np.max((self.moreX,self.moreY))
					#
					self.more = self.more if self.more<=20 else 20

				indexEmit = np.arange(particleEmit.shape[0])
				# Sort the Emittance from large to small:
				indexEmit = indexEmit[np.argsort(-1*particleEmit[:])]
				# Get each index's corresponding weight:
				weightSorted = self.g4blBeamMat[:,11][indexEmit]
				# Get the cumulative sum of their weight, from the largest entry to the last:
				weightCumSum = np.cumsum(weightSorted)
				# We will discard the 0.1*(self.more+5) percent of the particles:
				indexKept = indexEmit[weightCumSum>=(self.more+10)*1e-3*self.sumWeight] \
					if (self.more+10)*1e-3*self.sumWeight>weightCumSum[0] else indexEmit[1:]
				# Keep only the portion that has smaller emittance.
				self.g4blBeamMat = self.g4blBeamMat[indexKept,:]
				# Break if there are no more particles left:
				if self.data.shape[0]: break
				# Update the statistics of the new beam:
				getStat()
				getDist()
				getTWISS()
				GaussNewton()
				getTWISS()
				self.moreX, self.moreY = self.varMatDetX/self.rmsEmitX**2, self.varMatDetY/self.rmsEmitY**2
				iteration += 1

			#
			if self.data.shape[0]==0 :
				self.logFileCont += '%s %.5e %.5e %.5e %.5e %.5e %.5e %.5e %.5e\n'\
				                         %(0,0,0,0,0,0,0,0,0) + '0'
				self.logFileHandle.write(self.logFileCont)
				self.logFileHandle.close()
				return None
			else:
				# correctEmit is whether the particles are in the acceptance:
				self.correctEmit = np.all([(self.gammaX*(self.data[:,0]/1000)**2+2*self.alphaX*(self.data[:,0]/1000)*(self.xPrimeInit)+self.betaX*(self.xPrimeInit)**2)<=self.rmsEmitX*6,
				                           (self.gammaY*(self.data[:,1]/1000)**2+2*self.alphaY*(self.data[:,1]/1000)*(self.yPrimeInit)+self.betaY*(self.yPrimeInit)**2)<=self.rmsEmitY*6],
				                          axis=0)
				# Get the weight of the particles in the acceptance:
				self.correctParticles = self.data[self.correctEmit,:]
				self.correctWeight = np.sum(self.correctParticles[:,11])
				#
				self.logFileCont += '%s %.5e %.5e %.5e %.5e %.5e %.5e %.5e %.5e\n'\
				                         %(iteration,float(self.xMean),self.betaX,self.alphaX,float(np.sqrt(self.varMatDetX)),
				                           float(self.yMean),self.betaY,self.alphaY,float(np.sqrt(self.varMatDetY)))
				self.logFileCont += '%i'%(int(self.correctWeight))
				self.logFileHandle.write(self.logFileCont)
				self.logFileHandle.close()
				#
				return dict(mean=np.array([self.xMean, self.yMean]),
				            TWISS=np.array([[self.betaX,self.alphaX],[self.betaY,self.alphaY]]),
				            TWISSX=np.array([self.betaX,self.alphaX]),
				            TWISSY=np.array([self.betaY,self.alphaY]),
				            betaX = self.betaX,
				            betaY = self.betaY,
				            alphaX = self.alphaX,
				            alphaY = self.alphaY)

		def recoverBeam():
			del self.xStat, self.yStat, self.sumWeight, self.xPrime, self.yPrime, self.xMean, self.yMean
			del self.xPrimeMean, self.yPrimeMean, self.xSigma, self.ySigma, self.xPrimeSigma, self.yPrimeSigma
			del self.xRho, self.yRho, self.varMatX, self.varMatY, self.betaX, self.alphaX, self.gammaX
			del self.betaY, self.alphaY, self.gammaY, self.xxPrimeCount, self.yyPrimeCount, self.xCenter, self.yCenter
			del self.correctEmit, self.correctParticles, self.correctWeight
			del self.g4blBeamMat, self.xPrimeInit, self.yPrimeInit

		reduceEmitReturn = reduceEmit()
		recoverBeam()
		return reduceEmitReturn

	def generateGaussian(self,beamProperty):
		""" generateGaussian(x):
			Generate a 2-D array, of size N-by-12. For G4Beamline simulation.
			x is a List. x[0]=1-D array, of length 2. Mean of the U and U' distribution. [mu_U,mu_{U'}]
						 x[1]=2-D array, 2-by-2. Sigma matrix of U and U' distribution
						 x[2]=1-D array, of length 2. Mean of the V and V' distribution. [mu_V,mu_{V'}]
						 x[3]=2-D array, 2-by-2. Sigma matrix of V and V' distribution
						 x[4]=1-D array, of length 2. Upper and lower limits of the momentum distribution. [P_low,P_upper]
						 x[5]=1-D array, of length 2. Upper and lower limits of the time distribution. [t_low,t_upper]
						 x[6]=scalar, the number of particles in the beam.
						 x[7]=scaler or string, the PDGid of the particle, or the name of the particle
		"""
		# A dictionary, of particle PDGid
		PDGid = {'e+':-11, 'e-':11, 'anti_nu_e':-12, 'anti_nu_mu':-14, 'anti_nu_tau':-16, 'gamma':22, 'gluon':21, 'kaon+':321,'kaon-':-321,
		         'kaon0':311, 'mu+':-13, 'mu-':13, 'neutron':2112, 'nu_e':12, 'nu_mu':14, 'nu_tau':16, 'pi+':211, 'pi-':-211, 'pi0':111,
		         'proton':2212, 'rho+':213, 'rho-':-213, 'rho0':113, 'tau+':-15, 'tau-':15}
		# Nead at least one particle:
		if beamProperty[6]<=0:
			raise ValueError('Number of particles should be positive!')
		beamArray = np.zeros((beamProperty[6],12))
		# Get the PDGid of the particle from the dictionary
		try:
			particleID = int(beamProperty[7])
		except ValueError:
			try:
				particleID = PDGid[beamProperty[7]]
			except KeyError as e:
				print e
				print 'Can not find the PDGid of '+beamProperty[7]
		beamArray[:,7] = particleID
		# EventID, TrackID, ParentID, Weight
		beamArray[:,8] = np.arange(beamProperty[6])+1
		beamArray[:,9], beamArray[:,10], beamArray[:,11] = 0,0,1
		# Get the momentum of each particle:
		np.random.seed()
		beamMom = beamProperty[4][0]+(beamProperty[4][1]-beamProperty[4][0])*np.random.random(beamProperty[6])
		# Transverse phase space:
		np.random.seed()
		xPhaseSpace = np.random.multivariate_normal(beamProperty[0],beamProperty[1],beamProperty[6])
		np.random.seed()
		yPhaseSpace = np.random.multivariate_normal(beamProperty[2],beamProperty[3],beamProperty[6])
		# P_x, P_y, P_z:
		beamArray[:,5] = beamMom*np.sqrt(1/(xPhaseSpace[:,1]**2+yPhaseSpace[:,1]**2+1))
		beamArray[:,4] = beamArray[:,5]*yPhaseSpace[:,1]
		beamArray[:,3] = beamArray[:,5]*xPhaseSpace[:,1]
		# X,Y,Z
		beamArray[:,0] = xPhaseSpace[:,0]*1e3
		beamArray[:,1] = yPhaseSpace[:,0]*1e3
		# Get the time structure of the beam:
		np.random.seed()
		beamArray[:,6] = beamProperty[5][0]+(beamProperty[5][1]-beamProperty[5][0])*np.random.random(beamProperty[6])
		#
		return beamArray

	def genError(self,msg,expt=None):

		class beamError(Exception):
			def __init__(self,msg):
				Exception.__init__(self)
				self.value = msg
			def __str__(self):
				return self.value

		if expt == None:
			raise beamError(msg)
		else:
			raise expt,msg

	def getBeam(self):
		"""synonym for getData()"""
		return self.getData()

	def getData(self):
		"""Return the beam data ndarray if already set; return 0 if not."""
		if self.dataSet==1:
			return self.data
		else:
			print '---Error: The beam has not been set yet!---'
			return 0

	def getStat(self,coordinates='all'):
		"""
		Get the statistics of the current beam if the beam has been set and generated/loaded.
		coordinates = 'all','xx','yy','xy','tp'
		xx = x - x' phase space, return = [mu_array_of_x_xprime, covariance_matrix_of_x_xprime]
		yy = y - y' phase space, return = [mu_array_of_y_yprime, covariance_matrix_of_y_yprime]
		xy = x - y  real space,  return = [mu_array_of_x_y, covariance_matrix_of_x_y]
		tp = t - p  longitudinal space, return = [mu_array_of_t_p, covariance_of_t_p]
		all = list all, return = a list of the lists of all the spaces with len(return)=4
		"""
		def calCov(u,v,weight):
			'''Calculate the variance matrix of two vectors, taking the weights into account.
			'''
			uMean = np.sum(u*weight)/np.sum(weight)
			vMean = np.sum(v*weight)/np.sum(weight)
			uvCov = np.array([[np.sum((u-uMean)**2*weight)/np.sum(weight),np.sum((u-uMean)*(v-vMean)*weight)/np.sum(weight)],
			                  [np.sum((u-uMean)*(v-vMean)*weight)/np.sum(weight),np.sum((v-vMean)**2*weight)/np.sum(weight)]])
			return [np.array([uMean,vMean]),uvCov]

		if self.dataSet:
			if coordinates.lower() == 'all':
				# self.x(y)StatVec is the parameter vector for Gauss-Newton algorithm;
				# [uMean,uSigma,uPrimeMean,uPrimeSigma,uRho]
				self.xStat = calCov(self.x,self.xPrime,self.weight)
				self.xStatVec = [self.xStat[0][0],np.sqrt(self.xStat[1][0,0]),self.xStat[0][1],np.sqrt(self.xStat[1][1,1]),
				                 self.xStat[1][0,1]/(np.sqrt(self.xStat[1][0,0])*np.sqrt(self.xStat[1][1,1]))]
				self.yStat = calCov(self.y,self.yPrime,self.weight)
				self.yStatVec = [self.yStat[0][0],np.sqrt(self.yStat[1][0,0]),self.yStat[0][1],np.sqrt(self.yStat[1][1,1]),
				                 self.yStat[1][0,1]/(np.sqrt(self.yStat[1][0,0])*np.sqrt(self.yStat[1][1,1]))]
				self.xyStat = calCov(self.x,self.y,self.weight)
				self.tpStat = calCov(self.t,self.P,self.weight)
				return [self.xStat,self.yStat,self.xyStat,self.tpStat]
			elif coordinates.lower() == 'xx':
				self.xStat = calCov(self.x,self.xPrime,self.weight)
				return self.xStat
			elif coordinates.lower() == 'yy':
				self.yStat = calCov(self.y,self.yPrime,self.weight)
				return self.yStat
			elif coordinates.lower() == 'xy':
				self.xyStat = calCov(self.x,self.y,self.weight)
				return self.xyStat
			elif coordinates.lower() == 'tp':
				self.tpStat = calCov(self.t,self.P,self.weight)
				return self.tpStat
			else:
				print '---Error: Cannot recognize the argument:"',coordinates,'" getStat() failed.---'
				print '---Warning: return is None. ---'
				return None
		else:
			print '---Error: The beam is not set yet! ---'
			print '---Warning: return is None. ---'
			return None

	def loadBeam(self,beamData):
		"""Let user define the beam. This is typically done by, for example, if you have a G4Beamline beam:
		beamCase = beam()
		beamCase.loadBeam(beamCase.loadtxt_fast('YourG4BLFile',NumberOfBeginningComments))
		"""
		self.data = beamData
		self.xPrime = self.data[:,3]/self.data[:,5]
		self.yPrime = self.data[:,4]/self.data[:,5]
		self.P = np.sqrt(self.data[:,3]**2+self.data[:,4]**2+self.data[:,5]**2)
		self.sumWeight = np.sum(self.data[:,11])
		self.weight = self.data[:,11]
		self.x, self.y, self.z = self.data[:,0], self.data[:,1], self.data[:,2]
		self.t = self.data[:,6]
		self.dataSet = 1

	def loadtxt_fast(self,filename, skiprows=0, delimiter=' '):

		def iter_func():
			with open(filename,'r') as infile:
				for ii in range(skiprows):
					next(infile)
				skip = 0
				for line in infile:
					line = ' '.join(line.split()).split(delimiter)
					for item in line:
						yield float(item)
				self.rowlength = len(line)
		data = np.fromiter(iter_func(), dtype=float)
		data = data.reshape((-1,self.rowlength))
		return data

	def setX(self,x_list):
		try:
			if len(x_list[0])==2 and x_list[0].shape[0] == 2 and x_list[1].shape == (2,2):
				self.X[0] = x_list[0]
				self.X[1] = x_list[1]
				self.Xset = 1
			else:
				print '''---Error: Wrong parameters for X-X' phase space statistics. setX failed. ---'''
				self.Xset = 0
		except AttributeError as e:
			print e
			print 'Error in obtaining the shape of mu_x and sigma_x'

	def setY(self,y_list):
		try:
			if len(y_list[0])==2 and y_list[0].shape[0] == 2 and y_list[1].shape == (2,2):
				self.Y[0] = y_list[0]
				self.Y[1] = y_list[1]
				self.Yset = 1
		except AttributeError as e:
			print e
			print 'Error in obtaining the shape of mu_y and sigma_y'

	def setT(self,tlimits):
		try:
			if tlimits.shape[0] == 2:
				self.T = tlimits
				self.Tset = 1
		except AttributeError as e:
			print e
			print 'Error in obtaining the upper and lower limits of the time distribution'

	def setP(self,plimits):
		try:
			if plimits.shape[0] == 2:
				self.Plimits = plimits
				self.Pset = 1
		except AttributeError as e:
			print e
			print 'Error in obtaining the upper and lower limits of the momentum distribution'

	def setWeight(self,weightArray=1):
		"""
		Set the weights of the beam to a constant if weightArray = Constant,
		or set the weights of the beam to an array if weightArray = list or ndarray.
		"""
		if not self.dataSet:
			print '---Warning: The beam array has not been set yet. Cannot set weights.---'
			return
		try:
			if len(weightArray) >1:
				try:
					weightArrayCopy = np.array(weightArray)+0
				except TypeError:
					print '---Error: Cannot convert the weights to a ndarray ---\n---Warning: Weights NOT set.---'
					return
				if len(weightArrayCopy) != self.data.shape[0]:
					print '---Error: The length of the weight array is not the same with the number of particles ---'
					print '---Warning: The weights of the beam cannot be set.---'
					return
				else:
					self.data[:,11] = weightArrayCopy
					self.Weightset = 1
		except TypeError:
			try:
				weightArrayCopy = float(weightArray)
				self.data[:,11] = weightArrayCopy
				self.Weightset = 1
			except TypeError:
				print '---Error: Cannot convert the weight to a constant ---\n---Warning: Weights NOT set.---'
				return

	def setnumber(self,number):
		try:
			self.numberOfparticles = int(number)
			self.numberSet = 1
		except Exception as e:
			print e
			print 'Unable to set the number of particles.'

	def setPDGid(self,PDGid):
		self.PDGid = PDGid
		self.PDGset = 1

	def genBeam(self):
		self.whichSet = [self.Xset,self.Yset,self.Tset,self.Pset,self.numberSet,self.PDGset]
		self.whatSet = ['Horizontal Phase Space','Vertical Phase Space','Time distribution',
		                'Momentum Distribution','Number of particles','The PDGid of the particle']
		self.whatSet = [self.whatSet[ii] for ii in range(6) if self.whichSet[ii]==0]
		if len(self.whatSet)==0:
			self.data = self.generateGaussian(self.X+self.Y+[self.Plimits,self.T,self.numberOfparticles,self.PDGid])
		else:
			WarningString = 'Warning! '+reduce(lambda x,y: x+' and '+y, self.whatSet)+' are not set! The default values are used (which I think you won\'t like)'
			print WarningString
			self.data = self.generateGaussian(self.X+self.Y+[self.Plimits,self.T,self.numberOfparticles,self.PDGid])
		self.xPrime = self.data[:,3]/self.data[:,5]
		self.yPrime = self.data[:,4]/self.data[:,5]
		self.P = np.sqrt(self.data[:,3]**2+self.data[:,4]**2+self.data[:,5]**2)
		self.dataSet = 1

	def writeBeam(self,beamFileName):
		# Write the beam data to the file:
		np.savetxt(beamFileName,self.data,fmt=['%.4f',]*7+['%i',]*4+['%.2f'],header=self.beamFileHeader)

	def splitBeam(self):
		"""
		If the beam data has weight info for each of the particles, this function will split the particles up based on
		their weights.
		Notice:
			1. Each of the particles will be split by floor(N)+1 or floor(N) times, where N is the particle's weight, based on the
			fractional part of N.
			Thus, a particle with weight=11.2 will have a 20% probability to be split by 12 times, and 80% prob. by 11 times.
			on the other hand, a particle with weight=0.002 will only have 0.2% probability to be split by 1 time, or nothing.
			2. Since it's better not to have the same eventID for two particles, the eventID will be changed to 1-NumberOfParticles
			serially with increment=1.
		Usage:
			beamCase.splitBeam()
		Return:
			None.
		Effect:
			beamCase.data will be changed.
		"""
		if self.dataSet:
			oldNumParticles = self.data.shape[0]
			# A random array for the comparison with the fractional weight of each particle:
			particleWeightRand = np.random.random(oldNumParticles)
			# Use true-or-false array instead of 1 and 0 array:
			weightAddorNot = ((self.data[:,11]-np.floor(self.data[:,11])) >= particleWeightRand)
			newWeight = weightAddorNot + self.data[:,11]
			newWeight = newWeight.astype(int)
			# Method 1
			dataNew = np.repeat(self.data,newWeight,axis=0)
			dataNew[:,8] = np.arange(dataNew.shape[0])+1
			dataNew[:,11] = 1

			# Update the data and then get the new x', y' and P.
			self.loadBeam(dataNew)
		else:
			# If the beam data is not set yet, raise an exception
			self.genError('---Error: Trying to split the beam without setting it ---')