import logging
from typing import Union
import numpy as np
from numpy.linalg import solve
from .. import CoupledSystem
from ..response.utilities import Broadening, GaussianBroadening
from ..types import Array
from ..units import eV_to_au
from .base_calculator import BaseCalculator
from ..compute_config import get_compute_config
logger = logging.getLogger(__name__)
[docs]
class PolarizabilityCalculator(BaseCalculator):
"""Instances of this class enable the calculation of correlation energy
and spectrum of a coupled system, all polarizable objects in which have an
internal representation in the form of a polarizability tensor.
Spectra calculated with this calculator derive broadening from
the underlying polarizable units. Any additional or separate broadening
is not added.
Parameters
----------
coupled_system
Coupled system for which to carry out calculations.
imaginary_frequencies
Frequency grid along imaginary axis used for correlation energy
calculation; eV by default, optionally atomic units (see :attr:`units`).
units
`eV` to specify energies in eV 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.
Examples
--------
The following code snippet shows the calculation of the correlation energy
for a simple demo system:
>>> import numpy as np
>>> from strongcoca import CoupledSystem, PolarizableUnit
>>> from strongcoca.response import build_random_casida
>>> from strongcoca.calculators import PolarizabilityCalculator
>>>
>>> # construct coupled system
>>> response = build_random_casida(n_states=2)
>>> pu1 = PolarizableUnit(response, [0, 0, 0])
>>> pu2 = PolarizableUnit(response, [2, 0, 0])
>>> cs = CoupledSystem([pu1, pu2])
>>>
>>> # set up calculator
>>> calc = PolarizabilityCalculator(cs, np.linspace(0, 10, 100))
>>>
>>> # ... and compute the correlation energy
>>> calc.get_correlation_energy()
-0.72084000362...
"""
def __init__(self,
coupled_system: CoupledSystem,
imaginary_frequencies: Array,
units: str = 'eV',
name: str = 'PolarizabilityCalculator') -> None:
super().__init__(coupled_system, broadening=Broadening(), name=name)
imaginary_frequencies = np.asarray(imaginary_frequencies)
if units == 'eV':
imaginary_frequencies = imaginary_frequencies * eV_to_au
elif units != 'au':
raise ValueError(f"units has to be 'eV' or 'au', not '{units}'")
self._ifreq_w = imaginary_frequencies
self._K_nn: Union[np.ndarray, None] = None
# Check frequency grid
difreq_w = self._ifreq_w[1:] - self._ifreq_w[:-1] # type: ignore
self._difreq = difreq_w[0]
if not np.allclose(difreq_w, self._difreq):
raise ValueError('Frequency grid needs to be equally spaced.')
def _build_coupling_matrix(self) -> np.ndarray:
"""Builds coupling matrix.
Coupling matrix is formed of blocks:
Each block is a 3x3 dipole-dipole tensor and there is NxN blocks
in total (N: number of units in coupled system).
The 3x3 blocks on the diagonal are zero.
"""
Ni = len(self._coupled_system)
N3 = 3 * Ni
# Collect positions: (Ni, 3)
pos_iv = np.array([pu._position for pu in self._coupled_system])
# Pairwise displacement vectors R[i,j] = pos[j] - pos[i], shape (Ni, Ni, 3)
R_ijv = pos_iv[np.newaxis, :] - pos_iv[:, np.newaxis]
# Pairwise distances, shape (Ni, Ni).
# Set diagonal to 1.0 to avoid division by zero; those blocks are zeroed later.
r_ij = np.linalg.norm(R_ijv, axis=-1)
np.fill_diagonal(r_ij, 1.0)
# Dipole-dipole tensor for all pairs: T[i,j] = I/r³ - 3 R_ij R_ij^T / r⁵
# Shape: (Ni, Ni, 3, 3)
T_ijvw = (np.eye(3) / r_ij[:, :, np.newaxis, np.newaxis] ** 3
- 3.0 * (R_ijv[:, :, :, np.newaxis] * R_ijv[:, :, np.newaxis, :])
/ r_ij[:, :, np.newaxis, np.newaxis] ** 5)
# Zero diagonal blocks (self-interaction)
T_ijvw[np.arange(Ni), np.arange(Ni)] = 0.0
# Reshape (Ni, Ni, 3, 3) → (N3, N3): index [i,j,v,w] → [3i+v, 3j+w]
return T_ijvw.transpose(0, 2, 1, 3).reshape(N3, N3).astype(complex)
def _calculate_correlation_energy(self) -> float:
"""Returns the correlation energy of the coupled system in atomic units.
"""
Ni = len(self._coupled_system)
if Ni == 0:
return 0.0
K_nn = self.coupling_matrix
W = len(self._ifreq_w)
# Collect per-unit polarizabilities on CPU: calling response objects is
# inherently sequential Python and cannot be moved to GPU.
dm_wNvv = np.empty((W, Ni, 3, 3), dtype=complex)
for i, pu in enumerate(self._coupled_system):
dm_wNvv[:, i] = pu._get_dynamic_polarizability_imaginary_frequency(self._ifreq_w)
compute_config = get_compute_config()
gpu_backend = compute_config.backend
if gpu_backend != 'none': # pragma: no cover
# Fused GPU path: matmul → D = I−chi_K → slogdet → trace.
# K_nn is transferred once; chi_K is never moved back to CPU.
# Frequency chunking keeps peak VRAM at O(W_chunk × N²).
from .gpu import energy_trace as _gpu_energy_trace
trace_logD_w, trace_chiK_w = _gpu_energy_trace(
dm_wNvv, K_nn, gpu_backend, compute_config.precision)
else:
# Vectorized CPU path: single batched matmul then slogdet.
K_block = K_nn.reshape(Ni, 3, 3 * Ni)
chi_K_wnn = (dm_wNvv @ K_block).reshape(W, 3 * Ni, 3 * Ni)
D_wnn = np.eye(3 * Ni)[None] - chi_K_wnn
sign_w, logabsdet_w = np.linalg.slogdet(D_wnn)
trace_logD_w = np.log(sign_w) + logabsdet_w
# tr(chi_K†) = conj(tr(chi_K))
trace_chiK_w = np.diagonal(chi_K_wnn, axis1=-2, axis2=-1).sum(axis=-1)
integrand_w = trace_logD_w + trace_chiK_w.conj()
integral = np.sum(integrand_w) * self._difreq
energy: float = float(np.real(integral)) / (2 * np.pi)
return energy
@property
def coupling_matrix(self) -> np.ndarray:
r""" The coupling matrix
.. math::
\boldsymbol{T}_{ij} =
\frac{3\boldsymbol{r}^{(ij)}(\boldsymbol{r}^{(ij)})^\text{T}}{\left|r^{(ij)}\right|^5}
- \frac{1}{\left|r^{(ij)}\right|^3}
as a matrix where the rows and columns each correspond to the system index
:math:`i` or :math:`j` *and* the Cartesian direction.
This quantity is calculated on first use and then buffered.
"""
if self._K_nn is None:
self._K_nn = self._build_coupling_matrix()
return self._K_nn
def _get_dynamic_polarizability(self, frequencies: Array) -> np.ndarray:
if any(isinstance(pu.broadening, GaussianBroadening)
for pu in self._coupled_system):
raise NotImplementedError(
'Gaussian broadening is not supported in PolarizabilityCalculator. '
'See issue #42.')
freq_w = np.array(frequencies)
Ni = len(self._coupled_system)
# Collect per-unit polarizabilities for ALL frequencies at once.
# Calling _get_dynamic_polarizability with all frequencies is as fast as
# calling it with one frequency due to internal vectorization, so collecting
# upfront avoids repeating the call W times (once per chunk in the old code).
dm_wNvv = np.empty((len(freq_w), Ni, 3, 3), dtype=complex)
for i, pu in enumerate(self._coupled_system):
dm_wNvv[:, i] = pu._get_dynamic_polarizability(freq_w)
compute_config = get_compute_config()
gpu_backend = compute_config.backend
if gpu_backend != 'none': # pragma: no cover
# GPU path: fused A construction + solve on GPU, with GPU-memory-aware
# frequency chunking. A matrix never transferred to CPU.
from .gpu import spectrum_solve as _gpu_spectrum_solve
return _gpu_spectrum_solve(dm_wNvv, self.coupling_matrix, gpu_backend,
compute_config.precision)
# CPU path: chunk over frequencies to keep A matrix within RAM budget.
dm_wvv = np.zeros((len(freq_w), 3, 3), dtype=complex)
to_MiB = 1024 ** -2
mem_limit = compute_config.max_solve_mem / to_MiB
syssize = 3 ** 2 * Ni ** 2 * 8 # size of A matrix per frequency in bytes
chunksize = max(1, int(mem_limit) // syssize)
for indices in np.array_split(np.arange(len(freq_w)),
(len(freq_w) + chunksize - 1) // chunksize):
dm_wvv[indices] = self._get_dynamic_polarizability_chunk(dm_wNvv[indices])
return dm_wvv
def _get_dynamic_polarizability_chunk(self, dm_wNvv: np.ndarray) -> np.ndarray:
"""CPU path: build A and solve for a pre-collected frequency chunk.
Parameters
----------
dm_wNvv
Per-unit polarizabilities for this chunk, shape ``(W, Ni, 3, 3)``.
"""
Ni = len(self._coupled_system)
K_nn = self.coupling_matrix
K_block = K_nn.reshape(Ni, 3, 3 * Ni)
W = len(dm_wNvv)
N3 = 3 * Ni
# Build A = I + chi_K by looping over units with all W frequencies batched.
# This makes Ni BLAS calls of shape (W, 3, 3) @ (3, N3) instead of W*Ni
# calls of shape (1, 3, 3) @ (3, N3), which is significantly faster.
A_wnn = np.tile(np.eye(N3, dtype=complex), (W, 1, 1))
for i in range(Ni):
n0 = i * 3
A_wnn[:, n0:n0 + 3, :] += dm_wNvv[:, i] @ K_block[i]
rhs_wnv = np.ascontiguousarray(dm_wNvv.reshape(W, N3, 3))
red_wnv = solve(A_wnn, rhs_wnv)
return red_wnv.reshape(W, Ni, 3, 3).sum(axis=1)
def _get_dynamic_polarizability_imaginary_frequency(
self, frequencies: Array) -> np.ndarray:
raise NotImplementedError()
