diff --git a/tracking/monitors/monitors.py b/tracking/monitors/monitors.py index c82a2f4708e37dcc41e5fd223262b86bc2e75bed..cf139f7e8ed252bd1683077ecdebdb436b0ea3c0 100644 --- a/tracking/monitors/monitors.py +++ b/tracking/monitors/monitors.py @@ -295,7 +295,8 @@ class PhaseSpaceMonitor(Monitor): bunch_number : int Bunch to monitor mp_number : int or float - Number of macroparticle in the phase space to save. + Number of macroparticle in the phase space to save. If less than the + total number of macroparticles, a random fraction of the bunch is saved. file_name : string, optional Name of the HDF5 where the data will be stored. Must be specified the first time a subclass of Monitor is instancied and must be None @@ -346,6 +347,34 @@ class PhaseSpaceMonitor(Monitor): object_to_save : Bunch or Beam object """ self.track_bunch_data(object_to_save) + + def to_buffer(self, bunch): + """ + Save data to buffer. + + Parameters + ---------- + bunch : Bunch object + """ + self.time[self.buffer_count] = self.track_count + + if len(bunch.alive) != self.mp_number: + index = np.arange(len(bunch.alive)) + samples_meta = random.sample(list(index), self.mp_number) + samples = sorted(samples_meta) + else: + samples = slice(None) + + self.alive[:, self.buffer_count] = bunch.alive[samples] + for i, dim in enumerate(bunch): + self.particles[:, i, self.buffer_count] = bunch.particles[dim][samples] + + self.buffer_count += 1 + + if self.buffer_count == self.buffer_size: + self.write() + self.buffer_count = 0 + class BeamMonitor(Monitor):