def getWarmTunes(self, x, qx0, qy0, groupsDict): self.setQuads(x, groupsDict) kqd = self.lattice.getNodeForName("s11qd11").getParam("kq") kqf = self.lattice.getNodeForName("s11qd12").getParam("kq") kWarm1 = 1.0139780 * kqd * 1.3 / 1.76 kWarm2 = 1.0384325 * kqf * 1.3 / 1.76 kqd = self.lattice.getNodeForName("s52qd11").setParam("kq", kWarm1) kqf = self.lattice.getNodeForName("s52qd12").setParam("kq", kWarm2) matrix_lattice = TEAPOT_MATRIX_Lattice(self.lattice, self.bunch) (mux, b, a) = matrix_lattice.getRingTwissDataX() (muy, b, a) = matrix_lattice.getRingTwissDataY() qx1, qy1 = mux[-1][1], muy[-1][1] print("{} {}".format(qx1, qy1)) print("{} {}\n".format(qx0, qy0)) dqx = 1 - qx1 / qx0 dqy = 1 - qy1 / qy0 metric = dqx**2 + dqy**2 return metric
def set_bunch(off,lattice,bunch): bunch.deleteAllParticles() x_offset = off[0] xp_offset = off[1] y_offset = off[2] yp_offset = off[3] matrix_lattice = TEAPOT_MATRIX_Lattice(lattice,bunch) (muX, arrPosAlphaX, arrPosBetaX) = matrix_lattice.getRingTwissDataX() (muY, arrPosAlphaY, arrPosBetaY) = matrix_lattice.getRingTwissDataY() (arrDispX,arrDispPrimeX) = matrix_lattice.getRingDispersionDataX() (arrDispY,arrDispPrimeY) = matrix_lattice.getRingDispersionDataY() alpax= arrPosAlphaX[0][1] betax = arrPosBetaX[0][1] alpay= arrPosAlphaY[0][1] betay = arrPosBetaY[0][1] # ini. 2D match distribution twissX = TwissContainer(alpha = alpax, beta = betax, emittance = emittance_x/2) twissY = TwissContainer(alpha = alpay, beta = betay, emittance = emittance_y/2) dist = KVDist2D(twissX,twissY) for i in range(NPIC): (x,xp,y,yp) = dist.getCoordinates() bunch.addParticle(x+x_offset,xp+xp_offset,y+y_offset,yp+yp_offset,0.0,0.0)
def find_elements(self,count_el, find_el): matrix_lattice = TEAPOT_MATRIX_Lattice(self.lattice,self.bunch) (muX, arrPosAlphaX, arrPosBetaX) = matrix_lattice.getRingTwissDataX() (muY, arrPosAlphaY, arrPosBetaY) = matrix_lattice.getRingTwissDataY() OrbitX, OrbitY = orbit(self.lattice,self.bunch).get_orbit() bpm_x = [] bpm_y = [] eps_length = 1e-6 m = 0 pos_old = pos = 0.0 nodes = self.lattice.getNodes() for node in nodes: if node.getType() == count_el: m = m + 1 if node.getType() == find_el: for j in range(len(arrPosBetaX)): if (round(self.lattice.getNodePositionsDict()[node][1],4)==round(arrPosBetaX[j][0],4)): pos_old = pos pos = self.lattice.getNodePositionsDict()[node][1] if(abs(pos_old-pos) > eps_length): bpm_x.append([muX[j][1], arrPosBetaX[j][1], arrPosAlphaX[j][1], pos,m,node.getName(),OrbitX[j][1],muX[-1][1]]) bpm_y.append([muY[j][1], arrPosBetaY[j][1], arrPosAlphaY[j][1], pos,m,node.getName(),OrbitY[j][1],muY[-1][1]]) return bpm_x,bpm_y
def calculate_betas(teapot_lattice, bunch): matrix_lattice = TEAPOT_MATRIX_Lattice(teapot_lattice, bunch) TwissDataX, TwissDataY = matrix_lattice.getRingTwissDataX( ), matrix_lattice.getRingTwissDataY() return np.transpose(TwissDataX[-1]), np.transpose(TwissDataY[-1])
def twiss_throughout(lattice, bunch): """Get Twiss parameters throughout lattice. Parameters ---------- lattice : TEAPOT_Lattice A periodic lattice to track with. bunch : Bunch Test bunch to perform tracking. Returns ------- ndarray Columns are: [s, nux, nuy, alpha_x, alpha_x, beta_x, beta_y] """ # Extract Twiss parameters from one turn transfer matrix matrix_lattice = TEAPOT_MATRIX_Lattice(lattice, bunch) twiss_x = matrix_lattice.getRingTwissDataX() twiss_y = matrix_lattice.getRingTwissDataY() # Unpack and convert to ndarrays (nux, alpha_x, beta_x), (nuy, alpha_y, beta_y) = twiss_x, twiss_y nux, alpha_x, beta_x = np.array(nux), np.array(alpha_x), np.array(beta_x) nuy, alpha_y, beta_y = np.array(nuy), np.array(alpha_y), np.array(beta_y) # Merge into one array s = nux[:, 0] nux, alpha_x, beta_x = nux[:, 1], alpha_x[:, 1], beta_x[:, 1] nuy, alpha_y, beta_y = nuy[:, 1], alpha_y[:, 1], beta_y[:, 1] return np.vstack([s, nux, nuy, alpha_x, alpha_y, beta_x, beta_y]).T
def getTunesDeviation(self, x, qx0, qy0, groupsDict): kw1 = self.lattice.getNodeForName("s52qd11").getParam("kq") kw2 = self.lattice.getNodeForName("s52qd12").getParam("kq") print("warm quads are {} {}".format(kw1, kw2)) self.setQuads(x, groupsDict) kw1 = self.lattice.getNodeForName("s52qd11").getParam("kq") kw2 = self.lattice.getNodeForName("s52qd12").getParam("kq") print("warm quads are {} {}".format(kw1, kw2)) matrix_lattice = TEAPOT_MATRIX_Lattice(self.lattice, self.bunch) (mux, b, a) = matrix_lattice.getRingTwissDataX() (muy, b, a) = matrix_lattice.getRingTwissDataY() qx1, qy1 = mux[-1][1], muy[-1][1] #print("{} {}".format(qx1,qy1)) #print("{} {}\n".format(qx0,qy0)) dqx = 1 - qx1 / qx0 dqy = 1 - qy1 / qy0 metric = dqx**2 + dqy**2 return metric
def get_orbit(self): z0 = self.offset() matrix_lattice = TEAPOT_MATRIX_Lattice(self.lattice,self.bunch) OrbitX, OrbitY = matrix_lattice.getRingOrbit(z0) return OrbitX, OrbitY
def readtwiss_teapot(self, lattice, bunch): beamline = Optics() matrix_lattice = TEAPOT_MATRIX_Lattice(lattice, bunch) (arrmuX, arrPosAlphaX, arrPosBetaX) = matrix_lattice.getRingTwissDataX() (arrmuY, arrPosAlphaY, arrPosBetaY) = matrix_lattice.getRingTwissDataY() (DispersionX, DispersionXP) = matrix_lattice.getRingDispersionDataX() (DispersionY, DispersionYP) = matrix_lattice.getRingDispersionDataY() nodes = lattice.getNodes() for node in nodes: for j in range(len(arrPosBetaX)): if (round(lattice.getNodePositionsDict()[node][1], 4) == round(arrPosBetaX[j][0], 4)): muX = arrmuX[j][1] betaX = arrPosBetaX[j][1] alphaX = arrPosAlphaX[j][1] dx = DispersionX[j][1] dmux = DispersionXP[j][1] muY = arrmuY[j][1] betaY = arrPosBetaY[j][1] alphaY = arrPosAlphaY[j][1] dmuy = DispersionYP[j][1] if node.getType() == "quad teapot": k1l = node.getParam("kq") * node.getLength() else: k1l = 0.0 if node.getType() == "bend teapot": angle = node.getParam("theta") else: angle = 0.0 beamline.add(1) j = len(beamline) - 1 beamline[j] = Twiss() beamline[j].data['keyword'] = node.getName() beamline[j].data['marker'] = node.getType() beamline[j].data['s'] = round( lattice.getNodePositionsDict()[node][1], 4) beamline[j].data['L'] = node.getLength() beamline[j].data['alfx'] = alphaX beamline[j].data['alfy'] = alphaY beamline[j].data['betx'] = betaX beamline[j].data['bety'] = betaY beamline[j].data['Dx'] = dx beamline[j].data['Dpx'] = dmux beamline[j].data['mux'] = muX beamline[j].data['muy'] = muY beamline[j].data['angle'] = angle beamline[j].data['k1'] = k1l return beamline
def lattice_function(self, lattice, bunch, variables, constraints): self.variables = variables self.constraints = constraints nv = np.array(self.variables).shape[0] nc = np.array(self.constraints).shape[0] if not nv == nc: print "Number of variables and constraints must be equal" print "Stop." sys.exit(0) matrix_lattice = TEAPOT_MATRIX_Lattice(lattice, bunch) (muX, arrPosAlphaX, arrPosBetaX) = matrix_lattice.getRingTwissDataX() self.tuneX = muX[-1][1] nodes = lattice.getNodes() mux_kick = [] betax_kick = [] alphax_kick = [] for n in range(nv): for node in nodes: if node.getName() == self.variables[n][0][0]: for j in range(len(arrPosBetaX)): if (round(lattice.getNodePositionsDict()[node][1], 4) == round(arrPosBetaX[j][0], 4)): mux_kick.append(muX[j][1]) betax_kick.append(arrPosBetaX[j][1]) alphax_kick.append(arrPosAlphaX[j][1]) self.variables[n].insert(3, [ muX[j][1], arrPosBetaX[j][1], arrPosAlphaX[j][1] ]) mux_bpm = [] betax_bpm = [] alphax_bpm = [] for n in range(nc): for node in nodes: if node.getName() == self.constraints[n][0][0]: for j in range(len(arrPosBetaX)): if (round(lattice.getNodePositionsDict()[node][1], 4) == round(arrPosBetaX[j][0], 4)): mux_bpm.append(muX[j][1]) betax_bpm.append(arrPosBetaX[j][1]) alphax_bpm.append(arrPosAlphaX[j][1]) self.constraints[n].insert(3, [ muX[j][1], arrPosBetaX[j][1], arrPosAlphaX[j][1] ])
def offset(self): matrix_lattice = TEAPOT_MATRIX_Lattice(self.lattice,self.bunch) OTM = np.zeros((6, 6)) kickOTM = np.zeros((6)) mt = matrix_lattice.getOneTurnMatrix() for i in range(6): for j in range(6): OTM[i][j] = mt.get(i,j) for i in range(6): kickOTM[i] = mt.get(i,6) #z0 fulfill: z0 = Mz0 with M as one turn matrix # Mz + k = z => z = (I-M)^-1 k z0, _, _, _ = np.linalg.lstsq(np.eye(6, dtype=int) - OTM,kickOTM) return z0
def hyperSurface(self, theta, qx0, qy0, alpha): dk = self.normalize(theta, alpha, self.quadDict) print("hyper-surface computation") self.setQuads(dk, self.quadDict) matrix_lattice = TEAPOT_MATRIX_Lattice(self.lattice, self.bunch) (mux, bTmp, aTmp) = matrix_lattice.getRingTwissDataX() (muy, bTmp, aTmp) = matrix_lattice.getRingTwissDataY() qx, qy = mux[-1][1], muy[-1][1] print("current tunes are {} {}".format(qx, qy)) out = np.sqrt((qx - qx0)**2 + (qy - qy0)**2) print("the distance from set tunes hor:{} ver:{}".format( qx - qx0, qy - qy0)) print("out is {}\n".format(out)) return -out
def periodicFun(self, theta, qx0, qy0): self.setQuads(theta) matrix_lattice = TEAPOT_MATRIX_Lattice(self.lattice, self.bunch) TwissDataX, TwissDataY = matrix_lattice.getRingTwissDataX( ), matrix_lattice.getRingTwissDataY() betaX, betaY = self.getBetaBpm(TwissDataX, TwissDataY) metric_x = 1.0 - np.min(betaX) / np.max(betaX) metric_y = 1.0 - np.min(betaY) / np.max(betaY) arr = [metric_x, metric_y] if f_max - f_min > f_max: metric = 4 * f_max else: metric = np.sum(arr) return metric
def calculateBeat(lattice, bunch, sVar, betaX0, betaY0): matrix_lattice = TEAPOT_MATRIX_Lattice(lattice, bunch) TwissDataX, TwissDataY = matrix_lattice.getRingTwissDataX( ), matrix_lattice.getRingTwissDataY() betaX, betaY = np.transpose(TwissDataX[-1]), np.transpose( TwissDataY[-1]) qx, qy = np.transpose(TwissDataX[0])[-1][-1], np.transpose( TwissDataY[0])[-1][-1] print("current tunes are {} {}".format(qx, qy)) betX = np.interp(x=sVar, xp=betaX[0], fp=betaX[1]) betY = np.interp(x=sVar, xp=betaY[0], fp=betaY[1]) beatX = (betX - betaX0) beatY = (betY - betaY0) n = len(beatX) fftBeatX = np.abs(np.fft.rfft(beatX)) / n fftBeatY = np.abs(np.fft.rfft(beatY)) / n return fftBeatX, fftBeatY
def twiss_at_entrance(lattice, mass, energy): """Get 2D Twiss parameters at lattice entrance. Parameters ---------- lattice : TEAPOT_Lattice A periodic lattice to track with. mass, energy : float Particle mass [GeV/c^2] and kinetic energy [GeV]. Returns ------- alpha_x, alpha_y, beta_x, beta_y : float 2D Twiss parameters at lattice entrance. """ bunch, params_dict = initialize_bunch(mass, energy) matrix_lattice = TEAPOT_MATRIX_Lattice(lattice, bunch) _, arrPosAlphaX, arrPosBetaX = matrix_lattice.getRingTwissDataX() _, arrPosAlphaY, arrPosBetaY = matrix_lattice.getRingTwissDataY() alpha_x, alpha_y = arrPosAlphaX[0][1], arrPosAlphaY[0][1] beta_x, beta_y = arrPosBetaX[0][1], arrPosBetaY[0][1] return alpha_x, alpha_y, beta_x, beta_y
def transfer_matrix(lattice, mass, energy): """Shortcut to get transfer matrix from periodic lattice. Parameters ---------- lattice : TEAPOT_Lattice A periodic lattice to track with. mass, energy : float Particle mass [GeV/c^2] and kinetic energy [GeV]. Returns ------- M : ndarray, shape (4, 4) Transverse transfer matrix. """ bunch, params_dict = initialize_bunch(mass, energy) matrix_lattice = TEAPOT_MATRIX_Lattice(lattice, bunch) one_turn_matrix = matrix_lattice.oneTurnMatrix M = np.zeros((4, 4)) for i in range(4): for j in range(4): M[i, j] = one_turn_matrix.get(i, j) return M
def getPositions(self): matrix_lattice = TEAPOT_MATRIX_Lattice(self.lattice, self.bunch) TwissDataX = matrix_lattice.getRingTwissDataX() sArr = [x[0] for x in TwissDataX[0]] return sArr
def readtwiss_teapot(self,lattice, bunch): beamline=Optics() matrix_lattice = TEAPOT_MATRIX_Lattice(lattice,bunch) (arrmuX, arrPosAlphaX, arrPosBetaX) = matrix_lattice.getRingTwissDataX() (arrmuY, arrPosAlphaY, arrPosBetaY) = matrix_lattice.getRingTwissDataY() (DispersionX, DispersionXP) = matrix_lattice.getRingDispersionDataX() (DispersionY, DispersionYP) = matrix_lattice.getRingDispersionDataY() nodes = lattice.getNodes() matrixNodes = matrix_lattice.getNodes() idx = 0 Md = np.identity(6) # Matrix of downstream element for node in nodes: for j in range(len(arrPosBetaX)): if (round(lattice.getNodePositionsDict()[node][1],4)==round(arrPosBetaX[j][0],4)): muX = arrmuX[j][1] betaX = arrPosBetaX[j][1] alphaX = arrPosAlphaX[j][1] dx = DispersionX[j][1] dmux = DispersionXP[j][1] muY = arrmuY[j][1] betaY = arrPosBetaY[j][1] alphaY = arrPosAlphaY[j][1] dmuy = DispersionYP[j][1] if node.getType() == "quad teapot": k1l = node.getParam("kq")*node.getLength() else: if node.getType()=="multipole teapot": k1l = node.getParam("kls")[1] else: k1l = 0.0 if node.getType() == "bend teapot": angle = node.getParam("theta") else: angle = 0.0 beamline.add(1) j=len(beamline)-1 beamline[j]=Twiss() name = node.getName() beamline[j].data['keyword']=name beamline[j].data['marker']=node.getType() beamline[j].data['s']=round(lattice.getNodePositionsDict()[node][1],4) beamline[j].data['L']=node.getLength() beamline[j].data['alfx']=alphaX beamline[j].data['alfy']=alphaY beamline[j].data['betx']=betaX beamline[j].data['bety']=betaY beamline[j].data['Dx']=dx beamline[j].data['Dpx']=dmux beamline[j].data['mux']=muX beamline[j].data['muy']=muY beamline[j].data['angle']=angle beamline[j].data['k1']=k1l matName = matrixNodes[idx].getName() M = np.identity(6) while name in matName and idx < len(matrixNodes)-1: matNode = matrixNodes[idx] matName = matNode.getName() mt = matNode.getMatrix() idx+=1 for index in range(36): Md[index//6,index%6] = mt.get(index//6,index%6) M = np.dot(M[:],Md) beamline[j].data['map'] = M return beamline
def hyperBetaBeat(self, rhobeg=5e-5, A=0.1, qx0=0.1234, qy0=0.1234, k=tuple(), patternDict=dict(), quadsDict=dict(), betaX0=1.0, betaY0=1.0): #if not self.quadDict: self.findElements(patternDict, quadsDict) nQuadGroups = len(self.quadDict) self.lattice.getNodeForName("s52qd11").setParam("kq", k[0]) self.lattice.getNodeForName("s52qd12").setParam("kq", k[1]) kw1 = self.lattice.getNodeForName("s52qd11").getParam("kq") kw2 = self.lattice.getNodeForName("s52qd12").getParam("kq") print("warm quads are {} {}".format(kw1, kw2)) k1 = self.lattice.getNodeForName("s11qd11").getParam("kq") k2 = self.lattice.getNodeForName("s11qd12").getParam("kq") print("cold quads are {} {}".format(k1, k2)) theta = np.zeros(nQuadGroups) optionsDict = {'rhobeg': rhobeg, 'maxiter': 2000, 'disp': True} arg = (qx0, qy0) vec = som(self.getTunesDeviation, theta, method="COBYLA", options=optionsDict, args=arg) self.setQuads(vec.x) kd = self.lattice.getNodeForName("s11qd11").getParam("kq") kf = self.lattice.getNodeForName("s11qd12").getParam("kq") print("sollution found\n") matrix_lattice = TEAPOT_MATRIX_Lattice(self.lattice, self.bunch) TwissDataX, TwissDataY = matrix_lattice.getRingTwissDataX( ), matrix_lattice.getRingTwissDataY() betaX, betaY = np.transpose(TwissDataX[-1]), np.transpose( TwissDataY[-1]) L = self.lattice.getLength() sVar = np.linspace(0, L, len(betaX0)) betxSmpl = np.interp(x=sVar, xp=betaX[0], fp=betaX[1]) betySmpl = np.interp(x=sVar, xp=betaY[0], fp=betaY[1]) beat_x = (betxSmpl - betaX0) beat_y = (betySmpl - betaY0) n = len(beat_x) beta_x_fft = np.abs(np.fft.rfft(beat_x)) / n beta_y_fft = np.abs(np.fft.rfft(beat_y)) / n periodicHarmX = [x for i, x in enumerate(beta_x_fft) if not i % 6] periodicHarmY = [x for i, x in enumerate(beta_y_fft) if not i % 6] AperiodicHarmX = [x for i, x in enumerate(beta_x_fft) if i % 6] AperiodicHarmY = [x for i, x in enumerate(beta_y_fft) if i % 6] maxPosX = np.argmax(beta_x_fft) maxPosY = np.argmax(beta_y_fft) return kd, kf, maxPosX, maxPosY, beta_x_fft[maxPosX], beta_y_fft[ maxPosY], np.sum(periodicHarmX) / np.sum(AperiodicHarmX), np.sum( periodicHarmY) / np.sum(AperiodicHarmY)
paramsDict["dy"] = 0.006 setDict["elementtype"] = "quad" setDict["ringline"] = "ring" ESet = AddErrorSet(teapot_latt, positioni, positionf, setDict, paramsDict, seed_value=50) #ESet = AddErrorSet(lattice, positioni, positionf, setDict, paramsDict) # Random #---------------------------------------------ORBIT ERRORS------------------------------------------- # get lattice function from transfer matrices matrix_lattice = TEAPOT_MATRIX_Lattice(teapot_latt, b) (muX, arrPosAlphaX, arrPosBetaX) = matrix_lattice.getRingTwissDataX() (muY, arrPosAlphaY, arrPosBetaY) = matrix_lattice.getRingTwissDataY() #----------------------------------Bunch-Distribusion------------------------------------------------ # machted beam emittance_x, emittance_y = 10**(-8), 10**(-8 ) # should be checked and corrected twissX = TwissContainer(arrPosAlphaX[0][1], arrPosBetaX[0][1], emittance_x) twissY = TwissContainer(arrPosAlphaY[0][1], arrPosBetaY[0][1], emittance_y) dE = 0 * 10**(-13) n = 2048 # num of particles dist = GaussDist2D(twissX, twissY) for i in range(n): x, px, y, py = dist.getCoordinates()
lattice = TEAPOT_Lattice("lattice") #lattice.readMADX("fodo_thin.seq","fodo") lattice.readMADX("FODO_x6.seq", "fodo") #lattice.readMADX("cryring.madx","cryring") #lattice.readMADX("sis100_full_thin_fix.seq","sis100ring") lattice.setUseRealCharge(useCharge=1) #---------------------------------------------SPLIT LONG ELEMENTS-------------------------------------- for node in lattice.getNodes(): if node.getLength() > 1.0: node.setnParts(int(node.getLength() // 1 + 1)) #------------------------------------------------------------------------------------------------------ matrix_lattice = TEAPOT_MATRIX_Lattice(lattice, b) TwissDataX0, TwissDataY0 = matrix_lattice.getRingTwissDataX( ), matrix_lattice.getRingTwissDataY() beta_x0 = np.transpose(TwissDataX0[-1]) beta_y0 = np.transpose(TwissDataY0[-1]) #---------------------------------------------------- ERRORS------------------------------------------- # WE INTRODUCE MISALIGNEMENT IN THE QUADRUPOLES; dx, dy = HOR AND VER DISPLACEMENT OF QUADRUPOLES setDict = {} paramsDict = {} positioni = 0.0 positionf = lattice.getLength()
for node in latt.getNodes(): if node.getType()=="quad teapot": print(node.getParam("kq")) #------------------------------------------------------------------------------------------------------ matrix_lattice = TEAPOT_MATRIX_Lattice(lattice,b) TwissDataX0,TwissDataY0 = matrix_lattice.getRingTwissDataX(),matrix_lattice.getRingTwissDataY() beta_x0 = np.transpose(TwissDataX0[-1]) beta_y0 = np.transpose(TwissDataY0[-1]) ''' #------------------------------------------------------------------------------------------------------ matrix_latt = TEAPOT_MATRIX_Lattice(latt,b) TwissDataX,TwissDataY = matrix_latt.getRingTwissDataX(),matrix_latt.getRingTwissDataY() beta_x = np.transpose(TwissDataX[-1]) beta_y = np.transpose(TwissDataY[-1]) #------------------------------------------------------------------------------------------------------ plt.figure() #plt.plot(beta_x0[0],beta_x0[1]) #plt.plot(beta_y0[0],beta_y0[1], ls="--") plt.plot(beta_x[0],beta_x[1]) plt.plot(beta_y[0],beta_y[1], ls="--")