From 72f900eabb95dec6ef5aca3556be243c255e0c4f Mon Sep 17 00:00:00 2001 From: Ivan Maliyov Date: Mon, 16 Dec 2024 18:54:30 +0100 Subject: [PATCH] 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 | 2 +- src/perturbopy/postproc/__init__.py | 3 +- .../postproc/calc_modes/dyna_run.py | 148 +++++++++++- src/perturbopy/postproc/dbs/units_dict.py | 3 +- .../postproc/utils/spectra_generate_pulse.py | 25 +- .../postproc/utils/spectra_plots.py | 71 +++++- .../postproc/utils/spectra_trans_abs.py | 226 ++++++++++++++++++ 7 files changed, 462 insertions(+), 16 deletions(-) create mode 100644 src/perturbopy/postproc/utils/spectra_trans_abs.py diff --git a/setup.cfg b/setup.cfg index d83649a..617006d 100644 --- a/setup.cfg +++ b/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 diff --git a/src/perturbopy/postproc/__init__.py b/src/perturbopy/postproc/__init__.py index ed01368..67bdda2 100644 --- a/src/perturbopy/postproc/__init__.py +++ b/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 diff --git a/src/perturbopy/postproc/calc_modes/dyna_run.py b/src/perturbopy/postproc/calc_modes/dyna_run.py index 6567a43..36e1519 100644 --- a/src/perturbopy/postproc/calc_modes/dyna_run.py +++ b/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 diff --git a/src/perturbopy/postproc/dbs/units_dict.py b/src/perturbopy/postproc/dbs/units_dict.py index 825ba56..e2231f9 100644 --- a/src/perturbopy/postproc/dbs/units_dict.py +++ b/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 diff --git a/src/perturbopy/postproc/utils/spectra_generate_pulse.py b/src/perturbopy/postproc/utils/spectra_generate_pulse.py index 934cb2f..8bcf6b0 100644 --- a/src/perturbopy/postproc/utils/spectra_generate_pulse.py +++ b/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, diff --git a/src/perturbopy/postproc/utils/spectra_plots.py b/src/perturbopy/postproc/utils/spectra_plots.py index ac5b3f1..d707ac3 100644 --- a/src/perturbopy/postproc/utils/spectra_plots.py +++ b/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 diff --git a/src/perturbopy/postproc/utils/spectra_trans_abs.py b/src/perturbopy/postproc/utils/spectra_trans_abs.py new file mode 100644 index 0000000..763847f --- /dev/null +++ b/src/perturbopy/postproc/utils/spectra_trans_abs.py @@ -0,0 +1,226 @@ +""" +Utils for ultrafast spectroscopy: transient absorption calculations +based on electron-phonon and hole-phonon dynamics simulations. +""" + +import numpy as np + +from .memory import get_size +from .timing import TimingGroup +from .constants import energy_conversion_factor + + +def gaussian_delta(x, mu, sig): + """ + Gaussian delta-function normalized to 1. + + Parameters + ---------- + x : array + Array of x-values. + mu : float + Mean of the Gaussian. + sig : float + Standard deviation of the Gaussian. + """ + + return np.exp(-0.5 * ((x - mu) / sig)**2) / (sig * np.sqrt(2 * np.pi)) + + +def compute_trans_abs(elec_dyna_run, + hole_dyna_run, + de_grid=0.02, + eta=0.02, + save_npy=True, + ): + """ + Compute the transient absorption spectrum from the electron and hole dynamics simulations. + The data is saved in the current directory as numpy binary files (.npy). + + Parameters + ---------- + + elec_dyna_run : DynaRun + Object containing the results of the electron dynamics simulation. + + hole_dyna_run : DynaRun + Object containing the results of the hole dynamics simulation. + + de_grid : float + Energy grid spacing for the transient absorption spectrum. Affects the calculation time. + + eta : float + Broadening parameter for the Gaussian delta functions applied to the energy grid. + + save_npy : bool + Save the data as numpy binary files (.npy). + + Returns + ------- + + time_grid : np.ndarray + Time grid for the transient absorption spectrum. + + trans_abs_energy_grid : np.ndarray + Energy grid for the transient absorption spectrum. + + dA_elec : np.ndarray + Electron contribution to the transient absorption spectrum. Shape (num_energy_points, num_steps). + + dA_hole : np.ndarray + Hole contribution to the transient absorption spectrum. Shape (num_energy_points, num_steps). + """ + + trun = TimingGroup('trans. abs.') + trun.add('total', level=3).start() + + # Check the consistency of the two DynaRun objects + + # 1. Check the k-grids. q-grids can be different + if not np.alltrue(elec_dyna_run.boltz_kdim == hole_dyna_run.boltz_kdim): + raise ValueError('The k-grids of the electron and hole simulations are different.') + + # 2. Check the simulation times + if elec_dyna_run[1].num_steps != hole_dyna_run[1].num_steps: + raise ValueError('The number of time steps of the electron and hole simulations are different.') + if not np.allclose(elec_dyna_run[1].time_step, hole_dyna_run[1].time_step): + raise ValueError('The time steps of the electron and hole simulations are different.') + + # 3. Check if the HDF5 cdyna files are still open + if not elec_dyna_run._cdyna_file: + raise ValueError('The electron dynamics HDF5 file (prefix_cdyna.h5) is closed.') + if not hole_dyna_run._cdyna_file: + raise ValueError('The hole dynamics HDF5 file (prefix_cdyna.h5) is closed.') + + time_step = elec_dyna_run[1].time_step + num_steps = elec_dyna_run[1].num_steps + time_grid = np.arange(0, num_steps * time_step, time_step) + + # Energy arrays + elec_energy_array = elec_dyna_run._energies + hole_energy_array = hole_dyna_run._energies + + # K-point arrays + elec_kpoint_array = elec_dyna_run._kpoints + hole_kpoint_array = hole_dyna_run._kpoints + + # Convert from Ry to eV + elec_energy_array *= energy_conversion_factor('Ry', 'eV') + hole_energy_array *= energy_conversion_factor('Ry', 'eV') + + # Number of bands + elec_nband = len(elec_dyna_run.bands) + hole_nband = len(hole_dyna_run.bands) + + # Band quantities + VBM = np.max(hole_energy_array.ravel()) + CBM = np.min(elec_energy_array.ravel()) + bandgap = CBM - VBM + + energy_max = np.max(elec_energy_array.ravel()) + energy_min = np.min(hole_energy_array.ravel()) + + # Energy grid for the transient absorption spectrum, affect the calculation time + trans_abs_energy_grid = np.arange(bandgap, energy_max - energy_min, de_grid) + num_energy_points = trans_abs_energy_grid.shape[0] + print(f"{'Number of energy points':<30}: {num_energy_points}\n") + + # Find the interection between the electron and hole k-points + print('Finding intersect k-points...') + with trun.add('intersect kpoints') as t: + ekidx, hkidx = np.intersect1d( + elec_kpoint_array.view('float64,float64,float64').reshape(-1), + hole_kpoint_array.view('float64,float64,float64').reshape(-1), + return_indices=True)[1:] + + assert ekidx.size == hkidx.size, 'The number of intersect k-points is different.' + + num_intersect_kpoints = ekidx.size + + # Compute all the Gaussian deltas on the intersect grid + print('Computing Gaussian deltas...') + with trun.add('compute deltas') as t: + ALL_DELTAS = np.zeros((elec_nband, hole_nband, num_intersect_kpoints, num_energy_points)) + + # Here, we vectorize over the k-points (ekidx, hkidx), the most expensive operation + # Computation can still be vectorized for the energy points + for iband in range(elec_nband): + for jband in range(hole_nband): + for ienergy in range(num_energy_points): + # Factor of two from spin. + # TODO: add transient dipoles here + ALL_DELTAS[iband, jband, :, ienergy] = 2.0 * \ + gaussian_delta(elec_energy_array[ekidx, iband] - + hole_energy_array[hkidx, jband], + trans_abs_energy_grid[ienergy], eta) + + get_size(ALL_DELTAS, 'ALL_DELTAS', dump=True) + print('') + + # Compute the transient absorption spectrum first for electrons and holes separately + print('Computing transient absorption spectrum...') + with trun.add('compute trans. abs.') as t: + + # Store all the occupations at once in numpy arrays + # If the memory consumption is too high, one has to setup an explicit time loop for transient absorption + e_occs_time_array = np.zeros((num_steps, num_intersect_kpoints, elec_nband)) + h_occs_time_array = np.zeros((num_steps, num_intersect_kpoints, hole_nband)) + + # number of dynamics run + ndyna = 1 + for istep in range(num_steps): + + # HDF5 dataset names + elec_dset_name = f'dynamics_run_{ndyna}/snap_t_{istep}' + hole_dset_name = f'dynamics_run_{ndyna}/snap_t_{istep}' + + # Read the occupations + # We read them directly from the HDF5 files. Usually, this data is not stored in memory. + elec_occs = elec_dyna_run._cdyna_file[elec_dset_name][()] + # Minus sign for transient absorption, store only on the intersect grid, hence ekidx + elec_occs = -elec_occs[ekidx, :] + + hole_occs = 1.0 - hole_dyna_run._cdyna_file[hole_dset_name][()] + hole_occs = -hole_occs[hkidx, :] + + # Second index of elec_occs is band, one can select specific bands here + e_occs_time_array[istep, :, :] = elec_occs[:, :] + h_occs_time_array[istep, :, :] = hole_occs[:, :] + + get_size(e_occs_time_array, 'e_occs_time_array', dump=True) + get_size(h_occs_time_array, 'h_occs_time_array', dump=True) + print('') + + # Compute the transient absorption spectrum + dA_elec = np.zeros((num_energy_points, num_steps)) + dA_hole = np.zeros((num_energy_points, num_steps)) + for iband in range(elec_nband): + for jband in range(hole_nband): + + # Sum over intersect k-points + dA_elec[:, :] += (np.tensordot(e_occs_time_array[:, :, iband], + ALL_DELTAS[iband, jband, :, :], axes=((1), (0))) / num_intersect_kpoints).T + + dA_hole[:, :] += (np.tensordot(h_occs_time_array[:, :, jband], + ALL_DELTAS[iband, jband, :, :], axes=((1), (0))) / num_intersect_kpoints).T + + # Save the data + print('Saving the data...') + with trun.add('save data') as t: + pump_energy = elec_dyna_run.pump_pulse.pump_energy + ending = f'_Epump_{pump_energy:.3f}' + # Total transient absorption + np.save(f'trans_abs_dA{ending}', dA_elec + dA_hole) + # Electron contribution + np.save(f'trans_abs_dA_elec{ending}', dA_elec) + # Hole contribution + np.save(f'trans_abs_dA_hole{ending}', dA_hole) + # Time grid + np.save(f'trans_abs_T{ending}', time_grid) + # Energy grid + np.save(f'trans_abs_E{ending}', trans_abs_energy_grid) + + trun.timings['total'].stop() + print(trun) + + return time_grid, trans_abs_energy_grid, dA_elec, dA_hole