Rework WakePotential

Now based on two time arrays!
The bunch profile is computed from its own time array and interpolated into the wake function time array.
The input wake functions must be uniformly sampled!
Reference computation now gives loss factor and kick factor.
Tested with Wlong and Wxdip/Wydip -> OK
Untested for Wxquad/Wyquad.

tracking/wakepotential.py

@@ -16,14 +16,20 @@ from mbtrack2.tracking.element import Element
    
 class WakePotential(Element):
     """
-    Compute a wake potential from wake functions by performing a convolution 
-    with a bunch charge profile.
+    Compute a wake potential from uniformly sampled wake functions by 
+    performing a convolution with a bunch charge profile.
+    
+    Two different time bases are used. The first one is controled by the n_bin
+    parameter and is used to compute the bunch profile. Then the bunch profile
+    is interpolated on the wake function time base which is used to perform the
+    convolution to get the wake potential.
     
     Parameters
     ----------
     ring : Synchrotron object
     wakefield : Wakefield object
-        Wakefield object which contains the wake functions to be used.
+        Wakefield object which contains the wake functions to be used. The wake
+        functions must be uniformly sampled!
     n_bin : int, optional
         Number of bins for constructing the longitudinal bunch profile.
         
@@ -40,17 +46,24 @@ class WakePotential(Element):
     -------
     charge_density(bunch)
         Compute bunch charge density profile in [1/s].
-    dipole_moment(bunch, plane)
+    dipole_moment(bunch, plane, tau0)
         Return the dipole moment of the bunch computed on the same time array 
-        as the bunch profile.
+        as the wake function.
+    prepare_wakefunction(wake_type)
+        Prepare the wake function of a given wake_type to be used for the wake
+        potential computation. 
     get_wakepotential(bunch, wake_type)
-        Return the wake potential computed on the same time array as the bunch 
-        profile.
+        Return the wake potential computed on the wake function time array 
+        limited to the bunch profile.
     track(bunch)
         Tracking method for the element.
-    plot_last_wake(wake_type, plot_rho=True, plot_dipole=False)
+    plot_last_wake(wake_type)
         Plot the last wake potential of a given type computed during the last
         call of the track method.
+    reference_loss(bunch)
+        Calculate the loss factor and kick factor from the wake potential and 
+        compare it to a reference value assuming a Gaussian bunch computed in 
+        the frequency domain.
         
     """
     
@@ -105,16 +118,21 @@ class WakePotential(Element):
         if len(self.tau) != len(self.rho):
             self.tau = np.append(self.tau, self.tau[-1] + self.dtau)
             
