#!/usr/bin/env python
# Copyright 2014-2021 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Author: Tianyu Zhu <zhutianyu1991@gmail.com>
#
"""
Periodic k-point spin-restricted random phase approximation
(direct RPA/dRPA in chemistry) with N^4 scaling
Method:
Main routines are based on GW-AC method described in:
T. Zhu and G.K.-L. Chan, J. Chem. Theory. Comput. 17, 727-741 (2021)
X. Ren et al., New J. Phys. 14, 053020 (2012)
"""
import numpy as np
import scipy.linalg.blas as blas
import time
from pyscf import lib
from pyscf.lib import logger, temporary_env
from pyscf.ao2mo._ao2mo import r_e2
from pyscf.ao2mo.incore import _conc_mos
from pyscf.pbc import df, dft, scf
from pyscf.pbc.mp.kmp2 import get_nocc, get_nmo, get_frozen_mask
from fcdmft.gw.minimax.minimax_rpa import get_frequency_minimax_grid
from fcdmft.ac.grids import _get_scaled_legendre_roots
from fcdmft.gw.pbc.krgw_ac import _mo_occ_frozen, _mo_energy_frozen, _mo_frozen, get_rho_response, \
get_rho_response_head, get_rho_response_wing, get_qij
from fcdmft.rpa.pbc.rpa import get_idx_metal
from fcdmft.utils.kpts import get_kconserv_ria_efficient
from mpi4py import MPI
rank = MPI.COMM_WORLD.Get_rank()
size = MPI.COMM_WORLD.Get_size()
comm = MPI.COMM_WORLD
einsum = lib.einsum
[docs]
def kernel(rpa, mo_energy, mo_coeff, nw=None, with_e_hf=None):
"""RPA correlation and total energy
Parameters
----------
rpa : KRPA
rpa object
mo_energy : double array
molecular orbital energies
mo_coeff : double ndarray
molecular orbital coefficients
Lpq : double array, optional
density fitting 3-center integral in MO basis, by default None
nw : int, optional
number of frequency point on imaginary axis, by default None
with_e_hf : float, optional
extra input HF energy, by default None
Returns
-------
e_tot : float
RPA total energy
e_hf : float
HF energy (exact exchange for given mo_coeff)
e_corr : float
RPA correlation energy
"""
# Compute HF exchange energy (EXX)
dm = rpa._scf.make_rdm1()
if with_e_hf is None:
rhf = scf.KRHF(rpa.mol, rpa.kpts, exxdiv=rpa._scf.exxdiv)
rhf.verbose = 0
if hasattr(rpa._scf, 'sigma'):
rhf = scf.addons.smearing_(rhf, sigma=rpa._scf.sigma, method=rpa._scf.smearing_method)
rhf.with_df = rpa._scf.with_df
with temporary_env(rpa.with_df, verbose=0), temporary_env(rhf.mol, verbose=0):
e_hf = rhf.energy_elec(dm)[0]
e_hf += rpa._scf.energy_nuc()
else:
e_hf = with_e_hf
if rank == 0 and rpa.verbose >= logger.DEBUG:
logger.debug(rpa, f' Setting EXX energy explicitly to {e_hf}')
is_metal = hasattr(rpa._scf, 'sigma')
# Turn off FC for metals and outcore
if is_metal and rpa.fc:
if rank == 0:
logger.warn(rpa, 'FC not available for metals - setting rpa.fc to False')
rpa.fc = False
if rpa.outcore and rpa.fc:
if rank == 0:
logger.warn(rpa, 'FC not available for outcore - setting rpa.fc to False')
rpa.fc = False
# Grids for integration on imaginary axis
freqs, wts = rpa.get_grids(nw=nw, mo_energy=mo_energy)
# Compute RPA correlation energy
if rpa.outcore:
if is_metal:
e_corr = get_rpa_ecorr_outcore_metal(rpa, freqs, wts)
else:
e_corr = get_rpa_ecorr_outcore(rpa, freqs, wts)
else:
e_corr = get_rpa_ecorr(rpa, freqs, wts)
# Compute total energy
e_tot = e_hf + e_corr
if rank == 0 and rpa.verbose >= logger.DEBUG:
logger.debug(rpa, f' RPA total energy = {e_tot}')
logger.debug(rpa, f' EXX energy = {e_hf}, RPA corr energy = {e_corr}')
return e_tot, e_hf, e_corr
[docs]
def rho_accum_inner(Pi, eia, omega, Lov, alpha=0.0, fia=None):
"""Get contribution to response function from current occupied-virtual block.
Parameters
----------
Pi : complex 2d array
density-density response function, will be overwritten
eia : double 2d array
occupied-virtual orbital energy difference
omega : double
real position of imaginary frequency
Lov : complex 3d array
occupied-virtual block of three-center density-fitting matrix in MO
alpha : float, optional
prefactor, by default 0.0
fia : double 2d array, optional
occupied-virtual occupation number difference, by default None
"""
naux, nocc, nvir = Lov.shape
try:
import numexpr as ne
if fia is None:
eia = ne.evaluate('eia / (omega**2 + eia**2)')
else:
eia = ne.evaluate('eia * fia / (omega**2 + eia**2)')
Pia = ne.evaluate('Lov * eia').reshape(naux, nocc * nvir)
except ImportError:
# plain numpy, numexpr not available
if fia is None:
eia = eia / (omega**2 + eia**2)
else:
eia = eia * fia / (omega**2 + eia**2)
Pia = (Lov * eia).reshape(naux, nocc * nvir)
# The following call to blas.zgemm may be replaced with
# Pi += alpha * np.einsum('Pia, Qia -> PQ', Pia, Lov.conj(), optimize=True)
# with a moderate performance hit.
# zgemm is complex matrix multiplication. A wrapper is included in SciPy.
# C <- alpha * op(A) @ op(B) + beta * C
blas.zgemm(
alpha=alpha,
a=Lov.reshape(naux, nocc * nvir).T,
b=Pia.T,
trans_a=2, # take conjugate transpose of A (this gives Lov.conj())
trans_b=0, # B is Pia.T
beta=1.0,
c=Pi.T, # Pi.T += alpha * Lov.conj() @ Pia.T
overwrite_c=True,
)
return
[docs]
def get_rpa_ecorr(rpa, freqs, wts):
"""Compute RPA correlation energy.
Parameters
----------
rpa : KRPA
rpa object
freqs : double 1d array
frequency grid
wts : double 1d array
weight of grids
Returns
-------
e_corr : double
correlation energy
"""
mo_coeff = np.array(_mo_frozen(rpa, rpa._scf.mo_coeff))
mo_energy = np.array(_mo_energy_frozen(rpa, rpa._scf.mo_energy))
mo_occ = np.array(_mo_occ_frozen(rpa, rpa._scf.mo_occ))
nocc = rpa.nocc
nmo = rpa.nmo
nvir = nmo - nocc
nao = rpa._scf.mo_coeff[0].shape[0]
nkpts = rpa.nkpts
kpts = rpa.kpts
nw = len(freqs)
mydf = rpa.with_df
# possible kpts shift center
kscaled = rpa.mol.get_scaled_kpts(kpts)
kscaled -= kscaled[0]
is_metal = hasattr(rpa._scf, 'sigma')
segsize = nkpts // size
if rank >= size - (nkpts - segsize * size):
start = rank * segsize + rank - (size - (nkpts - segsize * size))
stop = min(nkpts, start + segsize + 1)
else:
start = rank * segsize
stop = min(nkpts, start + segsize)
if rpa.fc:
qij, q_abs, nq_pts = rpa.get_q_mesh(mo_energy, mo_coeff)
e_corr = 0j
# Precompute k-conservation table
# Given k-point indices (kL, i), kconserv_table[kshift,i] contains
# the index j that satisfies momentum conservation,
# (k(i) - k(j) - k(kL)) \dot a = 2n\pi
# i.e.
# - ki + kj + kL = G
kconserv_table = get_kconserv_ria_efficient(rpa.mol, kpts)
cderiarr = mydf.cderi_array()
for kL in range(start, stop):
# Lij: (ki, L, i, j) for looping every kL
if is_metal:
Lij = []
else:
Lij = None
# kidx: save kj that conserves with kL and ki (-ki+kj+kL=G)
# kidx_r: save ki that conserves with kL and kj (-ki+kj+kL=G)
# kidx = np.zeros((nkpts),dtype=np.int64)
# kidx_r = np.zeros((nkpts),dtype=np.int64)
for i, kpti in enumerate(kpts):
j = kconserv_table[kL, i]
kptj = kpts[j]
kconserv = -kscaled[i] + kscaled[j] + kscaled[kL]
assert np.linalg.norm(np.round(kconserv) - kconserv) < 1e-12 # kidx[i] = j
# kidx_r[j] = i
logger.debug(rpa, f'Read Lpq (kL: {kL+1} / {nkpts}, ki: {i}, kj: {j} @ Rank {rank})')
# Read (L|pq) and ao2mo transform to (L|ij)
# support uneqaul naux on different k points
Lpq = cderiarr.load(kpti, kptj)
if Lpq.shape[-1] == (nao * (nao + 1)) // 2:
Lpq = lib.unpack_tril(Lpq).reshape(-1, nao**2)
else:
Lpq = Lpq.reshape(-1, nao**2)
Lpq = Lpq.astype(np.complex128)
tao = []
ao_loc = None
moij, ijslice = _conc_mos(mo_coeff[i], mo_coeff[j])[2:]
naux = Lpq.shape[0]
if not is_metal:
if Lij is None:
Lij = np.zeros((nkpts, naux, nocc, nvir), dtype=np.complex128)
ijslice = (0, nocc, nmo + nocc, 2 * nmo)
r_e2(Lpq, moij, ijslice, tao, ao_loc, out=Lij[i])
else:
# Only (nocc+nfrac, nfrac+nvir) block of Lpq is needed
# This is consistent with the new get_rho_response_metal implementation
idx_occ_i, idx_frac_i, idx_vir_i = get_idx_metal(mo_occ[i])
idx_occ_j, idx_frac_j, idx_vir_j = get_idx_metal(mo_occ[j])
nocc_i = len(idx_occ_i)
nfrac_i = len(idx_frac_i)
nocc_j = len(idx_occ_j)
nfrac_j = len(idx_frac_j)
nvir_j = len(idx_vir_j)
ijslice = (0, nocc_i + nfrac_i, nmo + nocc_j, 2 * nmo)
Lij.append(r_e2(Lpq, moij, ijslice, tao, ao_loc).reshape(naux, nocc_i + nfrac_i, nfrac_j + nvir_j))
for w in range(nw):
if is_metal:
Pi = get_rho_response_metal(freqs[w], mo_energy, mo_occ, Lij, kconserv_table[kL])
else:
Pi = get_rho_response(freqs[w], mo_energy, Lij, kconserv_table[kL])
if kL == 0 and rpa.fc:
for iq in range(nq_pts):
# head Pi_00
Pi_00 = get_rho_response_head(freqs[w], mo_energy, qij[iq])
Pi_00 = 4.0 * np.pi / np.linalg.norm(q_abs[iq]) ** 2 * Pi_00
# wings Pi_P0
Pi_P0 = get_rho_response_wing(freqs[w], mo_energy, Lij, qij[iq])
Pi_P0 = np.sqrt(4.0 * np.pi) / np.linalg.norm(q_abs[iq]) * Pi_P0
# assemble Pi
Pi_fc = np.zeros((naux + 1, naux + 1), dtype=Pi.dtype)
Pi_fc[0, 0] = Pi_00
Pi_fc[0, 1:] = Pi_P0.conj()
Pi_fc[1:, 0] = Pi_P0
Pi_fc[1:, 1:] = Pi
e_corr += get_rpa_ecorr_w(Pi_fc, wts[w])
else:
e_corr += get_rpa_ecorr_w(Pi, wts[w])
e_corr = e_corr.real
e_corr = comm.allreduce(e_corr, op=MPI.SUM)
e_corr *= 1.0 / (2.0 * np.pi) / nkpts
return e_corr
[docs]
def get_rpa_ecorr_outcore(rpa, freqs, wts):
"""Low-memory routine to compute RPA correlation energy.
Parameters
----------
rpa : KRPA
rpa object
freqs : double 1d array
frequency grid
wts : double 1d array
weight of grids
Returns
-------
e_corr : double
correlation energy
"""
mo_coeff = np.array(_mo_frozen(rpa, rpa._scf.mo_coeff))
mo_energy = np.array(_mo_energy_frozen(rpa, rpa._scf.mo_energy))
nocc = rpa.nocc
nmo = rpa.nmo
nao = rpa._scf.mo_coeff[0].shape[0]
nkpts = rpa.nkpts
kpts = rpa.kpts
nw = len(freqs)
mydf = rpa.with_df
# possible kpts shift center
kscaled = rpa.mol.get_scaled_kpts(kpts)
kscaled -= kscaled[0]
segsize = nkpts // size
if rank >= size - (nkpts - segsize * size):
start = rank * segsize + rank - (size - (nkpts - segsize * size))
stop = min(nkpts, start + segsize + 1)
else:
start = rank * segsize
stop = min(nkpts, start + segsize)
if rpa.fc:
qij, q_abs, nq_pts = rpa.get_q_mesh(mo_energy, mo_coeff)
e_corr = 0j
# Precompute k-conservation table
# Given k-point indices (kL, i), kconserv_table[kshift,i] contains
# the index j that satisfies momentum conservation,
# (k(i) - k(j) - k(kL)) \dot a = 2n\pi
# i.e.
# - ki + kj + kL = G
kconserv_table = get_kconserv_ria_efficient(rpa.mol, kpts)
cderiarr = mydf.cderi_array()
for kL in range(start, stop):
Pi = None
nseg = nocc // rpa.segsize + 1
for iseg in range(nseg):
orb_start = iseg * rpa.segsize
orb_end = min((iseg + 1) * rpa.segsize, nocc)
if orb_end == orb_start:
continue
norb_this_iter = orb_end - orb_start
# Lij: (ki, L, i, j) for looping every kL
# kidx: save kj that conserves with kL and ki (-ki+kj+kL=G)
# kidx_r: save ki that conserves with kL and kj (-ki+kj+kL=G)
# kidx = np.zeros((nkpts),dtype=np.int64)
# kidx_r = np.zeros((nkpts),dtype=np.int64)
for i, kpti in enumerate(kpts):
j = kconserv_table[kL, i]
kptj = kpts[j]
# Find (ki,kj) that satisfies momentum conservation with kL
kconserv = -kscaled[i] + kscaled[j] + kscaled[kL]
assert np.linalg.norm(np.round(kconserv) - kconserv) < 1e-12
logger.debug(rpa, f'Read Lpq (kL: {kL+1} / {nkpts}, ki: {i}, kj: {j} @ Rank {rank})')
# Read (L|pq) and ao2mo transform to (L|ij)
# support uneqaul naux on different k points
Lpq = cderiarr.load(kpti, kptj)
if Lpq.shape[-1] == (nao * (nao + 1)) // 2:
Lpq = lib.unpack_tril(Lpq).reshape(-1, nao**2)
else:
Lpq = Lpq.reshape(-1, nao**2)
Lpq = Lpq.astype(np.complex128)
naux = Lpq.shape[0]
tao = []
ao_loc = None
moij, ijslice = _conc_mos(mo_coeff[i], mo_coeff[j])[2:]
ijslice = (orb_start, orb_end, nmo + nocc, 2 * nmo)
Lij_slice = r_e2(Lpq, moij, ijslice, tao, ao_loc)
Lij_slice = Lij_slice.reshape(naux, norb_this_iter, nmo - nocc)
if Pi is None:
Pi = np.zeros((nw, naux, naux), dtype=np.complex128)
# Find ka that conserves with ki and kL (-ki+ka+kL=G)
a_inner = kconserv_table[kL, i]
eia = mo_energy[i][orb_start:orb_end, None] - mo_energy[a_inner][None, nocc:]
for w in range(nw):
rho_accum_inner(Pi[w], eia, freqs[w], Lij_slice, alpha=4.0 / nkpts)
for w in range(nw):
e_corr += get_rpa_ecorr_w(Pi[w], wts[w])
e_corr = e_corr.real
e_corr = comm.allreduce(e_corr, op=MPI.SUM)
e_corr *= 1.0 / (2.0 * np.pi) / nkpts
return e_corr
[docs]
def get_rpa_ecorr_w(Pi_w, wts_w):
"""Get contribution to RPA correlation energy from a single frequency.
Parameters
----------
Pi_w : complex 2d array
density-density response function at a single frequency
wts_w : double
weights of the frequency
Returns
-------
e_corr : double
correlation energy
"""
# First, compute ec_w = Tr(Pi_w) + |log(det(I-Pi_w))|
ec_w = np.trace(Pi_w)
# The following two lines are equivalent to
# Pi_w = np.eye(naux) - Pi_w
blas.zdscal(-1.0, Pi_w.ravel(), overwrite_x=1)
np.fill_diagonal(Pi_w, np.diagonal(Pi_w) + 1.0)
ec_w += np.linalg.slogdet(Pi_w)[1]
#e_corr = 1.0 / (2.0 * np.pi) / nkpts * ec_w * wts_w
e_corr = ec_w * wts_w
return e_corr
[docs]
class KRPA(lib.StreamObject):
def __init__(self, mf, frozen=None):
self.mol = mf.mol # mol object
self._scf = mf # mean-field object
self.verbose = self.mol.verbose # verbose level
self.stdout = self.mol.stdout # standard output
self.max_memory = mf.max_memory # max memory in MB
# options
self.frozen = frozen # frozen orbital options
self.grids_alg = 'legendre' # algorithm to generate grids
self.outcore = False # low-memory routine
self.segsize = 50 # number of orbitals in one segment for outcore
self.fc = True # finite-size correction
self.fc_grid = False # grids for finite-size correction
# don't modify the following attributes, they are not input options
self._nocc = None # number of occupied orbitals
self._nmo = None # number of orbitals (exclude frozen orbitals)
self.kpts = mf.kpts # k-points
self.nkpts = len(self.kpts) # number of k-points
self.mo_energy = np.array(mf.mo_energy, copy=True) # orbital energy
self.mo_coeff = np.array(mf.mo_coeff, copy=True) # orbital coefficient
self.mo_occ = np.array(mf.mo_occ, copy=True) # occupation number
self.e_corr = None # correlation energy
self.e_hf = None # Hartree-Fock energy
self.e_tot = None # total energy
# KRPA must use GDF integrals
if getattr(mf, 'with_df', None):
self.with_df = mf.with_df
else:
raise NotImplementedError
self._keys.update(['with_df'])
return
[docs]
def dump_flags(self, verbose=None):
log = logger.Logger(self.stdout, self.verbose)
log.info('')
log.info('******** %s ********', self.__class__)
log.info('method = %s', self.__class__.__name__)
nocc = self.nocc
nvir = self.nmo - nocc
nkpts = self.nkpts
log.info(f'RPA nocc = {nocc}, nvir = {nvir}, nkpts = {nkpts}')
if self.frozen is not None:
log.info(f'frozen orbitals = {str(self.frozen)}')
log.info('grid type = %s', self.grids_alg)
log.info('outcore mode = %s', self.outcore)
if self.outcore is True:
log.info('outcore segment size = %d', self.segsize)
log.info('RPA finite size corrections = %s', self.fc)
log.info('')
return
@property
def nocc(self):
frozen_mask = get_frozen_mask(self)
nkpts = len(self._scf.mo_energy)
nelec = 0.0
for k in range(nkpts):
nelec += np.sum(self._scf.mo_occ[k][frozen_mask[k]])
nelec = int(nelec / nkpts)
return nelec // 2
@nocc.setter
def nocc(self, n):
self._nocc = n
@property
def nmo(self):
frozen_mask = get_frozen_mask(self)
return len(self._scf.mo_energy[0][frozen_mask[0]])
@nmo.setter
def nmo(self, n):
self._nmo = n
get_nocc = get_nocc
get_nmo = get_nmo
get_frozen_mask = get_frozen_mask
[docs]
def kernel(self, mo_energy=None, mo_coeff=None, nw=None, with_e_hf=None):
"""RPA correlation and total energy
Calculated total energy, HF energy and RPA correlation energy
are stored in self.e_tot, self.e_hf, self.e_corr
Parameters
----------
mo_energy : double array
molecular orbital energies
mo_coeff : double ndarray
molecular orbital coefficients
nw : int, optional
number of frequency point on imaginary axis, by default None
with_e_hf : float, optional
If given, overrides the HF energy computation.
Returns
-------
e_tot : float
RPA total energy
e_hf : float
HF energy (exact exchange for given mo_coeff)
e_corr : float
RPA correlation energy
"""
if mo_coeff is None:
mo_coeff = _mo_frozen(self, self._scf.mo_coeff)
if mo_energy is None:
mo_energy = _mo_energy_frozen(self, self._scf.mo_energy)
cput0 = (time.process_time(), time.perf_counter())
if rank == 0:
self.dump_flags()
self.e_tot, self.e_hf, self.e_corr = kernel(self, mo_energy, mo_coeff, nw=nw, with_e_hf=with_e_hf)
if rank == 0:
logger.timer(self, 'RPA', *cput0)
return self.e_tot, self.e_hf, self.e_corr
[docs]
def get_grids(self, alg=None, nw=None, mo_energy=None):
"""Generate grids for integration.
Parameters
----------
alg : str, optional
algorithm for generating grids, by default None
nw : int, optional
number of grids, by default None
mo_energy : double 2d array, optional
orbital energy, used for minimax grids, by default None
Returns
-------
freqs : double 1d array
frequency grid
wts : double 1d array
weight of grids
"""
if alg is None:
alg = self.grids_alg
if mo_energy is None:
mo_energy = _mo_energy_frozen(self, self._scf.mo_energy)
if alg == 'legendre':
nw = 40 if nw is None else nw
freqs, wts = _get_scaled_legendre_roots(nw)
elif alg == 'minimax':
nw = 30 if nw is None else nw
E_max = 0
E_min = 1e10
for k in range(self.nkpts):
E_max = max(E_max, mo_energy[k, -1] - mo_energy[k, 0])
E_min = min(E_min, mo_energy[k, self.nocc] - mo_energy[k, self.nocc - 1])
E_range = E_max / E_min
freqs, wts = get_frequency_minimax_grid(nw, E_range)
else:
raise NotImplementedError('Grids algorithm not implemented!')
return freqs, wts
[docs]
def get_q_mesh(self, mo_energy, mo_coeff):
"""Get q-mesh for finite size correction.
Equation 39-42 in doi.org/10.1021/acs.jctc.0c00704
Parameters
----------
mo_energy : double 2d array
orbital energy
mo_coeff : double 3d array
coefficient from AO to MO
Returns
-------
qij : double 1d array
q-mesh grids
q_abs : double 1d array
absolute positions of q-mesh grids
nq_pts : init
number of q-mesh grids
"""
nocc = self.nocc
nmo = self.nmo
nkpts = self.nkpts
# Set up q mesh for q->0 finite size correction
if not rpa.fc_grid:
q_pts = np.array([1e-3, 0, 0], dtype=np.double).reshape(1, 3)
else:
Nq = 3
q_pts = np.zeros(shape=[Nq**3 - 1, 3], dtype=np.double)
for i in range(Nq):
for j in range(Nq):
for k in range(Nq):
if i == 0 and j == 0 and k == 0:
continue
else:
q_pts[i * Nq**2 + j * Nq + k - 1, 0] = k * 5e-4
q_pts[i * Nq**2 + j * Nq + k - 1, 1] = j * 5e-4
q_pts[i * Nq**2 + j * Nq + k - 1, 2] = i * 5e-4
nq_pts = len(q_pts)
q_abs = rpa.mol.get_abs_kpts(q_pts)
# Get qij = 1/sqrt(Omega) * < psi_{ik} | e^{iqr} | psi_{ak-q} > at q: (nkpts, nocc, nvir)
qij = np.zeros(shape=[nq_pts, nkpts, nocc, nmo - nocc], dtype=np.complex128)
if not rpa.fc_grid:
for k in range(nq_pts):
qij[k] = get_qij(rpa, q_abs[k], mo_energy, mo_coeff)
else:
segsize = nq_pts // size
if rank >= size - (nq_pts - segsize * size):
start = rank * segsize + rank - (size - (nq_pts - segsize * size))
stop = min(nq_pts, start + segsize + 1)
else:
start = rank * segsize
stop = min(nq_pts, start + segsize)
for k in range(start, stop):
qij[k] = get_qij(rpa, q_abs[k], mo_energy, mo_coeff)
comm.Barrier()
qij_gather = comm.gather(qij)
if rank == 0:
for i in range(1, size):
qij += qij_gather[i]
comm.Barrier()
qij = comm.bcast(qij, root=0)
return qij, q_abs, nq_pts
if __name__ == '__main__':
import os
from pyscf.pbc import gto, dft, scf
from pyscf.pbc.lib import chkfile
# This test takes a few minutes
# Test on diamond
cell = gto.Cell()
cell.build(
unit='angstrom',
a="""
0.000000 1.783500 1.783500
1.783500 0.000000 1.783500
1.783500 1.783500 0.000000
""",
atom='C 1.337625 1.337625 1.337625; C 2.229375 2.229375 2.229375',
dimension=3,
max_memory=12000,
verbose=5,
pseudo='gth-pbe',
basis='gth-dzv',
precision=1e-12,
)
kpts = cell.make_kpts([3, 1, 1], scaled_center=[0, 0, 0])
gdf = df.RSGDF(cell, kpts)
gdf_fname = 'gdf_ints_311.h5'
gdf._cderi_to_save = gdf_fname
if not os.path.isfile(gdf_fname):
gdf.build()
chkfname = 'diamond_311.chk'
if os.path.isfile(chkfname):
kmf = scf.KRHF(cell, kpts).rs_density_fit()
kmf.with_df = gdf
kmf.with_df._cderi = gdf_fname
data = chkfile.load(chkfname, 'scf')
kmf.__dict__.update(data)
else:
kmf = scf.KRHF(cell, kpts).rs_density_fit()
kmf.with_df = gdf
kmf.with_df._cderi = gdf_fname
kmf.conv_tol = 1e-12
kmf.chkfile = chkfname
kmf.kernel()
rpa = KRPA(kmf)
rpa.kernel()
assert abs(rpa.e_corr - -0.20723389722097715) < 1e-6
assert abs(rpa.e_tot - -10.716348738655793) < 1e-6
rpa = KRPA(kmf)
rpa.fc = False
rpa.kernel()
assert abs(rpa.e_corr - -0.1852772037535004) < 1e-6
assert abs(rpa.e_tot - -10.694392044197565) < 1e-6
rpa = KRPA(kmf)
rpa.outcore = True
rpa.segsize = 2
rpa.kernel()
assert abs(rpa.e_corr - -0.1852772037535004) < 1e-6
assert abs(rpa.e_tot - -10.694392044197565) < 1e-6
# Test on Na (metallic)
cell = gto.Cell()
cell.build(
unit='angstrom',
a="""
-2.11250000000000 2.11250000000000 2.11250000000000
2.11250000000000 -2.11250000000000 2.11250000000000
2.11250000000000 2.11250000000000 -2.11250000000000
""",
atom="""
Na 0.00000 0.00000 0.00000
""",
dimension=3,
max_memory=126000,
verbose=5,
pseudo='gth-pade',
basis='gth-dzvp-molopt-sr',
precision=1e-10,
)
kpts = cell.make_kpts([2, 2, 1], scaled_center=[0, 0, 0])
gdf = df.RSGDF(cell, kpts)
gdf_fname = 'gdf_ints_221.h5'
gdf._cderi_to_save = gdf_fname
if not os.path.isfile(gdf_fname):
gdf.build()
chkfname = 'na_221.chk'
if os.path.isfile(chkfname):
kmf = dft.KRKS(cell, kpts).rs_density_fit()
kmf = scf.addons.smearing_(kmf, sigma=5e-3, method='fermi')
kmf.xc = 'lda'
kmf.with_df = gdf
kmf.with_df._cderi = gdf_fname
data = chkfile.load(chkfname, 'scf')
kmf.__dict__.update(data)
else:
kmf = dft.KRKS(cell, kpts).rs_density_fit()
kmf = scf.addons.smearing_(kmf, sigma=5e-3, method='fermi')
kmf.xc = 'lda'
kmf.with_df = gdf
kmf.with_df._cderi = gdf_fname
kmf.conv_tol = 1e-12
kmf.chkfile = chkfname
kmf.kernel()
rpa = KRPA(kmf)
rpa.kernel()
assert abs(rpa.e_corr - -0.04031792880477294) < 1e-6
assert abs(rpa.e_tot - -47.60693294927457) < 1e-6
rpa = KRPA(kmf)
rpa.outcore = True
rpa.segsize = 3
rpa.kernel()
assert abs(rpa.e_corr - -0.04031792880477294) < 1e-6
assert abs(rpa.e_tot - -47.60693294927457) < 1e-6
print('passed tests!')