diff --git a/tracking/monitors/monitors.py b/tracking/monitors/monitors.py
index 1d66539280b71ec7c5c08db334155c1c2c345d0c..1c61966aa0b2730851399723d63c6bdf4f177a4c 100644
--- a/tracking/monitors/monitors.py
+++ b/tracking/monitors/monitors.py
@@ -17,6 +17,7 @@ from mbtrack2.tracking.particles import Bunch, Beam
from mbtrack2.tracking.rf import CavityResonator
from scipy.interpolate import interp1d
from abc import ABCMeta
+from scipy.fft import rfft, rfftfreq
from mpi4py import MPI
class Monitor(Element, metaclass=ABCMeta):
@@ -812,7 +813,8 @@ class WakePotentialMonitor(Monitor):
class TuneMonitor(Monitor):
"""
- Monitor tunes in horizontal, vertical, and longitudinal plane.
+ Monitor tunes and the Fourier transform (using FFT algorithm) of the
+ osciallation in horizontal, vertical, and longitudinal plane.
Parameters
----------
@@ -822,8 +824,14 @@ class TuneMonitor(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.
+ Number of macro-particles to be used for tune and FFT computation.
+ This number cannot exceed mp_number.
+ save_tune : bool, optional
+ If True, tune data is saved.
+ save_fft : bool, optional
+ If True, FFT data is saved.
+ n_fft : int or float, optional
+ The number of points used for FFT computation.
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
@@ -846,19 +854,22 @@ class TuneMonitor(Monitor):
Methods
-------
track(bunch):
- Save tune data.
+ Save tune and/or FFT 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):
+ def __init__(self, ring, bunch_number, mp_number, sample_size, save_tune=True,
+ save_fft=False, n_fft=10000, 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), "tune_spread":(3, buffer_size,)}
- dict_file = {"tune":(3, total_size), "tune_spread":(3, total_size,)}
+
+ dict_buffer = {"tune":(3, buffer_size), "tune_spread":(3, buffer_size,),
+ "fft":(3, n_fft//2+1, buffer_size)}
+ dict_file = {"tune":(3, total_size), "tune_spread":(3, total_size,),
+ "fft":(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)
@@ -874,11 +885,17 @@ class TuneMonitor(Monitor):
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
+ self.save_tune = save_tune
+ self.save_fft = save_fft
+ self.save_every = save_every
+
+ if self.save_fft is True :
+ self.n_fft = n_fft
+ self.fourier_save = np.zeros((3, self.n_fft//2+1, buffer_size))
+
def track(self, object_to_save):
"""
Save tune data.
@@ -909,10 +926,9 @@ class TuneMonitor(Monitor):
self.save_count += 1
- if self.track_count > 0:
- if self.track_count % self.save_every == 0:
- self.to_buffer(bunch)
- self.save_count = 0
+ 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
@@ -921,10 +937,19 @@ class TuneMonitor(Monitor):
A method to hold saved data before writing it to the output file.
"""
- mean, spread = self.get_tune(bunch)
+
self.time[self.buffer_count] = self.track_count
- self.tune[:, self.buffer_count] = mean
- self.tune_spread[:, self.buffer_count] = spread
+
+ if self.save_tune is True:
+ mean, spread = self.get_tune(bunch)
+ self.tune[:, self.buffer_count] = mean
+ self.tune_spread[:, self.buffer_count] = spread
+
+ if self.save_fft is True:
+ fx, fy, ftau = self.get_fft()
+ self.fourier_save[0,:,self.buffer_count] = fx
+ self.fourier_save[1,:,self.buffer_count] = fy
+ self.fourier_save[2,:,self.buffer_count] = ftau
self.buffer_count += 1
@@ -940,12 +965,18 @@ class TuneMonitor(Monitor):
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]["tune_spread"][:,
- self.write_count * self.buffer_size:(self.write_count+1) *
- self.buffer_size] = self.tune_spread
+ if self.save_tune is True:
+ 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]["tune_spread"][:,
+ self.write_count * self.buffer_size:(self.write_count+1) *
+ self.buffer_size] = self.tune_spread
+
+ if self.save_fft is True:
+ self.file[self.group_name]["fft"][:,:,
+ self.write_count * self.buffer_size:(self.write_count+1) *
+ self.buffer_size] = self.fourier_save
self.file.flush()
self.write_count += 1
@@ -989,6 +1020,27 @@ class TuneMonitor(Monitor):
return (mean, spread)
+ def get_fft(self):
+ """
+ Compute the Fourier transform (using FFT algorithm) of the
+ osciallation in horizontal, vertical, and longitudinal plane.
+
+ Returns
+ -------
+ fourier_x_avg, fourier_y_avg, fourier_tau_avg : ndarray
+ The average of the transformed input in each plane.
+
+ """
+ fourier_x = rfft(self.x, n=self.n_fft)
+ fourier_y = rfft(self.y, n=self.n_fft)
+ fourier_tau = rfft(self.tau, n=self.n_fft)
+
+ fourier_x_avg = np.mean(abs(fourier_x),axis=0)
+ fourier_y_avg = np.mean(abs(fourier_y),axis=0)
+ fourier_tau_avg = np.mean(abs(fourier_tau),axis=0)
+
+ return (fourier_x_avg, fourier_y_avg, fourier_tau_avg)
+
class CavityMonitor(Monitor):
"""
Monitor a CavityResonator object and save attributes (mean, std, emit and current).
@@ -1051,5 +1103,4 @@ class CavityMonitor(Monitor):
self.to_buffer(cavity)
else:
raise TypeError("cavity should be a CavityResonator object.")
- self.track_count += 1
-
\ No newline at end of file
+ self.track_count += 1
diff --git a/tracking/monitors/plotting.py b/tracking/monitors/plotting.py
index f2f247c2f832306677d645ad2f21df998d0e8878..066fc8af1d1001c405a0f658c2448d48b5f74370 100644
--- a/tracking/monitors/plotting.py
+++ b/tracking/monitors/plotting.py
@@ -13,6 +13,7 @@ import matplotlib as mpl
import seaborn as sns
import h5py as hp
import random
+from scipy.fft import rfftfreq
def plot_beamdata(filename, dataset, dimension=None, stat_var=None, x_var="time"):
"""
@@ -492,7 +493,9 @@ def plot_wakedata(filename, bunch_number, wake_type="Wlong", start=0,
elif streak_plot is True:
return fig2
-def plot_tunedata(filename, bunch_number):
+def plot_tunedata(filename, bunch_number, ring=None, plot_tune=True, plot_fft=False,
+ dimension='x', min_tune=0, max_tune=0.5, min_turn=None,
+ max_turn=None, streak_plot=True, profile_plot=False):
"""
Plot data recorded by TuneMonitor.
@@ -503,6 +506,24 @@ def plot_tunedata(filename, bunch_number):
bunch_number : int
Bunch to plot. This has to be identical to 'bunch_number' parameter in
'BunchMonitor' object.
+ ring : Synchrotron object, optional
+ The ring configuration that is used in TuneMonitor. If None, the default
+ value of the revolution period and the revolution frequency are used,
+ which are 1.183 us and 0.845 MHz, respectively.
+ plot_tune : bool, optional
+ If True, tune data is plotted.
+ plot_fft : bool, optional
+ If True, FFT data is plotted.
+ dimension : {'x', 'y', 's'}
+ Option to plot FFT data in horizontal, vertical, or longitudinal plane.
+ min_tune, max_tune : int or float, optional
+ The minimum and the maximum tune values to plot FFT data.
+ min_turn, max_turn : int or float, optional
+ The minimum and the maximum number of turns to plot FFT data.
+ streak_plot : bool, optional
+ If True, the FFT data is plotted as a streak plot.
+ bunch_profile : bool, optional.
+ If True, the FFT data is plotted as line profiles.
Return
------
@@ -515,20 +536,93 @@ def plot_tunedata(filename, bunch_number):
group = "TuneData_{0}".format(bunch_number) # Data group of the HDF5 file
time = file[group]["time"]
- tune = file[group]["tune"]
- tune_spread = file[group]["tune_spread"]
-
- fig1, ax1 = plt.subplots()
- ax1.errorbar(x=time[1:], y=tune[0,1:], yerr=tune_spread[0,1:])
- ax1.errorbar(x=time[1:], y=tune[1,1:], yerr=tune_spread[1,1:])
- ax1.set_xlabel("Turn number")
- ax1.set_ylabel("Transverse tunes")
- plt.legend(["x","y"])
- fig2, ax2 = plt.subplots()
- ax2.errorbar(x=time[1:], y=tune[2,1:], yerr=tune_spread[2,1:])
- ax2.set_xlabel("Turn number")
- ax2.set_ylabel("Synchrotron tune")
+ if plot_tune is True:
+ tune = file[group]["tune"]
+ tune_spread = file[group]["tune_spread"]
+
+ fig1, ax1 = plt.subplots()
+ ax1.errorbar(x=time[1:], y=tune[0,1:], yerr=tune_spread[0,1:])
+ ax1.errorbar(x=time[1:], y=tune[1,1:], yerr=tune_spread[1,1:])
+ ax1.set_xlabel("Turn number")
+ ax1.set_ylabel("Transverse tunes")
+ plt.legend(["x","y"])
+
+ fig2, ax2 = plt.subplots()
+ ax2.errorbar(x=time[1:], y=tune[2,1:], yerr=tune_spread[2,1:])
+ ax2.set_xlabel("Turn number")
+ ax2.set_ylabel("Synchrotron tune")
+
+ if plot_fft is True:
+ if ring is None:
+ T0 = 1.183e-06
+ f0 = 0.845e6
+ else:
+ T0 = ring.T0
+ f0 = ring.f0
+
+ n_freq = file[group]['fft'].shape[1]
+ freq = rfftfreq((n_freq-1)*2, T0)
+ tune_fft = freq / f0
+
+ dimension_dict = {'x':0, 'y':1, 's':2}
+ axis = dimension_dict[dimension]
+
+ fourier_save = file[group]['fft'][axis]
+
+ if max_turn is None:
+ max_turn = time[-1]
+ if min_turn is None:
+ min_turn = time[1]
+
+ min_tune_iloc = np.where(tune_fft >= min_tune)[0][0]
+ max_tune_iloc = np.where(tune_fft <= max_tune)[0][-1]
+ save_every = int(time[1] - time[0])
+ min_turn_iloc = min_turn // save_every
+ max_turn_iloc = max_turn // save_every
+
+ if streak_plot is True:
+ fig3, ax3 = plt.subplots()
+ cmap = plt.get_cmap('Blues')
+
+ c = ax3.imshow(np.transpose(np.log(
+ fourier_save[min_tune_iloc:max_tune_iloc+1,
+ min_turn_iloc:max_turn_iloc+1])),
+ cmap=cmap, origin='lower' , aspect='auto',
+ extent=[min_tune, max_tune, min_turn, max_turn])
+ ax3.set_xlabel('$Q_{}$'.format(dimension))
+ ax3.set_ylabel("Turns")
+ cbar = fig3.colorbar(c, ax=ax3)
+ cbar.set_label("log FFT amplitude")
+
+ if profile_plot is True:
+ fig4, ax4 = plt.subplots()
+ ax4.plot(tune_fft[min_tune_iloc:max_tune_iloc+1],
+ fourier_save[min_tune_iloc:max_tune_iloc+1,
+ min_turn_iloc:max_turn_iloc+1])
+ ax4.set_xlabel('$Q_{}$'.format(dimension))
+ ax4.set_ylabel("FFT amplitude")
+ ax4.legend(time[min_turn_iloc:max_turn_iloc+1])
+
file.close()
- return (fig1, fig2)
\ No newline at end of file
+
+ if plot_tune is True and plot_fft is True:
+ if streak_plot is True and profile_plot is True:
+ return fig1, fig2, fig3, fig4
+ elif streak_plot is True:
+ return fig1, fig2, fig3
+ elif profile_plot is True:
+ return fig1, fig2, fig4
+
+ elif plot_tune is True:
+ return fig1, fig2
+
+ elif plot_fft is True:
+ if streak_plot is True and profile_plot is True:
+ return fig3, fig4
+ elif streak_plot is True:
+ return fig3
+ elif profile_plot is True:
+ return fig4
+
\ No newline at end of file