Add energy_loss method and some minor changes

A method to compare energy loss from a wake potential to a loss factor.
Minor changes :
- sigma method in Ring object has been edited.
- add WakePotential in __init__ in tracking module.

tracking/__init__.py

@@ -15,4 +15,5 @@ from mbtrack2.tracking.element import (Element, LongitudinalMap,
                                        TransverseMap, SynchrotronRadiation)
 from mbtrack2.tracking.aperture import (CircularAperture, ElipticalAperture,
                                         RectangularAperture)
+from mbtrack2.tracking.wakepotential import WakePotential
                                        

tracking/synchrotron.py

@@ -227,14 +227,14 @@ class Synchrotron:
     def sigma(self):
         """RMS beam size at equilibrium"""
         sigma = np.zeros((4,))
-        sigma[0] = (self.emit[0]*self.optics.local_beta[0] +
-                    self.optics.local_dispersion[0]**2*self.sigma_delta)**0.5
-        sigma[1] = (self.emit[0]*self.optics.local_alpha[0] +
-                    self.optics.local_dispersion[1]**2*self.sigma_delta)**0.5
-        sigma[2] = (self.emit[1]*self.optics.local_beta[1] +
-                    self.optics.local_dispersion[2]**2*self.sigma_delta)**0.5
-        sigma[3] = (self.emit[1]*self.optics.local_alpha[1] +
-                    self.optics.local_dispersion[3]**2*self.sigma_delta)**0.5
+        sigma[0] = (self.emit[0]*self.optics.beta(10)[0] +
+                    self.optics.dispersion(10)[0]**2*self.sigma_delta)**0.5
+        sigma[1] = (self.emit[0]*self.optics.alpha(10)[0] +
+                    self.optics.dispersion(10)[1]**2*self.sigma_delta)**0.5
+        sigma[2] = (self.emit[1]*self.optics.beta(10)[1] +
+                    self.optics.dispersion(10)[2]**2*self.sigma_delta)**0.5
+        sigma[3] = (self.emit[1]*self.optics.alpha(10)[1] +
+                    self.optics.dispersion(10)[3]**2*self.sigma_delta)**0.5
         return sigma
     
     def synchrotron_tune(self, Vrf):

tracking/wakepotential.py

@@ -9,6 +9,7 @@ effects, both longitudinal and transverse.
 
 import numpy as np
 import matplotlib.pyplot as plt
+import pandas as pd
 from scipy import signal
 from scipy.interpolate import interp1d
 from mbtrack2.tracking.element import Element
@@ -251,3 +252,69 @@ class WakePotential(Element):
             plt.legend()
             
         return fig
+    
+    def energy_loss(self, bunch, compare_to='theory'):
+        """
+        Calculate the energy loss from the wake potential and can be compared
+        to a theoretical value or a numerical value.
+
+        Parameters
+        ----------
+        bunch : Bunch object
+        compare_to : {"theory", "num"}, optional
+            The method of calculating the reference energy loss. Use 'theory' for
+            a theorytecal value computed from the analytical loss factor [1], 
+            or use 'num' for a numerical result obtained from loss_factor method
+            in Impedance class.
+
+        Returns
+        -------
+        energy_loss_data : DataFrame
+            An output showing the enegy loss compared to the reference value.
+            
+        Reference
+        ---------
+        [1] A. W. Chao and M. Tigner, "Handbook of Accelerator Physics and
+            Engineering", 3rd printing, pp.239
+
+        """
+        
+        if np.where(self.types=="Wlong")[0].size == 0:
+            raise TypeError("No longitudinal wake function data.")
+                
+        else:
+            Wp = self.get_wakepotential(bunch, "Wlong")
+            Eloss = -1*np.trapz(Wp*self.rho, self.tau)*bunch.charge
+            
+            res = self.wakefield
+            sigma = bunch['tau'].std()
+            if compare_to == 'theory':
+                Eloss_0 = 0.5*res.wr*res.Rs*1/res.Q *np.exp(-1*res.wr**2*sigma**2) \
+                    * bunch.charge
+            elif compare_to == 'num':
+                Eloss_0 = res.Zlong.loss_factor(sigma)*bunch.charge
+            else:
+                raise KeyError("Reference method not correct. Enter \"theory\" "
+                               "or \"num\" ")
+            
+            delta_Eloss = (Eloss-Eloss_0) / Eloss_0 *100
+            
+            column = ['Eloss', 'Eloss_ref', '% error']
+            energy_loss_data = pd.DataFrame(np.reshape([Eloss, Eloss_0, delta_Eloss], (1,3)), 
+                                            columns=column)
+        return energy_loss_data
+            
+            
+        
+                
+                           
+            
+    
+    
+    
+    
+    
+    
+    
+    
+    
\ No newline at end of file