Add option to monitor and plot FFT data

tracking/monitors/monitors.py

@@ -17,6 +17,7 @@ from mbtrack2.tracking.particles import Bunch, Beam
 from mbtrack2.tracking.rf import CavityResonator
 from scipy.interpolate import interp1d
 from abc import ABCMeta
+from scipy.fft import rfft, rfftfreq
 from mpi4py import MPI
 
 class Monitor(Element, metaclass=ABCMeta):
@@ -812,7 +813,8 @@ class WakePotentialMonitor(Monitor):
 
 class TuneMonitor(Monitor):
     """
-    Monitor tunes in horizontal, vertical, and longitudinal plane. 
+    Monitor tunes and the Fourier transform (using FFT algorithm) of the 
+    osciallation in horizontal, vertical, and longitudinal plane. 
     
     Parameters
     ----------
@@ -822,8 +824,14 @@ class TuneMonitor(Monitor):
     mp_number : int or float
         Total number of macro-particles in the bunch.
     sample_size : int or float
-        Number of macro-particles to be used for tune computation. This number
-        needs not to exceed mp_number.
+        Number of macro-particles to be used for tune and FFT computation.
+        This number cannot exceed mp_number.
+    save_tune : bool, optional
+        If True, tune data is saved. 
+    save_fft : bool, optional
+        If True, FFT data is saved.    
+    n_fft : int or float, optional
+        The number of points used for FFT computation.
     file_name : string, optional
         Name of the HDF5 where the data will be stored. Must be specified
         the first time a subclass of Monitor is instancied and must be None
@@ -846,19 +854,22 @@ class TuneMonitor(Monitor):
     Methods
     -------
     track(bunch):
-        Save tune data.
+        Save tune and/or FFT data.
     
     """
     
-    def __init__(self, ring, bunch_number, mp_number, sample_size, file_name=None, 
-                 save_every=10, buffer_size=5, total_size=10, mpi_mode=True):
+    def __init__(self, ring, bunch_number, mp_number, sample_size, save_tune=True,
+                 save_fft=False, n_fft=10000, file_name=None, save_every=10, 
+                 buffer_size=5, total_size=10, mpi_mode=True):
         
         self.ring = ring
         self.bunch_number = bunch_number
         group_name = "TuneData_" + str(self.bunch_number)
