Save mean_incoherent and std_incoherent in BunchSpectrumMonitor

tracking/monitors/monitors.py

@@ -1105,10 +1105,9 @@ class BunchSpectrumMonitor(Monitor):
     
     """
     
-    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"):
+    def __init__(self, ring, bunch_number, mp_number, sample_size, n_fft, 
+                 save_every, buffer_size, total_size, file_name=None, 
+                 mpi_mode=True, dim="all"):
         
         self.n_fft = int(n_fft)
         self.sample_size = int(sample_size)
@@ -1126,6 +1125,11 @@ class BunchSpectrumMonitor(Monitor):
         elif dim == "y":
             self.track_dict = {"y":0}
             self.mean_index = [False, False, True, False, False, False]
+        elif dim == "xy" or dim == "yx":
+            self.track_dict = {"x":0,"y":1}
+            self.mean_index = [True, False, True, False, False, False]
+        else:
+            raise ValueError("dim is not correct.")
         
         self.size_list = len(self.track_dict)
         
@@ -1134,9 +1138,13 @@ class BunchSpectrumMonitor(Monitor):
         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)}
+                       "coherent":(3, n_fft//2+1, buffer_size),
+                       "mean_incoherent":(3,buffer_size),
+                       "std_incoherent":(3,buffer_size)}
         dict_file = {"incoherent":(3, n_fft//2+1, total_size),
-                       "coherent":(3, n_fft//2+1, total_size)}
+                        "coherent":(3, n_fft//2+1, total_size),
+                        "mean_incoherent":(3,total_size),
+                        "std_incoherent":(3,total_size)}
         
         self.monitor_init(group_name, save_every, buffer_size, total_size,
                           dict_buffer, dict_file, file_name, mpi_mode)
@@ -1229,7 +1237,10 @@ class BunchSpectrumMonitor(Monitor):
         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,:,:])
+            incoherent, mean_incoherent, std_incoherent = self.get_incoherent_spectrum(self.positions[value,:,:])
+            self.incoherent[self.store_dict[key],:,self.buffer_count] = incoherent
+            self.mean_incoherent[self.store_dict[key],self.buffer_count] = mean_incoherent
+            self.std_incoherent[self.store_dict[key],self.buffer_count] = std_incoherent
             self.coherent[self.store_dict[key],:,self.buffer_count] = self.get_coherent_spectrum(self.mean[value])
         
         self.buffer_count += 1
@@ -1249,6 +1260,12 @@ class BunchSpectrumMonitor(Monitor):
         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]["mean_incoherent"][:, 
+                self.write_count * self.buffer_size:(self.write_count+1) * 
+                self.buffer_size] = self.mean_incoherent
+        self.file[self.group_name]["std_incoherent"][:, 
+                self.write_count * self.buffer_size:(self.write_count+1) * 
+                self.buffer_size] = self.std_incoherent
         self.file[self.group_name]["coherent"][:,:, 
                 self.write_count * self.buffer_size:(self.write_count+1) * 
                 self.buffer_size] = self.coherent
@@ -1265,12 +1282,23 @@ class BunchSpectrumMonitor(Monitor):
         -------
         incoherent : array
             Bunch incoherent spectrum.
+        mean_incoherent : float
+            Mean frequency of the maximum of each individual particle spectrum
+            in [Hz].
+        std_incoherent : float
+            Standard deviation of the frequency of the maximum of each 
+            individual particle spectrum in [Hz].
 
         """
         fourier = rfft(positions, n=self.n_fft)
-        incoherent = np.mean(np.abs(fourier), axis=0)
-        
-        return incoherent
+        fourier_abs = np.abs(fourier)
+        max_array = np.argmax(fourier_abs,axis=1)
+        freq_array = self.frequency_samples[max_array]
+        mean_incoherent = np.mean(freq_array)
+        std_incoherent = np.std(freq_array)
+        incoherent = np.mean(fourier_abs, axis=0)
+        
+        return incoherent, mean_incoherent, std_incoherent
     
     def get_coherent_spectrum(self, mean):
         """