Cumulant spectral (#83)

* add cumulant

* add cumulant pp

* all the format issues are fixed.

* add sto yml file

* Update src/perturbopy/postproc/calc_modes/spectral_cumulant.py

Co-authored-by: Sergei Kliavinek <klyavinekss@gmail.com>

* Update src/perturbopy/postproc/calc_modes/spectral_cumulant.py

Co-authored-by: Sergei Kliavinek <klyavinekss@gmail.com>

* update

* add shape checking

---------

Co-authored-by: Sergei Kliavinek <klyavinekss@gmail.com>

src/perturbopy/postproc/__init__.py

@@ -3,6 +3,7 @@
 """
 from .calc_modes.calc_mode import CalcMode
 from .calc_modes.bands import Bands
+from .calc_modes.spectral_cumulant import SpectralCumulant
 from .calc_modes.spins import Spins
 from .calc_modes.phdisp import Phdisp
 from .calc_modes.ephmat import Ephmat

src/perturbopy/postproc/calc_modes/calc_mode.py

@@ -45,8 +45,9 @@ def __init__(self, pert_dict):
 
         """
 
-        # Extract calculation mode name from pert_dict
+        # Extract calculation mode name and prefix from pert_dict
         self.calc_mode = pert_dict['input parameters']['after conversion'].pop('calc_mode')
+        self.prefix = pert_dict['input parameters']['after conversion'].pop('prefix')
 
         # Extract basic data from pert_dict
         self.alat = pert_dict['basic data']['alat']

src/perturbopy/postproc/calc_modes/dyna_run.py

@@ -34,7 +34,7 @@ def __init__(self, cdyna_file, tet_file, pert_dict):
             Dictionary containing the inputs and outputs from the dynamics-run calculation.
 
         """
-        
+
         self.timings = TimingGroup("dynamics-run")
 
         super().__init__(pert_dict)
@@ -175,7 +175,7 @@ def extract_steady_drift_vel(self, dyna_pp_yaml_path):
             raise FileNotFoundError(f'File {dyna_pp_yaml_path} not found')
 
         dyna_pp_dict = open_yaml(dyna_pp_yaml_path)
-        
+
         vels = dyna_pp_dict['dynamics-pp']['velocity']
         concs = dyna_pp_dict['dynamics-pp']['concentration']
 
@@ -185,9 +185,9 @@ def extract_steady_drift_vel(self, dyna_pp_yaml_path):
         steady_conc = []
 
         for irun, dynamics_run in self._data.items():
-            
+
             step_number += dynamics_run.num_steps
-            
+
             if np.allclose(dynamics_run.efield, np.array([0.0, 0.0, 0.0])):
                 steady_drift_vel.append(None)
                 steady_conc.append(concs[step_number])

src/perturbopy/postproc/calc_modes/spectral_cumulant.py

@@ -0,0 +1,132 @@
+import numpy as np
+import h5py as h5
+import matplotlib.pyplot as plt
+from perturbopy.postproc.calc_modes.calc_mode import CalcMode
+from perturbopy.io_utils.io import open_yaml, open_hdf5, close_hdf5
+import os
+
+
+class SpectralCumulant(CalcMode):
+    """
+    Class representation of a Perturbo cumulant calculation
+
+    Parameters
+    ----------
+    prefix : str
+        Prefix for the HDF5 file containing spectral function data.
+
+    Attributes
+    ----------
+    temp_array : numpy.ndarray
+        Array of temperatures corresponding to the spectral data.
+    freq_array : numpy.ndarray
+        Array of energy value in electron volts (eV).
+    freq_step : float
+        Energy step size for the energy grid in eV.
+    Akw : numpy.ndarray
+        Spectral function data, indexed by k-point, band, temperature, and ω.
+    """
+    def __init__(self, spectral_file, pert_dict):
+        """
+        Constructor method
+
+        Parameters
+        ----------
+        spectral_file : dict
+            Dictionary for the prefix_spectral_cumulant.h5
+        pert_dict : dict
+            Dictionary containing the inputs from the spectral-cum calculation.
+
+        """
+        super().__init__(pert_dict)
+        if self.calc_mode != 'spectral-cum':
+            raise ValueError('Calculation mode for a SpectralCumulantCalcMode object should be "spectral-cum"')
+
+        Akw = spectral_file['spectral_functions']
+        w_lower = np.asarray(spectral_file['w_lower_index'])
+        w_upper = np.asarray(spectral_file['w_upper_index'])
+        freq_step = np.asarray(spectral_file['wfreq_step_eV'])
+        self.temp_array = np.asanyarray(spectral_file['temperatures'])
+        self.freq_array = np.arange(w_lower, w_upper + 1) * freq_step
+        self.freq_step = freq_step
+        Akw_np = []
+        for key in Akw.keys():
+            Akw_np.append(np.asarray(Akw[key]))
+        close_hdf5(spectral_file)
+
+        self.Akw = np.asarray(Akw_np)
+
+    @classmethod
+    def from_hdf5_yaml(cls, spectral_path, yaml_path='pert_output.yml'):
+        """
+        Class method to create a SpectralCumulantCalcMode object from the HDF5 file and YAML file
+        generated by a Perturbo calculation
+
+        Parameters
+        ----------
+        spectral_path : str
+           Path to the HDF5 file generated by a spectral-cum calculation
+        yaml_path : str, optional
+           Path to the YAML file generated by a spectral-cum calculation
+
+        Returns
+        -------
+
+        """
+
+        if not os.path.isfile(spectral_path):
+            raise FileNotFoundError(f'File {spectral_path} not found')
+        if not os.path.isfile(yaml_path):
+            raise FileNotFoundError(f'File {yaml_path} not found')
+
+        spectral_file = open_hdf5(spectral_path)
+        yaml_dict = open_yaml(yaml_path)
+
+        return cls(spectral_file, yaml_dict)
+
+    def plot_Aw(self, ax, ik=0, it=0, ib=0):
+        """
+        Plots the spectral function A(ω) for a given k-point, temperature, and band.
+
+        Parameters
+        ----------
+        ax : matplotlib.axes.Axes
+            Axis on which to plot the spectral function.
+        ik : int, optional
+            Index of the k-point in the grid. Default is 0.
+        it : int, optional
+            Index of the temperature in the temperature array. Default is 0.
+        ib : int, optional
+            Index of the band. Default is 0.
+
+        Returns
+        -------
+        matplotlib.axes.Axes
+            The axis object with the plotted spectral function.
+
+        Notes
+        -----
+        The spectral function is normalized before plotting. The plot includes
+        labels, legends, and adjusted aesthetics for better visualization.
+        """
+        if it > len(self.temp_array):
+            raise ValueError('Temperature index is out of range')
+        if ib > self.Akw[0].shape[0]:
+            raise ValueError('Band index is out of range')
+        if ik > len(self.Akw):
+            raise ValueError('k-point index is out of range')
+        A0w = self.Akw[ik][ib, it, :]
+        freq_step = self.freq_step
+        freq_array = self.freq_array
+        # normalize
+        A0w = A0w / (np.sum(A0w) * freq_step)
+        # plot
+        ax.plot(freq_array, A0w, lw=2, label=f'T={ int(self.temp_array[it])} K')
+        ax.legend(fontsize=18)
+        plt.ylim([0, np.max(A0w) * 1.1])
+        plt.xlabel(r'$\omega-\epsilon_{nk}$ (eV)', fontsize=20)
+        plt.ylabel(r'A($\omega$) (eV$^{-1}$)', fontsize=20)
+        plt.xticks(fontsize=20)
+        plt.yticks(fontsize=20)
+        plt.tight_layout()
+        return ax