class Ring(object):
    r""" Class containing the general properties of the synchrotron that are
    independent of the RF system or the beam.

    The index :math:`n` denotes time steps, :math:`k` ring segments/sections
    and :math:`i` momentum compaction orders.

    Parameters
    ----------
    ring_length : float (opt: float array [n_sections])
        Length [m] of the n_sections ring segments of the synchrotron.
        An RF station, a synchrotron radiation kick, and/or an impedance kick
        can be included at the end of each ring section.
    alpha_0 : float (opt: float array/matrix [n_sections, n_turns+1])
        Momentum compaction factor of zeroth order :math:`\alpha_{0,k,i}` [1];
        can be input as single float or as a program of (n_turns + 1) turns
        (should be of the same size as synchronous_data).
        In case of higher order momentum compaction, check the
        documentation for the inputs: alpha_order, alpha_1, alpha_2
    synchronous_data : float (opt: float array/matrix [n_sections, n_turns+1])
        Design synchronous particle momentum (default) [eV], kinetic or
        total energy [eV] or bending field [T] on the design orbit.
        Input for each RF section :math:`p_{s,k,n}`.
        Can be input as a single constant float, or as a
        program of (n_turns + 1) turns. In case of several sections without
        acceleration, input: [[momentum_section_1], [momentum_section_2],
        etc.]. In case of several sections with acceleration, input:
        [momentum_program_section_1, momentum_program_section_2, etc.]. Can
        be input also as a tuple of time and momentum, see also
        'cycle_time' and 'PreprocessRamp'
    Particle : class
        A Particle-based class defining the primary, synchronous particle (mass
        and charge) that is reference for the momentum/energy in the ring.
    n_turns : int
        Optional: Number of turns :math:`n` [1] to be simulated.
        If a synchronous_data program is passed as a tuple (see below),
        the number of turns will be overwritten depending on the length in time
        of the program
    synchronous_data_type : str
        Optional: Choice of 'synchronous_data' type; can be 'momentum'
        (default), 'total energy', 'kinetic energy' or 'bending field'
        (requires bending_radius to be defined)
    bending_radius : float
        Optional: Radius [m] of the bending magnets,
        required if 'bending field' is set for the synchronous_data_type
    n_sections : int
        Optional: number of ring sections/segments; default is 1
    alpha_1 : float (opt: float array/matrix [n_sections, n_turns+1])
        Momentum compaction factor of first order
        :math:`\alpha_{1,k,i}` [1]; can be input as single float or as a
        program of (n_turns + 1) turns (should be of the same size as
        synchronous_data and alpha_0).
    alpha_2 : float (opt: float array/matrix [n_sections, n_turns+1])
        Optional : Momentum compaction factor of second order
        :math:`\alpha_{2,k,i}` [1]; can be input as single float or as a
        program of (n_turns + 1) turns (should be of the same size as
        synchronous_data and alpha_0).
    RingOptions : class
        Optional : A RingOptions-based class with default options to check the
        input and initialize the momentum program for the simulation.
        This object defines the interpolation scheme, plotting options, etc.
        The options for this object can be adjusted and passed to the Ring
        object.

    Attributes
    ----------
    ring_circumference : float
        Circumference of the synchrotron. Sum of ring segment lengths,
        :math:`C = \sum_k L_k` [m]
    ring_radius : float
        Radius of the synchrotron, :math:`R = C/(2 \pi)` [m]
    bending_radius : float
        Bending radius in dipole magnets, :math:`\rho` [m]
    alpha_order : int
        Highest order of momentum compaction (as defined by the input). Can
        be 0,1,2.
    alpha_0 : float matrix [n_sections, n_turns+1]
        Zeroth order momentum compaction factor
        :math:`\alpha_{0,k,n}`
    alpha_1 : float matrix [n_sections, n_turns+1]
        First order momentum compaction factor
        :math:`\alpha_{1,k,n}`
    alpha_2 : float matrix [n_sections, n_turns+1]
        Second order momentum compaction factor
        :math:`\alpha_{2,k,n}`
    eta_0 : float matrix [n_sections, n_turns+1]
        Zeroth order slippage factor :math:`\eta_{0,k,n} = \alpha_{0,k,n} -
        \frac{1}{\gamma_{s,k,n}^2}` [1]
    eta_1 : float matrix [n_sections, n_turns+1]
        First order slippage factor :math:`\eta_{1,k,n} =
        \frac{3\beta_{s,k,n}^2}{2\gamma_{s,k,n}^2} + \alpha_{1,k,n} -
        \alpha_{0,k,n}\eta_{0,k,n}` [1]
    eta_2 : float matrix [n_sections, n_turns+1]
        Second order slippage factor :math:`\eta_{2,k,n} =
        -\frac{\beta_{s,k,n}^2\left(5\beta_{s,k,n}^2-1\right)}
        {2\gamma_{s,k,n}^2} + \alpha_{2,k,n} - 2\alpha_{0,k,n}\alpha_{1,k,n}
        + \frac{\alpha_{1,k,n}}{\gamma_{s,k,n}^2} + \alpha_{0,k}^2\eta_{0,k,n}
        - \frac{3\beta_{s,k,n}^2\alpha_{0,k,n}}{2\gamma_{s,k,n}^2}` [1]
    momentum : float matrix [n_sections, n_turns+1]
        Synchronous relativistic momentum on the design orbit :math:`p_{s,k,n}`
    beta : float matrix [n_sections, n_turns+1]
        Synchronous relativistic beta program for each segment of the
        ring :math:`\beta_{s,k}^n = \frac{1}{\sqrt{1
        + \left(\frac{m}{p_{s,k,n}}\right)^2} }` [1]
    gamma : float matrix [n_sections, n_turns+1]
        Synchronous relativistic gamma program for each segment of the ring
        :math:`\gamma_{s,k,n} = \sqrt{ 1
        + \left(\frac{p_{s,k,n}}{m}\right)^2 }` [1]
    energy : float matrix [n_sections, n_turns+1]
        Synchronous total energy program for each segment of the ring
        :math:`E_{s,k,n} = \sqrt{ p_{s,k,n}^2 + m^2 }` [eV]
    kin_energy : float matrix [n_sections, n_turns+1]
        Synchronous kinetic energy program for each segment of the ring
        :math:`E_{s,kin} = \sqrt{ p_{s,k,n}^2 + m^2 } - m` [eV]
    delta_E : float matrix [n_sections, n_turns+1]
        Derivative of synchronous total energy w.r.t. time, for all sections,
        :math:`: \quad E_{s,k,n+1}- E_{s,k,n}` [eV]
    t_rev : float array [n_turns+1]
        Revolution period turn by turn.
        :math:`T_{0,n} = \frac{C}{\beta_{s,n} c}` [s]
    f_rev : float array [n_turns+1]
        Revolution frequency :math:`f_{0,n} = \frac{1}{T_{0,n}}` [Hz]
    omega_rev : float array [n_turns+1]
        Revolution angular frequency :math:`\omega_{0,n} = 2\pi f_{0,n}` [1/s]
    cycle_time : float array [n_turns+1]
        Cumulative cycle time, turn by turn, :math:`t_n = \sum_n T_{0,n}` [s].
        Possibility to extract cycle parameters at these moments using
        'parameters_at_time'.
    RingOptions : RingOptions()
        The RingOptions is kept as an attribute of the Ring object for further
        usage.

    Examples
    --------
    >>> # To declare a single-section synchrotron at constant energy:
    >>> # Particle type Proton
    >>> from beam.beam import Proton
    >>> from input_parameters.ring import Ring
    >>>
    >>> n_turns = 10
    >>> C = 26659
    >>> alpha_0 = 3.21e-4
    >>> momentum = 450e9
    >>> ring = Ring(C, alpha_0, momentum, Proton(), n_turns)
    >>>
    >>>
    >>> # To declare a two section synchrotron at constant energy and
    >>> # higher-order momentum compaction factors; particle Electron:
    >>> from beam.beam import Electron
    >>> from input_parameters.ring import Ring
    >>>
    >>> n_turns = 10
    >>> C = [13000, 13659]
    >>> alpha_0 = [[3.21e-4], [2.89e-4]]  # or [3.21e-4, 2.89e-4]
    >>> alpha_1 = [[2.e-5], [1.e-5]]  # or [2.e-5, 1.e-5]
    >>> alpha_2 = [[5.e-7], [5.e-7]]  # or [5.e-7, 5.e-7]
    >>> momentum = 450e9
    >>> ring = Ring(C, alpha_0, momentum, Electron(), n_turns,
    >>>             alpha_1=alpha_1, alpha_2=alpha_2)

    """

    def __init__(self, ring_length, alpha_0, synchronous_data, Particle,
                 n_turns=1, synchronous_data_type='momentum',
                 bending_radius=None, n_sections=1, alpha_1=None, alpha_2=None,
                 RingOptions=RingOptions()):

        # Conversion of initial inputs to expected types
        self.n_turns = int(n_turns)
        self.n_sections = int(n_sections)

        # Ring length and checks
        self.ring_length = np.array(ring_length, ndmin=1, dtype=float)
        self.ring_circumference = np.sum(self.ring_length)
        self.ring_radius = self.ring_circumference/(2*np.pi)

        if bending_radius is not None:
            self.bending_radius = float(bending_radius)
        else:
            self.bending_radius = bending_radius

        if self.n_sections != len(self.ring_length):
            #InputDataError
            raise RuntimeError("ERROR in Ring: Number of sections and ring " +
                               "length size do not match!")

        # Primary particle mass and charge used for energy calculations
        self.Particle = Particle

        # Keeps RingOptions as an attribute
        self.RingOptions = RingOptions

        # Reshaping the input synchronous data to the adequate format and
        # get back the momentum program from RingOptions
        self.momentum = RingOptions.reshape_data(
            synchronous_data,
            self.n_turns,
            self.n_sections,
            input_to_momentum=True,
            synchronous_data_type=synchronous_data_type,
            mass=self.Particle.mass,
            charge=self.Particle.charge,
            circumference=self.ring_circumference,
            bending_radius=self.bending_radius)

        # Updating the number of turns in case it was changed after ramp
        # interpolation
        if self.momentum.shape[1] != (self.n_turns+1):
            self.n_turns = self.momentum.shape[1] - 1
            warnings.warn("WARNING in Ring: The number of turns for the " +
                          "simulation was changed by passing a momentum " +
                          "program.")

        # Derived from momentum
        self.beta = np.sqrt(1/(1 + (self.Particle.mass/self.momentum)**2))
        self.gamma = np.sqrt(1 + (self.momentum/self.Particle.mass)**2)
        self.energy = np.sqrt(self.momentum**2 + self.Particle.mass**2)
        self.kin_energy = np.sqrt(self.momentum**2 + self.Particle.mass**2) - \
            self.Particle.mass
        self.delta_E = np.diff(self.energy, axis=1)
        self.t_rev = np.dot(self.ring_length, 1/(self.beta*c))
        self.cycle_time = np.cumsum(self.t_rev)  # Always starts with zero
        self.f_rev = 1/self.t_rev
        self.omega_rev = 2*np.pi*self.f_rev

        # Momentum compaction, checks, and derived slippage factors
        if RingOptions.t_start is None:
            interp_time = self.cycle_time
        else:
            interp_time = self.cycle_time+RingOptions.t_start

        self.alpha_0 = RingOptions.reshape_data(
            alpha_0, self.n_turns, self.n_sections,
            interp_time=interp_time)
        self.alpha_order = 0

        if alpha_1 is not None:
            self.alpha_1 = RingOptions.reshape_data(
                alpha_1, self.n_turns, self.n_sections,
                interp_time=interp_time)
            self.alpha_order = 1
        else:
            # Filling alpha_1 with zeros
            # This can be removed when the BLonD assembler is in place
            # to avoid high order momentum compaction programs filled
            # with zeros (should be propagated in RFStation.__init__())
            self.alpha_1 = np.zeros(self.alpha_0.shape)

        if alpha_2 is not None:
            self.alpha_2 = RingOptions.reshape_data(
                alpha_2, self.n_turns, self.n_sections,
                interp_time=interp_time)
            self.alpha_order = 2
        else:
            # Filling alpha_2 with zeros
            # This can be removed when the BLonD assembler is in place
            # to avoid high order momentum compaction programs filled
            # with zeros (should be propagated in RFStation.__init__())
            self.alpha_2 = np.zeros(self.alpha_0.shape)

        # Slippage factor derived from alpha, beta, gamma
        self.eta_generation()

    def eta_generation(self):
        """ Function to generate the slippage factors (zeroth, first, and
        second orders, see [1]_) from the momentum compaction and the
        relativistic beta and gamma program through the cycle.

        References
        ----------
        .. [1] "Accelerator Physics," S. Y. Lee, World Scientific,
                Third Edition, 2012.
        """

        for i in range(self.alpha_order+1):
            getattr(self, '_eta' + str(i))()

        # Fill unused eta arrays with zeros
        # This can be removed when the BLonD assembler is in place
        # to avoid high order momentum compaction programs filled
        # with zeros (should be propagated in RFStation.__init__())
        for i in range(self.alpha_order+1, 3):
            setattr(self, "eta_%s" % i, np.zeros([self.n_sections,
                                                  self.n_turns+1]))

    def _eta0(self):
        """ Function to calculate the zeroth order slippage factor eta_0 """

        self.eta_0 = np.empty([self.n_sections, self.n_turns+1])
        for i in range(0, self.n_sections):
            self.eta_0[i] = self.alpha_0[i] - self.gamma[i]**(-2.)

    def _eta1(self):
        """ Function to calculate the first order slippage factor eta_1 """

        self.eta_1 = np.empty([self.n_sections, self.n_turns+1])
        for i in range(0, self.n_sections):
            self.eta_1[i] = 3*self.beta[i]**2/(2*self.gamma[i]**2) + \
                self.alpha_1[i] - self.alpha_0[i]*self.eta_0[i]

    def _eta2(self):
        """ Function to calculate the second order slippage factor eta_2 """

        self.eta_2 = np.empty([self.n_sections, self.n_turns+1])
        for i in range(0, self.n_sections):
            self.eta_2[i] = - self.beta[i]**2*(5*self.beta[i]**2 - 1) / \
                (2*self.gamma[i]**2) + self.alpha_2[i] - 2*self.alpha_0[i] *\
                self.alpha_1[i] + self.alpha_1[i] / self.gamma[i]**2 + \
                self.alpha_0[i]**2*self.eta_0[i] - 3*self.beta[i]**2 * \
                self.alpha_0[i]/(2*self.gamma[i]**2)

    def parameters_at_time(self, cycle_moments):
        """ Function to return various cycle parameters at a specific moment in
        time. The cycle time is defined to start at zero in turn zero.

        Parameters
        ----------
        cycle_moments : float array
            Moments of time at which cycle parameters are to be calculated [s].

        Returns
        -------
        parameters : dictionary
            Contains 'momentum', 'beta', 'gamma', 'energy', 'kin_energy',
            'f_rev', 't_rev'. 'omega_rev', 'eta_0', and 'delta_E' interpolated
            to the moments contained in the 'cycle_moments' array

        """

        parameters = {}
        parameters['momentum'] = np.interp(cycle_moments, self.cycle_time,
                                           self.momentum[0])
        parameters['beta'] = np.interp(cycle_moments, self.cycle_time,
                                       self.beta[0])
        parameters['gamma'] = np.interp(cycle_moments, self.cycle_time,
                                        self.gamma[0])
        parameters['energy'] = np.interp(cycle_moments, self.cycle_time,
                                         self.energy[0])
        parameters['kin_energy'] = np.interp(cycle_moments, self.cycle_time,
                                             self.kin_energy[0])
        parameters['f_rev'] = np.interp(cycle_moments, self.cycle_time,
                                        self.f_rev)
        parameters['t_rev'] = np.interp(cycle_moments, self.cycle_time,
                                        self.t_rev)
        parameters['omega_rev'] = np.interp(cycle_moments, self.cycle_time,
                                            self.omega_rev)
        parameters['eta_0'] = np.interp(cycle_moments, self.cycle_time,
                                        self.eta_0[0])
        parameters['delta_E'] = np.interp(cycle_moments,
                                          self.cycle_time[1:],
                                          np.diff(self.energy[0]))

        return parameters class RFStation(object):
    r""" Class containing all the RF parameters for all the RF systems in one
    ring segment or RF station.

    **How to use RF programs:**

    * For 1 RF system and constant values of V, h, or phi, input a single value
    * For 1 RF system and varying values of V, h, or phi, input an array of
      n_turns+1 values
    * For several RF systems and constant values of V, h, or phi, input lists
      of single values
    * For several RF systems and varying values of V, h, or phi, input lists
      of arrays of n_turns+1 values
    * For pre-processing, pass a list of times-voltages, times-harmonics,
      and/or times-phases for **each** RF system as a tuple
      ((time_1, voltage_1), (time_2, voltage_2), ...)

    Optional: RF frequency other than the design frequency. In this case, need
    to use a beam phase loop for correct RF phase!

    Optional: empty RFStation (e.g. for machines with synchrotron radiation);
    input voltage as 0.

    The index :math:`n` denotes time steps, :math:`l` the index of the RF
    systems in the section.

    **N.B. for negative eta the RF phase has to be shifted by Pi w.r.t the time
    reference.**

    Parameters
    ----------
    Ring : class
        A Ring type class
    harmonic : float (opt: float array/matrix, tuple of float array/matrix)
        Harmonic number of the RF system, :math:`h_{l,n}` [1]. For input
        options, see above
    voltage : float (opt: float array/matrix, tuple of float array/matrix)
        RF cavity voltage as seen by the beam, :math:`V_{l,n}` [V]. For input
        options, see above
    phi_rf_d : float (opt: float array/matrix, tuple of float array/matrix)
        Programmed/designed RF cavity phase,
        :math:`\phi_{d,l,n}` [rad]. For input options, see above
    n_rf : int
        Optional, Number of harmonic rf systems in the section :math:`l`.
        Becomes mandatory for several rf systems.
    section_index : int
        Optional, In case of several sections in the Ring object, this
        specifies after which section the rf station is located (to get the
        right momentum program etc.). Value should be in the range
        1..Ring.n_sections
    omega_rf : float (opt: float array/matrix)
        Optional, Sets the rf angular frequency program that does not follow
        the harmonic condition. For input options, see above.
    phi_noise : float (opt: float array/matrix)
        Optional, programmed RF cavity phase noise, :math:`\phi_{N,l,n}` [rad].
        Added to all RF systems in the station. For input options, see above
    phi_modulation : class (opt: iterable of classes)
        A PhaseModulation type class (or iterable of classes)
    RFStationOptions : class
        Optionnal, A RFStationOptions-based class defining smoothing,
        interpolation, etc. options for harmonic, voltage, and/or
        phi_rf_d programme to be interpolated to a turn-by-turn programme

    Attributes
    ----------
    counter : int
        Counter of the current simulation time step; defined as a list in
        order to be passed by reference
    section_index : int
        Unique index :math:`k` of the RF station the present class is defined
        for. Input in the range 1..n_sections (see
        :py:class:`input_parameters.ring.Ring`).
        Inside the code, indices 0..n_sections-1 are used.
    Particle : class
        Inherited from
        :py:attr:`input_parameters.ring.Ring.Particle`
    n_turns : int
        Inherited from
        :py:attr:`input_parameters.ring.Ring.n_turns`
    ring_circumference : float
        Inherited from
        :py:attr:`input_parameters.ring.Ring.ring_circumference`
    section_length : float
        Length :math:`L_k` of the RF section; inherited from
        :py:attr:`input_parameters.ring.Ring.ring_length`
    length_ratio : float
        Fractional RF section length :math:`L_k/C`
    t_rev : float array [n_turns+1]
        Inherited from
        :py:attr:`input_parameters.ring.Ring.t_rev`
    momentum : float array [n_turns+1]
        Momentum program of the present RF section; inherited from
        :py:attr:`input_parameters.ring.Ring.momentum`
    beta : float array [n_turns+1]
        Relativistic beta of the present RF section; inherited from
        :py:attr:`input_parameters.ring.Ring.beta`
    gamma : float array [n_turns+1]
        Relativistic gamma of the present RF section; inherited from
        :py:attr:`input_parameters.ring.Ring.gamma`
    energy : float array [n_turns+1]
        Total energy of the present RF section; inherited from
        :py:attr:`input_parameters.ring.Ring.energy`
    delta_E : float array [n_turns+1]
        Time derivative of total energy of the present section; inherited from
        :py:attr:`input_parameters.ring.Ring.delta_E`
    alpha_order : int
        Inherited from
        :py:attr:`input_parameters.ring.Ring.alpha_order`
    charge : int
        Inherited from
        :py:attr:`beam.Particle.charge`
    alpha_0 : float array [n_turns+1]
        Zeroth order momentum compaction factor of the present section;
        inherited from
        :py:attr:`input_parameters.ring.Ring.alpha_0`
    alpha_1 : float array [n_turns+1]
        First order momentum compaction factor of the present section;
        inherited from
        :py:attr:`input_parameters.ring.Ring.alpha_1`
    alpha_2 : float array [n_turns+1]
        Second order momentum compaction factor of the present section;
        inherited from
        :py:attr:`input_parameters.ring.Ring.alpha_2`
    eta_0 : float array [n_turns+1]
        Zeroth order slippage factor of the present section; inherited from
        :py:attr:`input_parameters.ring.Ring.eta_0`
    eta_1 : float array [n_turns+1]
        First order slippage factor of the present section; inherited from
        :py:attr:`input_parameters.ring.Ring.eta_1`
    eta_2 : float array [n_turns+1]
        Second order slippage factor of the present section; inherited from
        :py:attr:`input_parameters.ring.Ring.eta_2`
    sign_eta_0 : float array
        Sign of the eta_0 array
    harmonic : float matrix [n_rf, n_turns+1]
        Harmonic number for each rf system,
        :math:`h_{l,n}` [1]
    voltage : float matrix [n_rf, n_turns+1]
        Actual rf voltage of each harmonic system,
        :math:`V_{rf,l,n}` [V]
    empty : bool
        Flag to specify if the RFStation is empty
    phi_rf_d : float matrix [n_rf, n_turns+1]
        Designed rf cavity phase of each harmonic system,
        :math:`\phi_{d,l,n}` [rad]
    phi_rf : float matrix [n_rf, n_turns+1]
        Actual RF cavity phase of each harmonic system used for tracking,
        :math:`\phi_{rf,l,n}` [rad]. Initially the same as the designed phase.
    omega_rf_d : float matrix [n_rf, n_turns+1]
        Design RF angular frequency of the RF systems in the station
        :math:`\omega_{d,l,n} = \frac{h_{l,n} \beta_{l,n} c}{R_{s,n}}` [Hz]
    omega_rf : float matrix [n_rf, n_turns+1]
        Actual RF angular frequency of the RF systems in the station
        :math:`\omega_{rf,l,n} = \frac{h_{l,n} \beta_{l,n} c}{R_{s,n}}` [Hz].
        Initially the same as the designed angular frequency.
    phi_noise : None or float matrix [n_rf, n_turns+1]
        Programmed cavity phase noise for each RF harmonic.
    phi_modulation : None or float matrix [n_rf, n_turns+1]
        Programmed cavity phase modulation for each RF harmonic.
    dphi_rf : float matrix [n_rf]
        Accumulated RF phase error of each harmonic system
        :math:`\Delta \phi_{rf,l,n}` [rad]
    t_rf : float matrix [n_rf, n_turns+1]
        RF period :math:`\frac{2 \pi}{\omega_{rf,l,n}}` [s]
    phi_s : float array [n_turns+1]
        Synchronous phase for this section (if not empty), calculated in
        :py:func:`input_parameters.rf_parameters.calculate_phi_s`
    Q_s : float array [n_turns+1]
        Synchrotron tune for this section (if not empty), calculated in
        :py:func:`input_parameters.rf_parameters.calculate_Q_s`
    omega_s0 : float array [n_turns+1]
        Central synchronous angular frequency corresponding to Q_s (single
        harmonic, no intensity effects)
        :math:`\omega_{s,0} = Q_s \omega_{\text{rev}}` [1/s], where
        :math:`\omega_{\text{rev}}` is defined in
        :py:class:`input_parameters.ring.Ring`)
    RFStationOptions : RFStationOptions()
        The RFStationOptions is kept as an attribute of the RFStationg object
        for further usage.


    Examples
    --------
    >>> # To declare a double-harmonic RF system for protons:
    >>>
    >>> n_turns = 10
    >>> C = 26659
    >>> alpha = 3.21e-4
    >>> momentum = 450e9
    >>> ring = Ring(n_turns, C, alpha, momentum)
    >>> rf_station = RFStation(ring, [35640, 71280], [6e6, 6e5], [0, 0], 2)

    """

    def __init__(self, Ring, harmonic, voltage, phi_rf_d, n_rf=1,
                 section_index=1, omega_rf=None, phi_noise=None,
                 phi_modulation=None, RFStationOptions=RFStationOptions()):

        # Different indices
        self.counter = [int(0)]
        self.section_index = int(section_index - 1)
        if self.section_index < 0 \
                or self.section_index > Ring.n_sections - 1:
            raise RuntimeError("ERROR in RFStation: section_index out of" +
                               " allowed range!")
        self.n_rf = int(n_rf)

        # Imported from Ring
        self.Particle = Ring.Particle
        self.n_turns = Ring.n_turns
        self.ring_circumference = Ring.ring_circumference
        self.section_length = Ring.ring_length[self.section_index]
        self.length_ratio = float(self.section_length/self.ring_circumference)
        self.t_rev = Ring.t_rev
        self.momentum = Ring.momentum[self.section_index]
        self.beta = Ring.beta[self.section_index]
        self.gamma = Ring.gamma[self.section_index]
        self.energy = Ring.energy[self.section_index]
        self.delta_E = Ring.delta_E[self.section_index]
        self.alpha_order = Ring.alpha_order
        self.charge = self.Particle.charge

        # The order alpha_order used here can be replaced by Ring.alpha_order
        # when the assembler can differentiate the cases 'simple' and 'exact'
        # for the drift
        alpha_order = 2
        for i in range(alpha_order+1):
            dummy = getattr(Ring, 'eta_' + str(i))
            setattr(self, "eta_%s" % i, dummy[self.section_index])
            dummy = getattr(Ring, 'alpha_' + str(i))
            setattr(self, "alpha_%s" % i, dummy[self.section_index])
        self.sign_eta_0 = np.sign(self.eta_0)

        # Reshape input rf programs
        # Reshape design harmonic
        self.harmonic = RFStationOptions.reshape_data(harmonic,
                                                      self.n_turns,
                                                      self.n_rf,
                                                      Ring.cycle_time,
                                                      Ring.RingOptions.t_start)
        # Reshape design voltage
        self.voltage = RFStationOptions.reshape_data(voltage,
                                                     self.n_turns,
                                                     self.n_rf,
                                                     Ring.cycle_time,
                                                     Ring.RingOptions.t_start)

        # Checking if the RFStation is empty
        if np.sum(self.voltage) == 0:
            self.empty = True
        else:
            self.empty = False

        # Reshape design phase
        self.phi_rf_d = RFStationOptions.reshape_data(phi_rf_d,
                                                      self.n_turns,
                                                      self.n_rf,
                                                      Ring.cycle_time,
                                                      Ring.RingOptions.t_start)

        # Calculating design rf angular frequency
        if omega_rf is None:
            self.omega_rf_d = 2.*np.pi*self.beta*c*self.harmonic / \
                (self.ring_circumference)
        else:
            self.omega_rf_d = RFStationOptions.reshape_data(
                omega_rf,
                self.n_turns,
                self.n_rf,
                Ring.cycle_time,
                Ring.RingOptions.t_start)

        # Reshape phase noise
        if phi_noise is not None:
            self.phi_noise = RFStationOptions.reshape_data(
                phi_noise,
                self.n_turns,
                self.n_rf,
                Ring.cycle_time,
                Ring.RingOptions.t_start)
        else:
            self.phi_noise = None
            
        if phi_modulation is not None:
            
            try:
                iter(phi_modulation)
            except TypeError:
                phi_modulation = [phi_modulation]
            
            dPhi = np.zeros([self.n_rf, self.n_turns+1])
            dOmega = np.zeros([self.n_rf, self.n_turns+1])
            for pMod in phi_modulation:
                system = np.where(self.harmonic[:,0] == pMod.harmonic)[0]
                if len(system) == 0:
                    raise ValueError("No matching harmonic in phi_modulation")
                elif len(system) > 1:
                    raise RuntimeError("""Phase modulation not yet 
                                       implemented with multiple systems 
                                       at the same harmonic.""")
                else:
                    system = system[0]
                    
                pMod.calc_modulation()
                pMod.calc_delta_omega((Ring.cycle_time, self.omega_rf_d[system]))
                dPhiInput, dOmegaInput =  pMod.extend_to_n_rf(self.harmonic[:,0])
                dPhi += RFStationOptions.reshape_data(dPhiInput,
                                                     self.n_turns,
                                                     self.n_rf,
                                                     Ring.cycle_time,
                                                     Ring.RingOptions.t_start)
                dOmega += RFStationOptions.reshape_data(dOmegaInput,
                                                       self.n_turns,
                                                       self.n_rf,
                                                       Ring.cycle_time,
                                                       Ring.RingOptions.t_start)
                
                
            
            self.phi_modulation = (dPhi, dOmega)
        else:
            
            self.phi_modulation = None

        # Copy of the desing rf programs in the one used for tracking
        # and that can be changed by feedbacks
        self.phi_rf = np.array(self.phi_rf_d)
        self.dphi_rf = np.zeros(self.n_rf)
        self.omega_rf = np.array(self.omega_rf_d)
        self.t_rf = 2*np.pi / self.omega_rf

        # From helper functions
        if not self.empty:
            self.phi_s = calculate_phi_s(self, self.Particle)
            self.Q_s = calculate_Q_s(self, self.Particle)
            self.omega_s0 = self.Q_s*Ring.omega_rev

    def eta_tracking(self, beam, counter, dE):
        r"""Function to calculate the slippage factor as a function of the
        energy offset :math:`\Delta E` of the particle. The slippage factor
        of the :math:`i` th order is :math:`\eta(\delta) = \sum_{i}(\eta_i \,
        \delta^i) = \sum_{i} \left(\eta_i \, \left[ \frac{\Delta E}
        {\beta_s^2 E_s} \right]^i \right)`

        """

        if self.alpha_order == 0:
            return self.eta_0[counter]
        else:
            eta = 0
            delta = dE/(beam.beta**2 * beam.energy)
            for i in range(self.alpha_order+1):
                eta_i = getattr(self, 'eta_' + str(i))[counter]
                eta += eta_i * (delta**i)
            return eta class RingAndRFTracker(object):
    r""" Class taking care of basic particle coordinate tracking for a given
    RF station and the part of the ring until the next station, see figure.

    .. image:: ring_and_RFstation.png
        :align: center
        :width: 600
        :height: 600

    The time step is fixed to be one turn, but the tracking can consist of
    multiple RingAndRFTracker objects. In this case, the user should make sure
    that the lengths of the stations sum up exactly to the circumference or use
    the FullRingAndRF object in order to let the code pre-process the
    parameters. Each RF station may contain several RF harmonic systems which
    are considered to be in the same location. First, the energy kick of the RF
    station is applied, and then the particle arrival time to the next station
    is updated. The change in RF phase, voltage, and frequency due to control
    loops is tracked as well.

    Parameters
    ----------
    RFStation : class
        A RFStation type class
    counter : [int]
        Inherited from
        :py:attr:`input_parameters.rf_parameters.RFStation.counter`
    length_ratio : float
        Inherited from
        :py:attr:`input_parameters.ring.Ring.length_ratio`
    section_length : float
        Inherited from
        :py:attr:`input_parameters.ring.Ring.section_length`
    t_rev : float
        Inherited from
        :py:attr:`input_parameters.ring.Ring.t_rev`
    n_rf : float
        Inherited from
        :py:attr:`input_parameters.rf_parameters.RFStation.n_rf`
    beta : float
        Inherited from
        :py:attr:`input_parameters.ring.Ring.beta`
    charge : float
        Inherited from
        :py:attr:`input_parameters.ring.Ring.Particle.charge`
    harmonic : float array
        Inherited from
        :py:attr:`input_parameters.rf_parameters.RFStation.harmonic`
    voltage : float array
        Inherited from
        :py:attr:`input_parameters.rf_parameters.RFStation.voltage`
    phi_noise : float array
        Inherited from
        :py:attr:`input_parameters.rf_parameters.RFStation.phi_noise`
    phi_modulation : 2-tuple of float array
        Inherited from
        :py:attr:`input_parameters.rf_parameters.RFStation.phi_modulation`
    phi_rf : float array
        Inherited from
        :py:attr:`input_parameters.rf_parameters.RFStation.phi_rf`
    phi_s : float array
        Inherited from
        :py:attr:`input_parameters.rf_parameters.RFStation.phi_s`
    alpha_0 : float array
        Inherited from
        :py:attr:`input_parameters.ring.Ring.alpha_0`
    alpha_1 : float array
        Inherited from
        :py:attr:`input_parameters.ring.Ring.alpha_1`
    alpha_2 : float array
        Inherited from
        :py:attr:`input_parameters.ring.Ring.alpha_2`
    eta_0 : float array
        Inherited from
        :py:attr:`input_parameters.ring.Ring.eta_0`
    eta_1 : float array
        Inherited from
        :py:attr:`input_parameters.ring.Ring.eta_1`
    eta_2 : float array
        Inherited from
        :py:attr:`input_parameters.ring.Ring.eta_2`
    alpha_order : float array
        Inherited from
        :py:attr:`input_parameters.ring.Ring.alpha_order`
    acceleration_kick : float array
        Inherited from
        :py:attr:`input_parameters.ring.Ring.delta_E`
        and multiplied by -1
    Beam : class
        A Beam type class
    solver : str
        Type of solver used for the drift equation; use 'simple' for 1st order
        approximation and 'exact' for exact solver
    BeamFeedback : class (optional)
        A BeamFeedback type class, beam-based feedback on RF frequency;
        default is None
    NoiseFeedback : class (optional)
        A NoiseFeedback type class, bunch-length feedback on RF noise;
        default is None
    periodicity : bool (optional)
        Option to switch periodic solver on/off; default is False (off)
    interpolation : bool (optional)
        Option to use sliced and interpolated voltage for the kicker; default
        is False

    """

    def __init__(self, RFStation, Beam, solver='simple', BeamFeedback=None,
                 NoiseFeedback=None, CavityFeedback=None, periodicity=False,
                 interpolation=False, Profile=None, TotalInducedVoltage=None):

        # Set up logging
        # self.logger = logging.getLogger(__class__.__name__)
        # self.logger.info("Class initialized")

        # Imports from RF parameters
        self.rf_params = RFStation
        self.counter = RFStation.counter
        self.length_ratio = RFStation.length_ratio
        self.section_length = RFStation.section_length
        self.t_rev = RFStation.t_rev
        self.n_rf = RFStation.n_rf
        self.beta = RFStation.beta
        self.charge = RFStation.Particle.charge
        self.harmonic = RFStation.harmonic
        self.voltage = RFStation.voltage
        self.phi_noise = RFStation.phi_noise
        self.phi_modulation = RFStation.phi_modulation
        self.phi_rf = RFStation.phi_rf
        if not self.rf_params.empty:
            self.phi_s = RFStation.phi_s
        self.omega_rf = RFStation.omega_rf
        self.alpha_0 = RFStation.alpha_0
        self.alpha_1 = RFStation.alpha_1
        self.alpha_2 = RFStation.alpha_2
        self.eta_0 = RFStation.eta_0
        self.eta_1 = RFStation.eta_1
        self.eta_2 = RFStation.eta_2
        self.alpha_order = RFStation.alpha_order
        self.acceleration_kick = - RFStation.delta_E

        # Other imports
        self.beam = Beam
        self.solver = str(solver)
        if self.solver not in ['simple', 'exact', 'legacy']:
            # SolverError
            raise RuntimeError("ERROR in RingAndRFTracker: Choice of" +
                               " longitudinal solver not recognised!")
        if self.alpha_order > 1:  # Force exact solver for higher orders of eta
            self.solver = 'exact'
        self.solver = self.solver.encode(encoding='utf_8')

        # Options
        self.beamFB = BeamFeedback
        self.noiseFB = NoiseFeedback
        self.cavityFB = CavityFeedback
        try:
            self.periodicity = bool(periodicity)
        except:
            # PeriodicityError
            raise RuntimeError("ERROR in RingAndRFTracker: Choice of" +
                               " periodicity not recognised!")
        try:
            self.interpolation = bool(interpolation)
        except:
            # InterpolationError
            raise RuntimeError("ERROR in RingAndRFTracker: Choice of" +
                               " interpolation not recognised!")
        self.profile = Profile
        self.totalInducedVoltage = TotalInducedVoltage
        if (self.interpolation is True) and (self.profile is None):
            # ProfileError
            raise RuntimeError("ERROR in RingAndRFTracker: Please specify a" +
                               " Profile object to use the interpolation option")
        if (self.cavityFB is not None) and (self.profile is None):
            # ProfileError
            raise RuntimeError("ERROR in RingAndRFTracker: Please specify a" +
                               " Profile object to use the CavityFeedback class")
        if (self.rf_params.empty is True) and (self.periodicity is True):
            # PeriodicityError
            raise RuntimeError("ERROR in RingAndRFTracker: Empty RFStation" +
                               " with periodicity not yet implemented!")
        if (self.cavityFB is not None) and (self.interpolation is False):
            self.interpolation = True
            warnings.warn('Setting interpolation to TRUE')
            # self.logger.warning("Setting interpolation to TRUE")

    def kick(self, beam_dt, beam_dE, index):
        """Function updating the particle energy due to the RF kick in a given
        RF station. The kicks are summed over the different harmonic RF systems
        in the station. The cavity phase can be shifted by the user via
        phi_offset. The main RF (harmonic[0]) has by definition phase = 0 at
        time = 0 below transition. The phases of all other RF systems are
        defined w.r.t.\ to the main RF. The increment in energy is given by the
        discrete equation of motion:

        .. math::
            \Delta E^{n+1} = \Delta E^n + \sum_{k=0}^{n_{\mathsf{rf}}-1}{e V_k^n \\sin{\\left(\omega_{\mathsf{rf,k}}^n \\Delta t^n + \phi_{\mathsf{rf,k}}^n \\right)}} - (E_s^{n+1} - E_s^n)

        """

        # voltage_kick = np.ascontiguousarray(self.charge*self.voltage[:, index])
        # omegarf_kick = np.ascontiguousarray(self.omega_rf[:, index])
        # phirf_kick = np.ascontiguousarray(self.phi_rf[:, index])
        bm.kick(beam_dt, beam_dE, self.voltage[:, index],
                self.omega_rf[:, index], self.phi_rf[:, index],
                self.charge, self.n_rf, self.acceleration_kick[index])

    def drift(self, beam_dt, beam_dE, index):
        """Function updating the particle arrival time to the RF station
        (drift). If only the zeroth order slippage factor is given, 'simple'
        and 'exact' solvers are available. The 'simple' solver is somewhat
        faster. Otherwise, the solver is automatically 'exact' and calculates
        the frequency slippage up to second order. The corresponding equations
        are (nb: the n indices correspond to the turn number):

        .. math::
            \\Delta t^{n+1} = \\Delta t^{n} + \\frac{L}{C} T_0^{n+1} \\left[ \\left(1+\\sum_{i=0}^{2}{\\alpha_i\\left(\\delta^{n+1}\\right)^{i+1}}\\right)   \\frac{1+\\left(\\Delta E/E_s\\right)^{n+1}}{1+\\delta^{n+1}}    - 1\\right] \quad \\text{(exact)}

        .. math::
            \\Delta t^{n+1} = \\Delta t^{n} + \\frac{L}{C} T_0^{n+1} \\left(\\frac{1}{1 - \\eta(\\delta^{n+1})\\delta^{n+1}} - 1\\right) \quad \\text{(legacy)}

        .. math::
            \\Delta t^{n+1} = \\Delta t^{n} + \\frac{L}{C} T_0^{n+1}\\eta_0\\delta^{n+1} \quad \\text{(simple)}

        The relative momentum needs to be calculated from the relative energy
        and is obtained as follow:

        .. math::
            \\delta = \\sqrt{1+\\beta_s^{-2}\\left[\\left(\\frac{\\Delta E}{E_s}\\right)^2 + 2\\frac{\\Delta E}{E_s}\\right]} - 1 \quad \\text{(exact)}

        .. math::
            \\delta = \\frac{\\Delta E}{\\beta_s^2 E_s} \quad \\text{(simple, legacy)}

        """
        bm.drift(beam_dt, beam_dE, self.solver, self.t_rev[index],
                 self.length_ratio, self.alpha_order, self.eta_0[index],
                 self.eta_1[index], self.eta_2[index], self.alpha_0[index],
                 self.alpha_1[index], self.alpha_2[index],
                 self.rf_params.beta[index], self.rf_params.energy[index])

    def rf_voltage_calculation(self):
        """Function calculating the total, discretised RF voltage seen by the
        beam at a given turn. Requires a Profile object.

        """
        voltages = np.ascontiguousarray(self.voltage[:, self.counter[0]])
        omega_rf = np.ascontiguousarray(self.omega_rf[:, self.counter[0]])
        phi_rf = np.ascontiguousarray(self.phi_rf[:, self.counter[0]])
        # TODO: test with multiple harmonics, think about 800 MHz OTFB
        if self.cavityFB:
            self.rf_voltage = voltages[0] * self.cavityFB.V_corr * \
                bm.sin(omega_rf[0]*self.profile.bin_centers +
                       phi_rf[0] + self.cavityFB.phi_corr) + \
                bm.rf_volt_comp(voltages[1:], omega_rf[1:], phi_rf[1:],
                                self.profile.bin_centers)
        else:
            self.rf_voltage = bm.rf_volt_comp(voltages, omega_rf, phi_rf,
                                              self.profile.bin_centers)

    def track(self):
        """Tracking method for the section. Applies first the kick, then the
        drift. Calls also RF/beam feedbacks if applicable. Updates the counter
        of the corresponding RFStation class and the energy-related variables
        of the Beam class.

        """
        turn = self.counter[0]

        # Add phase noise directly to the cavity RF phase
        if self.phi_noise is not None:
            if self.noiseFB is not None:
                self.phi_rf[:, turn] += \
                    self.noiseFB.x * self.phi_noise[:, turn]
            else:
                self.phi_rf[:, turn] += \
                    self.phi_noise[:, turn]

        # Add phase modulation directly to the cavity RF phase
        if self.phi_modulation is not None:
            self.phi_rf[:, turn] += \
                self.phi_modulation[0][:, turn]
            self.omega_rf[:, turn] += \
                self.phi_modulation[1][:, turn]

        # Determine phase loop correction on RF phase and frequency
        if self.beamFB is not None and turn >= self.beamFB.delay:
            self.beamFB.track()

        # Update the RF phase of all systems for the next turn
        # Accumulated phase offset due to beam phase loop or frequency offset
        self.rf_params.dphi_rf += 2.*np.pi*self.rf_params.harmonic[:,turn+1]* \
                                  (self.rf_params.omega_rf[:,turn+1] -
                                   self.rf_params.omega_rf_d[:,turn+1]) / \
                                  self.rf_params.omega_rf_d[:,turn+1]

        # Total phase offset
        self.rf_params.phi_rf[:,turn+1] += self.rf_params.dphi_rf

        if self.periodicity:

            # Distinguish the particles inside the frame from the particles on
            # the right-hand side of the frame.
            self.indices_right_outside = \
                np.where(self.beam.dt > self.t_rev[turn + 1])[0]
            self.indices_inside_frame = \
                np.where(self.beam.dt < self.t_rev[turn + 1])[0]

            if len(self.indices_right_outside) > 0:
                # Change reference of all the particles on the right of the
                # current frame; these particles skip one kick and drift
                self.beam.dt[self.indices_right_outside] -= \
                    self.t_rev[turn + 1]
                # Synchronize the bunch with the particles that are on the
                # RHS of the current frame applying kick and drift to the
                # bunch
                # After that all the particles are in the new updated frame
                self.insiders_dt = np.ascontiguousarray(
                    self.beam.dt[self.indices_inside_frame])
                self.insiders_dE = np.ascontiguousarray(
                    self.beam.dE[self.indices_inside_frame])
                self.kick(self.insiders_dt, self.insiders_dE, turn)
                self.drift(self.insiders_dt, self.insiders_dE, turn+1)
                self.beam.dt[self.indices_inside_frame] = self.insiders_dt
                self.beam.dE[self.indices_inside_frame] = self.insiders_dE
                # Check all the particles on the left of the just updated
                # frame and apply a second kick and drift to them with the
                # previous wave after having changed reference.
                self.indices_left_outside = np.where(self.beam.dt < 0)[0]

            else:
                self.kick(self.beam.dt, self.beam.dE, turn)
                self.drift(self.beam.dt, self.beam.dE, turn + 1)
                # Check all the particles on the left of the just updated
                # frame and apply a second kick and drift to them with the
                # previous wave after having changed reference.
                self.indices_left_outside = np.where(self.beam.dt < 0)[0]

            if len(self.indices_left_outside) > 0:
                left_outsiders_dt = np.ascontiguousarray(
                    self.beam.dt[self.indices_left_outside])
                left_outsiders_dE = np.ascontiguousarray(
                    self.beam.dE[self.indices_left_outside])
                left_outsiders_dt += self.t_rev[turn+1]
                self.kick(left_outsiders_dt, left_outsiders_dE, turn)
                self.drift(left_outsiders_dt, left_outsiders_dE, turn+1)
                self.beam.dt[self.indices_left_outside] = left_outsiders_dt
                self.beam.dE[self.indices_left_outside] = left_outsiders_dE

        else:

            if self.rf_params.empty is False:
                if self.interpolation:
                    self.rf_voltage_calculation()
                    if self.totalInducedVoltage is not None:
                        self.total_voltage = self.rf_voltage \
                            + self.totalInducedVoltage.induced_voltage
                    else:
                        self.total_voltage = self.rf_voltage

                    bm.linear_interp_kick(dt=self.beam.dt, dE=self.beam.dE,
                                          voltage=self.total_voltage,
                                          bin_centers=self.profile.bin_centers,
                                          charge=self.beam.Particle.charge,
                                          acceleration_kick=self.acceleration_kick[turn])

                    # self.drift(self.beam.dt, self.beam.dE, turn + 1)

                    # bm.LIKick_n_drift(self.beam.dt,
                    #                   self.beam.dE,
                    #                   self.total_voltage,
                    #                   self.profile.bin_centers,
                    #                   self.beam.Particle.charge,
                    #                   self.acceleration_kick[turn],
                    #                   self.solver,
                    #                   self.t_rev[turn],
                    #                   self.length_ratio,
                    #                   self.alpha_order,
                    #                   self.eta_0[turn],
                    #                   self.eta_1[turn],
                    #                   self.eta_2[turn],
                    #                   self.rf_params.beta[turn],
                    #                   self.rf_params.energy[turn])
                else:
                    self.kick(self.beam.dt, self.beam.dE, turn)

            self.drift(self.beam.dt, self.beam.dE, turn + 1)

        # Updating the beam synchronous momentum etc.
        self.beam.beta = self.rf_params.beta[turn+1]
        self.beam.gamma = self.rf_params.gamma[turn+1]
        self.beam.energy = self.rf_params.energy[turn+1]
        self.beam.momentum = self.rf_params.momentum[turn+1]

        # Increment by one the turn counter
        self.counter[0] += 1 def total_voltage(RFsection_list, harmonic = 'first'):
    """
    Total voltage from all the RF stations and systems in the ring.
    To be generalized.
    """
    
    n_sections = len(RFsection_list)
    
    #: *Sums up only the voltage of the first harmonic RF, 
    #: taking into account relative phases*
    if harmonic == 'first':
        Vcos = RFsection_list[0].voltage[0]*np.cos(RFsection_list[0].phi_rf[0])
        Vsin = RFsection_list[0].voltage[0]*np.sin(RFsection_list[0].phi_rf[0])
        if n_sections > 1:
            for i in range(1, n_sections):
                print(RFsection_list[i].voltage[0])
                Vcos += RFsection_list[i].voltage[0]*np.cos(RFsection_list[i].phi_rf[0])
                Vsin += RFsection_list[i].voltage[0]*np.sin(RFsection_list[i].phi_rf[0])
        Vtot = np.sqrt(Vcos**2 + Vsin**2)
        return Vtot
    
    #: *To be implemented*
    elif harmonic == "all":
        return 0

    else:
        warnings.filterwarnings("once")
        warnings.warn("WARNING: In total_voltage, harmonic choice not recognize!") class Beam(object):
    """Class containing the beam properties.

    This class containes the beam coordinates (dt, dE) and the beam properties.

    The beam coordinate 'dt' is defined as the particle arrival time to the RF
    station w.r.t. the reference time that is the sum of turns. The beam
    coordiate 'dE' is defined as the particle energy offset w.r.t. the
    energy of the synchronous particle.

    The class creates a beam with zero dt and dE, see distributions to match
    a beam with respect to the RF and intensity effects.

    Parameters
    ----------
    Ring : Ring
        Used to import different quantities such as the mass and the energy.
    n_macroparticles : int
        total number of macroparticles.
    intensity : float
        total intensity of the beam (in number of charge).

    Attributes
    ----------
    mass : float
        mass of the particle [eV].
    charge : int
        integer charge of the particle [e].
    beta : float
        relativistic velocity factor [].
    gamma : float
        relativistic mass factor [].
    energy : float
        energy of the synchronous particle [eV].
    momentum : float
        momentum of the synchronous particle [eV].
    dt : numpy_array, float
        beam arrival times with respect to synchronous time [s].
    dE : numpy_array, float
        beam energy offset with respect to the synchronous particle [eV].
    mean_dt : float
        average beam arrival time [s].
    mean_dE : float
        average beam energy offset [eV].
    sigma_dt : float
        standard deviation of beam arrival time [s].
    sigma_dE : float
        standard deviation of beam energy offset [eV].
    intensity : float
        total intensity of the beam in number of charges [].
    n_macroparticles : int
        total number of macroparticles in the beam [].
    ratio : float
        ratio intensity per macroparticle [].
    n_macroparticles_lost : int
        number of macro-particles marked as 'lost' [].
    id : numpy_array, int
        unique macro-particle ID number; zero if particle is 'lost'.

    See Also
    ---------
    distributions.matched_from_line_density:
        match a beam with a given bunch profile.
    distributions.matched_from_distribution_function:
        match a beam with a given distribution function in phase space.

    Examples
    --------
    >>> from input_parameters.ring import Ring
    >>> from beam.beam import Beam
    >>>
    >>> n_turns = 10
    >>> C = 100
    >>> eta = 0.03
    >>> momentum = 26e9
    >>> ring = Ring(n_turns, C, eta, momentum, 'proton')
    >>> n_macroparticle = 1e6
    >>> intensity = 1e11
    >>>
    >>> my_beam = Beam(ring, n_macroparticle, intensity)
    """

    def __init__(self, Ring, n_macroparticles, intensity):

        self.Particle = Ring.Particle
        self.beta = Ring.beta[0][0]
        self.gamma = Ring.gamma[0][0]
        self.energy = Ring.energy[0][0]
        self.momentum = Ring.momentum[0][0]
        self.dt = np.zeros([int(n_macroparticles)], dtype=bm.precision.real_t)
        self.dE = np.zeros([int(n_macroparticles)], dtype=bm.precision.real_t)
        self.mean_dt = 0.
        self.mean_dE = 0.
        self.sigma_dt = 0.
        self.sigma_dE = 0.
        self.intensity = float(intensity)
        self.n_macroparticles = int(n_macroparticles)
        self.ratio = self.intensity/self.n_macroparticles
        self.id = np.arange(1, self.n_macroparticles + 1, dtype=int)
        # For MPI
        self.n_total_macroparticles_lost = 0
        self.n_total_macroparticles = n_macroparticles
        self.is_splitted = False
        self._sumsq_dt = 0.
        self._sumsq_dE = 0.

    @property
    def n_macroparticles_lost(self):
        '''Number of lost macro-particles, defined as @property.

        Returns
        -------
        n_macroparticles_lost : int
            number of macroparticles lost.

        '''

        return len(np.where(self.id == 0)[0])

    @property
    def n_macroparticles_alive(self):
        '''Number of transmitted macro-particles, defined as @property.

        Returns
        -------
        n_macroparticles_alive : int
            number of macroparticles not lost.

        '''

        return self.n_macroparticles - self.n_macroparticles_lost

    def eliminate_lost_particles(self):
        """Eliminate lost particles from the beam coordinate arrays
        """

        indexalive = np.where(self.id != 0)[0]
        if len(indexalive) > 0:
            self.dt = np.ascontiguousarray(
                self.dt[indexalive], dtype=bm.precision.real_t)
            self.dE = np.ascontiguousarray(
                self.dE[indexalive], dtype=bm.precision.real_t)
            self.n_macroparticles = len(self.dt)
            self.id = np.arange(1, self.n_macroparticles + 1, dtype=int)
        else:
            # AllParticlesLost
            raise RuntimeError("ERROR in Beams: all particles lost and" +
                               " eliminated!")

    def statistics(self):
        '''
        Calculation of the mean and standard deviation of beam coordinates,
        as well as beam emittance using different definitions.
        Take no arguments, statistics stored in

        - mean_dt
        - mean_dE
        - sigma_dt
        - sigma_dE
        '''

        # Statistics only for particles that are not flagged as lost
        itemindex = np.where(self.id != 0)[0]
        # itemindex = bm.where(self.id, 0)
        self.mean_dt = bm.mean(self.dt[itemindex])
        self.sigma_dt = bm.std(self.dt[itemindex])
        self._sumsq_dt = np.dot(self.dt[itemindex], self.dt[itemindex])
        # self.min_dt = np.min(self.dt[itemindex])
        # self.max_dt = np.max(self.dt[itemindex])

        self.mean_dE = bm.mean(self.dE[itemindex])
        self.sigma_dE = bm.std(self.dE[itemindex])
        self._sumsq_dE = np.dot(self.dE[itemindex], self.dE[itemindex])

        # self.min_dE = np.min(self.dE[itemindex])
        # self.max_dE = np.max(self.dE[itemindex])

        # R.m.s. emittance in Gaussian approximation
        self.epsn_rms_l = np.pi*self.sigma_dE*self.sigma_dt  # in eVs

    def losses_separatrix(self, Ring, RFStation):
        '''Beam losses based on separatrix.

        Set to 0 all the particle's id not in the separatrix anymore.

        Parameters
        ----------
        Ring : Ring
            Used to call the function is_in_separatrix.
        RFStation : RFStation
            Used to call the function is_in_separatrix.
        '''

        itemindex = np.where(is_in_separatrix(Ring, RFStation, self,
                                              self.dt, self.dE) == False)[0]

        if itemindex.size != 0:
            self.id[itemindex] = 0

    def losses_longitudinal_cut(self, dt_min, dt_max):
        '''Beam losses based on longitudinal cuts.

        Set to 0 all the particle's id with dt not in the interval
        (dt_min, dt_max).

        Parameters
        ----------
        dt_min : float
            minimum dt.
        dt_max : float
            maximum dt.
        '''

        itemindex = np.where((self.dt - dt_min)*(dt_max - self.dt) < 0)[0]

        if itemindex.size != 0:
            self.id[itemindex] = 0

    def losses_energy_cut(self, dE_min, dE_max):
        '''Beam losses based on energy cuts, e.g. on collimators.

        Set to 0 all the particle's id with dE not in the interval (dE_min, dE_max).

        Parameters
        ----------
        dE_min : float
            minimum dE.
        dE_max : float
            maximum dE.
        '''

        itemindex = np.where((self.dE - dE_min)*(dE_max - self.dE) < 0)[0]

        if itemindex.size != 0:
            self.id[itemindex] = 0

    def losses_below_energy(self, dE_min):
        '''Beam losses based on lower energy cut.

        Set to 0 all the particle's id with dE below dE_min.

        Parameters
        ----------
        dE_min : float
            minimum dE.
        '''

        itemindex = np.where((self.dE - dE_min) < 0)[0]

        if itemindex.size != 0:
            self.id[itemindex] = 0

    def add_particles(self, new_particles):
        '''
        Method to add array of new particles to beam object
        New particles are given id numbers sequential from last id of this beam

        Parameters
        ----------
        new_particles : array-like
            (2, n) array of (dt, dE) for new particles
        '''

        try:
            newdt = new_particles[0]
            newdE = new_particles[1]
            if len(newdt) != len(newdE):
                raise blExcept.ParticleAdditionError(
                    "new_particles must have equal number of time and energy coordinates")
        except TypeError:
            raise blExcept.ParticleAdditionError(
                "new_particles shape must be (2, n)")

        nNew = len(newdt)

        self.id = np.concatenate((self.id, np.arange(self.n_macroparticles + 1,
                                                     self.n_macroparticles
                                                     + nNew + 1, dtype=int)))
        self.n_macroparticles += nNew

        self.dt = np.concatenate((self.dt, newdt))
        self.dE = np.concatenate((self.dE, newdE))

    def add_beam(self, other_beam):
        '''
        Method to add the particles from another beam to this beam
        New particles are given id numbers sequential from last id of this beam
        Particles with id == 0 keep id == 0 and are included in addition

        Parameters
        ----------
        other_beam : blond beam object
        '''

        if not isinstance(other_beam, type(self)):
            raise TypeError("add_beam method requires a beam object as input")

        self.dt = np.concatenate((self.dt, other_beam.dt))
        self.dE = np.concatenate((self.dE, other_beam.dE))

        counter = itl.count(self.n_macroparticles + 1)
        newids = np.zeros(other_beam.n_macroparticles)

        for i in range(other_beam.n_macroparticles):
            if other_beam.id[i]:
                newids[i] = next(counter)
            else:
                next(counter)

        self.id = np.concatenate((self.id, newids))
        self.n_macroparticles += other_beam.n_macroparticles

    def __iadd__(self, other):
        '''
        Initialisation of in place addition calls add_beam(other) if other
        is a blond beam object, calls add_particles(other) otherwise

        Parameters
        ----------
        other : blond beam object or (2, n) array
        '''

        if isinstance(other, type(self)):
            self.add_beam(other)
            return self
        else:
            self.add_particles(other)
            return self

    def split(self, random=False, fast=False):
        '''
        MPI ONLY ROUTINE: Splits the beam equally among the workers for
        MPI processing.

        Parameters
        ----------
        random : boolean
            Shuffle the beam before splitting, to be used with the
            approximation methonds.
        fast : boolean
            If true, it assumes that every worker has already a copy of the
            beam so only the particle ids are distributed.
            If false, all the coordinates are distributed by the master to all
            the workers.
        '''

        if not bm.mpiMode():
            raise RuntimeError(
                'ERROR: Cannot use this routine unless in MPI Mode')

        from ..utils.mpi_config import worker
        if worker.isMaster and random:
            import random
            random.shuffle(self.id)
            if fast == False:
                self.dt = self.dt[self.id-1]
                self.dE = self.dE[self.id-1]

        self.id = worker.scatter(self.id)
        if fast:
            self.dt = np.ascontiguousarray(self.dt[self.id-1])
            self.dE = np.ascontiguousarray(self.dE[self.id-1])
        else:
            self.dt = worker.scatter(self.dt)
            self.dE = worker.scatter(self.dE)

        assert (len(self.dt) == len(self.dE) and len(self.dt) == len(self.id))

        self.n_macroparticles = len(self.dt)
        self.is_splitted = True

    def gather(self, all=False):
        '''
        MPI ONLY ROUTINE: Gather the beam coordinates to the master or all workers.

        Parameters
        ----------
        all : boolean
            If true, every worker will get a copy of the whole beam coordinates.
            If false, only the master will gather the coordinates.
        '''
        if not bm.mpiMode():
            raise RuntimeError(
                'ERROR: Cannot use this routine unless in MPI Mode')
        from ..utils.mpi_config import worker

        if all:
            self.dt = worker.allgather(self.dt)
            self.dE = worker.allgather(self.dE)
            self.id = worker.allgather(self.id)
            self.is_splitted = False
        else:
            self.dt = worker.gather(self.dt)
            self.dE = worker.gather(self.dE)
            self.id = worker.gather(self.id)
            if worker.isMaster:
                self.is_splitted = False

        self.n_macroparticles = len(self.dt)

    def gather_statistics(self, all=False):
        '''
        MPI ONLY ROUTINE: Gather beam statistics.

        Parameters
        ----------
        all : boolean
            if true, all workers will gather the beam stats.
            If false, only the master will get the beam stats.
        '''
        if not bm.mpiMode():
            raise RuntimeError(
                'ERROR: Cannot use this routine unless in MPI Mode')

        from ..utils.mpi_config import worker
        if all:

            self.mean_dt = worker.allreduce(
                np.array([self.mean_dt]), operator='mean')[0]

            self.mean_dE = worker.allreduce(
                np.array([self.mean_dE]), operator='mean')[0]

            self.n_total_macroparticles_lost = worker.allreduce(
                np.array([self.n_macroparticles_lost]), operator='sum')[0]

            # self.n_total_macroparticles_alive = worker.allreduce(
            # np.array([self.n_macroparticles_alive]), operator='sum')[0]

            self.sigma_dt = worker.allreduce(
                np.array([self._sumsq_dt]), operator='sum')[0]
            self.sigma_dt = np.sqrt(
                self.sigma_dt/(self.n_total_macroparticles -
                               self.n_total_macroparticles_lost)
                - self.mean_dt**2)

            self.sigma_dE = worker.allreduce(
                np.array([self._sumsq_dE]), operator='sum')[0]
            self.sigma_dE = np.sqrt(
                self.sigma_dE/(self.n_total_macroparticles -
                               self.n_total_macroparticles_lost)
                - self.mean_dE**2)

            # self.sigma_dt = worker.allreduce(
            #     np.array([self.mean_dt, self.sigma_dt, self.n_macroparticles_alive]),
            #     operator='std')[0]

            # self.sigma_dE = worker.allreduce(
            #     np.array([self.mean_dE, self.sigma_dE,
            #               self.n_macroparticles_alive]),
            #     operator='std')[0]

        else:
            self.mean_dt = worker.reduce(
                np.array([self.mean_dt]), operator='mean')[0]

            self.mean_dE = worker.reduce(
                np.array([self.mean_dE]), operator='mean')[0]

            self.n_total_macroparticles_lost = worker.reduce(
                np.array([self.n_macroparticles_lost]), operator='sum')[0]

            self.sigma_dt = worker.reduce(
                np.array([self._sumsq_dt]), operator='sum')[0]
            self.sigma_dt = np.sqrt(
                self.sigma_dt/(self.n_total_macroparticles -
                               self.n_total_macroparticles_lost)
                - self.mean_dt**2)

            self.sigma_dE = worker.reduce(
                np.array([self._sumsq_dE]), operator='sum')[0]
            self.sigma_dE = np.sqrt(
                self.sigma_dE/(self.n_total_macroparticles -
                               self.n_total_macroparticles_lost)
                - self.mean_dE**2)

            # self.sigma_dt = worker.reduce(
            #     np.array([self.mean_dt, self.sigma_dt,
            #               self.n_macroparticles_alive]),
            #     operator='std')[0]

            # self.sigma_dE = worker.reduce(
            #     np.array([self.mean_dE, self.sigma_dE,
            #               self.n_macroparticles_alive]),
            #     operator='std')[0]

    def gather_losses(self, all=False):
        '''
        MPI ONLY ROUTINE: Gather beam losses.

        Parameters
        ----------
        all : boolean
            if true, all workers will gather the beam stats.
            If false, only the master will get the beam stats.
        '''
        if not bm.mpiMode():
            raise RuntimeError(
                'ERROR: Cannot use this routine unless in MPI Mode')

        from ..utils.mpi_config import worker

        if all:
            temp = worker.allgather(np.array([self.n_macroparticles_lost]))
            self.n_total_macroparticles_lost = np.sum(temp)
        else:
            temp = worker.gather(np.array([self.n_macroparticles_lost]))
            self.n_total_macroparticles_lost = np.sum(temp) class Proton(Particle):
    """ Implements a proton `Particle`.
    """

    def __init__(self):

        Particle.__init__(self, m_p*c**2/e, 1) def bigaussian(Ring, RFStation, Beam, sigma_dt, sigma_dE = None, seed = None,
               reinsertion = False):
    r"""Function generating a Gaussian beam both in time and energy 
    coordinates. Fills Beam.dt and Beam.dE arrays.
    
    Parameters
    ---------- 
    Ring : class
        A Ring type class
    RFStation : class
        An RFStation type class
    Beam : class
        A Beam type class
    sigma_dt : float
        R.m.s. extension of the Gaussian in time
    sigma_dE : float (optional)
        R.m.s. extension of the Gaussian in energy; default is None and will
        match the energy coordinate according to bucket height and sigma_dt
    seed : int (optional)
        Fixed seed to have a reproducible distribution
    reinsertion : bool (optional)
        Re-insert particles that are generated outside the separatrix into the
        bucket; default in False
    
    """
    
    warnings.filterwarnings("once")
    if Ring.n_sections > 1:
        warnings.warn("WARNING in bigaussian(): the usage of several" +
                      " sections is not yet implemented. Ignoring" +
                      " all but the first!")
    if RFStation.n_rf > 1:
        warnings.warn("WARNING in bigaussian(): the usage of multiple RF" +
                      " systems is not yet implemented. Ignoring" +
                      " higher harmonics!")
    
    counter = RFStation.counter[0]
    
    harmonic = RFStation.harmonic[0,counter]
    energy = RFStation.energy[counter]
    beta = RFStation.beta[counter]
    omega_rf = RFStation.omega_rf[0,counter] 
    phi_s = RFStation.phi_s[counter]
    phi_rf = RFStation.phi_rf[0,counter]
    eta0 = RFStation.eta_0[counter]
    
    # RF wave is shifted by Pi below transition
    if eta0<0:
        phi_rf -= np.pi
    
    # Calculate sigma_dE from sigma_dt using single-harmonic Hamiltonian
    if sigma_dE == None:
        voltage = RFStation.charge* \
                  RFStation.voltage[0,counter]
        eta0 = RFStation.eta_0[counter]
        
        phi_b = omega_rf*sigma_dt + phi_s
        sigma_dE = np.sqrt( voltage * energy * beta**2  
            * (np.cos(phi_b) - np.cos(phi_s) + (phi_b - phi_s) * np.sin(phi_s)) 
            / (np.pi * harmonic * np.fabs(eta0)) )
                
    Beam.sigma_dt = sigma_dt
    Beam.sigma_dE = sigma_dE
    
    # Generate coordinates
    np.random.seed(seed)
    
    Beam.dt = sigma_dt*np.random.randn(Beam.n_macroparticles).astype(dtype=bm.precision.real_t, order='C', copy=False) + \
        (phi_s - phi_rf)/omega_rf
    Beam.dE = sigma_dE * \
        np.random.randn(Beam.n_macroparticles).astype(
            dtype=bm.precision.real_t, order='C')
    
    # Re-insert if necessary
    if reinsertion == True:
        
        itemindex = np.where(is_in_separatrix(Ring, 
            RFStation, Beam, Beam.dt, Beam.dE) == False)[0]
         
        while itemindex.size != 0:

            Beam.dt[itemindex] = sigma_dt*np.random.randn(itemindex.size).astype(dtype=bm.precision.real_t, order='C', copy=False) \
                + (phi_s - phi_rf)/omega_rf

            Beam.dE[itemindex] = sigma_dE * \
                np.random.randn(itemindex.size).astype(
                    dtype=bm.precision.real_t, order='C')
            itemindex = np.where(is_in_separatrix(Ring,
                                                  RFStation, Beam, Beam.dt, Beam.dE) == False)[0] class CutOptions(object):
    r"""
    This class groups all the parameters necessary to slice the phase space
    distribution according to the time axis, apart from the array collecting
    the profile which is defined in the constructor of the class Profile below.

    Parameters
    ----------
    cut_left : float
        Left edge of the slicing (optional). A default value will be set if
        no value is given.
    cut_right : float
        Right edge of the slicing (optional). A default value will be set
        if no value is given.
    n_slices : int
        Optional input parameters, corresponding to the number of
        :math:`\sigma_{RMS}` of the Beam to slice (this will overwrite
        any input of cut_left and cut_right).
    n_sigma : float
        defines the left and right extremes of the profile in case those are
        not given explicitly
    cuts_unit : str
        the unit of cut_left and cut_right, it can be seconds 's' or radians
        'rad'
    RFSectionParameters : object
        RFSectionParameters[0][0] is necessary for the conversion from radians
        to seconds if cuts_unit = 'rad'. RFSectionParameters[0][0] is the value
        of omega_rf of the main harmonic at turn number 0

    Attributes
    ----------
    cut_left : float
    cut_right : float
    n_slices : int
    n_sigma : float
    cuts_unit : str
    RFSectionParameters : object
    edges : float array
        contains the edges of the slices
    bin_centers : float array
        contains the centres of the slices

    Examples
    --------
    >>> from input_parameters.ring import Ring
    >>> from input_parameters.rf_parameters import RFStation
    >>> self.ring = Ring(n_turns = 1, ring_length = 100,
    >>> alpha = 0.00001, momentum = 1e9)
    >>> self.rf_params = RFStation(Ring=self.ring, n_rf=1, harmonic=[4620],
    >>>                  voltage=[7e6], phi_rf_d=[0.])
    >>> CutOptions = profileModule.CutOptions(cut_left=0, cut_right=2*np.pi,
    >>> n_slices = 100, cuts_unit='rad', RFSectionParameters=self.rf_params)

    """

    def __init__(self, cut_left=None, cut_right=None, n_slices=100,
                 n_sigma=None, cuts_unit='s', RFSectionParameters=None):
        """
        Constructor
        """

        if cut_left is not None:
            self.cut_left = float(cut_left)
        else:
            self.cut_left = cut_left

        if cut_right is not None:
            self.cut_right = float(cut_right)
        else:
            self.cut_right = cut_right

        self.n_slices = int(n_slices)

        if n_sigma is not None:
            self.n_sigma = float(n_sigma)
        else:
            self.n_sigma = n_sigma

        self.cuts_unit = str(cuts_unit)

        self.RFParams = RFSectionParameters

        if self.cuts_unit == 'rad' and self.RFParams is None:
            # CutError
            raise RuntimeError('You should pass an RFParams object to ' +
                               'convert from radians to seconds')
        if self.cuts_unit != 'rad' and self.cuts_unit != 's':
            # CutError
            raise RuntimeError('cuts_unit should be "s" or "rad"')

        self.edges = np.zeros(n_slices + 1, dtype=bm.precision.real_t, order='C')
        self.bin_centers = np.zeros(n_slices, dtype=bm.precision.real_t, order='C')

    def set_cuts(self, Beam=None):
        """
        Method to set self.cut_left, self.cut_right, self.edges and
        self.bin_centers attributes.
        The frame is defined by :math:`n\sigma_{RMS}` or manually by the user.
        If not, a default frame consisting of taking the whole bunch +5% of the
        maximum distance between two particles in the bunch will be taken
        in each side of the frame.
        """

        if self.cut_left is None and self.cut_right is None:

            if self.n_sigma is None:
                dt_min = Beam.dt.min()
                dt_max = Beam.dt.max()
                self.cut_left = dt_min - 0.05 * (dt_max - dt_min)
                self.cut_right = dt_max + 0.05 * (dt_max - dt_min)
            else:
                mean_coords = np.mean(Beam.dt)
                sigma_coords = np.std(Beam.dt)
                self.cut_left = mean_coords - self.n_sigma*sigma_coords/2
                self.cut_right = mean_coords + self.n_sigma*sigma_coords/2

        else:

            self.cut_left = float(self.convert_coordinates(self.cut_left,
                                                           self.cuts_unit))
            self.cut_right = float(self.convert_coordinates(self.cut_right,
                                                            self.cuts_unit))

        self.edges = np.linspace(self.cut_left, self.cut_right,
                                 self.n_slices + 1).astype(dtype=bm.precision.real_t, order='C', copy=False)
        self.bin_centers = (self.edges[:-1] + self.edges[1:])/2
        self.bin_size = (self.cut_right - self.cut_left) / self.n_slices

    def track_cuts(self, Beam):
        """
        Track the slice frame (limits and slice position) as the mean of the
        bunch moves.
        Requires Beam statistics!
        Method to be refined!
        """

        delta = Beam.mean_dt - 0.5*(self.cut_left + self.cut_right)

        self.cut_left += delta
        self.cut_right += delta
        self.edges += delta
        self.bin_centers += delta

    def convert_coordinates(self, value, input_unit_type):
        """
        Method to convert a value from 'rad' to 's'.
        """

        if input_unit_type == 's':
            return value

        elif input_unit_type == 'rad':
            return value /\
                self.RFParams.omega_rf[0, self.RFParams.counter[0]]

    def get_slices_parameters(self):
        """
        Reuturn all the computed parameters.
        """
        return self.n_slices, self.cut_left, self.cut_right, self.n_sigma, \
            self.edges, self.bin_centers, self.bin_size
 class Profile(object):
    """
    Contains the beam profile and related quantities including beam spectrum,
    profile derivative.

    Parameters
    ----------

    Beam : object
        Beam from which the profile has to be calculated
    CutOptions : object
        Options for profile cutting (see above)
    FitOptions : object
        Options to get profile position and length (see above)
    FilterOptions : object
        Options to set a filter (see above)
    OtherSlicesOptions : object
        All remaining options, like smooth histogram and direct
        slicing (see above)

    Attributes
    ----------

    Beam : object
    n_slices : int
        number of slices to be used
    cut_left : float
        left extreme of the profile
    cut_right : float
        right extreme of the profile
    n_sigma : float
        defines the left and right extremes of the profile in case those are
        not given explicitly
    edges : float array
        contains the edges of the slices
    bin_centers : float array
        contains the centres of the slices
    bin_size : float
        lenght of one bin (or slice)
    n_macroparticles : float array
        contains the histogram (or profile); its elements are real if the
        smooth histogram tracking is used
    beam_spectrum : float array
        contains the spectrum of the beam (arb. units)
    beam_spectrum_freq : float array
        contains the frequencies on which the spectrum is computed [Hz]
    operations : list
        contains all the methods to be called every turn, like slice track,
        fitting, filtering etc.
    bunchPosition : float
        profile position [s]
    bunchLength : float
        profile length [s]
    filterExtraOptions : unknown (see above)

    Examples
    --------

    >>> n_slices = 100
    >>> CutOptions = profileModule.CutOptions(cut_left=0,
    >>>       cut_right=self.ring.t_rev[0], n_slices = n_slices, cuts_unit='s')
    >>> FitOptions = profileModule.FitOptions(fit_option='gaussian',
    >>>                                        fitExtraOptions=None)
    >>> filter_option = {'pass_frequency':1e7,
    >>>    'stop_frequency':1e8, 'gain_pass':1, 'gain_stop':2,
    >>>    'transfer_function_plot':False}
    >>> FilterOptions = profileModule.FilterOptions(filterMethod='chebishev',
    >>>         filterExtraOptions=filter_option)
    >>> OtherSlicesOptions = profileModule.OtherSlicesOptions(smooth=False,
    >>>                             direct_slicing = True)
    >>> self.profile4 = profileModule.Profile(my_beam, CutOptions = CutOptions,
    >>>                     FitOptions= FitOptions,
    >>>                     FilterOptions=FilterOptions,
    >>>                     OtherSlicesOptions = OtherSlicesOptions)

    """

    def __init__(self, Beam,
                 CutOptions=CutOptions(),
                 FitOptions=FitOptions(),
                 FilterOptions=FilterOptions(),
                 OtherSlicesOptions=OtherSlicesOptions()):
        """
        Constructor
        """

        # Copy of CutOptions object to be usef for reslicing
        self.cut_options = CutOptions

        # Define bins
        CutOptions.set_cuts(Beam)

        # Import (reference) Beam
        self.Beam = Beam

        # Get all computed parameters from CutOptions
        self.set_slices_parameters()

        # Initialize profile array as zero array
        self.n_macroparticles = np.zeros(self.n_slices, dtype=bm.precision.real_t, order='C')

        # Initialize beam_spectrum and beam_spectrum_freq as empty arrays
        self.beam_spectrum = np.array([], dtype=bm.precision.real_t, order='C')
        self.beam_spectrum_freq = np.array([], dtype=bm.precision.real_t, order='C')

        if OtherSlicesOptions.smooth:
            self.operations = [self._slice_smooth]
        else:
            self.operations = [self._slice]

        if FitOptions.fit_option is not None:
            self.fit_option = FitOptions.fit_option
            self.bunchPosition = 0.0
            self.bunchLength = 0.0
            if FitOptions.fit_option == 'gaussian':
                self.operations.append(self.apply_fit)
            elif FitOptions.fit_option == 'rms':
                self.operations.append(self.rms)
            elif FitOptions.fit_option == 'fwhm':
                self.operations.append(self.fwhm)

        if FilterOptions.filterMethod == 'chebishev':
            self.filterExtraOptions = FilterOptions.filterExtraOptions
            self.operations.append(self.apply_filter)

        if OtherSlicesOptions.direct_slicing:
            self.track()

    def set_slices_parameters(self):
        self.n_slices, self.cut_left, self.cut_right, self.n_sigma, \
            self.edges, self.bin_centers, self.bin_size = \
            self.cut_options.get_slices_parameters()

    def track(self):
        """
        Track method in order to update the slicing along with the tracker.
        The kwargs are currently only needed to forward the reduce kw argument
        needed for the MPI version.
        """

        for op in self.operations:
            op()

    def _slice(self):
        """
        Constant space slicing with a constant frame.
        """
        bm.slice(self.Beam.dt, self.n_macroparticles, self.cut_left,
                 self.cut_right)

        if bm.mpiMode():
            self.reduce_histo()

    def reduce_histo(self, dtype=np.uint32):
        if not bm.mpiMode():
            raise RuntimeError(
                'ERROR: Cannot use this routine unless in MPI Mode')

        from ..utils.mpi_config import worker

        if self.Beam.is_splitted:
            # Convert to uint32t for better performance
            self.n_macroparticles = self.n_macroparticles.astype(dtype, order='C')

            worker.allreduce(self.n_macroparticles)

            # Convert back to float64
            self.n_macroparticles = self.n_macroparticles.astype(dtype=bm.precision.real_t, order='C', copy=False)

        
    def scale_histo(self):
        if not bm.mpiMode():
            raise RuntimeError(
                'ERROR: Cannot use this routine unless in MPI Mode')

        from ..utils.mpi_config import worker
        if self.Beam.is_splitted:
            bm.mul(self.n_macroparticles, worker.workers, self.n_macroparticles)

    def _slice_smooth(self, reduce=True):
        """
        At the moment 4x slower than _slice but smoother (filtered).
        """
        bm.slice_smooth(self.Beam.dt, self.n_macroparticles, self.cut_left,
                        self.cut_right)

        if bm.mpiMode():
            self.reduce_histo(dtype=np.float64)

    def apply_fit(self):
        """
        It applies Gaussian fit to the profile.
        """

        if self.bunchLength == 0:
            p0 = [max(self.n_macroparticles), np.mean(self.Beam.dt),
                  np.std(self.Beam.dt)]
        else:
            p0 = [max(self.n_macroparticles), self.bunchPosition,
                  self.bunchLength/4]

        self.fitExtraOptions = ffroutines.gaussian_fit(self.n_macroparticles,
                                                       self.bin_centers, p0)
        self.bunchPosition = self.fitExtraOptions[1]
        self.bunchLength = 4*self.fitExtraOptions[2]

    def apply_filter(self):
        """
        It applies Chebishev filter to the profile.
        """
        self.n_macroparticles = ffroutines.beam_profile_filter_chebyshev(
            self.n_macroparticles, self.bin_centers, self.filterExtraOptions)

    def rms(self):
        """
        Computation of the RMS bunch length and position from the line
        density (bunch length = 4sigma).
        """

        self.bunchPosition, self.bunchLength = ffroutines.rms(
            self.n_macroparticles, self.bin_centers)

    def rms_multibunch(self, n_bunches, bunch_spacing_buckets, bucket_size_tau,
                       bucket_tolerance=0.40):
        """
        Computation of the bunch length (4sigma) and position from RMS.
        """

        self.bunchPosition, self.bunchLength = ffroutines.rms_multibunch(
            self.n_macroparticles, self.bin_centers, n_bunches,
            bunch_spacing_buckets, bucket_size_tau, bucket_tolerance)

    def fwhm(self, shift=0):
        """
        Computation of the bunch length and position from the FWHM
        assuming Gaussian line density.
        """

        self.bunchPosition, self.bunchLength = ffroutines.fwhm(
            self.n_macroparticles, self.bin_centers, shift)

    def fwhm_multibunch(self, n_bunches, bunch_spacing_buckets,
                        bucket_size_tau, bucket_tolerance=0.40, shift=0):
        """
        Computation of the bunch length and position from the FWHM
        assuming Gaussian line density for multibunch case.
        """

        self.bunchPosition, self.bunchLength = ffroutines.fwhm_multibunch(
            self.n_macroparticles, self.bin_centers, n_bunches,
            bunch_spacing_buckets, bucket_size_tau, bucket_tolerance, shift)

    def beam_spectrum_freq_generation(self, n_sampling_fft):
        """
        Frequency array of the beam spectrum
        """

        self.beam_spectrum_freq = bm.rfftfreq(n_sampling_fft, self.bin_size)

    def beam_spectrum_generation(self, n_sampling_fft):
        """
        Beam spectrum calculation
        """

        self.beam_spectrum = bm.rfft(self.n_macroparticles, n_sampling_fft)

    def beam_profile_derivative(self, mode='gradient'):
        """
        The input is one of the three available methods for differentiating
        a function. The two outputs are the bin centres and the discrete
        derivative of the Beam profile respectively.*
        """

        x = self.bin_centers
        dist_centers = x[1] - x[0]

        if mode == 'filter1d':
            derivative = ndimage.gaussian_filter1d(
                self.n_macroparticles, sigma=1, order=1, mode='wrap') / \
                dist_centers
        elif mode == 'gradient':
            derivative = np.gradient(self.n_macroparticles, dist_centers)
        elif mode == 'diff':
            derivative = np.diff(self.n_macroparticles) / dist_centers
            diffCenters = x[0:-1] + dist_centers/2
            derivative = np.interp(x, diffCenters, derivative)
        else:
            # ProfileDerivativeError
            raise RuntimeError('Option for derivative is not recognized.')

        return x, derivative
 class FitOptions(object):
    """
    This class defines the method to be used turn after turn to obtain the
    position and length of the bunch profile.

    Parameters
    ----------

    fit_method : string
        Current options are 'gaussian',
        'fwhm' (full-width-half-maximum converted to 4 sigma gaussian bunch)
        and 'rms'. The methods 'gaussian' and 'rms' give both 4 sigma.
    fitExtraOptions : unknown
        For the moment no options can be passed into fitExtraOptions

    Attributes
    ----------

    fit_method : string
    fitExtraOptions : unknown
    """

    def __init__(self, fit_option=None, fitExtraOptions=None):
        """
        Constructor
        """

        self.fit_option = str(fit_option)
        self.fitExtraOptions = fitExtraOptions
 class BunchMonitor(object):

    ''' Class able to save bunch data into h5 file. Use 'buffer_time' to select 
        the frequency of saving to file in number of turns.
        If in the constructor a Profile object is passed, that means that one
        wants to save the gaussian-fit bunch length as well (obviously the 
        Profile object has to have the fit_option set to 'gaussian').
    '''

    def __init__(self, Ring, RFParameters, Beam, filename,
                 buffer_time=None,
                 Profile=None, PhaseLoop=None, LHCNoiseFB=None):

        self.filename = filename
        self.n_turns = Ring.n_turns
        self.i_turn = 0
        self.buffer_time = buffer_time
        if buffer_time == None:
            self.buffer_time = self.n_turns
        self.rf_params = RFParameters
        self.beam = Beam
        self.profile = Profile
        if self.profile:
            if self.profile.fit_option != None:
                self.fit_option = True
            else:
                self.fit_option = False
        else:
            self.fit_option = False
        self.PL = PhaseLoop
        self.LHCNoiseFB = LHCNoiseFB

        # Initialise data and save initial state
        self.init_data(self.filename, (self.n_turns + 1,))

        # Track at initialisation
        self.track()

    def track(self):

        self.beam.statistics()

        # Write buffer with i_turn = RFcounter - 1
        self.write_buffer()

        # Synchronise to i_turn = RFcounter
        self.i_turn += 1

        if self.i_turn > 0 and (self.i_turn % self.buffer_time) == 0:
            self.open()
            self.write_data(self.h5file['Beam'], (self.n_turns + 1,))
            self.close()
            self.init_buffer()

    def init_data(self, filename, dims):

        # Prepare data
        self.beam.statistics()

        # Open file
        self.h5file = hp.File(filename + '.h5', 'w')
        self.h5file.require_group('Beam')

        # Create datasets and write first data points
        h5group = self.h5file['Beam']

        h5group.create_dataset("n_macroparticles_alive", shape=dims,
                               dtype='f',
                               compression="gzip", compression_opts=9)
        h5group["n_macroparticles_alive"][0] = self.beam.n_macroparticles_alive

        h5group.create_dataset("mean_dt", shape=dims, dtype='f',
                               compression="gzip", compression_opts=9)
        h5group["mean_dt"][0] = self.beam.mean_dt

        h5group.create_dataset("mean_dE", shape=dims, dtype='f',
                               compression="gzip", compression_opts=9)
        h5group["mean_dE"][0] = self.beam.mean_dE

        h5group.create_dataset("sigma_dt", shape=dims, dtype='f',
                               compression="gzip", compression_opts=9)
        h5group["sigma_dt"][0] = self.beam.sigma_dt

        h5group.create_dataset("sigma_dE", shape=dims, dtype='f',
                               compression="gzip", compression_opts=9)
        h5group["sigma_dE"][0] = self.beam.sigma_dE

        h5group.create_dataset("epsn_rms_l", shape=dims, dtype='f',
                               compression="gzip", compression_opts=9)
        h5group["epsn_rms_l"][0] = self.beam.epsn_rms_l

        if self.fit_option == True:

            h5group.create_dataset("bunch_length", shape=dims,
                                   dtype='f',
                                   compression="gzip", compression_opts=9)
            h5group["bunch_length"][0] = self.profile.bunchLength

        if self.PL:

            h5group.create_dataset("PL_omegaRF", shape=dims,
                                   dtype=np.float64,
                                   compression="gzip", compression_opts=9)
            h5group["PL_omegaRF"][0] = self.rf_params.omega_rf[0, 0]

            h5group.create_dataset("PL_phiRF", shape=dims,
                                   dtype='f',
                                   compression="gzip", compression_opts=9)
            h5group["PL_phiRF"][0] = self.rf_params.phi_rf[0, 0]

            h5group.create_dataset("PL_bunch_phase", shape=dims,
                                   dtype='f',
                                   compression="gzip", compression_opts=9)
            h5group["PL_bunch_phase"][0] = self.PL.phi_beam

            h5group.create_dataset("PL_phase_corr", shape=dims,
                                   dtype='f',
                                   compression="gzip", compression_opts=9)
            h5group["PL_phase_corr"][0] = self.PL.dphi

            h5group.create_dataset("PL_omegaRF_corr", shape=dims,
                                   dtype='f',
                                   compression="gzip", compression_opts=9)
            h5group["PL_omegaRF_corr"][0] = self.PL.domega_rf

            h5group.create_dataset("SL_dphiRF", shape=dims,
                                   dtype='f',
                                   compression="gzip", compression_opts=9)
            h5group["SL_dphiRF"][0] = self.rf_params.dphi_rf[0]

            h5group.create_dataset("RL_drho", shape=dims,
                                   dtype='f',
                                   compression="gzip", compression_opts=9)
            h5group["RL_drho"][0] = self.PL.drho

        if self.LHCNoiseFB:

            h5group.create_dataset("LHC_noise_FB_factor", shape=dims,
                                   dtype='f',
                                   compression="gzip", compression_opts=9)
            h5group["LHC_noise_FB_factor"][0] = self.LHCNoiseFB.x

            h5group.create_dataset("LHC_noise_FB_bl", shape=dims,
                                   dtype='f',
                                   compression="gzip", compression_opts=9)
            h5group["LHC_noise_FB_bl"][0] = self.LHCNoiseFB.bl_meas

            if self.LHCNoiseFB.bl_meas_bbb != None:

                h5group.create_dataset("LHC_noise_FB_bl_bbb",
                                       shape=(self.n_turns + 1,
                                              len(self.LHCNoiseFB.bl_meas_bbb)),
                                       dtype='f', compression="gzip",
                                       compression_opts=9)
                h5group["LHC_noise_FB_bl_bbb"][0,
                                               :] = self.LHCNoiseFB.bl_meas_bbb[:]

        # Close file
        self.close()

        # Initialise buffer for next turn
        self.init_buffer()

    def init_buffer(self):

        self.b_np_alive = np.zeros(self.buffer_time)
        self.b_mean_dt = np.zeros(self.buffer_time)
        self.b_mean_dE = np.zeros(self.buffer_time)
        self.b_sigma_dt = np.zeros(self.buffer_time)
        self.b_sigma_dE = np.zeros(self.buffer_time)
        self.b_epsn_rms = np.zeros(self.buffer_time)

        if self.fit_option == True:

            self.b_bl = np.zeros(self.buffer_time)

        if self.PL:

            self.b_PL_omegaRF = np.zeros(self.buffer_time)
            self.b_PL_phiRF = np.zeros(self.buffer_time)
            self.b_PL_bunch_phase = np.zeros(self.buffer_time)
            self.b_PL_phase_corr = np.zeros(self.buffer_time)
            self.b_PL_omegaRF_corr = np.zeros(self.buffer_time)
            self.b_SL_dphiRF = np.zeros(self.buffer_time)
            self.b_RL_drho = np.zeros(self.buffer_time)

        if self.LHCNoiseFB:

            self.b_LHCnoiseFB_factor = np.zeros(self.buffer_time)
            self.b_LHCnoiseFB_bl = np.zeros(self.buffer_time)
            if self.LHCNoiseFB.bl_meas_bbb != None:
                self.b_LHCnoiseFB_bl_bbb = np.zeros((self.buffer_time,
                                                     len(self.LHCNoiseFB.bl_meas_bbb)))

    def write_buffer(self):

        i = self.i_turn % self.buffer_time

        self.b_np_alive[i] = self.beam.n_macroparticles_alive
        self.b_mean_dt[i] = self.beam.mean_dt
        self.b_mean_dE[i] = self.beam.mean_dE
        self.b_sigma_dt[i] = self.beam.sigma_dt
        self.b_sigma_dE[i] = self.beam.sigma_dE
        self.b_epsn_rms[i] = self.beam.epsn_rms_l

        if self.fit_option == True:

            self.b_bl[i] = self.profile.bunchLength

        if self.PL:

            self.b_PL_omegaRF[i] = self.rf_params.omega_rf[0, self.i_turn]
            self.b_PL_phiRF[i] = self.rf_params.phi_rf[0, self.i_turn]
            self.b_PL_bunch_phase[i] = self.PL.phi_beam
            self.b_PL_phase_corr[i] = self.PL.dphi
            self.b_PL_omegaRF_corr[i] = self.PL.domega_rf
            self.b_SL_dphiRF[i] = self.rf_params.dphi_rf[0]
            self.b_RL_drho[i] = self.PL.drho

        if self.LHCNoiseFB:

            self.b_LHCnoiseFB_factor[i] = self.LHCNoiseFB.x
            self.b_LHCnoiseFB_bl[i] = self.LHCNoiseFB.bl_meas
            if self.LHCNoiseFB.bl_meas_bbb != None:
                self.b_LHCnoiseFB_bl_bbb[i, :] = self.LHCNoiseFB.bl_meas_bbb[:]

    def write_data(self, h5group, dims):

        i1 = self.i_turn - self.buffer_time
        i2 = self.i_turn

        h5group.require_dataset("n_macroparticles_alive", shape=dims,
                                dtype='f')
        h5group["n_macroparticles_alive"][i1:i2] = self.b_np_alive[:]

        h5group.require_dataset("mean_dt", shape=dims, dtype='f')
        h5group["mean_dt"][i1:i2] = self.b_mean_dt[:]

        h5group.require_dataset("mean_dE", shape=dims, dtype='f')
        h5group["mean_dE"][i1:i2] = self.b_mean_dE[:]

        h5group.require_dataset("sigma_dt", shape=dims, dtype='f')
        h5group["sigma_dt"][i1:i2] = self.b_sigma_dt[:]

        h5group.require_dataset("sigma_dE", shape=dims, dtype='f')
        h5group["sigma_dE"][i1:i2] = self.b_sigma_dE[:]

        h5group.require_dataset("epsn_rms_l", shape=dims, dtype='f')
        h5group["epsn_rms_l"][i1:i2] = self.b_epsn_rms[:]

        if self.fit_option == True:

            h5group.require_dataset("bunch_length", shape=dims,
                                    dtype='f')
            h5group["bunch_length"][i1:i2] = self.b_bl[:]

        if self.PL:

            h5group.require_dataset("PL_omegaRF", shape=dims,
                                    dtype=np.float64)
            h5group["PL_omegaRF"][i1:i2] = self.b_PL_omegaRF[:]

            h5group.require_dataset("PL_phiRF", shape=dims,
                                    dtype='f')
            h5group["PL_phiRF"][i1:i2] = self.b_PL_phiRF[:]

            h5group.require_dataset("PL_bunch_phase", shape=dims,
                                    dtype='f')
            h5group["PL_bunch_phase"][i1:i2] = self.b_PL_bunch_phase[:]

            h5group.require_dataset("PL_phase_corr", shape=dims,
                                    dtype='f')
            h5group["PL_phase_corr"][i1:i2] = self.b_PL_phase_corr[:]

            h5group.require_dataset("PL_omegaRF_corr", shape=dims,
                                    dtype='f')
            h5group["PL_omegaRF_corr"][i1:i2] = self.b_PL_omegaRF_corr[:]

            h5group.require_dataset("SL_dphiRF", shape=dims,
                                    dtype='f')
            h5group["SL_dphiRF"][i1:i2] = self.b_SL_dphiRF[:]

            h5group.require_dataset("RL_drho", shape=dims,
                                    dtype='f')
            h5group["RL_drho"][i1:i2] = self.b_RL_drho[:]

        if self.LHCNoiseFB:

            h5group.require_dataset("LHC_noise_FB_factor", shape=dims,
                                    dtype='f')
            h5group["LHC_noise_FB_factor"][i1:i2] = self.b_LHCnoiseFB_factor[:]

            h5group.require_dataset("LHC_noise_FB_bl", shape=dims,
                                    dtype='f')
            h5group["LHC_noise_FB_bl"][i1:i2] = self.b_LHCnoiseFB_bl[:]

            if self.LHCNoiseFB.bl_meas_bbb != None:
                h5group.require_dataset("LHC_noise_FB_bl_bbb", shape=(self.n_turns + 1,
                                                                      len(self.LHCNoiseFB.bl_meas_bbb)),
                                        dtype='f')
                h5group["LHC_noise_FB_bl_bbb"][i1:i2,
                                               :] = self.b_LHCnoiseFB_bl_bbb[:, :]

    def open(self):
        self.h5file = hp.File(self.filename + '.h5', 'r+')
        self.h5file.require_group('Beam')

    def close(self):
        self.h5file.close()
