Add WakePotential class

Adding a class for tracking wake potential

tracking/element.py

@@ -9,9 +9,13 @@ included in the tracking.
 """
 
 import numpy as np
+import pandas as pd
 from abc import ABCMeta, abstractmethod
 from functools import wraps
 from tracking.particles import Beam
+from scipy import signal
+import matplotlib.pyplot as plt
+from scipy.interpolate import interp1d
 
 class Element(metaclass=ABCMeta):
     """
@@ -95,6 +99,203 @@ class LongitudinalMap(Element):
         bunch["delta"] -= self.ring.U0 / self.ring.E0
         bunch["tau"] -= self.ring.ac * self.ring.T0 * bunch["delta"]
         
+class WakePotential(Element):
+    """
+    Resonator model based wake potential calculation for one turn.
+    
+    Parameters
+    ----------
+    ring : Synchrotron object.
+    Q_factor : float, optional
+        Resonator quality factor. The default value is 1.
+    f_res : float, optional
+        Resonator resonance frequency in [Hz]. The default value is 10e9 Hz.
+    R_shunt : float, optional
+        Resonator shunt impedance in [Ohm]. The default value is 100 Ohm.
+    n_bin : int, optional
+        Number of bins for constructing the longitudinal bunch profile.
+        The default is 65.
+        
+    Attributes
+    ----------
+    rho : array of shape (n_bin, )
+        Bunch charge density profile.
+    tau : array of shape (n_bin + time_extra, )
+        Time array starting from the head of the bunch until the wake tail 
+        called timestop.
+        
+        The length of time_extra is determined by the last position of the
+        bunch time_bunch[-1], timestop, and the mean bin width of the bunch
+        profile mean_bin_size as
+            len(time_extra) = (timestop - time_bunch[-1]) / mean_bin_size
+    W_long : array of shape (n_bin + time_extra, )
+        Wakefunction profile.
+    W_p : array of shape (n_bin + time_extra, )
+        Wake potential profile.
+    wp : array of shape (mp_number, )
+        Wake potential exerted on each macro-particle.
+    
+    Methods
+    -------
+    charge_density(bunch, n_bin=65)
+        Calculate bunch charge density.
+    plot(self, var, plot_rho=True)
+        Plotting wakefunction or wake potential.
+    track(bunch)
+        Tracking method for the element.
+        
+    """
+    
+    def __init__(self, ring, Q_factor=1, f_res=10e9, R_shunt=100, n_bin=65):       
+        self.ring = ring
+        self.n_bin = n_bin
+        
+        self.Q_factor = Q_factor
+        self.omega_res = 2*np.pi*f_res
+        self.R_shunt = R_shunt
+        
+        if Q_factor >= 0.5:
+            self.Q_factor_p = np.sqrt(self.Q_factor**2 - 0.25)
+            self.omega_res_p = (self.omega_res*self.Q_factor_p)/self.Q_factor
+        else:
+            self.Q_factor_pp = np.sqrt(0.25 - self.Q_factor**2)
+            self.omega_res_p = (self.omega_res*self.Q_factor_pp)/self.Q_factor
+            
+    def charge_density(self, bunch, n_bin):
+        self.bins = bunch.binning(n_bin=self.n_bin)
+        self.bin_data = self.bins[2]
+        self.bin_size = self.bins[0].length
+        
+        self.rho = bunch.charge_per_mp*self.bin_data/ \
+            (self.bin_size*bunch.charge)
+    
+    def init_timestop(self):
+        self.timestop = round(np.log(1000)/self.omega_res*2*self.Q_factor, 12)
+    
+    def time_array(self):        
+        time_bunch = self.bins[0].mid
+        mean_bin_size = np.mean(self.bin_size)
+        time_extra = np.arange(start = time_bunch[-1]+mean_bin_size, 
+                               stop = self.timestop, step = mean_bin_size)
+        
+        self.tau = np.concatenate((time_bunch,time_extra))
+     
+    def long_wakefunction(self):
+        w_list = []
+        if self.Q_factor >= 0.5:
+            for t in self.tau: 
+                if t >= 0:
+                    w_long = -(self.omega_res*self.R_shunt/self.Q_factor)*\
+                        np.exp(-self.omega_res*t/(2*self.Q_factor))*\
+                            (np.cos(self.omega_res_p*t)-\
+                             np.sin(self.omega_res_p*t)/(2*self.Q_factor_p))
+                else: 
+                    w_long = 0
+                    
+                w_list.append(w_long)
+                        
+        elif self.Q_factor < 0.5:
+            for t in self.tau: 
+                if t >= 0:
+                    w_long = -(self.omega_res*self.R_shunt/self.Q_factor)*\
+                        np.exp(-self.omega_res*t/(2*self.Q_factor))*\
+                            (np.cosh(self.omega_res_p*t)-\
+                             np.sinh(self.omega_res_p*t)/(2*self.Q_factor_pp))
+                else:
+                    w_long = 0
+                    
+                w_list.append(w_long)
+                
+        self.W_long = np.array(w_list)
+                
+    def wake_potential(self): 
+        self.W_p = signal.convolve(self.W_long*1e-12, self.rho, mode="same") 
+    
+    def plot(self, var, plot_rho=True):
+        """
+        Plotting wakefunction or wake potential.
+
+        Parameters
+        ----------
+        var : {'W_p', 'W_long' }
+            If 'W_p', the wake potential is plotted. 
+            If 'W_long', the wakefunction is plotted.            
+        plot_rho : bool, optional
+            Overlay the bunch charge density profile on the plot.
+            The default is True.
+
+        Returns
+        -------
+        fig : Figure
+            Figure object with the plot on it.
+
+        """
+        
+        fig, ax = plt.subplots()
+            
+        if var == "W_p":
+            ax.plot(self.tau*1e12, self.W_p*1e-12)
+            
+            ax.set_xlabel("$\\tau$ (ps)")
+            ax.set_ylabel("W$_p$ (V/pC)")
+    
+        elif var == "W_long":
+            ax.plot(self.tau*1e12, self.W_long*1e-12)
+            ax.set_xlabel("$\\tau$ (ps)")
+            ax.set_ylabel("W$_{||}$ ($\\Omega$/ps)")
+            
+        if plot_rho is True:
+            rho_array = np.array(self.rho)
+            rho_rescaled = rho_array/max(rho_array)*max(self.W_p)
+            
+            ax.plot(self.bins[0].mid*1e12, rho_rescaled*1e-12)
+            
+        else:
+            pass
+            
+        return fig
+    
+    def check_wake_tail(self):
+        """
+        Checking whether the full wakefunction is obtained by the calculated
+        initial timestop.
+
+        """
+        
+        ratio = np.abs(min(self.W_long) / self.W_long[-6:-1])
+        while any(ratio < 1000):
+            # Initial timestop is too short. 
+            # Extending timestop by 50 ps and recalculating."
+            self.timestop += 50e-12      
+            self.time_array()
+            self.long_wakefunction()
+            ratio = np.abs(min(self.W_long) / self.W_long[-6:-1])
+        
+    @Element.parallel
+    def track(self, bunch):
+        """
+        Tracking method for the element.
+        No bunch to bunch interaction, so written for Bunch objects and
+        @Element.parallel is used to handle Beam objects.
+
+        Parameters
+        ----------
+        bunch : Bunch or Beam object.
+        
+        """
+
+        self.charge_density(bunch, n_bin = self.n_bin)
+        self.init_timestop()
+        self.time_array()
+        self.long_wakefunction()
+        self.check_wake_tail()
+        self.wake_potential()
+        
+        f = interp1d(self.tau, self.W_p, fill_value = 0, bounds_error = False)
+        self.wp = f(bunch["tau"])
+        
+        bunch["delta"] += self.wp * bunch.charge / self.ring.E0
+        
 class SynchrotronRadiation(Element):
     """
     Element to handle synchrotron radiation, radiation damping and quantum