Skip to content

Commit 0014135

Browse files
giadarolgiadarol
committed
Reformatted example
1 parent 4ef3b30 commit 0014135

File tree

1 file changed

+191
-114
lines changed

1 file changed

+191
-114
lines changed

‎examples/pyheadtail_interface/002_headtailInstability.py‎

Lines changed: 191 additions & 114 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,8 @@
11
import time
22
import numpy as np
33
from scipy import constants
4-
protonMass = constants.value('proton mass energy equivalent in MeV')*1E6
4+
5+
protonMass = constants.value("proton mass energy equivalent in MeV") * 1e6
56
from scipy.stats import linregress
67
from scipy.signal import hilbert
78
from matplotlib import pyplot as plt
@@ -20,44 +21,46 @@
2021

2122
context = xo.ContextCpu(omp_num_threads=1)
2223

23-
nTurn = 3000 #int(1E4)
24-
bunch_intensity = 1.8E11
25-
n_macroparticles = int(1e5) #int(1E4)
26-
energy = 7E3 # [GeV]
27-
gamma = energy*1E9/protonMass
28-
betar = np.sqrt(1-1/gamma**2)
29-
normemit = 1.8E-6
24+
nTurn = 3000 # int(1E4)
25+
bunch_intensity = 1.8e11
26+
n_macroparticles = int(1e5) # int(1E4)
27+
energy = 7e3 # [GeV]
28+
gamma = energy * 1e9 / protonMass
29+
betar = np.sqrt(1 - 1 / gamma ** 2)
30+
normemit = 1.8e-6
3031
beta_x = 68.9
3132
beta_y = 70.34
3233
Q_x = 0.31
3334
Q_y = 0.32
34-
chroma = -5. #10.0
35-
sigma_4t = 1.2E-9
36-
sigma_z = sigma_4t/4.0*constants.c
35+
chroma = -5.0 # 10.0
36+
sigma_4t = 1.2e-9
37+
sigma_z = sigma_4t / 4.0 * constants.c
3738
momentumCompaction = 3.483575072011584e-04
38-
eta = momentumCompaction-1.0/gamma**2
39-
voltage = 12.0E6
39+
eta = momentumCompaction - 1.0 / gamma ** 2
40+
voltage = 12.0e6
4041
h = 35640
41-
p0=constants.m_p*betar*gamma*constants.c
42-
Q_s=np.sqrt(constants.e*voltage*eta*h/(2*np.pi*betar*constants.c*p0))
42+
p0 = constants.m_p * betar * gamma * constants.c
43+
Q_s = np.sqrt(constants.e * voltage * eta * h / (2 * np.pi * betar * constants.c * p0))
4344
circumference = 26658.883199999
44-
averageRadius = circumference/(2*np.pi)
45-
sigma_delta = Q_s*sigma_z/(averageRadius*eta)
46-
beta_s = sigma_z/sigma_delta
47-
emit_s = 4*np.pi*sigma_z*sigma_delta*p0/constants.e # eVs for PyHEADTAIL
45+
averageRadius = circumference / (2 * np.pi)
46+
sigma_delta = Q_s * sigma_z / (averageRadius * eta)
47+
beta_s = sigma_z / sigma_delta
48+
emit_s = 4 * np.pi * sigma_z * sigma_delta * p0 / constants.e # eVs for PyHEADTAIL
4849

4950
n_slices_wakes = 200
5051
limit_z = 3 * sigma_z
51-
wakefile = 'wakes/wakeforhdtl_PyZbase_Allthemachine_7000GeV_B1_2021_TeleIndex1_wake.dat'
52+
wakefile = "wakes/wakeforhdtl_PyZbase_Allthemachine_7000GeV_B1_2021_TeleIndex1_wake.dat"
5253
slicer_for_wakefields = UniformBinSlicer(n_slices_wakes, z_cuts=(-limit_z, limit_z))
53-
waketable = WakeTable(wakefile, ['time', 'dipole_x', 'dipole_y', 'quadrupole_x', 'quadrupole_y'])
54+
waketable = WakeTable(
55+
wakefile, ["time", "dipole_x", "dipole_y", "quadrupole_x", "quadrupole_y"]
56+
)
5457
wake_field = WakeField(slicer_for_wakefields, waketable)
5558