-    def dipole_moment(self, bunch, plane):
+        self.tau_mean = np.mean(self.tau)
+        self.tau -= self.tau_mean
+            
+    def dipole_moment(self, bunch, plane, tau0):
         """
         Return the dipole moment of the bunch computed on the same time array 
-        as the bunch profile.
+        as the wake function.
 
         Parameters
         ----------
         bunch : Bunch object
         plane : str
             Plane on which the dipole moment is computed, "x" or "y".
+        tau0 : array
+            Time array on which the dipole moment will be interpolated, in [s].
 
         Returns
         -------
@@ -128,7 +146,7 @@ class WakePotential(Element):
         dipole = dipole/self.profile
         dipole[np.isnan(dipole)] = 0
         
-        # Add N values to get same size as tau/rho
+        # Add N values to get same size as tau/profile
         if self.n_bin % 2 == 0:
             N = int(self.n_bin/2)
             dipole = np.append(dipole, np.zeros(N))
@@ -138,13 +156,98 @@ class WakePotential(Element):
             dipole = np.append(dipole, np.zeros(N))
             dipole = np.insert(dipole, 0, np.zeros(N+1))
             
-        setattr(self, "dipole_" + plane, dipole)
-        return dipole
+        # Interpole on tau0 to get the same size as W0
+        interp_dipole = interp1d(self.tau, dipole, fill_value=0, 
+                             bounds_error=False)
+        dipole0 = interp_dipole(tau0)
+            
+        setattr(self, "dipole_" + plane, dipole0)
+        return dipole0
+    
+    
+    def prepare_wakefunction(self, wake_type):
+        """
+        Prepare the wake function of a given wake_type to be used for the wake
+        potential computation. 
+        
+        The new time array keeps the same sampling time as given in the 
+        WakeFunction definition but is restricted to the bunch profile time 
+        array.
+
+        Parameters
+        ----------
+        wake_type : str
+            Type of the wake function to prepare: "Wlong", "Wxdip", ...
+
+        Returns
+        -------
+        tau0 : array
+            Time base of the wake function in [s].
+        dtau0 : float
+            Difference between two points of the wake function time base in 
+            [s].
+        W0 : array
+            Wake function array in [V/C] or [V/C/m].
+
+        """
+        
+        try:
+            tau0 = getattr(self, "tau0_" + wake_type)
+            dtau0 = getattr(self, "dtau0_" + wake_type)
+            W0 = getattr(self, "W0_" + wake_type)
+        except AttributeError:
+            tau0 = np.array(getattr(self.wakefield, wake_type).data.index)
+            dtau0 = tau0[1] - tau0[0]
+            W0 = np.array(getattr(self.wakefield, wake_type).data["real"])
+        
+        # Keep only the wake function on the rho window
+        ind = np.all([min(self.tau[0], 0) < tau0, max(self.tau[-1], 0) > tau0],
+                     axis=0)
+        tau0 = tau0[ind]
+        W0 = W0[ind]
+        
+        # Check the wake function window for assymetry
+        assym = (np.abs(tau0[-1]) - np.abs(tau0[0])) / dtau0
+        n_assym = int(np.floor(assym))
+        if np.floor(assym) > 1:
+            
+            # add at head
+            if np.abs(tau0[-1]) >  np.abs(tau0[0]):
+                tau0 = np.arange(tau0[0] - dtau0*n_assym, 
+                                 tau0[-1] + dtau0, 
+                                 dtau0)
+                n_to_add = len(tau0) - len(W0)
+                W0 = np.insert(W0, 0, np.zeros(n_to_add))
+                
+            # add at tail
+            elif np.abs(tau0[0]) >  np.abs(tau0[-1]):
+                tau0 = np.arange(tau0[0], 
+                                 tau0[-1] + dtau0*(n_assym+1), 
+                                 dtau0)
+                n_to_add = len(tau0) - len(W0)
+                W0 = np.insert(W0, 0, np.zeros(n_to_add))
+                
+        # Check is the wf is shorter than rho then add zeros
+        if (tau0[0] > self.tau[0]) or (tau0[-1] < self.tau[-1]):
+            n = max(int(np.ceil((tau0[0] - self.tau[0])/dtau0)),
+                    int(np.ceil((self.tau[-1] - tau0[-1])/dtau0)))
+            
+            tau0 = np.arange(tau0[0] - dtau0*n, 
+                             tau0[-1] + dtau0*(n+1), 
+                             dtau0)
+            W0 = np.insert(W0, 0, np.zeros(n))
+            n_to_add = len(tau0) - len(W0)
+            W0 = np.insert(W0, len(W0), np.zeros(n_to_add))
+            
+        setattr(self, "tau0_" + wake_type, tau0)
+        setattr(self, "dtau0_" + wake_type, dtau0)
+        setattr(self, "W0_" + wake_type, W0)
+        return (tau0, dtau0, W0)
         
     def get_wakepotential(self, bunch, wake_type):
         """
-        Return the wake potential computed on the same time array as the bunch 
-        profile.
+        Return the wake potential computed on the wake function time array 
+        limited to the bunch profile.
 
         Parameters
         ----------
