Impedance : energy_loss method for ImpedanceModel + fix on beam_loss_factor

ImpedanceModel -> new method to compute beam and bunch energy loss
beam_loss_factor -> drop 0 from data only if it is there, avoid bug if it is not

collective_effects/tools.py

@@ -385,7 +385,10 @@ def beam_loss_factor(impedance, frequency, spectrum, ring):
         negative_index = impedance.data.index*-1
         negative_data = impedance.data.set_index(negative_index)
         negative_data["imag"] = -1*negative_data["imag"]
-        negative_data = negative_data.drop(0)
+        try:
+            negative_data = negative_data.drop(0)
+        except KeyError:
+            pass
         
         all_data = impedance.data.append(negative_data)
         all_data = all_data.sort_index()

collective_effects/wakefield.py

@@ -738,4 +738,67 @@ class ImpedanceModel(Element):
         summary = pd.DataFrame(loss_array.T, index=attr_list, 
                                columns=list_components)
         return summary
+    
+    def energy_loss(self, sigma, M, bunch_spacing, I, n_points=10e6):
+        """
+        Compute the beam and bunch loss factor and energy losses for each type 
+        of element in the model assuming Gaussian bunches and constant spacing
+        between bunches.
+
+        Parameters
+        ----------
+        sigma : float
+            RMS bunch length in [s].
+        M : int
+            Number of bunches in the beam.
+        bunch_spacing : float
+            Time between two bunches in [s].
+        I : float
+            Total beam current in [A].
+        n_points : float, optional
+            Number of points used in the frequency spectrums.
+
+        Returns
+        -------
+        summary : Dataframe
+            Contains the beam and bunch loss factor and energy loss for the 
+            full model and for each type of different component.
+
+        """
+        
+        fmax = self.sum.Zlong.data.index.max()
+        fmin = self.sum.Zlong.data.index.min()
+        
+        Q = I*self.ring.T0/M
+        
+        if fmin >= 0:
+            fmin = -1*fmax
+        f = np.linspace(fmin, fmax, int(n_points))
+
+        beam_spect = tools.beam_spectrum(f, M, bunch_spacing, sigma= sigma)
+        
+        bunch_spect = tools.Gaussian_bunch_spectrum(f, sigma)
+        
+        attr_list = self.sum_names
+        
+        loss_array = np.zeros((len(attr_list),2))
+        
+        for i, attr in enumerate(attr_list):
+            try:
+                impedance = getattr(getattr(self, attr), "Zlong")
+                loss_array[i,0] = tools.beam_loss_factor(impedance, f, beam_spect, self.ring)
+                loss_array[i,1] = tools.beam_loss_factor(impedance, f, bunch_spect, self.ring)
+            except AttributeError:
+                pass
+
+        loss_array = loss_array*1e-12
+        summary = pd.DataFrame(loss_array, index=attr_list, 
+                               columns=["loss factor (beam) [V/pC]", "loss factor (bunch) [V/pC]"])
+        
+        summary["P (beam) [W]"] = summary["loss factor (beam) [V/pC]"]*1e12*Q**2/(self.ring.T0)
+        summary["P (bunch) [W]"] = summary["loss factor (bunch) [V/pC]"]*1e12*Q**2/(self.ring.T0)*M
+                
+        return summary
+        
+        
     
\ No newline at end of file