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

+# -*- 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