@@ -159,24 +262,27 @@ class WakePotential(Element):
 
         """
 
-        tau0 = getattr(self.wakefield, wake_type).data.index
-        W0 = getattr(self.wakefield, wake_type).data["real"]
-        interp_func_W0 = interp1d(tau0, W0, fill_value=0, 
+        (tau0, dtau0, W0) = self.prepare_wakefunction(wake_type)
+        
+        interp_profile = interp1d(self.tau, self.rho, fill_value=0, 
                                      bounds_error=False)
-        W = interp_func_W0(self.tau - np.mean(self.tau))
+        
+        profile0 = interp_profile(tau0)
+        
         if wake_type == "Wlong" or wake_type == "Wxquad" or wake_type == "Wyquad":
-            Wp = signal.convolve(self.rho, W*-1, mode='same')*self.dtau
+            Wp = signal.convolve(profile0, W0*-1, mode='same')*dtau0
         elif wake_type == "Wxdip":
-            dipole = self.dipole_moment(bunch, "x")
-            Wp = signal.convolve(self.rho*dipole, W*-1, mode='same')*self.dtau
+            dipole0 = self.dipole_moment(bunch, "x", tau0)
+            Wp = signal.convolve(profile0*dipole0, W0*-1, mode='same')*dtau0
         elif wake_type == "Wydip":
-            dipole = self.dipole_moment(bunch, "y")
-            Wp = signal.convolve(self.rho*dipole, W*-1, mode='same')*self.dtau
+            dipole0 = self.dipole_moment(bunch, "y", tau0)
+            Wp = signal.convolve(profile0*dipole0, W0*-1, mode='same')*dtau0
         else:
             raise ValueError("This type of wake is not taken into account.")
         
+        setattr(self,"profile0_" + wake_type, profile0)
         setattr(self, wake_type, Wp)
-        return Wp
+        return tau0, Wp
     
     @Element.parallel
     def track(self, bunch):
@@ -193,8 +299,8 @@ class WakePotential(Element):
 
         self.charge_density(bunch)
         for wake_type in self.types:
-            Wp = self.get_wakepotential(bunch, wake_type)
-            f = interp1d(self.tau, Wp, fill_value = 0, bounds_error = False)
+            tau0, Wp = self.get_wakepotential(bunch, wake_type)
+            f = interp1d(tau0 + self.tau_mean, Wp, fill_value = 0, bounds_error = False)
             Wp_interp = f(bunch["tau"])
             if wake_type == "Wlong":
                 bunch["delta"] += Wp_interp * bunch.charge / self.ring.E0
@@ -209,7 +315,8 @@ class WakePotential(Element):
                 bunch["yp"] += (bunch["y"] * Wp_interp * bunch.charge 
                                 / self.ring.E0)
                 
-    def plot_last_wake(self, wake_type, plot_rho=True, plot_dipole=False):
+    def plot_last_wake(self, wake_type, plot_rho=True, plot_dipole=False, 
+                       plot_wake_function=True):
         """
         Plot the last wake potential of a given type computed during the last
         call of the track method.
@@ -222,6 +329,8 @@ class WakePotential(Element):
             Plot the normalised bunch profile. The default is True.
         plot_dipole : bool, optional
             Plot the normalised dipole moment. The default is False.
+        plot_wake_function : bool, optional
+            Plot the normalised wake function. The default is True.
 
         Returns
         -------
@@ -234,84 +343,86 @@ class WakePotential(Element):
                   "Wydip" : r"$W_{p,y}^{D} (V/pC)$",
                   "Wxquad" : r"$W_{p,x}^{Q} (V/pC/m)$",
                   "Wyquad" : r"$W_{p,y}^{Q} (V/pC/m)$"}
-        fig, ax = plt.subplots()
+        
         Wp = getattr(self, wake_type)
-        ax.plot(self.tau*1e12, Wp*1e-12, label=labels[wake_type])
+        tau0 = getattr(self, "tau0_" + wake_type)
+        
+        fig, ax = plt.subplots()
+        ax.plot(tau0*1e12, Wp*1e-12, label=labels[wake_type])
         ax.set_xlabel("$\\tau$ (ps)")
         ax.set_ylabel(labels[wake_type])
 
         if plot_rho is True:
