Merge branch 'Empty-bunch'

tracking/monitors/monitors.py

@@ -218,27 +218,32 @@ class Monitor(Element, metaclass=ABCMeta):
         except ValueError:
             pass
         
-    def track_bunch_data(self, object_to_save):
+    def track_bunch_data(self, object_to_save, check_empty=False):
         """
         Track method to use when saving bunch data.
         
         Parameters
         ----------
         object_to_save : Beam or Bunch
+        check_emptiness: bool
+            If True, check if the bunch is empty. If it is, then do nothing.
         """
         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])
+                    bunch = object_to_save[object_to_save.mpi.bunch_num]
                 else:
-                    self.to_buffer(object_to_save[self.bunch_number])
+                    bunch = object_to_save[self.bunch_number]
             elif isinstance(object_to_save, Bunch):
-                self.to_buffer(object_to_save)
+                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)
+                
         self.track_count += 1
-        
             
 class BunchMonitor(Monitor):
     """
@@ -672,7 +677,7 @@ class ProfileMonitor(Monitor):
         ----------
         object_to_save : Bunch or Beam object
         """        
-        self.track_bunch_data(object_to_save)
+        self.track_bunch_data(object_to_save, check_empty=True)
         
 class WakePotentialMonitor(Monitor):
     """
@@ -1004,8 +1009,11 @@ class BunchSpectrumMonitor(Monitor):
             raise TypeError("object_to_save should be a Beam or Bunch object.")
         
         if skip is False:
-            for key, value in self.track_dict.items():
-                self.positions[value, :, self.save_count] = bunch[key][self.index_sample]
+            try:
+                for key, value in self.track_dict.items():
+                    self.positions[value, :, self.save_count] = bunch[key][self.index_sample]
+            except IndexError:
+                self.positions[value, :, self.save_count] = np.nan
             
             self.mean[:, self.save_count] = bunch.mean[self.mean_index]
             

tracking/parallel.py

@@ -166,14 +166,22 @@ class Mpi:
             
             dim = dimensions[i]
             n = n_bin[i]
-            bins, sorted_index, profile, center = bunch.binning(dimension=dim, n_bin=n)
             
-            self.__setattr__(dim + "_center", np.empty((len(beam), len(center)), dtype=np.float64))
+            if len(bunch) != 0:
+                bins, sorted_index, profile, center = bunch.binning(dimension=dim, n_bin=n)
+            else:
+                sorted_index = None
+                profile = np.zeros((n-1,),dtype=np.int64)
+                center = np.zeros((n-1,),dtype=np.float64)
+                if beam.filling_pattern[self.bunch_num] is True:
+                    beam.update_filling_pattern()
+                    beam.update_distance_between_bunches()
+               
+            self.__setattr__(dim + "_center", np.empty((self.size, n-1), dtype=np.float64))
             self.comm.Allgather([center,  self.MPI.DOUBLE], [self.__getattribute__(dim + "_center"), self.MPI.DOUBLE])
             
-            self.__setattr__(dim + "_profile", np.empty((len(beam), len(profile)), dtype=np.int64))
+            self.__setattr__(dim + "_profile", np.empty((self.size, n-1), dtype=np.int64))
             self.comm.Allgather([profile,  self.MPI.INT64_T], [self.__getattribute__(dim + "_profile"), self.MPI.INT64_T])
             
             self.__setattr__(dim + "_sorted_index", sorted_index)
-    
-    
\ No newline at end of file
+                    
\ No newline at end of file

tracking/particles.py

@@ -76,6 +76,8 @@ class Bunch:
         Number of particles in the bunch.
     current : float
         Bunch current in [A].
+    is_empty : bool
+        Return True if the bunch is empty.
     mean : array of shape (6,)
         Mean position of alive particles for each coordinates.
     std : array of shape (6,)
@@ -192,6 +194,11 @@ class Bunch:
     @current.setter
     def current(self, value):
         self.charge_per_mp = value * self.ring.T0 / self.__len__()
+        
+    @property
+    def is_empty(self):
+        """Return True if the bunch is empty."""
+        return ~np.any(self.alive)
     
     @property    
     def mean(self):

tracking/rf.py

@@ -226,10 +226,20 @@ class CavityResonator():
                         sorted_index = beam.mpi.tau_sorted_index
                 else:
                     # no mpi -> get bunch profile for current bunch
-                    (bins, sorted_index, profile, center) = bunch.binning()
-                    bin_length = center[1]-center[0]
-                    charge_per_mp = bunch.charge_per_mp
-                    self.bunch_index = index
+                    if len(bunch) != 0:
+                        (bins, sorted_index, profile, center) = bunch.binning()
+                        bin_length = center[1]-center[0]
+                        charge_per_mp = bunch.charge_per_mp
+                        self.bunch_index = index
+                    else:
+                        # Update filling pattern
+                        beam.update_filling_pattern()
+                        beam.update_distance_between_bunches()
+                        # save beam phasor value
+                        self.beam_phasor_record[index] = self.beam_phasor
+                        # phasor decay to be at t=0 of the next bunch
+                        self.phasor_decay(self.ring.T1, ref_frame="beam")
+                        continue
                 
                 energy_change = bunch["tau"]*0
                 

tracking/wakepotential.py

@@ -65,6 +65,10 @@ class WakePotential(Element):
         Calculate the loss factor and kick factor from the wake potential and 
         compare it to a reference value assuming a Gaussian bunch computed in 
         the frequency domain.
+    check_sampling()
+        Check if the wake function sampling is uniform.
+    reduce_sampling(factor)
+        Reduce wake function samping by an integer factor.
         
     """
     
