Add ProfileMonitor and misc on monitors.py

Add ProfileMonitor to save bunch profile
Add a track_bunch_data to Monitor class to group the code of track methods from BunchMonitor and PhaseSpaceMonitor
Implement BeamMonitor for the case without mpi
Remove hard coded h=416 ! in BeamMonitor

tracking/monitors/monitors.py

@@ -46,6 +46,8 @@ class Monitor(Element, metaclass=ABCMeta):
         Close the HDF5 file shared by all Monitor subclass, must be called 
         by at least an instance of a Montior subclass at the end of the 
         tracking.
+    track_bunch_data(object_to_save)
+        Track method to use when saving bunch data.
     """
     
     _file_name_storage = []
@@ -201,6 +203,28 @@ class Monitor(Element, metaclass=ABCMeta):
         except ValueError:
             pass
         
+    def track_bunch_data(self, object_to_save):
+        """
+        Track method to use when saving bunch data.
+        
+        Parameters
+        ----------
+        object_to_save : Beam or Bunch
+        """
+        if self.track_count % self.save_every == 0:
+            if isinstance(object_to_save, Beam):
+                if (object_to_save.mpi_switch == True):
+                    if object_to_save.mpi.bunch_num == self.bunch_number:
+                        self.to_buffer(object_to_save[object_to_save.mpi.bunch_num])
+                else:
+                    self.to_buffer(object_to_save[self.bunch_number])
+            elif isinstance(object_to_save, Bunch):
+                self.to_buffer(object_to_save)
+            else:                            
+                raise TypeError("object_to_save should be a Beam or Bunch object.")
+        self.track_count += 1
+        
+            
 class BunchMonitor(Monitor):
     """
     Monitor a single bunch and save attributes (mean, std, emit and current).
@@ -257,18 +281,7 @@ class BunchMonitor(Monitor):
         ----------
         object_to_save : Bunch or Beam object
         """        
-        if self.track_count % self.save_every == 0:
-            if isinstance(object_to_save, Beam):
-                if (object_to_save.mpi_switch == True):
-                    if object_to_save.mpi.bunch_num == self.bunch_number:
-                        self.to_buffer(object_to_save[object_to_save.mpi.bunch_num])
-                else:
-                    self.to_buffer(object_to_save[self.bunch_number])
-            elif isinstance(object_to_save, Bunch):
-                self.to_buffer(object_to_save)
-            else:                            
-                raise TypeError("object_to_save should be a Beam or Bunch object.")
-        self.track_count += 1
+        self.track_bunch_data(object_to_save)
 
             
 class PhaseSpaceMonitor(Monitor):
@@ -330,19 +343,7 @@ class PhaseSpaceMonitor(Monitor):
         ----------
         object_to_save : Bunch or Beam object
         """        
-        
-        if self.track_count % self.save_every == 0:
-            if isinstance(object_to_save, Beam):
-                if (object_to_save.mpi_switch == True):
-                    if object_to_save.mpi.bunch_num == self.bunch_number:
-                        self.to_buffer(object_to_save[object_to_save.mpi.bunch_num])
-                else:
-                    self.to_buffer(object_to_save[self.bunch_number])
-            elif isinstance(object_to_save, Bunch):
-                self.to_buffer(object_to_save)
-            else:                            
-                raise TypeError("object_to_save should be a Beam or Bunch object.")
-        self.track_count += 1
+        self.track_bunch_data(object_to_save)
 
             
 class BeamMonitor(Monitor):
@@ -377,25 +378,22 @@ class BeamMonitor(Monitor):
         Save data    
     """
     
-    def __init__(self, file_name=None, save_every=5, buffer_size=500, 
+    def __init__(self, h, file_name=None, save_every=5, buffer_size=500, 
                  total_size=2e4, mpi_mode=True):
         
         group_name = "Beam"
-        dict_buffer = {}
-        dict_file = {}
+        dict_buffer = {"mean" : (6, h, buffer_size), 
+                       "std" : (6, h, buffer_size),
+                       "emit" : (3, h, buffer_size),
+                       "current" : (h, buffer_size)}
+        dict_file = {"mean" : (6, h, total_size), 
+                       "std" : (6, h, total_size),
+                       "emit" : (3, h, total_size),
+                       "current" : (h, total_size)}
+        
         self.monitor_init(group_name, save_every, buffer_size, total_size,
                           dict_buffer, dict_file, file_name, mpi_mode)
                     
-        self.mean = np.zeros((6, self.buffer_size), dtype=float)
-        self.std = np.zeros((6, self.buffer_size), dtype=float)
-        self.emit = np.zeros((3, self.buffer_size), dtype=float)
-        self.current = np.zeros((self.buffer_size,), dtype=float)
-        
-        self.g.require_dataset("mean", (6, 416, self.total_size,), dtype=float)
-        self.g.require_dataset("std", (6, 416, self.total_size,), dtype=float)
-        self.g.require_dataset("emit", (3, 416, self.total_size,), dtype=float)
-        self.g.require_dataset("current", (416, self.total_size,), dtype=float)
-                    
     def track(self, beam):
         """
         Save data
