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 gimme(*detuners):
trans_map = TransverseMap(s, alpha_x, beta_x, D_x, alpha_y, beta_y,
D_y, Q_x, Q_y, *detuners)
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 create_transverse_map(self, chromaticity_on=True,
amplitude_detuning_on=True):
detuners = []
if chromaticity_on:
detuners.append(Chromaticity(self.Qp_x, self.Qp_y))
if amplitude_detuning_on:
detuners.append(
AmplitudeDetuning(self.app_x, self.app_y, self.app_xy)
)
self.transverse_map = TransverseMap(
self.s,
self.alpha_x, self.beta_x, self.D_x,
self.alpha_y, self.beta_y, self.D_y,
self.Q_x, self.Q_y, detuners, printer=self._printer)
yoffset = 0e-4
bunch.x += xoffset
bunch.y += yoffset
afile = open('bunch', 'wb')
pickle.dump(bunch, afile)
afile.close()
# ===================================
# CREATE TRANVERSE AND LONGITUDINAL MAPS
# ==================================
scale_factor = 2*bunch.p0 # for detuning coefficients
transverse_map = TransverseMap(s, alpha_x, beta_x, D_x, alpha_y, beta_y, D_y, Q_x, Q_y,
[Chromaticity(Qp_x, Qp_y),
AmplitudeDetuning(app_x*scale_factor, app_y*scale_factor, app_xy*scale_factor)])
longitudinal_map = LinearMap([alpha], circumference, Q_s)
# ======================================================================
# SET UP ACCELERATOR MAP AND START TRACKING
# ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
t0 = time.clock()
#reload object from file
file2 = open('bunch', 'rb')
bunch = pickle.load(file2)
import numpy as np
import pickle
import csv
import NAFFlib as pnf
from scipy.constants import c
from PyHEADTAIL.trackers.detuners import Chromaticity, AmplitudeDetuning
from PyHEADTAIL.trackers.transverse_tracking import TransverseMap
from PyHEADTAIL.trackers.simple_long_tracking import RFSystems, LinearMap
import matplotlib.pyplot as plt
n_turns = 1000 # 1000 turns are enough for the NAFF algorithm to compute the tune
# 1. CREATE ONE TURN MAP
transverse_map = TransverseMap(
pp.s, pp.alpha_x, pp.beta_x, pp.D_x, pp.alpha_y, pp.beta_y, pp.D_y, pp.Q_x,
pp.Q_y, [
Chromaticity(pp.Qp_x, pp.Qp_y),
AmplitudeDetuning(pp.app_x, pp.app_y, pp.app_xy)
])
longitudinal_map = LinearMap([pp.alpha], pp.circumference, pp.Q_s)
one_turn_map = [transverse_map[0]] + [longitudinal_map]
# 2. LOAD OBJECTS FROM FILES, BUNCH AND NOISE KICKS
bfile = open('input/bunch', 'rb')
bunch = pickle.load(bfile)
bfile.close()
# calculate initial actions
Jx = 1 / 2 * ((1 + pp.alpha_x**2) / pp.beta_x * bunch.x**2 +
2 * pp.alpha_x * bunch.x * bunch.xp + pp.beta_x * bunch.xp**2)
def _construct_transverse_map(
self, optics_mode=None,
circumference=None, n_segments=None, s=None, name=None,
alpha_x=None, beta_x=None, D_x=None,
alpha_y=None, beta_y=None, D_y=None,
accQ_x=None, accQ_y=None,
Qp_x=None, Qp_y=None, app_x=None, app_y=None, app_xy=None,
other_detuners=None, use_cython=None):
if optics_mode == 'smooth':
if circumference is None:
raise ValueError('circumference has to be specified '
'if optics_mode = "smooth"')
if n_segments is None:
raise ValueError('n_segments has to be specified '
'if optics_mode = "smooth"')
if s is not None:
raise ValueError('s vector should not be provided '
'if optics_mode = "smooth"')
s = (np.arange(0, n_segments + 1) * circumference / n_segments)
alpha_x = 0.*s
beta_x = 0.*s+beta_x
D_x = 0.*s+D_x
alpha_y = 0.*s
beta_y = 0.*s+beta_y
D_y = 0.*s+D_y
elif optics_mode == 'non-smooth':
if circumference is not None:
raise ValueError('circumference should not be provided '
'if optics_mode = "non-smooth"')
if n_segments is not None:
raise ValueError('n_segments should not be provided '
'if optics_mode = "non-smooth"')
if s is None:
raise ValueError('s has to be specified '
'if optics_mode = "smooth"')
else:
raise ValueError('optics_mode not recognized')
detuners = []
if Qp_x != 0 or Qp_y != 0:
detuners.append(Chromaticity(Qp_x, Qp_y))
if app_x != 0 or app_y != 0 or app_xy != 0:
detuners.append(AmplitudeDetuning(app_x, app_y, app_xy))
detuners += other_detuners
self.transverse_map = TransverseMap(
s=s,
alpha_x=alpha_x, beta_x=beta_x, D_x=D_x,
alpha_y=alpha_y, beta_y=beta_y, D_y=D_y,
accQ_x=accQ_x, accQ_y=accQ_y, detuners=detuners)
self.circumference = s[-1]
self.transverse_map.n_segments = len(s)-1
if name is None:
self.transverse_map.name = ['P_%d' % ip for ip in range(len(s)-1)]
self.transverse_map.name.append('end_ring')
else:
self.transverse_map.name = name
for i_seg, m in enumerate(self.transverse_map):
m.i0 = i_seg
m.i1 = i_seg+1
m.s0 = self.transverse_map.s[i_seg]
m.s1 = self.transverse_map.s[i_seg+1]
m.name0 = self.transverse_map.name[i_seg]
m.name1 = self.transverse_map.name[i_seg+1]
m.beta_x0 = self.transverse_map.beta_x[i_seg]
m.beta_x1 = self.transverse_map.beta_x[i_seg+1]
m.beta_y0 = self.transverse_map.beta_y[i_seg]
m.beta_y1 = self.transverse_map.beta_y[i_seg+1]
# insert transverse map in the ring
for m in self.transverse_map:
self.one_turn_map.append(m)
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(bunch, map_):
for i in range(n_turns):
for m in map_:
m.track(bunch)
def generate_bunch(n_macroparticles, alpha_x, alpha_y, beta_x, beta_y,
alpha_0, Q_s, R):
intensity = 1.05e11
sigma_z = 0.059958
gamma = 3730.26
eta = alpha_0 - 1. / gamma**2
gamma_t = 1. / np.sqrt(alpha_0)
p0 = np.sqrt(gamma**2 - 1) * m_p * c
beta_z = eta * R / 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()
return bunch
# In[4]:
# Basic parameters.
n_turns = 3
n_segments = 1
n_macroparticles = 10
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
# 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]:
# CASE I
# With amplitude detuning (python implementation)
# EXPECTED TUNE SPREADS AT THE GIVEN SETTINGS ARE 5e-4 FOR HORIZONTAL
# AND VERTICAL.
bunch = generate_bunch(n_macroparticles, alpha_x_inj, alpha_y_inj,
beta_x_inj, beta_y_inj, alpha_0, Q_s, R)
ampl_det = AmplitudeDetuning.from_octupole_currents_LHC(i_focusing=400,
i_defocusing=-400)
trans_map = TransverseMap(s, alpha_x, beta_x, D_x, alpha_y, beta_y, D_y,
Q_x, Q_y, [ampl_det])
trans_one_turn = [m for m in trans_map]
map_ = trans_one_turn
track(bunch, map_)
# In[7]:
# CASE II
# With first order Chromaticity (python implementation)
bunch = generate_bunch(n_macroparticles, alpha_x_inj, alpha_y_inj,
beta_x_inj, beta_y_inj, alpha_0, Q_s, R)
chroma = Chromaticity(Qp_x=[6], Qp_y=[3])
trans_map = TransverseMap(s, alpha_x, beta_x, D_x, alpha_y, beta_y, D_y,
Q_x, Q_y, [chroma])
trans_one_turn = [m for m in trans_map]
map_ = trans_one_turn
track(bunch, map_)
# In[8]:
# CASE III
# With higher order Chromaticity (python implementation)
bunch = generate_bunch(n_macroparticles, alpha_x_inj, alpha_y_inj,
beta_x_inj, beta_y_inj, alpha_0, Q_s, R)
chroma = Chromaticity(Qp_x=[6., 4e4], Qp_y=[3., 0., 2e8])
trans_map = TransverseMap(s, alpha_x, beta_x, D_x, alpha_y, beta_y, D_y,
Q_x, Q_y, [chroma])
trans_one_turn = [m for m in trans_map]
map_ = trans_one_turn
track(bunch, map_)
def run():
# HELPERS
def read_all_data(bfile, sfile, pfile):
bunchdata = hp.File(bfile + '.h5')
slicedata = hp.File(sfile + '.h5')
particledata = hp.File(pfile + '.h5part')
# Bunchdata
bdata = bunchdata['Bunch']
n_turns = len(bdata['mean_x'])
_ = np.empty(n_turns)
for key in list(bdata.keys()):
_[:] = bdata[key][:]
# Slicedata
sdata = slicedata['Slices']
sbdata = slicedata['Bunch']
n_turns = len(sbdata['mean_x'])
_ = np.empty(n_turns)
for key in list(sbdata.keys()):
_[:] = sbdata[key][:]
n_slices, n_turns = sdata['mean_x'].shape
_ = np.empty((n_slices, n_turns))
for key in list(sdata.keys()):
_[:, :] = sdata[key][:, :]
# Particledata
pdata = particledata['Step#0']
n_particles = len(pdata['x'])
n_steps = len(list(particledata.keys()))
_ = np.empty(n_particles)
for i in range(n_steps):
step = 'Step#%d' % i
for key in list(particledata[step].keys()):
_[:] = particledata[step][key][:]
bunchdata.close()
slicedata.close()
particledata.close()
def read_n_plot_data(bfile, sfile, pfile):
bunchdata = hp.File(bfile + '.h5')
slicedata = hp.File(sfile + '.h5')
particledata = hp.File(pfile + '.h5part')
fig = plt.figure(figsize=(16, 16))
ax1 = fig.add_subplot(311)
ax2 = fig.add_subplot(312)
ax3 = fig.add_subplot(313)
ax1.plot(bunchdata['Bunch']['mean_x'][:])
ax2.plot(slicedata['Slices']['mean_x'][:, :])
ax3.plot(particledata['Step#0']['x'][:])
#ax2.plot(slicedata[])
plt.show()
bunchdata.close()
slicedata.close()
particledata.close()
def generate_bunch(n_macroparticles, alpha_x, alpha_y, beta_x, beta_y,
alpha_0, Q_s, R):
intensity = 1.05e11
sigma_z = 0.059958
gamma = 3730.26
eta = alpha_0 - 1. / gamma**2
gamma_t = 1. / np.sqrt(alpha_0)
p0 = np.sqrt(gamma**2 - 1) * m_p * c
beta_z = eta * R / 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
) # WITH OR WITHOUT 4 PIjQuery202047649151738733053_1414145430832?
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 = 2
n_segments = 5
n_macroparticles = 500
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
# ##### Things tested: - Instantiation of the three monitors BunchMonitor, SliceMonitor, ParticleMonitor. - dump(beam) method for all the three. - read data from file. Plot example data from Bunch-, Slice- and Particle-Monitors. - SliceMonitor: does it handle/request slice_sets correctly? - Buffers are on for Bunch- and SliceMonitors. Look at one of the files in hdfview to check the units, attributes, ...
# 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]:
# Instantiate BunchMonitor, SliceMonitor and ParticleMonitor and dump data to file.
bunch = generate_bunch(n_macroparticles, alpha_x_inj, alpha_y_inj,
beta_x_inj, beta_y_inj, alpha_0, Q_s, R)
trans_map = TransverseMap(s, alpha_x, beta_x, D_x, alpha_y, beta_y, D_y,
Q_x, Q_y)
# Slicer config for SliceMonitor.
unibin_slicer = UniformBinSlicer(n_slices=10, n_sigma_z=None, z_cuts=None)
# Monitors
bunch_filename = 'bunch_mon'
slice_filename = 'slice_mon'
particle_filename = 'particle_mon'
bunch_monitor = BunchMonitor(filename=bunch_filename,
n_steps=n_turns,
parameters_dict={'Q_x': Q_x},
write_buffer_every=20)
slice_monitor = SliceMonitor(filename=slice_filename,
n_steps=n_turns,
slicer=unibin_slicer,
parameters_dict={'Q_x': Q_x},
write_buffer_every=20)
particle_monitor = ParticleMonitor(filename=particle_filename,
stride=10,
parameters_dict={'Q_x': Q_x})
arrays_dict = {}
map_ = trans_map
for i in range(n_turns):
for m_ in map_:
m_.track(bunch)
bunch_monitor.dump(bunch)
slice_monitor.dump(bunch)
slice_set_pmon = bunch.get_slices(unibin_slicer)
arrays_dict.update({
'slidx': slice_set_pmon.slice_index_of_particle,
'zz': bunch.z
})
particle_monitor.dump(bunch, arrays_dict)
read_all_data(bunch_filename, slice_filename, particle_filename)
os.remove(bunch_filename + '.h5')
os.remove(slice_filename + '.h5')
os.remove(particle_filename + '.h5part')
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
def generate_one_turn_map(original_map, list_of_optics_elements,
list_of_PyHT_objects, optics_file,
phase_x_col, beta_x_col, phase_y_col, beta_y_col,
circumference, alpha_x=None, beta_x=None, D_x=None,
alpha_y=None, beta_y=None, D_y=None, accQ_x=None,
accQ_y=None, Qp_x=None, Qp_y=None, app_x=0, app_y=0,
app_xy=0, other_detuners=[], use_cython=False):
with open(optics_file) as f:
content = f.readlines()
content = [x.strip() for x in content]
element_coordinates = np.zeros((len(list_of_optics_elements),5))
for j, element in enumerate(list_of_optics_elements):
element_name = element[0]
line_found = False
for i, l in enumerate(content):
s = l.split()
if s[0] == ('"' + element_name + '"'):
line_found = True
element_coordinates[j,0] = j
element_coordinates[j,1] = float(s[phase_x_col])
element_coordinates[j,2] = float(s[beta_x_col])
element_coordinates[j,3] = float(s[phase_y_col])
element_coordinates[j,4] = float(s[beta_y_col])
if line_found == False:
raise ValueError('Element ' + element_name + ' not found from the optics file ' + optics_file)
element_coordinates[:,3] = element_coordinates[:,3]*accQ_x/accQ_y
map_locations = []
for j, element in enumerate(list_of_optics_elements):
if element[1] == 'x':
map_locations.append(element_coordinates[j,1])
elif element[1] == 'y':
map_locations.append(element_coordinates[j,3])
else:
raise ValueError('Unknown plane')
map_locations = np.array(map_locations)
order = np.argsort(map_locations)
s = []
for i, j in enumerate(order):
if i == 0:
if map_locations[j]/accQ_x*circumference != 0.:
s.append(0.)
s.append(map_locations[j]/accQ_x*circumference)
s.append(circumference)
s = np.array(s)
alpha_x = 0.*s
beta_x = 0.*s+beta_x
D_x = 0.*s+D_x
alpha_y = 0.*s
beta_y = 0.*s+beta_y
D_y = 0.*s+D_y
detuners = []
if any(np.atleast_1d(Qp_x) != 0) or \
any(np.atleast_1d(Qp_y) != 0):
detuners.append(Chromaticity(Qp_x, Qp_y))
if app_x != 0 or app_y != 0 or app_xy != 0:
detuners.append(AmplitudeDetuning(app_x, app_y, app_xy))
detuners += other_detuners
transverse_map = TransverseMap(
s=s,
alpha_x=alpha_x, beta_x=beta_x, D_x=D_x,
alpha_y=alpha_y, beta_y=beta_y, D_y=D_y,
accQ_x=accQ_x, accQ_y=accQ_y, detuners=detuners)
transverse_map.n_segments = len(s)-1
transverse_map.name = ['P_%d' % ip for ip in range(len(s)-1)]
transverse_map.name.append('end_ring')
for i_seg, m in enumerate(transverse_map):
m.i0 = i_seg
m.i1 = i_seg+1
m.s0 = transverse_map.s[i_seg]
m.s1 = transverse_map.s[i_seg+1]
m.name0 = transverse_map.name[i_seg]
m.name1 = transverse_map.name[i_seg+1]
m.beta_x0 = transverse_map.beta_x[i_seg]
m.beta_x1 = transverse_map.beta_x[i_seg+1]
m.beta_y0 = transverse_map.beta_y[i_seg]
m.beta_y1 = transverse_map.beta_y[i_seg+1]
# insert transverse map in the ring
one_turn_map = []
for m in original_map:
if type(m) != PyHEADTAIL.trackers.transverse_tracking.TransverseSegmentMap:
one_turn_map.append(m)
for i, m in enumerate(transverse_map):
one_turn_map.append(m)
if i < len(order):
element_idx = order[i]
one_turn_map.append(list_of_PyHT_objects[element_idx])
return one_turn_map