Wakefield : Bugfix on ComplexData addition + Addition on Wakefields + Improve ImpedanceModel

Bugfix on ComplexData addition : sort data before interpolation
Wakefield class : Add static method to add wakefields
ImpedanceModel class : add methods to handle more easily the wakefields. Maybe can be improved more by using a single Wakefield of each type ?
ImpedanceModel class : add plot_area and summary. summary to be improved more

collective_effects/wakefield.py

@@ -11,9 +11,12 @@ import warnings
 import re
 import pandas as pd
 import numpy as np
+import matplotlib.pyplot as plt
+import collective_effects.tools as tools
 from scipy.interpolate import interp1d
 from tracking.element import Element
-
+from copy import deepcopy
+from scipy.integrate import trapz
 
 class ComplexData:
     """
@@ -79,12 +82,14 @@ class ComplexData:
         
         initial_data = pd.concat([self.data, data_to_concat], sort=True)
         initial_data = initial_data[~initial_data.index.duplicated(keep='first')]
+        initial_data = initial_data.sort_index()
 
         data_to_add = pd.concat(
                         [structure_to_add.data,
                          self.data.index.to_frame().set_index(index_name)],
                         sort=True)
         data_to_add = data_to_add[~data_to_add.index.duplicated(keep='first')]
+        data_to_add = data_to_add.sort_index()
 
         if method == 'common':
             max_variable = min(structure_to_add.data.index.max(),
@@ -345,6 +350,10 @@ class Wakefield:
         Add Impedance/WakeFunction to Wakefield.
     list_to_attr(structure_list)
         Add list of Impedance/WakeFunction to Wakefield.
+    add_wakefields(wake1, beta1, wake2, beta2)
+        Add two Wakefields taking into account beta functions.
+    add_several_wakefields(wakefields, beta)
+        Add a list of Wakefields taking into account beta functions.
     """
 
     def __init__(self, structure_list=None, name=None):
@@ -396,6 +405,78 @@ class Wakefield:
         """
         return np.array([comp for comp in dir(self) if re.match(r'[Z]', comp)])
     
+    @staticmethod
+    def add_wakefields(wake1, beta1, wake2, beta2):
+        """
+        Add two Wakefields taking into account beta functions.
+
+        Parameters
+        ----------
+        wake1 : Wakefield
+            Wakefield to add.
+        beta1 : array of shape (2,)
+            Beta functions at wake1 postion.
+        wake2 : Wakefield
+            Wakefield to add.
+        beta2 : array of shape (2,)
+            Beta functions at wake2 postion.
+
+        Returns
+        -------
+        wake_sum : Wakefield
+            Wakefield with summed components.
+
+        """
+        wake_sum = deepcopy(wake1)
+        for component_name1 in wake1.impedance_components:
+            impedance1 = getattr(wake_sum, component_name1)
+            weight1 = ((beta1[0] ** impedance1.power_x) * 
+                      (beta1[1] ** impedance1.power_y))
+            setattr(wake_sum, component_name1, weight1*impedance1)
+            
+        for component_name2 in wake2.impedance_components: 
+            impedance2 = getattr(wake2, component_name2)
+            weight2 = ((beta2[0] ** impedance2.power_x) * 
+                      (beta2[1] ** impedance2.power_y))
+            try:
+                impedance1 = getattr(wake_sum, component_name2)
+                setattr(wake_sum, component_name2, impedance1 +
+                        weight2*impedance2)
+            except AttributeError:
+                setattr(wake_sum, component_name2, weight2*impedance2)
+
+        return wake_sum
+    
+    @staticmethod
+    def add_several_wakefields(wakefields, beta):
+        """
+        Add a list of Wakefields taking into account beta functions.
+        
+        Parameters
+        ----------
+        wakefields : list of Wakefield
+            Wakefields to add.
+        beta : array of shape (len(wakefields), 2)
+            Beta functions at the Wakefield postions.
+
+        Returns
+        -------
+        wake_sum : Wakefield
+            Wakefield with summed components..
+
+        """
+        if len(wakefields) == 1:
+            return wakefields[0]
+        elif len(wakefields) > 1:
+            wake_sum = Wakefield.add_wakefields(wakefields[0], beta[:,0],
+                                     wakefields[1], beta[:,1])
+            for i in range(len(wakefields) - 2):
+                wake_sum = Wakefield.add_wakefields(wake_sum, [1 ,1], 
+                                         wakefields[i+2], beta[:,i+2])
+            return wake_sum
+        else:
+            raise ValueError("Error in input.")
+        
     
 class ImpedanceModel(Element):
     """
