diff --git a/tracking/monitors/__init__.py b/tracking/monitors/__init__.py
index 2f94fe639ee6b2337461aa8155f57816b4c4e70b..1d186dae35b7e1759984464db4dacf838da6fc71 100644
--- a/tracking/monitors/__init__.py
+++ b/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,
diff --git a/tracking/monitors/monitors.py b/tracking/monitors/monitors.py
index 0382055537a4508fcd02f6ad24e30184c7224aa7..c8674387184df2827c9bea3da85222d76e4beedd 100644
--- a/tracking/monitors/monitors.py
+++ b/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