Add save and load methods for Bunch and Beam

Save and load methods allow to save Bunch and Beam object and reload them later on.
Load method for Beam allow to choose if mpi is used or not.
Initial modification from Watanyu.

mbtrack2/tracking/particles.py

@@ -7,6 +7,7 @@ import numpy as np
 import matplotlib.pyplot as plt
 import seaborn as sns
 import pandas as pd
+import h5py as hp
 from scipy.constants import c, m_e, m_p, e
 
 class Particle:
@@ -57,6 +58,15 @@ class Bunch:
         Can be set to False to gain a speed increase.
     alive : bool, optional
         If False, the bunch is defined as empty.
+    load_from_file : str, optional
+        Name of bunch save file generated by save method.
+        If None, the bunch is initialized using the input parameters.
+        Otherwise, the bunch from the file is loaded and the other inputs are 
+        ignored.
+        Default is None.
+    load_suffix : str or int, optional
+        Suffix to the group name used to load the data from the HDF5 file.
+        Default is None.
         
     Attributes
     ----------
@@ -88,7 +98,12 @@ class Bunch:
         Initialize bunch particles with 6D gaussian phase space.
     plot_phasespace(x_var="tau", y_var="delta", plot_type="j")
         Plot phase space.
-        
+    save(file_name)
+        Save bunch object data (6D phase space, current, and state) in an HDF5 
+        file format.
+    load(file_name)
+        Load data from a HDF5 file recorded by Bunch save method.
+    
     References
     ----------
     [1] Wiedemann, H. (2015). Particle accelerator physics. 4th edition. 
@@ -96,7 +111,7 @@ class Bunch:
     """
     
     def __init__(self, ring, mp_number=1e3, current=1e-3, track_alive=True,
-                 alive=True):
+                 alive=True, load_from_file=None, load_suffix=None):
         
         self.ring = ring
         if not alive:
@@ -113,11 +128,13 @@ class Bunch:
         
         self.particles = np.zeros(self.mp_number, self.dtype)
         self.track_alive = track_alive
-        
         self.alive = np.ones((self.mp_number,),dtype=bool)
         self.current = current
         if not alive:
             self.alive = np.zeros((self.mp_number,),dtype=bool)
+            
+        if load_from_file is not None:
+            self.load(load_from_file, load_suffix, track_alive)
         
     def __len__(self):
         """Return the number of alive particles"""
@@ -391,6 +408,87 @@ class Bunch:
             raise ValueError("Plot type not recognised.")
             
         return fig
+    
+    def save(self, file_name, suffix=None, mpi_comm=None):
+        """
+        Save bunch object data (6D phase space, current, and state) in an HDF5 
+        file format.
+        
+        The output file is named as "<file_name>.hdf5".
+
+        Parameters
+        ----------
+        file_name : str
+            Name of the output file.
+        suffix : str or int, optional
+            Suffix to the group name used to save the data into the HDF5 file.
+            If None, it is not used.
+        mpi_comm : MPI communicator, optional
+            For internal use if mpi is used in Beam objects.
+            Default is None.
+        """
+        
+        if mpi_comm is None:
+            f = hp.File(file_name + ".hdf5", "a", libver="earliest")
+        else:
+            f = hp.File(file_name+ ".hdf5", "a", libver='earliest', 
+                        driver='mpio', comm=mpi_comm)
+        
+        if suffix is None:
+            group_name = "Bunch"
+        else:
+            group_name = "Bunch_" + str(suffix)
+            
+        g = f.create_group(group_name)
+        g.create_dataset("alive", (self.mp_number,), dtype=bool)
+        g.create_dataset("phasespace", (self.mp_number, 6), dtype=float)
+        g.create_dataset("current", (1,), dtype=float)
+        
+        f[group_name]["alive"][:] = self.alive
+        f[group_name]["current"][:] = self.current
+        for i, dim in enumerate(self):
+            f[group_name]["phasespace"][:,i] = self.particles[dim]
+        
+        f.close()
+        
+    def load(self, file_name, suffix=None, track_alive=True):
+        """
+        Load data from a HDF5 file recorded by Bunch save method.
+
+        Parameters
+        ----------
+        file_name : str
+            Path to the HDF5 file where the Bunch data is stored.
+        suffix : str or int, optional
+            Suffix to the group name used to load the data from the HDF5 file.
+            If None, it is not used.
+        track_alive : bool, optional
+            If False, the code no longer take into account alive/dead particles.
+            Should be set to True if element such as apertures are used.
+            Can be set to False to gain a speed increase.
+        """
+        
+        f = hp.File(file_name, "r")
+        
+        if suffix is None:
+            group_name = "Bunch"
+        else:
+            group_name = "Bunch_" + str(suffix)
+
+        self.mp_number = len(f[group_name]['alive'][:])
+        
+        for i, dim in enumerate(self):
+            self.particles[dim] = f[group_name]["phasespace"][:,i]
+        
+        self.alive = f[group_name]['alive'][:]
+        if f[group_name]['current'][:][0] != 0:
+            self.current = f[group_name]['current'][:][0]
+        else:
+            self.charge_per_mp = 0
+            
+        self.track_alive = track_alive
+        
+        f.close()
         
 class Beam:
     """
