Skip to content

Commit

Permalink
Merge pull request #8 from EuropeanSpallationSource/sync_w_python2_ve…
Browse files Browse the repository at this point in the history 
…rsion

Sync latest changes in the python2 version
  • Loading branch information
@juanfem
juanfem authored Sep 6, 2024
2 parents c10d921 + f92abe9 commit 88ac765
Show file tree
Hide file tree
Showing 32 changed files with 3,245 additions and 2,881 deletions.
2 changes: 1 addition & 1 deletion .pre-commit-config.yaml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@ repos:
hooks:
- id: end-of-file-fixer
- id: trailing-whitespace
exclude: '^tests/examples/SNS_Linac/linac_errors/.*|^tests/examples/AccLattice_Tests/.*|^tests/examples/Optimization/.*'
exclude: '^tests/examples/SNS_Linac/linac_errors/.*|^tests/examples/AccLattice_Tests/.*|^tests/examples/Optimization/.*|^.*\.dat'


- repo: https://github.com/ambv/black
Expand Down
134 changes: 82 additions & 52 deletions py/orbit/py_linac/lattice/LinacAccNodes.py
100644 → 100755
View file
Original file line number Diff line number Diff line change
Expand Up @@ -308,13 +308,15 @@ def __init__(self, name="quad"):
self.setType("linacQuad")

def fringeIN(node, paramsDict):
# B*rho = 3.335640952*momentum [T*m] if momentum in GeV/c
# B*rho = 3.335640952*momentum/charge [T*m] if momentum in GeV/c
usageIN = node.getUsage()
if not usageIN:
return
bunch = paramsDict["bunch"]
charge = bunch.charge()
momentum = bunch.getSyncParticle().momentum()
kq = node.getParam("dB/dr") / (3.335640952 * momentum)
# ---- The charge sign will be accounted for inside tracking module functions.
kq = node.getParam("dB/dr") / bunch.B_Rho()
poleArr = node.getParam("poles")
klArr = node.getParam("kls")
skewArr = node.getParam("skews")
Expand All @@ -329,12 +331,15 @@ def fringeIN(node, paramsDict):
TPB.multpfringeIN(bunch, pole, k, skew)

