Added plotting scripts for new spin calc modes

src/perturbopy/postproc/__init__.py

@@ -3,10 +3,13 @@
 """
 from .calc_modes.calc_mode import CalcMode
 from .calc_modes.bands import Bands
+from .calc_modes.spins import Spins
 from .calc_modes.phdisp import Phdisp
 from .calc_modes.ephmat import Ephmat
+from .calc_modes.ephmat_spin import EphmatSpin
 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_pp import DynaPP
 

src/perturbopy/postproc/calc_modes/ephmat_spin.py

@@ -0,0 +1,189 @@
+import numpy as np
+from perturbopy.postproc.calc_modes.calc_mode import CalcMode
+from perturbopy.postproc.dbs.units_dict import UnitsDict
+from perturbopy.postproc.dbs.recip_pt_db import RecipPtDB
+from perturbopy.postproc.utils.plot_tools import plot_dispersion, plot_recip_pt_labels, plot_vals_on_bands
+
+
+class EphmatSpin(CalcMode):
+    """
+    Class representation of a Perturbo ephmat_spin calculation.
+
+    Attributes
+    ----------
+    kpt : RecipPtDB
+       Database for the k-points used in the ephmat calculation, containing N points.
+    
+    qpt : RecipPtDB
+       Database for the q-points used in the ephmat calculation, containing M points.
+    
+    phdisp : UnitsDict
+       Database for the phonon energies computed by the ephmat calculation. The keys are
+       the phonon mode, and the values are an array (of length M) containing the energies at each q-point
+       with units phdisp.units
+    
+    ephmat_spin : UnitsDict
+       Database for the e-ph spin flip matrix elements computed by the ephmat_spin calculation. The keys are
+       the phonon mode, and the values are an array (of length NxM) where element (n, m)
+       is the e-ph spin flip matrix element (units ephmat.units) between an electron at k-point n and phonon at q-point m
+    
+    defpot : UnitsDict
+       Database for the deformation potentials computed by the phdisp calculation. The keys are
+       the phonon mode, and the values are an array (of length NxM) where element (n, m)
+       is the deformation potential (units defpot.units) of an electron at k-point n and phonon at q-point m.
+
+    """
+
+    def __init__(self, pert_dict):
+        """
+        Constructor method
+
+        Parameters
+        ----------
+        pert_dict : dict
+            Dictionary containing the inputs and outputs from the ephmat calculation.
+
+        """
+        super().__init__(pert_dict)
+
+        if self.calc_mode != 'ephmat_spin':
+            raise ValueError('Calculation mode for a EphmatSpin object should be "ephmat_spin"')
+
+        phdisp_units = self._pert_dict['ephmat_spin'].pop('phonon energy units')
+        defpot_units = self._pert_dict['ephmat_spin'].pop('deformation potential units')
+        ephmat_units = self._pert_dict['ephmat_spin'].pop('e-ph matrix elements units')
+        nmode = self._pert_dict['ephmat_spin'].pop('number of phonon modes')
+
+        kpath_units = self._pert_dict['ephmat_spin'].pop('k-path coordinate units')
+        kpath = np.array(self._pert_dict['ephmat_spin'].pop('k-path coordinates'))
+        kpoint_units = self._pert_dict['ephmat_spin'].pop('k-point coordinate units')
+        kpoint = np.array(self._pert_dict['ephmat_spin'].pop('k-point coordinates'))
+
+        qpath_units = self._pert_dict['ephmat_spin'].pop('q-path coordinate units')
+        qpath = np.array(self._pert_dict['ephmat_spin'].pop('q-path coordinates'))
+        qpoint_units = self._pert_dict['ephmat_spin'].pop('q-point coordinate units')
+        qpoint = np.array(self._pert_dict['ephmat_spin'].pop('q-point coordinates'))
+
+        ephmat_dat = self._pert_dict['ephmat_spin'].pop('phonon mode')
+
+        self.kpt = RecipPtDB.from_lattice(kpoint, kpoint_units, self.lat, self.recip_lat, kpath, kpath_units)
+        self.qpt = RecipPtDB.from_lattice(qpoint, qpoint_units, self.lat, self.recip_lat, qpath, qpath_units)
+
+        phdisp = {}
+        defpot = {}
+        ephmat = {}
+
+        N = len(self.kpt.path)
+        M = len(self.qpt.path)
+
+        for phidx in ephmat_dat.keys():
+            phdisp[phidx] = ephmat_dat[phidx].pop('phonon energy')
+            defpot[phidx] = np.array(ephmat_dat[phidx].pop('deformation potential')).reshape(N, M)
+            ephmat[phidx] = np.array(ephmat_dat[phidx].pop('e-ph matrix elements')).reshape(N, M)
+
+        self.phdisp = UnitsDict.from_dict(phdisp, phdisp_units)
+        self.defpot = UnitsDict.from_dict(defpot, defpot_units)
+        self.ephmat = UnitsDict.from_dict(ephmat, ephmat_units)
+
+    def plot_phdisp(self, ax, show_qpoint_labels=True, **kwargs):
+        """
+        Method to plot the phonon dispersion.
+
+        Parameters
+        ----------
+        ax : matplotlib.axes.Axes
+           Axis on which to plot the phdisp.
+
+        energy_window : tuple of int, optional
+           The range of band energies to be shown on the y-axis.
+
+        show_qpoint_labels : bool, optional
+           If true, the q-point labels stored in the labels attribute will be shown on the plot. Default true.
+
+        Returns
+        -------
+        ax: matplotlib.axes.Axes
+           Axis with the plotted bands.
+
+        """
+        ax = plot_dispersion(ax, self.qpt.path, self.phdisp, self.phdisp.units, **kwargs)
+
+        if show_qpoint_labels:
+            ax = plot_recip_pt_labels(ax, self.qpt.labels, self.qpt.points, self.qpt.path)
+
+        return ax
+
+    def plot_defpot(self, ax, kpoint_idx=0, show_qpoint_labels=True, **kwargs):
+        """
+        Method to plot the phonon dispersion.
+
+        Parameters
+        ----------
+        ax : matplotlib.axes.Axes
+           Axis on which to plot the phdisp.
+
+        kpoint_idx : int, optional
+            Index of the k-point to plot the deformation potentials for. Deformation potentials will be plotted along q-points, at this k-point
+            By default, it will be the first k-point.
+
+        energy_window : tuple of int, optional
+           The range of band energies to be shown on the y-axis.
+
+        show_qpoint_labels : bool, optional
+           If true, the q-point labels stored in the labels attribute will be shown on the plot. Default true.
+
+        Returns
+        -------
+        ax: matplotlib.axes.Axes
+           Axis with the plotted bands.
+
+        """
+
+        values = {}
+
+        for key, val in self.defpot.items():
+            values[key] = self.defpot[key][kpoint_idx, :]
+
+        ax = plot_vals_on_bands(ax, self.qpt.path, self.phdisp, self.phdisp.units, values=values, label=r'$\Phi$', **kwargs)
+
+        if show_qpoint_labels:
+            ax = plot_recip_pt_labels(ax, self.qpt.labels, self.qpt.points, self.qpt.path)
+
+        return ax
+
+    def plot_ephmat(self, ax, kpoint_idx=0, show_qpoint_labels=True, **kwargs):
+        """
+        Method to plot the phonon dispersion.
+
+        Parameters
+        ----------
+        ax : matplotlib.axes.Axes
+           Axis on which to plot the phdisp.
+
+        kpoint_idx : int, optional
+            Index of the k-point to plot the e-ph elements for. E-ph elements will be plotted along q-points, at this k-point
+            By default, it will be the first k-point.
+        energy_window : tuple of int, optional
+           The range of band energies to be shown on the y-axis.
+
+        show_qpoint_labels : bool, optional
+           If true, the q-point labels stored in the labels attribute will be shown on the plot. Default true.
+
+        Returns
+        -------
+        ax: matplotlib.axes.Axes
+           Axis with the plotted bands.
+
+        """
+
+        values = {}
+
+        for key, val in self.ephmat.items():
+            values[key] = self.ephmat[key][kpoint_idx, :]
+
+        ax = plot_vals_on_bands(ax, self.qpt.path, self.phdisp, self.phdisp.units, values=values, label=r'$|g flip|$', **kwargs)
+
+        if show_qpoint_labels:
+            ax = plot_recip_pt_labels(ax, self.qpt.labels, self.qpt.points, self.qpt.path)
+
+        return ax

src/perturbopy/postproc/calc_modes/imsigma_spin.py

@@ -0,0 +1,97 @@
+import numpy as np
+from perturbopy.postproc.calc_modes.calc_mode import CalcMode
+from perturbopy.postproc.dbs.units_dict import UnitsDict
+from perturbopy.postproc.dbs.recip_pt_db import RecipPtDB
+from perturbopy.postproc.utils.constants import hbar
+
+
+class ImsigmaSpin(CalcMode):
+    """
+    Class representation of a Perturbo imsigma_spin calculation.
+
+    Attributes
+    ----------
+    kpt : RecipPtDB
+       Database for the k-points used in the imsigma calculation, containing N points.
+    
+    bands : UnitsDict
+       Database for the band energies computed by the imsigma calculation. The keys are
+       the band index, and the values are an array (of length N) containing the energies at each k-point
+       with units bands.units
+
+    temper : UnitsDict
+        Dictionary of temperatures used in each configuration. The keys give the configuration number,
+        and the values are floats giving the temperature (with units temper.units)
+    
+    chem_pot : UnitsDict
+        Dictionary of chemical potentials used in each configuration. The keys
+        give the configuration number, and the values are floats giving the
+        chemical potential (with units chem_pot.units)
+    
+    imsigma_flip : UnitsDict
+        Dictionary of spin flip imaginary self-energies computed for each configuration. The top level keys are the
+        configuration number, and the second level keys are the band index. The values are arrays of length N giving the
+        imaginary self-energies along all the k-points at that band index for the configuration. Units are in imsigma.units.
+
+    imsigma_flip_mode : UnitsDict
+        Dictionary of spin flip imaginary self-energies resolved by phonon mode computed. The top level keys are the
+        configuration number, and the second level keys are the band index. The third level keys are
+        the phonon mode. Finally,the values are arrays of length N giving the imaginary self-energies along all the k-points
+        due to the given phonon mode at that band index for the configuration. Units are in imsigma_mode.units.
+    
+    """
+
+    def __init__(self, pert_dict):
+        """
+        Constructor method
+
+        Parameters
+        ----------
+        pert_dict : dict
+            Dictionary containing the inputs and outputs from the imsigma calculation.
+
+        """
+        super().__init__(pert_dict)
+
+        if self.calc_mode.split('-')[0] != 'imsigma_spin':
+            raise ValueError('Calculation mode for an ImsigmaSpinCalcMode object should be "imsigma_spin"')
+
+        kpoint_units = self._pert_dict['imsigma_spin'].pop('k-point coordinate units')
+        num_kpoints = self._pert_dict['imsigma_spin'].pop('number of k-points')
+        kpoint = np.array(self._pert_dict['imsigma_spin'].pop('k-point coordinates'))
+        self.kpt = RecipPtDB.from_lattice(kpoint, kpoint_units, self.lat, self.recip_lat)
+
+        energy_units = self._pert_dict['imsigma_spin'].pop('energy units')
+        num_bands = self._pert_dict['imsigma_spin'].pop('number of bands')
+        energies_dict = self._pert_dict['imsigma_spin']['energy'].pop('band index')
+        self.bands = UnitsDict.from_dict(energies_dict, energy_units)
+
+        num_config = self._pert_dict['imsigma_spin'].pop('number of configurations')
+        config_dat = self._pert_dict['imsigma_spin'].pop('configuration index')
+        num_modes = self._pert_dict['imsigma_spin'].pop('number of phonon modes')
+
+        self.temper = UnitsDict(units=self._pert_dict['imsigma_spin'].pop('temperature units'))
+        self.chem_pot = UnitsDict(units=self._pert_dict['imsigma_spin'].pop('chemical potential units'))
+        self.imsigma_flip = UnitsDict(units=self._pert_dict['imsigma_spin'].pop('Im(Sigma) units'))
+        self.imsigma_flip_mode = UnitsDict(units=self.imsigma_flip.units)
+
+
+        for config_idx in config_dat.keys():
+
+            self.imsigma_flip[config_idx] = {}
+            self.imsigma_flip_mode[config_idx] = {}
+            self.temper[config_idx] = config_dat[config_idx].pop('temperature')
+            self.chem_pot[config_idx] = config_dat[config_idx].pop('chemical potential')
+
+            imsigma_dat = config_dat[config_idx].pop('band index')
+
+            for mode in np.arange(1, num_modes + 1):
+                self.imsigma_flip_mode[config_idx][mode] = {}
+
+            for band_index in imsigma_dat.keys():
+
+                self.imsigma_flip[config_idx][band_index] = np.array(imsigma_dat[band_index]['Im(Sigma)']['total'])
+
+                for mode in np.arange(1, num_modes + 1):
+                    self.imsigma_flip_mode[config_idx][mode][band_index] = np.array(imsigma_dat[band_index]['Im(Sigma)']['phonon mode (total)'][mode])
+