diff --git a/setup.cfg b/setup.cfg index 617006d..73a5e3b 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, W504 +ignore = E114, E127, E124, E221, W293, E721, W503, F841, F401, E202, W504, E731 diff --git a/src/perturbopy/postproc/__init__.py b/src/perturbopy/postproc/__init__.py index fb521c0..56eb6d2 100644 --- a/src/perturbopy/postproc/__init__.py +++ b/src/perturbopy/postproc/__init__.py @@ -12,7 +12,7 @@ from .calc_modes.trans import Trans from .calc_modes.imsigma import Imsigma from .calc_modes.imsigma_spin import ImsigmaSpin -from .calc_modes.dyna_run import DynaRun +from .calc_modes.dyna_run import DynaRun, PumpPulse from .calc_modes.dyna_pp import DynaPP from .dbs.units_dict import UnitsDict diff --git a/src/perturbopy/postproc/calc_modes/dyna_run.py b/src/perturbopy/postproc/calc_modes/dyna_run.py index 36e1519..06e55d2 100644 --- a/src/perturbopy/postproc/calc_modes/dyna_run.py +++ b/src/perturbopy/postproc/calc_modes/dyna_run.py @@ -1,5 +1,6 @@ -import numpy as np import os +import numpy as np +import warnings from perturbopy.postproc.calc_modes.calc_mode import CalcMode from perturbopy.postproc.dbs.recip_pt_db import RecipPtDB from perturbopy.postproc.calc_modes.dyna_indiv_run import DynaIndivRun @@ -7,11 +8,21 @@ from perturbopy.postproc.utils.timing import Timing, TimingGroup from perturbopy.postproc.dbs.units_dict import UnitsDict +# for plotting +import matplotlib.pyplot as plt +from perturbopy.postproc.utils.spectra_plots import find_fwhm +from perturbopy.postproc.utils.plot_tools import plotparams +plt.rcParams.update(plotparams) + 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 @@ -33,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': @@ -95,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(): @@ -109,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 @@ -141,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): """ @@ -197,15 +225,24 @@ 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" + + text += f"{'Read snaps at init':>30}: {self.read_snaps}\n" + text += f"{'Number of k-points':>30}: {self.kpt.points_cart.shape[1]}\n" + nband = self._pert_dict['input parameters']['after conversion']['band_max'] - \ + self._pert_dict['input parameters']['after conversion']['band_min'] + 1 + text += f"{'Number of bands':>30}: {nband}\n" + text += f"{'Hole':>30}: {self._pert_dict['input parameters']['after conversion']['hole']}\n" + text += "\n" + + text += f"{'Number of runs':>30}: {self.num_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" + text += f"{'Electric field (V/cm)':>30}: {dynamics_run.efield}\n" return text @@ -245,6 +282,117 @@ def extract_steady_drift_vel(self, dyna_pp_yaml_path): return steady_drift_vel, steady_conc + @staticmethod + def to_cdyna_h5(prefix, band_structure_ryd, snap_array, time_step_fs, + path='.', new=True, overwrite=False, num_runs=1): + """ + Write the dynamics data into the prefix_cdyna.h5 HDF5 file. + Follows the script format of the Perturbo dynamics-run calculation. + Currently implemented as a static method. The idea is that if we need to write an HDF5 cdyna + file, it is forcfully different from an existing one, therefore, one or more + attributes are different from the original cdyna file. + + Parameters + ---------- + + prefix : str + Prefix of the HDF5 file. Same as prefix in any Perturbo calculation. + + band_structure_ryd : np.ndarray + Band structure in Rydberg units. Shape: num_kpoints x num_bands + + snap_array : np.ndarray + Array of carrier occupations. Shape: (num_steps, num_kpoints, num_bands) + + time_step_fs : float or np.ndarray + Time step in fs. If float, the same time step is used for all runs. + If np.ndarray, the time steps for each run are specified. + + path : str + Path to the directory where the HDF5 file is saved. Default is the current directory. + + new : bool + Flag to indicate if the file is new. Default is True. + + overwrite : bool + Flag to indicate if the file is overwritten. Default is False. + + num_runs : int + Number of dynamics_run runs. Default is 1. + """ + + trun = TimingGroup('write cdyna') + trun.add('total', level=3).start() + + if not new: + raise NotImplementedError("Only new files are supported") + if num_runs != 1: + raise NotImplementedError("Only num_runs=1 is supported") + + if isinstance(time_step_fs, (float, np.floating)): + time_step_fs = np.array([time_step_fs]) + if time_step_fs.shape[0] != num_runs: + raise ValueError("The number of time steps must be equal to num_runs") + + if path != '.': + os.makedirs(path, exist_ok=True) + + new_cdyna_filename = os.path.join(path, f'{prefix}_cdyna.h5') + + if os.path.isfile(new_cdyna_filename): + if overwrite: + warnings.warn(f'File {new_cdyna_filename} already exists. Overwriting it.') + else: + raise FileExistsError(f'File {new_cdyna_filename} already exists. Set overwrite=True to overwrite it.') + + if snap_array.ndim != 3: + raise ValueError("snap_array must be 3D. Shape: (num_steps, num_kpoints, num_bands)") + else: + num_steps_read, num_kpoints_read, num_bands_read = snap_array.shape + print("Occupations shape to write:") + print(f"num_steps: {num_steps_read}, num_kpoints: {num_kpoints_read}, num_bands: {num_bands_read}") + + if num_kpoints_read < num_steps_read or num_kpoints_read < num_bands_read: + warnings.warn("The number of k-points of snap array to write is smaller " + "than the number of time steps or bands.") + + new_cdyna_file = open_hdf5(new_cdyna_filename, 'w') + + new_cdyna_file.create_dataset('band_structure_ryd', data=band_structure_ryd) + new_cdyna_file['band_structure_ryd'].attrs['ryd2ev'] = 13.605698066 + + new_cdyna_file.create_dataset('num_runs', data=num_runs) + + for irun in range(1, num_runs + 1): + + # In Perturbo, in dynamics_run_1, snap_t_ start from 0, where snap_t_0 is the initial state + # In dynamics_run_2 and onwards, snap_t_ start from + if irun == 1: + offset = 0 + + if snap_array.shape[0] == 1: + # duplicate the initial state + snap_array = np.vstack((snap_array, snap_array)) + num_steps = snap_array.shape[0] - 1 + + else: + offset = 1 + num_steps = snap_array.shape[0] + + dyn_str = f'dynamics_run_{irun}' + new_cdyna_file.create_group(dyn_str) + new_cdyna_file[dyn_str].create_dataset('num_steps', data=num_steps) + new_cdyna_file[dyn_str].create_dataset('time_step_fs', data=time_step_fs[irun - 1]) + + for itime in range(snap_array.shape[0]): + new_cdyna_file[dyn_str].create_dataset(f'snap_t_{itime + offset}', data=snap_array[itime, :, :]) + + close_hdf5(new_cdyna_file) + + trun.timings['total'].stop() + print(f'{"Total time to write cdyna":>30}: {trun.timings["total"].total_runtime} s') + print() + class PumpPulse(): """ @@ -256,15 +404,24 @@ class PumpPulse(): pump_energy : float Energy of the pump pulse excitation in eV. - energy_broadening : float - Energy broadening of the pump pulse excitation in eV. + spectral_width_fwhm : float + Energy broadening FWHM of the pump pulse excitation in eV. + + pump_duration_fwhm : float + Duration FWHM of the pump pulse excitation in fs. num_steps : int Number of steps in the pump pulse excitation. - time_step : float + pump_factor : float + Factor for the pump pulse excitation. + + pump_time_step : float Time step of the pump pulse excitation in fs. + num_bands : int + Number of bands in the pump pulse excitation. Tailored for the Perturbo dynamics-run calculation. + num_kpoints : int Number of k-points in the pump pulse excitation. Tailored for the Perturbo dynamics-run calculation. @@ -278,6 +435,12 @@ class PumpPulse(): hole : bool Flag to indicate if the pump pulse excitation is for the hole. Must be the same as in the ultrafast simulation. + + time_profile : np.ndarray + Array of time profile for the pump pulse excitation. First column is time in fs, second column is the time profile. + + energy_profile : np.ndarray + Array of energy profile for the pump pulse excitation. First column is energy in eV, second column is the energy profile """ def __init__(self, pump_dict): @@ -298,14 +461,20 @@ def __init__(self, pump_dict): 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.spectral_width_fwhm = pump_dict['pump_spectral_width_fwhm'] + self.spectral_width_fwhm_units = pump_dict['pump_spectral_width_fwhm units'] + + self.pump_duration_fwhm = pump_dict['pump_duration_fwhm'] + self.pump_duration_fwhm_units = pump_dict['pump_duration_fwhm 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.pump_factor = pump_dict['pump_factor'] + + self.pump_time_step = pump_dict['pump_time_step'] + self.pump_time_step_units = pump_dict['pump_time_step units'] + self.num_bands = pump_dict['num_bands'] self.num_kpoints = pump_dict['num_kpoints'] self.carrier_number_array = \ @@ -318,16 +487,75 @@ def __init__(self, pump_dict): self.hole = pump_dict['hole'] + # arrays of time and energy profile + self.time_profile = np.array(pump_dict['time_profile']) + self.energy_profile = np.array(pump_dict['energy_profile']) + 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" + carrier = 'hole' if self.hole else 'electron' + title = f'Pump pulse parameters ({carrier})' + text = f'{title:*^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" + text += f"{'Pump factor':>30}: {self.pump_factor:.3e}\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"{f'Time step ({self.pump_time_step_units})':>30}: {self.pump_time_step}\n" text += f"{'Hole':>30}: {self.hole}\n" return text + + def plot_time_profile(self, ax=None): + """ + Plot the pump pulse time profile. + """ + + if self.time_profile is None: + warnings.warn('No time profile found for the pump pulse') + return None + + if ax is None: + fig, ax = plt.subplots(1, 1, figsize=(10, 8)) + + # Find the FWHM and half-maximum of the time profile + time_left_FWHM, time_right_FWHM, time_half_max = find_fwhm(self.time_profile[:, 0], self.time_profile[:, 1]) + + ax.plot(self.time_profile[:, 0], self.time_profile[:, 1]) + ax.plot([time_left_FWHM, time_right_FWHM], [time_half_max, time_half_max], marker='o', color='tab:red', lw=3, label='FWHM') + ax.set_xlabel('Time (fs)', fontsize=24) + ax.set_ylabel('Time Gaussian', fontsize=24) + ax.set_title(f'Time profile (FWHM = {time_right_FWHM - time_left_FWHM:.4f} fs)') + ax.legend() + ax.grid() + + return ax + + def plot_energy_profile(self, ax=None): + """ + Plot the pump pulse energy profile. + """ + + if self.energy_profile is None: + warnings.warn('No energy profile found for the pump pulse') + return None + + if ax is None: + fig, ax = plt.subplots(1, 1, figsize=(10, 8)) + + # Find the FWHM and half-maximum of the energy profile + energy_left_FWHM, energy_right_FWHM, energy_half_max = find_fwhm(self.energy_profile[:, 0], self.energy_profile[:, 1]) + + ax.plot(self.energy_profile[:, 0], self.energy_profile[:, 1]) + ax.plot([energy_left_FWHM, energy_right_FWHM], [energy_half_max, energy_half_max], marker='o', color='tab:red', lw=3, label='FWHM') + ax.axvline(self.pump_energy, color='gray', lw=3, label='Pump energy') + ax.set_xlabel('Energy (eV)', fontsize=24) + ax.set_ylabel('Energy Gaussian', fontsize=24) + ax.set_title(f'Energy profile (FWHM = {energy_right_FWHM - energy_left_FWHM:.4f} eV)') + ax.legend(loc='upper right') + ax.grid() + + return ax diff --git a/src/perturbopy/postproc/dbs/units_dict.py b/src/perturbopy/postproc/dbs/units_dict.py index e2231f9..7b97d75 100644 --- a/src/perturbopy/postproc/dbs/units_dict.py +++ b/src/perturbopy/postproc/dbs/units_dict.py @@ -4,7 +4,8 @@ class UnitsDict(dict): """ - This is a class representation of a set of physical quantities with units organized in a dictionary. + Modified dictionary class to store physical quantities with units. + Class representation of a set of physical quantities with units organized in a dictionary. The physical quantities may either be arrays or floats. Attributes diff --git a/src/perturbopy/postproc/utils/spectra_generate_pulse.py b/src/perturbopy/postproc/utils/spectra_generate_pulse.py index 8bcf6b0..265d6c0 100644 --- a/src/perturbopy/postproc/utils/spectra_generate_pulse.py +++ b/src/perturbopy/postproc/utils/spectra_generate_pulse.py @@ -11,7 +11,7 @@ from perturbopy.io_utils.io import open_hdf5, close_hdf5 -from matplotlib.animation import FuncAnimation +from perturbopy.postproc import PumpPulse from .memory import get_size from .timing import TimingGroup @@ -20,19 +20,13 @@ from . import spectra_plots -def gaussian(x, mu, sig, hole_nband, elec_nband): - """ - Gaussian function normalized to unity max occupation - """ - - return np.exp(-0.5 * ((x - mu) / sig)**2) / (hole_nband * elec_nband) - - def sigma_from_fwhm(fwhm): """ Comupte Gaussian sigma from the Full Width at Half Maximum (FWHM) parameter: .. math:: f(t) = A / ( sigma * sqrt(2pi)) exp(-(t-t0)^2 / (2sigma^2)) + + For fwhm = 1, sigma = 0.42466. """ sigma = 1.0 / (2.0 * np.sqrt(2.0 * np.log(2.0))) * fwhm @@ -40,11 +34,22 @@ def sigma_from_fwhm(fwhm): return sigma +def fwhm_from_sigma(sigma): + """ + Comupte FWHM from Gaussian sigma. + """ + + fwhm = 2.0 * np.sqrt(2.0 * np.log(2.0)) * sigma + + return fwhm + + def delta_occs_pulse_coef(t, dt, tw, sigma): """ Additional occupation due to the pulse excitation. - Assuming the Gaussian pulse shape, the occupation is increasing in - time according to the error function. + To compute precisely the additional occupation for dt, + we use the integral of the Gaussian function from t - dt to t, + which is the difference of two error functions. Parameters ---------- @@ -58,7 +63,7 @@ def delta_occs_pulse_coef(t, dt, tw, sigma): Time window during which the pump pulse will be applied. sigma : float - Gaussian sigma broadening. + Gaussian sigma for pulse duration. Returns ------- @@ -74,7 +79,7 @@ def delta_occs_pulse_coef(t, dt, tw, sigma): return 0.5 * numer / denom -def gaussian_excitation(pump_file, occs_amplitude, time_grid, time_window, pulse_fwhm, time_step, +def gaussian_excitation(pump_file, occs_amplitude, time_grid, time_window, pump_duration, time_step, hole=False, finite_width=True): """ The Gaussian pulse excitation for the Perturbo dynamics. @@ -96,7 +101,7 @@ def gaussian_excitation(pump_file, occs_amplitude, time_grid, time_window, pulse time_window : float The time window for the pump pulse (fs). During this time window, the pump pulse will be active. - pulse_fwhm : float + pump_duration : float The full width at half maximum of the pump pulse (fs). time_step : float @@ -127,7 +132,7 @@ def gaussian_excitation(pump_file, occs_amplitude, time_grid, time_window, pulse num_steps = time_grid.shape[0] # Compute the Gaussian sigma from the full width at half maximum (FWHM) - sigma = sigma_from_fwhm(pulse_fwhm) + sigma = sigma_from_fwhm(pump_duration) # Number of k-points and bands num_k, num_bands = occs_amplitude.shape @@ -143,10 +148,10 @@ def gaussian_excitation(pump_file, occs_amplitude, time_grid, time_window, pulse delta_occs_array = np.zeros((num_bands, num_steps, num_k)) # A pre-factor for the pump pulse, num_steps steps - pulse_coef = delta_occs_pulse_coef(time_grid[:], time_step, time_window, sigma) + pump_time_profile = delta_occs_pulse_coef(time_grid[:], time_step, time_window, sigma) # t - time, j - band, k - k-point - delta_occs_array[:, :, :] = np.einsum('t,jk->ktj', pulse_coef[:], occs_amplitude[:, :]) + delta_occs_array[:, :, :] = np.einsum('t,jk->ktj', pump_time_profile[:], occs_amplitude[:, :]) # Print the size of the array since it can be large: num_bands * num_steps * num_k get_size(delta_occs_array, f'delta(fnk(t)) {carrier}', dump=True) @@ -173,18 +178,25 @@ def setup_pump_pulse(elec_pump_pulse_path, hole_pump_pulse_path, elec_dyna_run, hole_dyna_run, pump_energy, pump_time_step=1.0, - pump_fwhm=20.0, - pump_energy_broadening=0.040, + pump_duration_fwhm=20.0, + pump_spectral_width_fwhm=0.090, pump_time_window=50.0, + pump_factor=0.3, finite_width=True, - animate=True - ): + animate=True, + plot_scale=1.0, + cnum_check=True, + tr_dipoles_sqr=None): """ Setup the Gaussian pump pulse excitation for electrons and holes. Write into the pump_pulse.h5 HDF5 file. We use raw k-point and energy arrays as read from the HDF5 files for efficiency. All energies in eV, k points in crystal coordinates. + Pulse duration and energy broadening FWHM are related, typically, + pump_spectral_width_fwhm = 1.8 / pump_duration_fwhm. + However, here, we leave the possibility to set them independently. + Parameters ---------- elec_pump_pulse_path : str @@ -208,11 +220,11 @@ def setup_pump_pulse(elec_pump_pulse_path, hole_pump_pulse_path, Time step in fs for the pump pulse generation. Note: the pump_time_step MUST match the one used in the dynamics run in Perturbo! - pump_fwhm : float, optional + pump_duration_fwhm : float, optional The full width at half maximum of the pump pulse (fs). - pump_energy_broadening : float, optional - Energy broadening of the pump pulse (eV). + pump_spectral_width_fwhm : float, optional + Energy broadening FWHM of the pump pulse (eV). pump_time_window : float, optional The time window for the pump pulse (fs). @@ -220,20 +232,30 @@ def setup_pump_pulse(elec_pump_pulse_path, hole_pump_pulse_path, The number of the snapshots in the pump_pulse.h5 file will be equal to pump_time_window / pump_time_step. + pump_factor : float, optional + The amplitude of the pump pulse. The maximum occupation will be pump_factor. + finite_width : bool, optional If True, the pulse is finite in time. If False, the pulse is a step function and occs_amplitude will be set as initial occupation. - """ - print(f"{'PUMP PULSE PARAMETERS':*^70}") - print(f"{'Pump pulse energy (eV):':>40} {pump_energy:.3f}") - print(f"{'Pump pulse energy broadening (eV):':>40} {pump_energy_broadening:.4f}") - print(f"{'Finite width:':>40} {finite_width}") - if finite_width: - print(f"{'Pump pulse time step (fs):':>40} {pump_time_step:.3f}") - print(f"{'Pump pulse FWHM (fs):':>40} {pump_fwhm:.3f}") - print(f"{'Pump pulse time window (fs):':>40} {pump_time_window:.3f}") - print("") + animate : bool, optional + If True, animate the pump pulse excitation. + + plot_scale : float + Scale factor for the scatter object sizes for carrier occupations. + Default is 1.0. + + cnum_check : bool, optional + If True, create a cnum_check folder with a prefix_cdyna.h5 file with the total + occupation summed over time steps. Run dynamics-pp with Perturbo (without dynamics-run) + to get the accurate total carrier number in the cnum_check folder (link the epr file there). + + tr_dipoles_sqr: np.ndarray, optional + Transition dipoles squared, valence-to-conduction for each k-point and band. + Currently, the k-grid for dipoles must match the one for electrons. + Experimental feature. + """ # Check electron and hole time steps elec_dyna_time_step = elec_dyna_run[1].time_step @@ -280,12 +302,20 @@ def setup_pump_pulse(elec_pump_pulse_path, hole_pump_pulse_path, print(f"{'Band gap (Eg) (eV):':>40} {bandgap:.4f}") print("") + # Currently, the number of k-points of the transition dipoles + # must be the same as the number of k-points of the electron energy array + # In a future development, we will interpolate the transition dipoles on the fly + if tr_dipoles_sqr is not None: + elec_nbard_tr_dip, hole_nband_tr_dip, num_k_tr_dip = tr_dipoles_sqr.shape + + if elec_nbard_tr_dip != elec_nband or hole_nband_tr_dip != hole_nband \ + or num_k_tr_dip != elec_kpoint_array.shape[0]: + raise ValueError(f'Wrong shape of the transition dipoles array: {tr_dipoles_sqr.shape}\n' + f'Expected shape: ({elec_nband}, {hole_nband}, {elec_kpoint_array.shape[0]})') + elec_occs_amplitude = np.zeros_like(elec_energy_array) hole_occs_amplitude = np.zeros_like(hole_energy_array) - # A pre-factor for the pump pulse - pump_factor = 1.0 - print('Computing the optically excited occupations...') t_occs = TimingGroup('Setup occupations') t_occs.add('total', level=3).start() @@ -300,6 +330,9 @@ def setup_pump_pulse(elec_pump_pulse_path, hole_pump_pulse_path, hole_kpoint_array.view('float64,float64,float64').reshape(-1), return_indices=True)[1:] + # Convert the pump energy FWHM to sigma + pump_energy_broadening_sigma = sigma_from_fwhm(pump_spectral_width_fwhm) + # Iteratre over electron and hole bands # Even though, the for loops are inefficient in Python, # the number of bands is rarely more than 10, therefore @@ -308,8 +341,14 @@ def setup_pump_pulse(elec_pump_pulse_path, hole_pump_pulse_path, for jband in range(hole_nband): # diff. between elec and hole for a given elec branch iband and hole branch jband delta = pump_factor * \ - gaussian(elec_energy_array[ekidx, iband] - hole_energy_array[hkidx, jband], - pump_energy, pump_energy_broadening, hole_nband, elec_nband) + spectra_plots.gaussian(elec_energy_array[ekidx, iband] - hole_energy_array[hkidx, jband], + pump_energy, pump_energy_broadening_sigma) + + # Multiply by the transition dipoles squared if provided + # Transition dipoles are assumed to be on the same k-point grid + # as the electron energy array + if tr_dipoles_sqr is not None: + delta *= tr_dipoles_sqr[iband, jband, ekidx] # Only for the intersected k points, we add the delta occupation elec_occs_amplitude[ekidx, iband] += delta @@ -327,6 +366,11 @@ def setup_pump_pulse(elec_pump_pulse_path, hole_pump_pulse_path, print(f"{'Difference:':>40} {sum(elec_occs_amplitude.ravel()) - sum(hole_occs_amplitude.ravel()):.4e}") print("") + # Create the energy profile, only for reference + en_grid = np.linspace(pump_energy - 3 * pump_energy_broadening_sigma, + pump_energy + 3 * pump_energy_broadening_sigma, 200) + en_profile = spectra_plots.gaussian(en_grid, pump_energy, pump_energy_broadening_sigma) + # Write to HDF5 file, replacing snap_t_1 # Electron if os.path.exists(elec_pump_pulse_path): @@ -341,11 +385,7 @@ def setup_pump_pulse(elec_pump_pulse_path, hole_pump_pulse_path, 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 + optional_params = np.ones(10, dtype=float) * (-9999) for h5f in [elec_pump_pulse_file, hole_pump_pulse_file]: h5f.create_group('pump_pulse_snaps') @@ -353,27 +393,37 @@ def setup_pump_pulse(elec_pump_pulse_path, hole_pump_pulse_path, h5f.create_dataset('pump_energy', data=pump_energy) h5f['pump_energy'].attrs['units'] = 'eV' - h5f.create_dataset('pump_energy_broadening', data=pump_energy_broadening) - h5f['pump_energy_broadening'].attrs['units'] = 'eV' + h5f.create_dataset('pump_spectral_width_fwhm', data=pump_spectral_width_fwhm) + h5f['pump_spectral_width_fwhm'].attrs['units'] = 'eV' - h5f['optional_params'] = optional_params + h5f.create_dataset('pump_factor', data=pump_factor) + h5f.create_dataset('optional_params', data=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}') + h5f.create_dataset('pump_time_step', data=pump_time_step) + h5f.create_dataset('pump_duration_fwhm', data=pump_duration_fwhm) + + pump_time_profile = delta_occs_pulse_coef(time_grid[:], pump_time_step, + pump_time_window, sigma_from_fwhm(pump_duration_fwhm)) + h5f.create_dataset('time_profile', data=np.c_[time_grid, pump_time_profile]) else: + pump_time_profile = None 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.create_dataset('pump_time_step', data=elec_dyna_time_step) + h5f.create_dataset('pump_duration_fwhm', data=elec_dyna_time_step) + h5f.create_dataset('time_profile', data=np.c_[0, 1]) + h5f['pump_duration_fwhm'].attrs['units'] = 'fs' h5f['time_window'].attrs['units'] = 'fs' - h5f['time_step'].attrs['units'] = 'fs' + h5f['pump_time_step'].attrs['units'] = 'fs' + + h5f.create_dataset('energy_profile', data=np.c_[en_grid, en_profile]) + h5f['energy_profile'].attrs['units'] = 'eV' + h5f['time_profile'].attrs['units'] = 'fs' elec_pump_pulse_file.create_dataset('num_bands', data=elec_nband) hole_pump_pulse_file.create_dataset('num_bands', data=hole_nband) @@ -383,21 +433,17 @@ def setup_pump_pulse(elec_pump_pulse_path, hole_pump_pulse_path, 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, time_grid, - pump_time_window, pump_fwhm, pump_time_step, + pump_time_window, pump_duration_fwhm, pump_time_step, hole=False, finite_width=finite_width) hole_delta_occs_array = \ gaussian_excitation(hole_pump_pulse_file, hole_occs_amplitude, time_grid, - pump_time_window, pump_fwhm, pump_time_step, + pump_time_window, pump_duration_fwhm, pump_time_step, hole=True, finite_width=finite_width) # Close @@ -407,11 +453,76 @@ def setup_pump_pulse(elec_pump_pulse_path, hole_pump_pulse_path, print('\nPlotting...') spectra_plots.plot_occ_ampl(elec_occs_amplitude, elec_kpoint_array, elec_energy_array, hole_occs_amplitude, hole_kpoint_array, - hole_energy_array, pump_energy) + hole_energy_array, pump_energy, plot_scale=plot_scale) if animate and finite_width: print('\nAnimating...') spectra_plots.animate_pump_pulse(pump_time_step, elec_delta_occs_array, elec_kpoint_array, elec_energy_array, hole_delta_occs_array, hole_kpoint_array, hole_energy_array, - pump_energy) + pump_energy, plot_scale=plot_scale) + + # Setup the pump pulse object + pump_dict = { + 'pump_energy': pump_energy, + 'pump_energy units': 'eV', + 'pump_spectral_width_fwhm': pump_spectral_width_fwhm, + 'pump_spectral_width_fwhm units': 'eV', + 'pump_duration_fwhm': pump_duration_fwhm, + 'pump_duration_fwhm units': 'fs', + 'num_steps': num_steps, + 'pump_factor': pump_factor, + 'pump_time_step': pump_time_step, + 'pump_time_step units': 'fs', + 'num_bands': elec_nband, + 'num_kpoints': elec_kpoint_array.shape[0], + 'optional_params': optional_params, + 'hole': None, + 'time_profile': np.c_[time_grid, pump_time_profile], + 'energy_profile': np.c_[en_grid, en_profile], + # These quantities will be computed at the perturbo.x stage + 'pump pulse carrier number': None, + 'carrier_number units': None, + } + + pump_pulse = PumpPulse(pump_dict) + + # Create the cnum_check folder + if cnum_check: + + if finite_width: + total_occ_elec = np.sum(elec_delta_occs_array, axis=1) + total_occ_hole = np.sum(hole_delta_occs_array, axis=1) + + else: + total_occ_elec = elec_occs_amplitude + total_occ_hole = hole_occs_amplitude + + elec_cnum_check_path = os.path.join(os.path.dirname(elec_pump_pulse_path), 'cnum_check') + hole_cnum_check_path = os.path.join(os.path.dirname(hole_pump_pulse_path), 'cnum_check') + + # Format the occupations for cdyna.h5 writing + total_occ_elec_write = total_occ_elec.copy() + # Transpose: num_band, num_k -> num_k, num_band + total_occ_elec_write = total_occ_elec_write.T + # Add the dummy time axis: 1, num_k, num_band + total_occ_elec_write = np.expand_dims(total_occ_elec_write, axis=0) + elec_dyna_run.to_cdyna_h5(elec_dyna_run.prefix, + elec_dyna_run._energies, + total_occ_elec_write, + elec_dyna_run[1].time_step, + path=elec_cnum_check_path, + overwrite=True) + + total_occ_hole_write = total_occ_hole.copy() + total_occ_hole_write = total_occ_hole_write.T + total_occ_hole_write = np.expand_dims(total_occ_hole_write, axis=0) + total_occ_hole_write = 1.0 - total_occ_hole_write + hole_dyna_run.to_cdyna_h5(hole_dyna_run.prefix, + hole_dyna_run._energies, + total_occ_hole_write, + hole_dyna_run[1].time_step, + path=hole_cnum_check_path, + overwrite=True) + + return pump_pulse diff --git a/src/perturbopy/postproc/utils/spectra_plots.py b/src/perturbopy/postproc/utils/spectra_plots.py index d707ac3..ec91871 100644 --- a/src/perturbopy/postproc/utils/spectra_plots.py +++ b/src/perturbopy/postproc/utils/spectra_plots.py @@ -7,7 +7,8 @@ import numpy as np import matplotlib.pyplot as plt import warnings -from scipy import special +from scipy.interpolate import Akima1DInterpolator +from scipy.optimize import brentq from perturbopy.io_utils.io import open_hdf5, close_hdf5 @@ -22,11 +23,85 @@ # Plotting parameters from .plot_tools import plotparams + plt.rcParams.update(plotparams) +def gaussian(x, mu, sigma): + """ + Gaussian function. Not normalized. + """ + + return np.exp(-0.5 * ((x - mu) / sigma)**2) + + +def find_fwhm(x, y, num_interp_points=2000): + """ + Find the Full Width at Half Maximum (FWHM) for a single prominent peak y(x), + using an Akima1DInterpolator to create a smooth spline, and root-finding + (brentq) for a precise half-max crossing. + + Parameters + ---------- + x : array_like + 1D array of x-values (assumed sorted in ascending order). + y : array_like + 1D array of y-values corresponding to x. + num_interp_points : int, optional + Number of points for evaluating the Akima spline; default is 2000. + + Returns + ------- + x_left : float + x-value at the left FWHM crossing. + + x_right : float + x-value at the right FWHM crossing. + + half_max : float + Half-maximum value of the peak. + """ + + # Create an Akima spline over the original data + akima = Akima1DInterpolator(x, y) + + # Evaluate on a fine grid to locate the approximate peak + x_fine = np.linspace(x[0], x[-1], num_interp_points) + y_fine = akima(x_fine) + + # 1. Find approximate peak index on the fine grid + i_max = np.argmax(y_fine) + x_max = x_fine[i_max] + y_max = y_fine[i_max] + half_max = y_max / 2.0 + + # f(x) = akima(x) - half_max + func = lambda xx: akima(xx) - half_max + + # --- Locate x_left by scanning left from the peak for sign change --- + x_left = x[0] # fallback in case no crossing is found + for i in range(i_max, 0, -1): + if y_fine[i - 1] - half_max < 0 < y_fine[i] - half_max or \ + y_fine[i - 1] - half_max > 0 > y_fine[i] - half_max: + # We found a bracket where f() changes sign + x_left = brentq(func, x_fine[i - 1], x_fine[i]) + break + + # --- Locate x_right by scanning right from the peak for sign change --- + x_right = x[-1] # fallback in case no crossing is found + for i in range(i_max, len(x_fine) - 1): + if y_fine[i] - half_max < 0 < y_fine[i + 1] - half_max or \ + y_fine[i] - half_max > 0 > y_fine[i + 1] - half_max: + # We found a bracket where f() changes sign + x_right = brentq(func, x_fine[i], x_fine[i + 1]) + break + + # The y-values at x_left and x_right are both half_max + return x_left, x_right, half_max + + def plot_occ_ampl(e_occs, elec_kpoint_array, elec_energy_array, - h_occs, hole_kpoint_array, hole_energy_array, pump_energy): + h_occs, hole_kpoint_array, hole_energy_array, pump_energy, plot_scale=1e3): """ Plot occupation amplitude. Currently hardcoded to the section kz = 0. @@ -52,6 +127,9 @@ def plot_occ_ampl(e_occs, elec_kpoint_array, elec_energy_array, pump_energy: float Pump energy in eV. + + plot_scale : float + Scale factor for the scatter object sizes. """ # find where kz == 0 @@ -60,15 +138,17 @@ def plot_occ_ampl(e_occs, elec_kpoint_array, elec_energy_array, # Plot electron occupations idx = np.where((elec_kpoint_array[:, 1] == 0) & (elec_kpoint_array[:, 0] == 0)) + e_occs_max = np.max(e_occs[idx, :].ravel()) for i in range(np.size(elec_energy_array, axis=1)): ax.scatter(elec_kpoint_array[idx, 2], elec_energy_array[idx, i], s=0.5, c='black', alpha=0.5) - ax.scatter(elec_kpoint_array[idx, 2], elec_energy_array[idx, i], s=e_occs[idx, i][0] * 1000 + 1e-10, c='red', alpha=0.5) + ax.scatter(elec_kpoint_array[idx, 2], elec_energy_array[idx, i], s=e_occs[idx, i][0] / e_occs_max * plot_scale + 1e-10, c='red', alpha=0.5) # Plot hole occupations idx = np.where((hole_kpoint_array[:, 1] == 0) & (hole_kpoint_array[:, 0] == 0)) + h_occs_max = np.max(h_occs[idx, :].ravel()) for i in range(np.size(hole_energy_array, axis=1)): ax.scatter(hole_kpoint_array[idx, 2], hole_energy_array[idx, i], s=0.5, c='black', alpha=0.5) - ax.scatter(hole_kpoint_array[idx, 2], hole_energy_array[idx, i], s=h_occs[idx, i][0] * 1000 + 1e-10, c='red', + ax.scatter(hole_kpoint_array[idx, 2], hole_energy_array[idx, i], s=h_occs[idx, i][0] / h_occs_max * plot_scale + 1e-10, c='red', alpha=0.5) fsize = 16 @@ -82,7 +162,8 @@ def plot_occ_ampl(e_occs, elec_kpoint_array, elec_energy_array, def animate_pump_pulse(time_step, elec_delta_occs_array, elec_kpoint_array, elec_energy_array, hole_delta_occs_array, hole_kpoint_array, hole_energy_array, - pump_energy): + pump_energy, + plot_scale=1e3): """ Animate the pump pulse excitation for electrons and holes. Defines fig and ax, initializes scatter objects for electron and hole occupations, and calls update_scatter. @@ -113,6 +194,9 @@ def animate_pump_pulse(time_step, pump_energy: float Pump energy in eV, used only in title. + + plot_scale : float + Scale factor for the scatter object sizes. """ fig, ax = plt.subplots(1, 1, figsize=(9, 6)) @@ -142,7 +226,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, hole_delta_occs_array), + elec_delta_occs_array, hole_delta_occs_array, plot_scale), interval=100, repeat=False) # Save the animation to gif @@ -153,7 +237,7 @@ def animate_pump_pulse(time_step, def update_scatter(anim_time, ax, time_step, idx_elec, idx_hole, elec_scat_list, hole_scat_list, - elec_delta_occs_array, hole_delta_occs_array): + elec_delta_occs_array, hole_delta_occs_array, plot_scale=1e3): """ Animate the pump pulse excitation for electrons and holes. @@ -186,16 +270,21 @@ def update_scatter(anim_time, ax, time_step, idx_elec, idx_hole, hole_delta_occs_array : numpy.ndarray Array of hole occupation changes, similar to elec_delta_occs_array. + + plot_scale : float + Scale factor for the scatter object sizes. """ elec_num_bands = elec_delta_occs_array.shape[0] hole_num_bands = hole_delta_occs_array.shape[0] + e_occs_max = np.max(elec_delta_occs_array[:, :, idx_elec].ravel()) for i in range(elec_num_bands): - elec_scat_list[i].set_sizes(elec_delta_occs_array[i, anim_time, idx_elec].flatten() * 2e4 + 1e-10) + elec_scat_list[i].set_sizes(elec_delta_occs_array[i, anim_time, idx_elec].flatten() / e_occs_max * plot_scale + 1e-10) + h_occs_max = np.max(hole_delta_occs_array[:, :, idx_hole].ravel()) for i in range(hole_num_bands): - hole_scat_list[i].set_sizes(hole_delta_occs_array[i, anim_time, idx_hole].flatten() * 2e4 + 1e-10) + hole_scat_list[i].set_sizes(hole_delta_occs_array[i, anim_time, idx_hole].flatten() / h_occs_max * plot_scale + 1e-10) suffix = ax.get_title().split(';')[0] ax.set_title(f'{suffix}; Time: {anim_time * time_step:.2f} fs') @@ -264,3 +353,92 @@ def plot_trans_abs_map(ax, time_grid, energy_grid, trans_abs, cbar.set_label(r'$\Delta A$ (arb. units)', fontsize=fsize) return ax + + +def plot_occs_on_bands(ax, kpath, bands, first_el_band_idx, + pump_energy, pump_spectral_width_fwhm, + scale=0.2): + """ + Plot the occupations created by a pump pulse on the bands. + Takes ax with the bands plotted and fills the bands with the occupations. + Occupations are set up with a Gaussian for each valence-conduction pair. + + Parameters + ---------- + + ax : matplotlib.axes.Axes + Axis for plotting the bands with occupations. + + kpath : numpy.ndarray + K-path for the bands, same as the one used for plotting the bands. + + bands : numpy.ndarray + Bands for the material. Shape: (num_bands, num_kpoints). Same as the one used for plotting the bands. + + first_el_band_idx : int + Index of the first electron (conduction) band. + + pump_energy : float + Pump energy in eV. + + pump_spectral_width_fwhm : float + Pump energy broadening full width at half maximum (FWHM) in eV. + + scale : float + Scale factor for the occupation amplitude. For plotting purposes. + + Returns + ------- + + ax : matplotlib.axes.Axes + Axis with the bands filled with the occupations. + """ + + # Total number of bands + nband = len(bands) + + # Total number of k-points + npoint = len(bands[1]) + + # Occupation amplitude for each bands + occs_amplitude_bands = np.zeros((nband, npoint)) + + elec_band_list = list(range(first_el_band_idx, nband)) + hole_band_list = list(range(1, first_el_band_idx)) + + elec_nband = len(elec_band_list) + hole_nband = len(hole_band_list) + + # First, compute the occupations for every valence-conduction pair + for i_elec_band in elec_band_list: + for i_hole_band in hole_band_list: + elec_band = bands[i_elec_band] + hole_band = bands[i_hole_band] + + occ = gaussian(elec_band - hole_band, + pump_energy, pump_spectral_width_fwhm) + + occs_amplitude_bands[i_elec_band - 1, :] += occ + occs_amplitude_bands[i_hole_band - 1, :] += occ + + max_occ = np.max(occs_amplitude_bands.ravel()) + factor = scale / max_occ + occs_amplitude_bands *= factor + + # Plot the occupations on bands with fill_between + # use kpath + for iband in range(1, nband): + + band = bands[iband] + + occ_top = band + occs_amplitude_bands[iband - 1, :] + occ_bottom = band - occs_amplitude_bands[iband - 1, :] + + if iband >= first_el_band_idx: + color = 'red' + else: + color = 'blue' + + ax.fill_between(kpath, occ_top, occ_bottom, color=color) + + return ax diff --git a/src/perturbopy/postproc/utils/spectra_trans_abs.py b/src/perturbopy/postproc/utils/spectra_trans_abs.py index 763847f..7ec7192 100644 --- a/src/perturbopy/postproc/utils/spectra_trans_abs.py +++ b/src/perturbopy/postproc/utils/spectra_trans_abs.py @@ -4,6 +4,7 @@ """ import numpy as np +import warnings from .memory import get_size from .timing import TimingGroup @@ -30,9 +31,10 @@ def gaussian_delta(x, mu, sig): def compute_trans_abs(elec_dyna_run, hole_dyna_run, de_grid=0.02, + energy_grid_max=None, eta=0.02, save_npy=True, - ): + tr_dipoles_sqr=None): """ 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). @@ -49,6 +51,10 @@ def compute_trans_abs(elec_dyna_run, de_grid : float Energy grid spacing for the transient absorption spectrum. Affects the calculation time. + energy_grid_max : float + Maximum probe energy for the spectrum. If None, the maximum band energy difference (condunction-valence) is used. + Minimum probe energy is always the bandgap. + eta : float Broadening parameter for the Gaussian delta functions applied to the energy grid. @@ -69,6 +75,11 @@ def compute_trans_abs(elec_dyna_run, dA_hole : np.ndarray Hole contribution to the transient absorption spectrum. Shape (num_energy_points, num_steps). + + tr_dipoles_sqr : np.ndarray + Transition dipoles squared, valence-to-conduction for each k-point and band. + Currently, the k-grid for dipoles must match the one for electrons. + Experimental feature. """ trun = TimingGroup('trans. abs.') @@ -77,7 +88,7 @@ def compute_trans_abs(elec_dyna_run, # 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): + if not np.all(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 @@ -117,13 +128,20 @@ def compute_trans_abs(elec_dyna_run, 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_band_max = np.max(elec_energy_array.ravel()) + energy_band_min = np.min(hole_energy_array.ravel()) + + if energy_grid_max is None: + energy_grid_max = energy_band_max - energy_band_min + elif energy_grid_max < bandgap: + warnings.warn(f'The maximum probe energy specified ({energy_grid_max}) is smaller than the bandgap ({bandgap}).' + f' The maximum probe energy is set to max(cond)-min(val) {energy_band_max - energy_band_min}.') + energy_grid_max = energy_band_max - energy_band_min # Energy grid for the transient absorption spectrum, affect the calculation time - trans_abs_energy_grid = np.arange(bandgap, energy_max - energy_min, de_grid) + trans_abs_energy_grid = np.arange(bandgap, energy_grid_max, de_grid) num_energy_points = trans_abs_energy_grid.shape[0] - print(f"{'Number of energy points':<30}: {num_energy_points}\n") + print(f"{'Number of energy grid points':>30}: {num_energy_points}\n") # Find the interection between the electron and hole k-points print('Finding intersect k-points...') @@ -137,6 +155,19 @@ def compute_trans_abs(elec_dyna_run, num_intersect_kpoints = ekidx.size + # Currently, the number of k-points of the transition dipoles + # must be the same as the number of k-points of the electron energy array + # In a future development, we will interpolate the transition dipoles on the fly + if tr_dipoles_sqr is not None: + elec_nband_tr_dip, hole_nband_tr_dip, num_k_tr_dip = tr_dipoles_sqr.shape + + if elec_nband_tr_dip != elec_nband or hole_nband_tr_dip != hole_nband \ + or num_k_tr_dip != elec_kpoint_array.shape[0]: + raise ValueError(f'Wrong shape of the transition dipoles array: {tr_dipoles_sqr.shape}') + else: + print('\nTransition dipoles are provided.') + print(f"{'Tr. dip. shape (num. elec. bands, num. hole bands, num. k points)':>30}: {tr_dipoles_sqr.shape}\n") + # Compute all the Gaussian deltas on the intersect grid print('Computing Gaussian deltas...') with trun.add('compute deltas') as t: @@ -148,12 +179,16 @@ def compute_trans_abs(elec_dyna_run, 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 * \ + ALL_DELTAS[iband, jband, :, ienergy] = \ gaussian_delta(elec_energy_array[ekidx, iband] - hole_energy_array[hkidx, jband], trans_abs_energy_grid[ienergy], eta) + if tr_dipoles_sqr is not None: + ALL_DELTAS[iband, jband, :, ienergy] *= 0.5 * tr_dipoles_sqr[iband, jband, ekidx] + else: + ALL_DELTAS[iband, jband, :, ienergy] *= 2.0 + get_size(ALL_DELTAS, 'ALL_DELTAS', dump=True) print('') @@ -208,7 +243,7 @@ def compute_trans_abs(elec_dyna_run, 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}' + ending = f'_Epump_{pump_energy:.4f}' # Total transient absorption np.save(f'trans_abs_dA{ending}', dA_elec + dA_hole) # Electron contribution diff --git a/src/perturbopy/yml_files/files_description.yml b/src/perturbopy/yml_files/files_description.yml index 365d356..09b64e9 100644 --- a/src/perturbopy/yml_files/files_description.yml +++ b/src/perturbopy/yml_files/files_description.yml @@ -158,7 +158,8 @@ prefix_popu.h5: pump_pulse.h5: type: HDF5 - description: "[OPTIONAL] HDF5 file containing the pump pulse information, used if pump_pulse=.true." + description: "[OPTIONAL] HDF5 file containing the pump pulse information, used if pump_pulse=.true.. The name of the pump pulse file can be specified with the pump_pulse_fname input parameter." + format example: mydoc_dynamics.html#pump_pulse_h5_file prefix.trans_coef: type: text