From eba56ed77960ac2af981683a29e9fc4710a88812 Mon Sep 17 00:00:00 2001
From: Gamelin Alexis <gamelin@synchrotron-soleil.fr>
Date: Sat, 2 May 2020 12:10:11 +0200
Subject: [PATCH] Add ProfileMonitor and misc on monitors.py
Add ProfileMonitor to save bunch profile
Add a track_bunch_data to Monitor class to group the code of track methods from BunchMonitor and PhaseSpaceMonitor
Implement BeamMonitor for the case without mpi
Remove hard coded h=416 ! in BeamMonitor
---
tracking/monitors/monitors.py | 262 +++++++++++++++++++++++++++-------
1 file changed, 212 insertions(+), 50 deletions(-)
diff --git a/tracking/monitors/monitors.py b/tracking/monitors/monitors.py
index af79efa..14ea801 100644
--- a/tracking/monitors/monitors.py
+++ b/tracking/monitors/monitors.py
@@ -46,6 +46,8 @@ class Monitor(Element, metaclass=ABCMeta):
Close the HDF5 file shared by all Monitor subclass, must be called
by at least an instance of a Montior subclass at the end of the
tracking.
+ track_bunch_data(object_to_save)
+ Track method to use when saving bunch data.
"""
_file_name_storage = []
@@ -200,6 +202,28 @@ class Monitor(Element, metaclass=ABCMeta):
self.file.close()
except ValueError:
pass
+
+ def track_bunch_data(self, object_to_save):
+ """
+ Track method to use when saving bunch data.
+
+ Parameters
+ ----------
+ object_to_save : Beam or Bunch
+ """
+ if self.track_count % self.save_every == 0:
+ if isinstance(object_to_save, Beam):
+ if (object_to_save.mpi_switch == True):
+ if object_to_save.mpi.bunch_num == self.bunch_number:
+ self.to_buffer(object_to_save[object_to_save.mpi.bunch_num])
+ else:
+ self.to_buffer(object_to_save[self.bunch_number])
+ elif isinstance(object_to_save, Bunch):
+ self.to_buffer(object_to_save)
+ else:
+ raise TypeError("object_to_save should be a Beam or Bunch object.")
+ self.track_count += 1
+
class BunchMonitor(Monitor):
"""
@@ -257,18 +281,7 @@ class BunchMonitor(Monitor):
----------
object_to_save : Bunch or Beam object
"""
- if self.track_count % self.save_every == 0:
- if isinstance(object_to_save, Beam):
- if (object_to_save.mpi_switch == True):
- if object_to_save.mpi.bunch_num == self.bunch_number:
- self.to_buffer(object_to_save[object_to_save.mpi.bunch_num])
- else:
- self.to_buffer(object_to_save[self.bunch_number])
- elif isinstance(object_to_save, Bunch):
- self.to_buffer(object_to_save)
- else:
- raise TypeError("object_to_save should be a Beam or Bunch object.")
- self.track_count += 1
+ self.track_bunch_data(object_to_save)
class PhaseSpaceMonitor(Monitor):
@@ -329,20 +342,8 @@ class PhaseSpaceMonitor(Monitor):
Parameters
----------
object_to_save : Bunch or Beam object
- """
-
- if self.track_count % self.save_every == 0:
- if isinstance(object_to_save, Beam):
- if (object_to_save.mpi_switch == True):
- if object_to_save.mpi.bunch_num == self.bunch_number:
- self.to_buffer(object_to_save[object_to_save.mpi.bunch_num])
- else:
- self.to_buffer(object_to_save[self.bunch_number])
- elif isinstance(object_to_save, Bunch):
- self.to_buffer(object_to_save)
- else:
- raise TypeError("object_to_save should be a Beam or Bunch object.")
- self.track_count += 1
+ """
+ self.track_bunch_data(object_to_save)
class BeamMonitor(Monitor):
@@ -377,24 +378,21 @@ class BeamMonitor(Monitor):
Save data
"""
- def __init__(self, file_name=None, save_every=5, buffer_size=500,
+ def __init__(self, h, file_name=None, save_every=5, buffer_size=500,
total_size=2e4, mpi_mode=True):
group_name = "Beam"
- dict_buffer = {}
- dict_file = {}
+ dict_buffer = {"mean" : (6, h, buffer_size),
+ "std" : (6, h, buffer_size),
+ "emit" : (3, h, buffer_size),
+ "current" : (h, buffer_size)}
+ dict_file = {"mean" : (6, h, total_size),
+ "std" : (6, h, total_size),
+ "emit" : (3, h, total_size),
+ "current" : (h, total_size)}
+
self.monitor_init(group_name, save_every, buffer_size, total_size,
dict_buffer, dict_file, file_name, mpi_mode)
-
- self.mean = np.zeros((6, self.buffer_size), dtype=float)
- self.std = np.zeros((6, self.buffer_size), dtype=float)
- self.emit = np.zeros((3, self.buffer_size), dtype=float)
- self.current = np.zeros((self.buffer_size,), dtype=float)
-
- self.g.require_dataset("mean", (6, 416, self.total_size,), dtype=float)
- self.g.require_dataset("std", (6, 416, self.total_size,), dtype=float)
- self.g.require_dataset("emit", (3, 416, self.total_size,), dtype=float)
- self.g.require_dataset("current", (416, self.total_size,), dtype=float)
def track(self, beam):
"""
@@ -408,13 +406,13 @@ class BeamMonitor(Monitor):
if (beam.mpi_switch == True):
self.to_buffer(beam[beam.mpi.bunch_num], beam.mpi.bunch_num)
else:
- raise NotImplementedError
+ self.to_buffer_no_mpi(beam)
self.track_count += 1
def to_buffer(self, bunch, bunch_num):
"""
- Save data to buffer.
+ Save data to buffer, if mpi is being used.
Parameters
----------
@@ -423,20 +421,41 @@ class BeamMonitor(Monitor):
"""
self.time[self.buffer_count] = self.track_count
- self.mean[:, self.buffer_count] = bunch.mean
- self.std[:, self.buffer_count] = bunch.std
- self.emit[:, self.buffer_count] = bunch.emit
- self.current[self.buffer_count] = bunch.current
+ self.mean[:, bunch_num, self.buffer_count] = bunch.mean
+ self.std[:, bunch_num, self.buffer_count] = bunch.std
+ self.emit[:, bunch_num, self.buffer_count] = bunch.emit
+ self.current[bunch_num, self.buffer_count] = bunch.current
self.buffer_count += 1
if self.buffer_count == self.buffer_size:
self.write(bunch_num)
self.buffer_count = 0
+
+ def to_buffer_no_mpi(self, beam):
+ """
+ Save data to buffer, if mpi is not being used.
+
+ Parameters
+ ----------
+ beam : Beam object
+ """
+
+ self.time[self.buffer_count] = self.track_count
+ self.mean[:, :, self.buffer_count] = beam.bunch_mean
+ self.std[:, :, self.buffer_count] = beam.bunch_std
+ self.emit[:, :, self.buffer_count] = beam.bunch_emit
+ self.current[:, self.buffer_count] = beam.bunch_current
+
+ self.buffer_count += 1
+
+ if self.buffer_count == self.buffer_size:
+ self.write_no_mpi()
+ self.buffer_count = 0
def write(self, bunch_num):
"""
- Write data from buffer to the HDF5 file.
+ Write data from buffer to the HDF5 file, if mpi is being used.
Parameters
----------
@@ -447,21 +466,164 @@ class BeamMonitor(Monitor):
self.file[self.group_name]["mean"][:, bunch_num,
self.write_count*self.buffer_size:(self.write_count+1) *
- self.buffer_size] = self.mean
+ self.buffer_size] = self.mean[:, bunch_num, :]
self.file[self.group_name]["std"][:, bunch_num,
self.write_count*self.buffer_size:(self.write_count+1) *
- self.buffer_size] = self.std
+ self.buffer_size] = self.std[:, bunch_num, :]
self.file[self.group_name]["emit"][:, bunch_num,
self.write_count*self.buffer_size:(self.write_count+1) *
- self.buffer_size] = self.emit
+ self.buffer_size] = self.emit[:, bunch_num, :]
self.file[self.group_name]["current"][bunch_num,
+ self.write_count*self.buffer_size:(self.write_count+1) *
+ self.buffer_size] = self.current[bunch_num, :]
+
+ self.file.flush()
+ self.write_count += 1
+
+ def write_no_mpi(self):
+ """
+ Write data from buffer to the HDF5 file, if mpi is not being used.
+ """
+
+ 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]["mean"][:, :,
+ self.write_count*self.buffer_size:(self.write_count+1) *
+ self.buffer_size] = self.mean
+
+ self.file[self.group_name]["std"][:, :,
+ self.write_count*self.buffer_size:(self.write_count+1) *
+ self.buffer_size] = self.std
+
+ self.file[self.group_name]["emit"][:, :,
+ self.write_count*self.buffer_size:(self.write_count+1) *
+ self.buffer_size] = self.emit
+
+ self.file[self.group_name]["current"][:,
self.write_count*self.buffer_size:(self.write_count+1) *
self.buffer_size] = self.current
self.file.flush()
self.write_count += 1
+
+
+class ProfileMonitor(Monitor):
+ """
+ Monitor a single bunch and save bunch profiles.
+
+ Parameters
+ ----------
+ bunch_number : int
+ Bunch to monitor.
+ dimensions : str or list of str, optional
+ Dimensions to save.
+ n_bin : int or list of int, optional
+ Number of bin to use in each dimension.
+ 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 data is 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(object_to_save)
+ Save data
+ """
+
+ def __init__(self, bunch_number, dimensions="tau", n_bin=75, file_name=None,
+ save_every=5, buffer_size=500, total_size=2e4, mpi_mode=True):
-
\ No newline at end of file
+ self.bunch_number = bunch_number
+ group_name = "ProfileData_" + str(self.bunch_number)
+
+ if isinstance(dimensions, str):
+ self.dimensions = [dimensions]
+ else:
+ self.dimensions = dimensions
+
+ if isinstance(n_bin, int):
+ self.n_bin = np.ones((len(self.dimensions),), dtype=int)*n_bin
+ else:
+ self.n_bin = n_bin
+
+ dict_buffer = {}
+ dict_file = {}
+ for index, dim in enumerate(dimensions):
+ dict_buffer.update({dim : (self.n_bin[index], buffer_size)})
+ dict_buffer.update({dim + "_bin" : (self.n_bin[index] + 1, buffer_size)})
+ dict_file.update({dim : (self.n_bin[index], total_size)})
+ dict_file.update({dim + "_bin" : (self.n_bin[index] + 1, 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
+
+ def to_buffer(self, bunch):
+ """
+ Save data to buffer.
+
+ Parameters
+ ----------
+ bunch : Bunch object
+ """
+
+ self.time[self.buffer_count] = self.track_count
+ for index, dim in enumerate(self.dimensions):
+ bins, sorted_index, profile = bunch.binning(dim, self.n_bin[index])
+ bin_array = np.append(np.array(bins.left), bins.right[-1])
+ profile_array = np.array(profile)
+ self.__getattribute__(dim + "_bin")[:, self.buffer_count] = bin_array
+ self.__getattribute__(dim)[:, self.buffer_count] = profile_array
+
+ 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 the HDF5 file."""
+
+ self.file[self.group_name]["time"][self.write_count*self.buffer_size:(
+ self.write_count+1)*self.buffer_size] = self.time
+
+ for dim in self.dimensions:
+ self.file[self.group_name][dim][:,
+ self.write_count * self.buffer_size:(self.write_count+1) *
+ self.buffer_size] = self.__getattribute__(dim)
+ self.file[self.group_name][dim + "_bin"][:,
+ self.write_count * self.buffer_size:(self.write_count+1) *
+ self.buffer_size] = self.__getattribute__(dim + "_bin")
+
+ self.file.flush()
+ self.write_count += 1
+
+ def track(self, object_to_save):
+ """
+ Save data.
+
+ Parameters
+ ----------
+ object_to_save : Bunch or Beam object
+ """
+ self.track_bunch_data(object_to_save)
\ No newline at end of file
--
GitLab