import numpy as np
import pyqmc
import pyqmc.energy
import pyqmc.gpu as gpu
[docs]class Ewald:
r"""
The Ewald summation computes the Coulomb energy of a periodic arrangement of charges.
The sum (of pair interactions) is separated into real-space (short range) and reciprocal-space (long range) sums, each of which converges quickly.
The separation is determined by the parameter :math:`\alpha`
The Ewald separation is:
.. math:: E_{\rm Coulomb} = E_{\rm real\ space} + E_{\rm reciprocal\ space} + E_{\rm self} + E_{\rm charged}
.. math:: E_{\rm real\ space} = \frac{1}{2} {\sum_{\vec{n}}}^\dagger \sum_{i=1}^N \sum_{j=1}^N q_i q_j \frac{{\rm erfc}(\alpha |\vec{x}_{ij}+\vec{n}|)}{|\vec{x}_{ij}+\vec{n}|}
:math:`\qquad\qquad` (The :math:`{\sum}^\dagger` means to exclude the self-terms, i.e. :math:`\vec{n}=0, \, i=j`)
.. math:: E_{\rm reciprocal\ space} = \frac{4\pi}{V} \frac{1}{2} \sum_{k \ne 0} \frac{1}{k^2} e^{-\frac{k^2}{4\alpha^2}} \left| \sum_{i=1}^N q_i e^{-i\vec{k}\cdot\vec{x}_i} \right|^2
.. math:: E_{\rm self} = -\frac{\alpha}{\sqrt{\pi}} \sum_{i=1}^N q_i^2
\qquad
E_{\rm charged} = -\frac{\pi}{2V\alpha^2} \left| \sum_{i=1}^N q_i \right|^2
The self energy corrects for a self-interaction included in the reciprocal-space term, and the charged-system correction is only necessary for systems with nonzero net charge.
In our implementation, the parts are further split into electron-electron, electron-ion, and ion-ion contributions. We use lower-case summation indices for electrons, and upper case for ions.
For ease of notation (and reading the code), let pair distances be denoted by
.. math:: r_{ijn} = |\vec{x}_{ij}+\vec{n}|
.. math:: r_{iIn} = |\vec{x}_{iI}+\vec{n}|
.. math:: r_{IJn} = |\vec{x}_{IJ}+\vec{n}|
Real space terms, arranged to sum over each pair only once:
.. math:: E_{\rm real\ space}^{\text{ion-ion}} = \sum_{\vec{n}} \sum_{I<J}^{N_{ion}} Z_I Z_J \frac{{\rm erfc}(\alpha r_{IJn})}{r_{IJn}}
+ \frac{1}{2} \sum_{I=1}^{N_{ion}} Z_I^2 C_{\rm self\ image}
.. math:: E_{\rm real\ space}^{ee} = \sum_{\vec{n}} \sum_{i<j}^{N_e} \frac{{\rm erfc}(\alpha r_{ijn})}{r_{ijn}}
+ \frac{N_e}{2} C_{\rm self\ image}
.. math:: E_{\rm real\ space}^{e\text{-ion}} = {\sum_{\vec{n}}} \sum_{i=1}^{N_e} \sum_{I=1}^{N_{ion}} -Z_I \frac{{\rm erfc}(\alpha r_{iIn})}{r_{iIn}}
where the interactions between a particle and its own image in other cells is represented by the sum
.. math:: C_{\rm self\ image} = \frac{1}{2} \sum_{\vec{n} \ne 0} \frac{{\rm erfc}(\alpha |\vec{n}|)}{|\vec{n}|}
Reciprocal space terms, summing over :math:`\vec{G}>0` -- reciprocal lattice vectors in positive octants. In other words, only one of :math:`\vec{G}` and :math:`-\vec{G}` is included in the sum, and :math:`\vec{G}=0` is not included.
.. math:: E_{\rm reciprocal\ space}^{\text{ion-ion}} = \sum_{\vec{G} > 0 } W_G \left| \sum_{I=1}^{N_{ion}} Z_I e^{-i\vec{G}\cdot\vec{x}_I} \right|^2
.. math:: E_{\rm reciprocal\ space}^{ee} = \sum_{\vec{G}>0} W_G \left| \sum_{i=1}^{N_e} e^{-i\vec{k}\cdot\vec{x}_i} \right|^2
.. math:: E_{\rm reciprocal\ space}^{e\text{-ion}} = \sum_{\vec{G}>0} W_G {\rm Re} \left[ 2 \sum_{i=1}^{N_e} \sum_{I=1}^{N_{ion}} -Z_I e^{-i\vec{k}\cdot\vec{x}_i} e^{i\vec{k}\cdot\vec{x}_I} \right]
where `gweight` is a factor that doesn't depend on the coordinates:
.. math:: W_G = \frac{4\pi}{V |\vec{G}|^2} e^{- \frac{|\vec{G}|^2}{ 4\alpha^2}}
Self energy:
.. math:: E_{\rm self}^{e} = - \frac{\alpha}{\sqrt{\pi}} N_e
\qquad
E_{\rm self}^{\rm ion} = - \frac{\alpha}{\sqrt{\pi}} \sum_{I=1}^{N_{ion}} Z_I^2
Charged-system energy:
.. math:: E_{\rm charged}^{ee} = - \frac{\pi}{2V\alpha^2} N_e^2
\qquad
E_{\rm charged}^{e\text{-ion}} = \frac{\pi}{2V\alpha^2} 2 N_e \sum_{I=1}^{N_{ion}} Z_I
.. math:: E_{\rm charged}^{\text{ion-ion}} = - \frac{\pi}{2V\alpha^2} \left[ \sum_{I=1}^{N_{ion}} Z_I^2 + 2 \sum_{I<J}^{N_{ion}} Z_I Z_J \right]
"""
def __init__(self, cell, ewald_gmax=200, nlatvec=1):
"""
:parameter cell: pyscf Cell object (simulation cell)
:parameter int ewald_gmax: how far to take reciprocal sum; probably never needs to be changed.
:parameter int nlatvec: how far to take real-space sum; probably never needs to be changed.
"""
self.nelec = np.array(cell.nelec)
self.atom_coords = cell.atom_coords()
self.atom_charges = gpu.cp.asarray(cell.atom_charges())
self.latvec = cell.lattice_vectors()
self.set_lattice_displacements(nlatvec)
self.set_up_reciprocal_ewald_sum(ewald_gmax)
[docs] def set_lattice_displacements(self, nlatvec):
"""
Generates list of lattice-vector displacements to add together for real-space sum
:parameter int nlatvec: sum goes from `-nlatvec` to `nlatvec` in each lattice direction.
"""
XYZ = np.meshgrid(*[np.arange(-nlatvec, nlatvec + 1)] * 3, indexing="ij")
xyz = np.stack(XYZ, axis=-1).reshape((-1, 3))
self.lattice_displacements = gpu.cp.asarray(np.dot(xyz, self.latvec))
[docs] def set_up_reciprocal_ewald_sum(self, ewald_gmax):
r"""
Determine parameters for Ewald sums.
:math:`\alpha` determines the partitioning of the real and reciprocal-space parts.
We define a weight `gweight` for the part of the reciprocal-space sums that doesn't depend on the coordinates:
.. math:: W_G = \frac{4\pi}{V |\vec{G}|^2} e^{- \frac{|\vec{G}|^2}{ 4\alpha^2}}
:parameter int ewald_gmax: max number of reciprocal lattice vectors to check away from 0
"""
cellvolume = np.linalg.det(self.latvec)
recvec = np.linalg.inv(self.latvec).T
# Determine alpha
smallestheight = np.amin(1 / np.linalg.norm(recvec, axis=1))
self.alpha = 5.0 / smallestheight
# Determine G points to include in reciprocal Ewald sum
gptsXpos = gpu.cp.meshgrid(
gpu.cp.arange(1, ewald_gmax + 1),
*[gpu.cp.arange(-ewald_gmax, ewald_gmax + 1)] * 2,
indexing="ij"
)
zero = gpu.cp.asarray([0])
gptsX0Ypos = gpu.cp.meshgrid(
zero,
gpu.cp.arange(1, ewald_gmax + 1),
gpu.cp.arange(-ewald_gmax, ewald_gmax + 1),
indexing="ij",
)
gptsX0Y0Zpos = gpu.cp.meshgrid(
zero, zero, gpu.cp.arange(1, ewald_gmax + 1), indexing="ij"
)
gs = zip(
*[
select_big(x, cellvolume, recvec, self.alpha)
for x in (gptsXpos, gptsX0Ypos, gptsX0Y0Zpos)
]
)
self.gpoints, self.gweight = [gpu.cp.concatenate(x, axis=0) for x in gs]
self.set_ewald_constants(cellvolume)
[docs] def set_ewald_constants(self, cellvolume):
r"""
Compute Ewald constants (independent of particle positions): self energy and charged-system energy. Here we compute the combined terms. These terms are independent of the convergence parameters `gmax` and `nlatvec`, but do depend on the partitioning parameter :math:`\alpha`.
We define two constants, `squareconst`, the coefficient of the squared charges,
and `ijconst`, the coefficient of the pairs:
.. math:: C_{ij} = - \frac{\pi}{V\alpha^2}
.. math:: C_{\rm square} = - \frac{\alpha}{\sqrt{\pi}} - \frac{\pi}{2V\alpha^2}
= - \frac{\alpha}{\sqrt{\pi}} - \frac{C_{ij}}{2}
The Ewald object doesn't retain information about the configurations, including number of electrons, so the electron constants are defined as functions of :math:`N_e`.
Self plus charged-system energy:
.. math:: E_{\rm self+charged}^{ee} = N_e C_{\rm square} + \frac{N_e(N_e-1)}{2} C_{ij}
.. math:: E_{\rm self+charged}^{e\text{-ion}} = - N_e \sum_{I=1}^{N_{ion}} Z_I C_{ij}
.. math:: E_{\rm self+charged}^{\text{ion-ion}} = \sum_{I=1}^{N_{ion}} Z_I^2 C_{\rm square} + \sum_{I<J}^{N_{ion}} Z_I Z_J C_{ij}
We also compute contributions from a single electron, to separate the Ewald sum by electron.
.. math:: E_{\rm self+charged}^{\rm single} = C_{\rm square} + \frac{N_e-1}{2} C_{ij} - \sum_{I=1}^{N_{ion}} Z_I C_{ij}
.. math:: E_{\rm self+charged}^{\text{single-test}} = C_{\rm square} - \sum_{I=1}^{N_{ion}} Z_I C_{ij}
"""
self.i_sum = np.sum(self.atom_charges)
ii_sum2 = np.sum(self.atom_charges ** 2)
ii_sum = (self.i_sum ** 2 - ii_sum2) / 2
self.ijconst = -np.pi / (cellvolume * self.alpha ** 2)
self.squareconst = -self.alpha / np.sqrt(np.pi) + self.ijconst / 2
self.ii_const = ii_sum * self.ijconst + ii_sum2 * self.squareconst
self.e_single_test = -self.i_sum * self.ijconst + self.squareconst
self.ion_ion = self.ewald_ion()
# XC correction not used, so we can compare to other codes
# rs = lambda ne: (3 / (4 * np.pi) / (ne * cellvolume)) ** (1 / 3)
# cexc = 0.36
# xc_correction = lambda ne: cexc / rs(ne)
def ee_const(self, ne):
return ne * (ne - 1) / 2 * self.ijconst + ne * self.squareconst
def ei_const(self, ne):
return -ne * self.i_sum * self.ijconst
def e_single(self, ne):
return (
0.5 * (ne - 1) * self.ijconst - self.i_sum * self.ijconst + self.squareconst
)
[docs] def ewald_ion(self):
r"""
Compute ion contribution to Ewald sums. Since the ions don't move in our calculations, the ion-ion term only needs to be computed once.
Note: We ignore the constant term :math:`\frac{1}{2} \sum_{I} Z_I^2 C_{\rm self\ image}` in the real-space ion-ion sum corresponding to the interaction of an ion with its own image in other cells.
The real-space part:
.. math:: E_{\rm real\ space}^{\text{ion-ion}} = \sum_{\vec{n}} \sum_{I<J}^{N_{ion}} Z_I Z_J \frac{{\rm erfc}(\alpha |\vec{x}_{IJ}+\vec{n}|)}{|\vec{x}_{IJ}+\vec{n}|}
The reciprocal-space part:
.. math:: E_{\rm reciprocal\ space}^{\text{ion-ion}} = \sum_{\vec{G} > 0 } W_G \left| \sum_{I=1}^{N_{ion}} Z_I e^{-i\vec{G}\cdot\vec{x}_I} \right|^2
:returns: ion-ion component of Ewald sum
:rtype: float
"""
# Real space part
if len(self.atom_charges) == 1:
ion_ion_real = 0
else:
dist = pyqmc.distance.MinimalImageDistance(self.latvec)
ion_distances, ion_inds = dist.dist_matrix(self.atom_coords[np.newaxis])
ion_distances = gpu.cp.asarray(ion_distances)
rvec = ion_distances[:, :, np.newaxis, :] + self.lattice_displacements
r = gpu.cp.linalg.norm(rvec, axis=-1)
charge_ij = gpu.cp.prod(self.atom_charges[np.asarray(ion_inds)], axis=1)
ion_ion_real = gpu.cp.einsum(
"j,ijk->", charge_ij, gpu.erfc(self.alpha * r) / r
)
# Reciprocal space part
GdotR = gpu.cp.dot(self.gpoints, gpu.cp.asarray(self.atom_coords.T))
self.ion_exp = gpu.cp.dot(gpu.cp.exp(1j * GdotR), self.atom_charges)
ion_ion_rec = gpu.cp.dot(self.gweight, gpu.cp.abs(self.ion_exp) ** 2)
ion_ion = ion_ion_real + ion_ion_rec
return ion_ion
def _real_cij(self, dists):
dists = gpu.cp.asarray(dists)
r = gpu.cp.zeros(dists.shape[:-1])
cij = gpu.cp.zeros(r.shape)
for ld in self.lattice_displacements:
r[:] = np.linalg.norm(dists + ld, axis=-1)
cij += gpu.erfc(self.alpha * r) / r
return cij
[docs] def ewald_electron(self, configs):
r"""
Compute the Ewald sum for e-e and e-ion
Note: We ignore the constant term :math:`\frac{N_e}{2} C_{\rm self\ image}` in the real-space e-e sum corresponding to the interaction of an electron with its own image in other cells.
Real space e-e:
.. math:: E_{\rm real\ space}^{ee} = \sum_{\vec{n}} \sum_{i<j}^{N_e} \frac{{\rm erfc}(\alpha r_{ijn})}{r_{ijn}}
Real space e-i:
.. math:: E_{\rm real\ space}^{e\text{-ion}} = {\sum_{\vec{n}}} \sum_{i=1}^{N_e} \sum_{I=1}^{N_{ion}} -Z_I \frac{{\rm erfc}(\alpha r_{iIn})}{r_{iIn}}
Reciprocal space e-e:
.. math:: E_{\rm reciprocal\ space}^{ee} = \sum_{\vec{G}>0} W_G \left| \sum_{i=1}^{N_e} e^{-i\vec{k}\cdot\vec{x}_i} \right|^2
Reciprocal space e-i:
.. math:: E_{\rm reciprocal\ space}^{e\text{-ion}} = \sum_{\vec{G}>0} W_G {\rm Re} \left[ 2 \sum_{i=1}^{N_e} \sum_{I=1}^{N_{ion}} -Z_I e^{-i\vec{k}\cdot\vec{x}_i} e^{i\vec{k}\cdot\vec{x}_I} \right]
:parameter configs: electron positions (walkers)
:type configs: (nconf, nelec, 3) PeriodicConfigs object
:returns:
* ee: electron-electron part
* ei: electron-ion part
:rtype: float, float
"""
nconf, nelec, ndim = configs.configs.shape
# Real space electron-ion part
# ei_distances shape (elec, conf, atom, dim)
ei_distances = configs.dist.dist_i(self.atom_coords, configs.configs)
ei_cij = self._real_cij(ei_distances)
ei_real_separated = gpu.cp.einsum("k,ijk->ji", -self.atom_charges, ei_cij)
# Real space electron-electron part
ee_real_separated = gpu.cp.zeros((nconf, nelec))
if nelec > 1:
ee_distances, ee_inds = configs.dist.dist_matrix(configs.configs)
ee_cij = self._real_cij(ee_distances)
for ((i, j), val) in zip(ee_inds, ee_cij.T):
ee_real_separated[:, i] += val
ee_real_separated[:, j] += val
ee_real_separated /= 2
ee_recip, ei_recip = self.reciprocal_space_electron(configs.configs)
ee = ee_real_separated.sum(axis=1) + ee_recip
ei = ei_real_separated.sum(axis=1) + ei_recip
return ee, ei
def reciprocal_space_electron(self, configs):
# Reciprocal space electron-electron part
e_GdotR = gpu.cp.einsum("hik,jk->hij", gpu.cp.asarray(configs), self.gpoints)
sum_e_sin = gpu.cp.sin(e_GdotR).sum(axis=1)
sum_e_cos = gpu.cp.cos(e_GdotR).sum(axis=1)
ee_recip = gpu.cp.dot(sum_e_sin ** 2 + sum_e_cos ** 2, self.gweight)
## Reciprocal space electron-ion part
coscos_sinsin = -self.ion_exp.real * sum_e_cos - self.ion_exp.imag * sum_e_sin
ei_recip = 2 * gpu.cp.dot(coscos_sinsin, self.gweight)
return ee_recip, ei_recip
def reciprocal_space_electron_separated(self, configs):
# Reciprocal space electron-electron part
e_GdotR = np.einsum("hik,jk->hij", gpu.cp.asarray(configs), self.gpoints)
e_sin = np.sin(e_GdotR)
e_cos = np.cos(e_GdotR)
sinsin = e_sin.sum(axis=1, keepdims=True) * e_sin
coscos = e_cos.sum(axis=1, keepdims=True) * e_cos
ee_recip = np.dot(coscos + sinsin - 0.5, self.gweight)
## Reciprocal space electron-ion part
coscos_sinsin = -self.ion_exp.real * e_cos + self.ion_exp.imag * e_sin
ei_recip = np.dot(coscos_sinsin, self.gweight)
return ee_recip, ei_recip
def save_separated(self, ee_recip, ei_recip, ee_real, ei_real):
# Combine parts
self.ei_separated = ei_real + 2 * ei_recip
self.ee_separated = ee_real + 1 * ee_recip
self.ewalde_separated = self.ei_separated + self.ee_separated
nelec = ee_recip.shape[1]
### Add back the 0.5 that was subtracted earlier
ee = self.ee_separated.sum(axis=1) + nelec / 2 * self.gweight.sum()
ei = self.ei_separated.sum(axis=1)
return ee, ei
[docs] def energy(self, configs):
r"""
Compute Coulomb energy for a set of configs.
.. math:: E_{\rm Coulomb} &= E_{\rm real+reciprocal}^{ee}
+ E_{\rm self+charged}^{ee}
\\&+ E_{\rm real+reciprocal}^{e\text{-ion}}
+ E_{\rm self+charged}^{e\text{-ion}}
\\&+ E_{\rm real+reciprocal}^{\text{ion-ion}}
+ E_{\rm self+charged}^{\text{ion-ion}}
:parameter configs: electron positions (walkers)
:type configs: (nconf, nelec, 3) PeriodicConfigs object
:returns:
* ee: electron-electron part
* ei: electron-ion part
* ii: ion-ion part
:rtype: float, float, float
"""
nelec = configs.configs.shape[1]
ee, ei = self.ewald_electron(configs)
ee += self.ee_const(nelec)
ei += self.ei_const(nelec)
ii = self.ion_ion + self.ii_const
return gpu.asnumpy(ee), gpu.asnumpy(ei), gpu.asnumpy(ii)
[docs] def energy_separated(self, configs):
"""
Compute Coulomb energy contribution from each electron (does not include ion-ion energy).
NOTE: energy() needs to be called first to update the separated energy values
:parameter configs: electron positions (walkers)
:type configs: (nconf, nelec, 3) PeriodicConfigs object
:returns: energies
:rtype: (nelec,) array
"""
raise NotImplementedError("ewalde_separated is currently not computed anywhere")
nelec = configs.configs.shape[1]
return self.e_single(nelec) + self.ewalde_separated
def select_big(gpts, cellvolume, recvec, alpha):
gpoints = gpu.cp.einsum("j...,jk->...k", gpts, recvec) * 2 * np.pi
gsquared = gpu.cp.einsum("...k,...k->...", gpoints, gpoints)
gweight = 4 * np.pi * gpu.cp.exp(-gsquared / (4 * alpha ** 2))
gweight /= cellvolume * gsquared
bigweight = gweight > 1e-10
return gpoints[bigweight], gweight[bigweight]