def trackDesignBunch(self,
bunch_in,
paramsDict=None,
actionContainer=None,
index_start=-1,
index_stop=-1):
"""
This will track the design bunch through the linac and set up RF Cavities times of
arrivals.
"""
if (actionContainer == None):
actionContainer = AccActionsContainer("Design Bunch Tracking")
if (paramsDict == None): paramsDict = {}
bunch = Bunch()
bunch_in.copyEmptyBunchTo(bunch)
bunch.getSyncParticle().time(0.)
paramsDict["bunch"] = bunch
def trackDesign(localParamsDict):
node = localParamsDict["node"]
node.trackDesign(localParamsDict)
actionContainer.addAction(trackDesign, AccActionsContainer.BODY)
self.trackActions(actionContainer, paramsDict, index_start, index_stop)
actionContainer.removeAction(trackDesign, AccActionsContainer.BODY)
return bunch
def addTrMatrxGenNodes(self,
accLattice,
node_or_nodes,
place=MarkerLinacNode.ENTRANCE):
"""
Adds the LinacTrMatrixGenNode to the nodes as child nodes.
"""
nodes = []
if (type(node_or_nodes) in [tuple, list]):
for node in node_or_nodes:
nodes.append(node)
else:
nodes.append(node_or_nodes)
#-----------------------------
for node in nodes:
trMatrxGenNode = LinacTrMatrixGenNode(self,
node.getName() + ":trMatrx")
node.addChildNode(trMatrxGenNode, place)
#----- set up the position of the TrMatrix nodes
actions = AccActionsContainer()
def accNodeExitAction(paramsDict):
"""
Nonbound function. Sets the position of the TrMatrix nodes.
"""
node = paramsDict["node"]
if (isinstance(node, LinacTrMatrixGenNode)):
pos = paramsDict["path_length"]
node.setPosition(pos)
actions.addAction(accNodeExitAction, AccNode.EXIT)
accLattice.trackActions(actions)
self.init()
return self.trMatrxNodes
def getNodeForNameFromWholeLattice(accLattice, name):
"""
Returns the accelerator node or an array of nodes with the same name.
This function could be replaced later by the method in the AccLattice class.
"""
paramsDict = {}
actions = AccActionsContainer()
nodes = []
def accNodeExitAction(paramsDict):
"""
Non-bound function. Finds the node in the lattice
with the specified name.
positions. This is a closure (well, maybe not
exactly). It uses external objects.
"""
node = paramsDict["node"]
if (node.getName() == name):
nodes.append(node)
actions.addAction(accNodeExitAction, AccNode.EXIT)
accLattice.trackActions(actions, paramsDict)
if (len(nodes) == 1):
return nodes[0]
elif (len(nodes) == 0):
return None
else:
return nodes
def __init__(self, teapot_lattice, bunch, name=None):
MATRIX_Lattice.__init__(self, name)
if (isinstance(teapot_lattice, TEAPOT_Lattice) != True):
orbitFinalize(
"Constructor orbit.teapot.TEAPOT_MATRIX_Lattice needs the TEAPOT_Lattice instance."
)
#memorize the TEAPOT LATTICE
if (name == None):
name = teapot_lattice.getName()
self.setName(name)
self.teapot_lattice = teapot_lattice
self.bunch = Bunch()
self.lost_bunch = Bunch()
bunch.copyEmptyBunchTo(self.bunch)
bunch.copyEmptyBunchTo(self.lost_bunch)
self.matrixGenerator = MatrixGenerator()
#----------make MATRIX lattice from TEAPOT
def twissAction(paramsDict):
node = paramsDict["node"]
bunch = paramsDict["bunch"]
active_index = node.getActivePartIndex()
n_parts = node.getnParts()
length = node.getLength(active_index)
if (isinstance(node, BaseTEAPOT) == True
and isinstance(node, RingRFTEAPOT) == False):
self.matrixGenerator.initBunch(bunch)
node.track(paramsDict)
#bunch is ready
matrixNode = BaseMATRIX(node.getName() + "_" +
str(active_index))
matrixNode.addParam("matrix_parent_node", node)
matrixNode.addParam("matrix_parent_node_type", node.getType())
matrixNode.addParam("matrix_parent_node_n_nodes", n_parts)
matrixNode.addParam("matrix_parent_node_active_index",
active_index)
matrixNode.setLength(length)
self.matrixGenerator.calculateMatrix(bunch,
matrixNode.getMatrix())
#print "============= name=",matrixNode.getName(),
#print " type=",matrixNode.getParam("matrix_parent_node_type"),
#print " L=",matrixNode.getLength()
self.addNode(matrixNode)
if (isinstance(node, RingRFTEAPOT) == True):
rf_node = RingRFTEAPOT(node.getName())
rf_node.setParamsDict(node.getParamsDict().copy())
self.addNode(rf_node)
accContainer = AccActionsContainer()
accContainer.addAction(twissAction, AccActionsContainer.BODY)
paramsDict = {}
paramsDict["bunch"] = self.bunch
paramsDict["lostbunch"] = self.lost_bunch
paramsDict["position"] = 0.
paramsDict["useCharge"] = self.teapot_lattice.getUseRealCharge()
self.teapot_lattice.trackActions(accContainer, paramsDict)
self.makeOneTurnMatrix()
self.initialize()
def getTransportMatrixArray(self, teapot_lattice):
accContainer = AccActionsContainer()
accContainer.addAction(self._twissAction, AccActionsContainer.BODY)
paramsDict = {}
paramsDict["bunch"] = self.b
paramsDict["position"] = 0.
self.matrixArr = []
self.index = 0
teapot_lattice.trackActions(accContainer, paramsDict)
return self.matrixArr
def trackDesignBunch(self, bunch, paramsDict = {}, actionContainer = None):
"""
It tracks the bunch through the AccNodeBunchTracker instance.
"""
if(actionContainer == None): actionContainer = AccActionsContainer("Design Bunch Tracking")
paramsDict["bunch"] = bunch
def trackDesign(paramsDict):
node = paramsDict["node"]
node.trackDesign(paramsDict)
actionContainer.addAction(trackDesign, AccActionsContainer.BODY)
self.trackActions(actionContainer,paramsDict)
actionContainer.removeAction(trackDesign, AccActionsContainer.BODY)
def trackBunch(self, bunch, paramsDict={}, actionContainer=None):
"""
It tracks the bunch through the lattice.
"""
if (actionContainer == None):
actionContainer = AccActionsContainer("Bunch Tracking")
paramsDict["bunch"] = bunch
def track(paramsDict):
node = paramsDict["node"]
node.track(paramsDict)
actionContainer.addAction(track, AccActionsContainer.BODY)
self.trackActions(actionContainer, paramsDict)
actionContainer.removeAction(track, AccActionsContainer.BODY)
def initialize(self):
"""
Method. Initializes the lattice and child node structures.
"""
res_dict = {}
for node in self.__children:
if(res_dict.has_key(node)):
msg = "The AccLattice class instance should not have duplicate nodes!"
msg = msg + os.linesep
msg = msg + "Method initialize():"
msg = msg + os.linesep
msg = msg + "Name of node=" + node.getName()
msg = msg + os.linesep
msg = msg + "Type of node=" + node.getType()
msg = msg + os.linesep
orbitFinalize(msg)
else:
res_dict[node] = None
node.initialize()
del res_dict
paramsDict = {}
actions = AccActionsContainer()
d = [0.]
posn = {}
def accNodeExitAction(paramsDict):
"""
Nonbound function. Sets lattice length and node
positions. This is a closure (well, maybe not
exactly). It uses external objects.
"""
node = paramsDict["node"]
parentNode = paramsDict["parentNode"]
if(isinstance(parentNode, AccLattice)):
posBefore = d[0]
d[0] += node.getLength()
posAfter = d[0]
posn[node]=(posBefore, posAfter)
actions.addAction(accNodeExitAction, AccNode.EXIT)
self.trackActions(actions, paramsDict)
self.__length = d[0]
self.__childPositions = posn
self.__isInitialized = True
def trackBunch(self,
bunch,
paramsDict={},
actionContainer=None,
index_start=-1,
index_stop=-1):
"""
It tracks the bunch through the lattice. Indexes index_start and index_stop are inclusive.
"""
if (actionContainer == None):
actionContainer = AccActionsContainer("Bunch Tracking")
paramsDict["bunch"] = bunch
paramsDict["useCharge"] = self.useCharge
def track(paramsDict):
node = paramsDict["node"]
node.track(paramsDict)
actionContainer.addAction(track, AccActionsContainer.BODY)
self.trackActions(actionContainer, paramsDict, index_start, index_stop)
actionContainer.removeAction(track, AccActionsContainer.BODY)
def trackBunch(self,
bunch,
paramsDict=None,
actionContainer=None,
index_start=-1,
index_stop=-1):
"""
It tracks the bunch through the lattice.
"""
if (actionContainer == None):
actionContainer = AccActionsContainer("Bunch Tracking")
if (paramsDict == None): paramsDict = {}
paramsDict["bunch"] = bunch
bunch.getSyncParticle().time(0.)
def track(paramsDict):
node = paramsDict["node"]
node.track(paramsDict)
actionContainer.addAction(track, AccActionsContainer.BODY)
self.trackActions(actionContainer, paramsDict, index_start, index_stop)
actionContainer.removeAction(track, AccActionsContainer.BODY)
def getAllNodesInLattice(accLattice):
"""
Returns the array with all nodes on all levels (even sub-child).
"""
nodes = []
paramsDict = {}
actions = AccActionsContainer()
def accNodeEntranceAction(paramsDict):
"""
Non-bound function. Add the node to the nodes array
at the entrance inside the node.
This is a closure (well, maybe not
exactly). It uses external objects.
"""
node = paramsDict["node"]
nodes.append(node)
actions.addAction(accNodeEntranceAction, AccNode.ENTRANCE)
accLattice.trackActions(actions, paramsDict)
return nodes
def getNodePosDictForWholeLattice(accLattice):
"""
Returns the dict[node] = (posStart,posEnd) for all nodes (not only for the firts level).
This function could be replaced later by the method in the AccLattice class.
"""
paramsDict = {}
actions = AccActionsContainer()
posStartDict = {}
posStopDict = {}
def accNodeEntranceAction(paramsDict):
"""
Non-bound function. Sets node's end positions.
This is a closure (well, maybe not exactly). .
"""
node = paramsDict["node"]
pos = paramsDict["path_length"]
posStartDict[node] = pos
def accNodeExitAction(paramsDict):
"""
Non-bound function. Sets node's end positions.
This is a closure (well, maybe not exactly). .
"""
node = paramsDict["node"]
pos = paramsDict["path_length"]
posStopDict[node] = pos
actions.addAction(accNodeEntranceAction, AccNode.ENTRANCE)
actions.addAction(accNodeExitAction, AccNode.EXIT)
accLattice.trackActions(actions, paramsDict)
posStartStopDict = {}
for key in posStartDict:
pos_start = posStartDict[key]
pos_end = posStopDict[key]
posStartStopDict[key] = (pos_start, pos_end)
return posStartStopDict
def getAllMagnetsInLattice(accLattice):
"""
Returns the array with all magnets including magnets on all levels (e.g correctors)
"""
magnets = []
paramsDict = {}
actions = AccActionsContainer()
def accNodeExitAction(paramsDict):
"""
Non-bound function. Finds the node in the lattice
with the specified name.
positions. This is a closure (well, maybe not
exactly). It uses external objects.
"""
node = paramsDict["node"]
if (isinstance(node, LinacMagnetNode)):
magnets.append(node)
actions.addAction(accNodeExitAction, AccNode.EXIT)
accLattice.trackActions(actions, paramsDict)
return magnets
def getChromaticitiesXY(self):
"""
Calculates chromaticities for X,Y planes for the whole ring
"""
(tuneX, tmp0, tmp1) = self.getRingTwissDataX()
(tuneY, tmp0, tmp1) = self.getRingTwissDataY()
tuneX = tuneX[1][-1]
tuneY = tuneY[1][-1]
self.matrixGenerator.initBunchChromCoeff(self.bunch)
#track bunch through the TEAPOT nodes
def twissAction(paramsDict):
node = paramsDict["node"]
if (isinstance(node, BaseTEAPOT) == True
and isinstance(node, RingRFTEAPOT) == False):
node.track(paramsDict)
accContainer = AccActionsContainer()
accContainer.addAction(twissAction, AccActionsContainer.BODY)
paramsDict = {}
paramsDict["bunch"] = self.bunch
self.teapot_lattice.trackActions(accContainer, paramsDict)
res_coeff = self.matrixGenerator.calcChromCoeff(self.bunch)
(coeff_x_dE, coeff_xp_dE, coeff_y_dE, coeff_yp_dE) = res_coeff
momentum = self.bunch.getSyncParticle().momentum()
mass = self.bunch.getSyncParticle().mass()
Ekin = self.bunch.getSyncParticle().kinEnergy()
chromX = -(momentum /
(2 * math.sin(2 * math.pi * tuneX))) * (coeff_x_dE +
coeff_xp_dE)
chromX = (momentum / (mass + Ekin)) * chromX
chromY = -(momentum /
(2 * math.sin(2 * math.pi * tuneY))) * (coeff_y_dE +
coeff_yp_dE)
chromY = (momentum / (mass + Ekin)) * chromY
return (chromX / (2 * math.pi), chromY / (2 * math.pi))
def initialize(self):
AccLattice.initialize(self)
#set up ring length for RF nodes
ringRF_Node = RingRFTEAPOT()
bunchwrap_Node = BunchWrapTEAPOT()
for node in self.getNodes():
if (node.getType() == ringRF_Node.getType()):
node.getParamsDict()["ring_length"] = self.getLength()
paramsDict = {}
actions = AccActionsContainer()
def accSetWrapLengthAction(paramsDict):
"""
Nonbound function. Sets lattice length for wrapper nodes
"""
node = paramsDict["node"]
bunchwrap_node = BunchWrapTEAPOT()
if (node.getType() == bunchwrap_Node.getType()):
node.getParamsDict()["ring_length"] = self.getLength()
actions.addAction(accSetWrapLengthAction, AccNode.EXIT)
self.trackActions(actions, paramsDict)
actions.removeAction(accSetWrapLengthAction, AccNode.EXIT)
def getResultsDict():
bunch = Bunch()
bunch_init.copyEmptyBunchTo(bunch)
#set up design
accLattice.trackDesignBunch(bunch)
#track through the lattice START SCL with 95.610
rf_gaps_eKin_phases_dict = {}
paramsDict = {
"test_pos": 95.605984,
"count": 0,
"rf_gap_dict": rf_gaps_eKin_phases_dict
}
actionContainer = AccActionsContainer("Bunch Tracking")
def action_exit(paramsDict):
node = paramsDict["node"]
length = node.getLength()
pos = paramsDict["test_pos"] + length
paramsDict["test_pos"] = pos
if (isinstance(paramsDict["parentNode"], AccLattice)):
if (isinstance(node, BaseRF_Gap)):
bunch_inner = paramsDict["bunch"]
eKin_out = bunch_inner.getSyncParticle().kinEnergy() * 1.0e+3
phase = makePhaseNear(
node.getGapPhase() * 180. / math.pi - 180., 0.)
rf_gaps_eKin_phases_dict = paramsDict["rf_gap_dict"]
rf_gaps_eKin_phases_dict[node] = [eKin_out, phase]
#print "debug eKin out=",eKin_out
actionContainer.addAction(action_exit, AccActionsContainer.EXIT)
accLattice.trackBunch(bunch,
paramsDict=paramsDict,
actionContainer=actionContainer)
return rf_gaps_eKin_phases_dict
elem1_1.setnParts(2)
elem1_1_1 = AccNode("el-1-1-1")
elem1_1_2 = AccNode("el-1-1-2")
elem1_1_3 = AccNode("el-1-1-3")
elem1_1_4 = AccNode("el-1-1-4")
elem1.addChildNode(elem1_1, AccNode.ENTRANCE)
elem1_1.addChildNode(elem1_1_1, AccNode.ENTRANCE)
elem1_1.addChildNode(elem1_1_2, AccNode.BODY, 0)
elem1_1.addChildNode(elem1_1_3, AccNode.BODY, 1)
elem1_1.addChildNode(elem1_1_4, AccNode.EXIT)
elem1_2 = AccNode("el-1-2")
elem2.addChildNode(elem1_2, AccNode.EXIT)
acts = AccActionsContainer()
def Blanks(n):
s = ""
for i in xrange(n):
s += " "
return s
nLevel = [0]
nElems = [0]
def funcEntrance(paramsDict):
nLevel[0] += 1
z0 = 0.
dE = 0.
bunch.addParticle(x0, xp0, y0, yp0, z0, dE)
bunch_init = Bunch()
bunch.copyBunchTo(bunch_init)
#set up design
accLattice.trackDesignBunch(bunch_init)
print "Design tracking completed."
#track through the lattice
paramsDict = {"old_pos": -1., "count": 0, "pos_step": 0.1}
actionContainer = AccActionsContainer("Test Design Bunch Tracking")
pos_start = 0.
file_out = open("pyorbit_hebt_distorted_trajectory.dat", "w")
st = " Node position x[mm] xp[mrad] y[mm] yp[mrad] "
file_out.write(st + "\n")
print st
def action_entrance(paramsDict):
node = paramsDict["node"]
bunch = paramsDict["bunch"]
pos = paramsDict["path_length"]
if (paramsDict["old_pos"] == pos): return
bunch_in = bunch_gen.getBunch(nParticles=100000,
distributorClass=WaterBagDist3D)
#bunch_in = bunch_gen.getBunch(nParticles = 100000, distributorClass = GaussDist3D)
#bunch_in = bunch_gen.getBunch(nParticles = 10000, distributorClass = KVDist3D)
print "Bunch Generation completed."
#set up design
accLattice.trackDesignBunch(bunch_in)
print "Design tracking completed."
#track through the lattice
paramsDict = {"old_pos": -1., "count": 0, "pos_step": 0.01}
actionContainer = AccActionsContainer("Bunch Tracking")
pos_start = 0.
twiss_analysis = BunchTwissAnalysis()
file_out = open("pyorbit_twiss_sizes_ekin_new_lattice.dat", "w")
s = " Node position "
s += " alphaX betaX emittX normEmittX"
s += " alphaY betaY emittY normEmittY"
s += " alphaZ betaZ emittZ emittZphiMeV"
s += " sizeX sizeY sizeZ_deg"
s += " eKin Nparts "
file_out.write(s + "\n")
print " N node position sizeX sizeY sizeZdeg eKin Nparts "
if (ind % size == rank):
bunch.addParticle(x, xp, y, yp, z, dE)
nParticlesGlobal = bunch.getSizeGlobal()
if (rank == 0):
print "total number of particles =", nParticlesGlobal
bunch.macroSize(macrosize)
#set up design
accLattice.trackDesignBunch(bunch)
paramsDict = {}
lost_parts_bunch = Bunch()
paramsDict["lostbunch"] = lost_parts_bunch
actionContainer = AccActionsContainer("Test Apertures")
twiss_analysis = BunchTwissAnalysis()
def action_entrance(paramsDict):
if (isinstance(paramsDict["parentNode"], AccLattice)):
node = paramsDict["node"]
pos = paramsDict["path_length"]
bunch = paramsDict["bunch"]
twiss_analysis.analyzeBunch(bunch)
x_rms = math.sqrt(
twiss_analysis.getTwiss(0)[1] *
twiss_analysis.getTwiss(0)[3]) * 1000.
y_rms = math.sqrt(
twiss_analysis.getTwiss(1)[1] *
start_ind = accLattice.getNodeIndex(quads[0])
stop_ind = accLattice.getNodeIndex(quads[len(quads) - 1])
accLattice.trackDesignBunch(bunch_tmp,
None,
None,
index_start=start_ind,
index_stop=stop_ind)
traj_x_function = Function()
traj_xp_function = Function()
#----- if you want more beam size points along the trajectory - change pos_step here
paramsDict = {"old_pos": -1., "count": 0, "pos_step": 0.1, "path_length": 0.}
actionContainer = AccActionsContainer("TEAPOT Bunch Tracking")
#----- start of the 1st quad
pos_start = -quads[0].getLength() / 2
def action_account(paramsDict):
node = paramsDict["node"]
name = node.getName()
bunchI = paramsDict["bunch"]
pos = paramsDict["path_length"]
if (paramsDict["old_pos"] == pos): return
if (paramsDict["old_pos"] + paramsDict["pos_step"] > pos): return
paramsDict["old_pos"] = pos
paramsDict["count"] += 1
traj_x_function.add(pos + pos_start, bunchI.x(0) * 1000.)
def TrackingBunch(accLattice, bunch, print_info=False):
#set up design
accLattice.trackDesignBunch(bunch)
#track through the lattice
paramsDict = {"old_pos": -1., "count": 0, "pos_step": 0.01}
actionContainer = AccActionsContainer("Bunch Tracking")
pos_start = 0.
twiss_analysis = BunchTwissAnalysis()
results_arr = []
def action_entrance(paramsDict):
node = paramsDict["node"]
bunch = paramsDict["bunch"]
pos = paramsDict["path_length"]
if (paramsDict["old_pos"] == pos): return
if (paramsDict["old_pos"] + paramsDict["pos_step"] > pos): return
paramsDict["old_pos"] = pos
paramsDict["count"] += 1
gamma = bunch.getSyncParticle().gamma()
beta = bunch.getSyncParticle().beta()
twiss_analysis.analyzeBunch(bunch)
x_rms = math.sqrt(
twiss_analysis.getTwiss(0)[1] *
twiss_analysis.getTwiss(0)[3]) * 1000.
y_rms = math.sqrt(
twiss_analysis.getTwiss(1)[1] *
twiss_analysis.getTwiss(1)[3]) * 1000.
z_rms = math.sqrt(
twiss_analysis.getTwiss(2)[1] *
twiss_analysis.getTwiss(2)[3]) * 1000.
z_to_phase_coeff = bunch_gen.getZtoPhaseCoeff(bunch)
z_rms_deg = z_to_phase_coeff * z_rms / 1000.0
nParts = bunch.getSizeGlobal()
(alphaX, betaX, emittX) = (twiss_analysis.getTwiss(0)[0],
twiss_analysis.getTwiss(0)[1],
twiss_analysis.getTwiss(0)[3] * 1.0e+6)
(alphaY, betaY, emittY) = (twiss_analysis.getTwiss(1)[0],
twiss_analysis.getTwiss(1)[1],
twiss_analysis.getTwiss(1)[3] * 1.0e+6)
(alphaZ, betaZ, emittZ) = (twiss_analysis.getTwiss(2)[0],
twiss_analysis.getTwiss(2)[1],
twiss_analysis.getTwiss(2)[3] * 1.0e+6)
norm_emittX = emittX * gamma * beta
norm_emittY = emittY * gamma * beta
#---- phi_de_emittZ will be in [pi*deg*MeV]
phi_de_emittZ = z_to_phase_coeff * emittZ
eKin = bunch.getSyncParticle().kinEnergy() * 1.0e+3
s_prt = " %5d %35s %4.5f " % (paramsDict["count"], node.getName(),
pos + pos_start)
s_prt += " %5.3f %5.3f %5.3f " % (x_rms, y_rms, z_rms_deg)
s_prt += " %10.6f %8d " % (eKin, nParts)
if (print_info): print s_prt
twiss_arr = [(alphaX, betaX, emittX, norm_emittX),
(alphaY, betaY, emittY, norm_emittY),
(alphaZ, betaZ, emittZ, phi_de_emittZ)]
rms_arr = [x_rms, y_rms, z_rms_deg]
results_arr.append(
[node.getName(), pos, rms_arr, twiss_arr, eKin, nParts])
def action_exit(paramsDict):
action_entrance(paramsDict)
actionContainer.addAction(action_entrance, AccActionsContainer.ENTRANCE)
actionContainer.addAction(action_exit, AccActionsContainer.EXIT)
accLattice.trackBunch(bunch,
paramsDict=paramsDict,
actionContainer=actionContainer)
return results_arr
print "dRFPhasem = ", dRFPhasem
#///////////////////////////////////////////////////////////
print "==============BEFORE============================"
b.dumpBunch()
print "=========================================="
#=====track action ============
def bodyAction(paramsDict):
node = paramsDict["node"]
node.track(paramsDict)
accContainer = AccActionsContainer()
accContainer.addAction(bodyAction, AccActionsContainer.BODY)
paramsDict = {}
paramsDict["bunch"] = b
lattice.trackActions(accContainer, paramsDict)
print "=============AFTER============================="
b.dumpBunch()
print "=========================================="
print "lattice length=", lattice.getLength()
print "beta=", b.getSyncParticle().beta()
print "TEAPOT time[sec]=", b.getSyncParticle().time()
print "SIMPLE time[sec]=", lattice.getLength() / (b.getSyncParticle().beta() *
2.99792458e+8)