Source code for strongcoca.response.excitations

import numpy as np

from .base import BaseResponse
from ..types import Array
from ..units import au_to_eV, eV_to_au, au_to_eA, eA_to_au
from .utilities import Broadening, NoArtificialBroadening, broaden




class Excitations(BaseResponse):
    """Objects of this class represent discrete excitation spectra.

    Parameters
    ----------
    energies
        Excitation energies; array with shape {n}.

        Units are eV by default, optionally atomic units (see :attr:`units`).
    transition_dipole_moments
        Transition dipole moments of the excitations; array with shape {n}x{3}.

        Units are eÅ by default, optionally atomic units (see :attr:`units`).
    broadening
        Artificial broadening used for the continuous response;
        defaults to no broadening.
    units
        `eVA` to specify energies in eV and transition_dipole_moments in eÅ
        or `au` to specify inputs in atomic units.

        This parameter determines whether conversion should be performed during
        initialization and has no effect on instance methods and variables.
    name
        Name of response.
    """
    def __init__(self,
                 energies: np.ndarray,
                 transition_dipole_moments: np.ndarray,
                 broadening: Broadening = NoArtificialBroadening(),
                 units: str = 'eVA',
                 name: str = 'Excitations') -> None:
        super().__init__(broadening=broadening, pbc=False, name=name)

        if units == 'eVA':
            energies = energies * eV_to_au
            transition_dipole_moments = transition_dipole_moments * eA_to_au
        elif units != 'au':
            raise ValueError(f"units has to be 'eVA' or 'au', not '{units}'")

        self._energies = energies
        self._transition_dipole_moments = transition_dipole_moments

    @property
    def energies(self) -> np.ndarray:
        """Excitation energies in units of eV; array with shape {n}."""
        return self._energies * au_to_eV  # type: ignore

    @property
    def transition_dipole_moments(self) -> np.ndarray:
        """Transition dipole moments of the excitations in units of eÅ;
        array with shape {n}x{3}.
        """
        return self._transition_dipole_moments * au_to_eA  # type: ignore

    @property
    def oscillator_strengths(self) -> np.ndarray:
        """Unitless oscillator strengths of the excitations;
        array with shape {n}.
        """
        osc_nv = self.oscillator_strength_vectors
        osc_n = np.average(osc_nv, axis=1)
        return osc_n  # type: ignore

    @property
    def oscillator_strength_vectors(self) -> np.ndarray:
        """Unitless oscillator strength vectors of the excitations;
        array with shape {n}x3.
        """
        omega_n = self._energies
        mu_nv = self._transition_dipole_moments
        osc_nv = 2 * omega_n[:, np.newaxis] * mu_nv**2
        return osc_nv  # type: ignore

    @property
    def oscillator_strength_tensors(self) -> np.ndarray:
        """Unitless oscillator strength tensors of the excitations;
        array with shape {n}x3x3.
        """
        omega_n = self._energies
        mu_nv = self._transition_dipole_moments
        osc_nvv = 2 * np.einsum('n,nx,ny->nxy', omega_n, mu_nv, mu_nv, optimize=True)
        return osc_nvv  # type: ignore



    def get_dipole_strength_function(self, frequencies: Array) -> np.ndarray:
        if isinstance(self._broadening, NoArtificialBroadening):
            raise ValueError(
                'get_dipole_strength_function requires broadening for discrete excitation '
                'spectra. Pass e.g. LorentzianBroadening(0.1) when constructing the '
                'response or calculator.')
        return super().get_dipole_strength_function(frequencies)


    def _get_dynamic_polarizability(self, frequencies: Array) -> np.ndarray:
        freq_w = np.asarray(frequencies)
        omega_I = self._energies
        osc_Ivv = self.oscillator_strength_tensors
        dm_wvv = broaden(omega_I, osc_Ivv, freq_w, broadening=self._broadening)
        return dm_wvv

    def _get_dynamic_polarizability_imaginary_frequency(
            self, frequencies: Array) -> np.ndarray:
        freq_w = np.asarray(frequencies)
        omega_I = self._energies
        osc_Ivv = self.oscillator_strength_tensors
        dm_wvv = broaden(omega_I, osc_Ivv, freq_w, freq_axis='imag')
        return dm_wvv