Initial commit

.gitignore

+*__pycache__*

CollectiveEffect/__init__.py

+# -*- coding: utf-8 -*-
+"""
+Created on Tue Jan 14 12:25:30 2020
+
+@author: gamelina
+"""
+

CollectiveEffect/instabilities.py

+# -*- coding: utf-8 -*-
+"""
+General calculations about instabilities
+
+@author: Alexis Gamelin
+@date: 15/01/2020
+"""
+
+from scipy.constants import c, m_e, e, pi, epsilon_0
+
+def MBI_threshold(ring, sigma, R, b):
+    """
+    Compute the microbunching instability (MBI) threshold for a bunched beam
+    considering the steady-state parallel plate model [1][2].
+    
+    Parameters
+    ----------
+    ring : Synchrotron object
+    sigma : float
+        RMS bunch length in [s]
+    R : float
+        dipole bending radius in [m]
+    b : float
+        vertical distance between the conducting parallel plates in [m]
+        
+    Returns
+    -------
+    I : float
+        MBI current threshold
+        
+    [1] : Y. Cai, "Theory of microwave instability and coherent synchrotron 
+    radiation in electron storage rings", SLAC-PUB-14561
+    [2] : D. Zhou, "Coherent synchrotron radiation and microwave instability in 
+    electron storage rings", PhD thesis, p112
+    """
+
+    Ia = 4*pi*epsilon_0*m_e*c**3/e # Alfven current
+    chi = sigma*(R/b**3)**(1/2) # Shielding paramter
+    xi = 0.5 + 0.34*chi
+    N = (ring.L0 * Ia * ring.ac * ring.gamma * ring.sigma_delta**2 * xi *
+        sigma**(1/3) / ( c * e * R**(1/3) ))
+    I = N*e/ring.T0
+    
+    return I
\ No newline at end of file

CollectiveEffect/resistive_wall.py

+# -*- coding: utf-8 -*-
+"""
+Define resistive wall elements based on the Wakefield class.
+
+@author: Alexis Gamelin
+@date: 14/01/2020
+"""
+
+import numpy as np
+from scipy.constants import mu_0, epsilon_0, c
+from CollectiveEffect.wakefield import Wakefield, Impedance
+
+def skin_depth(omega, rho, mu_r = 1, epsilon_r = 1):
+    """ General formula for the skin depth
+    
+    Parameters
+    ----------
+    omega: angular frequency in [Hz]
+    rho: resistivity in [ohm.m]
+    mu_r : relative magnetic permeability, optional
+    epsilon_r : relative electric permittivity, optional
+    
+    Returns
+    -------
+    delta : skin depth in [m]
+    
+    """
+    
+    delta = np.sqrt(2*rho/(np.abs(omega)*mu_r*mu_0))*np.sqrt(np.sqrt(1 + 
+                   (rho*np.abs(omega)*epsilon_r*epsilon_0)**2 ) 
+                    + rho*np.abs(omega)*epsilon_r*epsilon_0)
+    return delta
+    
+
+class CircularResistiveWall(Wakefield):
+    """
+    Resistive wall Wakefield element for a circular beam pipe, approximated
+    formulas from Eq. (2.77) of Chao book [1].
+    
+    Parameters
+    ----------
+    frequency: frequency points where the impedance will be evaluated in [Hz]
+    length: beam pipe length in [m]
+    rho: resistivity in [ohm.m]
+    radius: beam pipe radius in [m]
+    mu_r : relative magnetic permeability, optional
+    epsilon_r : relative electric permittivity, optional
+    """
+
+    def __init__(self, frequency, length, rho, radius, mu_r = 1, epsilon_r = 1):
+        super().__init__()
+        
+        omega = 2*np.pi*frequency
+        Z1 = length*(1 + np.sign(omega)*1j)*rho/(
+                2*np.pi*radius*skin_depth(omega,rho))
+        Z2 = c/omega*length*(1 + np.sign(omega)*1j)*rho/(
+                np.pi*radius**3*skin_depth(omega,rho))
+        
+        Zlong = Impedance(variable = frequency, function = Z1, impedance_type='long')
+        Zxdip = Impedance(variable = frequency, function = Z2, impedance_type='xdip')
+        Zydip = Impedance(variable = frequency, function = Z2, impedance_type='ydip')
+        
+        super().append_to_model(Zlong)
+        super().append_to_model(Zxdip)
+        super().append_to_model(Zydip)
\ No newline at end of file

CollectiveEffect/tools.py