@@ -408,13 +406,13 @@ class BeamMonitor(Monitor):
             if (beam.mpi_switch == True):
                 self.to_buffer(beam[beam.mpi.bunch_num], beam.mpi.bunch_num)
             else:
-                raise NotImplementedError
+                self.to_buffer_no_mpi(beam)
                     
         self.track_count += 1
         
     def to_buffer(self, bunch, bunch_num):
         """
-        Save data to buffer.
+        Save data to buffer, if mpi is being used.
         
         Parameters
         ----------
@@ -423,10 +421,10 @@ class BeamMonitor(Monitor):
         """
         
         self.time[self.buffer_count] = self.track_count
-        self.mean[:, self.buffer_count] = bunch.mean
-        self.std[:, self.buffer_count] = bunch.std
-        self.emit[:, self.buffer_count] = bunch.emit
-        self.current[self.buffer_count] = bunch.current
+        self.mean[:, bunch_num, self.buffer_count] = bunch.mean
+        self.std[:, bunch_num, self.buffer_count] = bunch.std
+        self.emit[:, bunch_num, self.buffer_count] = bunch.emit
+        self.current[bunch_num, self.buffer_count] = bunch.current
         
         self.buffer_count += 1
         
@@ -434,9 +432,30 @@ class BeamMonitor(Monitor):
             self.write(bunch_num)
             self.buffer_count = 0
             
