Coverage for strongcoca/calculators/polarizability_calculator.py: 100%
94 statements
« prev ^ index » next coverage.py v7.10.6, created at 2025-10-26 18:44 +0000
1import logging
3from typing import Union
4import numpy as np
5from numpy.linalg import solve
6from scipy.linalg import logm
8from .. import CoupledSystem
9from ..response.utilities import Broadening, GaussianBroadening
10from ..types import Array
11from ..units import eV_to_au
12from .base_calculator import BaseCalculator
13from .utilities import get_dipole_dipole_tensor
14from .. import env
16logger = logging.getLogger(__name__)
19class PolarizabilityCalculator(BaseCalculator):
20 """Instances of this class enable the calculation of correlation energy
21 and spectrum of a coupled system, all polarizable objects in which have an
22 internal representation in the form of a polarizability tensor.
24 Spectra calculated with this calculator derive broadening from
25 the underlying polarizable units. Any additional or separate broadening
26 is not added.
28 Parameters
29 ----------
30 coupled_system
31 Coupled system for which to carry out calculations.
32 imaginary_frequencies
33 Frequency grid along imaginary axis used for correlation energy
34 calculation; eV by default, optionally atomic units (see :attr:`units`).
35 units
36 `eV` to specify energies in eV or `au` to specify inputs in atomic units.
38 This parameter determines whether conversion should be performed during
39 initialization and has no effect on instance methods and variables.
40 name
41 Name of response.
43 Examples
44 --------
46 The following code snippet shows the calculation of the correlation energy
47 for a simple demo system:
49 >>> import numpy as np
50 >>> from strongcoca import CoupledSystem, PolarizableUnit
51 >>> from strongcoca.response import build_random_casida
52 >>> from strongcoca.calculators import PolarizabilityCalculator
53 >>>
54 >>> # construct coupled system
55 >>> response = build_random_casida(n_states=2)
56 >>> pu1 = PolarizableUnit(response, [0, 0, 0])
57 >>> pu2 = PolarizableUnit(response, [2, 0, 0])
58 >>> cs = CoupledSystem([pu1, pu2])
59 >>>
60 >>> # set up calculator
61 >>> calc = PolarizabilityCalculator(cs, np.linspace(0, 10, 100))
62 >>>
63 >>> # ... and compute the correlation energy
64 >>> calc.get_correlation_energy()
65 -0.72084000362...
66 """
68 def __init__(self,
69 coupled_system: CoupledSystem,
70 imaginary_frequencies: Array,
71 units: str = 'eV',
72 name: str = 'PolarizabilityCalculator') -> None:
73 super().__init__(coupled_system, broadening=Broadening(), name=name)
75 imaginary_frequencies = np.asarray(imaginary_frequencies)
76 if units == 'eV':
77 imaginary_frequencies = imaginary_frequencies * eV_to_au
78 elif units != 'au':
79 raise ValueError(f"units has to be 'eV' or 'au', not '{units}'")
81 self._ifreq_w = imaginary_frequencies
82 self._K_nn: Union[np.ndarray, None] = None
83 # Check frequency grid
84 difreq_w = self._ifreq_w[1:] - self._ifreq_w[:-1] # type: ignore
85 self._difreq = difreq_w[0]
86 if not np.allclose(difreq_w, self._difreq):
87 raise ValueError('Frequency grid needs to be equally spaced.')
89 def _build_coupling_matrix(self) -> np.ndarray:
90 """Builds coupling matrix.
92 Coupling matrix is formed of blocks:
93 Each block is a 3x3 dipole-dipole tensor and there is NxN blocks
94 in total (N: number of units in coupled system).
95 The 3x3 blocks on the diagonal are zero.
96 """
97 Ni = len(self._coupled_system)
98 K_nn = np.zeros((3 * Ni, 3 * Ni), dtype=complex)
99 for i, pu1 in enumerate(self._coupled_system):
100 # index range of the block
101 n0 = i * 3
102 n1 = n0 + 3
103 for j in range(i + 1, Ni):
104 # index range of the block
105 m0 = j * 3
106 m1 = m0 + 3
107 pu2 = self._coupled_system[j]
108 # todo: include cell and pbc
109 distance_v = pu2._position - pu1._position
110 K_vv = get_dipole_dipole_tensor(distance_v)
111 K_nn[n0:n1, m0:m1] = K_vv
112 K_nn[m0:m1, n0:n1] = K_vv.conj()
113 return K_nn
115 def _calculate_correlation_energy(self) -> float:
116 """Returns the correlation energy of the coupled system in atomic units.
117 """
119 # build coupling matrix
120 Ni = len(self._coupled_system)
121 K_nn = self.coupling_matrix
123 # Calculate the matrix product of
124 # 1) a block diagonal matrix formed of 3x3 polarizability blocks and
125 # 2) the coupling matrix (formed of 3x3 blocks too).
126 # The following does not construct the block diagonal matrix in full but
127 # does the matrix products block-by-block as the blocks are calculated.
128 chi_K_wnn = np.zeros((len(self._ifreq_w), 3 * Ni, 3 * Ni), dtype=complex)
129 for i, pu in enumerate(self._coupled_system):
130 # calculate polarizability block
131 dm_wvv = pu._get_dynamic_polarizability_imaginary_frequency(self._ifreq_w)
132 # index range of the block
133 n0 = i * 3
134 n1 = n0 + 3
135 # take matrix product with the coupling matrix
136 chi_K_wnn[:, n0:n1, :] = np.dot(dm_wvv, K_nn[n0:n1])
138 # carry out integration over imaginary frequency axis
139 integral = 0
140 # TODO: this frequency loop is slow; reformulate without it to speed up
141 for chi_K_nn in chi_K_wnn:
142 D_nn = np.eye(3 * Ni) - chi_K_nn
143 dint = np.trace(logm(D_nn) + chi_K_nn.T.conj())
144 integral += dint
145 integral *= self._difreq
146 energy = np.real(integral) / (2 * np.pi) # type: float
147 return energy
149 @property
150 def coupling_matrix(self) -> np.ndarray:
151 r""" The coupling matrix
153 .. math::
154 \boldsymbol{T}_{ij} =
155 \frac{3\boldsymbol{r}^{(ij)}(\boldsymbol{r}^{(ij)})^\text{T}}{\left|r^{(ij)}\right|^5}
156 - \frac{1}{\left|r^{(ij)}\right|^3}
158 as a matrix where the rows and columns each correspond to the system index
159 :math:`i` or :math:`j` *and* the Cartesian direction.
161 This quantity is calculated on first use and then buffered.
162 """
163 if self._K_nn is None:
164 self._K_nn = self._build_coupling_matrix()
165 return self._K_nn
167 def _get_dynamic_polarizability(self, frequencies: Array) -> np.ndarray:
168 if any(isinstance(pu.broadening, GaussianBroadening)
169 for pu in self._coupled_system):
170 raise NotImplementedError(
171 'Gaussian broadening is not supported in PolarizabilityCalculator. '
172 'See issue #42.')
174 # build coupling matrix
175 freq_w = np.array(frequencies)
176 dm_wvv = np.zeros((len(freq_w), 3, 3), dtype=complex)
178 # The linear solve is done frequency by frequency.
179 # For large systems, one should not do all frequencies at once
180 # as the intermediate A matrix takes up too much memory.
181 # Limit the matrix A to maxsize elements.
182 to_MiB = 1024 ** -2
183 mem_limit = env.get_float('MAX_SOLVE_MEM') / to_MiB
184 syssize = 3 ** 2 * len(self._coupled_system) ** 2 * 8 # Size of A matrix in bytes
185 chunksize = int(mem_limit) // syssize # Maximum number of frequencies at a time
186 if chunksize < 1:
187 # The system is so huge so even one frequency at a time violates our
188 # memory limit. We have to live with it.
189 chunksize = 1
191 # This loop will just iterate once for small enough systems
192 for indices in np.array_split(np.arange(0, len(freq_w)),
193 (len(frequencies) + chunksize - 1) // chunksize):
194 dm_wvv[indices] = self._get_dynamic_polarizability_chunk(freq_w[indices])
195 return dm_wvv
197 def _get_dynamic_polarizability_chunk(self, frequencies: np.ndarray) -> np.ndarray:
198 """ Calculate the dynamic polarizability for a subset of frequencies. """
199 Ni = len(self._coupled_system)
200 K_nn = self.coupling_matrix
201 freq_w = np.array(frequencies)
203 # Linear solve
204 rhs_wnv = np.zeros((len(freq_w), Ni * 3, 3), dtype=complex)
205 A_wnn = np.zeros((len(freq_w), Ni * 3, Ni * 3), dtype=complex)
206 A_wnn += np.eye(Ni * 3)
208 for i, pu in enumerate(self._coupled_system):
209 dm_wvv = pu._get_dynamic_polarizability(freq_w)
210 n0 = i * 3
211 n1 = n0 + 3
212 # Extract the right-hand side
213 rhs_wnv[:, n0:n1] = dm_wvv
214 # Extract the A matrix
215 A_wnn[:, n0:n1, :] += dm_wvv @ K_nn[n0:n1, :]
217 red_wnv = solve(A_wnn, rhs_wnv)
219 # Sum row blocks
220 dm_wvv = red_wnv.reshape((-1, Ni, 3, 3)).sum(axis=1)
221 return dm_wvv # type: ignore
223 def _get_dynamic_polarizability_imaginary_frequency(
224 self, frequencies: Array) -> np.ndarray:
225 raise NotImplementedError()