+# -*- coding: utf-8 -*-
+"""
+Tools and utilities functions for the Wakefield and impedances classes.
+
+@author: Alexis Gamelin
+@date: 14/01/2020
+"""
+
+import pandas as pd
+import numpy as np
+from scipy.integrate import trapz
+from CollectiveEffect.wakefield import Wakefield, Impedance, WakeFunction
+#from ..Tracking.synchrotron import Synchrotron
+
+def read_CST(file, impedance_type='long'):
+    """Read CST file format into an Impedance object."""
+    df = pd.read_csv(file, comment="#", header = None, sep = "\t", 
+                    names = ["Frequency","Real","Imaginary"])
+    df["Frequency"] = df["Frequency"]*1e9 
+    df.set_index("Frequency", inplace = True)
+    result = Impedance(variable = df.index,
+                       function = df["Real"] + 1j*df["Imaginary"],
+                       impedance_type=impedance_type)
+    return result
+
+def read_IW2D(file, impedance_type='long'):
+    """Read IW2D file format into an Impedance object."""
+    df = pd.read_csv(file, sep = ' ', header = None, 
+                     names = ["Frequency","Real","Imaginary"], skiprows=1)
+    df.set_index("Frequency", inplace = True)
+    result = Impedance(variable = df.index,
+                       function = df["Real"] + 1j*df["Imaginary"],
+                       impedance_type=impedance_type)
+    return result
+
+def loss_factor(wake, sigma):
+    """
+    Compute the loss factor or the kick factor of a Wakefield, Impedance or
+    WakeFunction for a Gaussian bunch. Formulas from Eq. (2) p239 and 
+    Eq.(7) p240 of [1].
+    
+    Parameters
+    ----------
+    wake: Wakefield, Impedance or WakeFunction object.
+    sigma: RMS bunch length in [s]
+    
+    Returns
+    -------
+    kloss: loss factor in [V/C] or kick factor in [V/C/m] depanding on the 
+    impedance type.
+        
+    [1] Handbook of accelerator physics and engineering, 3rd printing.
+    """
+    
+    if isinstance(wake, Impedance):
+        
+        positive_index = wake.data.index > 0
+        omega = 2*np.pi*wake.data.index[positive_index]
+        gaussian_spectrum = np.exp(-1*omega**2*sigma**2)
+        
+        if(wake.impedance_type == "long"):
+            kloss = trapz(x = omega, 
+                          y = 1/np.pi*wake.data["real"][positive_index]*gaussian_spectrum)
+        if(wake.impedance_type == "xdip" or wake.impedance_type == "ydip"):
+            kloss = trapz(x = omega, 
+                          y = 1/np.pi*wake.data["imag"][positive_index]*gaussian_spectrum)
+    if isinstance(wake, WakeFunction):
+        # To be implemented.
+        pass
+    
+    if isinstance(wake, Wakefield):
+        # To be implemented. Create field in wakefield with values ?
+        pass    
+    
+    return kloss
+
+def spectral_density(frequency, sigma, m = 1, mode="Hermite"):
+    """
+    Compute the spectral density of different modes for various values of the
+    head-tail mode number, based on Table 1 p238 of [1].
+    
+    Parameters
+    ----------
+    frequency : list or numpy array
+        sample points of the spectral density in [Hz]
+    sigma : float
+        RMS bunch length in [s]
+    m : int, optional
+        head-tail (or azimutal/synchrotron) mode number
+    mode: str, optional
+        type of the mode taken into account for the computation:
+        -"Hermite" modes for Gaussian bunches
+
+    Returns
+    -------
+    numpy array
+        
+    [1] Handbook of accelerator physics and engineering, 3rd printing.
+    """
+    
+    if mode == "Hermite":
+        return (2*np.pi*frequency*sigma)**(2*m)*np.exp(
+                -1*(2*np.pi*frequency*sigma)**2)
+    else:
+        raise NotImplementedError("Not implemanted yet.")
+    
+def effective_impedance(ring, imp, m, mu, sigma, xi=None, mode="Hermite"):
+    """
+    Compute the effective (longitudinal or transverse) impedance. 
+    Formulas from Eq. (1) and (2) p238 of [1].
+    
+    Parameters
+    ----------
+    ring : Synchrotron object
+    imp : Impedance object
+    mu : int
+        coupled bunch mode number, goes from 0 to (M-1) where M is the
+        number of bunches
+    m : int
+        head-tail (or azimutal/synchrotron) mode number
+    sigma : float
+        RMS bunch length in [s]
+    xi : float, optional
+        (non-normalized) chromaticity
+    mode: str, optional
+        type of the mode taken into account for the computation:
+        -"Hermite" modes for Gaussian bunches
+
+    Returns
+    -------
+    Zeff : float 
+        effective impedance in [ohm] or in [ohm/m] depanding on the impedance
+        type.
+        
+    [1] Handbook of accelerator physics and engineering, 3rd printing.
+    """
+    
+    if not isinstance(imp, Impedance):
+        raise TypeError("{} should be an Impedance object.".format(imp))
+        
+    fmin = imp.data.index.min()
+    fmax = imp.data.index.max()
+    if fmin > 0:
+        raise ValueError("A double sided impedance spectrum, with positive and"
+                         " negative frequencies, is needed to compute the "
+                         "effective impedance.")
+        
+    if mode == "Hermite":
+        def h(f):
+            return spectral_density(frequency=f, sigma=sigma, m=m,
+                                    mode="Hermite")
+    else:
+        raise NotImplementedError("Not implemanted yet.")
+    
+    M = 416
+    
+    pmax = fmax/(ring.f0 * M) - 50
+    pmin = fmin/(ring.f0 * M) + 50
+    
+    p = np.arange(pmin,pmax+1)
+    
+    if imp.impedance_type == "long":
+        fp = ring.f0*(p*M + mu + m*ring.tune[2])
+        fp = fp[np.nonzero(fp)] # Avoid division by 0
+        num = np.sum( imp(fp) * h(fp) / (fp*2*np.pi) )
+        den = np.sum( h(fp) )
+        Zeff = num/den
+        
+    elif imp.impedance_type == "xdip" or imp.impedance_type == "ydip":
+        if imp.impedance_type == "xdip":
+            tuneXY = ring.tune[0]
+            xi = ring.chro[0]
+        elif imp.impedance_type == "ydip":
+            tuneXY = ring.tune[1]
+            xi = ring.chro[1]
+        fp = ring.f0*(p*M + mu + tuneXY + m*ring.tuneS)
+        f_xi = xi/ring.eta*ring.f0
+        num = np.sum( imp(fp) * h(fp - f_xi) )
+        den = np.sum( h(fp - f_xi) )
+        Zeff = num/den
+    else:
+        raise TypeError("Effective impedance is only defined for long, xdip"
+                        " and ydip impedance type.")
+        
+    return Zeff

