Fix TuneMonitor

Fix bug when buffer_size=1 in TuneMonitor by switching to init buffer_count to 0. Now total_size should be 1 less.
Improve plot_tunedata

tracking/monitors/monitors.py

@@ -842,7 +842,8 @@ class TuneMonitor(Monitor):
     save_fft : bool, optional
         If True, FFT data is saved.    
     n_fft : int or float, optional
-        The number of points used for FFT computation.
+        The number of points used for FFT computation, if n_fft is bigger than
+        save_every zero-padding is applied.
     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
@@ -856,7 +857,7 @@ class TuneMonitor(Monitor):
         Total size of the save. The following relationships between the 
         parameters must exist: 
             total_size % buffer_size == 0
-            number of call to track / save_every == total_size    
+            number of call to track / save_every == total_size - 1   
     mpi_mode : bool, optional
         If True, open the HDF5 file in parallel mode, which is needed to
         allow several cores to write in the same file at the same time.
@@ -897,7 +898,7 @@ class TuneMonitor(Monitor):
         self.index_sample = sorted(random.sample(list(index), self.sample_size))
         
         self.save_count = 0
-        self.buffer_count = 1
+        self.buffer_count = 0
         
         self.save_tune = save_tune
         self.save_fft = save_fft
@@ -930,7 +931,7 @@ class TuneMonitor(Monitor):
         else:
             raise TypeError("object_to_save should be a Beam or Bunch object.")
         
-        if skip is not True:
+        if skip is False:
             self.x[:, self.save_count] = bunch["x"][self.index_sample]
             self.y[:, self.save_count] = bunch["y"][self.index_sample]
             self.tau[:, self.save_count] = bunch["tau"][self.index_sample]

tracking/monitors/plotting.py

@@ -493,7 +493,7 @@ def plot_wakedata(filename, bunch_number, wake_type="Wlong", start=0,
     elif streak_plot is True:
         return fig2
     
-def plot_tunedata(filename, bunch_number, ring=None, plot_tune=True, plot_fft=False,
+def plot_tunedata(filename, bunch_number, f0, 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):
     """
@@ -506,12 +506,10 @@ def plot_tunedata(filename, bunch_number, ring=None, plot_tune=True, plot_fft=Fa
     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.
+    f0 : float
+        Revolution frequency of the ring used for the tracking in [Hz].
     plot_tune : bool, optional
-        If True, tune data is plotted.
+        If True, tune data usinf NAFF is plotted.
     plot_fft : bool, optional
         If True, FFT data is plotted.
     dimension : {'x', 'y', 's'}
@@ -522,13 +520,13 @@ def plot_tunedata(filename, bunch_number, ring=None, plot_tune=True, plot_fft=Fa
         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.
+    profile_plot : bool, optional.
+        If True, the FFT data is plotted as line plots.
         
     Return
     ------
     fig : Figure
-        Figure object with the plot on it.
+        Figure objects with the plot on it.
 
     """
     
@@ -537,30 +535,29 @@ def plot_tunedata(filename, bunch_number, ring=None, plot_tune=True, plot_fft=Fa
     group = "TuneData_{0}".format(bunch_number)  # Data group of the HDF5 file
     time = file[group]["time"]
     
+    fig_to_return = []
+    
     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.errorbar(x=time, y=tune[0,:], yerr=tune_spread[0,:])
+        ax1.errorbar(x=time, y=tune[1,:], yerr=tune_spread[1,:])
         ax1.set_xlabel("Turn number")
         ax1.set_ylabel("Transverse tunes")
         plt.legend(["x","y"])
+        fig_to_return.append(fig1)
         
         fig2, ax2 = plt.subplots()        
-        ax2.errorbar(x=time[1:], y=tune[2,1:], yerr=tune_spread[2,1:])
+        ax2.errorbar(x=time, y=tune[2,:], yerr=tune_spread[2,:])
         ax2.set_xlabel("Turn number")
         ax2.set_ylabel("Synchrotron tune")
+        fig_to_return.append(fig2)
         
     if plot_fft is True:
-        if ring is None:
-            T0 = 1.183e-06
-            f0 = 0.845e6
-        else:
-            T0 = ring.T0
-            f0 = ring.f0
         
+        T0 = 1/f0
         n_freq = file[group]['fft'].shape[1]
         freq = rfftfreq((n_freq-1)*2, T0)
         tune_fft = freq / f0
@@ -573,12 +570,12 @@ def plot_tunedata(filename, bunch_number, ring=None, plot_tune=True, plot_fft=Fa
         if max_turn is None:
             max_turn = time[-1]
         if min_turn is None:
-            min_turn = time[1]
+            min_turn = time[0]
             
         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
+        min_turn_iloc = min_turn // save_every - 1
         max_turn_iloc = max_turn // save_every
         
     
@@ -595,6 +592,7 @@ def plot_tunedata(filename, bunch_number, ring=None, plot_tune=True, plot_fft=Fa
             ax3.set_ylabel("Turns")
             cbar = fig3.colorbar(c, ax=ax3)
             cbar.set_label("log FFT amplitude") 
+            fig_to_return.append(fig3)
             
         if profile_plot is True:
             fig4, ax4 = plt.subplots()
@@ -604,28 +602,11 @@ def plot_tunedata(filename, bunch_number, ring=None, plot_tune=True, plot_fft=Fa
             ax4.set_xlabel('$Q_{}$'.format(dimension))
             ax4.set_ylabel("FFT amplitude")
             ax4.legend(time[min_turn_iloc:max_turn_iloc+1])
+            fig_to_return.append(fig4)
             
     file.close()
-
    
-    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
+    return tuple(fig_to_return)
 
 def plot_cavitydata(filename, cavity_name, phasor="cavity", 
                     plot_type="bunch", bunch_number=0, turn=None):
@@ -740,4 +721,3 @@ def plot_cavitydata(filename, cavity_name, phasor="cavity",
     
     file.close()
     return fig
-    return fig