-
-        dict_buffer = {"tune":(3, buffer_size), "tune_spread":(3, buffer_size,)}
-        dict_file = {"tune":(3, total_size), "tune_spread":(3, total_size,)}
+        
+        dict_buffer = {"tune":(3, buffer_size), "tune_spread":(3, buffer_size,),
+                       "fft":(3, n_fft//2+1, buffer_size)}
+        dict_file = {"tune":(3, total_size), "tune_spread":(3, total_size,),
+                     "fft":(3, n_fft//2+1, total_size)}
         
         self.monitor_init(group_name, save_every, buffer_size, total_size,
                           dict_buffer, dict_file, file_name, mpi_mode)
@@ -874,11 +885,17 @@ class TuneMonitor(Monitor):
         index = np.arange(0, int(mp_number))
         self.index_sample = sorted(random.sample(list(index), self.sample_size))
         
-        self.save_every = save_every
         self.save_count = 0
-        
         self.buffer_count = 1
         
+        self.save_tune = save_tune
+        self.save_fft = save_fft
+        self.save_every = save_every
+        
+        if self.save_fft is True :
+            self.n_fft = n_fft 
+            self.fourier_save = np.zeros((3, self.n_fft//2+1, buffer_size))
+        
     def track(self, object_to_save):
         """
         Save tune data.
@@ -909,10 +926,9 @@ class TuneMonitor(Monitor):
             
             self.save_count += 1
             
-            if self.track_count > 0:
-                if self.track_count % self.save_every == 0:
-                    self.to_buffer(bunch)
-                    self.save_count = 0
+            if self.track_count > 0 and self.track_count % self.save_every == 0:
+                self.to_buffer(bunch)
+                self.save_count = 0
     
             self.track_count += 1          
         
@@ -921,10 +937,19 @@ class TuneMonitor(Monitor):
         A method to hold saved data before writing it to the output file.
 
         """
-        mean, spread = self.get_tune(bunch)
+        
         self.time[self.buffer_count] = self.track_count
-        self.tune[:, self.buffer_count] = mean
-        self.tune_spread[:, self.buffer_count] = spread
+        
+        if self.save_tune is True:
+            mean, spread = self.get_tune(bunch)
+            self.tune[:, self.buffer_count] = mean
+            self.tune_spread[:, self.buffer_count] = spread
+        
+        if self.save_fft is True:
+            fx, fy, ftau = self.get_fft()
+            self.fourier_save[0,:,self.buffer_count] = fx
+            self.fourier_save[1,:,self.buffer_count] = fy
+            self.fourier_save[2,:,self.buffer_count] = ftau
         
         self.buffer_count += 1
         
@@ -940,12 +965,18 @@ class TuneMonitor(Monitor):
         self.file[self.group_name]["time"][self.write_count*self.buffer_size:(
                     self.write_count+1)*self.buffer_size] = self.time
 
-        self.file[self.group_name]["tune"][:, 
-                self.write_count * self.buffer_size:(self.write_count+1) * 
-                self.buffer_size] = self.tune
-        self.file[self.group_name]["tune_spread"][:, 
-                self.write_count * self.buffer_size:(self.write_count+1) * 
-                self.buffer_size] = self.tune_spread
+        if self.save_tune is True:
+            self.file[self.group_name]["tune"][:, 
+                    self.write_count * self.buffer_size:(self.write_count+1) * 
+                    self.buffer_size] = self.tune
+            self.file[self.group_name]["tune_spread"][:, 
+                    self.write_count * self.buffer_size:(self.write_count+1) * 
+                    self.buffer_size] = self.tune_spread
+            
+        if self.save_fft is True:
+            self.file[self.group_name]["fft"][:,:, 
+                    self.write_count * self.buffer_size:(self.write_count+1) * 
+                    self.buffer_size] = self.fourier_save
             
         self.file.flush()
         self.write_count += 1
@@ -989,6 +1020,27 @@ class TuneMonitor(Monitor):
         
         return (mean, spread)
     
+    def get_fft(self):
+        """
+        Compute the Fourier transform (using FFT algorithm) of the 
+        osciallation in horizontal, vertical, and longitudinal plane. 
+
+        Returns
+        -------
+        fourier_x_avg, fourier_y_avg, fourier_tau_avg : ndarray
+            The average of the transformed input in each plane.
+
+        """
+        fourier_x = rfft(self.x, n=self.n_fft)
+        fourier_y = rfft(self.y, n=self.n_fft)
+        fourier_tau = rfft(self.tau, n=self.n_fft)
+        
+        fourier_x_avg =  np.mean(abs(fourier_x),axis=0)
+        fourier_y_avg =  np.mean(abs(fourier_y),axis=0)
+        fourier_tau_avg =  np.mean(abs(fourier_tau),axis=0)
+        
+        return (fourier_x_avg, fourier_y_avg, fourier_tau_avg)
+        
 class CavityMonitor(Monitor):
     """
     Monitor a CavityResonator object and save attributes (mean, std, emit and current).
@@ -1051,5 +1103,4 @@ class CavityMonitor(Monitor):
                 self.to_buffer(cavity)
             else:                            
                 raise TypeError("cavity should be a CavityResonator object.")
-        self.track_count += 1
-    
\ No newline at end of file
+        self.track_count += 1       

tracking/monitors/plotting.py

@@ -13,6 +13,7 @@ import matplotlib as mpl
 import seaborn as sns
 import h5py as hp
 import random
+from scipy.fft import rfftfreq
 
 def plot_beamdata(filename, dataset, dimension=None, stat_var=None, x_var="time"):
     """
@@ -492,7 +493,9 @@ def plot_wakedata(filename, bunch_number, wake_type="Wlong", start=0,
     elif streak_plot is True:
         return fig2
     
-def plot_tunedata(filename, bunch_number):
+def plot_tunedata(filename, bunch_number, ring=None, plot_tune=True, plot_fft=False,
+                  dimension='x', min_tune=0, max_tune=0.5, min_turn=None, 
+                  max_turn=None, streak_plot=True, profile_plot=False):
     """
     Plot data recorded by TuneMonitor.
     
@@ -503,6 +506,24 @@ def plot_tunedata(filename, bunch_number):
     bunch_number : int
         Bunch to plot. This has to be identical to 'bunch_number' parameter in 
         'BunchMonitor' object.
+    ring : Synchrotron object, optional
+        The ring configuration that is used in TuneMonitor. If None, the default
+        value of the revolution period and the revolution frequency are used,
+        which are 1.183 us and 0.845 MHz, respectively.
+    plot_tune : bool, optional
+        If True, tune data is plotted.
+    plot_fft : bool, optional
+        If True, FFT data is plotted.
+    dimension : {'x', 'y', 's'}
+        Option to plot FFT data in horizontal, vertical, or longitudinal plane.
+    min_tune, max_tune : int or float, optional
+        The minimum and the maximum tune values to plot FFT data.
+    min_turn, max_turn : int or float, optional
+        The minimum and the maximum number of turns to plot FFT data.
+    streak_plot : bool, optional
+        If True, the FFT data is plotted as a streak plot.
+    bunch_profile : bool, optional.
+        If True, the FFT data is plotted as line profiles.
         
     Return
     ------
@@ -515,20 +536,93 @@ def plot_tunedata(filename, bunch_number):
     
     group = "TuneData_{0}".format(bunch_number)  # Data group of the HDF5 file
     time = file[group]["time"]
-    tune = file[group]["tune"]
-    tune_spread = file[group]["tune_spread"]
-        
-    fig1, ax1 = plt.subplots()        
-    ax1.errorbar(x=time[1:], y=tune[0,1:], yerr=tune_spread[0,1:])
-    ax1.errorbar(x=time[1:], y=tune[1,1:], yerr=tune_spread[1,1:])
-    ax1.set_xlabel("Turn number")
-    ax1.set_ylabel("Transverse tunes")
-    plt.legend(["x","y"])
     
-    fig2, ax2 = plt.subplots()        
-    ax2.errorbar(x=time[1:], y=tune[2,1:], yerr=tune_spread[2,1:])
-    ax2.set_xlabel("Turn number")
-    ax2.set_ylabel("Synchrotron tune")
+    if plot_tune is True:
+        tune = file[group]["tune"]
+        tune_spread = file[group]["tune_spread"]
+            
+        fig1, ax1 = plt.subplots()        
+        ax1.errorbar(x=time[1:], y=tune[0,1:], yerr=tune_spread[0,1:])
+        ax1.errorbar(x=time[1:], y=tune[1,1:], yerr=tune_spread[1,1:])
+        ax1.set_xlabel("Turn number")
+        ax1.set_ylabel("Transverse tunes")
+        plt.legend(["x","y"])
+        
+        fig2, ax2 = plt.subplots()        
+        ax2.errorbar(x=time[1:], y=tune[2,1:], yerr=tune_spread[2,1:])
+        ax2.set_xlabel("Turn number")
+        ax2.set_ylabel("Synchrotron tune")
+        
+    if plot_fft is True:
+        if ring is None:
+            T0 = 1.183e-06
+            f0 = 0.845e6
+        else:
+            T0 = ring.T0
+            f0 = ring.f0
+        
+        n_freq = file[group]['fft'].shape[1]
+        freq = rfftfreq((n_freq-1)*2, T0)
+        tune_fft = freq / f0
+        
+        dimension_dict = {'x':0, 'y':1, 's':2}
+        axis = dimension_dict[dimension]
+        
+        fourier_save = file[group]['fft'][axis]
+        
+        if max_turn is None:
+            max_turn = time[-1]
+        if min_turn is None:
+            min_turn = time[1]
+            
+        min_tune_iloc = np.where(tune_fft >= min_tune)[0][0]
+        max_tune_iloc = np.where(tune_fft <= max_tune)[0][-1]
+        save_every = int(time[1] - time[0])
+        min_turn_iloc = min_turn // save_every
+        max_turn_iloc = max_turn // save_every
+        
     
+        if streak_plot is True:
+            fig3, ax3 = plt.subplots()
+            cmap = plt.get_cmap('Blues')
+        
+            c = ax3.imshow(np.transpose(np.log(
+                          fourier_save[min_tune_iloc:max_tune_iloc+1, 
+                                       min_turn_iloc:max_turn_iloc+1])),
+                          cmap=cmap, origin='lower' , aspect='auto',
+                          extent=[min_tune, max_tune, min_turn, max_turn])
+            ax3.set_xlabel('$Q_{}$'.format(dimension))
+            ax3.set_ylabel("Turns")
+            cbar = fig3.colorbar(c, ax=ax3)
+            cbar.set_label("log FFT amplitude") 
+            
+        if profile_plot is True:
+            fig4, ax4 = plt.subplots()
+            ax4.plot(tune_fft[min_tune_iloc:max_tune_iloc+1], 
+                     fourier_save[min_tune_iloc:max_tune_iloc+1,
+                                  min_turn_iloc:max_turn_iloc+1])
+            ax4.set_xlabel('$Q_{}$'.format(dimension))
+            ax4.set_ylabel("FFT amplitude")
+            ax4.legend(time[min_turn_iloc:max_turn_iloc+1])
+            
     file.close()
-    return (fig1, fig2)  
\ No newline at end of file
+    
+    if plot_tune is True and plot_fft is True:
+        if streak_plot is True and profile_plot is True:
+            return fig1, fig2, fig3, fig4
+        elif streak_plot is True:
+            return fig1, fig2, fig3
+        elif profile_plot is True:
+            return fig1, fig2, fig4
+        
+    elif plot_tune is True:
+        return fig1, fig2
+    
+    elif plot_fft is True:
+        if streak_plot is True and profile_plot is True:
+            return fig3, fig4
+        elif streak_plot is True:
+            return fig3
+        elif profile_plot is True:
+            return fig4
+        
\ No newline at end of file