CollectiveEffect/wakefield.py

+# -*- coding: utf-8 -*-
+"""
+This module defines general classes to describes wakefields, impedances and 
+wake functions. Based on impedance library by David Amorim.
+
+@author: David Amorin, Alexis Gamelin
+@date: 14/01/2020
+"""
+
+import warnings
+import re
+import pandas as pd
+import numpy as np
+from scipy.interpolate import interp1d
+
+class ComplexData:
+    """
+    Define a general data structure for a complex function based on a pandas 
+    DataFrame.
+
+    Parameters
+    ----------
+    variable : list, numpy array
+       contains the function variable values
+
+    function : list or numpy array of complex numbers
+        contains the values taken by the complex function
+    """
+
+    def __init__(self, variable=np.array([-1e15, 1e15]),
+                 function=np.array([0, 0])):
+        self.data = pd.DataFrame({'real': np.real(function),
+                                  'imag': np.imag(function)},
+                                 index=variable)
+        self.data.index.name = 'variable'
+
+    def add(self, structure_to_add, method='zero', interp_kind='cubic'):
+        """
+        Method to add two structures. If the data don't have the same length,
+        different cases are handled.
+        
+        Parameters
+        ----------
+        structure_to_add : ComplexData object, int, float or complex.
+            structure to be added.
+        method : str, optional
+            manage how the addition is done, possibilties are: 
+            -common: the common range of the index (variable) from the two
+                structures is used. The input data are cross-interpolated
+                so that the result data contains the points at all initial
+                points in the common domain.
+            -extrapolate: extrapolate the value of both ComplexData.
+            -zero: outside of the common range, the missing values are zero. In 
+                the common range, behave as "common" method.
+        interp_kind : str, optional
+            interpolation method which is passed to pandas and to
+            scipy.interpolate.interp1d.
+            
+        Returns
+        -------
+        ComplexData 
+            Contains the sum of the two inputs.
+        """
+
+        # Create first two new DataFrames merging the variable
+        # from the two input data
+
+        if isinstance(structure_to_add, (int, float, complex)):
+            structure_to_add = ComplexData(
+                                variable=self.data.index,
+                                function=structure_to_add * np.ones(len(self.data.index)))
+
+        initial_data = pd.concat(
+                        [self.data,
+                         structure_to_add.data.index.to_frame().set_index('variable')],
+                        sort=True)
+        initial_data = initial_data[~initial_data.index.duplicated(keep='first')]
+
+        data_to_add = pd.concat(
+                        [structure_to_add.data,
+                         self.data.index.to_frame().set_index('variable')],
+                        sort=True)
+        data_to_add = data_to_add[~data_to_add.index.duplicated(keep='first')]
+
+        if method == 'common':
+            max_variable = min(structure_to_add.data.index.max(),
+                               self.data.index.max())
+
+            min_variable = max(structure_to_add.data.index.min(),
+                               self.data.index.min())
+
+            initial_data = initial_data.interpolate(method=interp_kind)
+            mask = ((initial_data.index <= max_variable)
+                    & (initial_data.index >= min_variable))
+            initial_data = initial_data[mask]
+
+            data_to_add = data_to_add.interpolate(method=interp_kind)
+            mask = ((data_to_add.index <= max_variable)
+                    & (data_to_add.index >= min_variable))
+            data_to_add = data_to_add[mask]
+
+            result_structure = ComplexData()
+            result_structure.data = initial_data + data_to_add
+            return result_structure
+
+        if method == 'extrapolate':
+            print('Not there yet')
+            return self
+        
+        if method == 'zero':
+            max_variable = min(structure_to_add.data.index.max(),
+                               self.data.index.max())
+
+            min_variable = max(structure_to_add.data.index.min(),
+                               self.data.index.min())
+
+            mask = ((initial_data.index <= max_variable)
+                    & (initial_data.index >= min_variable))
+            initial_data[mask] = initial_data[mask].interpolate(method=interp_kind)
+
+            mask = ((data_to_add.index <= max_variable)
+                    & (data_to_add.index >= min_variable))
+            data_to_add[mask] = data_to_add[mask].interpolate(method=interp_kind)
+            
+            initial_data.replace(to_replace=np.nan, value=0, inplace=True)
+            data_to_add.replace(to_replace=np.nan, value=0, inplace=True)
+            
+            result_structure = ComplexData()
+            result_structure.data = initial_data + data_to_add
+            return result_structure
+
+    def __radd__(self, structure_to_add):
+        return self.add(structure_to_add, method='zero')
+
+    def __add__(self, structure_to_add):
+        return self.add(structure_to_add, method='zero')
+
+    def multiply(self, factor):
+        """
+        Multiply a data strucure with a float or an int.
+        If the multiplication is done with something else, throw a warning.
+        """
+        if isinstance(factor, (int, float)):
+            result_structure = ComplexData()
+            result_structure.data = self.data * factor
+            return result_structure
+        else:
+            warnings.warn(('The multiplication factor is not a float '
+                           'or an int.'), UserWarning)
+            return self
+
+    def __mul__(self, factor):
+        return self.multiply(factor)
+
+    def __rmul__(self, factor):
+        return self.multiply(factor)
+    
+    def __call__(self, values, interp_kind="cubic"):
+        """
+        Interpolation of the data by calling the class to have a function-like
+        behaviour.
+        
+        Parameters
+        ----------
+        values : list or numpy array of complex, int or float
+            values to be interpolated.
+        interp_kind : str, optional
+            interpolation method which is passed to scipy.interpolate.interp1d.
+            
+        Returns
+        -------
+        numpy array 
+            Contains the interpolated data.
+        """
+        real_func = interp1d(x = self.data.index, 
+                             y = self.data["real"], kind=interp_kind)
+        imag_func = interp1d(x = self.data.index, 
+                             y = self.data["imag"], kind=interp_kind)
+        return real_func(values) + 1j*imag_func(values)
+    
+class WakeFunction(ComplexData):
+    """
+    Define a WakeFunction object based on a ComplexData object.
+    """
+
+    def __init__(self,
+                 variable=np.array([-1e15, 1e15]),
+                 function=np.array([0, 0]), wake_type='long'):
+        super().__init__(variable, function)
+        pass
+    
+class Impedance(ComplexData):
+    """
+    Define an Impedance object based on a ComplexData object.
+    """
+
+    def __init__(self,
+                 variable=np.array([-1e15, 1e15]),
+                 function=np.array([0, 0]), impedance_type='long'):
+        super().__init__(variable, function)
+        self._impedance_type = impedance_type
+        self.data.index.name = 'frequency [Hz]'
+        self.initialize_impedance_coefficient()
+
+
+    def initialize_impedance_coefficient(self):
+        """
+        Define the impedance coefficients and the plane of the impedance that
+        are presents in attributes of the class.
+        """
+        table = self.impedance_name_and_coefficients_table()
+        
+        try:
+            component_coefficients = table[self.impedance_type].to_dict()
+        except KeyError:
+            print('Impedance type {} does not exist'.format(self.impedance_type))
+        
+        self.a = component_coefficients["a"]
+        self.b = component_coefficients["b"]
+        self.c = component_coefficients["c"]
+        self.d = component_coefficients["d"]
+        self.plane = component_coefficients["plane"]
+                    
+    def impedance_name_and_coefficients_table(self):
+        """
+        Return a table associating the human readbale names of an impedance
+        component and its associated coefficients and plane.
+        """
+
+        component_dict = {
+            'long': {'a': 0, 'b': 0, 'c': 0, 'd': 0, 'plane': 'z'},
+            'xcst': {'a': 0, 'b': 0, 'c': 0, 'd': 0, 'plane': 'x'},
+            'ycst': {'a': 0, 'b': 0, 'c': 0, 'd': 0, 'plane': 'y'},
+            'xdip': {'a': 1, 'b': 0, 'c': 0, 'd': 0, 'plane': 'x'},
+            'ydip': {'a': 0, 'b': 1, 'c': 0, 'd': 0, 'plane': 'y'},
+            'xydip': {'a': 0, 'b': 1, 'c': 0, 'd': 0, 'plane': 'x'},
+            'yxdip': {'a': 1, 'b': 0, 'c': 0, 'd': 0, 'plane': 'y'},
+            'xquad': {'a': 0, 'b': 0, 'c': 1, 'd': 0, 'plane': 'x'},
+            'yquad': {'a': 0, 'b': 0, 'c': 0, 'd': 1, 'plane': 'y'},
+            'xyquad': {'a': 0, 'b': 0, 'c': 0, 'd': 1, 'plane': 'x'},
+            'yxquad': {'a': 0, 'b': 0, 'c': 1, 'd': 0, 'plane': 'y'},
+            }
+
+        return pd.DataFrame(component_dict)
+
+    def add(self, structure_to_add, beta_x=1, beta_y=1, method='zero'):
+        """
+        Method to add two Impedance objects. The two structures are
+        compared so that the addition is self-consistent.
+        Beta functions can be precised as well.
+        """
+        if isinstance(structure_to_add, (int, float, complex)):
+            result = super().add(structure_to_add, method=method)
+        elif isinstance(structure_to_add, Impedance):
+            if (self.impedance_type == structure_to_add.impedance_type):
+                weight = (beta_x ** self.power_x) * (beta_y ** self.power_y)
+                result = super().add(structure_to_add * weight, method=method)
+            else:
+                warnings.warn(('The two Impedance objects do not have the '
+                               'same coordinates or plane or type. '
+                               'Returning initial Impedance object.'),
+                              UserWarning)
+                result = self
+                
+        impedance_to_return = Impedance(
+                                result.data.index,
+                                result.data.real.values + 1j*result.data.imag.values,
+                                self.impedance_type)    
+        return impedance_to_return
+
+    def __radd__(self, structure_to_add):
+        return self.add(structure_to_add)
+
+    def __add__(self, structure_to_add):
+        return self.add(structure_to_add)
+
+    def multiply(self, factor):
+        """
+        Multiply a Impedance object with a float or an int.
+        If the multiplication is done with something else, throw a warning.
+        """
+        result = super().multiply(factor)
+        impedance_to_return = Impedance(
+                                result.data.index,
+                                result.data.real.values + 1j*result.data.imag.values,
+                                self.impedance_type)    
+        return impedance_to_return
+
+    def __mul__(self, factor):
+        return self.multiply(factor)
+
+    def __rmul__(self, factor):
+        return self.multiply(factor)
+        
+    @property
+    def impedance_type(self):
+        return self._impedance_type
+    
+    @impedance_type.setter
+    def impedance_type(self, value):
+        self._impedance_type = value
+        self.initialize_impedance_coefficient()
+
+    @property
+    def power_x(self):
+        power_x = self.a/2 + self.c/2.
+        if self.plane == 'x':
+            power_x += 1./2.
+        return power_x
+
+    @property
+    def power_y(self):
+        power_y = self.b/2. + self.d/2.
+        if self.plane == 'y':
+            power_y += 1./2.
+        return power_y
+    
+class Wakefield:
+    """
+    Define a general Wakefield machine element.
+    It contains Imepdance or WakeFunction objects which defines the wakefield
+    and other informations: beta functions.
+    """
+
+    def __init__(self, betax=1, betay=1, structure_list=[]):
+        self.betax = betax
+        self.betay = betay
+        self.list_to_attr(structure_list)
+
+    def append_to_model(self, structure_to_add):
+        list_of_attributes = dir(self)
+        if isinstance(structure_to_add, Impedance):
+            attribute_name = "Z" + structure_to_add.impedance_type
+            if attribute_name in list_of_attributes:
+                raise ValueError("There is already a component of the type "
+                                 "{} in this element.".format(attribute_name))
+            else:
+                self.__setattr__(attribute_name, structure_to_add)
+        elif isinstance(structure_to_add, WakeFunction):
+            attribute_name = "W" + structure_to_add.impedance_type
+            if attribute_name in list_of_attributes:
+                raise ValueError("There is already a component of the type "
+                                 "{} in this element.".format(attribute_name))
+            else:
+                self.__setattr__(attribute_name, structure_to_add)
+        else:
+            raise ValueError("{} is not an Impedance nor a WakeFunction.".format(structure_to_add))
+    
+    def list_to_attr(self, structure_list):
+        for component in structure_list:
+            self.append_to_model(component)
+    
+    @property
+    def list_impedance_components(self):
+        """
+        Return the list of impedance components for the element,
+        based on the attributes names.
+        """
+        return np.array([comp for comp in dir(self) if re.match(r'[Z]', comp)])
+    
+    def add(self, element_to_add):
+        """
+        This function adds two Wakefield objects.
+        The different impedance components are weighted and added.
+        The result is a Wakefield object with all the components but the
+        beta function equal to 1.
+        """
+        
+        if not isinstance(element_to_add, Wakefield):
+            raise TypeError("You can only add Wakefield objects to other"
+                            "Wakefields objects.")
+        
+        list_of_summed_impedances = []
+        list_of_components_added = []
+        
+        betax1 = self.betax
+        betay1 = self.betay
+        betax2 = element_to_add.betax
+        betay2 = element_to_add.betay
+        
+        # We first go through element1 impedance components and add them
+        # to the element2 components. If the component is missing in element2,
+        # we then simply weight the element1 component and add it to the model.
+        for component_name in self.list_impedance_components:
+        
+            element1_impedance = self.__getattribute__(component_name)
+        
+            element1_weight = ((betax1 ** element1_impedance.power_x)
+                               * (betay1 ** element1_impedance.power_y))
+            
+            try:
+                element2_impedance = element_to_add.__getattribute__(component_name)
+                list_of_components_added.append(component_name)
+                element12_impedance_sum = element1_weight * element1_impedance
+                element12_impedance_sum = element12_impedance_sum.add(element2_impedance, betax2, betay2)
+            except:
+                element12_impedance_sum = element1_weight * element1_impedance
+        
+            list_of_summed_impedances.append(element12_impedance_sum)
+        
+        # We now go through the components which are unique to element2 and
+        # add them to the impedance model, with the beta function weighting
+        missing_components_list = list(set(element_to_add.list_impedance_components) -
+                                       set(list_of_components_added))    
+    
+        for component_name in missing_components_list:
+            
+            element2_impedance = element_to_add.__getattribute__(component_name)
+            element2_weight = ((betax2 ** element2_impedance.power_x)
+                               * (betay2 ** element2_impedance.power_y))
+            element12_impedance_sum = element2_weight * element2_impedance
+            
+            list_of_summed_impedances.append(element12_impedance_sum) 
+        
+        # Gather everything in an Wakefield object
+        sum_wakefield = Wakefield(betax=1., betay=1.,
+                              structure_list=list_of_summed_impedances)
+            
+        return sum_wakefield
+        
+    def __radd__(self, structure_to_add):
+        return self.add(structure_to_add)
+
+    def __add__(self, structure_to_add):
+        return self.add(structure_to_add)

