Add plot_tunedata and changes to TuneMonitor

Change track method to allow for Beam objects
If pynaff can not find the tune, the tune value is nan.
Nanmean and nanstd are used to avoid to bias the tune value.
! Untested if some particles are lost !

tracking/monitors/monitors.py

@@ -11,11 +11,11 @@ during tracking.
 import numpy as np
 import h5py as hp
 import PyNAFF as pnf
+import random
 from mbtrack2.tracking.element import Element
 from mbtrack2.tracking.particles import Bunch, Beam
 from abc import ABCMeta
-import random
-# from mpi4py import MPI
+from mpi4py import MPI
 
 class Monitor(Element, metaclass=ABCMeta):
     """
@@ -798,8 +798,8 @@ class TuneMonitor(Monitor):
         self.bunch_number = bunch_number
         group_name = "TuneData_" + str(self.bunch_number)
 
-        dict_buffer = {"tune":(3, buffer_size), "sigma_tune":(3, buffer_size,)}
-        dict_file = {"tune":(3, total_size), "sigma_tune":(3, total_size,)}
+        dict_buffer = {"tune":(3, buffer_size), "tune_spread":(3, buffer_size,)}
+        dict_file = {"tune":(3, total_size), "tune_spread":(3, total_size,)}
         
         self.monitor_init(group_name, save_every, buffer_size, total_size,
                           dict_buffer, dict_file, file_name, mpi_mode)
@@ -820,15 +820,30 @@ class TuneMonitor(Monitor):
         
         self.buffer_count = 1
         
-    def track(self, bunch):
+    def track(self, object_to_save):
         """
         Save tune data.
 
         Parameters
         ----------
-        bunch : Bunch object
+        object_to_save : Beam or Bunch object
 
         """
+        skip = False
+        if isinstance(object_to_save, Beam):
+            if (object_to_save.mpi_switch == True):
+                if object_to_save.mpi.bunch_num == self.bunch_number:
+                    bunch = object_to_save[object_to_save.mpi.bunch_num]
+                else:
+                    skip = True
+            else:
+                bunch = object_to_save[self.bunch_number]
+        elif isinstance(object_to_save, Bunch):
+            bunch = object_to_save
+        else:
+            raise TypeError("object_to_save should be a Beam or Bunch object.")
+        
+        if skip is not True:
             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]
@@ -837,20 +852,20 @@ class TuneMonitor(Monitor):
             
             if self.track_count > 0:
                 if self.track_count % self.save_every == 0:
-                self.get_tune(bunch)
-                self.to_buffer()
+                    self.to_buffer(bunch)
                     self.save_count = 0
     
             self.track_count += 1          
         
-    def to_buffer(self):
+    def to_buffer(self, bunch):
         """
         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] = self.mean_tune
-        self.sigma_tune[:, self.buffer_count] = self.tune_spread
+        self.tune[:, self.buffer_count] = mean
+        self.tune_spread[:, self.buffer_count] = spread
         
         self.buffer_count += 1
         
@@ -869,9 +884,9 @@ class TuneMonitor(Monitor):
         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]["sigma_tune"][:, 
+        self.file[self.group_name]["tune_spread"][:, 
                 self.write_count * self.buffer_size:(self.write_count+1) * 
-                self.buffer_size] = self.sigma_tune
+                self.buffer_size] = self.tune_spread
             
         self.file.flush()
         self.write_count += 1
@@ -895,41 +910,23 @@ class TuneMonitor(Monitor):
                 freq[i,0] = pnf.naff(self.x[i,:], turns=turn-1, nterms=1)[0][1] \
                                 / self.ring.T0
             except IndexError:
-                freq[i,0] = 0
+                freq[i,0] = np.nan
                
             try:
                 freq[i,1] = pnf.naff(self.y[i,:], turns=turn-1, nterms=1)[0][1] \
                                 / self.ring.T0
             except IndexError:
-                freq[i,1] = 0
+                freq[i,1] = np.nan
 
             try:              
                 freq[i,2] = pnf.naff(self.tau[i,:], turns=turn-1, nterms=1)[0][1] \
                                 / self.ring.T0
             except IndexError:
-                freq[i,2] = 0
-        
-        tune_single_particle = np.zeros((self.sample_size,3))
-        tune_single_particle[:,0] = freq[:,0] / self.ring.f0
-        tune_single_particle[:,1] = freq[:,1] / self.ring.f0 
-        tune_single_particle[:,2] = freq[:,2] / self.ring.f0 
-        
-        self.mean_tune = np.zeros((3,))
-        self.tune_spread = np.zeros((3,))
-        self.mean_tune[0] = tune_single_particle[:,0].mean() 
-        self.mean_tune[1] = tune_single_particle[:,1].mean() 
-        self.mean_tune[2] = tune_single_particle[:,2].mean() 
-        self.tune_spread[0] = tune_single_particle[:,0].std() 
-        self.tune_spread[1] = tune_single_particle[:,1].std() 
-        self.tune_spread[2] = tune_single_particle[:,2].std()
-        
-        
-        
-        
-        
-        
-        
-        
+                freq[i,2] = np.nan
         
+        tune_single_particle = freq / self.ring.f0
+        mean = np.nanmean(tune_single_particle, 0)
+        spread = np.nanstd(tune_single_particle, 0)
         
+        return (mean, spread)
     
\ No newline at end of file

tracking/monitors/plotting.py

@@ -482,3 +482,44 @@ def plot_wakedata(filename, bunch_number, wake_type="Wlong", start=0,
     elif streak_plot is True:
         return fig2
     
+def plot_tunedata(filename, bunch_number):
+    """
+    Plot data recorded by TuneMonitor.
+    
+    Parameters
+    ----------
+    filename : str 
+        Name of the HDF5 file that contains the data.
+    bunch_number : int
+        Bunch to plot. This has to be identical to 'bunch_number' parameter in 
+        'BunchMonitor' object.
+        
+    Return
+    ------
+    fig : Figure
+        Figure object with the plot on it.
+
+    """
+    
+    file = hp.File(filename, "r")
+    
+    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")
+    
+    file.close()
+    return (fig1, fig2)
+    
\ No newline at end of file