TuneMonitor added

tracking/monitors/__init__.py

@@ -9,7 +9,8 @@ from mbtrack2.tracking.monitors.monitors import (Monitor, BunchMonitor,
                                                  PhaseSpaceMonitor,
                                                  BeamMonitor,
                                                  ProfileMonitor,
-                                                 WakePotentialMonitor)
+                                                 WakePotentialMonitor,
+                                                 TuneMonitor)
 from mbtrack2.tracking.monitors.plotting import (plot_bunchdata, 
                                                  plot_phasespacedata,
                                                  plot_profiledata,

tracking/monitors/monitors.py

@@ -10,10 +10,12 @@ during tracking.
 
 import numpy as np
 import h5py as hp
+import PyNAFF as pnf
 from mbtrack2.tracking.element import Element
 from mbtrack2.tracking.particles import Bunch, Beam
 from abc import ABCMeta
-from mpi4py import MPI
+import random
+# from mpi4py import MPI
 
 class Monitor(Element, metaclass=ABCMeta):
     """
@@ -749,4 +751,185 @@ class WakePotentialMonitor(Monitor):
             self.to_buffer(wake_potential_to_save)
         self.track_count += 1
 
-            
\ No newline at end of file
+class TuneMonitor(Monitor):
+    """
+    Monitor tunes in horizontal, vertical, and longitudinal plane. 
+    
+    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 computation. This number
+        needs not to exceed mp_number.
+    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 tune 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    
+    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 tune data.
+    
+    """
+    
+    def __init__(self, ring, bunch_number, mp_number, sample_size, file_name=None, 
+                 save_every=10, buffer_size=5, total_size=10, mpi_mode=True):
+        
+        self.ring = ring
+        self.bunch_number = bunch_number
+        group_name = "TuneData_" + str(self.bunch_number)
+
+        dict_buffer = {"tune":(3, buffer_size), "sigma_tune":(3, buffer_size,)}
+        dict_file = {"tune":(3, total_size), "sigma_tune":(3, 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.sample_size = int(sample_size)
+        self.x = np.zeros((self.sample_size, save_every+1))
+        self.y = np.zeros((self.sample_size, save_every+1))
+        self.tau = np.zeros((self.sample_size, save_every+1))
+        
+        index = np.arange(0, int(mp_number))
+        self.index_sample = sorted(random.sample(list(index), self.sample_size))
+        
+        self.save_every = save_every
+        self.save_count = 0
+        
+        self.buffer_count = 1
+        
+    def track(self, bunch):
+        """
+        Save tune data.
+
+        Parameters
+        ----------
+        bunch : Bunch object
+
+        """
+        self.x[:, self.save_count] = bunch["x"][self.index_sample]
+        self.y[:, self.save_count] = bunch["y"][self.index_sample]
+        self.tau[:, self.save_count] = bunch["tau"][self.index_sample]
+        
+        self.save_count += 1
+        
+        if self.track_count > 0:
+            if self.track_count % self.save_every == 0:
+                self.get_tune(bunch)
+                self.to_buffer()
+                self.save_count = 0
+
+        self.track_count += 1              
+        
+    def to_buffer(self):
+        """
+        A method to hold saved data before writing it to the output file.
+
+        """
+        self.time[self.buffer_count] = self.track_count
+        self.tune[:, self.buffer_count] = self.mean_tune
+        self.sigma_tune[:, self.buffer_count] = self.tune_spread
+        
+        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]["tune"][:, 
+                self.write_count * self.buffer_size:(self.write_count+1) * 
+                self.buffer_size] = self.tune
+        self.file[self.group_name]["sigma_tune"][:, 
+                self.write_count * self.buffer_size:(self.write_count+1) * 
+                self.buffer_size] = self.sigma_tune
+            
+        self.file.flush()
+        self.write_count += 1
+
+    def get_tune(self, bunch):
+        """
+        Compute tune by using NAFF algorithm to indentify the fundamental
+        harmonic frequency of the particles' motion.
+
+        Parameters
+        ----------
+        bunch : Bunch object
+
+        """
+        
+        turn = self.save_every
+        freq = np.zeros((self.sample_size,3))
+        
+        for i in range(self.sample_size):
+            try:
+                freq[i,0] = pnf.naff(self.x[i,:], turns=turn-1, nterms=1)[0][1] \
+                                / self.ring.T0
+            except IndexError:
+                freq[i,0] = 0
+               
+            try:
+                freq[i,1] = pnf.naff(self.y[i,:], turns=turn-1, nterms=1)[0][1] \
+                                / self.ring.T0
+            except IndexError:
+                freq[i,1] = 0
+
+            try:              
+                freq[i,2] = pnf.naff(self.tau[i,:], turns=turn-1, nterms=1)[0][1] \
+                                / self.ring.T0
+            except IndexError:
+                freq[i,2] = 0
+        
+        tune_single_particle = np.zeros((self.sample_size,3))
+        tune_single_particle[:,0] = freq[:,0] / self.ring.f0
+        tune_single_particle[:,1] = freq[:,1] / self.ring.f0 
+        tune_single_particle[:,2] = freq[:,2] / self.ring.f0 
+        
+        self.mean_tune = np.zeros((3,))
+        self.tune_spread = np.zeros((3,))
+        self.mean_tune[0] = tune_single_particle[:,0].mean() 
+        self.mean_tune[1] = tune_single_particle[:,1].mean() 
+        self.mean_tune[2] = tune_single_particle[:,2].mean() 
+        self.tune_spread[0] = tune_single_particle[:,0].std() 
+        self.tune_spread[1] = tune_single_particle[:,1].std() 
+        self.tune_spread[2] = tune_single_particle[:,2].std()
+        
+        
+        
+        
+        
+        
+        
+        
+        
+        
+        
\ No newline at end of file