def fringeOUT(node, paramsDict):
# B*rho = 3.335640952*momentum/charge [T*m] if momentum in GeV/c
usageOUT = node.getUsage()
if not usageOUT:
return
bunch = paramsDict["bunch"]
charge = bunch.charge()
momentum = bunch.getSyncParticle().momentum()
kq = node.getParam("dB/dr") / (3.335640952 * momentum)
# ---- The charge sign will be accounted for inside tracking module functions
kq = node.getParam("dB/dr") / bunch.B_Rho()
poleArr = node.getParam("poles")
klArr = node.getParam("kls")
skewArr = node.getParam("skews")
Expand Down Expand Up @@ -393,27 +398,32 @@ def initialize(self):
for i in range(nParts):
self.setLength(lengthStep, i)
"""
#=============================================
# This is an old TEAPOT-like implementation
# of the Quad slicing.
#=============================================
lengthIN = (self.getLength()/(nParts - 1))/2.0
lengthOUT = (self.getLength()/(nParts - 1))/2.0
lengthStep = lengthIN + lengthOUT
self.setLength(lengthIN,0)
self.setLength(lengthOUT,nParts - 1)
for i in range(nParts-2):
self.setLength(lengthStep,i+1)
"""
#=============================================
# This is an old TEAPOT-like implementation
# of the Quad slicing.
#=============================================
lengthIN = (self.getLength()/(nParts - 1))/2.0
lengthOUT = (self.getLength()/(nParts - 1))/2.0
lengthStep = lengthIN + lengthOUT
self.setLength(lengthIN,0)
self.setLength(lengthOUT,nParts - 1)
for i in range(nParts-2):
self.setLength(lengthStep,i+1)
"""

def track(self, paramsDict):
"""
The Quad Combined Function TEAPOT class implementation
of the AccNode class track(probe) method.
"""
bunch = paramsDict["bunch"]
charge = bunch.charge()
momentum = bunch.getSyncParticle().momentum()
kq = self.getParam("dB/dr") / (3.335640952 * momentum)
# ---- The sign of dB/dr will be delivered to tracking module
# ---- functions as kq.
# ---- The charge sign will be accounted for inside tracking module
# ---- functions.
kq = self.getParam("dB/dr") / bunch.B_Rho()
nParts = self.getnParts()
index = self.getActivePartIndex()
length = self.getLength(index)
Expand Down Expand Up @@ -447,33 +457,33 @@ def track(self, paramsDict):
self.tracking_module.quad2(bunch, step / 4)
self.tracking_module.quad1(bunch, step / 4, kq)
"""
#=============================================
# This is an old TEAPOT-like implementation
# of the Quad tracking.
#=============================================
if(index == 0):
self.tracking_module.quad1(bunch, length, kq)
return
if(index > 0 and index < (nParts-1)):
self.tracking_module.quad2(bunch, length/2.0)
for i in range(len(poleArr)):
pole = poleArr[i]
kl = klArr[i]/(nParts - 1)
skew = skewArr[i]
TPB.multp(bunch,pole,kl,skew)
self.tracking_module.quad2(bunch, length/2.0)
self.tracking_module.quad1(bunch, length, kq)
return
if(index == (nParts-1)):
self.tracking_module.quad2(bunch, length)
for i in range(len(poleArr)):
pole = poleArr[i]
kl = klArr[i]*kq*length/(nParts - 1)
skew = skewArr[i]
TPB.multp(bunch,pole,kl,skew)
self.tracking_module.quad2(bunch, length)
self.tracking_module.quad1(bunch, length, kq)
"""
#=============================================
# This is an old TEAPOT-like implementation
# of the Quad tracking.
#=============================================
if(index == 0):
self.tracking_module.quad1(bunch, length, kq)
return
if(index > 0 and index < (nParts-1)):
self.tracking_module.quad2(bunch, length/2.0)
for i in range(len(poleArr)):
pole = poleArr[i]
kl = klArr[i]/(nParts - 1)
skew = skewArr[i]
TPB.multp(bunch,pole,kl,skew)
self.tracking_module.quad2(bunch, length/2.0)
self.tracking_module.quad1(bunch, length, kq)
return
if(index == (nParts-1)):
self.tracking_module.quad2(bunch, length)
for i in range(len(poleArr)):
pole = poleArr[i]
kl = klArr[i]*kq*length/(nParts - 1)
skew = skewArr[i]
TPB.multp(bunch,pole,kl,skew)
self.tracking_module.quad2(bunch, length)
self.tracking_module.quad1(bunch, length, kq)
"""
return

def getTotalField(self, z):
Expand Down Expand Up @@ -537,7 +547,17 @@ def fringeIN(node, paramsDict):
TPB.multpfringeIN(bunch, pole, kl, skew)
frinout = 1
TPB.wedgerotate(bunch, e, frinout)
TPB.wedgebendCF(bunch, e, inout, rho, len(poleArr), poleArr, klArr, skewArr, nParts - 1)
TPB.wedgebendCF(
bunch,
e,
inout,
rho,
len(poleArr),
poleArr,
klArr,
skewArr,
nParts - 1,
)
else:
if usageIN:
TPB.bendfringeIN(bunch, rho)
Expand All @@ -560,7 +580,17 @@ def fringeOUT(node, paramsDict):
nParts = paramsDict["parentNode"].getnParts()
if e != 0.0:
inout = 1
TPB.wedgebendCF(bunch, e, inout, rho, len(poleArr), poleArr, klArr, skewArr, nParts - 1)
TPB.wedgebendCF(
bunch,
e,
inout,
rho,
len(poleArr),
poleArr,
klArr,
skewArr,
nParts - 1,
)
if usageOUT:
frinout = 0
TPB.wedgerotate(bunch, -e, frinout)
Expand Down Expand Up @@ -697,11 +727,11 @@ def track(self, paramsDict):
length = self.getParam("effLength") / nParts
field = self.getParam("B")
bunch = paramsDict["bunch"]
charge = bunch.charge()
syncPart = bunch.getSyncParticle()
momentum = syncPart.momentum()
# dp/p = Q*c*B*L/p p in GeV/c c = 2.99792*10^8/10^9
# Q is used inside kick-method
kick = field * length * 0.299792 / momentum
kick = -field * charge * length * 0.299792 / momentum
self.tracking_module.kick(bunch, kick, 0.0, 0.0)


