[Fix] Monitor behaviour in multi-bunch

Fix Monitor.track_bunch_data method to check that only the bunch which bunch_num is specified can write on the file
Modify the behaviour of BunchMonitor, PhaseSpaceMonitor, ProfileMonitor
Make the same modification within WakePotentialMonitor.
Style change to BunchSpectrumMonitor

mbtrack2/tracking/monitors/monitors.py

@@ -226,20 +226,25 @@ class Monitor(Element, metaclass=ABCMeta):
         check_emptiness: bool
             If True, check if the bunch is empty. If it is, then do nothing.
         """
+        save = True
         if self.track_count % self.save_every == 0:
             
             if isinstance(object_to_save, Beam):
                 if (object_to_save.mpi_switch == True):
-                    bunch = object_to_save[object_to_save.mpi.bunch_num]
+                    if object_to_save.mpi.bunch_num == self.bunch_number:
+                        bunch = object_to_save[object_to_save.mpi.bunch_num]
+                    else:
+                        save = False
                 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 (check_empty == False) or (bunch.is_empty == False):
-                self.to_buffer(bunch)
+            
+            if save:
+                if (check_empty == False) or (bunch.is_empty == False):
+                    self.to_buffer(bunch)
                 
         self.track_count += 1
             
@@ -712,7 +717,7 @@ class WakePotentialMonitor(Monitor):
 
     Methods
     -------
-    track(wake_potential_to_save)
+    track(object_to_save, wake_potential_to_save)
         Save data.
     """
     
@@ -812,15 +817,31 @@ class WakePotentialMonitor(Monitor):
         self.file.flush()
         self.write_count += 1
                     
-    def track(self, wake_potential_to_save):
+    def track(self, object_to_save, wake_potential_to_save):
         """
         Save data.
         
         Parameters
         ----------
-        object_to_save : WakePotential object
-        """        
-        if self.track_count % self.save_every == 0:
+        object_to_save : Bunch or Beam object
+        wake_potential_to_save : WakePotential object
+        """
+        if isinstance(object_to_save, Beam):
+            if (object_to_save.mpi_switch == True):
+                if object_to_save.mpi.bunch_num == self.bunch_number:
+                    save = True
+                else:
+                    save = False
+            else:
+                raise NotImplementedError("WakePotentialMonitor for Beam " +
+                                          "objects is only available " +
+                                          "with MPI mode.")
+        elif isinstance(object_to_save, Bunch):
+            save = True
+        else:                            
+            raise TypeError("object_to_save should be a Beam or Bunch object.")
+            
+        if save and (self.track_count % self.save_every == 0):
             self.to_buffer(wake_potential_to_save)
         self.track_count += 1
     
@@ -992,13 +1013,13 @@ class BunchSpectrumMonitor(Monitor):
         object_to_save : Beam or Bunch object
 
         """
-        skip = False
+        save = True
         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
+                    save = False
             else:
                 bunch = object_to_save[self.bunch_number]
         elif isinstance(object_to_save, Bunch):
@@ -1006,7 +1027,7 @@ class BunchSpectrumMonitor(Monitor):
         else:
             raise TypeError("object_to_save should be a Beam or Bunch object.")
         
-        if skip is False:
+        if save:
             try:
                 for key, value in self.track_dict.items():
                     self.positions[value, :, self.save_count] = bunch[key][self.index_sample]