Python teapot.DriftTEAPOT類代碼示例

    def __init__(self, phaseLength, nMacrosMin, nBins,\
	    bunch, impeDict,\
	    name = "freq. dep. LImpedance node"):
        """
            Constructor. Creates the FreqDep_LImpedance-teapot element.
        """
        DriftTEAPOT.__init__(self, name)
        self.limpedance = LImpedance(phaseLength, nMacrosMin, nBins)
        self.setType("freq. dep. limpedance node")
        self.setLength(0.0)
        self.phaseLength = phaseLength
        self.nBins = nBins
        self.localDict = impeDict
        self.freq_tuple = self.localDict["freqs"]
        self.freq_range = (len(self.freq_tuple) - 1)
        self.z_tuple = self.localDict["z_imp"]
        self.c = consts.speed_of_light
        BetaRel = bunch.getSyncParticle().beta()
        Freq0 = (BetaRel * self.c) / self.phaseLength
        Z = []
        for n in range(self.nBins / 2 - 1):
            freq_mode = Freq0 * (n + 1)
            z_mode = interp(freq_mode, self.freq_range,\
                            self.freq_tuple, self.z_tuple)
            Z.append(z_mode)
        self.limpedance.assignImpedance(Z)

	def __init__(self, length, name = "base_rfnode"):
		"""
			Constructor. Creates Base RF Cavity TEAPOT element.
			It will never be called.
		"""
		DriftTEAPOT.__init__(self, name)
		self.setType("base rf node")
		self.setLength(0.0)

	def __init__(self, length, ma, density_fac, shape, a, b, c, d, angle, pos = 0, name = "collimator no name"):
		"""
		Constructor. Creates the Collimator TEAPOT element.
		"""
		DriftTEAPOT.__init__(self,name)
		self.collimator = Collimator(length,ma,density_fac,shape,a,b,c,d,angle,pos)
		self.setType("collimator teapot")
		self.setLength(length)

	def __init__(self, bunch, xbump, xpbump, ybump, ypbump, name = "bump"):
		"""
		Constructor. Creates the Bumpe TEAPOT element.
		"""
		DriftTEAPOT.__init__(self,name)
		self.simplebump = simpleBump(bunch, xbump, xpbump, ybump, ypbump);
		self.setType("Bump")
		self.setLength(0.0)

    def __init__(self, disp,\
	    name = "BendYF Displacement"):
        """
            Constructor. Creates BendDisplacementYF element.
        """
        DriftTEAPOT.__init__(self, name)
        self.setType("yf bend displacement node")
        self.setLength(0.0)
        self.disp = disp

	def __init__(self, a, b, pos = 0, c = 0, d = 0, name = "aperture"):
		DriftTEAPOT.__init__(self,name)
		self.shape = 3
		self.a = a
		self.b = b
		self.c = c
		self.d = d
		self.pos = pos
		self.Aperture = Aperture(self.shape, self.a, self.b, self.c, self.d, self.pos)

	def __init__(self, xmin, xmax, ymin, ymax, thick, name = "foil no name"):
		"""
		Constructor. Creates the Foil TEAPOT element.
		"""
		DriftTEAPOT.__init__(self,name)
		self.foil = Foil(xmin, xmax, ymin, ymax, thick)
		self.setType("foil teapot")
		self.setLength(0.0)
		# The user choice of scattering routine. Defualt (0) is full scatter
		self.scatterChoice = 0

	def __init__(self, name = "BPMSignal no name"):
		"""
		Constructor. Creates the StatLats TEAPOT element.
		"""
		DriftTEAPOT.__init__(self,name)
		self.bpm = BPMSignal()
		self.setType("BPMSignal")
		self.lattlength = 0.0
		self.setLength(0.0)
		self.position = 0.0

 def __init__(self, bx, by, ax, ay, ex, epx, l, zi, zf, ds, niters,
               resid, xrefi, yrefi, eulerai, eulerbi, eulergi, b, filename):
     """
     Constructor. Creates the FieldTracker TEAPOT element.
     """
     DriftTEAPOT.__init__(self,name)
     self.fieldtracker = FieldTracker(order, bx, by, ax, ay, ex, epx, l, zi, zf, ds, niters,
               resid, xrefi, yrefi, eulerai, eulerbi, eulergi, apflag, b)
     self.setType("fieldtracker teapot")
     self.setLength(l)

