PyHEADTAIL Quick Start Tutorial¶
This notebook explains how to set up a basic tracking simulation with the macro-particle code PyHEADTAIL (https://github.com/PyCOMPLETE/PyHEADTAIL) in Python 3.
Enjoy!! Adrian Oeftiger, 2022
import numpy as np
np.random.seed(42)
import matplotlib.pyplot as plt
%matplotlib inline
import seaborn as sbs
sbs.set_context('talk')
from scipy.constants import c, e, m_p
If PyHEADTAIL is not available yet, you can install PyHEADTAIL to your user directory via
$ pip install --user PyHEADTAIL
You can do this here in the notebook by typing
!pip install --user PyHEADTAIL.
So let's go and import it!
import PyHEADTAIL
To start with an already assembled one turn map, we will use the git repository PyCERNmachines where some basic machine setups for CERN synchrotrons are stored. Here we will use the CERN Proton Synchrotron at injection with LHC beam type set-up. If you do not have PyCERNmachines installed yet, you can get it from git (again prepend an exclamation mark "!" to start it from a notebook cell here):
$ git clone https://github.com/PyCOMPLETE/PyCERNmachines
# only do this once if you did not download it yet.
!git clone https://github.com/PyCOMPLETE/PyCERNmachines
!cd PyCERNmachines && echo && git describe
from PyCERNmachines.CERNmachines import PS
We want to set up the PS at injection plateau on fundamental harmonic $h=7$, hence the keyword machine_configuration='LHCbeam_h7'. One can initiate the PS with either linear or non-linear synchrotron motion via the keyword longitudinal_focusing.
The transverse bare tunes are set to $Q_x=6.23$ and $Q_y=6.24$.
Ekin = 1.4e9
gamma = 1 + Ekin * e / (m_p * c**2)
beta = np.sqrt(1 - gamma**-2)
Q_x = 6.23
Q_y = 6.24
machine = PS(n_segments=1, gamma=gamma, machine_configuration='LHCbeam_h7',
longitudinal_focusing='non-linear', Q_x=Q_x, Q_y=Q_y)
Furthermore, we only want to use the fundamental harmonic, while the PyCERNmachines.PS is by default launched with a second harmonic on $h=0$. Therefore we remove the second kick from the RFSystems internal list. Also, we are below transition, hence we want a $\pi$ phase offset for the fundamental.
machine.longitudinal_map.pop_kick(1)
machine.longitudinal_map.phi_offsets[0] += np.pi
Now let's generate a macro-particle distribution describing an LHC-type bunch. The machine offers two convenience functions:
generate_6D_Gaussian_bunch: a true 6D Gaussian distribution. In the longitudinal plane, the distribution is cut along the separatrix defined by theRFBucket.generate_6D_Gaussian_bunch_matched: a 4D transverse Gaussian distribution with a matched thermal (Gaussian-like) distribution in the longitudinal phase space.
n_macroparticles = 1000
intensity = 1.6e12 # in protons per bunch
epsn_x = epsn_y = 2.5e-6 # in [m rad]
sigma_z = 185e-9 / 4 * beta * c
bunch = machine.generate_6D_Gaussian_bunch_matched(
n_macroparticles=n_macroparticles, intensity=intensity,
epsn_x=epsn_x, epsn_y=epsn_y, sigma_z=sigma_z)
Now we have a bunch with a matched longitudinal distribution. We can plot the longitudinal phase space with the Hamiltonian contours of the RFBucket.
rfsystems = machine.longitudinal_map
rfbucket = rfsystems.get_bucket(bunch)
zz = np.linspace(*rfbucket.interval, num=1000)
dp_max = rfbucket.separatrix(0)
dpp = np.linspace(-dp_max*1.1, dp_max*1.1, num=100)
ZZ, DPP = np.meshgrid(zz, dpp)
HH = rfbucket.hamiltonian(ZZ, DPP)
plt.scatter(bunch.z, bunch.dp, marker='.', color='black')
plt.contour(ZZ, DPP, HH, 15, cmap=plt.get_cmap('coolwarm'))
plt.xlim(rfbucket.interval)
plt.ylim(-dp_max, dp_max)
plt.xlabel('longitudinal position $z$ [m]')
plt.ylabel('momentum deviation $\delta$ [1]');
The one turn map of the machine now includes 1 betatron tracking segment TransverseSegmentMap and the synchrotron tracking element RFSystems:
machine.one_turn_map
We are now ready to track the beam, simulating just the betatron motion. We will record the incoherent motion of the particles to show that they all move with the bare machine tune.
n_turns = 512
x_rec = np.empty((n_turns, bunch.macroparticlenumber), dtype=np.float64)
y_rec = np.empty((n_turns, bunch.macroparticlenumber), dtype=np.float64)
For the tracking, we execute the loop over all turns and over the machine transverse elements, which are stored in machine.transverse_map.
for i in range(n_turns):
for m in machine.transverse_map:
x_rec[i, :] = bunch.x
y_rec[i, :] = bunch.y
m.track(bunch)
Let's plot one sample particle motion:
plt.plot(x_rec[:, 0], label='x position')
plt.plot(y_rec[:, 0], label='y position')
plt.legend()
plt.xlim(500, 512)
plt.xlabel('turns')
plt.ylabel('transverse position $x,y$ [m]');
The incoherent tune spectra only show the bare machine tune peak:
freqs = 6 + np.fft.rfftfreq(n_turns)
fig, ax = plt.subplots(1, 2, figsize=(10, 5))
for x, y in zip(x_rec.T[:20], y_rec.T[:20]):
ax[0].plot(freqs, np.abs(np.fft.rfft(x)))
ax[1].plot(freqs, np.abs(np.fft.rfft(y)))
ax[0].set_title('horizontal incoherent spectrum $x$')
ax[0].set_xlabel('$Q_x$')
ax[1].set_title('vertical incoherent spectrum $y$')
ax[1].set_xlabel('$Q_y$');
The entire one turn map of the machine currently holds the RFSystems and the TransverseMapSegment for both longitudinal and transverse tracking. We can therefore run the full 6D tracking including the synchrotron motion by simply calling machine.track(bunch):
n_turns = 8192
z_rec = np.empty((n_turns, bunch.macroparticlenumber), dtype=np.float64)
for i in range(n_turns):
z_rec[i, :] = bunch.z
machine.track(bunch)
Now the incoherent longitudinal spectrum will show the tune depression with synchrotron amplitude due to the radio-frequency bucket non-linearities (remember we chose longitudinal_focusing='non-linear', one could also choose 'linear').
freqs = np.fft.rfftfreq(n_turns)
for z in z_rec.T[:50]:
plt.plot(freqs, np.abs(np.fft.rfft(z)))
plt.xlim(0, 0.005);
plt.title('longitudinal incoherent spectrum $z$')
plt.axvline(rfbucket.Q_s, color='red', ls='--')
plt.text(rfbucket.Q_s, plt.ylim()[1]*0.8, ' linear Q_s0')
plt.xlabel('$Q_s$');
To add up to $n$th order chromaticity in the transverse plane, one can provide an $n$ entry vector as a keyword argument to the machine at init time:
Qp_x = [-0.83 * Q_x] # linear horizontal chromaticity, in tune units
Qp_y = [-1.12 * Q_y] # linear vertical chromaticity, in tune units
machine = PS(n_segments=1, gamma=gamma, machine_configuration='LHCbeam_h7',
longitudinal_focusing='non-linear', Q_x=Q_x, Q_y=Q_y,
# chromaticity enters here:
Qp_x=Qp_x, Qp_y=Qp_y)
Generate another fresh bunch distribution:
# as before for the RF setup below transition energy:
machine.longitudinal_map.pop_kick(1)
machine.longitudinal_map.phi_offsets[0] += np.pi
bunch = machine.generate_6D_Gaussian_bunch_matched(
n_macroparticles=n_macroparticles, intensity=intensity,
epsn_x=epsn_x, epsn_y=epsn_y, sigma_z=sigma_z)
Repeating the 6D tracking exercise, ...
n_turns = 512
x_rec = np.empty((n_turns, bunch.macroparticlenumber), dtype=np.float64)
y_rec = np.empty((n_turns, bunch.macroparticlenumber), dtype=np.float64)
for i in range(n_turns):
x_rec[i, :] = bunch.x
y_rec[i, :] = bunch.y
machine.track(bunch)
... we now find the transverse tune spread around shifted tune values due to the linear chromaticity we plugged in:
freqs = 6 + np.fft.rfftfreq(n_turns)
fig, ax = plt.subplots(1, 2, figsize=(10, 5))
for x, y in zip(x_rec.T[:50], y_rec.T[:50]):
ax[0].plot(freqs, np.abs(np.fft.rfft(x)))
ax[1].plot(freqs, np.abs(np.fft.rfft(y)))
ax[0].set_title('horizontal incoherent spectrum $x$')
ax[0].set_xlabel('$Q_x$')
ax[0].set_xlim(6.19, 6.27)
ax[1].set_title('vertical incoherent spectrum $y$')
ax[1].set_xlabel('$Q_y$')
ax[1].set_xlim(6.2, 6.28);
So far so good, you now know how to set up a basic tracking simulation. You may want to go ahead and have a look at further code examples in the playgrounds repository https://github.com/PyCOMPLETE/PyHEADTAIL-playground. There you'll find out how to tune up your PyHEADTAIL simulation with e.g.
- wakefield kicks for impedance-driven instability studies
- frozen / self-consistent space charge models for longitudinal, transverse or 3D studies
- apertures for loss studies
- multi-harmonic synchrotron motion for RF gymnastics (e.g. PS triple splitting)
- different longitudinal distributions
- synchrotron radiation
- transverse feedback
- an introduction to PyHEADTAIL on the GPU
- ...
to name just a few. And... never hesitate to have a look inside the source code of the modules. :-)
Best of success for your studies!