Expand Down Expand Up @@ -742,11 +772,11 @@ def track(self, paramsDict):
length = self.getParam("effLength") / nParts
field = self.getParam("B")
bunch = paramsDict["bunch"]
charge = bunch.charge()
syncPart = bunch.getSyncParticle()
momentum = syncPart.momentum()
# dp/p = Q*c*B*L/p p in GeV/c, c = 2.99792*10^8/10^9
# Q is used inside kick-method
kick = field * length * 0.299792 / momentum
kick = field * charge * length * 0.299792 / momentum
self.tracking_module.kick(bunch, 0, kick, 0.0)


Expand Down Expand Up @@ -803,6 +833,7 @@ def track(self, paramsDict):
The Thick Kick class implementation of the AccNode class track(probe) method.
"""
bunch = paramsDict["bunch"]
charge = bunch.charge()
momentum = bunch.getSyncParticle().momentum()
Bx = self.getParam("Bx")
By = self.getParam("By")
Expand All @@ -812,9 +843,8 @@ def track(self, paramsDict):
# print "debug name =",self.getName()," Bx=",Bx," By=",By," L=",self.getLength(index)," index=",index
# ==========================================
# dp/p = Q*c*B*L/p p in GeV/c, c = 2.99792*10^8/10^9
# Q is used inside kick-method
kickY = Bx * length * 0.299792 / momentum
kickX = By * length * 0.299792 / momentum
kickY = +Bx * charge * length * 0.299792 / momentum
kickX = -By * charge * length * 0.299792 / momentum
self.tracking_module.drift(bunch, length / 2.0)
self.tracking_module.kick(bunch, kickX, kickY, 0.0)
self.tracking_module.drift(bunch, length / 2.0)
Expand Down
65 changes: 43 additions & 22 deletions py/orbit/py_linac/lattice/LinacFieldOverlappingNodes.py
100644 → 100755
View file
Original file line number Diff line number Diff line change
Expand Up @@ -288,6 +288,7 @@ def track(self, paramsDict):
index = self.getActivePartIndex()
part_length = self.getLength(index)
bunch = paramsDict["bunch"]
charge = bunch.charge()
syncPart = bunch.getSyncParticle()
eKin_in = syncPart.kinEnergy()
momentum = syncPart.momentum()
Expand All @@ -300,7 +301,10 @@ def track(self, paramsDict):
phase_shift = rfCavity.getPhase() - rfCavity.getDesignPhase()
phase = rfCavity.getFirstGapEtnrancePhase() + phase_shift
# ----------------------------------------
phase = math.fmod(frequency * (arrival_time - designArrivalTime) * 2.0 * math.pi + phase, 2.0 * math.pi)
phase = math.fmod(
frequency * (arrival_time - designArrivalTime) * 2.0 * math.pi + phase,
2.0 * math.pi,
)
if index == 0:
self.part_pos = self.z_min
self.gap_phase_vs_z_arr = [
Expand All @@ -314,21 +318,20 @@ def track(self, paramsDict):
Ep = self.getEzFiledInternal(zp, rfCavity, E0L, rf_ampl)
# ------- track through a quad
G = self.getTotalField((zm + z0) / 2)
GP = 0.0
dB_dz = 0.0
if self.useLongField == True:
GP = self.getTotalFieldDerivative((zm + z0) / 2)
dB_dz = self.getTotalFieldDerivative((zm + z0) / 2)
if abs(G) != 0.0:
kq = G / (3.335640952 * momentum)
kq = G / bunch.B_Rho()
# ------- track through a quad
step = part_length / 2
self.tracking_module.quad1(bunch, step / 4.0, kq)
self.tracking_module.quad2(bunch, step / 2.0)
self.tracking_module.quad1(bunch, step / 2.0, kq)
self.tracking_module.quad2(bunch, step / 2.0)
self.tracking_module.quad1(bunch, step / 4.0, kq)
if abs(GP) != 0.0:
kqP = GP / (3.335640952 * momentum)
self.tracking_module.quad3(bunch, step, kqP)
if abs(dB_dz) != 0.0:
self.tracking_module.quad3(bunch, step, dB_dz)
else:
self.tracking_module.drift(bunch, part_length / 2)
self.part_pos += part_length / 2
Expand All @@ -344,23 +347,30 @@ def track(self, paramsDict):
# s += " dE= %9.6f "%((eKin_out-eKin_in)*1000.)
# print s
# ---- this part is the debugging ---STOP---
self.cppGapModel.trackBunch(bunch, part_length / 2, Em, E0, Ep, frequency, phase + delta_phase + modePhase)
self.cppGapModel.trackBunch(
bunch,
part_length / 2,
Em,
E0,
Ep,
frequency,
phase + delta_phase + modePhase,
)
# ------- track through a quad
G = self.getTotalField((z0 + zp) / 2)
GP = 0.0
dB_dz = 0.0
if self.useLongField == True:
GP = self.getTotalFieldDerivative((z0 + zp) / 2)
dB_dz = self.getTotalFieldDerivative((z0 + zp) / 2)
if abs(G) != 0.0:
kq = G / (3.335640952 * momentum)
kq = G / bunch.B_Rho()
step = part_length / 2
self.tracking_module.quad1(bunch, step / 4.0, kq)
self.tracking_module.quad2(bunch, step / 2.0)
self.tracking_module.quad1(bunch, step / 2.0, kq)
self.tracking_module.quad2(bunch, step / 2.0)
self.tracking_module.quad1(bunch, step / 4.0, kq)
if abs(GP) != 0.0:
kqP = GP / (3.335640952 * momentum)
self.tracking_module.quad3(bunch, step, kqP)
if abs(dB_dz) != 0.0:
self.tracking_module.quad3(bunch, step, dB_dz)
else:
self.tracking_module.drift(bunch, part_length / 2)
# ---- advance the particle position
Expand Down Expand Up @@ -436,7 +446,10 @@ def trackDesign(self, paramsDict):
else:
first_gap_arr_time = rfCavity.getDesignArrivalTime()
# print "debug name=",self.getName()," delta_phase=",frequency*(arrival_time - first_gap_arr_time)*360.0," phase=",phase*180/math.pi
phase = math.fmod(frequency * (arrival_time - first_gap_arr_time) * 2.0 * math.pi + phase, 2.0 * math.pi)
phase = math.fmod(
frequency * (arrival_time - first_gap_arr_time) * 2.0 * math.pi + phase,
2.0 * math.pi,
)
# print "debug design name=",self.getName()," arr_time=",arrival_time," phase=",phase*180./math.pi," E0TL=",E0TL*1.0e+3," freq=",frequency
if index == 0:
self.part_pos = self.z_min
Expand Down Expand Up @@ -464,7 +477,15 @@ def trackDesign(self, paramsDict):
# s += " dE= %9.6f "%((eKin_out-eKin_in)*1000.)
# print s
# ---- this part is the debugging ---STOP---
self.cppGapModel.trackBunch(bunch, part_length / 2, Em, E0, Ep, frequency, phase + delta_phase + modePhase)
self.cppGapModel.trackBunch(
bunch,
part_length / 2,
Em,
E0,
Ep,
frequency,
phase + delta_phase + modePhase,
)
self.tracking_module.drift(bunch, part_length / 2)
# ---- advance the particle position
self.part_pos += part_length / 2
Expand Down Expand Up @@ -627,16 +648,17 @@ def track(self, paramsDict):
if index == 0:
self.z_value = -self.getLength() / 2
bunch = paramsDict["bunch"]
charge = bunch.charge()
momentum = bunch.getSyncParticle().momentum()
n_steps = int(length / self.z_step) + 1
z_step = length / n_steps
for z_ind in range(n_steps):
z = self.z_value + z_step * (z_ind + 0.5)
G = self.getTotalField(z)
GP = 0.0
dB_dz = 0.0
if self.useLongField == True:
GP = self.getTotalFieldDerivative(z)
kq = G / (3.335640952 * momentum)
dB_dz = self.getTotalFieldDerivative(z)
kq = G / bunch.B_Rho()
if abs(kq) == 0.0:
self.tracking_module.drift(bunch, z_step)
continue
Expand All @@ -646,9 +668,8 @@ def track(self, paramsDict):
self.tracking_module.quad1(bunch, z_step / 2.0, kq)
self.tracking_module.quad2(bunch, z_step / 2.0)
self.tracking_module.quad1(bunch, z_step / 4.0, kq)
if abs(GP) != 0.0:
kqP = GP / (3.335640952 * momentum)
self.tracking_module.quad3(bunch, z_step, kqP)
if abs(dB_dz) != 0.0:
self.tracking_module.quad3(bunch, z_step, dB_dz)
self.z_value += length

def getTotalField(self, z_from_center):
Expand Down
Loading

0 comments on commit 88ac765

Please sign in to comment.