Coverage for strongcoca / calculators / polarizability_calculator.py: 100%
92 statements
« prev ^ index » next coverage.py v7.13.4, created at 2026-07-25 16:26 +0000
1import logging
3from typing import Union
4import numpy as np
5from numpy.linalg import solve
6from .. import CoupledSystem
7from ..response.utilities import Broadening, GaussianBroadening
8from ..types import Array
9from ..units import eV_to_au
10from .base_calculator import BaseCalculator
11from ..compute_config import get_compute_config
13logger = logging.getLogger(__name__)
16class PolarizabilityCalculator(BaseCalculator):
17 """Instances of this class enable the calculation of correlation energy
18 and spectrum of a coupled system, all polarizable objects in which have an
19 internal representation in the form of a polarizability tensor.
21 Spectra calculated with this calculator derive broadening from
22 the underlying polarizable units. Any additional or separate broadening
23 is not added.
25 Parameters
26 ----------
27 coupled_system
28 Coupled system for which to carry out calculations.
29 imaginary_frequencies
30 Frequency grid along imaginary axis used for correlation energy
31 calculation; eV by default, optionally atomic units (see :attr:`units`).
32 units
33 `eV` to specify energies in eV or `au` to specify inputs in atomic units.
35 This parameter determines whether conversion should be performed during
36 initialization and has no effect on instance methods and variables.
37 name
38 Name of response.
40 Examples
41 --------
43 The following code snippet shows the calculation of the correlation energy
44 for a simple demo system:
46 >>> import numpy as np
47 >>> from strongcoca import CoupledSystem, PolarizableUnit
48 >>> from strongcoca.response import build_random_casida
49 >>> from strongcoca.calculators import PolarizabilityCalculator
50 >>>
51 >>> # construct coupled system
52 >>> response = build_random_casida(n_states=2)
53 >>> pu1 = PolarizableUnit(response, [0, 0, 0])
54 >>> pu2 = PolarizableUnit(response, [2, 0, 0])
55 >>> cs = CoupledSystem([pu1, pu2])
56 >>>
57 >>> # set up calculator
58 >>> calc = PolarizabilityCalculator(cs, np.linspace(0, 10, 100))
59 >>>
60 >>> # ... and compute the correlation energy
61 >>> calc.get_correlation_energy()
62 -0.72084000362...
63 """
65 def __init__(self,
66 coupled_system: CoupledSystem,
67 imaginary_frequencies: Array,
68 units: str = 'eV',
69 name: str = 'PolarizabilityCalculator') -> None:
70 super().__init__(coupled_system, broadening=Broadening(), name=name)
72 imaginary_frequencies = np.array(imaginary_frequencies, dtype=float)
73 if units == 'eV':
74 imaginary_frequencies = imaginary_frequencies * eV_to_au
75 elif units != 'au':
76 raise ValueError(f"units has to be 'eV' or 'au', not '{units}'")
78 self._ifreq_w = imaginary_frequencies
79 self._K_nn: Union[np.ndarray, None] = None
80 # Check frequency grid
81 difreq_w = self._ifreq_w[1:] - self._ifreq_w[:-1]
82 self._difreq = difreq_w[0]
83 if not np.allclose(difreq_w, self._difreq):
84 raise ValueError('Frequency grid needs to be equally spaced.')
86 def _build_coupling_matrix(self) -> np.ndarray:
87 """Builds coupling matrix.
89 Coupling matrix is formed of blocks:
90 Each block is a 3x3 dipole-dipole tensor and there is NxN blocks
91 in total (N: number of units in coupled system).
92 The 3x3 blocks on the diagonal are zero.
93 """
94 Ni = len(self._coupled_system)
95 N3 = 3 * Ni
97 # Collect positions: (Ni, 3)
98 pos_iv = np.array([pu._position for pu in self._coupled_system])
100 # Pairwise displacement vectors R[i,j] = pos[j] - pos[i], shape (Ni, Ni, 3)
101 R_ijv = pos_iv[np.newaxis, :] - pos_iv[:, np.newaxis]
103 # Pairwise distances, shape (Ni, Ni).
104 # Set diagonal to 1.0 to avoid division by zero; those blocks are zeroed later.
105 r_ij = np.linalg.norm(R_ijv, axis=-1)
106 np.fill_diagonal(r_ij, 1.0)
108 # Dipole-dipole tensor for all pairs: T[i,j] = I/r³ - 3 R_ij R_ij^T / r⁵
109 # Shape: (Ni, Ni, 3, 3)
110 T_ijvw = (np.eye(3) / r_ij[:, :, np.newaxis, np.newaxis] ** 3
111 - 3.0 * (R_ijv[:, :, :, np.newaxis] * R_ijv[:, :, np.newaxis, :])
112 / r_ij[:, :, np.newaxis, np.newaxis] ** 5)
114 # Zero diagonal blocks (self-interaction)
115 T_ijvw[np.arange(Ni), np.arange(Ni)] = 0.0
117 # Reshape (Ni, Ni, 3, 3) → (N3, N3): index [i,j,v,w] → [3i+v, 3j+w]
118 return T_ijvw.transpose(0, 2, 1, 3).reshape(N3, N3).astype(complex) # type: ignore
120 def _calculate_correlation_energy(self) -> float:
121 """Returns the correlation energy of the coupled system in atomic units.
122 """
124 Ni = len(self._coupled_system)
125 if Ni == 0:
126 return 0.0
127 K_nn = self.coupling_matrix
128 W = len(self._ifreq_w)
130 # Collect per-unit polarizabilities on CPU: calling response objects is
131 # inherently sequential Python and cannot be moved to GPU.
132 dm_wNvv = np.empty((W, Ni, 3, 3), dtype=complex)
133 for i, pu in enumerate(self._coupled_system):
134 dm_wNvv[:, i] = pu._get_dynamic_polarizability_imaginary_frequency(self._ifreq_w)
136 compute_config = get_compute_config()
137 gpu_backend = compute_config.backend
138 if gpu_backend != 'none': # pragma: no cover
139 # Fused GPU path: matmul → D = I−chi_K → slogdet → trace.
140 # K_nn is transferred once; chi_K is never moved back to CPU.
141 # Frequency chunking keeps peak VRAM at O(W_chunk × N²).
142 from .gpu import energy_trace as _gpu_energy_trace
143 trace_logD_w, trace_chiK_w = _gpu_energy_trace(
144 dm_wNvv, K_nn, gpu_backend, compute_config.precision)
145 else:
146 # Vectorized CPU path: single batched matmul then slogdet.
147 K_block = K_nn.reshape(Ni, 3, 3 * Ni)
148 chi_K_wnn = (dm_wNvv @ K_block).reshape(W, 3 * Ni, 3 * Ni)
149 D_wnn = np.eye(3 * Ni)[None] - chi_K_wnn
150 sign_w, logabsdet_w = np.linalg.slogdet(D_wnn)
151 trace_logD_w = np.log(sign_w) + logabsdet_w
152 # tr(chi_K†) = conj(tr(chi_K))
153 trace_chiK_w = np.diagonal(chi_K_wnn, axis1=-2, axis2=-1).sum(axis=-1)
155 integrand_w = trace_logD_w + trace_chiK_w.conj()
156 integral = np.sum(integrand_w) * self._difreq
157 energy: float = float(np.real(integral)) / (2 * np.pi)
158 return energy
160 @property
161 def coupling_matrix(self) -> np.ndarray:
162 r""" The coupling matrix
164 .. math::
165 \boldsymbol{T}_{ij} =
166 \frac{3\boldsymbol{r}^{(ij)}(\boldsymbol{r}^{(ij)})^\text{T}}{\left|r^{(ij)}\right|^5}
167 - \frac{1}{\left|r^{(ij)}\right|^3}
169 as a matrix where the rows and columns each correspond to the system index
170 :math:`i` or :math:`j` *and* the Cartesian direction.
172 This quantity is calculated on first use and then buffered.
173 """
174 if self._K_nn is None:
175 self._K_nn = self._build_coupling_matrix()
176 return self._K_nn
178 def _get_dynamic_polarizability(self, frequencies: Array) -> np.ndarray:
179 if any(isinstance(pu.broadening, GaussianBroadening)
180 for pu in self._coupled_system):
181 raise NotImplementedError(
182 'Gaussian broadening is not supported in PolarizabilityCalculator. '
183 'See issue #42.')
185 freq_w = np.array(frequencies)
186 Ni = len(self._coupled_system)
188 # Collect per-unit polarizabilities for ALL frequencies at once.
189 # Calling _get_dynamic_polarizability with all frequencies is as fast as
190 # calling it with one frequency due to internal vectorization, so collecting
191 # upfront avoids repeating the call W times (once per chunk in the old code).
192 dm_wNvv = np.empty((len(freq_w), Ni, 3, 3), dtype=complex)
193 for i, pu in enumerate(self._coupled_system):
194 dm_wNvv[:, i] = pu._get_dynamic_polarizability(freq_w)
196 compute_config = get_compute_config()
197 gpu_backend = compute_config.backend
198 if gpu_backend != 'none': # pragma: no cover
199 # GPU path: fused A construction + solve on GPU, with GPU-memory-aware
200 # frequency chunking. A matrix never transferred to CPU.
201 from .gpu import spectrum_solve as _gpu_spectrum_solve
202 return _gpu_spectrum_solve(dm_wNvv, self.coupling_matrix, gpu_backend,
203 compute_config.precision)
205 # CPU path: chunk over frequencies to keep A matrix within RAM budget.
206 dm_wvv = np.zeros((len(freq_w), 3, 3), dtype=complex)
207 to_MiB = 1024 ** -2
208 mem_limit = compute_config.max_solve_mem / to_MiB
209 syssize = 3 ** 2 * Ni ** 2 * 8 # size of A matrix per frequency in bytes
210 chunksize = max(1, int(mem_limit) // syssize)
211 for indices in np.array_split(np.arange(len(freq_w)),
212 (len(freq_w) + chunksize - 1) // chunksize):
213 dm_wvv[indices] = self._get_dynamic_polarizability_chunk(dm_wNvv[indices])
214 return dm_wvv
216 def _get_dynamic_polarizability_chunk(self, dm_wNvv: np.ndarray) -> np.ndarray:
217 """CPU path: build A and solve for a pre-collected frequency chunk.
219 Parameters
220 ----------
221 dm_wNvv
222 Per-unit polarizabilities for this chunk, shape ``(W, Ni, 3, 3)``.
223 """
224 Ni = len(self._coupled_system)
225 K_nn = self.coupling_matrix
226 K_block = K_nn.reshape(Ni, 3, 3 * Ni)
227 W = len(dm_wNvv)
228 N3 = 3 * Ni
230 # Build A = I + chi_K by looping over units with all W frequencies batched.
231 # This makes Ni BLAS calls of shape (W, 3, 3) @ (3, N3) instead of W*Ni
232 # calls of shape (1, 3, 3) @ (3, N3), which is significantly faster.
233 A_wnn = np.tile(np.eye(N3, dtype=complex), (W, 1, 1))
234 for i in range(Ni):
235 n0 = i * 3
236 A_wnn[:, n0:n0 + 3, :] += dm_wNvv[:, i] @ K_block[i]
238 rhs_wnv = np.ascontiguousarray(dm_wNvv.reshape(W, N3, 3))
239 red_wnv = solve(A_wnn, rhs_wnv)
240 return red_wnv.reshape(W, Ni, 3, 3).sum(axis=1) # type: ignore
242 def _get_dynamic_polarizability_imaginary_frequency(
243 self, frequencies: Array) -> np.ndarray:
244 raise NotImplementedError()