Machines/__init__.py

+# -*- coding: utf-8 -*-
+"""
+Created on Tue Jan 14 18:11:33 2020
+
+@author: gamelina
+"""
+

Machines/soleil.py

+# -*- coding: utf-8 -*-
+"""
+SOLEIL synchrotron parameters script.
+
+@author: Alexis Gamelin
+@date: 14/01/2020
+"""
+
+import numpy as np
+from Tracking.synchrotron import Synchrotron, Electron, Optics
+from Tracking.beam import Beam
+
+def soleil(mode = 'Uniform'):
+    """
+    
+    """
+    
+    h = 416
+    L = 3.540969742590899e+02
+    E0 = 2.75e9
+    particle = Electron()
+    ac = 4.16e-4
+    U0 = 1171e3
+    tau = np.array([0, 0, 3.27e-3])
+    tune = np.array([18.15687, 10.22824, 0.00502])
+    emit = np.array([4.5e-9, 4.5e-9*0.01])
+    sigma_0 = 8e-12
+    sigma_delta = 9e-4
+    
+    # mean values
+    beta = np.array([3, 1.3])
+    alpha = np.array([0, 0])
+    mean_val = Optics(beta, alpha)
+    
+    ring = Synchrotron(h, L, E0, particle, ac=ac, U0=U0, tau=tau,
+                       mean_optics=mean_val, emit=emit, tune=tune,
+                       sigma_delta=sigma_delta, sigma_0=sigma_0)
+    
+    if mode == 'Uniform':
+        beam = Beam(ring, )
+    elif mode == 'Hybrid':
+        pass
+    elif mode == 'Low alpha/25':
+        pass
+    elif mode == 'Low alpha/25':
+        pass
+    elif mode == '8 bunches':
+        pass
+    elif mode == 'Single':
+        pass
+    else:
+        raise ValueError("{} is not a correct operation mode.".format(mode))
+    
+    
+    return ring
\ No newline at end of file