@@ -447,6 +545,10 @@ class Beam:
         Call mpi_gather and switch off MPI parallelisation
     plot(var, option=None)
         Plot variables with respect to bunch number.
+    save(file_name)
+        Save beam object data in an HDF5 file format.
+    load(file_name, mpi)
+        Load data from a HDF5 file recorded by Beam save method.
     """
     
     def __init__(self, ring, bunch_list=None):
@@ -808,4 +910,77 @@ class Beam:
        
         return fig
         
+    def save(self, file_name):
+        """
+        Save beam object data in an HDF5 file format.
+        
+        The output file is named as "<file_name>.hdf5".
+
+        Parameters
+        ----------
+        file_name : str
+            Name of the output file.
+        """
+        if self.mpi_switch is True:
+            for i, bunch in enumerate(self):
+                if i in self.mpi.table[:,1]:
+                    if i == self.mpi.bunch_num:
+                        mp_number = self[self.mpi.bunch_num].mp_number
+                        self.mpi.comm.bcast(mp_number, root=self.mpi.rank)
+                        self[self.mpi.bunch_num].save(file_name, 
+                                                      self.mpi.bunch_num, 
+                                                      self.mpi.comm)
+                    else:
+                        mp_number = None
+                        mp_number = self.mpi.comm.bcast(mp_number, root=self.mpi.bunch_to_rank(i))
+                        f = hp.File(file_name+ ".hdf5", "a", libver='earliest', 
+                                    driver='mpio', comm=self.mpi.comm)
+                        group_name = "Bunch_" + str(i)
+                        g = f.create_group(group_name)
+                        g.create_dataset("alive", (mp_number,), dtype=bool)
+                        g.create_dataset("phasespace", (mp_number, 6), dtype=float)
+                        g.create_dataset("current", (1,), dtype=float)  
+                        f.close()
+                else:
+                    bunch.save(file_name, 
+                               i,
+                               self.mpi.comm)
+        else:
+            for i, bunch in enumerate(self):
+                bunch.save(file_name, i)
     
+    def load(self, file_name, mpi, track_alive=True):
+        """
+        Load data from a HDF5 file recorded by Beam save method.
+
+        Parameters
+        ----------
+        file_name : str
+            Path to the HDF5 file where the Beam data is stored.
+        mpi : bool
+            If True, only a single bunch is fully initialized on each core, the
+            other bunches are initialized with a single marco-particle.
+        track_alive : bool, optional
+            If False, the code no longer take into account alive/dead particles.
+            Should be set to True if element such as apertures are used.
+            Can be set to False to gain a speed increase.
+            The default is True.
+        """
+        if mpi is True:
+            self.__init__(self.ring)
+            f = hp.File(file_name, "r")
+            filling_pattern = []
+            for i in range(self.ring.h):
+                current = f["Bunch_" + str(i)]['current'][:][0]
+                filling_pattern.append(current)
+                    
+            self.init_beam(filling_pattern,
+                           mp_per_bunch=1,
+                           track_alive=track_alive,
+                           mpi=True)
+            self[self.mpi.bunch_num].load(file_name, self.mpi.bunch_num)
+        else:
+            for i, bunch in enumerate(self):
+                bunch.load(file_name, i, track_alive)
+            self.update_filling_pattern()
+            self.update_distance_between_bunches()
\ No newline at end of file

mbtrack2/tracking/wakepotential.py

@@ -77,6 +77,9 @@ class WakePotential(Element):
         self.ring = ring
         self.n_bin = n_bin
         self.check_sampling()
+        
+        # Suppress numpy warning for floating-point operations.
+        np.seterr(invalid='ignore')
             
     def charge_density(self, bunch):
         """