+    def to_buffer_no_mpi(self, beam):
+        """
+        Save data to buffer, if mpi is not being used.
+        
+        Parameters
+        ----------
+        beam : Beam object
+        """
+              
+        self.time[self.buffer_count] = self.track_count
+        self.mean[:, :, self.buffer_count] = beam.bunch_mean
+        self.std[:, :, self.buffer_count] = beam.bunch_std
+        self.emit[:, :, self.buffer_count] = beam.bunch_emit
+        self.current[:, self.buffer_count] = beam.bunch_current
+        
+        self.buffer_count += 1
+        
+        if self.buffer_count == self.buffer_size:
+            self.write_no_mpi()
+            self.buffer_count = 0
+
     def write(self, bunch_num):
         """
-        Write data from buffer to the HDF5 file.
+        Write data from buffer to the HDF5 file, if mpi is being used.
         
         Parameters
         ----------
@@ -447,17 +466,44 @@ class BeamMonitor(Monitor):
     
         self.file[self.group_name]["mean"][:, bunch_num, 
                  self.write_count*self.buffer_size:(self.write_count+1) * 
-                 self.buffer_size] = self.mean
+                 self.buffer_size] = self.mean[:, bunch_num, :]
                  
         self.file[self.group_name]["std"][:, bunch_num, 
                  self.write_count*self.buffer_size:(self.write_count+1) * 
-                 self.buffer_size] = self.std
+                 self.buffer_size] = self.std[:, bunch_num, :]
 
         self.file[self.group_name]["emit"][:, bunch_num, 
                  self.write_count*self.buffer_size:(self.write_count+1) * 
-                 self.buffer_size] = self.emit
+                 self.buffer_size] = self.emit[:, bunch_num, :]
 
         self.file[self.group_name]["current"][bunch_num, 
+                 self.write_count*self.buffer_size:(self.write_count+1) * 
+                 self.buffer_size] = self.current[bunch_num, :]
+                 
+        self.file.flush() 
+        self.write_count += 1
+        
+    def write_no_mpi(self):
+        """
+        Write data from buffer to the HDF5 file, if mpi is not being used.
+        """
+        
+        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]["mean"][:, :, 
+                 self.write_count*self.buffer_size:(self.write_count+1) * 
+                 self.buffer_size] = self.mean
+                 
+        self.file[self.group_name]["std"][:, :, 
+                 self.write_count*self.buffer_size:(self.write_count+1) * 
+                 self.buffer_size] = self.std
+
+        self.file[self.group_name]["emit"][:, :, 
+                 self.write_count*self.buffer_size:(self.write_count+1) * 
+                 self.buffer_size] = self.emit
+
+        self.file[self.group_name]["current"][:, 
                  self.write_count*self.buffer_size:(self.write_count+1) * 
                  self.buffer_size] = self.current
                  
@@ -465,3 +511,119 @@ class BeamMonitor(Monitor):
         self.write_count += 1
 
         
+class ProfileMonitor(Monitor):
+    """
+    Monitor a single bunch and save bunch profiles.
+    
+    Parameters
+    ----------
+    bunch_number : int
+        Bunch to monitor.
+    dimensions : str or list of str, optional
+        Dimensions to save.
+    n_bin : int or list of int, optional
+        Number of bin to use in each dimension.
+    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
+        the following times.
+    save_every : int or float, optional
+        Set the frequency of the save. The data is saved every save_every 
+        call of the montior.
+    buffer_size : int or float, optional
+        Size of the save buffer.
+    total_size : int or float, optional
+        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
+    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.
+        If False, open the HDF5 file in standard mode.
+
+    Methods
+    -------
+    track(object_to_save)
+        Save data
+    """
+    
+    def __init__(self, bunch_number, dimensions="tau", n_bin=75, file_name=None, 
+                 save_every=5, buffer_size=500, total_size=2e4, mpi_mode=True):
+        
+        self.bunch_number = bunch_number
+        group_name = "ProfileData_" + str(self.bunch_number)
+        
+        if isinstance(dimensions, str):
+            self.dimensions = [dimensions]
+        else:
+            self.dimensions = dimensions
+            
+        if isinstance(n_bin, int):
+            self.n_bin = np.ones((len(self.dimensions),), dtype=int)*n_bin
+        else:
+            self.n_bin = n_bin
+        
+        dict_buffer = {}
+        dict_file = {}
+        for index, dim in enumerate(dimensions):
+            dict_buffer.update({dim : (self.n_bin[index], buffer_size)})
+            dict_buffer.update({dim + "_bin" : (self.n_bin[index] + 1, buffer_size)})
+            dict_file.update({dim : (self.n_bin[index], total_size)})
+            dict_file.update({dim + "_bin" : (self.n_bin[index] + 1, total_size)})
+
+        self.monitor_init(group_name, save_every, buffer_size, total_size,
+                          dict_buffer, dict_file, file_name, mpi_mode)
+        
+        self.dict_buffer = dict_buffer
+        self.dict_file = dict_file
+        
+    def to_buffer(self, bunch):
+        """
+        Save data to buffer.
+        
+        Parameters
+        ----------
+        bunch : Bunch object
+        """
+
+        self.time[self.buffer_count] = self.track_count
+        for index, dim in enumerate(self.dimensions):
+            bins, sorted_index, profile = bunch.binning(dim, self.n_bin[index])
+            bin_array = np.append(np.array(bins.left), bins.right[-1])
+            profile_array = np.array(profile)
+            self.__getattribute__(dim + "_bin")[:, self.buffer_count] = bin_array
+            self.__getattribute__(dim)[:, self.buffer_count] = profile_array
+        
+        self.buffer_count += 1
+        
+        if self.buffer_count == self.buffer_size:
+            self.write()
+            self.buffer_count = 0
+            
+    def write(self):
+        """Write data from buffer to the HDF5 file."""
+        
+        self.file[self.group_name]["time"][self.write_count*self.buffer_size:(
+                    self.write_count+1)*self.buffer_size] = self.time
+
+        for dim in self.dimensions:
+            self.file[self.group_name][dim][:, 
+                    self.write_count * self.buffer_size:(self.write_count+1) * 
+                    self.buffer_size] = self.__getattribute__(dim)
+            self.file[self.group_name][dim + "_bin"][:, 
+                    self.write_count * self.buffer_size:(self.write_count+1) * 
+                    self.buffer_size] = self.__getattribute__(dim + "_bin")
+            
+        self.file.flush()
+        self.write_count += 1
+                    
+    def track(self, object_to_save):
+        """
+        Save data.
+        
+        Parameters
+        ----------
+        object_to_save : Bunch or Beam object
+        """        
+        self.track_bunch_data(object_to_save)
\ No newline at end of file