@@ -408,8 +489,40 @@ class ImpedanceModel(Element):
         Wakefields to add to the model.
     wakefiled_postions : list
         Longitudinal positions corresponding to the added Wakfields.
-        
     
+    Attributes
+    ----------
+    wakefields : list of Wakefield objects
+        Wakefields in the model.
+    positions : array
+        Wakefields positions.
+    names : array
+        Names of (unique) Wakefield objects.
+    sum : Wakefield
+        Sum of every Wakefield in the model.
+    sum_"name" : Wakefield
+        Sum of every Wakefield with the same "name".
+    sum_names : array
+        Names of attributes where the wakefields are summed by name.    
+    
+    Methods
+    -------
+    add(wakefield_list, wakefiled_postions)
+        Add elements to the model.
+    add_multiple_elements(wakefield, wakefiled_postions)
+        Add the same element at different locations to the model.
+    sum_elements()
+        Sum all wakefields into self.sum.
+    update_name_list()
+        Update self.names with uniques names of self.wakefields.
+    find_wakefield(name)
+        Return indexes of wakefields with the same name in self.wakefields.
+    sum_by_name(name)
+        Sum the elements with the same name in the model into sum_name.
+    sum_by_name_all(name)
+        Sum all the elements with the same name in the model into sum_name.
+    plot_area(Z_type="Zlong", component="real", sigma=None, attr_list=None)
+        Plot the contributions of different kind of wakefields.
     """
     
     def __init__(self, ring, wakefield_list=None, wakefiled_postions=None):
@@ -417,6 +530,8 @@ class ImpedanceModel(Element):
         self.optics = self.ring.optics
         self.wakefields = []
         self.positions = np.array([])
+        self.names = np.array([])
+        self.sum_names = np.array([])
         self.add(wakefield_list, wakefiled_postions)
         
     def track(self, beam):
@@ -430,27 +545,75 @@ class ImpedanceModel(Element):
         raise NotImplementedError
         
     def sum_elements(self):
+        """Sum all wakefields into self.sum"""
+        beta = self.optics.beta(self.positions)
+        self.sum = Wakefield.add_several_wakefields(self.wakefields, beta)
+        if "sum" not in self.sum_names:
+            self.sum_names = np.append(self.sum_names, "sum")
+    
+    def update_name_list(self):
+        """Update self.names with uniques names of self.wakefields."""
+        for wakefield in self.wakefields:
+            if wakefield.name is None:
+                continue
+            if wakefield.name not in self.names:
+                self.names = np.append(self.names, wakefield.name)
+                
+    def count_elements(self):
+        """Count number of each type of wakefield in the model."""
+        self.count = np.zeros(len(self.names))
+        for wakefield in self.wakefields:
+            if wakefield.name is not None:
+                self.count += (wakefield.name == self.names).astype(int)
+                
+    def find_wakefield(self, name):
         """
-        Sum all the elements in the model into sum_wakefield.
+        Return indexes of wakefields with the same name in self.wakefields.
+
+        Parameters
+        ----------
+        name : str
+            Wakefield name.
+
+        Returns
+        -------
+        index : list
+            Index of positions in self.wakefields.
+
         """
-        beta = self.optics.beta(self.positions)
-        list_impedance_components = ["Zlong","Zxdip","Zydip","Zxquad","Zyquad"]
-        sum_wakefield = Wakefield()
-        for component_name in list_impedance_components:
-            sum_imp = Impedance(variable=np.array([0,1]), 
-                                function=np.array([0, 0]),
-                                impedance_type=component_name[1:])
-            for index, wakefield in enumerate(self.wakefields):
-                try:
-                    impedance = wakefield.__getattribute__(component_name)
-                    weight = ((beta[0,index] ** impedance.power_x) * 
-                              (beta[1,index] ** impedance.power_y))
-                    sum_imp += weight*impedance
-                except AttributeError:
-                    pass
-            sum_wakefield.append_to_model(sum_imp)
-        self.sum_wakefield = sum_wakefield
+        index = []
+        for i, wakefield in enumerate(self.wakefields):
+            if wakefield.name == name:
+                index.append(i)
+        return index
+
+    def sum_by_name(self, name):
+        """
+        Sum the elements with the same name in the model into sum_name.
         
+        Parameters
+        ----------
+        name : str
+            Name of the wakefield to sum.
+        """
+        attribute_name = "sum_" + name
+        index = self.find_wakefield(name)
+        beta = self.optics.beta(self.positions[index])
+        wakes = []
+        for i in index:
+            wakes.append(self.wakefields[i])
+        wake_sum = Wakefield.add_several_wakefields(wakes, beta)
+        setattr(self, attribute_name, wake_sum)
+        if attribute_name not in self.sum_names:
+            self.sum_names = np.append(self.sum_names, attribute_name)
+            
+    def sum_by_name_all(self):
+        """
+        Sum all the elements with the same name in the model into sum_name.
+        """
+        for name in self.names:
+            self.sum_by_name(name)
+                    
     def add(self, wakefield_list, wakefiled_postions):
         """
         Add elements to the model.
