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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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
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()
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()
bunch.getSyncParticle().kinEnergy(energy)
print "Generate Lattice."
lattice = TEAPOT_Lattice("no_sc_lattice")
lattice.readMAD("./lattice/sis18_inj.lat","SIS18_MID")
offset = 0.#70e-3
#============Set BPMs into lattice===============
length = round(lattice.getLength()/12,4)
for i in range(12):
name = TeapotBPMSignalNode("BPM")
addTeapotDiagnosticsNode(lattice, i*length, name)
#============Set BPMs into lattice===============
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]
#============add particles to beam===============
# ini. 2D match distribution
twissX = TwissContainer(alpha = alpax, beta = betax, emittance = emittance_x/2)
twissY = TwissContainer(alpha = alpay, beta = betay, emittance = emittance_y/2)
# L+=node.getLength()
print("check sbends")
L = 0
period = 0
for node in lattice.getNodes():
if node.getType() == "bend teapot":
print("parameters. theta {} length {}".format(node.getParam("theta"),
node.getLength()))
loc = L + node.getLength() / 2
#print("loc = {} period loc = {}".format(loc, loc-LL*period/12))
period += 1
L += node.getLength()
print(LL)
#-----------------------------------------------
matrix_lattice = TEAPOT_MATRIX_Lattice(lattice, b)
for matrixNode in matrix_lattice.getNodes():
mt = matrixNode.getMatrix()
if "mh" in matrixNode.getName():
print("{}\n[{} {}]\n[{} {}]\n".format(matrixNode.getName(),
mt.get(0, 0), mt.get(0, 1),
mt.get(1, 0), mt.get(1, 1)))
mt = matrix_lattice.getOneTurnMatrix()
cos_phi_x = (mt.get(0, 0) + mt.get(1, 1)) / 2.0
cos_phi_y = (mt.get(2, 2) + mt.get(3, 3)) / 2.0
print("COS[MUX] ={}".format(cos_phi_x))
print("COS[MUY] ={}".format(cos_phi_y))
(muX0, arrPosAlphaX0, arrPosBetaX0) = matrix_lattice.getRingTwissDataX()
energy = 2
syncPart = b.getSyncParticle()
syncPart.kinEnergy(energy)
EE = syncPart.kinEnergy()
print(EE)
#---------------------------------------------Bunch init---------------------------------------------
print "Generate Lattice."
#---------------------------------------------Make a Teapot Lattice----------------------------------
lattice = TEAPOT_Lattice("lattice")
lattice.readMADX("fodo_thin.seq", "fodo")
lattice.setUseRealCharge(useCharge=1)
#-----------------------------------------------
matrix_lattice = TEAPOT_MATRIX_Lattice(lattice, b)
(muX0, arrPosAlphaX0, arrPosBetaX0) = matrix_lattice.getRingTwissDataX()
(muY0, arrPosAlphaY0, arrPosBetaY0) = matrix_lattice.getRingTwissDataY()
beamline0 = Optics().readtwiss_teapot(lattice, b)
solve = EnvelopeSolver(beamline0)
twiss0 = solve.match_twiss_matrix(12.5e-6, 12.5e-6, 0.0, 0.0)
#-----------------------------------------------
#---------------------------------------------Make a Teapot Lattice----------------------------------
print "INTRODUCE MISALIGNEMENT IN THE QUADRUPOLES"
#---------------------------------------------ORBIT ERRORS-------------------------------------------
# WE INTRODUCE MISALIGNEMENT IN THE QUADRUPOLES; dx, dy = HOR AND VER DISPLACEMENT OF QUADRUPOLES
setDict = {}
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
print ""
print "Start."
teapot_latt = teapot.TEAPOT_Lattice()
print "Read MAD."
teapot_latt.readMAD("../MAD/LATTICE", "RING")
ring_length = teapot_latt.getLength()
print "Lattice=", teapot_latt.getName(), " ring length [m] =", ring_length
Tkin = 1.0
bunch = Bunch()
bunch.getSyncParticle().kinEnergy(Tkin)
matrix_lattice = TEAPOT_MATRIX_Lattice(teapot_latt, bunch)
print "Lattice=", matrix_lattice.getName(
), " matrix lattice length [m] =", matrix_lattice.getLength()
nodes = matrix_lattice.getNodes()
for node in nodes:
if (isinstance(node, BaseMATRIX)):
print "node =", node.getName(), " dict=", node.getParamsDict(
)["matrix_parent_node"].getParamsDict()
else:
print "node =", node.getName(), " dict=", node.getParamsDict()
one_turn_matrix = matrix_lattice.getOneTurnMatrix()
print "=============one turn pyORBIT matrix for MATRIX lattice===="
printM(one_turn_matrix)
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)
print ("(%1d,%1d)=% 6.5e "%(i,j,m.get(i,j))),
print ""
print "Start."
teapot_latt = teapot.TEAPOT_Lattice()
print "Read MAD."
teapot_latt.readMAD("../MAD/LATTICE","RING")
ring_length = teapot_latt.getLength()
print "Lattice=",teapot_latt.getName()," ring length [m] =",ring_length
Tkin = 1.0
bunch = Bunch()
bunch.getSyncParticle().kinEnergy(Tkin)
matrix_lattice = TEAPOT_MATRIX_Lattice(teapot_latt,bunch)
print "Lattice=",matrix_lattice.getName()," matrix lattice length [m] =",matrix_lattice.getLength()
nodes = matrix_lattice.getNodes()
for node in nodes:
if(isinstance(node, BaseMATRIX)):
print "node =",node.getName()," dict=",node.getParamsDict()["matrix_parent_node"].getParamsDict()
else:
print "node =",node.getName()," dict=",node.getParamsDict()
one_turn_matrix = matrix_lattice.getOneTurnMatrix()
print "=============one turn pyORBIT matrix for MATRIX lattice===="
printM(one_turn_matrix)
print "=== ring parameters ==="
paramsDict["sample"] = "Gaussian" paramsDict["mean"] = 0.0 paramsDict["sigma"] = 1.0 #paramsDict["sample"] = "Uniform" #paramsDict["minimum"] = -1.0 #paramsDict["maximum"] = 1.0 #ESet = AddErrorSet(latt, positioni, positionf, setDict, paramsDict, seed_value=40) #ESet = AddErrorSet(latt, positioni, positionf, setDict, paramsDict) # Random #print "FIELD ERRORS IN THE QUADRUPOLES/MULTIPOLES ARE APPLIED" #------------------------------------------------------------------------------------------------------ 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]) #------------------------------------------------------------------------------------------------------ matrix_lattice = TEAPOT_MATRIX_Lattice(lattice, b) TwissDataX0, TwissDataY0 = matrix_lattice.getRingTwissDataX( ), matrix_lattice.getRingTwissDataY() beta_x0 = np.transpose(TwissDataX0[-1])