def addTeapotCollimatorNode(lattice, position, collimator_node):
	"""
	It will put one Teapot collimation node in the lattice 
	"""
	length_tolerance = 0.0001
	lattice.initialize()
	if(position > lattice.getLength() ):
		position = lattice.getLength();
		print "User-specified aperture position is larger than lattice length.  Resetting it to lattice length."
	collimator_node.setPosition(position);
	position_start = position
	position_stop = position + collimator_node.getLength()
	(node_start_ind,node_stop_ind,z,ind) = (-1,-1, 0., 0)
	for node in lattice.getNodes():
		if(position_start >= z and position_start <= z + node.getLength()):
			node_start_ind = ind
		if(position_stop >= z and position_stop <= z + node.getLength()):
			node_stop_ind = ind
		ind += 1
		z += node.getLength()
	#-------now we check that between start and end we have only non-modified drift elements
	#-------if the space charge was added first - that is a problem. The collimation should be added first.
	for node in lattice.getNodes()[node_start_ind:node_stop_ind+1]:
		#print "debug node=",node.getName()," type=",node.getType()," L=",node.getLength()
		if(not isinstance(node,DriftTEAPOT)):
			print "Non-drift node=",node.getName()," type=",node.getType()," L=",node.getLength()
			orbitFinalize("We have non-drift element at the place of the collimator! Stop!")
			#if(node.getNumberOfChildren() != 4):
			#print "Node=",node.getName()," type=",node.getType()," L=",node.getLength()," N children nodes=",node.getNumberOfChildren()
			#orbitFinalize("Drift element was modified with additional functionality (SC or something else)! Add collimation first! Stop!")
	# make array of nodes from collimator in the center and possible two drifts if their length is more than length_tolerance [m]
	nodes_new_arr = [collimator_node,]
	drift_node_start = lattice.getNodes()[node_start_ind]
	drift_node_stop = lattice.getNodes()[node_stop_ind]	
	#------now we will create two drift nodes: before the collimator and after
	#------if the length of one of these additional drifts less than length_tollerance [m] we skip this drift 
	if(position_start > lattice.getNodePositionsDict()[drift_node_start][0] +  length_tolerance):
		drift_node_start_new = DriftTEAPOT(drift_node_start.getName())
		drift_node_start_new.setLength(position_start - lattice.getNodePositionsDict()[drift_node_start][0])
		nodes_new_arr.insert(0,drift_node_start_new)
	if(position_stop < lattice.getNodePositionsDict()[drift_node_stop][1] - length_tolerance):
		drift_node_stop_new = DriftTEAPOT(drift_node_stop.getName())
		drift_node_stop_new.setLength(lattice.getNodePositionsDict()[drift_node_stop][1] - position_stop)
		nodes_new_arr.append(drift_node_stop_new)
	#------ now we will modify the lattice by replacing the found part with the new nodes
	lattice.getNodes()[node_start_ind:node_stop_ind+1] = nodes_new_arr
	# initialize the lattice
	lattice.initialize()

def addTeapotDiagnosticsNode(lattice, position, diagnostics_node):
	"""
	It will put one Teapot diagnostics node in the lattice 
	"""
	length_tollerance = 0.0001
	lattice.initialize()
	position_start = position
	position_stop = position + diagnostics_node.getLength()
	diagnostics_node.setPosition(position)
	diagnostics_node.setLatticeLength(lattice.getLength())
	(node_start_ind,node_stop_ind,z,ind) = (-1,-1, 0., 0)
	for node in lattice.getNodes():
		if(position_start >= z and position_start <= z + node.getLength()):
			node_start_ind = ind
		if(position_stop >= z and position_stop <= z + node.getLength()):
			node_stop_ind = ind
		ind += 1
		z += node.getLength()
	#-------now we check that between start and end we have only non-modified drift elements
	#-------if the space charge was added first - that is a problem. The collimation should be added first.
	for node in lattice.getNodes()[node_start_ind:node_stop_ind+1]:
		#print "debug node=",node.getName()," type=",node.getType()," L=",node.getLength()
		if(not isinstance(node,DriftTEAPOT)):
			print "Non-drift node=",node.getName()," type=",node.getType()," L=",node.getLength()
			orbitFinalize("We have non-drift element at the place of the diagnostics! Stop!")
			
	# make array of nodes from diagnostics in the center and possible two drifts if their length is more than length_tollerance [m]
	nodes_new_arr = [diagnostics_node,]
	drift_node_start = lattice.getNodes()[node_start_ind]
	drift_node_stop = lattice.getNodes()[node_stop_ind]	
	#------now we will create two drift nodes: before the diagnostics and after
	#------if the length of one of these additional drifts less than length_tollerance [m] we skip this drift 
	if(position_start > lattice.getNodePositionsDict()[drift_node_start][0] +  length_tollerance):
		drift_node_start_new = DriftTEAPOT(drift_node_start.getName())
		drift_node_start_new.setLength(position_start - lattice.getNodePositionsDict()[drift_node_start][0])
		nodes_new_arr.insert(0,drift_node_start_new)
	if(position_stop < lattice.getNodePositionsDict()[drift_node_stop][1] - length_tollerance):
		drift_node_stop_new = DriftTEAPOT(drift_node_stop.getName())
		drift_node_stop_new.setLength(lattice.getNodePositionsDict()[drift_node_stop][1] - position_stop)
		nodes_new_arr.append(drift_node_stop_new)
	#------ now we will modify the lattice by replacing the found part with the new nodes
	lattice.getNodes()[node_start_ind:node_stop_ind+1] = nodes_new_arr
	# initialize the lattice
	lattice.initialize()

#---PRINT Function for Matrix
def printM(m):
	print "----matrix--- size=",m.size()
	for i in xrange(m.size()[0]):
		for j in xrange(m.size()[1]):
			print ("(%1d,%1d)=% 6.5e "%(i,j,m.get(i,j))),
		print ""	
		

b = Bunch()
Ekin = 1.0 # in GeV
b.getSyncParticle().kinEnergy(Ekin)

#define TEAPOT drift
node0 = DriftTEAPOT("drift")
node0.setLength(1.0)

#define TEAPOT quad
node1 = QuadTEAPOT("quad")
node1.setLength(1.0)
node1.addParam("kq",0.5)
	
matrixGenerator = MatrixGenerator()

#========matrix for drift =====
m = Matrix(6,6)
matrixGenerator.initBunch(b)
node0.trackBunch(b)
matrixGenerator.calculateMatrix(b,m)	
print "drift matrix L=",node0.getLength()