56-
damping_time = 100000000 #33.
59+
damping_time = 100000000 # 33.
5760
damper = TransverseDamper(dampingrate_x=damping_time, dampingrate_y=damping_time)
5861
i_oct = 0.0000001
59-
detx_x = 1.4E5*i_oct/550.0 # from PTC with ATS optics, telescopic factor 1.0
60-
detx_y = -1.0E5*i_oct/550.0
62+
detx_x = 1.4e5 * i_oct / 550.0 # from PTC with ATS optics, telescopic factor 1.0
63+
detx_y = -1.0e5 * i_oct / 550.0
6164

6265
# expected octupole threshold with damper is 273A according to https://indico.cern.ch/event/902528/contributions/3798807/attachments/2010534/3359300/20200327_RunIII_stability_NMounet.pdf
6366
# expected growth rate with damper but without octupole is ~0.3 [$10^{-4}$/turn] (also according to Nicolas' presentation)
@@ -67,10 +70,21 @@
6770
########### PyHEADTAIL part ##############
6871

6972
particles = generate_Gaussian6DTwiss(
70-
macroparticlenumber = n_macroparticles, intensity = bunch_intensity, charge = constants.e, mass = constants.m_p,
71-
circumference = circumference, gamma = gamma,
72-
alpha_x = 0.0, alpha_y = 0.0, beta_x = beta_x, beta_y = beta_y, beta_z = beta_s,
73-
epsn_x = normemit, epsn_y = normemit, epsn_z=emit_s)
73+
macroparticlenumber=n_macroparticles,
74+
intensity=bunch_intensity,
75+
charge=constants.e,
76+
mass=constants.m_p,
77+
circumference=circumference,
78+
gamma=gamma,
79+
alpha_x=0.0,
80+
alpha_y=0.0,
81+
beta_x=beta_x,
82+
beta_y=beta_y,
83+
beta_z=beta_s,
84+
epsn_x=normemit,
85+
epsn_y=normemit,
86+
epsn_z=emit_s,
87+
)
7488

7589
x0 = np.copy(particles.x)
7690
px0 = np.copy(particles.xp)
@@ -79,22 +93,49 @@
7993
zeta0 = np.copy(particles.z)
8094
delta0 = np.copy(particles.dp)
8195

82-
print('PyHt size comp x',particles.sigma_x(),np.sqrt(normemit*beta_x/gamma/betar))
83-
print('PyHt size comp y',particles.sigma_y(),np.sqrt(normemit*beta_y/gamma/betar))
84-
print('PyHt size comp z',particles.sigma_z(),sigma_z)
85-
print('PyHt size comp delta',particles.sigma_dp(),sigma_delta)
96+
print(
97+
"PyHt size comp x", particles.sigma_x(), np.sqrt(normemit * beta_x / gamma / betar)
98+
)
99+
print(
100+
"PyHt size comp y", particles.sigma_y(), np.sqrt(normemit * beta_y / gamma / betar)
101+
)
102+
print("PyHt size comp z", particles.sigma_z(), sigma_z)
103+
print("PyHt size comp delta", particles.sigma_dp(), sigma_delta)
86104