@@ -74,6 +78,7 @@ class WakePotential(Element):
         self.n_types = len(self.wakefield.wake_components)
         self.ring = ring
         self.n_bin = n_bin
+        self.check_sampling()
             
     def charge_density(self, bunch):
         """
@@ -299,24 +304,25 @@ class WakePotential(Element):
         bunch : Bunch or Beam object.
         
         """
-
-        self.charge_density(bunch)
-        for wake_type in self.types:
-            tau0, Wp = self.get_wakepotential(bunch, wake_type)
-            f = interp1d(tau0 + self.tau_mean, Wp, fill_value = 0, bounds_error = False)
-            Wp_interp = f(bunch["tau"])
-            if wake_type == "Wlong":
-                bunch["delta"] += Wp_interp * bunch.charge / self.ring.E0
-            elif wake_type == "Wxdip":
-                bunch["xp"] += Wp_interp * bunch.charge / self.ring.E0
-            elif wake_type == "Wydip":
-                bunch["yp"] += Wp_interp * bunch.charge / self.ring.E0
-            elif wake_type == "Wxquad":
-                bunch["xp"] += (bunch["x"] * Wp_interp * bunch.charge 
-                                / self.ring.E0)
-            elif wake_type == "Wyquad":
-                bunch["yp"] += (bunch["y"] * Wp_interp * bunch.charge 
-                                / self.ring.E0)
+        
+        if len(bunch) != 0:
+            self.charge_density(bunch)
+            for wake_type in self.types:
+                tau0, Wp = self.get_wakepotential(bunch, wake_type)
+                f = interp1d(tau0 + self.tau_mean, Wp, fill_value = 0, bounds_error = False)
+                Wp_interp = f(bunch["tau"])
+                if wake_type == "Wlong":
+                    bunch["delta"] += Wp_interp * bunch.charge / self.ring.E0
+                elif wake_type == "Wxdip":
+                    bunch["xp"] += Wp_interp * bunch.charge / self.ring.E0
+                elif wake_type == "Wydip":
+                    bunch["yp"] += Wp_interp * bunch.charge / self.ring.E0
+                elif wake_type == "Wxquad":
+                    bunch["xp"] += (bunch["x"] * Wp_interp * bunch.charge 
+                                    / self.ring.E0)
+                elif wake_type == "Wyquad":
+                    bunch["yp"] += (bunch["y"] * Wp_interp * bunch.charge 
+                                    / self.ring.E0)
                 
     def plot_last_wake(self, wake_type, plot_rho=True, plot_dipole=False, 
                        plot_wake_function=True):
@@ -526,7 +532,38 @@ class WakePotential(Element):
                                  columns=column, 
                                  index=index)
         return loss_data
+
+    def check_sampling(self):
+        """
+        Check if the wake function sampling is uniform.
+
+        Raises
+        ------
+        ValueError
+
+        """
+        for wake_type in self.types:
+            idx = getattr(self.wakefield, wake_type).data.index
+            diff = idx[1:]-idx[:-1]
+            result = np.all(np.isclose(diff, diff[0], atol=1e-15))
+            if result is False:
+                raise ValueError("The wake function must be uniformly sampled.")
     
+    def reduce_sampling(self, factor):
+        """
+        Reduce wake function samping by an integer factor.
+        
+        Used to reduce computation time for long bunches.
+
+        Parameters
+        ----------
+        factor : int
+
+        """
+        for wake_type in self.types:
+            idx = getattr(self.wakefield, wake_type).data.index[::factor]
+            getattr(self.wakefield, wake_type).data = getattr(self.wakefield, wake_type).data.loc[idx]
+        self.check_sampling()