diff --git a/tracking/particles.py b/tracking/particles.py index 4e515376fd04bc3801b8036bb4ca700d6e4a50ad..1165ee0ca5024b94cfbcd78d0b73c7852d6404ac 100644 --- a/tracking/particles.py +++ b/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."""