87-
chromatic_detuner = ChromaticitySegment(dQp_x = chroma,dQp_y = 0.0)
88-
transverse_detuner = AmplitudeDetuningSegment(dapp_x=detx_x*p0, dapp_y=detx_x*p0, dapp_xy=detx_y*p0, dapp_yx=detx_y*p0, alpha_x=0.0, beta_x=beta_x,
89-
alpha_y=0.0, beta_y=beta_y)
105+
chromatic_detuner = ChromaticitySegment(dQp_x=chroma, dQp_y=0.0)
106+
transverse_detuner = AmplitudeDetuningSegment(
107+
dapp_x=detx_x * p0,
108+
dapp_y=detx_x * p0,
109+
dapp_xy=detx_y * p0,
110+
dapp_yx=detx_y * p0,
111+
alpha_x=0.0,
112+
beta_x=beta_x,
113+
alpha_y=0.0,
114+
beta_y=beta_y,
115+
)
90116
arc_transverse = TransverseSegmentMap(
91-
alpha_x_s0 = 0.0, beta_x_s0 = beta_x, D_x_s0 = 0.0, alpha_x_s1 = 0.0, beta_x_s1 = beta_x, D_x_s1 = 0.0,
92-
alpha_y_s0 = 0.0, beta_y_s0 = beta_y, D_y_s0 = 0.0, alpha_y_s1 = 0.0, beta_y_s1 = beta_y, D_y_s1 = 0.0,
93-
dQ_x = Q_x, dQ_y = Q_y,segment_detuners=[chromatic_detuner,transverse_detuner])
94-
arc_longitudinal = LinearMap(alpha_array=[momentumCompaction], circumference=circumference, Q_s=Q_s)
117+
alpha_x_s0=0.0,
118+
beta_x_s0=beta_x,
119+
D_x_s0=0.0,
120+
alpha_x_s1=0.0,
121+
beta_x_s1=beta_x,
122+
D_x_s1=0.0,
123+
alpha_y_s0=0.0,
124+
beta_y_s0=beta_y,
125+
D_y_s0=0.0,
126+
alpha_y_s1=0.0,
127+
beta_y_s1=beta_y,
128+
D_y_s1=0.0,
129+
dQ_x=Q_x,
130+
dQ_y=Q_y,
131+
segment_detuners=[chromatic_detuner, transverse_detuner],
132+
)
133+
arc_longitudinal = LinearMap(
134+
alpha_array=[momentumCompaction], circumference=circumference, Q_s=Q_s
135+
)
95136

96137
turns = np.arange(nTurn)
97-
x = np.zeros(nTurn,dtype=float)
138+
x = np.zeros(nTurn, dtype=float)
98139
for turn in range(nTurn):
99140
if turn == checkTurn:
100141
x1 = np.copy(particles.x)
@@ -113,81 +154,116 @@
113154
damper.track(particles)
114155
time3 = time.time()
115156
x[turn] = np.average(particles.x)
116-
if turn%1000 == 0:
117-
print(f'PyHt - turn {turn}: time for arc {time1-time0}s, for wake {time2-time1}s, for damper {time3-time2}s')
118-
x /= np.sqrt(normemit*beta_x/gamma/betar)
157+
if turn % 1000 == 0:
158+
print(
159+
f"PyHt - turn {turn}: time for arc {time1-time0}s, for wake {time2-time1}s, for damper {time3-time2}s"
160+
)
161+
x /= np.sqrt(normemit * beta_x / gamma / betar)
119162
plt.figure(0)
120-
plt.plot(turns,x,label=f'{i_oct}A')
163+
plt.plot(turns, x, label=f"{i_oct}A")
121164
iMin = 1000
122-
iMax = nTurn-1000
165+
iMax = nTurn - 1000
123166
if iMin >= iMax:
124167
iMin = 0
125168
iMax = nTurn
126169
ampl = np.abs(hilbert(x))
127-
b,a,r,p,stderr = linregress(turns[iMin:iMax],np.log(ampl[iMin:iMax]))
128-
plt.plot(turns,np.exp(a+b*turns),'--k',label=f'{1/b:.3E} turns')
129-
print(f'Growth rate {b*1E4} [$10^{-4}$/turn]')
130-
plt.title('PyHEADTAIL')
131-
plt.legend(loc='upper left')
132-
plt.xlabel('Turn')
133-
plt.ylabel('x [$\sigma_x$]')
170+
b, a, r, p, stderr = linregress(turns[iMin:iMax], np.log(ampl[iMin:iMax]))
171+
plt.plot(turns, np.exp(a + b * turns), "--k", label=f"{1/b:.3E} turns")
172+
print(f"Growth rate {b*1E4} [$10^{-4}$/turn]")
173+
plt.title("PyHEADTAIL")
174+
plt.legend(loc="upper left")
175+
plt.xlabel("Turn")
176+
plt.ylabel("x [$\sigma_x$]")
134177

135178
############ xsuite-PyHEADTAIL part (the WakeField instance is shared) ########################
136179

