diff --git a/tracking/monitors/__init__.py b/tracking/monitors/__init__.py
index ffff2a982845826ee31b220f00633f247934d3a6..3f1f9d7b02ae3b87e9dd7566b0360d5983f9b211 100644
--- a/tracking/monitors/__init__.py
+++ b/tracking/monitors/__init__.py
@@ -12,7 +12,8 @@ from mbtrack2.tracking.monitors.monitors import (Monitor, BunchMonitor,
WakePotentialMonitor,
TuneMonitor,
CavityMonitor,
- BunchSpectrumMonitor)
+ BunchSpectrumMonitor,
+ BeamSpectrumMonitor)
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 5b70c9bb35b038cc64607da042955c19062e7f05..8c2271c72db3d12c6eaf71803ee916e10badaef2 100644
--- a/tracking/monitors/monitors.py
+++ b/tracking/monitors/monitors.py
@@ -1089,6 +1089,15 @@ class BunchSpectrumMonitor(Monitor):
allow several cores to write in the same file at the same time.
If False, open the HDF5 file in standard mode.
+ Attributes
+ ----------
+ fft_resolution : float
+ Return the fft resolution in [Hz].
+ signal_resolution : float
+ Return the signal resolution in [Hz].
+ frequency_samples : array of float
+ Return the fft frequency samples in [Hz].
+
Methods
-------
track(bunch):
@@ -1146,9 +1155,37 @@ class BunchSpectrumMonitor(Monitor):
self.incoherent = np.zeros((3, self.n_fft//2+1, buffer_size))
self.coherent = np.zeros((3, self.n_fft//2+1, buffer_size))
+ self.file[self.group_name].create_dataset(
+ "freq", data=self.frequency_samples)
+
+ @property
+ def fft_resolution(self):
+ """
+ Return the fft resolution in [Hz].
+
+ It is defined as the sampling frequency over the number of samples.
+ """
+ return self.ring.f0/self.n_fft
+
+ @property
+ def signal_resolution(self):
+ """
+ Return the signal resolution in [Hz].
+
+ It is defined as the inverse of the signal length.
+ """
+ return 1/(self.ring.T0*self.save_every)
+
+ @property
+ def frequency_samples(self):
+ """
+ Return the fft frequency samples in [Hz].
+ """
+ return rfftfreq(self.n_fft, self.ring.T0)
+
def track(self, object_to_save):
"""
- Save positions and mean data.
+ Save spectrum data.
Parameters
----------
@@ -1227,7 +1264,7 @@ class BunchSpectrumMonitor(Monitor):
Returns
-------
incoherent : array
- The average of the transformed input in each plane.
+ Bunch incoherent spectrum.
"""
fourier = rfft(positions, n=self.n_fft)
@@ -1243,7 +1280,196 @@ class BunchSpectrumMonitor(Monitor):
Returns
-------
coherent : array
- The average of the transformed input in each plane.
+ Bunch coherent spectrum.
+
+ """
+ coherent = np.abs(rfft(mean, n=self.n_fft))
+
+ return coherent
+
+class BeamSpectrumMonitor(Monitor):
+ """
+ Monitor coherent beam spectrum.
+
+ Parameters
+ ----------
+ ring : Synchrotron object
+ 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 spectrum is computed saved 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.
+
+ Attributes
+ ----------
+ fft_resolution : float
+ Return the fft resolution in [Hz].
+ signal_resolution : float
+ Return the signal resolution in [Hz].
+ frequency_samples : array of float
+ Return the fft frequency samples in [Hz].
+
+ Methods
+ -------
+ track(bunch):
+ Save tune and/or FFT data.
+
+ """
+
+ def __init__(self, ring, 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.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
+ group_name = "BeamSpectrum"
+
+ dict_buffer = {"coherent":(3, n_fft//2+1, buffer_size)}
+ dict_file = {"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.mean = np.zeros((self.size_list, ring.h, save_every))
+ self.coherent = np.zeros((3, self.n_fft//2+1, buffer_size))
+
+ self.file[self.group_name].create_dataset(
+ "freq", data=self.frequency_samples)
+
+ @property
+ def fft_resolution(self):
+ """
+ Return the fft resolution in [Hz].
+
+ It is defined as the sampling frequency over the number of samples.
+ """
+ return self.ring.f1/self.n_fft
+
+ @property
+ def signal_resolution(self):
+ """
+ Return the signal resolution in [Hz].
+
+ It is defined as the inverse of the signal length.
+ """
+ return 1/(self.ring.T0*self.save_every)
+
+ @property
+ def frequency_samples(self):
+ """
+ Return the fft frequency samples in [Hz].
+ """
+ return rfftfreq(self.n_fft, self.ring.T1)
+
+ def track(self, beam):
+ """
+ Save mean data.
+
+ Parameters
+ ----------
+ beam : Beam object
+
+ """
+ if (beam.mpi_switch == True):
+ bunch_num = beam.mpi.bunch_num
+ bunch = beam[bunch_num]
+ self.mean[:, bunch_num, self.buffer_count] = bunch.mean[self.mean_index]
+ else:
+ self.mean[:, :, self.buffer_count] = beam.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(beam)
+ self.save_count = 0
+
+ self.track_count += 1
+
+ def to_buffer(self, beam):
+ """
+ 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():
+ if (beam.mpi_switch == True):
+ data_core = self.mean[value, beam.mpi.bunch_num, :]
+ full_data = beam.mpi.comm.allgather(data_core)
+ data = np.reshape(full_data, (-1), 'F')
+ else:
+ data = np.reshape(self.mean[value, :, :], (-1), 'F')
+ self.coherent[self.store_dict[key],:,self.buffer_count] = self.get_beam_spectrum(data)
+ 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]["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_beam_spectrum(self, mean):
+ """
+ Compute the beam coherent spectrum i.e. the absolute value of the FT
+ of the mean position of every bunch.
+
+ Returns
+ -------
+ coherent : array
+ The beam coherent spectrum.
"""
coherent = np.abs(rfft(mean, n=self.n_fft))