from __future__ import annotations
from warnings import warn
import itertools
import logging
from scipy.special import erfc
import numpy as np
from numba import jit
from dopyqo.colors import *
[docs]
def nuclear_repulsion_energy_ewald(
atom_positions: np.ndarray,
atomic_numbers: np.ndarray,
lattice_vectors: np.ndarray,
lattice_vectors_reciprocal: np.ndarray,
cell_volume: float,
gcut: float,
sigma: float | None = None,
rust_impl: bool = True,
) -> np.ndarray:
r"""Calculate the nuclear repulsion energy :math:`E` between all nuclei in a periodic lattice, i.e.,
.. math::
E = \frac{1}{2} \sum_I \sum_J \sum_{T}{}' \frac{Z_I Z_J}{\lvert R_I - R_J - T \rvert}
where :math:`'` denotes that the term :math:`T=0, I=J` is omitted, :math:`Z_I` are the charges
of the nuclei and :math:`R_I` their positions, and :math:`T` are all possible lattice translation
vectors. This is equivalent to
.. math::
E = \frac{1}{2} \sum_I Z_I \Phi_I(R_I)
where :math:`\Phi_I(r)` is the electrostatic potential generated by all periodic nuclei EXCEPT
the nucleus :math:`Z_I` at position :math:`R_I` (its periodic images are still included).
To calculate :math:`\Phi_I(r)` we define charge densities and solve the Poisson equation in
real- and reciprocal-space. The nuclei :math:`Z_I` are modelled by delta distributions at
positions :math:`R_I + T`. Additionally, Ewald's method is used, i.e., we add and subtract
Gaussian charge distributions at all :math:`R_I + T`, and subtract a homogeneous background
charge simulating the electron charges.
A short-ranged term including all delta distributions minus all Gaussian charges (except the
term where :math:`T=0` and :math:`I=J`) can be calculated in real-space. A long-ranged term
including all Gaussian charges at :math:`R_I + T` minus the constant background charge can be
calculated in reciprocal-space. The term describing the subtraction of the Gaussian charge at
:math:`R_I` is calculated in real-space and is the so-called self-energy term. The
:math:`G \to 0` limit in the reciprocal space calculation is also calculated separately and is
called the charged-system term. The value :math:`\sigma` determining the splitting between the
real-space and reciprocal-space term is calculated as done in Quantum ESPRESSO (see
PW/src/ewald.f90:
https://github.com/QEF/q-e/blob/de3035747f5d8f2ec9a67869827341ebb43f5b12/PW/src/ewald.f90)
We calculate :math:`E = E_S + E_L + E_{\text{self}} + E_{\text{charged}}` with:
Short-ranged term calculated in real-space:
.. math::
E_S = \frac{1}{2} \sum_{T}{}' \sum_I \sum_J
\frac{Z_I Z_J}{\lvert R_I - R_J - T \rvert}
\operatorname{erfc}\!\left(\lvert R_I - R_J - T \rvert \sqrt{\sigma}\right)
where :math:`'` denotes that the term :math:`T=0, I=J` is omitted.
Long-ranged term calculated in reciprocal-space:
.. math::
E_L = \frac{2\pi}{V} \sum_{k \neq 0} \sum_I \sum_J
\frac{Z_I Z_J}{\lvert k \rvert^2}
e^{i k \cdot (R_I - R_J)}
e^{-\lvert k \rvert^2 / (4\sigma)}
Self-energy from subtracting the Gaussian charge at position :math:`R_I` (in real-space):
.. math::
E_{\text{self}} = -\sum_I Z_I^2 \sqrt{\frac{\sigma}{\pi}}
Charged-system term from the :math:`G=0` term in the reciprocal-space sum:
.. math::
E_{\text{charged}} = -\frac{\pi}{2 V \sigma} \left(\sum_I Z_I\right)^2
Notes:
- When pseudopotentials are used, :math:`Z_I` are the valence charges, i.e. the charge of the
nucleus not modelled by the pseudopotential.
- This function should calculate the same Ewald energy as Quantum ESPRESSO, called
"ewald contribution" in the pw.x output file and "ewald" in the total_energy section in
the output-xml.
- Equivalent to the Ewald calculation in PySCF via pyscf.pbc.gto.cell.ewald.
References:
- Kittel, Appendix B: http://metal.elte.hu/~groma/Anyagtudomany/kittel.pdf
- PhysRevB.53.1814, equation (16): https://journals.aps.org/prb/pdf/10.1103/PhysRevB.53.1814
- https://courses.physics.illinois.edu/phys466/sp2013/projects/2003/Team2/ewald_text.htm
- S. W. de Leeuw et al. Proc. R. Soc. Lond. A 1980 373, 27-56, doi: 10.1098/rspa.1980.0135, equation (1.5)
- M. P. Allen and D. J. Tildesley "Computer Simulation of Solids", equation (5.20)
- http://micro.stanford.edu/mediawiki/images/4/46/Ewald_notes.pdf
- https://homepages.uc.edu/~becktl/tlb_ewald.pdf
- Quantum ESPRESSO source code in q-e/PW/src/ewald.f90
Args:
atom_positions (np.ndarray): Array of the coordinates :math:`R_I` of every atom. Shape (#atoms, 3)
atomic_numbers (np.ndarray): 1D array of atomic number of each atom. Shape (#atoms,)
lattice_vectors (np.ndarray): 2D array of lattice vectors. Each row is one lattice vector. Shape (3, 3)
lattice_vectors_reciprocal (np.ndarray): 2D array of reciprocal lattice vectors. Each row is one reciprocal lattice vector. Shape (3, 3)
cell_volume (float): Cell volume of the computational real space cell.
gcut (float): Cutoff of the reciprocal space sum in :math:`E_L`, i.e., :math:`\lvert k \rvert < G_{\text{cut}}`. In Quantum ESPRESSO the G-cutoff of the density is used.
sigma (float | None, optional): Value for :math:`\sigma` in the Ewald summation to split
real- and reciprocal-space sums. If None, computed to satisfy
.. math::
\sum_I 2 Z_I \sqrt{\frac{\sigma}{\pi}}
\operatorname{erfc}\!\left(
\sqrt{\frac{(4 G_{\text{cut}})^2}{4\sigma}}
\right) \leq 10^{-7}
Defaults to None.
Raises:
RuntimeError: If :math:`\sigma` could not be calculated.
Returns:
float: Ewald energy
"""
if rust_impl:
try:
import dopyqo_rs as calc_rs
except ImportError:
print(f"{ORANGE}Ewald warning: Could not import dopyqo_rs package. Falling back to python implementation.{RESET_COLOR}")
rust_impl = False # Set such that variable accurately represents if rust implementation is used
else: # No exception
logging.info("Using Rust implementation.")
return calc_rs.ewald(
atom_positions,
atomic_numbers,
lattice_vectors,
lattice_vectors_reciprocal,
cell_volume,
gcut,
)
if sigma is None:
sigma = 2.8
charge = np.sum(atomic_numbers)
gcutm = gcut**2
#
# choose sigma in order to have convergence in the sum over G
# upperbound is a safe upper bound for the error in the sum over G
#
while True:
if sigma <= 0.0:
raise RuntimeError("Optimal sigma for Ewald sum not found!")
upperbound = 2.0 * charge**2 * np.sqrt(sigma / np.pi) * erfc(np.sqrt(gcutm / 4.0 / sigma))
if upperbound > 1e-7:
sigma = sigma - 0.1
else:
break
logging.info("sigma %f", sigma)
e_short = 0.0 # real-space sum
converged_real = False
# n_max = 5
n_max_used = 0
alat = np.linalg.norm(lattice_vectors[0], ord=2)
t_vec_max_norm = 4.0 / np.sqrt(sigma) / alat
b_norms = np.linalg.norm(lattice_vectors_reciprocal, ord=2, axis=1)
n_max_x, n_max_y, n_max_z = (b_norms * t_vec_max_norm).astype(int) + 2
# Generate list or translation vectors ordered by their norm
# n_unordered = itertools.product(range(-n_max, n_max + 1), repeat=3)
logging.debug("Generating unordered translation vectors...")
t_vecs_unordered = [
[np.dot(lattice_vectors, np.array([nx, ny, nz])), [nx, ny, nz]]
for nx, ny, nz in itertools.product(
range(-n_max_x, n_max_x + 1),
range(-n_max_y, n_max_y + 1),
range(-n_max_z, n_max_z + 1),
)
# for nx, ny, nz in n_unordered
]
t_vecs = sorted(t_vecs_unordered, key=lambda x: np.linalg.norm(x[0], ord=2))
for t_vec, (nx, ny, nz) in t_vecs:
for i, pos_i in enumerate(atom_positions): # \sum_I
z_i = atomic_numbers[i]
for j, pos_j in enumerate(atom_positions): # \sum_J
if np.all(t_vec == 0.0) and i == j:
continue
z_j = atomic_numbers[j]
r = np.linalg.norm(pos_i - pos_j - t_vec, ord=2)
e_short += z_i * z_j * erfc(r * np.sqrt(sigma)) / r
# with this choice terms up to ZiZj*erfc(4) are counted (erfc(4)=1.5e-8)
if r >= 4 / np.sqrt(sigma):
# print(f"Stopping real-sum at {nx}/{ny}/{nz} (sigma={sigma})")
converged_real = True
n_max_used = np.max([np.abs(nx), np.abs(ny), np.abs(nz), n_max_used])
# break
# if converged_real:
# break
if converged_real:
logging.info("Real-space sum converged! Largest n was %i", n_max_used)
else:
logging.warning("Real-space sum NOT converged!")
e_short *= 1 / 2
e_long = 0.0 # reciprocal-space sum
# Estimate size of Miller indices grid using the cutoff-energy and reciprocal lattice vectors
# A simple, but not accurate enough, estimate would be gcutrho/|b_i| for i \in {1, 2, 3}
# Here we take the shape of the reciprocal lattice into account:
# max[ gcutrho/(|b_i| + \sum_{j \neq i} b_i . b_j ) ] for i \in {1, 2, 3}
# TODO: To use the same grid QE uses we have to read the Miller indices used for the density,
# which is probably saved in charge-density.dat.
# Alternatively check how pymatgen chooses the reciprocal grid in its Ewald calculation method
b1 = lattice_vectors_reciprocal[0]
b2 = lattice_vectors_reciprocal[1]
b3 = lattice_vectors_reciprocal[2]
mill_max = max(
[
round(abs(gcut / (np.linalg.norm(b1) + np.dot(b1, b2) + np.dot(b1, b3)))),
round(abs(gcut / (np.linalg.norm(b2) + np.dot(b2, b1) + np.dot(b2, b3)))),
round(abs(gcut / (np.linalg.norm(b3) + np.dot(b3, b1) + np.dot(b3, b2)))),
]
)
x, y, z = np.meshgrid(
np.arange(-mill_max, mill_max + 1),
np.arange(-mill_max, mill_max + 1),
np.arange(-mill_max, mill_max + 1),
indexing="ij",
)
mill_rho = np.stack((x.ravel(), y.ravel(), z.ravel()), axis=1)
k_vecs = np.einsum("ij, kj -> ki", lattice_vectors_reciprocal.T, mill_rho)
norms = np.linalg.norm(k_vecs, axis=1)
k_vecs = k_vecs[norms <= gcut]
for k_vec in k_vecs:
k2 = np.linalg.norm(k_vec, ord=2) ** 2
if np.isclose(k2, 0.0):
continue
structure_factor = np.sum(atomic_numbers * np.exp(1.0j * np.dot(atom_positions, k_vec)))
e_long += np.exp(-k2 / (4 * sigma)) * np.abs(structure_factor) ** 2 / k2
e_long *= 4.0 * np.pi / 2.0 / cell_volume
logging.debug("Reciprocal-space sum calculated!")
# Term arising from substracting the Gaussian charge distribution at R_I
# in real-space
e_self = -np.sum(atomic_numbers**2) * np.sqrt(sigma / np.pi)
# The analytic limit for G → 0 in the reciprocal sum
# The G=0 term in the reciprocal term is only partly cancelled by the uniform
# background charge. The remaining term is the following
# Also see Martin, 2020, "Electronic Structure" eq. (F.6) at the text afterwards
# which unfortunately misses a derivation
e_charged = -np.pi / cell_volume / sigma * np.sum(atomic_numbers) ** 2 / 2
# e_charged = -np.sum(atomic_numbers) ** 2 / sigma / 2 * np.pi / cell_volume
e_tot = e_short + e_long + e_self + e_charged
logging.info(
"e_short: %f, e_long: %f, e_self: %f, e_charged: %f, e_tot: %f",
e_short,
e_long,
e_self,
e_charged,
e_tot,
)
return e_tot
[docs]
def check_symmetry_one_body_matrix(matrix: np.ndarray):
"""Check if given matrix satisfies the symmetries of one-body matrices (hermitian)
Args:
matrix (np.ndarray): Matrix to check symmetries of
Returns:
bool: Bool specifying of symmetry is satisfied
"""
matrix_hermitian = matrix.T.conj()
allclose_hermitian = np.allclose(matrix, matrix_hermitian)
return allclose_hermitian
[docs]
def check_symmetry_two_body_matrix(matrix: np.ndarray):
"""Check if given matrix satisfies the symmetries of ERIs
Args:
matrix (np.ndarray): Matrix to check symmetries of
Returns:
tuple[bool, bool, bool]: Bools specifying if following symmetries are satisfied: swap symmetry, hermitian symmetry, hermitian+swap symmetry
"""
matrix_swap = matrix.swapaxes(0, 1).swapaxes(2, 3)
matrix_hermitian = matrix.T.conj()
matrix_hemitian_swap = matrix.T.conj().swapaxes(0, 1).swapaxes(2, 3)
allclose_swap = np.allclose(matrix, matrix_swap) # , rtol=1e-5, atol=1e-5)
allclose_hermitian = np.allclose(matrix, matrix_hermitian) # , rtol=1e-5, atol=1e-5)
allclose_hermitian_swap = np.allclose(matrix, matrix_hemitian_swap) # , rtol=1e-5, atol=1e-5)
allclose_ljki = None
allclose_ikjl = None
allclose_kilj = None
allclose_jlik = None
if np.allclose(matrix.real, matrix):
logging.info("Two-body matrix is real. Checking real-symmetries...")
allclose_hermitian = np.allclose(matrix, matrix.T) # ijkl = lkji
allclose_swap = np.allclose(matrix, matrix.transpose(1, 0, 3, 2)) # ijkl = jilk
allclose_ljki = np.allclose(matrix, matrix.transpose(3, 1, 2, 0)) # ijkl = ljki
allclose_ikjl = np.allclose(matrix, matrix.transpose(0, 2, 1, 3)) # ijkl = ikjl
allclose_hermitian_swap = np.allclose(matrix, matrix.T.transpose(1, 0, 3, 2)) # ijkl = klij
allclose_kilj = np.allclose(matrix, matrix.transpose(2, 0, 3, 1)) # ijkl = kilj
allclose_jlik = np.allclose(matrix, matrix.transpose(1, 3, 0, 2)) # ijkl = jlik
return (
allclose_swap,
allclose_hermitian,
allclose_hermitian_swap,
allclose_ljki,
allclose_ikjl,
allclose_kilj,
allclose_jlik,
)
[docs]
def iTj(p: np.ndarray, c_ip: np.ndarray) -> np.ndarray: # Calculates <i|T|j>
"""Calculate kinetic energy matrix in Hartree units in the Kohn-Sham basis
Args:
p (np.ndarray): Array of momentum vectors, shape (#waves, 3)
c_ip (np.ndarray): Array of coefficients describing the Kohn-Sham orbitals
in the plane wave basis, shape (#bands, #waves)
Returns:
np.ndarray: Kinetic energy matrix
"""
# Kinetic energy matrix in Hartree units
p_norm = p[:, 0] ** 2 + p[:, 1] ** 2 + p[:, 2] ** 2 # p^2
# 1/2 <i|p^2|j>
return 0.5 * np.einsum("ip, p, jp -> ij", c_ip.conjugate(), p_norm, c_ip)
def iUj(
p: np.ndarray,
c_ip: np.ndarray,
atom_positions: np.ndarray,
atomic_numbers: np.ndarray,
cell_volume: float,
) -> np.ndarray:
r"""Calculate nuclear interaction (:math:`U(r) = \sum_I Z_I / \lvert R_I - r \rvert`) matrix in Hartree units in the Kohn-Sham basis
Args:
p (np.ndarray): Array of momentum vectors
c_ip (np.ndarray): Array of coefficients describing the Kohn-Sham orbitals in the plane wave basis
atom_positions (np.ndarray): 2D array of positions of each atom. Shape (#atoms, 3)
atomic_numbers (np.ndarray): 1D array of atomic number of each ato. Shape (#atoms,)m
cell_volume (float): Cell volume of the computationen real space cell.
Returns:
np.ndarray: Nuclear interaction matrix
"""
# Nuclear interaction matrix in Hartree units
# Calculates <i|U|j>
# <p|U|q> = 4\pi / cell_volume \sum_I Z_I exp(-i (p-q) . R_I) 1/(q-p)²
# <i|U|j> = \sum_{p,q,p!=q} c_{p,i}^* c_{q,j} U_pq = 4\pi / cell_volume \sum_{p,q,p!=q} \sum_I Z_I c_{p,i}^* c_{q,j} exp(-i (p-q) . R_I) 1/(q-p)²
q = p
Z_I = atomic_numbers # shape (#atoms,)
R_I = atom_positions # shape (#atoms, 3)
q = p
p_minus_q = p[:, None] - q[None] # shape (#waves, #waves, 3)
# q[None] is of shape (1, #waves, 3) while p[:, None] is of shape (#waves, 1, 3)
# so p_minus_q[i,j] = p[i] - q[j]
p_minus_q = q[:, None] - p[None] # shape (#waves, #waves, 3)
p_minus_q_norm = np.linalg.norm(p_minus_q, ord=2, axis=2) # Has zeros on diagonal, shape (#waves, #waves)
p_minus_q_norm_filtered = p_minus_q_norm.copy() + np.eye(p_minus_q_norm.shape[0]) # Replace zeros on diagonal with ones
p_minus_q_dot_R = np.sum(p_minus_q[None] * R_I[:, None, None], axis=3) # sum over 3D-coordinates, shape (#atoms, #waves, #waves)
# \sum_I 4\pi/V e^{-i (G-G').R_I}
prefactor = 4 * np.pi / cell_volume * np.sum(Z_I[:, None, None] * np.exp(-1j * p_minus_q_dot_R), axis=0) # sum over atoms, shape (#waves, #waves)
one_over_p_minus_q_norm_squared = 1 / p_minus_q_norm_filtered**2
one_over_p_minus_q_norm_squared = one_over_p_minus_q_norm_squared.copy() - np.eye(one_over_p_minus_q_norm_squared.shape[0])
u_mat = prefactor * one_over_p_minus_q_norm_squared
assert (u_mat.diagonal() == 0.0).all()
return c_ip.conj() @ u_mat @ c_ip.T