Added link to the pump pulse HDF5 file

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

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

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

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