diff --git a/tracking/monitors/monitors.py b/tracking/monitors/monitors.py
index c8674387184df2827c9bea3da85222d76e4beedd..1cd781da2df3e508129c8e243b68c68ce3ca177d 100644
--- a/tracking/monitors/monitors.py
+++ b/tracking/monitors/monitors.py
@@ -11,11 +11,11 @@ during tracking.
import numpy as np
import h5py as hp
import PyNAFF as pnf
+import random
from mbtrack2.tracking.element import Element
from mbtrack2.tracking.particles import Bunch, Beam
from abc import ABCMeta
-import random
-# from mpi4py import MPI
+from mpi4py import MPI
class Monitor(Element, metaclass=ABCMeta):
"""
@@ -798,8 +798,8 @@ class TuneMonitor(Monitor):
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,)}
+ dict_buffer = {"tune":(3, buffer_size), "tune_spread":(3, buffer_size,)}
+ dict_file = {"tune":(3, total_size), "tune_spread":(3, total_size,)}
self.monitor_init(group_name, save_every, buffer_size, total_size,
dict_buffer, dict_file, file_name, mpi_mode)
@@ -820,37 +820,52 @@ class TuneMonitor(Monitor):
self.buffer_count = 1
- def track(self, bunch):
+ def track(self, object_to_save):
"""
Save tune data.
Parameters
----------
- bunch : Bunch object
+ object_to_save : Beam or 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
+ skip = False
+ if isinstance(object_to_save, Beam):
+ if (object_to_save.mpi_switch == True):
+ if object_to_save.mpi.bunch_num == self.bunch_number:
+ bunch = object_to_save[object_to_save.mpi.bunch_num]
+ else:
+ skip = True
+ else:
+ bunch = object_to_save[self.bunch_number]
+ elif isinstance(object_to_save, Bunch):
+ bunch = object_to_save
+ else:
+ raise TypeError("object_to_save should be a Beam or Bunch object.")
- 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
+ if skip is not True:
+ 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.to_buffer(bunch)
+ self.save_count = 0
+
+ self.track_count += 1
- def to_buffer(self):
+ def to_buffer(self, bunch):
"""
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] = self.mean_tune
- self.sigma_tune[:, self.buffer_count] = self.tune_spread
+ self.tune[:, self.buffer_count] = mean
+ self.tune_spread[:, self.buffer_count] = spread
self.buffer_count += 1
@@ -869,9 +884,9 @@ class TuneMonitor(Monitor):
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.file[self.group_name]["tune_spread"][:,
self.write_count * self.buffer_size:(self.write_count+1) *
- self.buffer_size] = self.sigma_tune
+ self.buffer_size] = self.tune_spread
self.file.flush()
self.write_count += 1
@@ -895,41 +910,23 @@ class TuneMonitor(Monitor):
freq[i,0] = pnf.naff(self.x[i,:], turns=turn-1, nterms=1)[0][1] \
/ self.ring.T0
except IndexError:
- freq[i,0] = 0
+ freq[i,0] = np.nan
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
+ freq[i,1] = np.nan
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()
-
-
-
-
-
-
-
-
+ freq[i,2] = np.nan
+ tune_single_particle = freq / self.ring.f0
+ mean = np.nanmean(tune_single_particle, 0)
+ spread = np.nanstd(tune_single_particle, 0)
-
\ No newline at end of file
+ return (mean, spread)
+
\ No newline at end of file
diff --git a/tracking/monitors/plotting.py b/tracking/monitors/plotting.py
index de90791b3d3b595a973b8398cfc691c24f909bcc..5b02c9bb2a8ae45c488a7e9e92bdd17e10fc63b7 100644
--- a/tracking/monitors/plotting.py
+++ b/tracking/monitors/plotting.py
@@ -481,4 +481,45 @@ def plot_wakedata(filename, bunch_number, wake_type="Wlong", start=0,
return fig
elif streak_plot is True:
return fig2
+
+def plot_tunedata(filename, bunch_number):
+ """
+ Plot data recorded by TuneMonitor.
+
+ Parameters
+ ----------
+ filename : str
+ Name of the HDF5 file that contains the data.
+ bunch_number : int
+ Bunch to plot. This has to be identical to 'bunch_number' parameter in
+ 'BunchMonitor' object.
+
+ Return
+ ------
+ fig : Figure
+ Figure object with the plot on it.
+
+ """
+
+ file = hp.File(filename, "r")
+
+ 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")
+
+ file.close()
+ return (fig1, fig2)
\ No newline at end of file