137-
if True: #Use the initial coordinates generated by PyHEADTAIL
138-
particles = PyHtXtParticles(circumference=circumference,particlenumber_per_mp=bunch_intensity/n_macroparticles,
139-
_context=context,
140-
q0 = 1,
141-
mass0 = protonMass,
142-
gamma0 = gamma,
143-
x=x0,
144-
px=px0,
145-
y=y0,
146-
py=py0,
147-
zeta=zeta0,
148-
delta=delta0
149-
)
150-
else: #Use new random coordinates with the same distribution
180+
if True: # Use the initial coordinates generated by PyHEADTAIL
181+
particles = PyHtXtParticles(
182+
circumference=circumference,
183+
particlenumber_per_mp=bunch_intensity / n_macroparticles,
184+
_context=context,
185+
q0=1,
186+
mass0=protonMass,
187+
gamma0=gamma,
188+
x=x0,
189+
px=px0,
190+
y=y0,
191+
py=py0,
192+
zeta=zeta0,
193+
delta=delta0,
194+
)
195+
else: # Use new random coordinates with the same distribution
151196
checkTurn = -1
152-
particles = PyHtXtParticles(circumference=circumference,particlenumber_per_mp=bunch_intensity/n_macroparticles,
153-
_context=context,
154-
q0 = 1,
155-
mass0 = protonMass,
156-
gamma0 = gamma,
157-
x=np.sqrt(normemit*beta_x/gamma/betar)*np.random.randn(n_macroparticles),
158-
px=np.sqrt(normemit/beta_x/gamma/betar)*np.random.randn(n_macroparticles),
159-
y=np.sqrt(normemit*beta_y/gamma/betar)*np.random.randn(n_macroparticles),
160-
py=np.sqrt(normemit/beta_y/gamma/betar)*np.random.randn(n_macroparticles),
161-
zeta=sigma_z*np.random.randn(n_macroparticles),
162-
delta=sigma_delta*np.random.randn(n_macroparticles)
163-
)
164-
165-
print('PyHtXt size comp x',particles.sigma_x(),np.sqrt(normemit*beta_x/gamma/betar))
166-
print('PyHtXt size comp y',particles.sigma_y(),np.sqrt(normemit*beta_y/gamma/betar))
167-
print('PyHtXt size comp z',particles.sigma_z(),sigma_z)
168-
print('PyHtXt size comp delta',particles.sigma_dp(),sigma_delta)
169-
170-
arc = xt.LinearTransferMatrix(_context=context,
171-
alpha_x_0 = 0.0, beta_x_0 = beta_x, disp_x_0 = 0.0,
172-
alpha_x_1 = 0.0, beta_x_1 = beta_x, disp_x_1 = 0.0,
173-
alpha_y_0 = 0.0, beta_y_0 = beta_y, disp_y_0 = 0.0,
174-
alpha_y_1 = 0.0, beta_y_1 = beta_y, disp_y_1 = 0.0,
175-
Q_x = Q_x, Q_y = Q_y,
176-
beta_s = beta_s, Q_s = -Q_s,
177-
chroma_x = chroma,chroma_y = 0.0,
178-
detx_x = detx_x,detx_y = detx_y,dety_y = detx_x, dety_x = detx_y,
179-
energy_ref_increment=0.0,energy_increment=0)
197+
particles = PyHtXtParticles(
198+
circumference=circumference,
199+
particlenumber_per_mp=bunch_intensity / n_macroparticles,
200+
_context=context,
201+
q0=1,
202+
mass0=protonMass,
203+
gamma0=gamma,
204+
x=np.sqrt(normemit * beta_x / gamma / betar)
205+
* np.random.randn(n_macroparticles),
206+
px=np.sqrt(normemit / beta_x / gamma / betar)
207+
* np.random.randn(n_macroparticles),
208+
y=np.sqrt(normemit * beta_y / gamma / betar)
209+
* np.random.randn(n_macroparticles),
210+
py=np.sqrt(normemit / beta_y / gamma / betar)
211+
* np.random.randn(n_macroparticles),
212+
zeta=sigma_z * np.random.randn(n_macroparticles),
213+
delta=sigma_delta * np.random.randn(n_macroparticles),
214+
)
215+
216+
print(
217+
"PyHtXt size comp x",
218+
particles.sigma_x(),
219+
np.sqrt(normemit * beta_x / gamma / betar),
220+
)
221+
print(
222+
"PyHtXt size comp y",
223+
particles.sigma_y(),
224+
np.sqrt(normemit * beta_y / gamma / betar),
225+
)
226+
print("PyHtXt size comp z", particles.sigma_z(), sigma_z)
227+
print("PyHtXt size comp delta", particles.sigma_dp(), sigma_delta)
228+
229+
arc = xt.LinearTransferMatrix(
230+
_context=context,
231+
alpha_x_0=0.0,
232+
beta_x_0=beta_x,
233+
disp_x_0=0.0,
234+
alpha_x_1=0.0,
235+
beta_x_1=beta_x,
236+
disp_x_1=0.0,
237+
alpha_y_0=0.0,
238+
beta_y_0=beta_y,
239+
disp_y_0=0.0,
240+
alpha_y_1=0.0,
241+
beta_y_1=beta_y,
242+
disp_y_1=0.0,
243+
Q_x=Q_x,
244+
Q_y=Q_y,
245+
beta_s=beta_s,
246+
Q_s=-Q_s,
247+
chroma_x=chroma,
248+
chroma_y=0.0,
249+
detx_x=detx_x,
250+
detx_y=detx_y,
251+
dety_y=detx_x,
252+
dety_x=detx_y,
253+
energy_ref_increment=0.0,
254+
energy_increment=0,
255+
)
180256

