Coverage for strongcoca/calculators/casida_calculator.py: 100%

1import logging 

2from typing import Optional 

3import numpy as np 

4from scipy.linalg import block_diag 

5from .base_calculator import BaseCalculator 

6from ..response.excitations import Excitations 

7from ..response.utilities import Broadening, NoArtificialBroadening 

8from .. import CoupledSystem 

9from .utilities import get_dipole_dipole_tensor 

10from ..types import Array 

11 

12 

13logger = logging.getLogger(__name__) 

14 

15 

16class CasidaCalculator(BaseCalculator): 

17 

18 """Instances of this class enable the calculation of correlation energy 

19 and spectrum of a coupled system, all polarizable objects in which have an 

20 internal representation in the form of a Casida matrix. 

21 

22 The calculations are done for discrete excitations, and continuous spectra 

23 calculated with this calculator use the given broadening. 

24 

25 Parameters 

26 ---------- 

27 coupled_system 

28 Coupled system for which to carry out calculations. 

29 broadening 

30 Artificial broadening used for the continuous response; 

31 defaults to no broadening. 

32 name 

33 Name of response. 

34 

35 Examples 

36 -------- 

37 

38 The following code snippet shows the calculation of the correlation energy 

39 for a simple demo system: 

40 

41 >>> from strongcoca import CoupledSystem, PolarizableUnit 

42 >>> from strongcoca.response import build_random_casida 

43 >>> from strongcoca.calculators import CasidaCalculator 

44 >>> 

45 >>> # construct coupled system 

46 >>> response = build_random_casida(n_states=2, name='Tiger', random_seed=42) 

47 >>> pu1 = PolarizableUnit(response, [0, 0, 0], name='Cheetah') 

48 >>> pu2 = PolarizableUnit(response, [3, 0, 0], name='Leopard') 

49 >>> cs = CoupledSystem([pu1, pu2]) 

50 >>> 

51 >>> # set up calculator 

52 >>> calc = CasidaCalculator(cs) 

53 >>> 

54 >>> # ... and compute the correlation energy 

55 >>> calc.get_correlation_energy() 

56 -0.050798240... 

57 """ 

58 

59 def __init__(self, 

60 coupled_system: CoupledSystem, 

61 broadening: Broadening = NoArtificialBroadening(), 

62 name: str = 'CasidaCalculator') -> None: 

63 super().__init__(coupled_system, broadening=broadening, name=name) 

64 self._excitations: Optional[Excitations] = None 

65 

66 def _discard_data(self) -> None: 

67 super()._discard_data() 

68 self._excitations = None 

69 

70 @property 

71 def excitations(self) -> Excitations: 

72 """Excitations as :class:`~strongcoca.response.excitations.Excitations`.""" 

73 self._verify_state() 

74 if self._excitations is None: 

75 self._excitations = self._build_casida_system() 

76 return self._excitations 

77 

78 def _calculate_correlation_energy(self) -> float: 

79 """Returns the correlation energy of the coupled system in atomic units. 

80 """ 

81 logger.debug('Entering CasidaCalculator._calculate_correlation_energy') 

82 energy = 0.5 * (np.sum(self.excitations._energies) - np.trace(self._D)) # type: float 

83 return energy 

84 

85 def _build_casida_system(self) -> Excitations: 

86 """Compile the Casida matrices representing the coupled system. 

87 """ 

88 logger.debug('Entering CasidaCalculator._build_casida_system') 

89 

90 for k, pu in enumerate(self._coupled_system): 

91 if not hasattr(pu.response, '_D_n') or \ 

92 not hasattr(pu.response, '_K_nn') or \ 

93 not hasattr(pu.response, '_mu_nv'): 

94 raise ValueError(f'Polarizable unit #{k} has no internal Casida representation.') 

95 

96 # The 'type: ignore' statements on the following lines are necessary 

97 # to apease mypy since it does not pick up that the loop above ensures 

98 # that the components of the Casida representation are in fact available. 

99 D = block_diag(*[np.diag(pu.response._D_n) # type: ignore 

100 for pu in self._coupled_system]) 

101 

102 K_blocks = [] 

103 for ii, pu1 in enumerate(self._coupled_system): 

104 row = [] 

105 mu1 = np.array([pu1.orientation.apply(m) 

106 for m in pu1.response._mu_nv]) # type: ignore 

107 for jj, pu2 in enumerate(self._coupled_system): 

108 if ii == jj: 

109 K_sub = pu1.response._K_nn # type: ignore 

110 else: 

111 # todo: include cell and pbc 

112 distance_vector = pu2._position - pu1._position 

113 mu2 = np.array([pu2.orientation.apply(m) 

114 for m in pu2.response._mu_nv]) # type: ignore 

115 K_sub = np.dot(mu1, np.dot(get_dipole_dipole_tensor(distance_vector), mu2.T)) 

116 if ii > jj: 

117 K_sub = K_sub.conj() 

118 row.append(K_sub) 

119 K_blocks.append(row) 

120 K = np.block(K_blocks) 

121 

122 C = D ** 2 + 2 * np.dot(np.sqrt(D), np.dot(K, np.sqrt(D))) 

123 U = block_diag(*[pu.response.U for pu in self._coupled_system]) # type: ignore 

124 C = np.dot(U.conj().T, np.dot(C, U)) 

125 self._D = np.sqrt(np.diag(np.diag(C))) 

126 eigval_I, eigvec_NI = np.linalg.eig(C) 

127 srt_I = np.argsort(eigval_I) 

128 eigval_I = eigval_I[srt_I] 

129 eigvec_NI = eigvec_NI[:, srt_I] 

130 

131 # Calculate transition dipole moments 

132 omega_I = np.sqrt(eigval_I) 

133 F_NI = np.dot(U, eigvec_NI) 

134 mu_vNN = [] 

135 for v in range(3): 

136 mu_NN = block_diag(*[np.diag(pu.response._mu_nv[:, v]) # type: ignore 

137 for pu in self._coupled_system]) 

138 mu_vNN.append(mu_NN) 

139 mu_Iv = np.einsum('vNN,NN,NI,I->Iv', mu_vNN, np.sqrt(D), F_NI, 

140 1 / np.sqrt(omega_I), 

141 optimize=True) 

142 return Excitations(omega_I, mu_Iv, broadening=self._broadening, 

143 name=self._name, units='au') 

144 

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

146 return self.excitations._get_dynamic_polarizability(frequencies) 

147 

148 def _get_dynamic_polarizability_imaginary_frequency( 

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

150 return self.excitations._get_dynamic_polarizability_imaginary_frequency( 

151 frequencies)