Transient absorption calculation and plotting (#91)

* Transient absorption calculation with numpy vectorization.

* More docs for pump pulse class and read more pump parameters from YAML.

* Transient absorption plotting

setup.cfg

@@ -1,3 +1,3 @@
 [flake8]
 max-line-length = 160
-ignore = E114, E127, E124, E221, W293, E721, W503, F841, F401, E202
+ignore = E114, E127, E124, E221, W293, E721, W503, F841, F401, E202, W504

src/perturbopy/postproc/__init__.py

@@ -17,7 +17,8 @@
 from .dbs.units_dict import UnitsDict
 from .dbs.recip_pt_db import RecipPtDB
 
-from .utils import constants, plot_tools, lattice, spectra_generate_pulse, timing
+from .utils import constants, plot_tools, lattice, spectra_generate_pulse, \
+    timing, spectra_trans_abs, spectra_plots
 
 import warnings
 import os

src/perturbopy/postproc/calc_modes/dyna_run.py

@@ -20,6 +20,15 @@ class DynaRun(CalcMode):
        Database for the band energies computed by the bands calculation.
     num_runs : int
         Number of separate simulations performed
+    _cdyna_file : h5py.File
+        HDF5 file containing the results of the dynamics-run calculation.
+        We cannot store in RAM the whole file, so we will access the data as needed.
+    _tet_file : h5py.File
+        HDF5 file containing the results of the setup calculation required before the dynamics-run calculation.
+    _kpoints : array
+        Raw array of k-points. Shape (num_kpoints, 3)
+    _energies : array
+        Raw array of band energies. Shape (num_kpoints, num_bands)
     _dat : dict
         Python dictionary of DynaIndivRun objects containing results from each simulation
     """
@@ -42,6 +51,19 @@ def __init__(self, cdyna_file, tet_file, pert_dict):
         if self.calc_mode != 'dynamics-run':
             raise ValueError('Calculation mode for a DynamicsRunCalcMode object should be "dynamics-run"')
 
+        self._cdyna_file = cdyna_file
+        self._tet_file = tet_file
+
+        self.pump_pulse = self._pert_dict['input parameters']['after conversion']['pump_pulse']
+
+        if self.pump_pulse:
+            if 'pump pulse' not in self._pert_dict['dynamics-run'].keys():
+                raise ValueError('pump_pulse was set to .true. but no "pump pulse" key'
+                                 ' was found in the dynamics-run section of the YAML file')
+
+            pump_dict = self._pert_dict['dynamics-run']['pump pulse']
+            self.pump_pulse = PumpPulse(pump_dict)
+
         kpoint = np.array(tet_file['kpts_all_crys_coord'][()])
 
         self.kpt = RecipPtDB.from_lattice(kpoint, "crystal", self.lat, self.recip_lat)
@@ -50,6 +72,10 @@ def __init__(self, cdyna_file, tet_file, pert_dict):
         energies_dict = {i + 1: np.array(energies[:, i]) for i in range(0, energies.shape[1])}
         self.bands = UnitsDict.from_dict(energies_dict, 'Ry')
 
+        # k- and q-point grids
+        self.boltz_kdim = np.array(self._pert_dict['input parameters']['after conversion']['boltz_kdim'])
+        self.boltz_qdim = np.array(self._pert_dict['input parameters']['after conversion']['boltz_qdim'])
+
         # Raw arrays
         self._kpoints = kpoint
         self._energies = energies
@@ -117,6 +143,23 @@ def from_hdf5_yaml(cls, cdyna_path, tet_path, yaml_path='pert_output.yml'):
 
         return cls(cdyna_file, tet_file, yaml_dict)
 
+    def close_hdf5_files(self):
+        """
+        Method to close the HDF5 files: _cdyna_file and _tet_file.
+        After the DynaRun object is created, the HDF5 files are kept open.
+        One has to close them manually.
+        """
+
+        # Check if cdyan_file is open
+        if self._cdyna_file:
+            print(f'Closing {self._cdyna_file.filename}')
+            close_hdf5(self._cdyna_file)
+
+        # Check if tet_file is open
+        if self._tet_file:
+            print(f'Closing {self._tet_file.filename}')
+            close_hdf5(self._tet_file)
+
     def __getitem__(self, index):
         """
         Method to index the DynamicsRunCalcMode object
@@ -149,20 +192,22 @@ def __len__(self):
 
         return self.num_runs
 
-    def get_info(self):
+    def __str__(self):
         """
-        Method to get overall information on dynamics runs
+        Method to print the dynamics run information
         """
 
-        print(f"\nThis simulation has {self.num_runs} runs")
+        title = f'dynamics-run: {self.prefix}'
+        text = f"{title:*^60}\n"
+        text += f"This simulation has {self.num_runs} runs\n"
 
         for irun, dynamics_run in self._data.items():
+            text += f"{'Dynamics run':>30}: {irun}\n"
+            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"
 
-            print(f"{'Dynamics run':>30}: {irun}")
-            print(f"{'Number of steps':>30}: {dynamics_run.num_steps}")
-            print(f"{'Time step (fs)':>30}: {dynamics_run.time_step}")
-            print(f"{'Electric field (V/cm)':>30}: {dynamics_run.efield}")
-            print("")
+        return text
 
     def extract_steady_drift_vel(self, dyna_pp_yaml_path):
         """
@@ -199,3 +244,90 @@ def extract_steady_drift_vel(self, dyna_pp_yaml_path):
                 steady_conc.append(concs[step_number])
 
         return steady_drift_vel, steady_conc
+
+
+class PumpPulse():
+    """
+    Class for pump pulse excitation.
+
+    Attributes
+    ----------
+
+    pump_energy : float
+        Energy of the pump pulse excitation in eV.
+
+    energy_broadening : float
+        Energy broadening of the pump pulse excitation in eV.
+
+    num_steps : int
+        Number of steps in the pump pulse excitation.
+
+    time_step : float
+        Time step of the pump pulse excitation in fs.
+
+    num_kpoints : int
+        Number of k-points in the pump pulse excitation. Tailored for the Perturbo dynamics-run calculation.
+
+    carrier_number_array : np.ndarray
+        Additional carrier number array for the pump pulse excitation.
+
+    optional_params : dict
+        Optional parameters for the pump pulse excitation.
+        Specific to the pulse shape. 10 parameters allocated.
+
+    hole : bool
+        Flag to indicate if the pump pulse excitation is for the hole.
+        Must be the same as in the ultrafast simulation.
+    """
+
+    def __init__(self, pump_dict):
+        """
+        Constructor method
+
+        Parameters
+        ----------
+        pump_dict : dict
+            Dictionary containing the pump pulse excitation parameters.
+        """
+
+        # TODO: use UnitsDict for entries in pump_dict
+        # CURRENTLY, UnitsDict seems to be ill-suited for floats with units
+        # input_dict = {'pump_energy': pump_dict['pump_energy']}
+        # self.pump_energy = UnitsDict.from_dict(input_dict, units='eV')
+
+        self.pump_energy = pump_dict['pump_energy']
+        self.pump_energy_units = pump_dict['pump_energy units']
+
+        self.energy_broadening = pump_dict['energy_broadening']
+        self.energy_broadening_units = pump_dict['energy_broadening units']
+
+        self.num_steps = pump_dict['num_steps']
+
+        self.time_step = pump_dict['time_step']
+        self.time_step_units = pump_dict['time_step units']
+
+        self.num_kpoints = pump_dict['num_kpoints']
+
+        self.carrier_number_array = \
+            np.array(pump_dict['pump pulse carrier number'])
+        self.carrier_number_units = pump_dict['carrier_number units']
+
+        # Optional parameters are specific to the pump pulse excitation
+        # 10 parameters allocated
+        self.optional_params = pump_dict['optional_params']
+
+        self.hole = pump_dict['hole']
+
+    def __str__(self):
+        """
+        Method to print the pump pulse excitation parameters.
+        """
+
+        text = 'Pump pulse excitation parameters:\n'
+        text += f"{'Pump energy':>30}: {self.pump_energy} {self.pump_energy_units}\n"
+        text += f"{'Energy broadening':>30}: {self.energy_broadening} {self.energy_broadening_units}\n"
+        text += f"{'Number of steps':>30}: {self.num_steps}\n"
+        text += f"{'Time step':>30}: {self.time_step} {self.time_step_units}\n"
+        text += f"{'Hole':>30}: {self.hole}\n"
+
+        return text

src/perturbopy/postproc/dbs/units_dict.py

@@ -10,6 +10,7 @@ class UnitsDict(dict):
     Attributes
     ----------
     data : dict
+       TODO: UnitsDict DOES NOT HAVE A data ATTRIBUTE
        Dictionary of floats or array_like types of physical quantities
     units : str {}
        The units of the physical quantities
@@ -45,5 +46,5 @@ def convert_lists_to_numpy(data):
         input_dict = convert_lists_to_numpy(input_dict)
         units_dict = cls(units)
         units_dict.update(input_dict)
-        
+
         return units_dict

src/perturbopy/postproc/utils/spectra_generate_pulse.py

@@ -1,5 +1,5 @@
 """
-Utils for ultrafast spectroscopy: pump pulse generation and post-processing
+Utils for ultrafast spectroscopy: pump pulse generation.
 """
 
 import os
@@ -265,8 +265,8 @@ def setup_pump_pulse(elec_pump_pulse_path, hole_pump_pulse_path,
     elec_energy_array *= energy_conversion_factor('Ry', 'eV')
     hole_energy_array *= energy_conversion_factor('Ry', 'eV')
 
-    VBM = max(hole_energy_array[:, -1])
-    CBM = min(elec_energy_array[:, -1])
+    VBM = np.max(hole_energy_array.ravel())
+    CBM = np.min(elec_energy_array.ravel())
     bandgap = CBM - VBM
 
     elec_nband = len(elec_dyna_run.bands)
@@ -340,6 +340,13 @@ def setup_pump_pulse(elec_pump_pulse_path, hole_pump_pulse_path,
     elec_pump_pulse_file = open_hdf5(elec_pump_pulse_path, mode='w')
     hole_pump_pulse_file = open_hdf5(hole_pump_pulse_path, mode='w')
 
+    # Optional pump pulse parameters specific to a given shape of pulse
+    optional_params = np.zeros(10, dtype=float)
+    optional_params[0] = pump_factor
+
+    if finite_width:
+        optional_params[1] = pump_fwhm
+
     for h5f in [elec_pump_pulse_file, hole_pump_pulse_file]:
         h5f.create_group('pump_pulse_snaps')
 
@@ -349,16 +356,21 @@ def setup_pump_pulse(elec_pump_pulse_path, hole_pump_pulse_path,
         h5f.create_dataset('pump_energy_broadening', data=pump_energy_broadening)
         h5f['pump_energy_broadening'].attrs['units'] = 'eV'
 
+        h5f['optional_params'] = optional_params
+
         h5f.create_dataset('finite_width', data=finite_width)
         if finite_width:
             h5f.create_dataset('time_window', data=pump_time_window)
             h5f.create_dataset('num_steps', data=num_steps)
             h5f.create_dataset('time_step', data=pump_time_step)
+            h5f.create_dataset('pulse_type', data=f'gaussian; FWHM: {pump_fwhm:.3f} fs; '
+                                                    f'pump_factor: {pump_factor:.3f}')
 
         else:
             h5f.create_dataset('time_window', data=elec_dyna_time_step)
             h5f.create_dataset('num_steps', data=1)
             h5f.create_dataset('time_step', data=elec_dyna_time_step)
+            h5f.create_dataset('pulse_type', data='step')
 
         h5f['time_window'].attrs['units'] = 'fs'
         h5f['time_step'].attrs['units'] = 'fs'
@@ -368,6 +380,13 @@ def setup_pump_pulse(elec_pump_pulse_path, hole_pump_pulse_path,
     elec_pump_pulse_file.create_dataset('num_kpoints', data=elec_kpoint_array.shape[0])
     hole_pump_pulse_file.create_dataset('num_kpoints', data=hole_kpoint_array.shape[0])
 
+    elec_pump_pulse_file.create_dataset('hole', data=0)
+    hole_pump_pulse_file.create_dataset('hole', data=1)
+
+    # Add carrier attribute to the pump_pulse_snaps group
+    elec_pump_pulse_file.create_dataset('carrier', data='electrons')
+    hole_pump_pulse_file.create_dataset('carrier', data='holes')
+
     # We save the delta_occs_array for animation. If it takes too much space, remove it.
     elec_delta_occs_array = \
         gaussian_excitation(elec_pump_pulse_file, elec_occs_amplitude,

src/perturbopy/postproc/utils/spectra_plots.py

@@ -12,6 +12,9 @@
 from perturbopy.io_utils.io import open_hdf5, close_hdf5
 
 from matplotlib.animation import FuncAnimation
+from mpl_toolkits.axes_grid1 import make_axes_locatable
+from matplotlib.cm import ScalarMappable
+from matplotlib.colors import Normalize
 
 from .memory import get_size
 from .timing import TimingGroup
@@ -139,8 +142,7 @@ def animate_pump_pulse(time_step,
     ani = FuncAnimation(fig, update_scatter, frames=elec_delta_occs_array.shape[1],
                         fargs=(ax, time_step, idx_elec, idx_hole,
                                elec_scat_list, hole_scat_list,
-                               elec_delta_occs_array, elec_kpoint_array, elec_energy_array,
-                               hole_delta_occs_array, hole_kpoint_array, hole_energy_array),
+                               elec_delta_occs_array, hole_delta_occs_array),
                         interval=100, repeat=False)
 
     # Save the animation to gif
@@ -197,3 +199,68 @@ def update_scatter(anim_time, ax, time_step, idx_elec, idx_hole,
 
     suffix = ax.get_title().split(';')[0]
     ax.set_title(f'{suffix}; Time: {anim_time * time_step:.2f} fs')
+
+
+def plot_trans_abs_map(ax, time_grid, energy_grid, trans_abs,
+                       num_contours=500, cmap=plt.cm.RdGy,
+                       vmin=None, vmax=None):
+    """
+    Plot the transient absorption map as a function of time (x-axis) and energy (y-axis).
+    Color represents the absorption intensity.
+
+    Parameters
+    ----------
+
+    ax : matplotlib.axes.Axes
+        Axis for plotting the transient absorption map.
+
+    time_grid : numpy.ndarray
+        Time grid for the transient absorption in fs.
+
+    energy_grid : numpy.ndarray
+        Energy grid for the transient absorption in eV.
+
+    trans_abs : numpy.ndarray
+        Transient absorption 2D array. Shape: (len(enery_grid), len(time_grid)).
+        Can be full, electron, or hole contributions.
+
+    num_contours : int
+        Number of contours to plot.
+
+    cmap : matplotlib.colors.Colormap
+        Colormap for the plot.
+
+    vmin : float
+        Minimum value for the transient absorption.
+
+    vmax : float
+        Maximum value for the transient absorption.
+
+    Returns
+    -------
+    ax : matplotlib.axes.Axes
+        Axis for plotting the transient absorption map.
+    """
+
+    time_mesh, energy_mesh = np.meshgrid(time_grid, energy_grid)
+
+    # Min and max values for plot
+    if vmin is None:
+        vmin = np.min(trans_abs)
+    if vmax is None:
+        vmax = np.max(trans_abs)
+
+    ax.contourf(time_mesh, energy_mesh, trans_abs, num_contours, cmap=cmap, vmin=vmin, vmax=vmax)
+
+    fsize = 24
+    ax.set_xlabel('Time delay (fs)', fontsize=fsize)
+    ax.set_ylabel('Energy (eV)', fontsize=fsize)
+
+    # Add colorbar as a separate axis
+    norm = Normalize(vmin=vmin, vmax=vmax)
+    divider = make_axes_locatable(ax)
+    cax = divider.append_axes('right', size='5%', pad=0.08)
+    cbar = plt.colorbar(ScalarMappable(cmap=cmap, norm=norm), cax=cax, orientation='vertical', format='%.1f')
+    cbar.set_label(r'$\Delta A$ (arb. units)', fontsize=fsize)
+
+    return ax