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)
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_tracking_with_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
with detuners
'''
adetuner = AmplitudeDetuning(1e-2, 5e-2, 1e-3)
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,
adetuner,
printer=SilentPrinter())
beam_p = self.create_bunch()
for s in pure_python_map:
s.track(beam_p)
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)
file2.close()
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)
Jy = 1 / 2 * ((1 + pp.alpha_y**2) / pp.beta_y * bunch.y**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(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_)
convert_hanarmonicities = True
if convert_hanarmonicities:
scale_factor = 2 * bunch.p0
app_x = pp.app_x * scale_factor
app_y = pp.app_y * scale_factor
app_xy = pp.app_xy * scale_factor
else:
app_x = pp.app_x
app_y = pp.app_y
app_xy = pp.app_xy
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(app_x, app_y, app_xy)])
longitudinal_map = LinearMap([pp.alpha], pp.circumference, pp.Q_s)
one_turn_map = [transverse_map[0]] + [longitudinal_map]
# 3. CREATE LISTS TO SAVE THE TBT DATA
X = []
Y = []
DP = []
# 4. SET UP THE ACCELERATOR AND START TRACKING
for i in range(n_turns):
# bunch.yp += ampKicks[i] * np.sin(2 * np.pi * 400e6 / (bunch.beta * c) * bunch.z)
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
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():
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.Qs))
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)
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)
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)
# print ('test NOT passed. No error raised!')
except ValueError as exc:
# print ('test passed.\n')
# print ('Error message:\n' + exc.message)
pass