import numpy as np
import os
from blond.input_parameters.ring import Ring
from blond.input_parameters.rf_parameters import RFStation
from blond.trackers.tracker import RingAndRFTracker
from blond.trackers.utilities import total_voltage
from blond.beam.beam import Beam, Proton
from blond.beam.distributions import bigaussian
from blond.beam.profile import CutOptions, Profile, FitOptions
from blond.monitors.monitors import BunchMonitor
from blond.plots.plot import Plot
# coding: utf8
# Copyright 2014-2017 CERN. This software is distributed under the
# terms of the GNU General Public Licence version 3 (GPL Version 3), 
# copied verbatim in the file LICENCE.md.
# In applying this licence, CERN does not waive the privileges and immunities 
# granted to it by virtue of its status as an Intergovernmental Organization or
# submit itself to any jurisdiction.
# Project website: http://blond.web.cern.ch/

'''
Example input for simulating a ring with multiple RF stations
No intensity effects

:Authors: **Helga Timko**
'''

from __future__ import division, print_function
this_directory = os.path.dirname(os.path.realpath(__file__)) + '/'


try:
    os.mkdir(this_directory + '../output_files')
except:
    pass
try:
    os.mkdir(this_directory + '../output_files/EX_04_fig')
except:
    pass

# Simulation parameters -------------------------------------------------------
# Bunch parameters
N_b = 1.e9           # Intensity
N_p = 10001          # Macro-particles
tau_0 = 0.4e-9          # Initial bunch length, 4 sigma [s]