@@ -462,10 +625,115 @@ class ImpedanceModel(Element):
         wakefiled_postions : list
             Longitudinal positions corresponding to the added Wakfields.
         """
-        
         if (wakefield_list is not None) and (wakefiled_postions is not None):
             for wakefield in wakefield_list:
                 self.wakefields.append(wakefield)
                 
             for position in wakefiled_postions:
-                self.positions = np.append(self.positions, position)
\ No newline at end of file
+                self.positions = np.append(self.positions, position)
+                
+        self.update_name_list()
+                
+    def add_multiple_elements(self, wakefield, wakefiled_postions):
+        """
+        Add the same element at different locations to the model.
+
+        Parameters
+        ----------
+        wakefield : Wakefield object
+            Wakefield to add to the model.
+        wakefiled_postions : list
+            Longitudinal positions corresponding to the added Wakfield.
+        """
+        for position in wakefiled_postions:
+            self.positions = np.append(self.positions, position)
+            self.wakefields.append(wakefield)
+            
+        self.update_name_list()
+            
+    def plot_area(self, Z_type="Zlong", component="real", sigma=None, 
+                  attr_list=None):
+        """
+        Plot the contributions of different kind of wakefields.
+
+        Parameters
+        ----------
+        Z_type : str, optional
+            Type of impedance to plot.
+        component : str, optional
+            Component to plot, can be "real" or "imag". 
+        sigma : float, optional
+            RMS bunch length in [s] to use for the spectral density. If equal
+            to None, the spectral density is not plotted.
+        attr_list : list or array of str, optional
+            Attributes to plot. 
+
+        """
+        if attr_list is None:
+            attr_list = self.sum_names[self.sum_names != "sum"]
+        
+        # sort plot by decresing area size        
+        area = np.zeros((len(attr_list),))
+        for index, attr in enumerate(attr_list):
+            sum_imp = getattr(getattr(self, attr), Z_type)
+            area[index] = trapz(sum_imp.data[component], sum_imp.data.index)
+        sorted_index = area.argsort()[::-1]
+        
+        # Init fig
+        fig = plt.figure()
+        ax = fig.gca()
+        zero_impedance = getattr(self.sum, Z_type)*0
+        total_imp = 0
+        legend = []
+        
+        if sigma is not None:
+            legend.append("Spectral density for sigma = " + str(sigma) + " s")
+        
+        # Main plot
+        for index in  sorted_index:
+            attr = attr_list[index]
+            # Set all impedances with common indexes using + zero_impedance
+            sum_imp = getattr(getattr(self, attr), Z_type) + zero_impedance
+            ax.fill_between(sum_imp.data.index, total_imp, 
+                            total_imp + sum_imp.data[component])
+            total_imp += sum_imp.data[component]
+            if attr_list is None:
+                legend.append(attr[4:])
+            else:
+                legend.append(attr)
+            
+        if sigma is not None:
+            spect = tools.spectral_density(zero_impedance.data.index, sigma)
+            spect = spect/spect.max()*total_imp.max()
+            ax.plot(sum_imp.data.index, spect, 'r', linewidth=2.5)
+        
+        ax.legend(legend)            
+        ax.set_xlabel("Frequency [Hz]")
+        
+        if Z_type == "Zlong":
+            ax.set_ylabel(Z_type + " [Ohm] - " + component + " part")
+        else:
+            ax.set_ylabel(Z_type + " [Ohm/m] - " + component + " part")
+            
+    def summary(self,  sigma, attr_list=None, list_components=None):
+        
+        if attr_list is None:
+            attr_list = self.sum_names
+        if list_components is None:
+            list_components = self.sum.impedance_components
+        
+        shape_array = (len(list_components), len(attr_list))
+        loss_array = np.zeros(shape_array)
+        
+        for i, attr in enumerate(attr_list):
+            for j, component_name in enumerate(list_components):
+                try:
+                    impedance = getattr(getattr(self, attr), component_name)
+                    loss_array[j,i] = tools.loss_factor(impedance, sigma)
+                except AttributeError:
+                        pass
+                    
+        summary = pd.DataFrame(loss_array.T, index=attr_list, 
+                               columns=list_components)
+        return summary
+    
\ No newline at end of file