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
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
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
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 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)
# ~ for i in range(1, p['n_macroparticles']): # ~ particle_output.AddNewParticle(i) # Update for turn -1 (pre tracking) particle_output.update(bunch, -1) #---------------------------------------------------- # Do some turns and dump particle information #---------------------------------------------------- print '\nnow start tracking...' for turn in range(p['turns_max']): Lattice.trackBunch(bunch, paramsDict) LinearRestoringForce(bunch, s['RestoringForce']) bunchtwissanalysis.analyzeBunch(bunch) # analyze twiss and emittance if turn in p['turns_print']: saveBunchAsMatfile(bunch, "output/mainbunch_%s"%(str(turn).zfill(6))) saveBunchAsMatfile(lostbunch, "lost/lostbunch_%s"%(str(turn).zfill(6))) output.save_to_matfile('output') output.update() particle_output.update(bunch, turn) # For last turn output particle dictionary if turn == (p['turns_max']-1): for i in range(0, p['n_macroparticles']): particle_output.print_particle(i) particle_output.print_all_particles()
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) x = bunchtwissanalysis.getCorrelation(0, 0) y = bunchtwissanalysis.getCorrelation(2, 2)
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 #tracker.spatialEps(0.00000000001) # we have to specify the number of time steps
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) if rank == 0: print("==========================") print("turn: " + str(i)) print("lost: " + str(lostbunch.getSize())) print("remain: " + str(b.getSize())) print("charge: " + str(b.macroSize()))
eKin = bunch.getSyncParticle().kinEnergy()*1.0e+3 s = " %5d %35s %4.5f %5.3f %5.3f %5.3f %5.3f %5.3f %5.3f %7.5f %10.6f %8d "%(paramsDict["count"],node.getName(),pos,x_rms,y_rms,z_rms,z_rms_deg,xp_rms,yp_rms,dE_rms,eKin,bunch.getSize()) file_out.write(s +"\n") print s #actionContainer.addAction(action_entrance, AccActionsContainer.ENTRANCE) actionContainer.addAction(action_exit, AccActionsContainer.EXIT) time_start = time.clock() accLattice.trackBunch(bunch_in, paramsDict = paramsDict, actionContainer = actionContainer) time_exec = time.clock() - time_start print "time[sec]=",time_exec file_out.close() eKin = bunch_in.getSyncParticle().kinEnergy()*1.0e+3 twiss_analysis.analyzeBunch(bunch_in) (alphaX,betaX,emittX) = (twiss_analysis.getTwiss(0)[0],twiss_analysis.getTwiss(0)[1],twiss_analysis.getTwiss(0)[3]) (alphaY,betaY,emittY) = (twiss_analysis.getTwiss(1)[0],twiss_analysis.getTwiss(1)[1],twiss_analysis.getTwiss(1)[3]) (alphaZ,betaZ,emittZ) = (twiss_analysis.getTwiss(2)[0],twiss_analysis.getTwiss(2)[1],twiss_analysis.getTwiss(2)[3]) print " ========= CCL4 exit PyORBIT Twiss =========== eKin=",eKin print " aplha beta emitt[mm*mrad] X= ( %6.4f , %6.4f , %6.4f ) "%(alphaX,betaX,emittX*1.0e+6) print " aplha beta emitt[mm*mrad] Y= ( %6.4f , %6.4f , %6.4f ) "%(alphaY,betaY,emittY*1.0e+6) print " aplha beta emitt[mm*MeV] Z= ( %6.4f , %6.4f , %6.4f ) "%(alphaZ,betaZ,emittZ*1.0e+6)
actionContainer.addAction(action_exit, AccActionsContainer.EXIT) time_start = time.clock() accLattice.trackBunch(bunch_in, paramsDict=paramsDict, actionContainer=actionContainer) time_exec = time.clock() - time_start print "time[sec]=", time_exec fl_loss_pos_coord_out.close() file_out.close() eKin = bunch_in.getSyncParticle().kinEnergy() * 1.0e+3 twiss_analysis.analyzeBunch(bunch_in) (alphaX, betaX, emittX) = (twiss_analysis.getTwiss(0)[0], twiss_analysis.getTwiss(0)[1], twiss_analysis.getTwiss(0)[3]) (alphaY, betaY, emittY) = (twiss_analysis.getTwiss(1)[0], twiss_analysis.getTwiss(1)[1], twiss_analysis.getTwiss(1)[3]) (alphaZ, betaZ, emittZ) = (twiss_analysis.getTwiss(2)[0], twiss_analysis.getTwiss(2)[1], twiss_analysis.getTwiss(2)[3]) print " ========= CCL4 exit PyORBIT Twiss =========== eKin=", eKin print " aplha beta emitt[mm*mrad] X= ( %6.4f , %6.4f , %6.4f ) " % ( alphaX, betaX, emittX * 1.0e+6) print " aplha beta emitt[mm*mrad] Y= ( %6.4f , %6.4f , %6.4f ) " % ( alphaY, betaY, emittY * 1.0e+6)
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)
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) start_ind = accLattice.getNodeIndex(quads[0]) stop_ind = accLattice.getNodeIndex(quads[len(quads)-1])
(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 #tracker.spatialEps(0.00000000001) # we have to specify the number of time steps
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
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) z_rms_deg = z_to_phase_coeff * z_rms / 1000.0 print "========== Bunch RMS sizes at BNCH02 entrance:" print "(Sx[mm],Sy[mm],Sz[deg]) = %5.3f %5.3f %5.3f " % (x_rms, y_rms,