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
Ejemplo n.º 2
0
 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
Ejemplo n.º 3
0
    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)
Ejemplo n.º 5
0
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)
Ejemplo n.º 6
0
    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)
Ejemplo n.º 7
0
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)
Ejemplo n.º 8
0
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_)
Ejemplo n.º 9
0
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')
Ejemplo n.º 10
0
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
Ejemplo n.º 12
0
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