diff --git a/tracking/monitors/__init__.py b/tracking/monitors/__init__.py
index 3ff31947a254f22aecfdd614e555637ef5b6462a..ffff2a982845826ee31b220f00633f247934d3a6 100644
--- a/tracking/monitors/__init__.py
+++ b/tracking/monitors/__init__.py
@@ -11,7 +11,8 @@ from mbtrack2.tracking.monitors.monitors import (Monitor, BunchMonitor,
ProfileMonitor,
WakePotentialMonitor,
TuneMonitor,
- CavityMonitor)
+ CavityMonitor,
+ BunchSpectrumMonitor)
from mbtrack2.tracking.monitors.plotting import (plot_bunchdata,
plot_phasespacedata,
plot_profiledata,
diff --git a/tracking/monitors/monitors.py b/tracking/monitors/monitors.py
index 0d3b9a4cde683c0c301c793ed69d36bbc70a226c..5b70c9bb35b038cc64607da042955c19062e7f05 100644
--- a/tracking/monitors/monitors.py
+++ b/tracking/monitors/monitors.py
@@ -1052,6 +1052,203 @@ class TuneMonitor(Monitor):
fourier_tau_avg = np.mean(abs(fourier_tau),axis=0)
return (fourier_x_avg, fourier_y_avg, fourier_tau_avg)
+
+class BunchSpectrumMonitor(Monitor):
+ """
+ Monitor the coherent and incoherent bunch spectrums.
+
+ Parameters
+ ----------
+ ring : Synchrotron object
+ bunch_number : int
+ Bunch to monitor
+ mp_number : int or float
+ Total number of macro-particles in the bunch.
+ sample_size : int or float
+ Number of macro-particles to be used for tune and FFT computation.
+ This number cannot exceed mp_number.
+ n_fft : int or float, optional
+ The number of points used for FFT computation, if n_fft is bigger than
+ save_every zero-padding is applied.
+ 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
+ the following times.
+ save_every : int or float, optional
+ Set the frequency of the save. The spectrums are computed every
+ save_every call of the montior.
+ buffer_size : int or float, optional
+ Size of the save buffer.
+ total_size : int or float, optional
+ Total size of the save. The following relationships between the
+ parameters must exist:
+ total_size % buffer_size == 0
+ number of call to track / save_every == total_size - 1
+ mpi_mode : bool, optional
+ If True, open the HDF5 file in parallel mode, which is needed to
+ allow several cores to write in the same file at the same time.
+ If False, open the HDF5 file in standard mode.
+
+ Methods
+ -------
+ track(bunch):
+ Save spectrum data.
+
+ """
+
+ def __init__(self, ring, bunch_number, mp_number, sample_size, n_fft=10000,
+ file_name=None, save_every=10,
+ buffer_size=5, total_size=10, mpi_mode=True,
+ dim="all"):
+
+ self.n_fft = int(n_fft)
+ self.sample_size = int(sample_size)
+ self.store_dict = {"x":0,"y":1,"tau":2}
+
+ if dim == "all":
+ self.track_dict = {"x":0,"y":1,"tau":2}
+ self.mean_index = [True, False, True, False, True, False]
+ elif dim == "tau":
+ self.track_dict = {"tau":0}
+ self.mean_index = [False, False, False, False, True, False]
+ elif dim == "x":
+ self.track_dict = {"x":0}
+ self.mean_index = [True, False, False, False, False, False]
+ elif dim == "y":
+ self.track_dict = {"y":0}
+ self.mean_index = [False, False, True, False, False, False]
+
+ self.size_list = len(self.track_dict)
+
+ self.ring = ring
+ self.bunch_number = bunch_number
+ group_name = "BunchSpectrum_" + str(self.bunch_number)
+
+ dict_buffer = {"incoherent":(3, n_fft//2+1, buffer_size),
+ "coherent":(3, n_fft//2+1, buffer_size)}
+ dict_file = {"incoherent":(3, n_fft//2+1, total_size),
+ "coherent":(3, n_fft//2+1, total_size)}
+
+ self.monitor_init(group_name, save_every, buffer_size, total_size,
+ dict_buffer, dict_file, file_name, mpi_mode)
+
+ self.dict_buffer = dict_buffer
+ self.dict_file = dict_file
+
+ self.save_count = 0
+
+ self.positions = np.zeros((self.size_list, self.sample_size, save_every+1))
+ self.mean = np.zeros((self.size_list, save_every+1))
+
+ index = np.arange(0, int(mp_number))
+ self.index_sample = sorted(random.sample(list(index), self.sample_size))
+
+ self.incoherent = np.zeros((3, self.n_fft//2+1, buffer_size))
+ self.coherent = np.zeros((3, self.n_fft//2+1, buffer_size))
+
+ def track(self, object_to_save):
+ """
+ Save positions and mean data.
+
+ Parameters
+ ----------
+ object_to_save : Beam or Bunch object
+
+ """
+ skip = False
+ if isinstance(object_to_save, Beam):
+ if (object_to_save.mpi_switch == True):
+ if object_to_save.mpi.bunch_num == self.bunch_number:
+ bunch = object_to_save[object_to_save.mpi.bunch_num]
+ else:
+ skip = True
+ else:
+ bunch = object_to_save[self.bunch_number]
+ elif isinstance(object_to_save, Bunch):
+ bunch = object_to_save
+ else:
+ 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]
+
+ self.mean[:, self.save_count] = bunch.mean[self.mean_index]
+
+ self.save_count += 1
+
+ if self.track_count > 0 and self.track_count % self.save_every == 0:
+ self.to_buffer(bunch)
+ self.save_count = 0
+
+ self.track_count += 1
+
+ def to_buffer(self, bunch):
+ """
+ A method to hold saved data before writing it to the output file.
+
+ """
+
+ self.time[self.buffer_count] = self.track_count
+
+ for key, value in self.track_dict.items():
+ self.incoherent[self.store_dict[key],:,self.buffer_count] = self.get_incoherent_spectrum(self.positions[value,:,:])
+ self.coherent[self.store_dict[key],:,self.buffer_count] = self.get_coherent_spectrum(self.mean[value])
+
+ self.buffer_count += 1
+
+ if self.buffer_count == self.buffer_size:
+ self.write()
+ self.buffer_count = 0
+
+ def write(self):
+ """
+ Write data from buffer to output file.
+
+ """
+ self.file[self.group_name]["time"][self.write_count*self.buffer_size:(
+ self.write_count+1)*self.buffer_size] = self.time
+
+ self.file[self.group_name]["incoherent"][:,:,
+ self.write_count * self.buffer_size:(self.write_count+1) *
+ self.buffer_size] = self.incoherent
+ self.file[self.group_name]["coherent"][:,:,
+ self.write_count * self.buffer_size:(self.write_count+1) *
+ self.buffer_size] = self.coherent
+
+ self.file.flush()
+ self.write_count += 1
+
+ def get_incoherent_spectrum(self, positions):
+ """
+ Compute the incoherent spectrum i.e. the average of the absolute value
+ of the FT of the position of every particule of the bunch.
+
+ Returns
+ -------
+ incoherent : array
+ The average of the transformed input in each plane.
+
+ """
+ fourier = rfft(positions, n=self.n_fft)
+ incoherent = np.mean(np.abs(fourier), axis=0)
+
+ return incoherent
+
+ def get_coherent_spectrum(self, mean):
+ """
+ Compute the coherent spectrum i.e. the absolute value of the FT of the
+ mean position of the bunch.
+
+ Returns
+ -------
+ coherent : array
+ The average of the transformed input in each plane.
+
+ """
+ coherent = np.abs(rfft(mean, n=self.n_fft))
+
+ return coherent
class CavityMonitor(Monitor):
"""