# Machine and RF parameters
C = 26658.883        # Machine circumference [m]
p_s = 450.e9         # Synchronous momentum [eV]
h = 35640            # Harmonic number
V1 = 2e6           # RF voltage, station 1 [eV]
V2 = 4e6           # RF voltage, station 1 [eV]
dphi = 0             # Phase modulation/offset
gamma_t = 55.759505  # Transition gamma
alpha = 1./gamma_t/gamma_t        # First order mom. comp. factor

# Tracking details
N_t = 2000           # Number of turns to track
dt_plt = 200         # Time steps between plots



# Simulation setup ------------------------------------------------------------
print("Setting up the simulation...")
print("")


# Define general parameters containing data for both RF stations
general_params = Ring([0.3*C, 0.7*C], [[alpha], [alpha]], 
                                   [p_s*np.ones(N_t+1), p_s*np.ones(N_t+1)], 
                                   Proton(), N_t, n_sections = 2)


# Define RF station parameters and corresponding tracker
beam = Beam(general_params, N_p, N_b)
rf_params_1 = RFStation(general_params, [h], [V1], [dphi],
                                  section_index=1)
long_tracker_1 = RingAndRFTracker(rf_params_1, beam)


rf_params_2 = RFStation(general_params, [h], [V2], [dphi],