def __init__(self, name="Monitor no name"):
"""
Constructor. Creates the aperutre element.
"""
NodeTEAPOT.__init__(self, name)
self.setType("monitor teapot")
self.twiss = BunchTwissAnalysis()
def __init__(self, file, order, nodispersion, emitnorm):
self.file_out = file
self.order = order
self.bunchtwissanalysis = BunchTwissAnalysis()
if (nodispersion == False):
self.dispterm = -1
else:
self.dispterm = 1
if (emitnorm == True):
self.emitnormterm = 1
else:
self.emitnormterm = -1
class MomentsSetMember:
"""
This class delivers the beam moments
"""
def __init__(self, file, order, nodispersion, emitnorm):
self.file_out = file
self.order = order
self.bunchtwissanalysis = BunchTwissAnalysis()
if (nodispersion == False):
self.dispterm = -1
else:
self.dispterm = 1
if (emitnorm == True):
self.emitnormterm = 1
else:
self.emitnormterm = -1
def writeMoments(self, s, bunch, lattlength=0):
sp = bunch.getSyncParticle()
time = sp.time()
if lattlength > 0:
time = sp.time() / (lattlength / (sp.beta() * speed_of_light))
self.bunchtwissanalysis.computeBunchMoments(bunch, self.order,
self.dispterm,
self.emitnormterm)
# if mpi operations are enabled, this section of code will
# determine the rank of the present node
rank = 0 # default is primary node
mpi_init = orbit_mpi.MPI_Initialized()
comm = orbit_mpi.mpi_comm.MPI_COMM_WORLD
if (mpi_init):
rank = orbit_mpi.MPI_Comm_rank(comm)
# only the primary node needs to output the calculated information
if (rank == 0):
self.file_out.write(str(s) + "\t" + str(time) + "\t")
for i in range(0, self.order + 1):
for j in range(0, i + 1):
self.file_out.write(
str(self.bunchtwissanalysis.getBunchMoment(i - j, j)) +
"\t")
self.file_out.write("\n")
def resetFile(self, file):
self.file_out = file
class StatLatsSetMember:
"""
This class delivers the statistical twiss parameters
"""
def __init__(self, file):
self.file_out = file
self.bunchtwissanalysis = BunchTwissAnalysis()
def writeStatLats(self, s, bunch, lattlength=0):
self.bunchtwissanalysis.analyzeBunch(bunch)
emitx = self.bunchtwissanalysis.getEmittance(0)
betax = self.bunchtwissanalysis.getBeta(0)
alphax = self.bunchtwissanalysis.getAlpha(0)
betay = self.bunchtwissanalysis.getBeta(1)
alphay = self.bunchtwissanalysis.getAlpha(1)
emity = self.bunchtwissanalysis.getEmittance(1)
dispersionx = self.bunchtwissanalysis.getDispersion(0)
ddispersionx = self.bunchtwissanalysis.getDispersionDerivative(0)
#dispersiony = self.bunchtwissanalysis.getDispersion(1, bunch)
#ddispersiony = self.bunchtwissanalysis.getDispersionDerivative(1, bunch)
sp = bunch.getSyncParticle()
time = sp.time()
if lattlength > 0:
time = sp.time() / (lattlength / (sp.beta() * speed_of_light))
# if mpi operations are enabled, this section of code will
# determine the rank of the present node
rank = 0 # default is primary node
mpi_init = orbit_mpi.MPI_Initialized()
comm = orbit_mpi.mpi_comm.MPI_COMM_WORLD
if (mpi_init):
rank = orbit_mpi.MPI_Comm_rank(comm)
# only the primary node needs to output the calculated information
if (rank == 0):
self.file_out.write(
str(s) + "\t" + str(time) + "\t" + str(emitx) + "\t" +
str(emity) + "\t" + str(betax) + "\t" + str(betay) + "\t" +
str(alphax) + "\t" + str(alphay) + "\t" + str(dispersionx) +
"\t" + str(ddispersionx) + "\n")
def closeStatLats(self):
self.file_out.close()
def resetFile(self, file):
self.file_out = file
def __init__(self,frequency = 402.5e+6): self.bunch_frequency = frequency #set H- mass #self.bunch.mass(0.9382723 + 2*0.000511) self.bunch = Bunch() self.bunch.getSyncParticle().kinEnergy(0.0025) self.init_coords = (0.,0.,0.,0.,0.,0.) self.bunch.mass(0.939294) self.bunch.charge(-1.0) self.c = 2.99792458e+8 # speed of light in m/sec self.beam_current = 38.0 # beam current in mA , design = 38 mA self.rf_wave_lenght = self.c/self.bunch_frequency self.si_e_charge = 1.6021773e-19 #---------------------------------------- self.twiss_analysis = BunchTwissAnalysis()
class MomentsSetMember:
"""
This class delivers the beam moments
"""
def __init__(self, file, order, nodispersion, emitnorm):
self.file_out = file
self.order = order
self.bunchtwissanalysis = BunchTwissAnalysis()
if(nodispersion == False):
self.dispterm = -1
else:
self.dispterm = 1
if(emitnorm == True):
self.emitnormterm = 1
else:
self.emitnormterm = -1
def writeMoments(self, s, bunch, lattlength = 0 ):
sp = bunch.getSyncParticle()
time = sp.time()
if lattlength > 0:
time = sp.time()/(lattlength/(sp.beta() * speed_of_light))
self.bunchtwissanalysis.computeBunchMoments(bunch, self.order, self.dispterm, self.emitnormterm)
# if mpi operations are enabled, this section of code will
# determine the rank of the present node
rank = 0 # default is primary node
mpi_init = orbit_mpi.MPI_Initialized()
comm = orbit_mpi.mpi_comm.MPI_COMM_WORLD
if (mpi_init):
rank = orbit_mpi.MPI_Comm_rank(comm)
# only the primary node needs to output the calculated information
if (rank == 0):
self.file_out.write(str(s) + "\t" + str(time) + "\t")
for i in range(0,self.order+1):
for j in range(0,i+1):
self.file_out.write(str(self.bunchtwissanalysis.getBunchMoment(i-j,j)) + "\t")
self.file_out.write("\n")
def resetFile(self, file):
self.file_out = file
class StatLatsSetMember: """ This class delivers the statistical twiss parameters """ def __init__(self, file): self.file_out = file self.bunchtwissanalysis = BunchTwissAnalysis() def writeStatLats(self, s, bunch, lattlength = 0): self.bunchtwissanalysis.analyzeBunch(bunch) emitx = self.bunchtwissanalysis.getEmittance(0) betax = self.bunchtwissanalysis.getBeta(0) alphax = self.bunchtwissanalysis.getAlpha(0) betay = self.bunchtwissanalysis.getBeta(1) alphay = self.bunchtwissanalysis.getAlpha(1) emity = self.bunchtwissanalysis.getEmittance(1) dispersionx = self.bunchtwissanalysis.getDispersion(0) ddispersionx = self.bunchtwissanalysis.getDispersionDerivative(0) #dispersiony = self.bunchtwissanalysis.getDispersion(1, bunch) #ddispersiony = self.bunchtwissanalysis.getDispersionDerivative(1, bunch) sp = bunch.getSyncParticle() time = sp.time() if lattlength > 0: time = sp.time()/(lattlength/(sp.beta() * speed_of_light)) # if mpi operations are enabled, this section of code will # determine the rank of the present node rank = 0 # default is primary node mpi_init = orbit_mpi.MPI_Initialized() comm = orbit_mpi.mpi_comm.MPI_COMM_WORLD if (mpi_init): rank = orbit_mpi.MPI_Comm_rank(comm) # only the primary node needs to output the calculated information if (rank == 0): self.file_out.write(str(s) + "\t" + str(time) + "\t" + str(emitx)+ "\t" + str(emity)+ "\t" + str(betax)+ "\t" + str(betay)+ "\t" + str(alphax)+ "\t" + str(alphay) + "\t" + str(dispersionx) + "\t" + str(ddispersionx) +"\n") def closeStatLats(self): self.file_out.close() def resetFile(self, file): self.file_out = file
def __init__(self, lattice, start_node=None, stop_node=None):
#---- change in the corrector field
self.deltaB = 0.001
#----------------------------------
self.lattice = lattice
self.start_node = start_node
self.stop_node = stop_node
self.bpm_node_arr = []
self.transvBPM_arr = []
self.dch_node_arr = []
self.dcv_node_arr = []
#-------------------------
self.quad_node_arr = []
self.quad_transvBPM_arr = []
#-------------------------
self._updateBPM_Nodes()
self._updateDC_Nodes(None, DCorrectorH)
self._updateDC_Nodes(None, DCorrectorV)
self._updateQuad_Nodes()
#-------------------------
self.twiss_analysis = BunchTwissAnalysis()
def __init__(self, file, order, nodispersion, emitnorm): self.file_out = file self.order = order self.bunchtwissanalysis = BunchTwissAnalysis() if(nodispersion == False): self.dispterm = -1 else: self.dispterm = 1 if(emitnorm == True): self.emitnormterm = 1 else: self.emitnormterm = -1
class BPMSignal: """ This class delivers the average value for coordinate x and y """ def __init__(self): self.bunchtwissanalysis = BunchTwissAnalysis() self.xAvg = 0.0 self.yAvg = 0.0 self.xpAvg = 0.0 self.ypAvg = 0.0 def analyzeSignal(self, bunch): self.bunchtwissanalysis.analyzeBunch(bunch) # if mpi operations are enabled, this section of code will # determine the rank of the present node rank = 0 # default is primary node mpi_init = orbit_mpi.MPI_Initialized() comm = orbit_mpi.mpi_comm.MPI_COMM_WORLD if (mpi_init): rank = orbit_mpi.MPI_Comm_rank(comm) # only the primary node needs to output the calculated information if (rank == 0): self.xAvg = self.bunchtwissanalysis.getAverage(0) self.xpAvg = self.bunchtwissanalysis.getAverage(1) self.yAvg = self.bunchtwissanalysis.getAverage(2) self.ypAvg = self.bunchtwissanalysis.getAverage(3) def getSignalX(self): return self.xAvg def getSignalXP(self): return self.xpAvg def getSignalY(self): return self.yAvg def getSignalYP(self): return self.ypAvg
class BPMSignal:
"""
This class delivers the average value for coordinate x and y
"""
def __init__(self):
self.bunchtwissanalysis = BunchTwissAnalysis()
self.xAvg = 0.0
self.yAvg = 0.0
self.xpAvg = 0.0
self.ypAvg = 0.0
def analyzeSignal(self, bunch):
self.bunchtwissanalysis.analyzeBunch(bunch)
# if mpi operations are enabled, this section of code will
# determine the rank of the present node
rank = 0 # default is primary node
mpi_init = orbit_mpi.MPI_Initialized()
comm = orbit_mpi.mpi_comm.MPI_COMM_WORLD
if (mpi_init):
rank = orbit_mpi.MPI_Comm_rank(comm)
# only the primary node needs to output the calculated information
if (rank == 0):
self.xAvg = self.bunchtwissanalysis.getAverage(0)
self.xpAvg = self.bunchtwissanalysis.getAverage(1)
self.yAvg = self.bunchtwissanalysis.getAverage(2)
self.ypAvg = self.bunchtwissanalysis.getAverage(3)
def getSignalX(self):
return self.xAvg
def getSignalXP(self):
return self.xpAvg
def getSignalY(self):
return self.yAvg
def getSignalYP(self):
return self.ypAvg
def bunchCentering(bunch):
"""
Bunch center after generating can have small deviation from the (0,0,0,0,0,0)
This function will force centering the bunch.
"""
twiss_analysis = BunchTwissAnalysis()
twiss_analysis.analyzeBunch(bunch)
#-----------------------------------------------
# let's center the beam
(x_avg, y_avg) = (twiss_analysis.getAverage(0),
twiss_analysis.getAverage(2))
(xp_avg, yp_avg) = (twiss_analysis.getAverage(1),
twiss_analysis.getAverage(3))
(z_avg, dE_avg) = (twiss_analysis.getAverage(4),
twiss_analysis.getAverage(5))
for part_id in range(bunch.getSize()):
bunch.x(part_id, bunch.x(part_id) - x_avg)
bunch.y(part_id, bunch.y(part_id) - y_avg)
bunch.xp(part_id, bunch.xp(part_id) - xp_avg)
bunch.yp(part_id, bunch.yp(part_id) - yp_avg)
bunch.z(part_id, bunch.z(part_id) - z_avg)
bunch.dE(part_id, bunch.dE(part_id) - dE_avg)
#-----------------------------------------------
return (x_avg, y_avg, dE_avg)
class MonitorTEAPOT(NodeTEAPOT):
"""
Monitor TEAPOT element.
"""
def __init__(self, name="Monitor no name"):
"""
Constructor. Creates the aperutre element.
"""
NodeTEAPOT.__init__(self, name)
self.setType("monitor teapot")
self.twiss = BunchTwissAnalysis()
def track(self, paramsDict):
"""
The bunchtuneanalysis-teapot class implementation of the AccNodeBunchTracker class track(probe) method.
"""
length = self.getLength(self.getActivePartIndex())
bunch = paramsDict["bunch"]
self.twiss.analyzeBunch(bunch)
self.addParam("xAvg", self.twiss.getAverage(0))
self.addParam("xpAvg", self.twiss.getAverage(1))
self.addParam("yAvg", self.twiss.getAverage(2))
self.addParam("ypAvg", self.twiss.getAverage(3))
def __init__(self, file):
self.file_out = file
self.bunchtwissanalysis = BunchTwissAnalysis()
def __init__(self): self.bunchtwissanalysis = BunchTwissAnalysis() self.xAvg = 0.0 self.yAvg = 0.0 self.xpAvg = 0.0 self.ypAvg = 0.0
def __init__(self, filename):
self.file_out = open(filename, "a")
self.bunchtwissanalysis = BunchTwissAnalysis()
print "Bunch Generation completed."
bnch02 = accLattice.getNodeForName("LI_MEBT1:BNCH02:Rg01")
bnch02_ind = accLattice.getNodeIndex(bnch02)
#set up design - arrival times at each fist gaps of all RF cavities
accLattice.trackDesignBunch(bunch_in)
print "Design tracking completed."
#real tracking of the bunch. Now we have the bunch at the BNCH02 entrance
accLattice.trackBunch(bunch_in, None, None, 0, bnch02_ind - 1)
#---- memorize initial coordinates of particles in PartAttributes
copyCoordsToInitCoordsAttr(bunch_in)
twiss_analysis = BunchTwissAnalysis()
bunch = bunch_in
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.
(alphaZ, betaZ, emittZ) = (twiss_analysis.getTwiss(2)[0],
twiss_analysis.getTwiss(2)[1],
twiss_analysis.getTwiss(2)[3])
gammaZ = (1.0 + alphaZ**2) / betaZ
dE_rms = math.sqrt(gammaZ * emittZ * 1.0e+6) * 1000.
z_to_phase_coeff = bunch_gen.getZtoPhaseCoeff(bunch)
sy = np.sqrt(betay * emitlimy)
spy = np.sqrt(gy * emitlimy)
for i in range(50):
for j in range(100):
b.addParticle(10 * sx * float(j) / 100 * cos(float(i) / 50 * pi / 2),
0,
10 * sx * float(j) / 100 * sin(float(i) / 50 * pi / 2),
0, 0, 0.000013229) #0.000013229
b.partAttrValue("ParticleIdNumber",
b.getSize() - 1, 0,
b.getSize() - 1)
"""--------------------------------- end ---------------------------------"""
paramsDict = {}
lostbunch = Bunch()
paramsDict["lostbunch"] = lostbunch
paramsDict["bunch"] = b
b.dumpBunch("./6D_initial_bunch_1p.dat")
ana = BunchTwissAnalysis()
for i in range(5000):
print("turn: " + str(i))
print("remaining: " + str(b.getSize()) + "\n")
lattice.trackBunch(b, paramsDict)
b.dumpBunch("./6D_final_bunch_1p.dat")
#bunch_in = bunch_gen.getBunch(nParticles = 20000, distributorClass = WaterBagDist3D)
bunch_in = bunch_gen.getBunch(nParticles = 20000, distributorClass = GaussDist3D)
#bunch_in = bunch_gen.getBunch(nParticles = 20000, distributorClass = KVDist3D)
print "Bunch Generation completed."
#set up design
accLattice.trackDesignBunch(bunch_in)
print "Design tracking completed."
#track through the lattice
paramsDict = {"test_pos":0.,"count":0}
actionContainer = AccActionsContainer("Test Design Bunch Tracking")
twiss_analysis = BunchTwissAnalysis()
print " N node position sizeX sizeY sizeZ sizeZdeg sizeXP sizeYP size_dE eKin Nparts"
file_out = open("pyorbit_scl_sizes_ekin.dat","w")
file_out.write(" N node position sizeX sizeY sizeZ sizeZdeg sizeXP sizeYP sizedE eKin Nparts \n")
def action_entrance(paramsDict):
if(isinstance(paramsDict["parentNode"],AccLattice)):
node = paramsDict["node"]
pos = paramsDict["test_pos"]
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]*twiss_analysis.getTwiss(1)[3])*1000.
z_rms = math.sqrt(twiss_analysis.getTwiss(2)[1]*twiss_analysis.getTwiss(2)[3])*1000.
z_rms_deg = bunch_gen.getZtoPhaseCoeff(bunch)*z_rms/1000.0
else:
continue
for rf_gap in rf_gaps:
etl0 = rf_gap.getParam("E0TL")
rf_gap.setParam("E0TL", etl0 * coeff)
print "Design tracking completed."
collimation_dict = {"MEBT": 16.0, "DTL": 12.5, "CCL": 15.0, "SCL": 38.0}
loss_part_pos_arr = []
#track through the lattice
paramsDict = {"test_pos": 0., "count": 0}
actionContainer = AccActionsContainer("Test Design Bunch Tracking")
twiss_analysis = BunchTwissAnalysis()
bunchExtremaCalculator = BunchExtremaCalculator()
(xMin, xMax, yMin, yMax, zMin,
zMax) = bunchExtremaCalculator.extremaXYZ(bunch_in)
print " N node position sizeX sizeY sizeZ sizeZdeg sizeXP sizeYP size_dE eKin Nparts"
file_out = open("pyorbit_scl_sizes_ekin.dat", "w")
file_out.write(
" N node position sizeX sizeY sizeZ sizeZdeg sizeXP sizeYP sizedE eKin Nparts \n"
)
fl_loss_pos_coord_out = open("loss_parts_pos_coord+1.dat", "w")
def action_entrance(paramsDict):
print " aplha beta emitt Y= ( %8.2f, %8.2f, %10.3e )"%(alphaY,betaY,emittY) print " aplha beta emitt Z= ( %8.2f, %8.2f, %10.3e )"%(alphaZ,betaZ,emittZ) twissX = TwissContainer(alphaX,betaX,emittX) twissY = TwissContainer(alphaY,betaY,emittY) twissZ = TwissContainer(alphaZ,betaZ,emittZ) distGen = GaussDist3D(twissX,twissY,twissZ) distGen = WaterBagDist3D(twissX,twissY,twissZ) distGen = KVDist3D(twissX,twissY,twissZ) n_parts = 1000 for i in range(n_parts): (x,xp,y,yp,z,dE) = distGen.getCoordinates() bunch_ini.addParticle(x,xp,y,yp,z,dE) twiss_analysis = BunchTwissAnalysis() twiss_analysis.analyzeBunch(bunch_ini) print "============Initial Bunch Twiss parameters ==================" print "X Twiss = ( %8.2f, %8.2f, %8.2f, %10.3e )"%twiss_analysis.getTwiss(0) print "Y Twiss = ( %8.2f, %8.2f, %8.2f, %10.3e )"%twiss_analysis.getTwiss(1) print "Z Twiss = ( %8.2f, %8.2f, %8.2f, %10.3e )"%twiss_analysis.getTwiss(2) 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. print "Initial Bunch (x_rms,y_rms,z_rms) = ( %6.3f, %6.3f, %6.3f )"%(x_rms,y_rms,z_rms) print "============================================================" bunch_tmp = Bunch() bunch_ini.copyBunchTo(bunch_tmp)
n_parts = 10000 for i in range(n_parts): (x,xp,y,yp,z,dE) = distGen.getCoordinates() b.addParticle(x,xp,y,yp,z,dE) b.compress() syncPart = b.getSyncParticle() syncPart.kinEnergy(TK) # copy the initial bunch to another to track through TEAPOT Quad b1 = Bunch() b.copyBunchTo(b1) twiss_analysis = BunchTwissAnalysis() twiss_analysis.analyzeBunch(b) print "============before==================" print "X Twiss =",twiss_analysis.getTwiss(0) print "Y Twiss =",twiss_analysis.getTwiss(1) print "Z Twiss =",twiss_analysis.getTwiss(2) #b.dumpBunch() G = 30.0 # [T/m] length = 0.1 # [m] fieldSource = FieldSource(G) tracker = RungeKuttaTracker(length) print "Tracker Entrance plane (a,b,c,d)=",tracker.entrancePlane() print "Tracker Exit plane (a,b,c,d)=",tracker.exitPlane() # the spatial eps is useless because we have quad field = 0 on axis for the syncPart
class TrajectoryCorrection:
"""
Class to perform correction of the trajectory in the linac type lattice
or in the part of it. To correct trajectories the Linac DCorrectorH and
DCorrectorV objects will be used. The trajectory is defined by BPMs nodes
that specifically defined for this trajectory correction. The correction
algorithm will build response matrix for correctors (assuming no tilt) and
apply the correctors field to provide the best correction to the target
values for each BPM. By default the goal for all BPMs is 0., but user can
redefine these values and set of BPMs. User also can redefine what
dipole correctors will be used.
"""
def __init__(self, lattice, start_node=None, stop_node=None):
#---- change in the corrector field
self.deltaB = 0.001
#----------------------------------
self.lattice = lattice
self.start_node = start_node
self.stop_node = stop_node
self.bpm_node_arr = []
self.transvBPM_arr = []
self.dch_node_arr = []
self.dcv_node_arr = []
#-------------------------
self.quad_node_arr = []
self.quad_transvBPM_arr = []
#-------------------------
self._updateBPM_Nodes()
self._updateDC_Nodes(None, DCorrectorH)
self._updateDC_Nodes(None, DCorrectorV)
self._updateQuad_Nodes()
#-------------------------
self.twiss_analysis = BunchTwissAnalysis()
def setStartStopNodes(self, start_node=None, stop_node=None):
self.start_node = start_node
self.stop_node = stop_node
self._updateBPM_Nodes()
self._updateDC_Nodes(None, DCorrectorH)
self._updateDC_Nodes(None, DCorrectorV)
self._updateQuad_Nodes()
def _getStartStopIndexes(self):
"""
Returns the start and stop indexes of the start and stop nodes
"""
start_ind = -1
if (self.start_node != None):
start_ind = self.lattice.getNodeIndex(self.start_node)
stop_ind = -1
if (self.stop_node != None):
stop_ind = self.lattice.getNodeIndex(self.stop_node)
return (start_ind, stop_ind)
def _returnFilteredNodes(self, nodes):
"""
Returns list of nodes between start and stop nodes.
"""
(start_ind, stop_ind) = self._getStartStopIndexes()
nodes_tmp = []
for node in nodes:
ind = self.lattice.getNodeIndex(node)
res = True
if (start_ind >= 0):
if (ind < start_ind):
res = False
if (stop_ind >= 0):
if (ind > stop_ind):
res = False
if (res):
nodes_tmp.append(node)
return nodes_tmp
def _updateBPM_Nodes(self, bpms=None):
"""
Updates BPM nodes between start_node and stop_node.
Returns array of bpm nodes.
"""
node_pos_dict = self.lattice.getNodePositionsDict()
self.cleanBPM_Nodes()
self.bpm_node_arr = []
bpm_nodes = []
if (bpms == None):
markers = self.lattice.getNodesOfClass(MarkerLinacNode)
for node in markers:
if (node.getName().find("BPM") >= 0):
bpm_nodes.append(node)
else:
bpm_nodes = bpms
self.bpm_node_arr = self._returnFilteredNodes(bpm_nodes)
for bpm in self.bpm_node_arr:
(start_pos, end_pos) = node_pos_dict[bpm]
transvBPM = TransverseBPM(self, bpm)
transvBPM.setPosition(start_pos)
bpm.addChildNode(transvBPM, AccNode.ENTRANCE)
self.transvBPM_arr.append(transvBPM)
return self.bpm_node_arr
def _updateDC_Nodes(self, nodes=None, class_type=None):
"""
Updates DCorrector nodes that we use for orbit correction
"""
if (class_type == None): return
dc_node_arr = None
if (class_type == DCorrectorH):
dc_node_arr = self.dch_node_arr
if (class_type == DCorrectorV):
dc_node_arr = self.dcv_node_arr
if (dc_node_arr == None): return
del dc_node_arr[:]
if (nodes == None):
dc_node_arr += self.lattice.getNodesOfClass(class_type)
else:
for node in nodes:
if (isinstance(node, class_type)):
dc_node_arr.append(node)
node_arr = self._returnFilteredNodes(dc_node_arr)
del dc_node_arr[:]
dc_node_arr += node_arr
return dc_node_arr
def _updateQuad_Nodes(self, nodes=None):
"""
Updates Quad nodes with TransverseBPM instances
"""
node_pos_dict = self.lattice.getNodePositionsDict()
self.cleanQuad_Nodes()
self.quad_node_arr = []
quad_arr = []
if (nodes == None):
quad_arr = self.lattice.getNodesOfClass(Quad)
else:
quad_arr = nodes
quad_arr = self._returnFilteredNodes(quad_arr)
for quad in quad_arr:
(start_pos, end_pos) = node_pos_dict[quad]
transvBPM = TransverseBPM(self, quad, "_entrance")
transvBPM.setPosition(start_pos)
quad.addChildNode(transvBPM, AccNode.ENTRANCE)
self.quad_transvBPM_arr.append(transvBPM)
transvBPM = TransverseBPM(self, quad, "_exit")
transvBPM.setPosition(end_pos)
quad.addChildNode(transvBPM, AccNode.EXIT)
self.quad_transvBPM_arr.append(transvBPM)
self.quad_node_arr = quad_arr
return self.quad_node_arr
def setBPMs(self, bpms):
"""
Sets custom bpm array for analysis.
In reality they can be any nodes.
Returns array of bpm nodes.
"""
return self._updateBPM_Nodes(bpms)
def getBPMs(self):
"""
Returns array of bpm nodes.
"""
return self.bpm_node_arr
def setDCHs(self, dchs):
"""
Sets the DCorrectorHs
"""
self._updateDC_Nodes(dchs, DCorrectorH)
return self.dch_node_arr
def getDCHs(self):
"""
Returns DCorrectorHs array
"""
return self.dch_node_arr
def setDCHVs(self, dchs):
"""
Sets the DCorrectorVs
"""
self._updateDC_Nodes(dchs, DCorrectorV)
return self.dcv_node_arr
def getDCVs(self):
"""
Returns DCorrectorVs array
"""
return self.dcv_node_arr
def setQuads(self, quads):
"""
Sets Quads
"""
self._updateQuad_Nodes(quads)
return self.quad_node_arr
def getQuads(self):
"""
Returns Quads array
"""
return self.quad_node_arr
def getTransverseBPMs(self):
return self.transvBPM_arr
def getQuadTransverseBPMs(self):
return self.quad_transvBPM_arr
def getTransverseBPMforBPM(self, bpm):
"""
Retuns the TransverseBPM instance for particular BPM.
"""
for child in bpm.getChildNodes(AccNode.ENTRANCE):
if (isinstance(child, TransverseBPM)):
return child
return None
def cleanBPM_Nodes(self):
"""
Removes TransverseBPM child nodes from BPM nodes
"""
for bpm in self.bpm_node_arr:
child_arr = bpm.getChildNodes(AccNode.ENTRANCE)
transvBPM = None
for child in child_arr:
if (isinstance(child, TransverseBPM)):
transvBPM = child
break
if (transvBPM != None):
child_arr.remove(transvBPM)
self.transvBPM_arr = []
def cleanQuad_Nodes(self):
"""
Removes TransverseBPM child nodes from Quad nodes
"""
for node in self.quad_node_arr:
for place in [AccNode.ENTRANCE, AccNode.EXIT]:
child_arr = node.getChildNodes(place)
transvBPM = None
for child in child_arr:
if (isinstance(child, TransverseBPM)):
transvBPM = child
break
if (transvBPM != None):
child_arr.remove(transvBPM)
self.quad_transvBPM_arr = []
def correctTrajectory(self, bunch_initial):
"""
This method will calculate and applyes fields to dipole correctors
to achive minimal beam deviation from the center at all BPMs.
LSQM method is used.
"""
bunch_in = self.makeOneParticleBunch(bunch_initial)
#--------------------------
bunch_init = Bunch()
bunch_in.copyEmptyBunchTo(bunch_init)
#---- add one particle to the bunch_init
x0 = 0.
xp0 = 0.
y0 = 0.
yp0 = 0.
z0 = 0.
dE = 0.
bunch_init.addParticle(x0, xp0, y0, yp0, z0, dE)
#----------------------------------
n_bpms = len(self.bpm_node_arr)
n_corrs_hor = len(self.dch_node_arr)
n_corrs_ver = len(self.dcv_node_arr)
horResponceMtrx = Matrix(n_bpms, n_corrs_hor)
verResponceMtrx = Matrix(n_bpms, n_corrs_ver)
self._calculatBPM_Matrix(bunch_init,
horResponceMtrx,
self.dch_node_arr,
axis=0)
self._calculatBPM_Matrix(bunch_init,
verResponceMtrx,
self.dcv_node_arr,
axis=1)
#printM(horResponceMtrx)
horResponceMtrxTr = horResponceMtrx.copy()
horResponceMtrxTr.transpose()
verResponceMtrxTr = verResponceMtrx.copy()
verResponceMtrxTr.transpose()
horAmtrx = (horResponceMtrxTr.mult(horResponceMtrx)).invert()
verAmtrx = (verResponceMtrxTr.mult(verResponceMtrx)).invert()
#printM(horAmtrx)
#print "det(horAmtrx) = ",horAmtrx.det()
horATmtrx = horAmtrx.mult(horResponceMtrxTr)
verATmtrx = verAmtrx.mult(verResponceMtrxTr)
#---- matrices for LSQM are ready - now use initial bunch to get
#---- the trajectory that should be flattened
bunch_init = Bunch()
bunch_in.copyBunchTo(bunch_init)
(bpm_value_hor_arr,
bpm_value_ver_arr) = self.calculateTrajectory(bunch_init)
horBPM_V = PhaseVector(len(bpm_value_hor_arr))
for ind in range(len(bpm_value_hor_arr)):
horBPM_V.set(ind, bpm_value_hor_arr[ind])
verBPM_V = PhaseVector(len(bpm_value_ver_arr))
for ind in range(len(bpm_value_ver_arr)):
verBPM_V.set(ind, bpm_value_ver_arr[ind])
dch_val_V = horATmtrx.mult(horBPM_V)
dcv_val_V = verATmtrx.mult(verBPM_V)
#print "debug =========== final DCH fields ================"
for ind in range(dch_val_V.size()):
dc = self.dch_node_arr[ind]
delta_field = dch_val_V.get(ind)
#print "debug corr =",dc.getName()," init field [T] =",dc.getParam("B")," delta [T] =",delta_field
dc.setParam("B", dc.getParam("B") - delta_field)
#print "debug =========== final DCV fields ================"
for ind in range(dcv_val_V.size()):
dc = self.dcv_node_arr[ind]
delta_field = dcv_val_V.get(ind)
#print "debug corr =",dc.getName()," init field [T] =",dc.getParam("B")," delta [T] =",delta_field
dc.setParam("B", dc.getParam("B") - delta_field)
def _calculatBPM_Matrix(self,
bunch_init,
responceMtrx,
dc_node_arr,
axis=None):
corr_field_arr = []
for dc_node in dc_node_arr:
corr_field_arr.append(dc_node.getParam("B"))
#---------------------------------------------
bpm_value_init_arr = self.calculateTrajectory(bunch_init)[axis]
for dc_node_ind in range(len(dc_node_arr)):
dc_node = dc_node_arr[dc_node_ind]
field = corr_field_arr[dc_node_ind] + self.deltaB
dc_node.setParam("B", field)
bpm_value_arr = self.calculateTrajectory(bunch_init)[axis]
for bpm_ind in range(len(self.transvBPM_arr)):
deltaVal = bpm_value_arr[bpm_ind] - bpm_value_init_arr[bpm_ind]
derivative = deltaVal / self.deltaB
responceMtrx.set(bpm_ind, dc_node_ind, derivative)
dc_node.setParam("B", corr_field_arr[dc_node_ind])
def calculateTrajectory(self, bunch_init, print_info=False):
"""
It tracks bunch with one particle through the lattice and
fills out the BPM data in self.transvBPM_arr
Returns arrays of x,y corrdinates of the bunch in
(pm_value_hor_arr,bpm_value_ver_arr) tuple.
"""
bunch = Bunch()
bunch_init.copyBunchTo(bunch)
(start_ind, stop_ind) = self._getStartStopIndexes()
self.lattice.trackDesignBunch(bunch, None, None, start_ind, stop_ind)
self.lattice.trackBunch(bunch, None, None, start_ind, stop_ind)
bpm_value_hor_arr = []
bpm_value_ver_arr = []
for transvBPM in self.transvBPM_arr:
val_hor = transvBPM.getCoordinates()[0]
val_ver = transvBPM.getCoordinates()[2]
bpm_value_hor_arr.append(val_hor)
bpm_value_ver_arr.append(val_ver)
#---- Start trajectory printing
if (print_info):
print " index name position[m] x[mm] y[mm] "
for ind in range(len(bpm_value_hor_arr)):
hor_val = bpm_value_hor_arr[ind] * 1000.
ver_val = bpm_value_ver_arr[ind] * 1000.
bpm_node = self.bpm_node_arr[ind]
transvBPM = self.transvBPM_arr[ind]
pos = transvBPM.getPosition()
print " %2d %20s %8.3f %+6.2f %+6.2f " % (
ind, bpm_node.getName(), pos, hor_val, ver_val)
#---- Stop trajectory printing
return (bpm_value_hor_arr, bpm_value_ver_arr)
def makeOneParticleBunch(self, bunch_init):
"""
Returns the bunch with one particle which coordinates
are equal to the average coordinates in the initial
bunch.
"""
bunch_out = Bunch()
bunch_init.copyEmptyBunchTo(bunch_out)
x_avg = 0.
xp_avg = 0.
y_avg = 0.
yp_avg = 0.
z_avg = 0.
dE_avg = 0.
nParts = bunch_init.getSizeGlobal()
if (nParts > 0):
self.twiss_analysis.analyzeBunch(bunch_init)
x_avg = self.twiss_analysis.getAverage(0)
xp_avg = self.twiss_analysis.getAverage(1)
y_avg = self.twiss_analysis.getAverage(2)
yp_avg = self.twiss_analysis.getAverage(3)
z_avg = self.twiss_analysis.getAverage(4)
dE_avg = self.twiss_analysis.getAverage(5)
#----------------------------------------
#---- add one particle to the bunch_out
bunch_out.addParticle(x_avg, xp_avg, y_avg, yp_avg, z_avg, dE_avg)
return bunch_out
# ID the particles
particleidnumber = ParticleIdNumber()
particleidnumber.addParticleIdNumbers(bunch) # Give them uniquely number ids
paramsDict["bunch"] = bunch
lostbunch = Bunch()
paramsDict["lostbunch"] = lostbunch
lostbunch.addPartAttr("LostParticleAttributes")
#----------------------------------------------
# Start simulation
#----------------------------------------------
#Prepare the analysis class that will look at emittances, etc.
bunchtwissanalysis = BunchTwissAnalysis()
# #----------------------------------------------------
# # Add space charge nodes
# #----------------------------------------------------
# ''
# # Make a SC solver using frozen potential
# print '\nSetting up the space charge calculations ...'
# sc_path_length_min = 0.00000001
# bunchtwissanalysis.analyzeBunch(bunch)
#
# # consider line density based on measurement provided in file
# density=np.loadtxt('density_2.dat', dtype=float)
# LineDensity = InterpolatedLineDensityProfile(-12.5,12.5,density.tolist())
#
# # simple Gaussian line density - sufficient to create a certain tune spread - density defined above not required
#b = Bunch()
#b.readBunch('./intensity_11/n_parts/initial_bunch.dat')
paramsDict = {}
lostbunch = Bunch()
paramsDict["lostbunch"] = lostbunch
paramsDict["bunch"] = b
#todo:
pre_list = './Files_shifted'
test_name = 'intensity_10'
b.dumpBunch(pre_list + "/initial_bunch_" + str(test_name) + ".dat")
ana = BunchTwissAnalysis()
f = open(pre_list + "/tracking_" + str(test_name) + ".txt", 'w')
for i in range(500):
ana.analyzeBunch(b)
if rank == 0:
f.write(
str(i) + " " + str(ana.getEmittanceNormalized(0)) + " " +
str(ana.getEmittanceNormalized(1)) + " " + str(ana.getBeta(0)) +
" " + str(ana.getBeta(1)) + " " + str(lostbunch.getSize()) + "\n")
lattice.trackBunch(b, paramsDict)
class SNS_Linac_BunchGenerator: """ Generates the pyORBIT SNS Linac Bunches using the Gauss distribution. Twiss parameters have the following units: x in [m], xp in [rad] and the X and Y emittances are un-normalized. The longitudinal emittance is in [GeV*m]. """ def __init__(self,frequency = 402.5e+6): self.bunch_frequency = frequency #set H- mass #self.bunch.mass(0.9382723 + 2*0.000511) self.bunch = Bunch() self.bunch.getSyncParticle().kinEnergy(0.0025) self.init_coords = (0.,0.,0.,0.,0.,0.) self.bunch.mass(0.939294) self.bunch.charge(-1.0) self.c = 2.99792458e+8 # speed of light in m/sec self.beam_current = 38.0 # beam current in mA , design = 38 mA self.rf_wave_lenght = self.c/self.bunch_frequency self.si_e_charge = 1.6021773e-19 #---------------------------------------- self.twiss_analysis = BunchTwissAnalysis() def setInitialCorrdsCenter(self,x0,xp0,y0,yp0,z0,dE0): self.init_coords = (x0,xp0,y0,yp0,z0,dE0) def getInitialCorrdsCenter(self): return self.init_coords def setParticleCharge(self,charge): """ Sets the particle charge H- => -1.0 and proton => +1.0 """ self.bunch.charge(charge) def getKinEnergy(self): """ Returns the kinetic energy in GeV """ return self.bunch.getSyncParticle().kinEnergy() def setKinEnergy(self, e_kin = 0.0025): """ Sets the kinetic energy in GeV """ self.bunch.getSyncParticle().kinEnergy(e_kin) def getZtoPhaseCoeff(self,bunch): """ Returns the coefficient to calculate phase in degrees from the z-coordinate. """ bunch_lambda = bunch.getSyncParticle().beta()*self.rf_wave_lenght phase_coeff = 360./bunch_lambda return phase_coeff def getBeamCurrent(self): """ Returns the beam currect in mA """ return self.beam_current def setBeamCurrent(self, current): """ Sets the beam currect in mA """ self.beam_current = current def getBunch(self, nParticles, twissX, twissY, twissZ, cut_off = -1.): """ Returns the pyORBIT bunch with particular number of particles. """ (x0,xp0,y0,yp0,z0,dE0) = self.init_coords comm = orbit_mpi.mpi_comm.MPI_COMM_WORLD rank = orbit_mpi.MPI_Comm_rank(comm) size = orbit_mpi.MPI_Comm_size(comm) data_type = mpi_datatype.MPI_DOUBLE main_rank = 0 bunch = Bunch() self.bunch.copyEmptyBunchTo(bunch) macrosize = (self.beam_current*1.0e-3/self.bunch_frequency) macrosize /= (math.fabs(bunch.charge())*self.si_e_charge) distributor = GaussDist3D(twissX,twissY,twissZ, cut_off) bunch.getSyncParticle().time(0.) for i in range(nParticles): (x,xp,y,yp,z,dE) = distributor.getCoordinates() (x,xp,y,yp,z,dE) = orbit_mpi.MPI_Bcast((x,xp,y,yp,z,dE),data_type,main_rank,comm) if(i%size == rank): bunch.addParticle(x+x0,xp+xp0,y+y0,yp+yp0,z+z0,dE+dE0) nParticlesGlobal = bunch.getSizeGlobal() bunch.macroSize(macrosize/nParticlesGlobal) return bunch def bunchAnalysis(self,bunch): #---- returns (x_rms,y_rms,z_rms_deg) self.twiss_analysis.analyzeBunch(bunch) gamma = bunch.getSyncParticle().gamma() beta = bunch.getSyncParticle().beta() x_rms = math.sqrt(self.twiss_analysis.getTwiss(0)[1]*self.twiss_analysis.getTwiss(0)[3])*1000. y_rms = math.sqrt(self.twiss_analysis.getTwiss(1)[1]*self.twiss_analysis.getTwiss(1)[3])*1000. z_rms = math.sqrt(self.twiss_analysis.getTwiss(2)[1]*self.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 (alphaX,betaX,emittX) = (self.twiss_analysis.getTwiss(0)[0],self.twiss_analysis.getTwiss(0)[1],self.twiss_analysis.getTwiss(0)[3]*1.0e+6) (alphaY,betaY,emittY) = (self.twiss_analysis.getTwiss(1)[0],self.twiss_analysis.getTwiss(1)[1],self.twiss_analysis.getTwiss(1)[3]*1.0e+6) (alphaZ,betaZ,emittZ) = (self.twiss_analysis.getTwiss(2)[0],self.twiss_analysis.getTwiss(2)[1],self.twiss_analysis.getTwiss(2)[3]*1.0e+6) norm_emittX = emittX*gamma*beta norm_emittY = emittY*gamma*beta #gammaZ = (1.0+alphaZ**2)/betaZ #dE_sigma = math.sqrt(gammaZ*emittZ) #print "debug dE_sigma [MeV]=",dE_sigma #---- phi_de_emittZ will be in [pi*deg*MeV] phi_de_emittZ = z_to_phase_coeff*emittZ return (x_rms,y_rms,z_rms_deg) def shiftPhase(self,bunch,delta_phi_deg): dz = - delta_phi_deg/self.getZtoPhaseCoeff(bunch) nParts = bunch.getSize() for ind in range(nParts): z = bunch.z(ind) bunch.z(ind,z+dz)
#----------------------------Add Space Charge nodes----------------------------------------------------
sc_path_length_min = 0.000001
sizeX = 32 #number of grid points in horizontal direction
sizeY = 32 #number of grid points in vertical direction
sizeZ = 16 #number of longitudinal slices
calc2p5d = SpaceChargeForceCalc2p5D(sizeX, sizeY, sizeZ)
scLatticeModifications.setSC2p5DAccNodes(lattice, sc_path_length_min, calc2p5d)
print "SC nodes appied to the lattice"
#--------------------------------Tracking-----------------------------------------
bunchtwissanalysis = BunchTwissAnalysis()
n_turns = 250
strIn = "_{}_{}_{}_{:.0e}_{}".format("gauss", n_particles, emitx,
total_macroSize, n_turns)
w = open("tune_scan_shift4{}.dat".format(strIn), "w")
#w = open("tune_scan_corrected{}.dat".format(strIn),"w")
w.write("qx,qy,emitx,emity\n")
for wp in tunes[:]:
setQuads(wp[0], wp[1])
matrix_lattice = TEAPOT_MATRIX_Lattice(lattice, b)
TwissDataX, TwissDataY = matrix_lattice.getRingTwissDataX(
), matrix_lattice.getRingTwissDataY()
print("{} {}".format(TwissDataX[0][-1][1], TwissDataY[0][-1][1]))
# num of particles = 1
dist = GaussDist2D(twissX, twissY)
#x,px,y,py = dist.getCoordinates()
#x,px,y,py = -0.00114903585122, 0.00124391282321, 0.00310204490361, 0.000305760909693
x, px, y, py = 0., 0., 0., 0.
b.addParticle(x, px, y, py, 0, 0)
#----------------------------------------------------------------------------------------------------
nodes = teapot_latt.getNodes()
position = [
teapot_latt.getNodePositionsDict()[node][0] for node in nodes
if node.getType() == 'monitor teapot'
] # s coordinates of BPMs
bunchtwissanalysis = BunchTwissAnalysis()
#----------------------------------------------------------------------------------------------------
def getAllBPMsInfo(beam, lattice_nodes):
lst = list()
for node in lattice_nodes:
node.trackBunch(beam)
if node.getType() == 'monitor teapot':
x = b.x(0)
y = b.y(0)
lst.append([x, y, px])
return lst
sc_path_length_min = 0.000001
sizeX = 32 #number of grid points in horizontal direction
sizeY = 32 #number of grid points in vertical direction
sizeZ = 16 #number of longitudinal slices
calc2p5d = SpaceChargeForceCalc2p5D(sizeX, sizeY, sizeZ)
scLatticeModifications.setSC2p5DAccNodes(lattice, sc_path_length_min, calc2p5d)
print "SC nodes appied to the lattice"
#--------------------------------Tracking-----------------------------------------
print("tracking")
x_space0 = [(b.x(i), b.px(i)) for i in range(n_particles)]
bunchtwissanalysis = BunchTwissAnalysis()
bunchtwissanalysis.analyzeBunch(b)
Ex = bunchtwissanalysis.getEmittance(0)
Ey = bunchtwissanalysis.getEmittance(1)
x = bunchtwissanalysis.getCorrelation(0, 0)
y = bunchtwissanalysis.getCorrelation(2, 2)
emitX, emitY, xMean, yMean, turn = [Ex], [Ey], [x], [y], [0]
n_turns = 250
for i in range(n_turns):
lattice.trackBunch(b, trackDict)
# if (i+1)%10==0 or i<9:
bunchtwissanalysis.analyzeBunch(b)
Ex = bunchtwissanalysis.getEmittance(0)
Ey = bunchtwissanalysis.getEmittance(1)
# Add tune analysis child node
#-----------------------------------------------------
parentnode_number = 97
parentnode = Lattice.getNodes()[parentnode_number]
Twiss_at_parentnode_entrance = Lattice.getNodes()[parentnode_number-1].getParamsDict()
tunes = TeapotTuneAnalysisNode("tune_analysis")
tunes.assignTwiss(*[Twiss_at_parentnode_entrance[k] for k in ['betax','alphax','etax','etapx','betay','alphay','etay','etapy']])
tunes.assignClosedOrbit(*[Twiss_at_parentnode_entrance[k] for k in ['orbitx','orbitpx','orbity','orbitpy']])
addTeapotDiagnosticsNodeAsChild(Lattice, parentnode, tunes)
# Define twiss analysis and output dictionary
#-----------------------------------------------------------------------
print '\n\t\tbunchtwissanalysis on MPI process: ', rank
bunchtwissanalysis = BunchTwissAnalysis() #Prepare the analysis class that will look at emittances, etc.
get_dpp = lambda b, bta: np.sqrt(bta.getCorrelation(5,5)) / (b.getSyncParticle().gamma()*b.mass()*b.getSyncParticle().beta()**2)
get_bunch_length = lambda b, bta: 4 * np.sqrt(bta.getCorrelation(4,4)) / (speed_of_light*b.getSyncParticle().beta())
get_eps_z = lambda b, bta: 1e9 * 4 * pi * bta.getEmittance(2) / (speed_of_light*b.getSyncParticle().beta())
output_file = 'output/output.mat'
output = Output_dictionary()
output.addParameter('turn', lambda: turn)
output.addParameter('epsn_x', lambda: bunchtwissanalysis.getEmittanceNormalized(0))
output.addParameter('epsn_y', lambda: bunchtwissanalysis.getEmittanceNormalized(1))
output.addParameter('eps_z', lambda: get_eps_z(bunch, bunchtwissanalysis))
output.addParameter('intensity', lambda: bunchtwissanalysis.getGlobalMacrosize())
output.addParameter('n_mp', lambda: bunchtwissanalysis.getGlobalCount())
output.addParameter('D_x', lambda: bunchtwissanalysis.getDispersion(0))
output.addParameter('D_y', lambda: bunchtwissanalysis.getDispersion(1))
output.addParameter('bunchlength', lambda: get_bunch_length(bunch, bunchtwissanalysis))
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
if sub.getType() == "SC2p5D":
count += 1
dis = dis + sub.getLengthOfSC()
pos_list.append(dis)
print(len(pos_list))
file_writer = open(pre_list + "/SC_KICKS" + test_name + ".txt", 'w')
for i in pos_list:
file_writer.write(str(i) + "\n")
file_writer.close()
b.dumpBunch(pre_list + "/initial_bunch_" + str(test_name) + ".dat")
ana = BunchTwissAnalysis()
f = open(pre_list + "/tracking_" + str(test_name) + ".txt", 'w')
for i in range(1000):
ana.analyzeBunch(b)
if rank == 0:
f.write(
str(i) + " " + str(ana.getEmittanceNormalized(0)) + " " +
str(ana.getEmittanceNormalized(1)) + " " + str(ana.getBeta(0)) +
" " + str(ana.getBeta(1)) + " " + str(lostbunch.getSize()) + "\n")
lattice.trackBunch(b, paramsDict)
def __init__(self):
self.bunchtwissanalysis = BunchTwissAnalysis()
self.xAvg = 0.0
self.yAvg = 0.0
self.xpAvg = 0.0
self.ypAvg = 0.0
#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 "
def action_entrance(paramsDict):
node = paramsDict["node"]
b.partAttrValue("ParticleIdNumber", n, 0, n)
f = open("./track_data/" + str(n) + ".txt", 'w')
f_list.append(f)
n += 1
"""--------------------------------- end ---------------------------------"""
paramsDict = {}
lostbunch = Bunch()
paramsDict["lostbunch"] = lostbunch
paramsDict["bunch"] = b
ana = BunchTwissAnalysis()
for i in range(1000):
if rank == 0:
print("==========================")
print("turn: " + str(i))
print("lost: " + str(lostbunch.getSize()))
print("remain: " + str(b.getSize()))
print("charge: " + str(b.macroSize()))
print("intensity: " + str(b.macroSize() * b.getSize() * size))
print("==========================")
lattice.trackBunch(b, paramsDict)
for i in range(5000):
ana.analyzeBunch(b)
distGen = WaterBagDist3D(twissX, twissY, twissZ)
distGen = KVDist3D(twissX, twissY, twissZ)
n_parts = 10000
for i in range(n_parts):
(x, xp, y, yp, z, dE) = distGen.getCoordinates()
b.addParticle(x, xp, y, yp, z, dE)
b.compress()
syncPart = b.getSyncParticle()
syncPart.kinEnergy(TK)
# copy the initial bunch to another to track through TEAPOT Quad
b1 = Bunch()
b.copyBunchTo(b1)
twiss_analysis = BunchTwissAnalysis()
twiss_analysis.analyzeBunch(b)
print "============before=================="
print "X Twiss =", twiss_analysis.getTwiss(0)
print "Y Twiss =", twiss_analysis.getTwiss(1)
print "Z Twiss =", twiss_analysis.getTwiss(2)
#b.dumpBunch()
G = 30.0 # [T/m]
length = 0.1 # [m]
fieldSource = FieldSource(G)
tracker = RungeKuttaTracker(length)
print "Tracker Entrance plane (a,b,c,d)=", tracker.entrancePlane()
print "Tracker Exit plane (a,b,c,d)=", tracker.exitPlane()
# the spatial eps is useless because we have quad field = 0 on axis for the syncPart
def __init__(self, filename): self.file_out = open(filename,"a") self.bunchtwissanalysis = BunchTwissAnalysis()
if frozen:
tunes.assignTwiss(*[Twiss_at_parentnode_entrance[k] for k in ['betax','alphax','etax','etapx','betay','alphay','etay','etapy']])
tunes.assignClosedOrbit(*[Twiss_at_parentnode_entrance[k] for k in ['orbitx','orbitpx','orbity','orbitpy']])
addTeapotDiagnosticsNodeAsChild(Lattice, parentnode, tunes)
#----------------------------------------------------
# Prepare a bunch object to store particle coordinates
#----------------------------------------------------
bunch_tmp = Bunch()
bunch.copyEmptyBunchTo(bunch_tmp)
bunch_tmp.addPartAttr('ParticlePhaseAttributes')
#----------------------------------------------------
# Define twiss analysis and output dictionary
#----------------------------------------------------
bunchtwissanalysis = BunchTwissAnalysis() #Prepare the analysis class that will look at emittances, etc.
get_dpp = lambda b, bta: np.sqrt(bta.getCorrelation(5,5)) / (b.getSyncParticle().gamma()*b.mass()*b.getSyncParticle().beta()**2)
get_bunch_length = lambda b, bta: 4 * np.sqrt(bta.getCorrelation(4,4)) / (speed_of_light*b.getSyncParticle().beta())
get_eps_z = lambda b, bta: 1e9 * 4 * pi * bta.getEmittance(2) / (speed_of_light*b.getSyncParticle().beta())
output = Output_dictionary()
output.addParameter('turn', lambda: turn)
output.addParameter('intensity', lambda: bunchtwissanalysis.getGlobalMacrosize())
output.addParameter('n_mp', lambda: bunchtwissanalysis.getGlobalCount())
output.addParameter('gamma', lambda: bunch.getSyncParticle().gamma())
output.addParameter('mean_x', lambda: bunchtwissanalysis.getAverage(0))
output.addParameter('mean_xp', lambda: bunchtwissanalysis.getAverage(1))
output.addParameter('mean_y', lambda: bunchtwissanalysis.getAverage(2))
output.addParameter('mean_yp', lambda: bunchtwissanalysis.getAverage(3))
output.addParameter('mean_z', lambda: bunchtwissanalysis.getAverage(4))
output.addParameter('mean_dE', lambda: bunchtwissanalysis.getAverage(5))
def __init__(self, file): self.file_out = file self.bunchtwissanalysis = BunchTwissAnalysis()
