From 6fe4d1a0cfe13869c44fab58dfd541174ccc3936 Mon Sep 17 00:00:00 2001
From: Gamelin Alexis <alexis.gamelin@synchrotron-soleil.fr>
Date: Wed, 17 Nov 2021 18:33:18 +0100
Subject: [PATCH] Allow for tracking with empty bunches

Check that the bunch is not empty before binning in Mpi.share_distributions. If it is then the beam filling pattern is updated and empty array are sent.
In CavityResonator.track, the tracking without mpi is modified to take into account the case with an empty bunch. Mpi case is fine as the beam filling pattern is updated in the call for Mpi.share_distributions.
In WakePotential.track, if the bunch is empty tracking does nothing.
---
 tracking/parallel.py      | 18 +++++++++++++-----
 tracking/rf.py            | 18 ++++++++++++++----
 tracking/wakepotential.py | 37 +++++++++++++++++++------------------
 3 files changed, 46 insertions(+), 27 deletions(-)

diff --git a/tracking/parallel.py b/tracking/parallel.py
index 5145c50..c35460f 100644
--- a/tracking/parallel.py
+++ b/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
diff --git a/tracking/rf.py b/tracking/rf.py
index 8274964..f5fc314 100644
--- a/tracking/rf.py
+++ b/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
                 
diff --git a/tracking/wakepotential.py b/tracking/wakepotential.py
index 71aec2b..b22b173 100644
--- a/tracking/wakepotential.py
+++ b/tracking/wakepotential.py
@@ -299,24 +299,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):
-- 
GitLab