import logging
from typing import Union
import numpy as np
from numpy.linalg import solve
from scipy.linalg import logm
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 .utilities import get_dipole_dipole_tensor
from .. import env
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)
K_nn = np.zeros((3 * Ni, 3 * Ni), dtype=complex)
for i, pu1 in enumerate(self._coupled_system):
# index range of the block
n0 = i * 3
n1 = n0 + 3
for j in range(i + 1, Ni):
# index range of the block
m0 = j * 3
m1 = m0 + 3
pu2 = self._coupled_system[j]
# todo: include cell and pbc
distance_v = pu2._position - pu1._position
K_vv = get_dipole_dipole_tensor(distance_v)
K_nn[n0:n1, m0:m1] = K_vv
K_nn[m0:m1, n0:n1] = K_vv.conj()
return K_nn
def _calculate_correlation_energy(self) -> float:
"""Returns the correlation energy of the coupled system in atomic units.
"""
# build coupling matrix
Ni = len(self._coupled_system)
K_nn = self.coupling_matrix
# Calculate the matrix product of
# 1) a block diagonal matrix formed of 3x3 polarizability blocks and
# 2) the coupling matrix (formed of 3x3 blocks too).
# The following does not construct the block diagonal matrix in full but
# does the matrix products block-by-block as the blocks are calculated.
chi_K_wnn = np.zeros((len(self._ifreq_w), 3 * Ni, 3 * Ni), dtype=complex)
for i, pu in enumerate(self._coupled_system):
# calculate polarizability block
dm_wvv = pu._get_dynamic_polarizability_imaginary_frequency(self._ifreq_w)
# index range of the block
n0 = i * 3
n1 = n0 + 3
# take matrix product with the coupling matrix
chi_K_wnn[:, n0:n1, :] = np.dot(dm_wvv, K_nn[n0:n1])
# carry out integration over imaginary frequency axis
integral = 0
# TODO: this frequency loop is slow; reformulate without it to speed up
for chi_K_nn in chi_K_wnn:
D_nn = np.eye(3 * Ni) - chi_K_nn
dint = np.trace(logm(D_nn) + chi_K_nn.T.conj())
integral += dint
integral *= self._difreq
energy = np.real(integral) / (2 * np.pi) # type: float
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.')
# build coupling matrix
freq_w = np.array(frequencies)
dm_wvv = np.zeros((len(freq_w), 3, 3), dtype=complex)
# The linear solve is done frequency by frequency.
# For large systems, one should not do all frequencies at once
# as the intermediate A matrix takes up too much memory.
# Limit the matrix A to maxsize elements.
to_MiB = 1024 ** -2
mem_limit = env.get_float('MAX_SOLVE_MEM') / to_MiB
syssize = 3 ** 2 * len(self._coupled_system) ** 2 * 8 # Size of A matrix in bytes
chunksize = int(mem_limit) // syssize # Maximum number of frequencies at a time
if chunksize < 1:
# The system is so huge so even one frequency at a time violates our
# memory limit. We have to live with it.
chunksize = 1
# This loop will just iterate once for small enough systems
for indices in np.array_split(np.arange(0, len(freq_w)),
(len(frequencies) + chunksize - 1) // chunksize):
dm_wvv[indices] = self._get_dynamic_polarizability_chunk(freq_w[indices])
return dm_wvv
def _get_dynamic_polarizability_chunk(self, frequencies: np.ndarray) -> np.ndarray:
""" Calculate the dynamic polarizability for a subset of frequencies. """
Ni = len(self._coupled_system)
K_nn = self.coupling_matrix
freq_w = np.array(frequencies)
# Linear solve
rhs_wnv = np.zeros((len(freq_w), Ni * 3, 3), dtype=complex)
A_wnn = np.zeros((len(freq_w), Ni * 3, Ni * 3), dtype=complex)
A_wnn += np.eye(Ni * 3)
for i, pu in enumerate(self._coupled_system):
dm_wvv = pu._get_dynamic_polarizability(freq_w)
n0 = i * 3
n1 = n0 + 3
# Extract the right-hand side
rhs_wnv[:, n0:n1] = dm_wvv
# Extract the A matrix
A_wnn[:, n0:n1, :] += dm_wvv @ K_nn[n0:n1, :]
red_wnv = solve(A_wnn, rhs_wnv)
# Sum row blocks
dm_wvv = red_wnv.reshape((-1, Ni, 3, 3)).sum(axis=1)
return dm_wvv # type: ignore
def _get_dynamic_polarizability_imaginary_frequency(
self, frequencies: Array) -> np.ndarray:
raise NotImplementedError()
