Add docstrings and small code style fixes

collective_effects/wakefield.py

@@ -71,14 +71,13 @@ class ComplexData:
         # 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(index_name)],
-                        sort=True)
+            structure_to_add = ComplexData(variable=self.data.index,
+                                           function=(structure_to_add * 
+                                                     np.ones(len(self.data.index))))
+                                
+        data_to_concat = structure_to_add.data.index.to_frame().set_index(index_name)
+        
+        initial_data = pd.concat([self.data, data_to_concat], sort=True)
         initial_data = initial_data[~initial_data.index.duplicated(keep='first')]
 
         data_to_add = pd.concat(
@@ -325,16 +324,41 @@ class Impedance(ComplexData):
     
 class Wakefield:
     """
-    Define a general Wakefield machine element.
-    It contains Imepdance or WakeFunction objects which defines the wakefield
-    and other informations: beta functions.
+    Defines a Wakefield which corresponds to a single physical element which 
+    produces different types of wakes, represented by their Impedance or 
+    WakeFunction objects.
+    
+    Parameters
+    ----------
+    structure_list : list, optional
+        list of Impedance/WakeFunction objects to add to the Wakefield.
+    name : str, optional
+    
+    Attributes
+    ----------
+    impedance_components : array of str
+        Impedance component names present for this element.
+        
+    Methods
+    -------
+    append_to_model(structure_to_add)
+        Add Impedance/WakeFunction to Wakefield.
+    list_to_attr(structure_list)
+        Add list of Impedance/WakeFunction to Wakefield.
     """
 
-    def __init__(self, structure_list=[], name=None):
+    def __init__(self, structure_list=None, name=None):
         self.list_to_attr(structure_list)
         self.name = name
 
     def append_to_model(self, structure_to_add):
+        """
+        Add Impedance/WakeFunction component to Wakefield.
+
+        Parameters
+        ----------
+        structure_to_add : Impedance or WakeFunction object
+        """
         list_of_attributes = dir(self)
         if isinstance(structure_to_add, Impedance):
             attribute_name = "Z" + structure_to_add.impedance_type
@@ -354,26 +378,46 @@ class Wakefield:
             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)
+        """
+         Add list of Impedance/WakeFunction components to Wakefield.
+
+        Parameters
+        ----------
+        structure_list : list of Impedance or WakeFunction objects.
+        """
+        if structure_list is not None:
+            for component in structure_list:
+                self.append_to_model(component)
     
     @property
-    def list_impedance_components(self):
+    def impedance_components(self):
         """
-        Return the list of impedance components for the element,
-        based on the attributes names.
+        Return an array of the impedance component names for the element.
         """
         return np.array([comp for comp in dir(self) if re.match(r'[Z]', comp)])
     
     
 class ImpedanceModel(Element):
+    """
+    Define the impedance model of the machine.
+    
+    Parameters
+    ----------
+    ring : Synchrotron object
+    wakefield_list : list of Wakefield objects
+        Wakefields to add to the model.
+    wakefiled_postions : list
+        Longitudinal positions corresponding to the added Wakfields.
+        
+    
+    """
     
-    def __init__(self, ring, wakefields_to_add=None, wakefileds_postions=None):
+    def __init__(self, ring, wakefield_list=None, wakefiled_postions=None):
         self.ring = ring
         self.optics = self.ring.optics
         self.wakefields = []
         self.positions = np.array([])
-        self.add(wakefields_to_add, wakefileds_postions)
+        self.add(wakefield_list, wakefiled_postions)
         
     def track(self, beam):
         """
@@ -386,7 +430,9 @@ class ImpedanceModel(Element):
         raise NotImplementedError
         
     def sum_elements(self):
-        
+        """
+        Sum all the elements in the model into sum_wakefield.
+        """
         beta = self.optics.beta(self.positions)
         list_impedance_components = ["Zlong","Zxdip","Zydip","Zxquad","Zyquad"]
         sum_wakefield = Wakefield()
@@ -405,11 +451,21 @@ class ImpedanceModel(Element):
             sum_wakefield.append_to_model(sum_imp)
         self.sum_wakefield = sum_wakefield
         
-    def add(self, wakefield_list, position_list):
+    def add(self, wakefield_list, wakefiled_postions):
+        """
+        Add elements to the model.
+
+        Parameters
+        ----------
+        wakefield_list : list of Wakefield objects
+            Wakefields to add to the model.
+        wakefiled_postions : list
+            Longitudinal positions corresponding to the added Wakfields.
+        """
         
-        if (wakefield_list is not None) and (position_list is not None):
+        if (wakefield_list is not None) and (wakefiled_postions is not None):
             for wakefield in wakefield_list:
                 self.wakefields.append(wakefield)
                 
-            for position in position_list:
+            for position in wakefiled_postions:
                 self.positions = np.append(self.positions, position)
\ No newline at end of file

tracking/particles.py

@@ -12,7 +12,8 @@ from tracking.parallel import Mpi
 from scipy.constants import c, m_e, m_p, e
 
 class Particle:
-    """ Define a particle object
+    """
+    Define a particle object
 
     Attributes
     ----------

tracking/synchrotron.py

 # -*- coding: utf-8 -*-
 """
-General elements
+Module where the synchrotron class is defined.
 
 @author: Alexis Gamelin
 @date: 15/01/2020
@@ -8,11 +8,66 @@ General elements
 
 import numpy as np
 from scipy.constants import c, e
-
         
 class Synchrotron:
-    """ Synchrotron class to store main properties. """
+    """
+    Synchrotron class to store main properties.
+    
+    Optional parameters are optional only if the Optics object passed to the 
+    class uses a loaded lattice.
     
+    Parameters
+    ----------
+    h : int
+        Harmonic number of the accelerator.
+    optics : Optics object
+        Object where the optic functions are stored.
+    particle : Particle object
+        Particle which is accelerated.
+    tau : array of shape (3,)
+        Horizontal, vertical and longitudinal damping times in [s].
+    sigma_delta : float
+        Equilibrium energy spread.
+    sigma_0 : float
+        Natural bunch length in [s].
+    emit : array of shape (2,)
+        Horizontal and vertical equilibrium emittance in [m.rad].
+    L : float, optional
+        Ring circumference in [m].
+    E0 : float, optional
+        Nominal (total) energy of the ring in [eV].
+    ac : float, optional
+        Momentum compaction factor.
+    tune : array of shape (2,), optional
+        Horizontal and vertical tunes.
+    chro : array of shape (2,), optional
+        Horizontal and vertical (non-normalized) chromaticities.
+    U0 : float, optional
+        Energy loss per turn in [eV].
+        
+    Attributes
+    ----------
+    T0 : float
+        Revolution time in [s].
+    f0 : float
+        Revolution frequency in [Hz].
+    omega0 : float
+        Angular revolution frequency in [Hz.rad]
+    f1 : float
+        Fundamental RF frequency in [Hz].
+    omega1 : float
+        Fundamental RF angular frequency in [Hz.rad].
+    k1 : float
+        Fundamental RF wave number in [m**-1].
+    gamma : float
+        Relativistic Lorentz gamma.
+    beta : float
+        Relativistic Lorentz beta.
+    eta : float
+        Momentum compaction.
+    sigma : array of shape (4,)
+        RMS beam size at equilibrium in [m].
+    """
     def __init__(self, h, optics, particle, **kwargs):
         self._h = h
         self.particle = particle
@@ -36,9 +91,7 @@ class Synchrotron:
         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
                 
     @property
     def h(self):