Tracking/__init__.py

+# -*- coding: utf-8 -*-
+"""
+Created on Tue Jan 14 18:11:33 2020
+
+@author: gamelina
+"""
+

Tracking/beam.py

+# -*- coding: utf-8 -*-
+"""
+Beam and bunch elements
+
+@author: Alexis Gamelin
+@date: 17/01/2020
+"""
+
+import numpy as np
+import numpy.ma as ma
+import warnings
+
+class Bunch:
+    """
+    Define a bunch
+    
+    Parameters
+    ----------
+    ring : Synchrotron object
+    mp_number : float, optional
+        Macro-particle number
+    current : float, optional
+        Bunch current in [A]
+
+    Attributes
+    ----------
+    mp_number : int
+        Macro-particle number
+    charge : float
+        Bunch charge in [C]
+    charge_per_mp : float
+        Charge per macro-particle in [C]
+    particle_number : int
+        Number of particles in the bunch
+    current : float
+        Bunch current in [A]
+        
+    Methods
+    ------
+    stats()
+        Compute the mean and the standard deviation of alive particles for each
+        coordinates.
+    init_gaussian(cov=None, mean=None, **kwargs)
+        Initialize bunch particles with 6D gaussian phase space.
+    """
+    
+    def __init__(self, ring, mp_number=1e3, current=1e-3):
+        dpart = np.dtype([("x",np.float64),("xp",np.float64),
+                          ("y",np.float64),("yp",np.float64),
+                          ("tau",np.float64),("delta",np.float64)])
+        self.dtype = dpart
+        particles = np.zeros((int(mp_number),6), dtype=np.float64)
+        self.particles = ma.masked_array(particles)
+        
+        self.ring = ring
+        self._mp_number = int(mp_number)
+        self.current = current
+        
+    def __len__(self):
+        """Return the number of alive particles"""
+        #mask = self.particles["alive"] == True
+        return len(self.particles)
+        
+    def __getitem__(self, i):
+        """Return the alive particle number i"""
+        return self.particles[i]
+    
+    def __setitem__(self, i, value):
+        """Set value to the alive particle number i"""
+        self.particles[i] = value
+    
+    def __iter__(self):
+        """Iterate over alive particles"""
+        return self.particles.__iter__()
+        
+    @property
+    def mp_number(self):
+        """Macro-particle number"""
+        return self._mp_number
+    
+    @mp_number.setter
+    def mp_number(self, value):
+        self._mp_number = int(value)
+        self.__init__(self.ring, value, self.charge)
+        
+    @property
+    def charge_per_mp(self):
+        """Charge per macro-particle [C]"""
+        return self._charge_per_mp
+    
+    @charge_per_mp.setter
+    def charge_per_mp(self, value):
+        self._charge_per_mp = value
+        
+    @property
+    def charge(self):
+        """Bunch charge in [C]"""
+        return self.__len__()*self.charge_per_mp
+    
+    @charge.setter
+    def charge(self, value):
+        self.charge_per_mp = value / self.__len__()
+    
+    @property
+    def particle_number(self):
+        """Particle number"""
+        return int(self.charge / np.abs(self.ring.particle.charge))
+    
+    @particle_number.setter
+    def particle_number(self, value):
+        self.charge_per_mp = value * self.ring.particle.charge / self.__len__()
+        
+    @property
+    def current(self):
+        """Bunch current [A]"""
+        return self.charge / self.ring.T0
+    
+    @current.setter
+    def current(self, value):
+        self.charge_per_mp = value * self.ring.T0 / self.__len__()
+        
+    def mean(self):
+        """
+        Compute the mean position of alive particles for each 
+        coordinates.
+        
+        Returns
+        -------
+        mean : numpy array
+            mean position of alive particles
+        """
+        mean = [[self[name].mean()] for name in self.dtype.names]
+        return np.array(mean)
+    
+    def std(self):
+        """
+        Compute the standard deviation of the position of alive 
+        particles for each coordinates.
+        
+        Returns
+        -------
+        std : numpy array
+            standard deviation of the position of alive particles
+        """
+        std = [[self[name].std()] for name in self.dtype.names]
+        return np.array(std)
+        
+    def emit(self):
+        """
+        Compute the bunch emittance for each plane [1].
+        Correct definition of long emit ?
+        
+        Returns
+        -------
+        emit : numpy array
+            bunch emittance
+        
+        References
+        ----------
+        [1] Wiedemann, H. (2015). Particle accelerator physics. 4th 
+        edition. Springer, Eq.(8.39) of p224.
+        """
+        emitX = (np.mean(self['x']**2)*np.mean(self['xp']**2) - 
+                 np.mean(self['x']*self['xp'])**2)**(0.5)
+        emitY = (np.mean(self['y']**2)*np.mean(self['yp']**2) - 
+                 np.mean(self['y']*self['yp'])**2)**(0.5)
+        emitS = (np.mean(self['tau']**2)*np.mean(self['delta']**2) - 
+                 np.mean(self['tau']*self['delta'])**2)**(0.5)
+        return np.array([emitX, emitY, emitS])
+        
+    def init_gaussian(self, cov=None, mean=None, **kwargs):
+        """
+        Initialize bunch particles with 6D gaussian phase space.
+        Covariance matrix is taken from [1].
+        
+        Parameters
+        ----------
+        cov : (6,6) array, optional
+            Covariance matrix of the bunch distribution
+        mean : (6,) array, optional
+            Mean of the bunch distribution
+        
+        References
+        ----------
+        [1] Wiedemann, H. (2015). Particle accelerator physics. 4th 
+        edition. Springer, Eq.(8.38) of p223.
+
+        """
+        if mean is None:
+            mean = np.zeros((6,))
+        
+        if cov is None:
+            sigma_0 = kwargs.get("sigma_0", self.ring.sigma_0)
+            sigma_delta = kwargs.get("sigma_delta", self.ring.sigma_delta)
+            optics = kwargs.get("optics", self.ring.mean_optics)
+            
+            cov = np.zeros((6,6))
+            cov[0,0] = self.ring.emit[0]*optics.beta[0]
+            cov[1,1] = self.ring.emit[0]*optics.gamma[0]
+            cov[0,1] = -1*self.ring.emit[0]*optics.alpha[0]
+            cov[1,0] = -1*self.ring.emit[0]*optics.alpha[0]
+            cov[2,2] = self.ring.emit[1]*optics.beta[1]
+            cov[3,3] = self.ring.emit[1]*optics.gamma[1]
+            cov[2,3] = -1*self.ring.emit[1]*optics.alpha[1]
+            cov[3,2] = -1*self.ring.emit[1]*optics.alpha[1]
+            cov[4,4] = sigma_0**2
+            cov[5,5] = sigma_delta**2
+            
+        values = np.random.multivariate_normal(mean, cov, size=self.mp_number)
+        self.particles["x"] = values[:,0]
+        self.particles["xp"] = values[:,1]
+        self.particles["y"] = values[:,2]
+        self.particles["yp"] = values[:,3]
+        self.particles["tau"] = values[:,4]
+        self.particles["delta"] = values[:,5]
+        
+class Beam:
+    """
+    Define a Beam object composed of several Bunch objects. 
+    
+    Parameters
+    ----------
+    ring : Synchrotron object
+    bunch_list : list of Bunch object, optional
+
+    Attributes
+    ----------
+    current : float
+        Total beam current in [A]
+    charge : float
+        Total bunch charge in [C]
+    particle_number : int
+        Total number of particle in the beam
+    filling_pattern : list of bool
+        Filling pattern of the beam
+        
+    Methods
+    ------
+    init_beam()
+    """
+    
+    def __init__(self, ring, bunch_list=None):
+        self.ring = ring
+        
+        if bunch_list is None:
+            self.init_beam(np.zeros((self.ring.h,1),dtype=bool))
+        else:
+            if (len(bunch_list) != self.ring.h):
+                raise ValueError(("The length of the bunch list is {} ".format(len(bunch_list)) + 
+                                  "but should be {}".format(self.ring.h)))
+            self.bunch_list = bunch_list
+            
+    def __len__(self):
+        """Return the number of bunches"""
+        length = 0
+        for bunch in self:
+            if bunch is not None:
+                length += 1
+        return length        
+    
+    def __getitem__(self, i):
+        """Return the bunch number i"""
+        return self.bunch_list.__getitem__(i)
+    
+    def __setitem__(self, i, value):
+        """Set value to the bunch number i"""
+        self.bunch_list.__setitem__(i, value)
+    
+    def __iter__(self):
+        """Iterate over the bunches"""
+        return self.bunch_list.__iter__()
+        
+    def init_beam(self, filling_pattern, current_per_bunch=1e-3, 
+                  mp_per_bunch=1e3):
+        """
+        Initialize beam with a given filling pattern but with uniform current
+        per bunch and marco-particle number per bunch.
+        
+        Parameters
+        ----------
+        filling_pattern : numpy array or list of length ring.h
+            Filling pattern of the beam,
+        charge_per_bunch : float, optional
+            Charge per bunch in [C]
+        mp_per_bunch : float, optional
+            Macro-particle number per bunch
+        """
+        
+        if (len(filling_pattern) != self.ring.h):
+            raise ValueError(("The length of filling pattern is {} ".format(len(filling_pattern)) + 
+                              "but should be {}".format(self.ring.h)))
+        
+        bunch_list = []
+        for value in filling_pattern:
+            if value == True:
+                bunch_list.append(Bunch(self.ring, mp_per_bunch, current_per_bunch))
+            elif value == False:
+                bunch_list.append(None)
+            else:
+                raise ValueError("{} should be True or False".format(value))
+        self.bunch_list = bunch_list
+        # init bunches => gauss
+        
+    @property
+    def current(self):
+        """Total beam current in [A]"""
+        current = 0
+        for bunch in self:
+            if bunch is not None:
+                current += bunch.current
+        return current
+    
+    @property
+    def charge(self):
+        """Total beam charge in [C]"""
+        charge = 0
+        for bunch in self:
+            if bunch is not None:
+                charge += bunch.charge
+        return charge
+    
+    @property
+    def particle_number(self):
+        """Total number of particle in the beam"""
+        particle_number = 0
+        for bunch in self:
+            if bunch is not None:
+                particle_number += bunch.particle_number
+        return particle_number
+    
+    @property
+    def filling_pattern(self):
+        """Filling pattern of the beam"""
+        filling_pattern = []
+        for bunch in self:
+            if isinstance(bunch, Bunch):
+                filling_pattern.append(True)
+            else:
+                filling_pattern.append(False)
+        return filling_pattern
+        
\ No newline at end of file