181257
turns = np.arange(nTurn)
182-
x = np.zeros(nTurn,dtype=float)
258+
x = np.zeros(nTurn, dtype=float)
183259
for turn in range(nTurn):
184260
if turn == checkTurn:
185-
print('x',particles.x-x1)
186-
print('px',particles.px-px1)
187-
print('y',particles.y-y1)
188-
print('py',particles.py-py1)
189-
print('z',particles.zeta-zeta1)
190-
print('delta',particles.delta-delta1)
261+
print("x", particles.x - x1)
262+
print("px", particles.px - px1)
263+
print("y", particles.y - y1)
264+
print("py", particles.py - py1)
265+
print("z", particles.zeta - zeta1)
266+
print("delta", particles.delta - delta1)
191267

192268
time0 = time.time()
193269
arc.track(particles)
@@ -197,24 +273,25 @@
197273
damper.track(particles)
198274
time3 = time.time()
199275
x[turn] = np.average(particles.x)
200-
if turn%1000 == 0:
201-
print(f'PyHtXt - turn {turn}: time for arc {time1-time0}s, for wake {time2-time1}s, for damper {time3-time2}s')
202-
x /= np.sqrt(normemit*beta_x/gamma/betar)
276+
if turn % 1000 == 0:
277+
print(
278+
f"PyHtXt - turn {turn}: time for arc {time1-time0}s, for wake {time2-time1}s, for damper {time3-time2}s"
279+
)
280+
x /= np.sqrt(normemit * beta_x / gamma / betar)
203281
plt.figure(1)
204-
plt.plot(turns,x,label=f'{i_oct}A')
282+
plt.plot(turns, x, label=f"{i_oct}A")
205283
iMin = 1000
206-
iMax = nTurn-1000
284+
iMax = nTurn - 1000
207285
if iMin >= iMax:
208286
iMin = 0
209287
iMax = nTurn
210288
ampl = np.abs(hilbert(x))
211-
b,a,r,p,stderr = linregress(turns[iMin:iMax],np.log(ampl[iMin:iMax]))
212-
plt.plot(turns,np.exp(a+b*turns),'--k',label=f'{1/b:.3E} turns')
213-
print(f'Growth rate {b*1E4} [$10^{-4}$/turn]')
214-
plt.title('xsuite-PyHEADTAIL')
215-
plt.legend(loc='upper left')
216-
plt.xlabel('Turn')
217-
plt.ylabel('x [$\sigma_x$]')
289+
b, a, r, p, stderr = linregress(turns[iMin:iMax], np.log(ampl[iMin:iMax]))
290+
plt.plot(turns, np.exp(a + b * turns), "--k", label=f"{1/b:.3E} turns")
291+
print(f"Growth rate {b*1E4} [$10^{-4}$/turn]")
292+
plt.title("xsuite-PyHEADTAIL")
293+
plt.legend(loc="upper left")
294+
plt.xlabel("Turn")
295+
plt.ylabel("x [$\sigma_x$]")
218296

219297
plt.show()
220-

0 commit comments

Comments
 (0)