Coverage for strongcoca / response / excitations.py: 100%

56 statements  

« prev     ^ index     » next       coverage.py v7.13.4, created at 2026-07-25 16:26 +0000

1import numpy as np 

2 

3from .base import BaseResponse 

4from ..types import Array 

5from ..units import au_to_eV, eV_to_au, au_to_eA, eA_to_au 

6from .utilities import Broadening, NoArtificialBroadening, broaden 

7 

8 

9class Excitations(BaseResponse): 

10 """Objects of this class represent discrete excitation spectra. 

11 

12 Parameters 

13 ---------- 

14 energies 

15 Excitation energies; array with shape {n}. 

16 

17 Units are eV by default, optionally atomic units (see :attr:`units`). 

18 transition_dipole_moments 

19 Transition dipole moments of the excitations; array with shape {n}x{3}. 

20 

21 Units are eÅ by default, optionally atomic units (see :attr:`units`). 

22 broadening 

23 Artificial broadening used for the continuous response; 

24 defaults to no broadening. 

25 units 

26 `eVA` to specify energies in eV and transition_dipole_moments in eÅ 

27 or `au` to specify inputs in atomic units. 

28 

29 This parameter determines whether conversion should be performed during 

30 initialization and has no effect on instance methods and variables. 

31 name 

32 Name of response. 

33 """ 

34 def __init__(self, 

35 energies: np.ndarray, 

36 transition_dipole_moments: np.ndarray, 

37 broadening: Broadening = NoArtificialBroadening(), 

38 units: str = 'eVA', 

39 name: str = 'Excitations') -> None: 

40 super().__init__(broadening=broadening, pbc=False, name=name) 

41 

42 energies = np.array(energies, dtype=float) 

43 transition_dipole_moments = np.array(transition_dipole_moments, dtype=float) 

44 

45 if units == 'eVA': 

46 energies *= eV_to_au 

47 transition_dipole_moments *= eA_to_au 

48 elif units != 'au': 

49 raise ValueError(f"units has to be 'eVA' or 'au', not '{units}'") 

50 

51 self._energies = energies 

52 self._transition_dipole_moments = transition_dipole_moments 

53 

54 @property 

55 def energies(self) -> np.ndarray: 

56 """Excitation energies in units of eV; array with shape {n}.""" 

57 return self._energies * au_to_eV 

58 

59 @property 

60 def transition_dipole_moments(self) -> np.ndarray: 

61 """Transition dipole moments of the excitations in units of eÅ; 

62 array with shape {n}x{3}. 

63 """ 

64 return self._transition_dipole_moments * au_to_eA 

65 

66 @property 

67 def oscillator_strengths(self) -> np.ndarray: 

68 """Unitless oscillator strengths of the excitations; 

69 array with shape {n}. 

70 """ 

71 osc_nv = self.oscillator_strength_vectors 

72 osc_n = np.average(osc_nv, axis=1) 

73 return osc_n # type: ignore 

74 

75 @property 

76 def oscillator_strength_vectors(self) -> np.ndarray: 

77 """Unitless oscillator strength vectors of the excitations; 

78 array with shape {n}x3. 

79 """ 

80 omega_n = self._energies 

81 mu_nv = self._transition_dipole_moments 

82 osc_nv = 2 * omega_n[:, np.newaxis] * mu_nv**2 

83 return osc_nv # type: ignore 

84 

85 @property 

86 def oscillator_strength_tensors(self) -> np.ndarray: 

87 """Unitless oscillator strength tensors of the excitations; 

88 array with shape {n}x3x3. 

89 """ 

90 omega_n = self._energies 

91 mu_nv = self._transition_dipole_moments 

92 osc_nvv = 2 * np.einsum('n,nx,ny->nxy', omega_n, mu_nv, mu_nv, optimize=True) 

93 return osc_nvv # type: ignore 

94 

95 def get_dipole_strength_function(self, frequencies: Array) -> np.ndarray: 

96 if isinstance(self._broadening, NoArtificialBroadening): 

97 raise ValueError( 

98 'get_dipole_strength_function requires broadening for discrete excitation ' 

99 'spectra. Pass e.g. LorentzianBroadening(0.1) when constructing the ' 

100 'response or calculator.') 

101 return super().get_dipole_strength_function(frequencies) 

102 

103 def _get_dynamic_polarizability(self, frequencies: Array) -> np.ndarray: 

104 freq_w = np.asarray(frequencies) 

105 omega_I = self._energies 

106 osc_Ivv = self.oscillator_strength_tensors 

107 dm_wvv = broaden(omega_I, osc_Ivv, freq_w, broadening=self._broadening) 

108 return dm_wvv 

109 

110 def _get_dynamic_polarizability_imaginary_frequency( 

111 self, frequencies: Array) -> np.ndarray: 

112 freq_w = np.asarray(frequencies) 

113 omega_I = self._energies 

114 osc_Ivv = self.oscillator_strength_tensors 

115 dm_wvv = broaden(omega_I, osc_Ivv, freq_w, freq_axis='imag') 

116 return dm_wvv