Tracking/synchrotron.py

+# -*- coding: utf-8 -*-
+"""
+General elements
+
+@author: Alexis Gamelin
+@date: 15/01/2020
+"""
+
+import numpy as np
+from scipy.constants import c, m_e, m_p, e
+
+class Particle:
+    """ Define a particle
+
+    Attributes
+    ----------
+    mass : float
+        total particle mass in [kg]
+    charge : float
+        electrical charge in [C]
+    E_rest : float
+        particle rest energy in [eV]
+    """
+    def __init__(self, mass, charge):
+        self.mass = mass
+        self.charge = charge
+
+    @property
+    def E_rest(self):
+        return self.mass * c ** 2 / e
+    
+class Electron(Particle):
+    """ Define an electron"""
+    def __init__(self):
+        super().__init__(m_e, -1*e)
+        
+class Proton(Particle):
+    """ Define a proton"""
+    def __init__(self):
+        super().__init__(m_p, e)
+
+class Optics:
+    """Handle optic functions"""
+    def __init__(self, beta, alpha):
+        self.beta = beta
+        self.alpha = alpha
+
+    @property
+    def gamma(self):
+        return (1 + self.alpha**2)/self.beta
+        
+class Synchrotron:
+    """ Synchrotron class to store main properties. """
+    def __init__(self, h, L, E0, particle, **kwargs):
+        self.particle = particle
+        self._h = h
+        self.L = L
+        self.E0 = E0 # Nominal (total) energy of the ring [eV]
+        
+        self.ac = kwargs.get('ac') # Momentum compaction factor
+        self.U0 = kwargs.get('U0') # Energy loss per turn [eV]
+        self.tau = kwargs.get('tau') # X/Y/S damping times [s]
+        self.sigma_delta = kwargs.get('sigma_delta') # Equilibrium energy spread
+        self.sigma_0 = kwargs.get('sigma_0') # Natural bunch length [s]
+        self.tune = kwargs.get('tune') # X/Y/S tunes
+        self.emit = kwargs.get('emit') # X/Y emittances in [m.rad]
+        self.chro = kwargs.get('chro') # X/Y (non-normalized) chromaticities
+        self.mean_optics = kwargs.get('mean_optics') # Optics object with mean values
+        
+    @property
+    def h(self):
+        """Harmonic number"""
+        return self._h
+    
+    @h.setter
+    def h(self,value):
+        self._h = value
+        self.L = self.L  # call setter
+        
+    @property
+    def L(self):
+        """Ring circumference [m]"""
+        return self._L
+    
+    @L.setter
+    def L(self,value):
+        self._L = value
+        self._T0 = self.L/c
+        self._f0 = 1/self.T0
+        self._omega0 = 2*np.pi*self.f0
+        self._f1 = self.h*self.f0
+        self._omega1 = 2*np.pi*self.f1
+        self._k1 = self.omega1/c
+        
+    @property
+    def T0(self):
+        """Revolution time [s]"""
+        return self._T0
+    
+    @T0.setter
+    def T0(self, value):
+        self.L = c*value
+        
+    @property
+    def f0(self):
+        """Revolution frequency [Hz]"""
+        return self._f0
+    
+    @f0.setter
+    def f0(self,value):
+        self.L = c/value
+        
+    @property
+    def omega0(self):
+        """Angular revolution frequency [Hz rad]"""
+        return self._omega0
+    
+    @omega0.setter
+    def omega0(self,value):
+        self.L = 2*np.pi*c/value
+        
+    @property
+    def f1(self):
+        """Fundamental RF frequency [Hz]"""
+        return self._f1
+    
+    @f1.setter
+    def f1(self,value):
+        self.L = self.h*c/value
+        
+    @property
+    def omega1(self):
+        """Fundamental RF angular frequency[Hz rad]"""
+        return self._omega1
+    
+    @omega1.setter
+    def omega1(self,value):
+        self.L = 2*np.pi*self.h*c/value
+        
+    @property
+    def k1(self):
+        """Fundamental RF wave number [m**-1]"""
+        return self._k1
+    
+    @k1.setter
+    def k1(self,value):
+        self.L = 2*np.pi*self.h/value
+    
+    @property
+    def gamma(self):
+        """Relativistic gamma"""
+        return self._gamma
+
+    @gamma.setter
+    def gamma(self, value):
+        self._gamma = value
+        self._beta = np.sqrt(1 - self.gamma**-2)
+        self._E0 = self.gamma*self.particle.mass*c**2/e
+
+    @property
+    def beta(self):
+        """Relativistic beta"""
+        return self._beta
+
+    @beta.setter
+    def beta(self, value):
+        self.gamma = 1/np.sqrt(1-value**2)
+        
+    @property
+    def E0(self):
+        """Nominal (total) energy of the ring [eV]"""
+        return self._E0
+    
+    @E0.setter
+    def E0(self, value):
+        self.gamma = value/(self.particle.mass*c**2/e)
+
+    @property
+    def eta(self):
+        """Momentum compaction"""
+        return self.ac - 1/(self.gamma**2)
\ No newline at end of file

__init__.py

+# -*- coding: utf-8 -*-
+"""
+Created on Tue Jan 14 18:11:33 2020
+
+@author: gamelina
+"""
+