+            profile0 = getattr(self,"profile0_" + wake_type)
+            profile_rescaled = profile0/max(profile0)*max(np.abs(Wp))
             rho_rescaled = self.rho/max(self.rho)*max(np.abs(Wp))
-            ax.plot(self.tau*1e12, rho_rescaled*1e-12, label=r"$\rho$ (a.u.)")
+            ax.plot(tau0*1e12, profile_rescaled*1e-12, label=r"$\rho$ interpolated (a.u.)")
+            ax.plot((self.tau + self.tau_mean)*1e12, rho_rescaled*1e-12, label=r"$\rho$ (a.u.)", linestyle='dashed')
+            plt.legend()
+            
+        if plot_wake_function is True:
+            W0 = getattr(self, "W0_" + wake_type)
+            W0_rescaled = W0/max(W0)*max(np.abs(Wp))
+            ax.plot(tau0*1e12, W0_rescaled*1e-12, label=r"$W_{function}$ (a.u.)")
             plt.legend()
             
         if plot_dipole is True:
             dipole = getattr(self, "dipole_" + wake_type[1])
             dipole_rescaled = dipole/max(dipole)*max(np.abs(Wp))
-            ax.plot(self.tau*1e12, dipole_rescaled*1e-12, label=r"Dipole moment (a.u.)")
+            ax.plot(tau0*1e12, dipole_rescaled*1e-12, label=r"Dipole moment (a.u.)")
             plt.legend()
             
         return fig
     
-    def energy_loss(self, bunch, compare_to='num'):
+    def reference_loss(self, bunch):
         """
-        Calculate the energy loss from the wake potential and compare it to a 
-        reference value assuming a Gaussian bunch.
+        Calculate the loss factor and kick factor from the wake potential and 
+        compare it to a reference value assuming a Gaussian bunch computed in 
+        the frequency domain.
 
         Parameters
         ----------
         bunch : Bunch object
-        compare_to : str, {"theory", "num"}, optional
-            The method for calculating the reference energy loss. Use 'theory' 
-            for the analytical loss factor of a resonator [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, p.239
+        loss_data : DataFrame
+            An output showing the loss/kick factors compared to the reference 
+            values.
 
         """
+        loss = []
+        loss_0 = []
+        delta_loss = []
+        index = []
+        for wake_type in self.types:
+            tau0, Wp = self.get_wakepotential(bunch, wake_type)
+            profile0 = getattr(self, "profile0_" + wake_type)
+            factorTD = -1*np.trapz(Wp*profile0, tau0)
             
-        if np.where(self.types=="Wlong")[0].size == 0:
-            raise TypeError("No longitudinal wake function data.")
+            if wake_type == "Wxdip":
+                factorTD = factorTD / bunch["x"].mean()
+            if wake_type == "Wydip":
+                factorTD = factorTD / bunch["y"].mean()
             
-        else:
-            Wp = self.get_wakepotential(bunch, "Wlong")
-            Eloss = -1*np.trapz(Wp*self.rho, self.tau)*bunch.charge
-            
-            res = self.wakefield
+            Z = getattr(self.wakefield, "Z" + wake_type[1:])
             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 = ['Energy loss [eV]', 'Reference energy loss [eV]', 'Relative error [%]']
-            energy_loss_data = pd.DataFrame(np.reshape([Eloss, Eloss_0, delta_Eloss], (1,3)), 
-                                            columns=column)
-        return energy_loss_data
-            
-            
-        
-                
-                           
-            
+            factorFD = Z.loss_factor(sigma)
             
+            loss.append(factorTD)
+            loss_0.append(factorFD)
+            delta_loss.append( (factorTD - factorFD) / factorFD *100 )
+            if wake_type == "Wlong":
+                index.append("Wlong [V/C]")
+            else:
+                index.append(wake_type + " [V/C/m]")
             
+            column = ['TD factor', 'FD factor', 'Relative error [%]']
             
+        loss_data = pd.DataFrame(np.array([loss, loss_0, delta_loss]).T, 
+                                 columns=column, 
+                                 index=index)
+        return loss_data