From a5daf052c3679a6e18b05d4d33cfbe4464148429 Mon Sep 17 00:00:00 2001
From: Gamelin Alexis <gamelin@synchrotron-soleil.fr>
Date: Fri, 20 Mar 2020 14:53:23 +0100
Subject: [PATCH] Monitor with open/close
Add a monitor module to save tracking data based on h5py
Bunch Monitor is used to get statistics from a specific bunch
PhaseSpaceMonitor is used to save the phase space of a specific bunch
For now the .hdf5 file is opened and closed at each write operation, this seems to cause some problems for mpi.
---
tracking/monitors.py | 169 +++++++++++++++++++++++++++++++++++++++++++
1 file changed, 169 insertions(+)
create mode 100644 tracking/monitors.py
diff --git a/tracking/monitors.py b/tracking/monitors.py
new file mode 100644
index 0000000..5d6f686
--- /dev/null
+++ b/tracking/monitors.py
@@ -0,0 +1,169 @@
+# -*- coding: utf-8 -*-
+"""
+This module defines the different monitor class which are used to save data
+from tracking.
+
+@author: Alexis Gamelin
+@date: 17/03/2020
+"""
+
+import numpy as np
+import h5py as hp
+from tracking.element import Element
+from tracking.particles import Bunch, Beam
+from abc import ABCMeta, abstractmethod
+
+class Monitor(Element, metaclass=ABCMeta):
+ """
+ """
+
+ file_name_storage = []
+
+ @property
+ def file_name(self):
+ try:
+ return self.file_name_storage[0]
+ except IndexError:
+ print("File name for monitors not set.")
+ raise ValueError
+
+ def monitor_init(self, group_name, save_every, total_size,
+ buffer_size, dict_buffer, dict_file, file_name=None):
+
+ if file_name is not None:
+ if len(self.file_name_storage) == 0:
+ self.file_name_storage.append(file_name + ".hdf5")
+ else:
+ raise ValueError("File name for monitors is already attributed.")
+
+ self.group_name = group_name
+ self.save_every = int(save_every)
+ self.total_size = int(total_size)
+ self.buffer_size = int(buffer_size)
+ if total_size % buffer_size != 0:
+ raise ValueError("total_size must be divisible by buffer_size.")
+ self.buffer_count = 0
+ self.write_count = 0
+ self.track_count = 0
+
+ for key, value in dict_buffer.items():
+ self.__setattr__(key,np.zeros(value))
+ self.time = np.zeros((self.buffer_size,), dtype=int)
+
+ with hp.File(self.file_name, "a", libver='latest') as f:
+ g = f.require_group(self.group_name)
+ g.require_dataset("time", (self.total_size,), dtype=int)
+ for key, value in dict_file.items():
+ g.require_dataset(key, value, dtype=float)
+
+
+ slice_dict = {}
+ for key, value in dict_file.items():
+ slice_dict[key] = []
+ for i in range(len(value)-1):
+ slice_dict[key].append(slice(None))
+ self.slice_dict = slice_dict
+
+ def write(self):
+ """
+ """
+
+ with hp.File(self.file_name, "a", libver='latest') as f:
+ f[self.group_name]["time"][self.write_count*self.buffer_size:(
+ self.write_count+1)*self.buffer_size] = self.time
+ for key, value in self.dict_buffer.items():
+ slice_list = list(self.slice_dict[key])
+ slice_list.append(slice(self.write_count*self.buffer_size,
+ (self.write_count+1)*self.buffer_size))
+ slice_tuple = tuple(slice_list)
+
+ f[self.group_name][key][slice_tuple] = self.__getattribute__(key)
+ self.write_count += 1
+
+ def to_buffer(self, object_to_save):
+ """
+ """
+ self.time[self.buffer_count] = self.track_count
+ for key, value in self.dict_buffer.items():
+ slice_list = list(self.slice_dict[key])
+ slice_list.append(self.buffer_count)
+ slice_tuple = tuple(slice_list)
+ self.__getattribute__(key)[slice_tuple] = object_to_save.__getattribute__(key)
+ self.buffer_count += 1
+
+ if self.buffer_count == self.buffer_size:
+ self.write()
+ self.buffer_count = 0
+
+class BunchMonitor(Monitor):
+ """
+ Monitor a
+ """
+
+ def __init__(self, bunch_number, file_name=None,
+ save_every=5, buffer_size = 500, total_size = 1e5):
+
+ self.bunch_number = bunch_number
+ group_name = "BunchData_" + str(self.bunch_number)
+ dict_buffer = {"mean":(6,buffer_size), "std":(6,buffer_size),
+ "emit":(3,buffer_size), "current":(buffer_size,)}
+ dict_file = {"mean":(6,total_size), "std":(6,total_size),
+ "emit":(3,total_size), "current":(total_size,)}
+ self.monitor_init(group_name, save_every, total_size,
+ buffer_size, dict_buffer, dict_file, file_name)
+
+ self.dict_buffer = dict_buffer
+ self.dict_file = dict_file
+
+ def track(self, object_to_save):
+ """
+ """
+ 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 PhaseSpaceMonitor(Monitor):
+ """
+ Monitor a
+ """
+
+ def __init__(self, bunch_number, mp_number, file_name=None,
+ save_every=1e4, buffer_size = 1, total_size = 100):
+
+ self.bunch_number = bunch_number
+ self.mp_number = int(mp_number)
+ group_name = "PhaseSpaceData_" + str(self.bunch_number)
+ dict_buffer = {"particles":(self.mp_number, 6, buffer_size)}
+ dict_file = {"particles":(self.mp_number, 6, total_size)}
+ self.monitor_init(group_name, save_every, total_size,
+ buffer_size, dict_buffer, dict_file, file_name)
+
+ self.dict_buffer = dict_buffer
+ self.dict_file = dict_file
+
+ def track(self, object_to_save):
+ """
+ """
+
+ 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
--
GitLab