Add a mpi mode to Beam.init_beam

This avoid to store a huge amount of data in all cores before calling mpi_init to share it to all the cores.

tracking/particles.py

@@ -524,7 +524,7 @@ class Beam:
         self._distance_between_bunches =  distance
         
     def init_beam(self, filling_pattern, current_per_bunch=1e-3, 
-                  mp_per_bunch=1e3, track_alive=True):
+                  mp_per_bunch=1e3, track_alive=True, mpi=False):
         """
         Initialize beam with a given filling pattern and marco-particle number 
         per bunch. Then initialize the different bunches with a 6D gaussian
@@ -548,12 +548,19 @@ class Beam:
             If False, the code no longer take into account alive/dead particles.
             Should be set to True if element such as apertures are used.
             Can be set to False to gain a speed increase.
+        mpi : bool, optional
+            If True, only a single bunch is fully initialized on each core, the
+            other bunches are initialized with a single marco-particle.
         """
         
         if (len(filling_pattern) != self.ring.h):
             raise ValueError(("The length of filling pattern is {} ".format(len(filling_pattern)) + 
                               "but should be {}".format(self.ring.h)))
         
+        if mpi is True:
+            mp_per_bunch_mpi = mp_per_bunch
+            mp_per_bunch = 1
+        
         filling_pattern = np.array(filling_pattern)
         bunch_list = []
         if filling_pattern.dtype == np.dtype("bool"):
@@ -577,8 +584,15 @@ class Beam:
         self.update_filling_pattern()
         self.update_distance_between_bunches()
         
-        for bunch in self.not_empty:
+        if mpi is True:
+            self.mpi_init()
+            current = self[self.mpi.rank_to_bunch(self.mpi.rank)].current
+            bunch =  Bunch(self.ring, mp_per_bunch_mpi, current, track_alive)
             bunch.init_gaussian()
+            self[self.mpi.rank_to_bunch(self.mpi.rank)] = bunch
+        else:
+            for bunch in self.not_empty:
+                bunch.init_gaussian()
     
     def update_filling_pattern(self):
         """Update the beam filling pattern."""