# Authors: Artem Pulkin, pyscf authors
"""
This module implements the G0W0 approximation on top of `pyscf.tdscf.rhf_slow` and `pyscf.tdscf.proxy` TD
implementations. Unlike `gw.py`, all integrals are stored in memory. Several variants of GW are available:
* (this module) `pyscf.gw_slow`: the molecular implementation;
* `pyscf.pbc.gw.gw_slow`: single-kpoint PBC (periodic boundary condition) implementation;
* `pyscf.pbc.gw.kgw_slow_supercell`: a supercell approach to PBC implementation with multiple k-points. Runs the
molecular code for a model with several k-points for the cost of discarding momentum conservation and using dense
instead of sparse matrixes;
* `pyscf.pbc.gw.kgw_slow`: a PBC implementation with multiple k-points;
"""
from pyscf.lib import einsum, direct_sum
from pyscf.lib import logger
import numpy
from scipy.optimize import newton, bisect
from itertools import product
# Convention for these modules:
# * IMDS contains routines for intermediates
# * kernel finds GW roots
# * GW provides a container
[docs]
class AbstractIMDS:
orb_dims = 1
def __init__(self, td, eri=None):
"""
GW intermediates interface.
Args:
td: a container with TD solution;
eri (PhysERI): a container with electron repulsion integrals;
"""
self.td = td
if eri is None:
self.eri = td.eri
else:
self.eri = eri
[docs]
def get_rhs(self, p):
"""
The right-hand side of the quasiparticle equation.
Args:
p (int, tuple): the orbital;
Returns:
Right-hand sides of the quasiparticle equation
"""
raise NotImplementedError
[docs]
def entire_space(self):
"""
The entire orbital space.
Returns:
An iterable of the entire orbital space.
"""
raise NotImplementedError
[docs]
def get_sigma_element(self, omega, p, **kwargs):
"""
The diagonal matrix element of the self-energy matrix.
Args:
omega (float): the energy value;
p (int, tuple): the orbital;
Returns:
The diagonal matrix element.
"""
raise NotImplementedError
[docs]
def quasiparticle_eq(self, p, **kwargs):
"""
The quasiparticle equation `f(omega) = 0`.
Args:
p (int, tuple): the orbital;
**kwargs: keyword arguments to `get_sigma_element`;
Returns:
A callable function of one parameter.
"""
rhs = self.get_rhs(p)
def quasiparticle_eq(omega):
return omega - self.get_sigma_element(omega, p, **kwargs).real - rhs
return quasiparticle_eq
[docs]
def initial_guess(self, p):
"""
Retrieves the initial guess for the quasiparticle energy for orbital `p`.
Args:
p (int, tuple): the orbital;
Returns:
The value of initial guess (float).
"""
raise NotImplementedError
[docs]
def corrected_moe(eri, p):
"""
Calculates the corrected orbital energy.
Args:
eri (PhysERI): a container with electron repulsion integrals;
p (int): orbital;
Returns:
The corrected orbital energy.
"""
moe = eri.mo_energy[p]
moc = eri.mo_coeff[:, p]
vk = - eri.ao2mo((
moc[:, numpy.newaxis],
eri.mo_coeff_full[:, :eri.nocc_full],
eri.mo_coeff_full[:, :eri.nocc_full],
moc[:, numpy.newaxis],
)).squeeze().trace()
mf = eri.model
v_mf = eri.squeeze(mf.get_veff() - mf.get_j())
v_mf = einsum("i,ij,j", moc.conj(), v_mf, moc)
return moe + vk - v_mf
[docs]
class IMDS(AbstractIMDS):
def __init__(self, td, eri=None):
"""
GW intermediates.
Args:
td: a container with TD solution;
eri: a container with electron repulsion integrals;
"""
super(IMDS, self).__init__(td, eri=eri)
# MF
self.nocc = self.eri.nocc
self.o, self.v = self.eri.mo_energy[:self.nocc], self.eri.mo_energy[self.nocc:]
# TD
self.td_xy = self.td.xy
self.td_e = self.td.e
self.tdm = self.construct_tdm()
def __getitem__(self, item):
return self.eri[item]
[docs]
def get_rhs(self, p):
# return self.eri.mo_energy[p]
return corrected_moe(self.eri, p)
[docs]
def construct_tdm(self):
td_xy = 2 * numpy.asarray(self.td_xy)
tdm_oo = einsum('vxia,ipaq->vxpq', td_xy, self["oovo"])
tdm_ov = einsum('vxia,ipaq->vxpq', td_xy, self["oovv"])
tdm_vv = einsum('vxia,ipaq->vxpq', td_xy, self["ovvv"])
if numpy.iscomplexobj(self["oovv"]):
tdm_vo = einsum('vxia,ipaq->vxpq', td_xy, self["ovvo"])
else:
tdm_vo = tdm_ov.swapaxes(2, 3).conj()
tdm = numpy.concatenate(
(
numpy.concatenate((tdm_oo, tdm_ov), axis=3),
numpy.concatenate((tdm_vo, tdm_vv), axis=3)
),
axis=2,
)
return tdm
[docs]
def get_sigma_element(self, omega, p, eta, vir_sgn=1):
tdm = self.tdm.sum(axis=1)
evi = direct_sum('v-i->vi', self.td_e, self.o)
eva = direct_sum('v+a->va', self.td_e, self.v)
sigma = numpy.sum(tdm[:, :self.nocc, p] ** 2 / (omega + evi - 1j * eta))
sigma += numpy.sum(tdm[:, self.nocc:, p] ** 2 / (omega - eva + vir_sgn * 1j * eta))
return sigma
[docs]
def initial_guess(self, p):
return self.eri.mo_energy[p]
@property
def entire_space(self):
return [numpy.arange(self.eri.nmo)]
[docs]
class LoggingFunction:
def __init__(self, m):
"""
A function number->number logging calls.
Args:
m (callable): an underlying method of a single number returning a number;
"""
self.m = m
self.__x__ = []
self.__y__ = []
@property
def x(self):
return numpy.asarray(self.__x__)
@property
def y(self):
return numpy.asarray(self.__y__)
def __call__(self, x):
y = self.m(x)
self.__x__.append(x)
self.__y__.append(y)
return y
[docs]
def plot_call_history(self, title=""):
"""
Plots calls to this function.
Args:
title (str): plot title;
"""
if len(self.x) > 1:
from matplotlib import pyplot
x = self.x.real
y = self.y.real
pyplot.scatter(x[1:], y[1:], marker='+', color="black", s=10)
pyplot.scatter(x[:1], y[:1], marker='+', color="red", s=50)
pyplot.axhline(y=0, color="grey")
pyplot.title(title + " ncalls: {:d}".format(len(self.x)))
pyplot.show()
[docs]
def kernel(imds, orbs=None, linearized=False, eta=1e-3, tol=1e-9, method="fallback"):
"""
Calculates GW energies.
Args:
imds (AbstractIMDS): GW intermediates;
orbs (Iterable): indexes of MO orbitals to correct;
linearized (bool): whether to apply a single-step linearized correction to energies instead of iterative
procedure;
eta (float): imaginary energy for the Green's function;
tol (float): tolerance for the search of zero;
method (str): 'bisect' finds roots no matter what but, potentially, wrong ones, 'newton' finding roots close to
the correct one but, potentially, failing during iterations, or 'fallback' using 'newton' and proceeding to
'bisect' in case of failure;
Returns:
Corrected orbital energies.
"""
if method not in ('newton', 'bisect', 'fallback'):
raise ValueError("Cannot recognize method='{}'".format(method))
# Check implementation consistency
_orbs = imds.entire_space
if not isinstance(_orbs, list) or not len(_orbs) == imds.orb_dims:
raise RuntimeError("The object returned by 'imds.entire_space' is not a list of length {:d}: {}".format(
imds.orb_dims,
repr(_orbs),
))
# Assign default value
if orbs is None:
orbs = _orbs
# Make sure it is a list
if not isinstance(orbs, list):
orbs = [orbs]
# Add missing dimensions
if len(orbs) < imds.orb_dims:
orbs = _orbs[:-len(orbs)] + orbs
shape = tuple(len(i) for i in orbs)
gw_energies = numpy.zeros(shape, dtype=float)
for i_p in product(*tuple(numpy.arange(i) for i in shape)):
p = tuple(i[j] for i, j in zip(orbs, i_p))
if imds.orb_dims == 1:
p = p[0]
if linearized:
raise NotImplementedError
# v_mf = imds.vmf
# vk = imds.vk
# de = 1e-6
# ep = imds.e_mf[p]
# # TODO: analytic sigma derivative
# sigma = imds.get_sigma_element(ep, p, eta).real
# dsigma = imds.get_sigma_element(ep + de, p, eta).real - sigma
# zn = 1.0 / (1 - dsigma / de)
# e = ep + zn * (sigma.real + vk[p] - v_mf[p])
# gw_energies[i_p] = e
else:
debug = LoggingFunction(imds.quasiparticle_eq(p, eta=eta))
if method == "newton":
try:
gw_energies[i_p] = newton(debug, imds.initial_guess(p), tol=tol, maxiter=100)
except Exception as e:
e.message = "When calculating root @p={} the following exception occurred:\n\n{}".format(
repr(p),
e.message,
)
debug.plot_call_history("Exception during Newton " + str(p))
raise
elif method == "bisect":
gw_energies[i_p] = bisect(debug, -100, 100, xtol=tol, maxiter=100)
elif method == "fallback":
try:
gw_energies[i_p] = newton(debug, imds.initial_guess(p), tol=tol, maxiter=100)
except RuntimeError:
logger.warn(imds.td._scf,
"Failed to converge with newton, using bisect on the interval [{:.3e}, {:.3e}]".format(
min(debug.x), max(debug.x),))
gw_energies[i_p] = bisect(debug, min(debug.x), max(debug.x), xtol=tol, maxiter=100)
return gw_energies
[docs]
class GW:
base_imds = IMDS
def __init__(self, td, eri=None):
"""
Performs GW calculation. Roots are stored in `self.mo_energy`.
Args:
td: a container with TD solution;
eri: a container with electron repulsion integrals;
"""
self.td = td
self.imds = self.base_imds(td, eri=eri)
self.mo_energy = None
self.orbs = None
self.method = "fallback"
self.eta = 1e-3
[docs]
def kernel(self):
"""
Calculates GW roots.
Returns:
GW roots.
"""
self.mo_energy = kernel(self.imds, orbs=self.orbs, method=self.method, eta=self.eta)
return self.mo_energy