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,