def test_wakefield_platesresonator(self):
'''
Track through a ParallelPlatesResonator wakefield
'''
Dx = np.append(np.linspace(0., 20., self.nsegments), [0])
# add some dispersion/alpha
lhc = m.LHC(n_segments=self.nsegments,
machine_configuration='450GeV',
app_x=1e-9,
app_y=2e-9,
app_xy=-1.5e-11,
chromaticity_on=False,
amplitude_detuning_on=True,
alpha_x=1.2 * np.ones(self.nsegments),
D_x=Dx,
printer=SilentPrinter())
self.n_macroparticles = 200000
bunch_cpu = self.create_lhc_bunch(lhc) #self.create_gaussian_bunch()
bunch_gpu = self.create_lhc_bunch(lhc) #self.create_gaussian_bunch()
n_slices = 50 #5
frequency = 8e8 #1e9
R_shunt = 23e3 # [Ohm]
Q = 1.
unif_bin_slicer = UniformBinSlicer(n_slices=n_slices, n_sigma_z=1)
#res = CircularResonator(R_shunt=R_shunt, frequency=frequency, Q=Q)
res = ParallelPlatesResonator(R_shunt=R_shunt,
frequency=frequency,
Q=Q,
printer=SilentPrinter())
wake_field = WakeField(unif_bin_slicer, res)
self.assertTrue(
self._track_cpu_gpu([wake_field], bunch_cpu, bunch_gpu),
'Tracking Wakefield CircularResonator CPU/GPU differs')
def test_wakefield_wakefile(self):
'''
Track an LHC bunch and a LHC wakefield
'''
wakefile = 'autoruntests/wake_table.dat' #'./wakeforhdtl_PyZbase_Allthemachine_450GeV_B1_LHC_inj_450GeV_B1.dat'
Qp_x, Qp_y = 1., 1.
Q_s = 0.0049
n_macroparticles = 10
intensity = 1e11
longitudinal_focusing = 'linear'
machine = m.LHC(n_segments=1,
machine_configuration='450GeV',
longitudinal_focusing=longitudinal_focusing,
Qp_x=[Qp_x],
Qp_y=[Qp_y],
Q_s=Q_s,
beta_x=[65.9756],
beta_y=[71.5255],
printer=SilentPrinter())
epsn_x = 3.5e-6
epsn_y = 3.5e-6
sigma_z = 1.56e-9 * c / 4.
np.random.seed(0)
bunch_cpu = machine.generate_6D_Gaussian_bunch(n_macroparticles,
intensity,
epsn_x,
epsn_y,
sigma_z=sigma_z)
np.random.seed(0)
bunch_gpu = machine.generate_6D_Gaussian_bunch(n_macroparticles,
intensity,
epsn_x,
epsn_y,
sigma_z=sigma_z)
n_slices_wakefields = 55
n_sigma_z_wakefields = 3
slicer_for_wakefields_cpu = UniformBinSlicer(
n_slices_wakefields, n_sigma_z=n_sigma_z_wakefields)
wake_components = [
'time', 'dipole_x', 'dipole_y', 'no_quadrupole_x',
'no_quadrupole_y', 'no_dipole_xy', 'no_dipole_yx'
]
wake_table_cpu = WakeTable(wakefile,
wake_components,
printer=SilentPrinter())
wake_field_cpu = WakeField(slicer_for_wakefields_cpu, wake_table_cpu)
# also checked for 100 turns!
self.assertTrue(
self._track_cpu_gpu([wake_field_cpu],
bunch_cpu,
bunch_gpu,
nturns=2),
'Tracking through WakeField(waketable) differs')
def test_RFQ_detuner(self):
'''
Test the RFQ as a detuning element in the transverse map
'''
bunch_cpu = self.create_gaussian_bunch()
bunch_gpu = self.create_gaussian_bunch()
detuner = RFQTransverseDetuner(v_2=1.5e9,
omega=341e5 * 2.5 * np.pi,
phi_0=-np.pi / 4,
beta_x_RFQ=200.,
beta_y_RFQ=99.)
transverse_map = tt.TransverseMap(self.s,
self.alpha_x,
self.beta_x,
self.Dx,
self.alpha_y,
self.beta_y,
self.Dy,
self.Qx,
self.Qy,
detuner,
printer=SilentPrinter())
self.assertTrue(
self._track_cpu_gpu(transverse_map, bunch_cpu, bunch_gpu),
'Transverse tracking with RFQDetuner CPU/GPU differs')
def gimme(D_x=None, D_y=None, *args, **kwargs):
# Parameters for transverse map.
s = np.arange(0, n_segments + 1) * C / n_segments
alpha_x_inj = 0.
alpha_y_inj = 0.
beta_x_inj = 66.0064
beta_y_inj = 71.5376
alpha_x = alpha_x_inj * np.ones(n_segments)
beta_x = beta_x_inj * np.ones(n_segments)
if D_x is None:
D_x = np.zeros(n_segments)
alpha_y = alpha_y_inj * np.ones(n_segments)
beta_y = beta_y_inj * np.ones(n_segments)
if D_y is None:
D_y = np.zeros(n_segments)
if 'detuners' in kwargs:
trans_map = TransverseMap(
s, alpha_x, beta_x, D_x, alpha_y, beta_y, D_y, Q_x, Q_y, kwargs['detuners'])
else:
trans_map = TransverseMap(
s, alpha_x, beta_x, D_x, alpha_y, beta_y, D_y, Q_x, Q_y,
printer=SilentPrinter())
long_map = LinearMap(alpha_0, C, Q_s)
bunch = generate_bunch(
n_macroparticles, alpha_x_inj, alpha_y_inj, beta_x_inj, beta_y_inj,
long_map)
return bunch, trans_map, long_map
def test_rf_bucket_distribution(self):
'''Tests the functionality of the rf-bucket matchor'''
#SPS Q20 flattop
nparticles = 100
h1 = 4620
h2 = 4 * 4620
V1 = 10e6
V2 = 1e6
dphi1 = 0
dphi2 = 0
alpha = 0.00308
p_increment = 0
long_map = RFSystems(self.circumference, [h1, h2], [V1, V2],
[dphi1, dphi2], [alpha],
self.gamma,
p_increment,
charge=self.charge,
mass=self.mass)
bucket = long_map.get_bucket(gamma=self.gamma)
bunch = gf.ParticleGenerator(nparticles,
1e11,
constants.e,
constants.m_p,
self.circumference,
self.gamma,
distribution_z=gf.RF_bucket_distribution(
bucket,
epsn_z=0.002,
printer=SilentPrinter())).generate()
def test_PSB_linear(self):
psb = m.PSB(n_segments=self.nsegments,
printer=SilentPrinter(),
machine_configuration='160MeV',
longitudinal_focusing='linear')
bunch_cpu = self.create_all1_bunch()
bunch_gpu = self.create_all1_bunch()
one_turn_map = psb.one_turn_map
self.assertTrue(
self._track_cpu_gpu(one_turn_map, bunch_cpu, bunch_gpu),
'Tracking through SPS Q26 injection one turn map CPU/GPU differs.')
def test_transverse_damper(self):
'''
Track through a transverse damper
'''
bunch_cpu = self.create_all1_bunch()
bunch_gpu = self.create_all1_bunch()
dampingrate_x = 0.01
dampingrate_y = 0.05
damp = TransverseDamper(dampingrate_x,
dampingrate_y,
printer=SilentPrinter())
self.assertTrue(self._track_cpu_gpu([damp], bunch_cpu, bunch_gpu),
'Tracking TransverseDamper CPU/GPU differs')
def test_transverse_track_with_Adetuning(self):
'''
Track the GPU bunch through a TransverseMap with amplitude detuning
Check if results on CPU and GPU are the same
Special test since it requires the PyCERNmachines module to create
the LHC beam. Skip if not available.
'''
# Use a real bunch (not all set to 1), otherwise the Action/detuning
# gets too big and the numbers blow up.
Dx = np.append(np.linspace(0., 20., self.nsegments), [0])
# add some dispersion/alpha
lhc = m.LHC(n_segments=self.nsegments,
machine_configuration='450GeV',
app_x=1e-9,
app_y=2e-9,
app_xy=-1.5e-11,
chromaticity_on=False,
amplitude_detuning_on=True,
alpha_x=1.2 * np.ones(self.nsegments),
D_x=Dx,
printer=SilentPrinter())
# create pure Python map, PyCERNmachine uses Cython.
adetuner = AmplitudeDetuning(lhc.app_x, lhc.app_y, lhc.app_xy)
transverse_map = tt.TransverseMap(lhc.s,
lhc.alpha_x,
lhc.beta_x,
lhc.D_x,
lhc.alpha_y,
lhc.beta_y,
lhc.D_y,
lhc.Q_x,
lhc.Q_y,
adetuner,
printer=SilentPrinter())
bunch_cpu = self.create_lhc_bunch(lhc)
bunch_gpu = self.create_lhc_bunch(lhc)
self.assertTrue(
self._track_cpu_gpu(transverse_map, bunch_cpu, bunch_gpu),
'Transverse tracking with Adetuning CPU/GPU differs')
def test_SPS_nonlinear(self):
sps = m.SPS(n_segments=self.nsegments,
printer=SilentPrinter(),
machine_configuration='Q26-injection',
Qp_x=2,
Qp_y=-2.2,
charge=e,
mass=m_p)
bunch_cpu = self.create_all1_bunch()
bunch_gpu = self.create_all1_bunch()
one_turn_map = sps.one_turn_map
self.assertTrue(
self._track_cpu_gpu(one_turn_map, bunch_cpu, bunch_gpu),
'Tracking through SPS Q26 injection one turn map CPU/GPU differs.')
def setUp(self):
np.random.seed(0)
self.nparticles = 1000
self.epsx = 0.5
self.intensity = 1e11
self.charge = constants.e
self.mass = constants.m_p
self.circumference = 99
self.gamma = 27.1
self.generator = gf.ParticleGenerator(
self.nparticles, self.intensity,
self.charge, self.mass, self.circumference, self.gamma,
distribution_x=gf.gaussian2D(0.5),alpha_x=-0.7, beta_x=4, D_x=0,
distribution_z=gf.gaussian2D(3.0),
printer=SilentPrinter())
self.beam = self.generator.generate()
def test_transverse_track_no_detuning(self):
'''
Track the GPU bunch through a TransverseMap with no detuning
Check if results on CPU and GPU are the same
'''
bunch_cpu = self.create_all1_bunch()
bunch_gpu = self.create_all1_bunch()
transverse_map = tt.TransverseMap(self.s,
self.alpha_x,
self.beta_x,
self.Dx,
self.alpha_y,
self.beta_y,
self.Dy,
self.Qx,
self.Qy,
printer=SilentPrinter())
self.assertTrue(
self._track_cpu_gpu(transverse_map, bunch_cpu, bunch_gpu),
'Transverse tracking w/o detuning CPU/GPU differs')
def test_tracking_without_detuner(self):
'''Tests whether the cython and pure python version of the
TransverseMap yield the same results (x,y,xp,yp) when tracking a beam
without detuners
'''
pure_python_map = pure_py.TransverseMap(self.s,
self.alpha_x,
self.beta_x,
self.Dx,
self.alpha_y,
self.beta_y,
self.Dy,
self.Qx,
self.Qy,
printer=SilentPrinter())
beam_p = self.create_bunch()
for s in pure_python_map:
s.track(beam_p)
def test_transverse_track_with_Cdetuning(self):
'''
Track the GPU bunch through a TransverseMap with chromaticity (detuning)
Check if results on CPU and GPU are the same
'''
bunch_cpu = self.create_all1_bunch()
bunch_gpu = self.create_all1_bunch()
detuner = Chromaticity(Qp_x=[5, 1], Qp_y=[7, 2])
transverse_map = tt.TransverseMap(self.s,
self.alpha_x,
self.beta_x,
self.Dx,
self.alpha_y,
self.beta_y,
self.Dy,
self.Qx,
self.Qy,
detuner,
printer=SilentPrinter())
self.assertTrue(
self._track_cpu_gpu(transverse_map, bunch_cpu, bunch_gpu),
'Transverse tracking with chromaticity CPU/GPU differs')
def run():
def track_n_save(bunch, map_):
mean_x = np.empty(n_turns)
mean_y = np.empty(n_turns)
sigma_z = np.empty(n_turns)
for i in xrange(n_turns):
mean_x[i] = bunch.mean_x()
mean_y[i] = bunch.mean_y()
sigma_z[i] = bunch.sigma_z()
for m_ in map_:
m_.track(bunch)
return mean_x, mean_y, sigma_z
def my_fft(data):
t = np.arange(len(data))
fft = np.fft.rfft(data)
fft_freq = np.fft.rfftfreq(t.shape[-1])
return fft_freq, np.abs(fft.real)
def generate_bunch(n_macroparticles, alpha_x, alpha_y, beta_x, beta_y,
linear_map):
intensity = 1.05e11
sigma_z = 0.059958
gamma = 3730.26
p0 = np.sqrt(gamma**2 - 1) * m_p * c
beta_z = (linear_map.eta(dp=0, gamma=gamma) *
linear_map.circumference / (2 * np.pi * linear_map.Q_s))
epsn_x = 3.75e-6 # [m rad]
epsn_y = 3.75e-6 # [m rad]
epsn_z = 4 * np.pi * sigma_z**2 * p0 / (beta_z * e)
bunch = generators.generate_Gaussian6DTwiss(
macroparticlenumber=n_macroparticles,
intensity=intensity,
charge=e,
gamma=gamma,
mass=m_p,
circumference=C,
alpha_x=alpha_x,
beta_x=beta_x,
epsn_x=epsn_x,
alpha_y=alpha_y,
beta_y=beta_y,
epsn_y=epsn_y,
beta_z=beta_z,
epsn_z=epsn_z)
# print ('bunch sigma_z=' + bunch.sigma_z())
return bunch
def track_n_show(bunch, slicer, map_woWakes, wake_field):
fig, ((ax1, ax2)) = plt.subplots(2, 1, figsize=(16, 16))
xp_diff = np.zeros(n_macroparticles)
for i in xrange(n_turns):
for m_ in map_woWakes:
m_.track(bunch)
# Dipole X kick.
if i == (n_turns - 1):
xp_old = bunch.xp.copy()
wake_field.track(bunch)
if i == (n_turns - 1):
xp_diff[:] = bunch.xp[:] - xp_old[:]
# Plot bunch.z vs. slice index of particle. Mark particles within
# z cuts in green.
nsigmaz_lbl = ' (nsigmaz =' + str(n_sigma_z) + ')'
slice_set = bunch.get_slices(slicer)
pidx = slice_set.particles_within_cuts
slidx = slice_set.slice_index_of_particle
z_cut_tail, z_cut_head = slice_set.z_cut_tail, slice_set.z_cut_head
# In[4]:
# Basic parameters.
n_turns = 10
n_segments = 1
n_macroparticles = 50
Q_x = 64.28
Q_y = 59.31
Q_s = 0.0020443
C = 26658.883
R = C / (2. * np.pi)
alpha_x_inj = 0.
alpha_y_inj = 0.
beta_x_inj = 66.0064
beta_y_inj = 71.5376
alpha_0 = [0.0003225]
#waketable filename
fn = os.path.join(os.path.dirname(__file__), 'wake_table.dat')
# In[5]:
# Parameters for transverse map.
s = np.arange(0, n_segments + 1) * C / n_segments
alpha_x = alpha_x_inj * np.ones(n_segments)
beta_x = beta_x_inj * np.ones(n_segments)
D_x = np.zeros(n_segments)
alpha_y = alpha_y_inj * np.ones(n_segments)
beta_y = beta_y_inj * np.ones(n_segments)
D_y = np.zeros(n_segments)
# In[6]:
# In[7]:
# In[8]:
# CASE TEST SETUP
trans_map = TransverseMap(s, alpha_x, beta_x, D_x, alpha_y, beta_y, D_y,
Q_x, Q_y)
long_map = LinearMap(alpha_0, C, Q_s)
bunch = generate_bunch(n_macroparticles, alpha_x_inj, alpha_y_inj,
beta_x_inj, beta_y_inj, long_map)
# In[9]:
# CASE I
# Transverse and long. tracking (linear), and wakes from WakeTable source.
# DIPOLE X, UniformBinSlicer
n_sigma_z = 2
n_slices = 15
uniform_bin_slicer = UniformBinSlicer(n_slices=n_slices,
n_sigma_z=n_sigma_z)
# Definition of WakeField as a composition of different sources.
wake_file_columns = [
'time', 'dipole_x', 'no_dipole_y', 'no_quadrupole_x',
'no_quadrupole_y', 'no_dipole_xy', 'no_dipole_yx'
]
table = WakeTable(fn,
wake_file_columns,
printer=SilentPrinter(),
warningprinter=SilentPrinter())
wake_field = WakeField(uniform_bin_slicer, table)
trans_map = [m for m in trans_map]
map_woWakes = trans_map + [long_map]
track_n_save(bunch, map_woWakes)
# In[10]:
# CASE II
# Transverse and long. tracking (linear), and wakes from WakeTable source.
# DIPOLE X, UniformChargeSlicer
n_sigma_z = 2
n_slices = 15
uniform_charge_slicer = UniformChargeSlicer(n_slices=n_slices,
n_sigma_z=n_sigma_z)
# Definition of WakeField as a composition of different sources.
wake_file_columns = [
'time', 'dipole_x', 'no_dipole_y', 'no_quadrupole_x',
'no_quadrupole_y', 'no_dipole_xy', 'no_dipole_yx'
]
table = WakeTable(fn,
wake_file_columns,
printer=SilentPrinter(),
warningprinter=SilentPrinter())
wake_field = WakeField(uniform_charge_slicer, table)
trans_map = [m for m in trans_map]
map_woWakes = trans_map + [long_map]
track_n_save(bunch, map_woWakes)
# In[11]:
# CASE III
# Transverse and long. tracking (linear), and wakes from WakeTable source.
# Quadrupole X, UniformChargeSlicer
n_sigma_z = 2
n_slices = 15
uniform_charge_slicer = UniformChargeSlicer(n_slices=n_slices,
n_sigma_z=n_sigma_z)
# Definition of WakeField as a composition of different sources.
wake_file_columns = [
'time', 'no_dipole_x', 'no_dipole_y', 'quadrupole_x',
'no_quadrupole_y', 'no_dipole_xy', 'no_dipole_yx'
]
table = WakeTable(fn,
wake_file_columns,
printer=SilentPrinter(),
warningprinter=SilentPrinter())
wake_field = WakeField(uniform_charge_slicer, table)
trans_map = [m for m in trans_map]
map_woWakes = trans_map + [long_map]
track_n_save(bunch, map_woWakes)
# In[12]:
# CASE IV
# Transverse and long. tracking (linear), and wakes from WakeTable source.
# Quadrupole X, UniformBinSlicer
n_sigma_z = 2
n_slices = 15
uniform_bin_slicer = UniformBinSlicer(n_slices=n_slices,
n_sigma_z=n_sigma_z)
# Definition of WakeField as a composition of different sources.
wake_file_columns = [
'time', 'no_dipole_x', 'no_dipole_y', 'quadrupole_x',
'no_quadrupole_y', 'no_dipole_xy', 'no_dipole_yx'
]
table = WakeTable(fn,
wake_file_columns,
printer=SilentPrinter(),
warningprinter=SilentPrinter())
wake_field = WakeField(uniform_bin_slicer, table)
trans_map = [m for m in trans_map]
map_woWakes = trans_map + [long_map]
track_n_save(bunch, map_woWakes)
# In[15]:
# CASE V
# Transverse and long. tracking (linear),
# Resonator circular
n_sigma_z = 2
n_slices = 15
uniform_bin_slicer = UniformBinSlicer(n_slices=n_slices,
n_sigma_z=n_sigma_z)
# Definition of WakeField as a composition of different sources.
reson_circ = CircularResonator(R_shunt=1e6, frequency=1e8, Q=1)
wake_field = WakeField(uniform_bin_slicer, reson_circ)
trans_map = [m for m in trans_map]
map_woWakes = trans_map + [long_map]
track_n_save(bunch, map_woWakes)
# In[16]:
# CASE V b.
# Transverse and long. tracking (linear),
# Several Resonators circular
n_sigma_z = 2
n_slices = 15
uniform_bin_slicer = UniformBinSlicer(n_slices=n_slices,
n_sigma_z=n_sigma_z)
# Definition of WakeField as a composition of different sources.
reson_circ = CircularResonator(R_shunt=1e6, frequency=1e8, Q=1)
reson_circ2 = CircularResonator(R_shunt=1e6, frequency=1e9, Q=0.8)
reson_circ3 = CircularResonator(R_shunt=5e6, frequency=1e6, Q=0.2)
wake_field = WakeField(uniform_bin_slicer, reson_circ, reson_circ2,
reson_circ3)
trans_map = [m for m in trans_map]
map_woWakes = trans_map + [long_map]
track_n_save(bunch, map_woWakes)
# In[17]:
# CASE V c.
# Transverse and long. tracking (linear),
# Resonator parallel_plates
n_sigma_z = 2
n_slices = 15
uniform_bin_slicer = UniformBinSlicer(n_slices=n_slices,
n_sigma_z=n_sigma_z)
# Definition of WakeField as a composition of different sources.
reson_para = ParallelPlatesResonator(R_shunt=1e6, frequency=1e8, Q=1)
wake_field = WakeField(uniform_bin_slicer, reson_para)
trans_map = [m for m in trans_map]
map_woWakes = trans_map + [long_map]
track_n_save(bunch, map_woWakes)
# In[18]:
# CASE V d.
# Transverse and long. tracking (linear),
# Resonator w. longitudinal wake
n_sigma_z = 2
n_slices = 15
uniform_bin_slicer = UniformBinSlicer(n_slices=n_slices,
n_sigma_z=n_sigma_z)
# Definition of WakeField as a composition of different sources.
reson = Resonator(R_shunt=1e6,
frequency=1e8,
Q=1,
Yokoya_X1=1,
Yokoya_X2=1,
Yokoya_Y1=1,
Yokoya_Y2=1,
switch_Z=True)
wake_field = WakeField(uniform_bin_slicer, reson)
trans_map = [m for m in trans_map]
map_woWakes = trans_map + [long_map]
track_n_save(bunch, map_woWakes)
# In[19]:
# CASE VI
# Transverse and long. tracking (linear),
# ResistiveWall circular
n_sigma_z = 2
n_slices = 15
uniform_bin_slicer = UniformBinSlicer(n_slices=n_slices,
n_sigma_z=n_sigma_z)
# Definition of WakeField as a composition of different sources.
resis_circ = CircularResistiveWall(pipe_radius=5e-2,
resistive_wall_length=C,
conductivity=1e6,
dt_min=1e-3)
wake_field = WakeField(uniform_bin_slicer, resis_circ)
trans_map = [m for m in trans_map]
map_woWakes = trans_map + [long_map]
track_n_save(bunch, map_woWakes)
# In[20]:
# CASE VI b.
# Transverse and long. tracking (linear),
# ResistiveWall parallel_plates
n_sigma_z = 2
n_slices = 15
uniform_bin_slicer = UniformBinSlicer(n_slices=n_slices,
n_sigma_z=n_sigma_z)
# Definition of WakeField as a composition of different sources.
resis_para = ParallelPlatesResistiveWall(pipe_radius=5e-2,
resistive_wall_length=C,
conductivity=1e6,
dt_min=1e-3)
wake_field = WakeField(uniform_bin_slicer, resis_para)
trans_map = [m for m in trans_map]
map_woWakes = trans_map + [long_map]
track_n_save(bunch, map_woWakes)
# In[21]:
# CASE VII.
# Transverse and long. tracking (linear),
# Pass mixture of WakeSources to define WakeField.
n_sigma_z = 2
n_slices = 15
uniform_bin_slicer = UniformBinSlicer(n_slices=n_slices,
n_sigma_z=n_sigma_z)
# Definition of WakeField as a composition of different sources.
resis_circ = CircularResistiveWall(pipe_radius=5e-2,
resistive_wall_length=C,
conductivity=1e6,
dt_min=1e-3)
reson_para = ParallelPlatesResonator(R_shunt=1e6, frequency=1e8, Q=1)
wake_file_columns = [
'time', 'dipole_x', 'dipole_y', 'quadrupole_x', 'quadrupole_y',
'dipole_xy', 'dipole_yx'
]
table = WakeTable(fn,
wake_file_columns,
printer=SilentPrinter(),
warningprinter=SilentPrinter())
wake_field = WakeField(uniform_bin_slicer, resis_circ, reson_para, table)
trans_map = [m for m in trans_map]
map_woWakes = trans_map + [long_map]
track_n_save(bunch, map_woWakes)
def run():
def track_n_save(bunch, map_):
mean_x = np.empty(n_turns)
mean_y = np.empty(n_turns)
sigma_z = np.empty(n_turns)
for i in xrange(n_turns):
mean_x[i] = bunch.mean_x()
mean_y[i] = bunch.mean_y()
sigma_z[i] = bunch.sigma_z()
for m_ in map_:
m_.track(bunch)
return mean_x, mean_y, sigma_z
def copy_bunch(bunch_source, bunch_target):
bunch_target.x = bunch_source.x.copy()
bunch_target.xp = bunch_source.xp.copy()
bunch_target.y = bunch_source.y.copy()
bunch_target.yp = bunch_source.yp.copy()
bunch_target.z = bunch_source.z.copy()
bunch_target.dp = bunch_source.dp.copy()
def my_fft(data):
t = np.arange(len(data))
fft = np.fft.rfft(data)
fft_freq = np.fft.rfftfreq(t.shape[-1])
return fft_freq, np.abs(fft.real)
def generate_bunch(n_macroparticles, alpha_x, alpha_y, beta_x, beta_y,
linear_map):
intensity = 1.05e11
sigma_z = 0.059958
gamma = 3730.26
gamma_t = 1. / np.sqrt(alpha_0)
p0 = np.sqrt(gamma**2 - 1) * m_p * c
beta_z = (linear_map.eta(dp=0, gamma=gamma) *
linear_map.circumference / (2 * np.pi * linear_map.Q_s))
epsn_x = 3.75e-6 # [m rad]
epsn_y = 3.75e-6 # [m rad]
epsn_z = 4 * np.pi * sigma_z**2 * p0 / (beta_z * e)
bunch = generators.generate_Gaussian6DTwiss(
macroparticlenumber=n_macroparticles,
intensity=intensity,
charge=e,
gamma=gamma,
mass=m_p,
circumference=C,
alpha_x=alpha_x,
beta_x=beta_x,
epsn_x=epsn_x,
alpha_y=alpha_y,
beta_y=beta_y,
epsn_y=epsn_y,
beta_z=beta_z,
epsn_z=epsn_z)
return bunch
# In[4]:
# Basic parameters.
n_turns = 3
n_segments = 5
n_macroparticles = 10
Q_x = 64.28
Q_y = 59.31
Q_s = 0.0020443
C = 26658.883
R = C / (2. * np.pi)
alpha_0 = [0.0003225]
# **Things tested**
# - Instantiation of a TransverseMap (and therewith of several
# TransverseSegmentMaps as we have more than 1 segment).
# - With and without detuners, i.e. instantiation of Chromaticity and
# AmplitudeDetuning DetunerCollections as well as the corresponding
# SegmentDetuners.
# - Are betatron tunes Q_{x,y} and detuning strengths correctly
# scaled to segment lengths?
# - TransverseSegmentMap.track(beam) method.
# - Check spectrum of beam centroid motion.
# - Betatron tune (implicitly checks the scaling to segment lengths)
# - If chromaticity and linear synchro motion are on: synchrotron sidebands?
# - If amplitude detuning is on and there is initial kick: decoherence?
# - Does dispersion work correctly?
# - Is exception risen when s[0] != 0 or s[-1] != C?
# - Get optics at injection.
# In[5]:
# TEST CASE SETUP
def gimme(D_x=None, D_y=None, *args, **kwargs):
# Parameters for transverse map.
s = np.arange(0, n_segments + 1) * C / n_segments
alpha_x_inj = 0.
alpha_y_inj = 0.
beta_x_inj = 66.0064
beta_y_inj = 71.5376
alpha_x = alpha_x_inj * np.ones(n_segments)
beta_x = beta_x_inj * np.ones(n_segments)
if D_x is None:
D_x = np.zeros(n_segments)
alpha_y = alpha_y_inj * np.ones(n_segments)
beta_y = beta_y_inj * np.ones(n_segments)
if D_y is None:
D_y = np.zeros(n_segments)
if 'detuners' in kwargs:
trans_map = TransverseMap(s, alpha_x, beta_x, D_x, alpha_y, beta_y,
D_y, Q_x, Q_y, kwargs['detuners'])
else:
trans_map = TransverseMap(s,
alpha_x,
beta_x,
D_x,
alpha_y,
beta_y,
D_y,
Q_x,
Q_y,
printer=SilentPrinter())
long_map = LinearMap(alpha_0, C, Q_s)
bunch = generate_bunch(n_macroparticles, alpha_x_inj, alpha_y_inj,
beta_x_inj, beta_y_inj, long_map)
return bunch, trans_map, long_map
# In[6]:
# CASE I
# Pure transverse tracking. Without detuners.
bunch, trans_map, _ = gimme()
map_ = trans_map
mean_x, mean_y, sigma_z = track_n_save(bunch, map_)
# In[7]:
# CASE II
# Without detuners. With linear synchrotron motion.
bunch, trans_map, long_map = gimme()
# This tests if TransverseMap is actually a sequence.
trans_one_turn = [m for m in trans_map]
map_ = trans_one_turn + [long_map]
mean_x, mean_y, sigma_z = track_n_save(bunch, map_)
# In[8]:
# CASE III
# With chromaticity in horizontal and vertical. With linear synchrotron motion.
chroma = Chromaticity(Qp_x=6, Qp_y=10)
bunch, trans_map, long_map = gimme(detuners=(chroma, ))
trans_one_turn = [m for m in trans_map]
map_ = trans_one_turn + [long_map]
mean_x, mean_y, sigma_z = track_n_save(bunch, map_)
# In[9]:
# CASE IV
# With amplitude detuning. With linear synchrotron motion. With initial kick.
ampl_det = AmplitudeDetuning.from_octupole_currents_LHC(i_focusing=200,
i_defocusing=-200)
bunch, trans_map, long_map = gimme(detuners=(ampl_det, ))
trans_one_turn = [m for m in trans_map]
map_ = trans_one_turn + [long_map]
bunch.x += 0.0003
bunch.y += 0.0005
mean_x, mean_y, sigma_z = track_n_save(bunch, map_)
# In[10]:
# CASE V
# With amplitude detuning and chromaticity. With linear synchrotron motion. With initial kick.
ampl_det = AmplitudeDetuning.from_octupole_currents_LHC(i_focusing=200,
i_defocusing=-200)
chroma = Chromaticity(Qp_x=10, Qp_y=6)
bunch, trans_map, long_map = gimme(detuners=(ampl_det, chroma))
trans_one_turn = [m for m in trans_map]
map_ = trans_one_turn + [long_map]
bunch.x += 0.0003
bunch.y += 0.0005
mean_x, mean_y, sigma_z = track_n_save(bunch, map_)
# In[11]:
# CASE VI
n_turns = 2
n_segments = 5
# Initial bunch distro to be used with and without dispersion
# for direct comparison.
bunch_orig, _, _ = gimme()
# With dispersion.
D_x = np.zeros(n_segments)
D_y = np.zeros(n_segments)
D_x[1] = 4.5
D_y[3] = 2.3
D_x[2] = 4.5
D_y[4] = 2.3
bunch_wD, trans_map_wD, _ = gimme(D_x, D_y)
copy_bunch(bunch_orig, bunch_wD)
# Add dispersion (manually for now).
bunch_wD.x += D_x[0] * bunch_wD.dp
bunch_wD.y += D_y[0] * bunch_wD.dp
x_i_wD = np.zeros((n_segments * n_turns, n_macroparticles))
y_i_wD = np.zeros((n_segments * n_turns, n_macroparticles))
for j in xrange(n_segments * n_turns):
x_i_wD[j, :] = bunch_wD.x
y_i_wD[j, :] = bunch_wD.y
trans_map_wD[j % n_segments].track(bunch_wD)
# Without dispersion.
bunch_woD, trans_map_woD, _ = gimme()
copy_bunch(bunch_orig, bunch_woD)
x_i_woD = np.zeros((n_segments * n_turns, n_macroparticles))
y_i_woD = np.zeros((n_segments * n_turns, n_macroparticles))
dp_i_woD = np.zeros((n_segments * n_turns, n_macroparticles))
for j in xrange(n_segments * n_turns):
x_i_woD[j, :] = bunch_woD.x
y_i_woD[j, :] = bunch_woD.y
dp_i_woD[j, :] = bunch_woD.dp
trans_map_woD[j % n_segments].track(bunch_woD)
# In[12]:
# Test how detuning parameters and betatron tunes are scaled
# for the TransverseSegmentMaps.
Qp_x = 8.
Qp_y = 10.
chroma = Chromaticity(Qp_x, Qp_y)
bunch, trans_map, _ = gimme(detuners=(chroma, ))
i = 1
# print 'Q_x'
# print Q_x / float(n_segments)
# print trans_map[i].dQ_x
# print 'Q_y'
# print Q_y / float(n_segments)
# print trans_map[i].dQ_y
# print 'Qp_x'
# print Qp_x / float(n_segments)
# print trans_map[i].segment_detuners[0].dQp_x
# print 'Qp_y'
# print Qp_y / float(n_segments)
# print trans_map[i].segment_detuners[0].dQp_y
app_x = 20.
app_y = 12.
app_xy = 30.
ampl_det = AmplitudeDetuning(app_x, app_y, app_xy)
bunch, trans_map, _ = gimme(detuners=(ampl_det, ))
# print 'app_x'
# print app_x / float(n_segments)
# print trans_map[i].segment_detuners[0].dapp_x
# print 'app_y'
# print app_y / float(n_segments)
# print trans_map[i].segment_detuners[0].dapp_y
# print 'app_xy'
# print app_xy / float(n_segments)
# print trans_map[i].segment_detuners[0].dapp_xy
# In[13]:
# Test if optics at injection are correctly returned.
s = np.arange(0, n_segments + 1) * C / n_segments
alpha_x_inj = 0.
alpha_y_inj = 0.
beta_x_inj = 66.0064
beta_y_inj = 71.5376
alpha_x = alpha_x_inj * np.ones(n_segments)
beta_x = beta_x_inj * np.ones(n_segments)
alpha_y = alpha_y_inj * np.ones(n_segments)
beta_y = beta_y_inj * np.ones(n_segments)
trans_map = TransverseMap(s,
alpha_x,
beta_x,
D_x,
alpha_y,
beta_y,
D_y,
Q_x,
Q_y,
printer=SilentPrinter())
inj_opt_dict = trans_map.get_injection_optics()
# In[14]:
# Test if exception is risen when s[0] != 0 or s[-1] != C
s = np.array([0., 4, 10, 12, C - 1])
try:
trans_map = TransverseMap(s,
alpha_x,
beta_x,
D_x,
alpha_y,
beta_y,
D_y,
Q_x,
Q_y,
printer=SilentPrinter())
# print ('test NOT passed. No error raised!')
except ValueError as exc:
# print ('test passed.\n')
# print ('Error message:\n' + exc.message)
pass