read_snaps option for DynaRun

src/perturbopy/postproc/calc_modes/dyna_run.py

@@ -19,6 +19,10 @@ class DynaRun(CalcMode):
     """
     Class representation of a Perturbo dynamics-run calculation.
 
+    TODO: better handling with snaps. Make it as property.
+    Print a sensible message if snaps is None and access is attempted.
+    Implement a getter and setter for snaps linked to HDF5 file.
+
     Attributes
     ----------
     kpt : RecipPtDB
@@ -40,19 +44,32 @@ class DynaRun(CalcMode):
         Python dictionary of DynaIndivRun objects containing results from each simulation
     """
 
-    def __init__(self, cdyna_file, tet_file, pert_dict):
+    def __init__(self, cdyna_file, tet_file, pert_dict, read_snaps=True):
         """
         Constructor method
 
         Parameters
         ----------
+        cdyna_file : h5py.File
+            HDF5 file containing the results of the dynamics-run calculation.
+            The file is kept open after the object is created.
+
+        tet_file : h5py.File
+            Tetrahedron HDF5 file obtained from the setup calculation.
+
         pert_dict : dict
-            Dictionary containing the inputs and outputs from the dynamics-run calculation.
+            Dictionary containing the information from the dynamics-run YAML file.
+
+        read_snaps : bool, optional
+            Flag to read the snapshots from the HDF5 file.
+            Reading the snapshots can be time-consuming and memory-intensive.
 
         """
 
         self.timings = TimingGroup("dynamics-run")
 
+        self.read_snaps = read_snaps
+
         super().__init__(pert_dict)
 
         if self.calc_mode != 'dynamics-run':
@@ -102,10 +119,14 @@ def __init__(self, cdyna_file, tet_file, pert_dict):
                 # a dynamics run must have at least one snap
                 numk, numb = cdyna_file[dyn_str]['snap_t_1'][()].shape
 
-                snap_t = np.zeros((numb, numk, num_steps), dtype=np.float64)
+                if read_snaps:
+                    print(f'Reading snapshots for {dyn_str}...')
+                    snap_t = np.zeros((numb, numk, num_steps), dtype=np.float64)
 
-                for itime in range(num_steps):
-                    snap_t[:, :, itime] = cdyna_file[dyn_str][f'snap_t_{itime + 1}'][()].T
+                    for itime in range(num_steps):
+                        snap_t[:, :, itime] = cdyna_file[dyn_str][f'snap_t_{itime + 1}'][()].T
+                else:
+                    snap_t = None
 
                 # Get E-field, which is only present if nonzero
                 if "efield" in cdyna_file[dyn_str].keys():
@@ -116,7 +137,7 @@ def __init__(self, cdyna_file, tet_file, pert_dict):
                 self._data[irun] = DynaIndivRun(num_steps, time_step, snap_t, time_units='fs', efield=efield)
 
     @classmethod
-    def from_hdf5_yaml(cls, cdyna_path, tet_path, yaml_path='pert_output.yml'):
+    def from_hdf5_yaml(cls, cdyna_path, tet_path, yaml_path='pert_output.yml', read_snaps=True):
         """
         Class method to create a DynamicsRunCalcMode object from the HDF5 file and YAML file
         generated by a Perturbo calculation
@@ -148,7 +169,7 @@ def from_hdf5_yaml(cls, cdyna_path, tet_path, yaml_path='pert_output.yml'):
         cdyna_file = open_hdf5(cdyna_path)
         tet_file = open_hdf5(tet_path)
 
-        return cls(cdyna_file, tet_file, yaml_dict)
+        return cls(cdyna_file, tet_file, yaml_dict, read_snaps=read_snaps)
 
     def close_hdf5_files(self):
         """
@@ -204,7 +225,7 @@ def __str__(self):
         Method to print the dynamics run information
         """
 
-        title = f'dynamics-run: {self.prefix}'
+        title = f' dynamics-run: {self.prefix} '
         text = f"{title:*^60}\n"
         text += f"This simulation has {self.num_runs} runs\n"
 
@@ -213,6 +234,7 @@ def __str__(self):
             text += f"{'Number of steps':>30}: {dynamics_run.num_steps}\n"
             text += f"{'Time step (fs)':>30}: {dynamics_run.time_step}\n"
             text += f"{'Electric field (V/cm)':>30}: {dynamics_run.efield}\n\n"
+            text += f"{'Read snaps at init:':>30}: {self.read_snaps}\n"
 
         return text
 
@@ -455,7 +477,7 @@ def __str__(self):
         Method to print the pump pulse excitation parameters.
         """
 
-        text = 'Pump pulse excitation parameters:\n'
+        text = f'{" Pump pulse parameters ":*^60}\n'
         text += f"{f'Pump energy ({self.pump_energy_units})':>30}: {self.pump_energy:.4f}\n"
         text += f"{f'Energy broadening FWHM ({self.spectral_width_fwhm_units})':>30}: {self.spectral_width_fwhm:.4f}\n"
         text += f"{f'Pulse duration FWHM ({self.pump_duration_fwhm_units})':>30}: {self.pump_duration_fwhm:.4f}\n"