def __init__(self, name = "no name"):
"""
Constructor. Creates the base MATRIX element.
"""
AccNodeBunchTracker.__init__(self,name)
self.setType("base matrix")
self.matrix = Matrix(7,7)
def getOneTurnMatrix(self,teapot_lattice): self.matrixArr = self.getTransportMatrixArray(teapot_lattice) final_matrix = Matrix(6,6) final_matrix.unit() for (s0,s1,matrix) in self.matrixArr: final_matrix = matrix.mult(final_matrix) return final_matrix
def getOneTurnMatrix(self, teapot_lattice):
self.matrixArr = self.getTransportMatrixArray(teapot_lattice)
final_matrix = Matrix(6, 6)
final_matrix.unit()
for (s0, s1, matrix) in self.matrixArr:
final_matrix = matrix.mult(final_matrix)
return final_matrix
def setM(arr):
n = len(arr)
m = Matrix(n, n)
for i in xrange(n):
for j in xrange(n):
m.set(i, j, arr[i][j])
return m
def __init__(self, name=None):
AccLattice.__init__(self, name)
self.oneTurnMatrix = Matrix(7, 7)
self.oneTurnMatrix.unit()
self.Matrix = Matrix(7, 7)
self.Matrix.unit()
def __init__(self, trMatricesController, name="TrMatrixGen"):
if (name == "TrMatrixGen"):
name += name + ":" + str(trMatricesController.getCount())
MarkerLinacNode.__init__(self, name)
self.trMatricesController = trMatricesController
self.trMtrxNode_ind = trMatricesController.getCount()
self.use_twiss_weight_x = 0
self.use_twiss_weight_y = 0
self.use_twiss_weight_z = 0
self.relativistic_beta = 0.
self.relativistic_gamma = 0.
#--------------------------------------
self.trMtrx = Matrix(7, 7)
#--------------------------------------
self.trMatricesController.addNode(self)
class BaseMATRIX(AccNodeBunchTracker):
""" The base abstract class of the BaseMATRIX accelerator lattice elements hierarchy. """
def __init__(self, name = "no name"):
"""
Constructor. Creates the base MATRIX element.
"""
AccNodeBunchTracker.__init__(self,name)
self.setType("base matrix")
self.matrix = Matrix(7,7)
def getMatrix(self):
"""
Returns the (7,7) Matrix for this transport AccNode.
"""
return self.matrix
def track(self, paramsDict):
"""
It is tracking the parameter dictionary (with bunch as a "bunch") through the BaseMATRIX element.
"""
bunch = paramsDict["bunch"]
self.matrix.track(bunch)
def _twissAction(self, paramsDict):
node = paramsDict["node"]
bunch = paramsDict["bunch"]
s_start = paramsDict["position"]
s_stop = s_start + node.getLength(node.getActivePartIndex())
paramsDict["position"] = s_stop
if (isinstance(node, BaseTEAPOT) == True
and isinstance(node, RingRFTEAPOT) == False):
#print "i=",self.index," ind=",node.getActivePartIndex()," name=",node.getName(),
#print " type=",node.getType()," L=",node.getLength(node.getActivePartIndex())
self.matrixGenerator.initBunch(bunch)
node.track(paramsDict)
transp_matrix = Matrix(6, 6)
self.matrixGenerator.calculateMatrix(bunch, transp_matrix)
self.matrixArr.append((s_start, s_stop, transp_matrix))
m = transp_matrix
#print "detX=", (m.get(0,0)*m.get(1,1) - m.get(0,1)*m.get(1,0))
#printM(transp_matrix)
self.index = self.index + 1
def __init__(self, name = None): AccLattice.__init__(self,name) self.oneTurmMatrix = Matrix(7,7) self.oneTurmMatrix.unit()
print ""
print "Start."
n = Nll()
k = .01
b = Bunch()
b.addParticle(1., 2., 3., 4., 5., 6.)
b.dumpBunch()
m = Matrix(6, 6)
m.unit()
m.set(0, 0, math.cos(k))
m.set(0, 1, math.sin(k))
m.set(1, 0, -math.sin(k))
m.set(1, 1, math.cos(k))
m.set(2, 2, math.cos(k))
m.set(2, 3, math.sin(k))
m.set(3, 2, -math.sin(k))
m.set(3, 3, math.cos(k))
printM(m)
for i in range(20000):
m.track(b)
n.TRACK_EXT(b)
qh04_pos_end = accLattice.getNodePositionsDict()[accLattice.getNodeForName(
"QH04")][1]
qv01_pos_start = accLattice.getNodePositionsDict()[accLattice.getNodeForName(
"QV01")][0]
qh02_pos_start = accLattice.getNodePositionsDict()[accLattice.getNodeForName(
"QH02")][0]
qv03_pos_start = accLattice.getNodePositionsDict()[accLattice.getNodeForName(
"QV03")][0]
qh04_pos_start = accLattice.getNodePositionsDict()[accLattice.getNodeForName(
"QH04")][0]
qv01_pos_center = (qv01_pos_start + qv01_pos_end) / 2.0
#---- qv01_3d coordinate transformation from the lattice system
transfCoordsMatrix = Matrix(4, 4)
transfCoordsMatrix.unit()
dist_from_qv01_center = 0.
transfCoordsMatrix.set(2, 3, -dist_from_qv01_center)
qv01_3d.transormfMatrix(transfCoordsMatrix)
#---- qh02_3d coordinate transformation from the lattice system
transfCoordsMatrix = Matrix(4, 4)
transfCoordsMatrix.unit()
dist_from_qv01_center = qh02_pos_end - qv01_pos_end
transfCoordsMatrix.set(2, 3, -dist_from_qv01_center)
qh02_3d.transormfMatrix(transfCoordsMatrix)
#---- qv03_3d coordinate transformation from the lattice system
transfCoordsMatrix = Matrix(4, 4)
transfCoordsMatrix.unit()
v1 = PhaseVector(5) setV(v) setV(v1) printV(v) printV(v1) v3 = v.add(1) printV(v3) v3 = v.add(v1) printV(v3) print "v*v1 =", v.mult(v1) print "-----------------------------------------------" #check matrix print "----mult ----" m = Matrix(5,4) m1 = Matrix(4,5) setM(m) setM(m1) m3 = m.mult(m1) printM(m3) print "----add -----" m = Matrix(4,4) m1 = Matrix(4,4) setM(m) setM(m1) m3 = m.add(m1) printM(m3) m4 = m3.invert()
def _makePolynomial(self):
nPoints = len(self.x_y_err_arr)
if (nPoints < (self.order + 1)):
self.order = nPoints - 1
#check if just one of errors is zero
infoZeroErr = 1.0
for [x, y, err] in self.x_y_err_arr:
infoZeroErr *= err
for i in range(nPoints):
[x, y, err] = self.x_y_err_arr[i]
sigma = 1.0
if (infoZeroErr != 0.):
sigma = 1.0 / (err * err)
self.x_y_err_arr[i][2] = sigma
#now make A matrix
aMatr = Matrix(nPoints, self.order + 1)
for i in range(nPoints):
for j in range(self.order + 1):
x = self.x_y_err_arr[i][0]
aMatr.set(i, j, math.pow(x, j))
aTCa = Matrix(self.order + 1, self.order + 1)
for i in range(self.order + 1):
for j in range(self.order + 1):
a = 0.
for k in range(nPoints):
sigma = self.x_y_err_arr[k][2]
a += aMatr.get(k, i) * sigma * aMatr.get(k, j)
aTCa.set(i, j, a)
#now the resuting coefficients and errors
aTCaI = aTCa.invert()
e = aTCaI.mult(aTCa)
if (aTCa == None):
print "python PolynomialFit: Problem with data."
for i in range(nPoints):
x = self.x_y_err_arr[i][0]
y = self.x_y_err_arr[i][1]
err = self.x_y_err_arr[i][2]
print " x,y,err = %12.5g %12.5g %12.5g " % (x, y, err)
print "Stop."
sys.exit(1)
coef_arr = [0.] * (self.order + 1)
err_arr = [0.] * (self.order + 1)
for i in range(self.order + 1):
err_arr[i] = math.sqrt(math.fabs(aTCaI.get(i, i)))
for i in range(self.order + 1):
coef_arr[i] = 0.
for j in range(self.order + 1):
for k in range(nPoints):
sigma = self.x_y_err_arr[k][2]
y = self.x_y_err_arr[k][1]
coef_arr[i] += aTCaI.get(i, j) * aMatr.get(k,
j) * sigma * y
# polinimial coefficients are found
self.polynomial.order(self.order)
for i in range(len(coef_arr)):
self.polynomial.coefficient(i, coef_arr[i])
# now let's calculate errors
if (infoZeroErr == 0.):
total_sigma = 0.
for k in range(nPoints):
x = self.x_y_err_arr[k][0]
y = self.x_y_err_arr[k][1]
total_sigma += (self.polynomial.value(x) - y)**2
total_sigma = math.sqrt(total_sigma / (nPoints - 2))
for i in range(len(err_arr)):
err_arr[i] *= total_sigma
# set the resulting coefficients and errors array
self.coef_err_arr = [coef_arr, err_arr]
def trackDispersionData(self,momentum,mass, disp, disp_p, direction = "x"):
"""
Returns the tuple ([(position, disp),...],[(position,disp_p),...] ).
The tracking starts from the values specified as the initial parameters.
The possible values for direction parameter "x" or "y".
"""
if(direction.lower() != "x" and direction.lower() != "y"):
orbitFinalize("Class orbit.matrix_lattice.MATRIX_Lattice, method trackDispersionData(...): direction should be x or y.")
#track dispersion
eps_length = 0.00001 # 10^-6 meter
dir_ind = 0
if(direction.lower() == "y"):
dir_ind = 2
track_m = Matrix(3,3)
track_m.unit()
track_m.set(2,0,0.)
track_m.set(2,1,0.)
track_m.set(2,2,1.)
track_v = PhaseVector(3)
#kinematics coefficient calculation
Etotal = math.sqrt(momentum**2 + mass**2)
beta = momentum/Etotal
gamma = Etotal/mass
Ekin = Etotal - mass
m_coeff = momentum*momentum/(mass + Ekin)
track_v.set(0,disp)
track_v.set(1,disp_p)
track_v.set(2,1.)
position = 0.
pos_arr = []
disp_arr = []
disp_p_arr = []
#put in array the initial dispersions
pos_arr.append(position)
disp_arr.append(track_v.get(0))
disp_p_arr.append(track_v.get(1))
position_old = position
disp_old = disp
#count = 0
for matrixNode in self.getNodes():
if(isinstance(matrixNode,BaseMATRIX) == True):
disp = track_v.get(0)
if(abs(position_old-position) > eps_length or abs(disp_old - disp) > eps_length):
pos_arr.append(position)
disp_arr.append(track_v.get(0))
disp_p_arr.append(track_v.get(1))
disp_old = disp
position_old = position
mt = matrixNode.getMatrix()
ind0 = 0+dir_ind
ind1 = 1+dir_ind
track_m.set(0,0,mt.get(ind0,ind0))
track_m.set(0,1,mt.get(ind0,ind1))
track_m.set(1,0,mt.get(ind1,ind0))
track_m.set(1,1,mt.get(ind1,ind1))
track_m.set(0,2,mt.get(ind0,5)*m_coeff)
track_m.set(1,2,mt.get(ind1,5)*m_coeff)
track_v = track_m.mult(track_v)
position = position + matrixNode.getLength()
pos_arr.append(position)
disp_arr.append(track_v.get(0))
disp_p_arr.append(track_v.get(1))
#pack the resulting tuple
graph_disp_arr = []
graph_disp_p_arr = []
for i in range(len(pos_arr)):
graph_disp_arr.append((pos_arr[i],disp_arr[i]))
graph_disp_p_arr.append((pos_arr[i],disp_p_arr[i]))
return (graph_disp_arr,graph_disp_p_arr)
def trackTwissData(self, alpha, beta, direction="x"):
"""
Returns the tuple ([(position, phase advance/2/pi),...], [(position, alpha),...],[(position,beta),...] ).
The tracking starts from the values specified as the initial parameters.
The possible values for direction parameter "x" or "y".
"""
if (direction.lower() != "x" and direction.lower() != "y"):
orbitFinalize(
"Class orbit.matrix_lattice.MATRIX_Lattice, method trackTwissData(...): direction should be x or y."
)
#track twiss
eps_length = 0.00001 # 10^-6 meter
dir_ind = 0
if (direction.lower() == "y"):
dir_ind = 2
gamma = (1.0 + alpha * alpha) / beta
track_m = Matrix(3, 3)
track_m.unit()
track_v = PhaseVector(3)
track_v.set(0, alpha)
track_v.set(1, beta)
track_v.set(2, gamma)
position = 0.
pos_arr = []
alpha_arr = []
beta_arr = []
#phi is for tune accumulation
phi = 0.
#phase advance
mu_arr = []
#count = 0
pos_arr.append(position)
mu_arr.append(phi)
alpha_arr.append(track_v.get(0))
beta_arr.append(track_v.get(1))
for matrixNode in self.getNodes():
mt = matrixNode.getMatrix()
ind0 = 0 + dir_ind
ind1 = 1 + dir_ind
track_m.set(
0, 0,
mt.get(ind0, ind0) * mt.get(ind1, ind1) +
mt.get(ind0, ind1) * mt.get(ind1, ind0))
track_m.set(0, 1, -mt.get(ind0, ind0) * mt.get(ind1, ind0))
track_m.set(0, 2, -mt.get(ind0, ind1) * mt.get(ind1, ind1))
track_m.set(1, 0, -2 * mt.get(ind0, ind0) * mt.get(ind0, ind1))
track_m.set(1, 1, mt.get(ind0, ind0) * mt.get(ind0, ind0))
track_m.set(1, 2, mt.get(ind0, ind1) * mt.get(ind0, ind1))
track_m.set(2, 0, -2 * mt.get(ind1, ind0) * mt.get(ind1, ind1))
track_m.set(2, 1, mt.get(ind1, ind0) * mt.get(ind1, ind0))
track_m.set(2, 2, mt.get(ind1, ind1) * mt.get(ind1, ind1))
alpha_0 = track_v.get(0)
beta_0 = track_v.get(1)
delta_phi = math.atan(
mt.get(ind0, ind1) /
(beta_0 * mt.get(ind0, ind0) - alpha_0 * mt.get(ind0, ind1)))
phi = phi + delta_phi
track_v = track_m.mult(track_v)
position = position + matrixNode.getLength()
if (isinstance(matrixNode, BaseMATRIX) == True):
#only the main nodes are used, the or-cases deal with markers with zero length
#print position, matrixNode.getParam("matrix_parent_node_active_index") == 1
#print position, matrixNode.getName(), track_v.get(1)
if matrixNode.getLength() > 0:
#print position, matrixNode.getName(), track_v.get(1)
pos_arr.append(position)
alpha_arr.append(track_v.get(0))
beta_arr.append(track_v.get(1))
mu_arr.append(phi / (2 * math.pi))
#count = count + 1
#pack the resulting tuple
tune = phi / (2 * math.pi)
graph_alpha_arr = []
graph_beta_arr = []
graph_mu_arr = []
#print count, len(pos_arr)
for i in range(len(pos_arr)):
graph_mu_arr.append((pos_arr[i], mu_arr[i]))
graph_alpha_arr.append((pos_arr[i], alpha_arr[i]))
graph_beta_arr.append((pos_arr[i], beta_arr[i]))
return (graph_mu_arr, graph_alpha_arr, graph_beta_arr)
class LinacTrMatrixGenNode(MarkerLinacNode):
"""
Linac Accelerator Nodes for Transport Matrices generation.
These nodes are using thethe Initial Coordinates particles Attrubutes.
Each node (if it is not the first one) calculates the transport matrix
between the previous node and itself.
The matrix is a 7x7 matrix that transforms the initial particles
coordinates to the final ones that are in the bu
"""
def __init__(self, trMatricesController, name="TrMatrixGen"):
if (name == "TrMatrixGen"):
name += name + ":" + str(trMatricesController.getCount())
MarkerLinacNode.__init__(self, name)
self.trMatricesController = trMatricesController
self.trMtrxNode_ind = trMatricesController.getCount()
self.use_twiss_weight_x = 0
self.use_twiss_weight_y = 0
self.use_twiss_weight_z = 0
self.relativistic_beta = 0.
self.relativistic_gamma = 0.
#--------------------------------------
self.trMtrx = Matrix(7, 7)
#--------------------------------------
self.trMatricesController.addNode(self)
def setInternalIndex(self, ind):
"""
Sets the index of the TrMatrxGenNode in the controller
"""
self.trMtrxNode_ind = ind
def getTrMatricesController(self):
"""
Returns the LinacTrMatricesContrioller that keeps the references to the TrMatrxGenNodes.
"""
return self.trMatricesController
def getTwissWeightUse(self):
"""
Returns (use_x,use,use_z) tuple where use_{} == 1 means the Twiss weights will be used.
"""
res_arr = [True, True, True]
if (self.use_twiss_weight_x == 0): res_arr[0] = False
if (self.use_twiss_weight_y == 0): res_arr[1] = False
if (self.use_twiss_weight_z == 0): res_arr[2] = False
return tuple(res_arr)
def setTwissWeightUse(self, use_twiss_weight_x, use_twiss_weight_y,
use_twiss_weight_z):
"""
Sets (use_x,use,use_z) tuple where use_{} == 1 means the Twiss weights will be used.
"""
self.use_twiss_weight_x = 0
self.use_twiss_weight_y = 0
self.use_twiss_weight_z = 0
if (use_twiss_weight_x == True): self.use_twiss_weight_x = 1
if (use_twiss_weight_y == True): self.use_twiss_weight_y = 1
if (use_twiss_weight_z == True): self.use_twiss_weight_z = 1
def track(self, paramsDict):
bunch = paramsDict["bunch"]
self.relativistic_beta = bunch.getSyncParticle().beta()
self.relativistic_gamma = bunch.getSyncParticle().gamma()
if (self.trMtrxNode_ind == 0):
self.trMtrx.unit()
copyCoordsToInitCoordsAttr(bunch)
else:
transportMtrxFromInitCoords(bunch, self.trMtrx,
self.use_twiss_weight_x,
self.use_twiss_weight_y,
self.use_twiss_weight_z)
def trackDesign(self, paramsDict):
"""
This method does nothing for the aperture case.
"""
pass
def getBeta(self):
"""
Returns relativistic beta at this node.
"""
return self.relativistic_beta
def getGamma(self):
"""
Returns relativistic gamma at this node.
"""
return self.relativistic_gamma
def getTransportMatrix(self):
"""
Return transport matrix (7x7).
"""
return self.trMtrx
def getDetXYZ(self, trMtrx=None):
"""
Returns the determinants of the transformations in (x,y,z) directions.
"""
if (trMtrx == None): trMtrx = self.trMtrx
det_x = trMtrx.get(0, 0) * trMtrx.get(1, 1) - trMtrx.get(
1, 0) * trMtrx.get(0, 1)
det_y = trMtrx.get(0 + 2, 0 + 2) * trMtrx.get(
1 + 2, 1 + 2) - trMtrx.get(1 + 2, 0 + 2) * trMtrx.get(
0 + 2, 1 + 2)
det_z = trMtrx.get(0 + 4, 0 + 4) * trMtrx.get(
1 + 4, 1 + 4) - trMtrx.get(1 + 4, 0 + 4) * trMtrx.get(
0 + 4, 1 + 4)
return (det_x, det_y, det_z)
def getNormDetXYZ(self):
"""
Returns the normalized determinants of the transformations in (x,y,z) directions.
"""
(node0, node1) = self.getTwoNodes()
beta_in = node0.getBeta()
beta_out = node1.getBeta()
gamma_in = node0.getGamma()
gamma_out = node1.getGamma()
gb_in = beta_in * gamma_in
gb_out = beta_out * gamma_out
(det_x, det_y, det_z) = self.getDetXYZ(self.trMtrx)
return ((gb_out / gb_in) * det_x, (gb_out / gb_in) * det_y,
(beta_in / beta_out) * det_z)
def getTwoNodes(self):
"""
Returns two LinacTrMatrixGenNode nodes. The transport matrix is between these nodes.
"""
node0 = self
if (self.trMtrxNode_ind > 0):
node0 = self.trMatricesController.getNode(0)
node1 = self.trMatricesController.getNode(self.trMtrxNode_ind)
return (node0, node1)
def printMatrix(self):
"""
Print the matrix.
"""
name0 = "None"
if (self.trMtrxNode_ind > 0):
name0 = self.trMatricesController.getNode(self.trMtrxNode_ind -
1).getName()
name1 = self.trMatricesController.getNode(
self.trMtrxNode_ind).getName()
print "----Transport matrix--- from name0=", name0, " to name1=", name1
m = self.trMtrx
for i in xrange(m.size()[0]):
for j in xrange(m.size()[1]):
print("m(" + str(i) + "," + str(j) + ")=" +
"%12.5g" % m.get(i, j) + " "),
print ""
stack_size = 8
elem_count = 0
lattice_arr = []
while ( s != ""):
if(len(s_arr) < stack_size):
s_arr.append(s)
else:
#======= we have the full stack of lines ===========
if(s_arr[1].find("Element transfer matrix:") >= 0):
res_arr = s_arr[0].split()
name = res_arr[len(res_arr) - 2][1:]
elem_count = elem_count + 1
#print "i=",elem_count," name=",name
matrix = Matrix(6,6)
for i in range(6):
res_arr = s_arr[i+2].split()
for j in range(6):
matrix.set(i,j, float(res_arr[j]))
#printM(matrix)
lattice_arr.append((name,matrix))
for it in range(stack_size-1):
s_arr[it] = s_arr[it+1]
s_arr[stack_size-1] = s
#print "debug s=",s
s = inF.readline()
inF.close()
mad_matrix = Matrix(6,6)
# Initial Bunch #------------------------------ nParticles = 10 b_in = Bunch() for i in range(nParticles): (x,xp,y,yp,z,dE) = (random.random(),random.random(),random.random(),random.random(),random.random(),random.random()) b_in.addParticle(x,xp,y,yp,z,dE) #----set up ids ParticleIdNumber.addParticleIdNumbers(b_in) b_out = Bunch() b_in.copyEmptyBunchTo(b_out) #-----set up initial transport matix, later we have to extract it from bunches mtrxA_init = Matrix(7,7) mtrxA_init.zero() for ix in range(6): mtrxA_init.set(ix,6,1.0*ix) mtrxA_init.set(6,6,1.0) for ix in range(6): mtrxA_init.set(ix,ix,1.*(ix+1)) printM(mtrxA_init, "Init M ") #----------------------------------------------------------------------------- #------ put particles from b_in bunch to b_out noise_level = 0.000001 for i in range(nParticles): coord_arr = (b_in.x(i),b_in.xp(i),b_in.y(i),b_in.yp(i),b_in.z(i),b_in.dE(i)) coord_arr_res = [0.,0.,0.,0.,0.,0.] for ix in range(6): coord_arr_res[ix] = mtrxA_init.get(ix,6)
E0TL_Z_RMS_arr.append([0.154 / 1000., 12.80])
E0TL_Z_RMS_arr.append([0.170 / 1000., 11.30])
E0TL_Z_RMS_arr.append([0.178 / 1000., 10.92])
E0TL_Z_RMS_arr.append([0.186 / 1000., 10.60])
E0TL_Z_RMS_arr.append([0.196 / 1000., 10.80])
n_points = len(E0TL_Z_RMS_arr)
use_twiss_weight_x = True
use_twiss_weight_y = True
use_twiss_weight_z = True
res_sizes_trMtrx_arr = []
for [E0TL, z_rms_exp] in E0TL_Z_RMS_arr:
trMtrx = Matrix(7, 7)
bnch02.setParam("E0TL", E0TL)
bnch02.setParam("E0L", E0TL)
bunch = Bunch()
bunch_in.copyBunchTo(bunch)
accLattice.trackBunch(bunch, None, None, bnch02_ind)
transportMtrxFromInitCoords(bunch, trMtrx, use_twiss_weight_x,
use_twiss_weight_y, use_twiss_weight_z)
twiss_analysis.analyzeBunch(bunch)
x_rms = math.sqrt(
twiss_analysis.getTwiss(0)[1] * twiss_analysis.getTwiss(0)[3]) * 1000.
y_rms = math.sqrt(
twiss_analysis.getTwiss(1)[1] * twiss_analysis.getTwiss(1)[3]) * 1000.
z_rms = math.sqrt(
twiss_analysis.getTwiss(2)[1] * twiss_analysis.getTwiss(2)[3]) * 1000.
z_to_phase_coeff = bunch_gen.getZtoPhaseCoeff(bunch)
printM(final_matrix) #=====trackBunch method for the whole lattice approach============ print "============= Matrix for the whole lattice ==========" transp_matrix = Matrix(6,6) matrixGenerator = MatrixGenerator() matrixGenerator.initBunch(matrixTracker.b) teapot_latt.trackBunch(matrixTracker.b) matrixGenerator.calculateMatrix(matrixTracker.b,transp_matrix) transormToMAD(transp_matrix,momentum,beta) printM(transp_matrix) #=====trackBunch method for each TEAPOT element approach============ print "Total number of TEAPOT Base Nodes=",len(teapot_latt.getNodes()) printNames = [] res_matrix = Matrix(6,6) res_matrix.unit() n_max = len(teapot_latt.getNodes()) #n_max = 74 length = 0. for i in range(n_max): node = teapot_latt.getNodes()[i] print "i=",i," name =",node.getName()," L=",length m = Matrix(6,6) matrixGenerator.initBunch(matrixTracker.b) node.trackBunch(matrixTracker.b) matrixGenerator.calculateMatrix(matrixTracker.b,m) #m_prt0 = m.copy() #transormToMAD(m_prt0,momentum,beta) #printM(m_prt0) res_matrix = m.mult(res_matrix)
def trackDispersionData(self, momentum, mass, disp, disp_p, direction="x"):
"""
Returns the tuple ([(position, disp),...],[(position,disp_p),...] ).
The tracking starts from the values specified as the initial parameters.
The possible values for direction parameter "x" or "y".
"""
if (direction.lower() != "x" and direction.lower() != "y"):
orbitFinalize(
"Class orbit.matrix_lattice.MATRIX_Lattice, method trackDispersionData(...): direction should be x or y."
)
#track dispersion
eps_length = 0.00001 # 10^-6 meter
dir_ind = 0
if (direction.lower() == "y"):
dir_ind = 2
track_m = Matrix(3, 3)
track_m.unit()
track_m.set(2, 0, 0.)
track_m.set(2, 1, 0.)
track_m.set(2, 2, 1.)
track_v = PhaseVector(3)
#kinematics coefficient calculation
Etotal = math.sqrt(momentum**2 + mass**2)
beta = momentum / Etotal
gamma = Etotal / mass
Ekin = Etotal - mass
m_coeff = momentum * momentum / (mass + Ekin)
track_v.set(0, disp)
track_v.set(1, disp_p)
track_v.set(2, 1.)
position = 0.
pos_arr = []
disp_arr = []
disp_p_arr = []
#put in array the initial dispersions
#count = 0
pos_arr.append(position)
disp_arr.append(track_v.get(0))
disp_p_arr.append(track_v.get(1))
for matrixNode in self.getNodes():
if (isinstance(matrixNode, BaseMATRIX) == True):
mt = matrixNode.getMatrix()
ind0 = 0 + dir_ind
ind1 = 1 + dir_ind
track_m.set(0, 0, mt.get(ind0, ind0))
track_m.set(0, 1, mt.get(ind0, ind1))
track_m.set(1, 0, mt.get(ind1, ind0))
track_m.set(1, 1, mt.get(ind1, ind1))
track_m.set(0, 2, mt.get(ind0, 5) * m_coeff)
track_m.set(1, 2, mt.get(ind1, 5) * m_coeff)
track_v = track_m.mult(track_v)
position = position + matrixNode.getLength()
#only the main nodes are used, the or-cases deal with markers with zero length
if (isinstance(matrixNode, BaseMATRIX) == True):
pos_arr.append(position)
disp_arr.append(track_v.get(0))
disp_p_arr.append(track_v.get(1))
#pack the resulting tuple
graph_disp_arr = []
graph_disp_p_arr = []
for i in range(len(pos_arr)):
graph_disp_arr.append((pos_arr[i], disp_arr[i]))
graph_disp_p_arr.append((pos_arr[i], disp_p_arr[i]))
return (graph_disp_arr, graph_disp_p_arr)
class MATRIX_Lattice(AccLattice):
"""
The subclass of the AccLattice class. Shell class for the BaseMATRIX nodes collection.
In the beginning the lattcie is empty.
"""
def __init__(self, name = None):
AccLattice.__init__(self,name)
self.oneTurmMatrix = Matrix(7,7)
self.oneTurmMatrix.unit()
def initialize(self):
"""
Method. Initializes the matrix lattice, child node structures, and calculates
the one turn matrix.
"""
AccLattice.initialize(self)
self.makeOneTurnMatrix()
def makeOneTurnMatrix(self):
"""
Calculates the one turn matrix.
"""
self.oneTurmMatrix.unit()
for matrixNode in self.getNodes():
if(isinstance(matrixNode,BaseMATRIX) == True):
self.oneTurmMatrix = matrixNode.getMatrix().mult(self.oneTurmMatrix)
return self.oneTurmMatrix
def getOneTurnMatrix(self):
"""
Returns the one turn matrix.
"""
return self.oneTurmMatrix
def getRingParametersDict(self,momentum,mass):
"""
Returns the dictionary with different ring parametrs
calculated from the one turn transport matrix.
"""
res_dict = {}
Etotal = math.sqrt(momentum**2 + mass**2)
beta = momentum/Etotal
gamma = Etotal/mass
Ekin = Etotal - mass
res_dict["momentum [GeV/c]"] = momentum
res_dict["mass [GeV]"] = mass
res_dict["Ekin [GeV]"] = Ekin
#longitudinal params
c = 2.99792458e+8
ring_length = self.getLength()
T = ring_length/(beta*c)
res_dict["period [sec]"] = T
res_dict["frequency [Hz]"] = 1./T
#transverse twiss parameters
mt = self.oneTurmMatrix
res_dict["fractional tune x"] = None
res_dict["fractional tune y"] = None
res_dict["alpha x"] = None
res_dict["alpha y"] = None
res_dict["beta x [m]"] = None
res_dict["beta y [m]"] = None
res_dict["gamma x [m^-1]"] = None
res_dict["gamma y [m^-1]"] = None
res_dict["dispersion x [m]"] = None
res_dict["dispersion y [m]"] = None
res_dict["dispersion prime x"] = None
res_dict["dispersion prime y"] = None
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
if(abs(cos_phi_x) >= 1.0 or abs(cos_phi_x) >= 1.0):
return res_dict
sign_x = +1.0
if(abs(mt.get(0,1)) != 0.): sign_x = mt.get(0,1)/abs(mt.get(0,1))
sign_y = +1.0
if(abs(mt.get(2,3)) != 0.): sign_y = mt.get(2,3)/abs(mt.get(2,3))
sin_phi_x = math.sqrt(1. - cos_phi_x*cos_phi_x)*sign_x
sin_phi_y = math.sqrt(1. - cos_phi_y*cos_phi_y)*sign_y
nux = math.acos(cos_phi_x)/(2*math.pi)*sign_x
nuy = math.acos(cos_phi_y)/(2*math.pi) *sign_y
res_dict["fractional tune x"] = nux
res_dict["fractional tune y"] = nuy
# alpha, beta, gamma
beta_x = mt.get(0,1)/sin_phi_x
beta_y = mt.get(2,3)/sin_phi_y
alpha_x = (mt.get(0,0) - mt.get(1,1))/(2*sin_phi_x)
alpha_y = (mt.get(2,2) - mt.get(3,3))/(2*sin_phi_y)
gamma_x = -mt.get(1,0)/sin_phi_x
gamma_y = -mt.get(3,2)/sin_phi_y
# dispersion and dispersion prime
m_coeff = momentum * momentum / Etotal
disp_x = m_coeff*(mt.get(0,5)*(1-mt.get(1,1))+mt.get(0,1)*mt.get(1,5))/(2-mt.get(0,0)-mt.get(1,1))
disp_y = m_coeff*(mt.get(2,5)*(1-mt.get(3,3))+mt.get(2,3)*mt.get(3,5))/(2-mt.get(2,2)-mt.get(3,3))
disp_pr_x = m_coeff*(mt.get(1,0)*mt.get(0,5)+mt.get(1,5)*(1-mt.get(0,0)))/(2-mt.get(0,0)-mt.get(1,1))
disp_pr_y = m_coeff*(mt.get(3,2)*mt.get(2,5)+mt.get(3,5)*(1-mt.get(2,2)))/(2-mt.get(2,2)-mt.get(3,3))
res_dict["alpha x"] = alpha_x
res_dict["alpha y"] = alpha_y
res_dict["beta x [m]"] = beta_x
res_dict["beta y [m]"] = beta_y
res_dict["gamma x [m^-1]"] = gamma_x
res_dict["gamma y [m^-1]"] = gamma_y
res_dict["dispersion x [m]"] = disp_x
res_dict["dispersion y [m]"] = disp_y
res_dict["dispersion prime x"] = disp_pr_x
res_dict["dispersion prime y"] = disp_pr_y
#more longitudinal params
termx = mt.get(4,0) * disp_x
termxp = mt.get(4,1) * disp_pr_x
termy = mt.get(4,2) * disp_y
termyp = mt.get(4,3) * disp_pr_y
termdE = mt.get(4,5) * beta * beta * Etotal
eta_ring = -(termx + termxp + termy + termyp + termdE)\
/ ring_length
res_dict["eta"] = eta_ring
alpha_p = eta_ring + 1. / (gamma * gamma)
sqarg = alpha_p
if(alpha_p < 0.):
sqarg = -alpha_p
gamma_trans = 1.0 / math.sqrt(sqarg)
res_dict["gamma transition"] = gamma_trans
res_dict["transition energy [GeV]"] = (gamma_trans - 1.0)*mass
res_dict["momentum compaction"] = alpha_p
return res_dict
def getRingTwissDataX(self,momentum,mass):
"""
Returns the tuple ([(position, phase advanceX)/2/pi,...],[(position, alphaX),...],[(position,betaX),...] ).
"""
res_dict = self.getRingParametersDict(momentum,mass)
alpha_x = res_dict["alpha x"]
beta_x = res_dict["beta x [m]"]
return self.trackTwissData(alpha_x,beta_x,"x")
def getRingTwissDataY(self,momentum,mass):
"""
Returns the tuple ([(position, phase advanceY/2/pi),...],[(position, alphaY),...],[(position,betaY),...] ).
"""
res_dict = self.getRingParametersDict(momentum,mass)
alpha_y = res_dict["alpha y"]
beta_y = res_dict["beta y [m]"]
return self.trackTwissData(alpha_y,beta_y,"y")
def trackTwissData(self, alpha, beta, direction = "x"):
"""
Returns the tuple ([(position, phase advance/2/pi),...], [(position, alpha),...],[(position,beta),...] ).
The tracking starts from the values specified as the initial parameters.
The possible values for direction parameter "x" or "y".
"""
if(direction.lower() != "x" and direction.lower() != "y"):
orbitFinalize("Class orbit.matrix_lattice.MATRIX_Lattice, method trackTwissData(...): direction should be x or y.")
#track twiss
eps_length = 0.00001 # 10^-6 meter
dir_ind = 0
if(direction.lower() == "y"):
dir_ind = 2
gamma = (1.0+alpha*alpha)/beta
track_m = Matrix(3,3)
track_m.unit()
track_v = PhaseVector(3)
track_v.set(0,alpha)
track_v.set(1,beta)
track_v.set(2,gamma)
position = 0.
pos_arr = []
alpha_arr = []
beta_arr = []
#phi is for tune accumulation
phi = 0.
#phase advance
mu_arr = []
#put in array the initial twiss
pos_arr.append(position)
mu_arr.append(phi)
alpha_arr.append(track_v.get(0))
beta_arr.append(track_v.get(1))
position_old = position
beta_old = beta
#count = 0
for matrixNode in self.getNodes():
if(isinstance(matrixNode,BaseMATRIX) == True):
beta = track_v.get(1)
if(abs(position_old-position) > eps_length or abs(beta_old - beta) > eps_length):
pos_arr.append(position)
alpha_arr.append(track_v.get(0))
beta_arr.append(track_v.get(1))
mu_arr.append(phi/(2*math.pi))
mt = matrixNode.getMatrix()
ind0 = 0+dir_ind
ind1 = 1+dir_ind
track_m.set(0,0,mt.get(ind0,ind0)*mt.get(ind1,ind1)+mt.get(ind0,ind1)*mt.get(ind1,ind0))
track_m.set(0,1,-mt.get(ind0,ind0)*mt.get(ind1,ind0))
track_m.set(0,2,-mt.get(ind0,ind1)*mt.get(ind1,ind1))
track_m.set(1,0,-2*mt.get(ind0,ind0)*mt.get(ind0,ind1))
track_m.set(1,1,mt.get(ind0,ind0)*mt.get(ind0,ind0))
track_m.set(1,2,mt.get(ind0,ind1)*mt.get(ind0,ind1))
track_m.set(2,0,-2*mt.get(ind1,ind0)*mt.get(ind1,ind1))
track_m.set(2,1,mt.get(ind1,ind0)*mt.get(ind1,ind0))
track_m.set(2,2,mt.get(ind1,ind1)*mt.get(ind1,ind1))
alpha_0 = track_v.get(0)
beta_0 = track_v.get(1)
delta_phi = math.atan( mt.get(ind0,ind1)/( beta_0*mt.get(ind0,ind0) - alpha_0*mt.get(ind0,ind1)))
phi = phi + delta_phi
track_v = track_m.mult(track_v)
position_old = position
beta_old = beta_0
position = position + matrixNode.getLength()
#count = count + 1
pos_arr.append(position)
alpha_arr.append(track_v.get(0))
beta_arr.append(track_v.get(1))
mu_arr.append(phi/(2*math.pi))
#pack the resulting tuple
tune = phi/(2*math.pi)
graph_alpha_arr = []
graph_beta_arr = []
graph_mu_arr = []
for i in range(len(pos_arr)):
graph_mu_arr.append((pos_arr[i], mu_arr[i]))
graph_alpha_arr.append((pos_arr[i],alpha_arr[i]))
graph_beta_arr.append((pos_arr[i],beta_arr[i]))
return (graph_mu_arr,graph_alpha_arr,graph_beta_arr)
def getRingDispersionDataX(self,momentum,mass):
"""
Returns the tuple ([(position, dispX),...],[(position,disp_pX),...] ).
"""
res_dict = self.getRingParametersDict(momentum,mass)
disp = res_dict["dispersion x [m]"]
disp_p = res_dict["dispersion prime x"]
return self.trackDispersionData(momentum,mass,disp, disp_p,"x")
def getRingDispersionDataY(self,momentum,mass):
"""
Returns the tuple ([(position, dispY),...],[(position,disp_pY),...] ).
"""
res_dict = self.getRingParametersDict(momentum,mass)
disp = res_dict["dispersion y [m]"]
disp_p = res_dict["dispersion prime y"]
return self.trackDispersionData(momentum,mass,disp, disp_p,"y")
def trackDispersionData(self,momentum,mass, disp, disp_p, direction = "x"):
"""
Returns the tuple ([(position, disp),...],[(position,disp_p),...] ).
The tracking starts from the values specified as the initial parameters.
The possible values for direction parameter "x" or "y".
"""
if(direction.lower() != "x" and direction.lower() != "y"):
orbitFinalize("Class orbit.matrix_lattice.MATRIX_Lattice, method trackDispersionData(...): direction should be x or y.")
#track dispersion
eps_length = 0.00001 # 10^-6 meter
dir_ind = 0
if(direction.lower() == "y"):
dir_ind = 2
track_m = Matrix(3,3)
track_m.unit()
track_m.set(2,0,0.)
track_m.set(2,1,0.)
track_m.set(2,2,1.)
track_v = PhaseVector(3)
#kinematics coefficient calculation
Etotal = math.sqrt(momentum**2 + mass**2)
beta = momentum/Etotal
gamma = Etotal/mass
Ekin = Etotal - mass
m_coeff = momentum*momentum/(mass + Ekin)
track_v.set(0,disp)
track_v.set(1,disp_p)
track_v.set(2,1.)
position = 0.
pos_arr = []
disp_arr = []
disp_p_arr = []
#put in array the initial dispersions
pos_arr.append(position)
disp_arr.append(track_v.get(0))
disp_p_arr.append(track_v.get(1))
position_old = position
disp_old = disp
#count = 0
for matrixNode in self.getNodes():
if(isinstance(matrixNode,BaseMATRIX) == True):
disp = track_v.get(0)
if(abs(position_old-position) > eps_length or abs(disp_old - disp) > eps_length):
pos_arr.append(position)
disp_arr.append(track_v.get(0))
disp_p_arr.append(track_v.get(1))
disp_old = disp
position_old = position
mt = matrixNode.getMatrix()
ind0 = 0+dir_ind
ind1 = 1+dir_ind
track_m.set(0,0,mt.get(ind0,ind0))
track_m.set(0,1,mt.get(ind0,ind1))
track_m.set(1,0,mt.get(ind1,ind0))
track_m.set(1,1,mt.get(ind1,ind1))
track_m.set(0,2,mt.get(ind0,5)*m_coeff)
track_m.set(1,2,mt.get(ind1,5)*m_coeff)
track_v = track_m.mult(track_v)
position = position + matrixNode.getLength()
pos_arr.append(position)
disp_arr.append(track_v.get(0))
disp_p_arr.append(track_v.get(1))
#pack the resulting tuple
graph_disp_arr = []
graph_disp_p_arr = []
for i in range(len(pos_arr)):
graph_disp_arr.append((pos_arr[i],disp_arr[i]))
graph_disp_p_arr.append((pos_arr[i],disp_p_arr[i]))
return (graph_disp_arr,graph_disp_p_arr)
def getSubLattice(self, index_start = -1, index_stop = -1,):
"""
It returns the new MATRIX_Lattice with children with indexes
between index_start and index_stop inclusive
"""
return self._getSubLattice(MATRIX_Lattice(),index_start,index_stop)
def trackBunch(self, bunch, paramsDict = {}, actionContainer = None):
"""
It tracks the bunch through the lattice.
"""
if(actionContainer == None): actionContainer = AccActionsContainer("Bunch Tracking")
paramsDict["bunch"] = bunch
def track(paramsDict):
node = paramsDict["node"]
node.track(paramsDict)
actionContainer.addAction(track, AccActionsContainer.BODY)
self.trackActions(actionContainer,paramsDict)
actionContainer.removeAction(track, AccActionsContainer.BODY)
def trackOrbit(self, z0):
"""
Returns the tuple ([(position, x),...], [(position, y),...] ).
The tracking starts from the values specified as the initial parameters.
z0 fulfill: z0 = Mz0 with M as one turn matrix
"""
eps_length = 0.00001 # 10^-6 meter
pos_arr = []
orbitX_arr = []
orbitY_arr = []
orbitXP_arr = []
orbitYP_arr = []
position = 0.
pos_arr.append(position)
track_o = Matrix(6, 6)
track_ov = PhaseVector(6)
track_ov.set(0, z0[0])
track_ov.set(1, z0[1])
track_ov.set(2, z0[2])
track_ov.set(3, z0[3])
track_ov.set(4, z0[4])
track_ov.set(5, z0[5])
orbitX_arr.append(track_ov.get(0))
orbitY_arr.append(track_ov.get(2))
orbitXP_arr.append(track_ov.get(1))
orbitYP_arr.append(track_ov.get(3))
position_old = position
for matrixNode in self.getNodes():
if (isinstance(matrixNode, BaseMATRIX) == True):
if (abs(position_old - position) > eps_length):
pos_arr.append(position)
orbitX_arr.append(track_ov.get(0))
orbitY_arr.append(track_ov.get(2))
orbitXP_arr.append(track_ov.get(1))
orbitYP_arr.append(track_ov.get(3))
mt = matrixNode.getMatrix()
for i in range(6):
for j in range(6):
track_o.set(i, j, mt.get(i, j))
track_ov = track_o.mult(track_ov)
for i in range(6):
tmp = track_ov.get(i)
track_ov.set(i, tmp + mt.get(i, 6))
position_old = position
position = position + matrixNode.getLength()
pos_arr.append(position)
orbitX_arr.append(track_ov.get(0))
orbitY_arr.append(track_ov.get(2))
orbitXP_arr.append(track_ov.get(1))
orbitYP_arr.append(track_ov.get(3))
#pack the resulting tuple
graph_orbitX_arr = []
graph_orbitY_arr = []
for i in range(len(pos_arr)):
graph_orbitX_arr.append(
(pos_arr[i], orbitX_arr[i], orbitXP_arr[i]))
graph_orbitY_arr.append(
(pos_arr[i], orbitY_arr[i], orbitYP_arr[i]))
return (graph_orbitX_arr, graph_orbitY_arr)
v1 = PhaseVector(5) setV(v) setV(v1) printV(v) printV(v1) v3 = v.add(1) printV(v3) v3 = v.add(v1) printV(v3) print "v*v1 =", v.mult(v1) print "-----------------------------------------------" #check matrix print "----mult ----" m = Matrix(5, 4) m1 = Matrix(4, 5) setM(m) setM(m1) m3 = m.mult(m1) printM(m3) print "----add -----" m = Matrix(4, 4) m1 = Matrix(4, 4) setM(m) setM(m1) m3 = m.add(m1) printM(m3) m4 = m3.invert()
def trackTwissData(self, alpha, beta, direction = "x"):
"""
Returns the tuple ([(position, phase advance/2/pi),...], [(position, alpha),...],[(position,beta),...] ).
The tracking starts from the values specified as the initial parameters.
The possible values for direction parameter "x" or "y".
"""
if(direction.lower() != "x" and direction.lower() != "y"):
orbitFinalize("Class orbit.matrix_lattice.MATRIX_Lattice, method trackTwissData(...): direction should be x or y.")
#track twiss
eps_length = 0.00001 # 10^-6 meter
dir_ind = 0
if(direction.lower() == "y"):
dir_ind = 2
gamma = (1.0+alpha*alpha)/beta
track_m = Matrix(3,3)
track_m.unit()
track_v = PhaseVector(3)
track_v.set(0,alpha)
track_v.set(1,beta)
track_v.set(2,gamma)
position = 0.
pos_arr = []
alpha_arr = []
beta_arr = []
#phi is for tune accumulation
phi = 0.
#phase advance
mu_arr = []
#put in array the initial twiss
pos_arr.append(position)
mu_arr.append(phi)
alpha_arr.append(track_v.get(0))
beta_arr.append(track_v.get(1))
position_old = position
beta_old = beta
#count = 0
for matrixNode in self.getNodes():
if(isinstance(matrixNode,BaseMATRIX) == True):
beta = track_v.get(1)
if(abs(position_old-position) > eps_length or abs(beta_old - beta) > eps_length):
pos_arr.append(position)
alpha_arr.append(track_v.get(0))
beta_arr.append(track_v.get(1))
mu_arr.append(phi/(2*math.pi))
mt = matrixNode.getMatrix()
ind0 = 0+dir_ind
ind1 = 1+dir_ind
track_m.set(0,0,mt.get(ind0,ind0)*mt.get(ind1,ind1)+mt.get(ind0,ind1)*mt.get(ind1,ind0))
track_m.set(0,1,-mt.get(ind0,ind0)*mt.get(ind1,ind0))
track_m.set(0,2,-mt.get(ind0,ind1)*mt.get(ind1,ind1))
track_m.set(1,0,-2*mt.get(ind0,ind0)*mt.get(ind0,ind1))
track_m.set(1,1,mt.get(ind0,ind0)*mt.get(ind0,ind0))
track_m.set(1,2,mt.get(ind0,ind1)*mt.get(ind0,ind1))
track_m.set(2,0,-2*mt.get(ind1,ind0)*mt.get(ind1,ind1))
track_m.set(2,1,mt.get(ind1,ind0)*mt.get(ind1,ind0))
track_m.set(2,2,mt.get(ind1,ind1)*mt.get(ind1,ind1))
alpha_0 = track_v.get(0)
beta_0 = track_v.get(1)
delta_phi = math.atan( mt.get(ind0,ind1)/( beta_0*mt.get(ind0,ind0) - alpha_0*mt.get(ind0,ind1)))
phi = phi + delta_phi
track_v = track_m.mult(track_v)
position_old = position
beta_old = beta_0
position = position + matrixNode.getLength()
#count = count + 1
pos_arr.append(position)
alpha_arr.append(track_v.get(0))
beta_arr.append(track_v.get(1))
mu_arr.append(phi/(2*math.pi))
#pack the resulting tuple
tune = phi/(2*math.pi)
graph_alpha_arr = []
graph_beta_arr = []
graph_mu_arr = []
for i in range(len(pos_arr)):
graph_mu_arr.append((pos_arr[i], mu_arr[i]))
graph_alpha_arr.append((pos_arr[i],alpha_arr[i]))
graph_beta_arr.append((pos_arr[i],beta_arr[i]))
return (graph_mu_arr,graph_alpha_arr,graph_beta_arr)
#lattice.addNode(elem3)
#-----------------------------
# Set TEAPOT nodes parameters
#-----------------------------
elem1.setLength(1.0)
elem2.setLength(1.0)
elem3.setLength(1.0)
elem2.setnParts(5)
elem2.addParam("kq", -0.5)
lattice.initialize()
print "lattice length=", lattice.getLength()
transp_matrix = Matrix(6, 6)
matrixGenerator = MatrixGenerator()
matrixGenerator.initBunch(b)
lattice.trackBunch(b)
matrixGenerator.calculateMatrix(b, transp_matrix)
def printM(m):
print "----matrix--- size=", m.size()
for i in xrange(m.size()[0]):
for j in xrange(m.size()[1]):
print("(%1d,%1d)= % 6.5e " % (i, j, m.get(i, j))),
print ""
momentum = matrixTracker.b.getSyncParticle().momentum() beta = matrixTracker.b.getSyncParticle().beta() gamma = matrixTracker.b.getSyncParticle().gamma() mass = matrixTracker.b.getSyncParticle().mass() print " etotal=", gamma * mass print " ekin=", (gamma - 1) * mass matrixArr = matrixTracker.getTransportMatrixArray(teapot_latt) print "n marix =", len(matrixArr) final_matrix = matrixTracker.getOneTurnMatrix(teapot_latt) transormToMAD(final_matrix, momentum, beta) printM(final_matrix) #=====trackBunch method for the whole lattice approach============ print "============= Matrix for the whole lattice ==========" transp_matrix = Matrix(6, 6) matrixGenerator = MatrixGenerator() matrixGenerator.initBunch(matrixTracker.b) teapot_latt.trackBunch(matrixTracker.b) matrixGenerator.calculateMatrix(matrixTracker.b, transp_matrix) transormToMAD(transp_matrix, momentum, beta) printM(transp_matrix) #=====trackBunch method for each TEAPOT element approach============ print "Total number of TEAPOT Base Nodes=", len(teapot_latt.getNodes()) printNames = [] res_matrix = Matrix(6, 6) res_matrix.unit() n_max = len(teapot_latt.getNodes()) #n_max = 74 length = 0.
class MATRIX_Lattice(AccLattice):
"""
The subclass of the AccLattice class. Shell class for the BaseMATRIX nodes collection.
In the beginning the lattcie is empty.
"""
def __init__(self, name=None):
AccLattice.__init__(self, name)
self.oneTurnMatrix = Matrix(7, 7)
self.oneTurnMatrix.unit()
self.Matrix = Matrix(7, 7)
self.Matrix.unit()
def initialize(self):
"""
Method. Initializes the matrix lattice, child node structures, and calculates
the one turn matrix.
"""
AccLattice.initialize(self)
self.makeOneTurnMatrix()
def makeOneTurnMatrix(self):
"""
Calculates the one turn matrix.
"""
self.oneTurnMatrix.unit()
eps_length = 0.00001 # 10^-6 meter
position = 0.0
position_old = position
for matrixNode in self.getNodes():
if (isinstance(matrixNode, BaseMATRIX) == True):
self.oneTurnMatrix = matrixNode.getMatrix().mult(
self.oneTurnMatrix)
"""
if(abs(position_old-position) > eps_length):
print position, self.oneTurnMatrix.get(0,6)
position_old = position
position = position + matrixNode.getLength()
"""
return self.oneTurnMatrix
def makeMatrix(self, pos):
"""
Calculates the one turn matrix.
"""
self.Matrix.unit()
position = 0.0
for matrixNode in self.getNodes():
if (isinstance(matrixNode, BaseMATRIX) == True and position < pos):
self.Matrix = matrixNode.getMatrix().mult(self.Matrix)
position = position + matrixNode.getLength()
#print position, self.Matrix.get(0,6)
return self.Matrix
def getOneTurnMatrix(self):
"""
Returns the one turn matrix.
"""
return self.oneTurnMatrix
def getRingParametersDict(self, momentum, mass):
"""
Returns the dictionary with different ring parametrs
calculated from the one turn transport matrix.
"""
res_dict = {}
Etotal = math.sqrt(momentum**2 + mass**2)
beta = momentum / Etotal
gamma = Etotal / mass
Ekin = Etotal - mass
res_dict["momentum [GeV/c]"] = momentum
res_dict["mass [GeV]"] = mass
res_dict["Ekin [GeV]"] = Ekin
#longitudinal params
c = 2.99792458e+8
ring_length = self.getLength()
T = ring_length / (beta * c)
res_dict["period [sec]"] = T
res_dict["frequency [Hz]"] = 1. / T
#transverse twiss parameters
mt = self.oneTurnMatrix
res_dict["fractional tune x"] = None
res_dict["fractional tune y"] = None
res_dict["alpha x"] = None
res_dict["alpha y"] = None
res_dict["beta x [m]"] = None
res_dict["beta y [m]"] = None
res_dict["gamma x [m^-1]"] = None
res_dict["gamma y [m^-1]"] = None
res_dict["dispersion x [m]"] = None
res_dict["dispersion y [m]"] = None
res_dict["dispersion prime x"] = None
res_dict["dispersion prime y"] = None
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
if (abs(cos_phi_x) >= 1.0 or abs(cos_phi_x) >= 1.0):
return res_dict
sign_x = +1.0
if (abs(mt.get(0, 1)) != 0.): sign_x = mt.get(0, 1) / abs(mt.get(0, 1))
sign_y = +1.0
if (abs(mt.get(2, 3)) != 0.): sign_y = mt.get(2, 3) / abs(mt.get(2, 3))
sin_phi_x = math.sqrt(1. - cos_phi_x * cos_phi_x) * sign_x
sin_phi_y = math.sqrt(1. - cos_phi_y * cos_phi_y) * sign_y
nux = math.acos(cos_phi_x) / (2 * math.pi) * sign_x
nuy = math.acos(cos_phi_y) / (2 * math.pi) * sign_y
res_dict["fractional tune x"] = nux
res_dict["fractional tune y"] = nuy
# alpha, beta, gamma
beta_x = mt.get(0, 1) / sin_phi_x
beta_y = mt.get(2, 3) / sin_phi_y
alpha_x = (mt.get(0, 0) - mt.get(1, 1)) / (2 * sin_phi_x)
alpha_y = (mt.get(2, 2) - mt.get(3, 3)) / (2 * sin_phi_y)
gamma_x = -mt.get(1, 0) / sin_phi_x
gamma_y = -mt.get(3, 2) / sin_phi_y
# dispersion and dispersion prime
m_coeff = momentum * momentum / Etotal
disp_x = m_coeff * (mt.get(0, 5) * (1 - mt.get(1, 1)) + mt.get(0, 1) *
mt.get(1, 5)) / (2 - mt.get(0, 0) - mt.get(1, 1))
disp_y = m_coeff * (mt.get(2, 5) * (1 - mt.get(3, 3)) + mt.get(2, 3) *
mt.get(3, 5)) / (2 - mt.get(2, 2) - mt.get(3, 3))
disp_pr_x = m_coeff * (mt.get(1, 0) * mt.get(0, 5) + mt.get(1, 5) *
(1 - mt.get(0, 0))) / (2 - mt.get(0, 0) -
mt.get(1, 1))
disp_pr_y = m_coeff * (mt.get(3, 2) * mt.get(2, 5) + mt.get(3, 5) *
(1 - mt.get(2, 2))) / (2 - mt.get(2, 2) -
mt.get(3, 3))
res_dict["alpha x"] = alpha_x
res_dict["alpha y"] = alpha_y
res_dict["beta x [m]"] = beta_x
res_dict["beta y [m]"] = beta_y
res_dict["gamma x [m^-1]"] = gamma_x
res_dict["gamma y [m^-1]"] = gamma_y
res_dict["dispersion x [m]"] = disp_x
res_dict["dispersion y [m]"] = disp_y
res_dict["dispersion prime x"] = disp_pr_x
res_dict["dispersion prime y"] = disp_pr_y
#more longitudinal params
termx = mt.get(4, 0) * disp_x
termxp = mt.get(4, 1) * disp_pr_x
termy = mt.get(4, 2) * disp_y
termyp = mt.get(4, 3) * disp_pr_y
termdE = mt.get(4, 5) * beta * beta * Etotal
eta_ring = -(termx + termxp + termy + termyp + termdE)\
/ ring_length
res_dict["eta"] = eta_ring
alpha_p = eta_ring + 1. / (gamma * gamma)
sqarg = alpha_p
if (alpha_p < 0.):
sqarg = -alpha_p
gamma_trans = 1.0 / math.sqrt(sqarg)
res_dict["gamma transition"] = gamma_trans
res_dict["transition energy [GeV]"] = (gamma_trans - 1.0) * mass
res_dict["momentum compaction"] = alpha_p
return res_dict
def getRingTwissDataX(self, momentum, mass):
"""
Returns the tuple ([(position, phase advanceX)/2/pi,...],[(position, alphaX),...],[(position,betaX),...] ).
"""
res_dict = self.getRingParametersDict(momentum, mass)
alpha_x = res_dict["alpha x"]
beta_x = res_dict["beta x [m]"]
return self.trackTwissData(alpha_x, beta_x, "x")
def getRingTwissDataY(self, momentum, mass):
"""
Returns the tuple ([(position, phase advanceY/2/pi),...],[(position, alphaY),...],[(position,betaY),...] ).
"""
res_dict = self.getRingParametersDict(momentum, mass)
alpha_y = res_dict["alpha y"]
beta_y = res_dict["beta y [m]"]
return self.trackTwissData(alpha_y, beta_y, "y")
def trackTwissData(self, alpha, beta, direction="x"):
"""
Returns the tuple ([(position, phase advance/2/pi),...], [(position, alpha),...],[(position,beta),...] ).
The tracking starts from the values specified as the initial parameters.
The possible values for direction parameter "x" or "y".
"""
if (direction.lower() != "x" and direction.lower() != "y"):
orbitFinalize(
"Class orbit.matrix_lattice.MATRIX_Lattice, method trackTwissData(...): direction should be x or y."
)
#track twiss
eps_length = 0.00001 # 10^-6 meter
dir_ind = 0
if (direction.lower() == "y"):
dir_ind = 2
gamma = (1.0 + alpha * alpha) / beta
track_m = Matrix(3, 3)
track_m.unit()
track_v = PhaseVector(3)
track_v.set(0, alpha)
track_v.set(1, beta)
track_v.set(2, gamma)
position = 0.
pos_arr = []
alpha_arr = []
beta_arr = []
#phi is for tune accumulation
phi = 0.
#phase advance
mu_arr = []
#count = 0
pos_arr.append(position)
mu_arr.append(phi)
alpha_arr.append(track_v.get(0))
beta_arr.append(track_v.get(1))
for matrixNode in self.getNodes():
mt = matrixNode.getMatrix()
ind0 = 0 + dir_ind
ind1 = 1 + dir_ind
track_m.set(
0, 0,
mt.get(ind0, ind0) * mt.get(ind1, ind1) +
mt.get(ind0, ind1) * mt.get(ind1, ind0))
track_m.set(0, 1, -mt.get(ind0, ind0) * mt.get(ind1, ind0))
track_m.set(0, 2, -mt.get(ind0, ind1) * mt.get(ind1, ind1))
track_m.set(1, 0, -2 * mt.get(ind0, ind0) * mt.get(ind0, ind1))
track_m.set(1, 1, mt.get(ind0, ind0) * mt.get(ind0, ind0))
track_m.set(1, 2, mt.get(ind0, ind1) * mt.get(ind0, ind1))
track_m.set(2, 0, -2 * mt.get(ind1, ind0) * mt.get(ind1, ind1))
track_m.set(2, 1, mt.get(ind1, ind0) * mt.get(ind1, ind0))
track_m.set(2, 2, mt.get(ind1, ind1) * mt.get(ind1, ind1))
alpha_0 = track_v.get(0)
beta_0 = track_v.get(1)
delta_phi = math.atan(
mt.get(ind0, ind1) /
(beta_0 * mt.get(ind0, ind0) - alpha_0 * mt.get(ind0, ind1)))
phi = phi + delta_phi
track_v = track_m.mult(track_v)
position = position + matrixNode.getLength()
if (isinstance(matrixNode, BaseMATRIX) == True):
#only the main nodes are used, the or-cases deal with markers with zero length
#print position, matrixNode.getParam("matrix_parent_node_active_index") == 1
#print position, matrixNode.getName(), track_v.get(1)
if matrixNode.getLength() > 0:
#print position, matrixNode.getName(), track_v.get(1)
pos_arr.append(position)
alpha_arr.append(track_v.get(0))
beta_arr.append(track_v.get(1))
mu_arr.append(phi / (2 * math.pi))
#count = count + 1
#pack the resulting tuple
tune = phi / (2 * math.pi)
graph_alpha_arr = []
graph_beta_arr = []
graph_mu_arr = []
#print count, len(pos_arr)
for i in range(len(pos_arr)):
graph_mu_arr.append((pos_arr[i], mu_arr[i]))
graph_alpha_arr.append((pos_arr[i], alpha_arr[i]))
graph_beta_arr.append((pos_arr[i], beta_arr[i]))
return (graph_mu_arr, graph_alpha_arr, graph_beta_arr)
def getRingDispersionDataX(self, momentum, mass):
"""
Returns the tuple ([(position, dispX),...],[(position,disp_pX),...] ).
"""
res_dict = self.getRingParametersDict(momentum, mass)
disp = res_dict["dispersion x [m]"]
disp_p = res_dict["dispersion prime x"]
return self.trackDispersionData(momentum, mass, disp, disp_p, "x")
def getRingDispersionDataY(self, momentum, mass):
"""
Returns the tuple ([(position, dispY),...],[(position,disp_pY),...] ).
"""
res_dict = self.getRingParametersDict(momentum, mass)
disp = res_dict["dispersion y [m]"]
disp_p = res_dict["dispersion prime y"]
return self.trackDispersionData(momentum, mass, disp, disp_p, "y")
def trackDispersionData(self, momentum, mass, disp, disp_p, direction="x"):
"""
Returns the tuple ([(position, disp),...],[(position,disp_p),...] ).
The tracking starts from the values specified as the initial parameters.
The possible values for direction parameter "x" or "y".
"""
if (direction.lower() != "x" and direction.lower() != "y"):
orbitFinalize(
"Class orbit.matrix_lattice.MATRIX_Lattice, method trackDispersionData(...): direction should be x or y."
)
#track dispersion
eps_length = 0.00001 # 10^-6 meter
dir_ind = 0
if (direction.lower() == "y"):
dir_ind = 2
track_m = Matrix(3, 3)
track_m.unit()
track_m.set(2, 0, 0.)
track_m.set(2, 1, 0.)
track_m.set(2, 2, 1.)
track_v = PhaseVector(3)
#kinematics coefficient calculation
Etotal = math.sqrt(momentum**2 + mass**2)
beta = momentum / Etotal
gamma = Etotal / mass
Ekin = Etotal - mass
m_coeff = momentum * momentum / (mass + Ekin)
track_v.set(0, disp)
track_v.set(1, disp_p)
track_v.set(2, 1.)
position = 0.
pos_arr = []
disp_arr = []
disp_p_arr = []
#put in array the initial dispersions
#count = 0
pos_arr.append(position)
disp_arr.append(track_v.get(0))
disp_p_arr.append(track_v.get(1))
for matrixNode in self.getNodes():
if (isinstance(matrixNode, BaseMATRIX) == True):
mt = matrixNode.getMatrix()
ind0 = 0 + dir_ind
ind1 = 1 + dir_ind
track_m.set(0, 0, mt.get(ind0, ind0))
track_m.set(0, 1, mt.get(ind0, ind1))
track_m.set(1, 0, mt.get(ind1, ind0))
track_m.set(1, 1, mt.get(ind1, ind1))
track_m.set(0, 2, mt.get(ind0, 5) * m_coeff)
track_m.set(1, 2, mt.get(ind1, 5) * m_coeff)
track_v = track_m.mult(track_v)
position = position + matrixNode.getLength()
#only the main nodes are used, the or-cases deal with markers with zero length
if (isinstance(matrixNode, BaseMATRIX) == True):
pos_arr.append(position)
disp_arr.append(track_v.get(0))
disp_p_arr.append(track_v.get(1))
#pack the resulting tuple
graph_disp_arr = []
graph_disp_p_arr = []
for i in range(len(pos_arr)):
graph_disp_arr.append((pos_arr[i], disp_arr[i]))
graph_disp_p_arr.append((pos_arr[i], disp_p_arr[i]))
return (graph_disp_arr, graph_disp_p_arr)
def getRingOrbit(self, z0):
"""
Returns the tuple ([(position, x] ).
"""
return self.trackOrbit(z0)
def trackOrbit(self, z0):
"""
Returns the tuple ([(position, x),...], [(position, y),...] ).
The tracking starts from the values specified as the initial parameters.
z0 fulfill: z0 = Mz0 with M as one turn matrix
"""
eps_length = 0.00001 # 10^-6 meter
pos_arr = []
orbitX_arr = []
orbitY_arr = []
orbitXP_arr = []
orbitYP_arr = []
position = 0.
pos_arr.append(position)
track_o = Matrix(6, 6)
track_ov = PhaseVector(6)
track_ov.set(0, z0[0])
track_ov.set(1, z0[1])
track_ov.set(2, z0[2])
track_ov.set(3, z0[3])
track_ov.set(4, z0[4])
track_ov.set(5, z0[5])
orbitX_arr.append(track_ov.get(0))
orbitY_arr.append(track_ov.get(2))
orbitXP_arr.append(track_ov.get(1))
orbitYP_arr.append(track_ov.get(3))
position_old = position
for matrixNode in self.getNodes():
if (isinstance(matrixNode, BaseMATRIX) == True):
if (abs(position_old - position) > eps_length):
pos_arr.append(position)
orbitX_arr.append(track_ov.get(0))
orbitY_arr.append(track_ov.get(2))
orbitXP_arr.append(track_ov.get(1))
orbitYP_arr.append(track_ov.get(3))
mt = matrixNode.getMatrix()
for i in range(6):
for j in range(6):
track_o.set(i, j, mt.get(i, j))
track_ov = track_o.mult(track_ov)
for i in range(6):
tmp = track_ov.get(i)
track_ov.set(i, tmp + mt.get(i, 6))
position_old = position
position = position + matrixNode.getLength()
pos_arr.append(position)
orbitX_arr.append(track_ov.get(0))
orbitY_arr.append(track_ov.get(2))
orbitXP_arr.append(track_ov.get(1))
orbitYP_arr.append(track_ov.get(3))
#pack the resulting tuple
graph_orbitX_arr = []
graph_orbitY_arr = []
for i in range(len(pos_arr)):
graph_orbitX_arr.append(
(pos_arr[i], orbitX_arr[i], orbitXP_arr[i]))
graph_orbitY_arr.append(
(pos_arr[i], orbitY_arr[i], orbitYP_arr[i]))
return (graph_orbitX_arr, graph_orbitY_arr)
def getSubLattice(
self,
index_start=-1,
index_stop=-1,
):
"""
It returns the new MATRIX_Lattice with children with indexes
between index_start and index_stop inclusive
"""
return self._getSubLattice(MATRIX_Lattice(), index_start, index_stop)
def trackBunch(self, bunch, paramsDict={}, actionContainer=None):
"""
It tracks the bunch through the lattice.
"""
if (actionContainer == None):
actionContainer = AccActionsContainer("Bunch Tracking")
paramsDict["bunch"] = bunch
def track(paramsDict):
node = paramsDict["node"]
node.track(paramsDict)
actionContainer.addAction(track, AccActionsContainer.BODY)
self.trackActions(actionContainer, paramsDict)
actionContainer.removeAction(track, AccActionsContainer.BODY)
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 _makePolynomial(self): nPoints = len(self.x_y_err_arr) if(nPoints < (self.order+1)): self.order = nPoints - 1 #check if just one of errors is zero infoZeroErr = 1.0 for [x,y,err] in self.x_y_err_arr: infoZeroErr *= err for i in range(nPoints): [x,y,err] = self.x_y_err_arr[i] sigma = 1.0 if(infoZeroErr != 0.): sigma = 1.0/(err*err) self.x_y_err_arr[i][2] = sigma #now make A matrix aMatr = Matrix(nPoints,self.order+1) for i in range(nPoints): for j in range(self.order+1): x = self.x_y_err_arr[i][0] aMatr.set(i,j,math.pow(x,j)) aTCa = Matrix(self.order+1,self.order+1) for i in range(self.order+1): for j in range(self.order+1): a = 0. for k in range(nPoints): sigma = self.x_y_err_arr[k][2] a += aMatr.get(k,i)*sigma*aMatr.get(k,j) aTCa.set(i,j,a) #now the resuting coefficients and errors aTCaI = aTCa.invert() e = aTCaI.mult(aTCa) if(aTCa == None): print "python PolynomialFit: Problem with data." for i in range(nPoints): x = self.x_y_err_arr[i][0] y = self.x_y_err_arr[i][1] err = self.x_y_err_arr[i][2] print " x,y,err = %12.5g %12.5g %12.5g "%(x,y,err) print "Stop." sys.exit(1) coef_arr = [0.]*(self.order+1) err_arr = [0.]*(self.order+1) for i in range(self.order+1): err_arr[i] = math.sqrt(math.fabs(aTCaI.get(i,i))) for i in range(self.order+1): coef_arr[i] = 0. for j in range(self.order+1): for k in range(nPoints): sigma = self.x_y_err_arr[k][2] y = self.x_y_err_arr[k][1] coef_arr[i] += aTCaI.get(i,j)*aMatr.get(k,j)*sigma*y # polinimial coefficients are found self.polynomial.order(self.order) for i in range(len(coef_arr)): self.polynomial.coefficient(i,coef_arr[i]) # now let's calculate errors if(infoZeroErr == 0.): total_sigma = 0. for k in range(nPoints): x = self.x_y_err_arr[k][0] y = self.x_y_err_arr[k][1] total_sigma += (self.polynomial.value(x)-y)**2 total_sigma = math.sqrt(total_sigma/(nPoints-2)) for i in range(len(err_arr)): err_arr[i] *= total_sigma # set the resulting coefficients and errors array self.coef_err_arr = [coef_arr,err_arr]
b_in_has_ms = b_in.hasPartAttr("macrosize")
for i in range(nParticles):
(x, xp, y, yp, z, dE) = (random.random(), random.random(), random.random(),
random.random(), random.random(), random.random())
b_in.addParticle(x, xp, y, yp, z, dE)
if (b_in_has_ms):
b_in.partAttrValue("macrosize", i, 0, random.random())
b_in.addParticle(0., 0., 0., 0., 0., 0.)
#----set up ids
ParticleIdNumber.addParticleIdNumbers(b_in)
b_out = Bunch()
b_in.copyEmptyBunchTo(b_out)
#-----set up initial transport matix, later we have to extract it from bunches
mtrxA_init = Matrix(7, 7)
mtrxA_init.zero()
for ix in range(6):
mtrxA_init.set(ix, 6, 1.0 * ix)
mtrxA_init.set(6, 6, 1.0)
for ix in range(6):
mtrxA_init.set(ix, ix, 1. * (ix + 1))
mtrxA_init.set(0, 1, 1.5 * .001)
mtrxA_init.set(2, 3, 1.6 * .001)
mtrxA_init.set(4, 5, 1.7 * .001)
mtrxA_init.set(1, 0, 1.4 * .001)
mtrxA_init.set(3, 2, 1.3 * .001)
mtrxA_init.set(5, 4, 1.2 * .001)
Ekin = 1.0 # in GeV
b.getSyncParticle().kinEnergy(Ekin)
#define TEAPOT drift
node0 = DriftTEAPOT("drift")
node0.setLength(1.0)
#define TEAPOT quad
node1 = QuadTEAPOT("quad")
node1.setLength(1.0)
node1.addParam("kq",0.5)
matrixGenerator = MatrixGenerator()
#========matrix for drift =====
m = Matrix(6,6)
matrixGenerator.initBunch(b)
node0.trackBunch(b)
matrixGenerator.calculateMatrix(b,m)
print "drift matrix L=",node0.getLength()
printM(m)
#========matrix for quad =====
matrixGenerator.initBunch(b)
node1.trackBunch(b)
matrixGenerator.calculateMatrix(b,m)
print "quad matrix L=",node1.getLength()," kq=",node1.getParam("kq")
printM(m)
#---- now let's prepare array of 4x4 (transport) #---- and 6x6 (Twiss params transport) matrices matrix_arr = [] for ind in range(1,g_func.getSize()): z0 = g_func.x(ind-1) z1 = g_func.x(ind) g0 = g_func.y(ind-1) g1 = g_func.y(ind) z_avg = (z1+z0)/2 g_avg = (g0+g1)/2 g_avg *= charge dL = z1-z0 #print "debug dL=",dL, " y= %12.3f "%y_avg #---- this is a transport 4x4 matrix for coordinates x,x',y,y' trMtrx = Matrix(4,4) k2 = abs(g_avg/Brho) kq = math.sqrt(k2) sin_kL = math.sin(kq*dL) cos_kL = math.cos(kq*dL) sinh_kL = math.sinh(kq*dL) cosh_kL = math.cosh(kq*dL) if(g_avg == 0.): trMtrx.set(0,0,1.) trMtrx.set(1,1,1.) trMtrx.set(2,2,1.) trMtrx.set(3,3,1.) trMtrx.set(0,1,dL) trMtrx.set(2,3,dL) else: if(g_avg > 0.):
