diff --git a/mbtrack2/tracking/particles.py b/mbtrack2/tracking/particles.py
index 49b9adf41b8c683d0a4f63fdfccde5d4b0170ef5..f086eaaa0790a2eb329f455cc995cb851bc4b8c7 100644
--- a/mbtrack2/tracking/particles.py
+++ b/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