From 395c1ef64045e2c130d7a1192df5611a3221dfc1 Mon Sep 17 00:00:00 2001
From: Gamelin Alexis <gamelin@synchrotron-soleil.fr>
Date: Tue, 10 Nov 2020 12:59:41 +0100
Subject: [PATCH] Add plot_tunedata and changes to TuneMonitor
Change track method to allow for Beam objects
If pynaff can not find the tune, the tune value is nan.
Nanmean and nanstd are used to avoid to bias the tune value.
! Untested if some particles are lost !
---
tracking/monitors/monitors.py | 95 +++++++++++++++++------------------
tracking/monitors/plotting.py | 41 +++++++++++++++
2 files changed, 87 insertions(+), 49 deletions(-)
diff --git a/tracking/monitors/monitors.py b/tracking/monitors/monitors.py
index c867438..1cd781d 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 de90791..5b02c9b 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
--
GitLab