#!/usr/bin/env python
# Copyright 2022-2023 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.
#
# Authors: Xing Zhang <zhangxing.nju@gmail.com>
#
from functools import reduce
import numpy as np
from pyscf import lib
from pyscf.lib import logger
from pyscf.pbc.lib import ktensor
from pyscf.pbc.lib.kpts import MORotationMatrix, KQuartets
from pyscf.pbc.df import GDF, RSGDF
from pyscf.pbc.mp.kmp2 import (
get_nocc,
padded_mo_coeff,
padding_k_idx,
)
from pyscf.pbc.cc import kintermediates_rhf_ksymm as imdk
from pyscf.pbc.cc.kccsd_rhf import (
RCCSD,
_get_epq,
_init_df_eris,
)
einsum = lib.einsum
[docs]
def update_amps(cc, t1, t2, eris):
kpts = cc.kpts
kqrts = cc.kqrts
rmat = cc.rmat
kconserv = cc.khelper.kconserv
nkpts, nocc, nvir = t1.shape
fock = eris.fock
mo_e_o = [e[:nocc] for e in eris.mo_energy]
mo_e_v = [e[nocc:] for e in eris.mo_energy]
# Get location of padded elements in occupied and virtual space
nonzero_opadding, nonzero_vpadding = padding_k_idx(cc, kind="split")
ki_ibz_bz = kpts.ibz2bz[np.arange(kpts.nkpts_ibz)]
fov = fock[:, :nocc, nocc:]
kconserv = cc.khelper.kconserv
Foo = imdk.cc_Foo(kpts, kqrts, t1, t2, eris, rmat)
Fvv = imdk.cc_Fvv(kpts, kqrts, t1, t2, eris, rmat)
Fov = imdk.cc_Fov(kpts, kqrts, t1, t2, eris, rmat)
Loo = imdk.Loo(kpts, kqrts, t1, t2, eris, rmat)
Lvv = imdk.Lvv(kpts, kqrts, t1, t2, eris, rmat)
Fov = Fov.todense()
# Move energy terms to the other side
for ki_ibz in range(kpts.nkpts_ibz):
ki = kpts.ibz2bz[ki_ibz]
Foo[ki][np.diag_indices(nocc)] -= mo_e_o[ki]
Fvv[ki][np.diag_indices(nvir)] -= mo_e_v[ki]
Loo[ki][np.diag_indices(nocc)] -= mo_e_o[ki]
Lvv[ki][np.diag_indices(nvir)] -= mo_e_v[ki]
t1new = ktensor.empty_like(t1)
t1 = t1.todense()
t2new = ktensor.empty_like(t2)
t2 = t2.todense()
# T1 equation
for ka_ibz in range(kpts.nkpts_ibz):
ki = ka = kpts.ibz2bz[ka_ibz]
t1new[ka] = fov[ka].conj()
t1new[ka] += -2. * einsum('kc,ka,ic->ia', fov[ki], t1[ka], t1[ki])
t1new[ka] += einsum('ac,ic->ia', Fvv[ka], t1[ki])
t1new[ka] += -einsum('ki,ka->ia', Foo[ki], t1[ka])
for i, kq in enumerate(kqrts.kqrts_ibz):
ki, kk, kc, kd = kq
ka = ki
Svovv = 2 * eris.vovv[ka, kk, kc] - eris.vovv[ka, kk, kd].transpose(0, 1, 3, 2)
tau_term_1 = t2[ki, kk, kc].copy()
if ki == kc and kk == kd:
tau_term_1 += einsum('ic,kd->ikcd', t1[ki], t1[kk])
fock = einsum('akcd,ikcd->ia', Svovv, tau_term_1)
for _, iop in kqrts.loop_stabilizer(i):
rmat_oo = rmat.oo[ka][iop]
rmat_vv = rmat.vv[ka][iop]
t1new[ka] += einsum('ia,im,ae->me', fock, rmat_oo, rmat_vv.conj())
for i, kq in enumerate(kqrts.kqrts_ibz):
kk, kl, ki, kd = kq
ka = ki
Sooov = 2 * eris.ooov[kk, kl, ki] - eris.ooov[kl, kk, ki].transpose(1, 0, 2, 3)
tau_term_1 = t2[kk, kl, ka].copy()
if kk == ka and kl == kc:
tau_term_1 += einsum('ka,lc->klac', t1[ka], t1[kc])
fock = -einsum('klic,klac->ia', Sooov, tau_term_1)
op_group = kqrts.stars_ops[i]
ka_prim = kpts.k2opk[ka, op_group]
mask = np.isin(ka_prim, ki_ibz_bz)
for iop, ka_p in zip(op_group[mask], ka_prim[mask]):
rmat_oo = rmat.oo[ka][iop]
rmat_vv = rmat.vv[ka][iop]
t1new[ka_p] += einsum('ia,im,ae->me', fock, rmat_oo, rmat_vv.conj())
for i, kq in enumerate(kqrts.kqrts_ibz):
ki, kk, ka, kc = kq
if ka == ki and kk == kc:
tau_term = 2 * t2[kk, ki, kk] - t2[ki, kk, kk].transpose(1, 0, 2, 3)
if ki == kk:
tau_term += einsum('ic,ka->kica', t1[ki], t1[ka])
fock = einsum('kc,kica->ia', Fov[kc], tau_term)
fock += einsum('akic,kc->ia', 2 * eris.voov[ka, kk, ki], t1[kc])
fock += einsum('kaic,kc->ia', -eris.ovov[kk, ka, ki], t1[kc])
for _, iop in kqrts.loop_stabilizer(i):
rmat_oo = rmat.oo[ka][iop]
rmat_vv = rmat.vv[ka][iop]
t1new[ka] += einsum('ia,im,ae->me', fock, rmat_oo, rmat_vv.conj())
for ki_ibz in range(kpts.nkpts_ibz):
ka = ki = kpts.ibz2bz[ki_ibz]
# Remove zero/padded elements from denominator
eia = _get_epq([0,nocc,ki,mo_e_o,nonzero_opadding],
[0,nvir,ka,mo_e_v,nonzero_vpadding],
fac=[1.0,-1.0])
t1new[ki] /= eia
# T2 equation
Loo = Loo.todense()
Lvv = Lvv.todense()
for i, kq in enumerate(kqrts.kqrts_ibz):
ki, kj, ka, kb = kq
t2new[ki, kj, ka] = eris.oovv[ki, kj, ka].conj()
mem_now = lib.current_memory()[0]
if (cc.incore_complete or
_memory_4d(cc, [nocc,]*4) + mem_now < cc.max_memory * .9):
Woooo = imdk.cc_Woooo(kpts, kqrts, t1, t2, eris, rmat)
else:
metadata = {'kpts': kpts, 'kqrts': kqrts, 'rmat': rmat,
'label': 'oooo', 'trans': 'ccnn',
'incore': False}
Woooo = ktensor.empty([nocc,]*4, dtype=t1.dtype, metadata=metadata)
Woooo = imdk.cc_Woooo(kpts, kqrts, t1, t2, eris, rmat, Woooo)
mem_now = lib.current_memory()[0]
if (not cc.ktensor_direct and
_memory_4d(cc, [nocc,]*4, False) + mem_now < cc.max_memory * .9):
Woooo = Woooo.todense()
def _t2_oooo(ki,kj,ka,kb):
t2new_tmp = 0
for kl in range(nkpts):
kk = kconserv[kj, kl, ki]
tau_term = t2[kk, kl, ka].copy()
if kl == kb and kk == ka:
tau_term += einsum('ic,jd->ijcd', t1[ka], t1[kb])
t2new_tmp += 0.5 * einsum('klij,klab->ijab', Woooo[kk, kl, ki], tau_term)
return t2new_tmp
for i, kq in enumerate(kqrts.kqrts_ibz):
ki, kj, ka, kb = kq
t2new_tmp = _t2_oooo(ki,kj,ka,kb)
t2new[ki, kj, ka] += t2new_tmp
if lib.isin_1d((kj,ki,kb,ka), kqrts.kqrts_ibz):
t2new[kj, ki, kb] += t2new_tmp.transpose(1, 0, 3, 2)
else:
t2new_tmp = _t2_oooo(kj,ki,kb,ka)
t2new[ki, kj, ka] += t2new_tmp.transpose(1, 0, 3, 2)
Woooo = None
add_vvvv_(cc, t2new, t1, t2, eris)
def _t2_voov1(ki,kj,ka,kb):
t2new_tmp = einsum('ac,ijcb->ijab', Lvv[ka], t2[ki, kj, ka])
t2new_tmp += einsum('ki,kjab->ijab', -Loo[ki], t2[ki, kj, ka])
kc = kj
tmp2 = np.asarray(eris.vovv[kc, ki, kb]).transpose(3, 2, 1, 0).conj() \
- einsum('kbic,ka->abic', eris.ovov[ka, kb, ki], t1[ka])
t2new_tmp += einsum('abic,jc->ijab', tmp2, t1[kj])
kk = kb
tmp2 = np.asarray(eris.ooov[kj, ki, kk]).transpose(3, 2, 1, 0).conj() \
+ einsum('akic,jc->akij', eris.voov[ka, kk, ki], t1[kj])
t2new_tmp -= einsum('akij,kb->ijab', tmp2, t1[kb])
return t2new_tmp
for i, kq in enumerate(kqrts.kqrts_ibz):
ki, kj, ka, kb = kq
t2new_tmp = _t2_voov1(ki,kj,ka,kb)
t2new[ki, kj, ka] += t2new_tmp
if lib.isin_1d((kj,ki,kb,ka), kqrts.kqrts_ibz):
t2new[kj, ki, kb] += t2new_tmp.transpose(1, 0, 3, 2)
else:
t2new_tmp = _t2_voov1(kj,ki,kb,ka)
t2new[ki, kj, ka] += t2new_tmp.transpose(1, 0, 3, 2)
mem_now = lib.current_memory()[0]
if (cc.incore_complete or
_memory_4d(cc, [nocc,nocc,nvir,nvir])*2 + mem_now < cc.max_memory*.9):
Wvoov = imdk.cc_Wvoov(kpts, kqrts, t1, t2, eris, rmat)
Wvovo = imdk.cc_Wvovo(kpts, kqrts, t1, t2, eris, rmat)
else:
metadata = {'kpts': kpts, 'kqrts': kqrts, 'rmat': rmat,
'trans': 'ccnn', 'incore': False}
Wvoov = ktensor.empty([nvir,nocc,nocc,nvir], dtype=t1.dtype,
metadata={**metadata, 'label':'voov'})
Wvovo = ktensor.empty([nvir,nocc,nvir,nocc], dtype=t1.dtype,
metadata={**metadata, 'label':'vovo'})
Wvoov = imdk.cc_Wvoov(kpts, kqrts, t1, t2, eris, rmat, Wvoov)
Wvovo = imdk.cc_Wvovo(kpts, kqrts, t1, t2, eris, rmat, Wvovo)
mem_now = lib.current_memory()[0]
if (not cc.ktensor_direct and
_memory_4d(cc, [nocc,nocc,nvir,nvir], False)*2 + mem_now < cc.max_memory*.9):
Wvoov = Wvoov.todense()
Wvovo = Wvovo.todense()
def _t2_voov2(ki,kj,ka,kb):
t2new_tmp = 0
for kk in range(nkpts):
kc = kconserv[ka, ki, kk]
tmp_voov = 2. * Wvoov[ka, kk, ki] - Wvovo[ka, kk, kc].transpose(0, 1, 3, 2)
t2new_tmp += einsum('akic,kjcb->ijab', tmp_voov, t2[kk, kj, kc])
kc = kconserv[ka, ki, kk]
t2new_tmp -= einsum('akic,kjbc->ijab', Wvoov[ka, kk, ki], t2[kk, kj, kb])
kc = kconserv[kk, ka, kj]
t2new_tmp -= einsum('bkci,kjac->ijab', Wvovo[kb, kk, kc], t2[kk, kj, ka])
return t2new_tmp
for i, kq in enumerate(kqrts.kqrts_ibz):
ki, kj, ka, kb = kq
t2new_tmp = _t2_voov2(ki,kj,ka,kb)
t2new[ki, kj, ka] += t2new_tmp
if lib.isin_1d((kj,ki,kb,ka), kqrts.kqrts_ibz):
t2new[kj, ki, kb] += t2new_tmp.transpose(1, 0, 3, 2)
else:
t2new_tmp = _t2_voov2(kj,ki,kb,ka)
t2new[ki, kj, ka] += t2new_tmp.transpose(1, 0, 3, 2)
Wvoov = Wvovo = None
for i, kq in enumerate(kqrts.kqrts_ibz):
ki, kj, ka, kb = kq
eia = _get_epq([0,nocc,ki,mo_e_o,nonzero_opadding],
[0,nvir,ka,mo_e_v,nonzero_vpadding],
fac=[1.0,-1.0])
ejb = _get_epq([0,nocc,kj,mo_e_o,nonzero_opadding],
[0,nvir,kb,mo_e_v,nonzero_vpadding],
fac=[1.0,-1.0])
eijab = eia[:, None, :, None] + ejb[:, None, :]
t2new[ki, kj, ka] /= eijab
return t1new, t2new
[docs]
def add_vvvv_(cc, Ht2, t1, t2, eris):
kpts = cc.kpts
kqrts = cc.kqrts
rmat = cc.rmat
nocc = cc.nocc
nmo = cc.nmo
nvir = nmo - nocc
nkpts = kpts.nkpts
kconserv = cc.khelper.kconserv
mem_now = lib.current_memory()[0]
if (not cc.incore_complete and
cc.direct and getattr(eris, 'Lpv', None) is not None):
def get_Wvvvv(ka, kb, kc):
Lpv = eris.Lpv
kd = kconserv[ka, kc, kb]
Lbd = (Lpv[kb,kd][:,nocc:] -
lib.einsum('Lkd,kb->Lbd', Lpv[kb,kd][:,:nocc], t1[kb]))
Wvvvv = lib.einsum('Lac,Lbd->abcd', Lpv[ka,kc][:,nocc:], Lbd)
Lbd = None
kcbd = lib.einsum('Lkc,Lbd->kcbd', Lpv[ka,kc][:,:nocc],
Lpv[kb,kd][:,nocc:])
Wvvvv -= lib.einsum('kcbd,ka->abcd', kcbd, t1[ka])
Wvvvv *= (1. / nkpts)
return Wvvvv
elif (cc.incore_complete or
_memory_4d(cc, [nvir,]*4) + mem_now < cc.max_memory * .9):
_Wvvvv = imdk.cc_Wvvvv(kpts, kqrts, t1, t2, eris, rmat)
mem_now = lib.current_memory()[0]
if (not cc.ktensor_direct and
_memory_4d(cc, [nvir,]*4, False) + mem_now < cc.max_memory * .9):
_Wvvvv = _Wvvvv.todense()
get_Wvvvv = lambda ka, kb, kc: _Wvvvv[ka, kb, kc]
else:
metadata = {'kpts': kpts, 'kqrts': kqrts, 'rmat': rmat,
'label': 'vvvv', 'trans': 'ccnn',
'incore': False}
_Wvvvv = ktensor.empty([nvir,]*4, dtype=t1.dtype, metadata=metadata)
_Wvvvv = imdk.cc_Wvvvv(kpts, kqrts, t1, t2, eris, rmat, _Wvvvv)
mem_now = lib.current_memory()[0]
if (not cc.ktensor_direct and
_memory_4d(cc, [nvir,]*4, False) + mem_now < cc.max_memory * .9):
_Wvvvv = _Wvvvv.todense()
get_Wvvvv = lambda ka, kb, kc: _Wvvvv[ka, kb, kc]
kakb, igroup = np.unique(kqrts.kqrts_ibz[:,2:], axis=0, return_inverse=True)
for i in range(np.amax(igroup) + 1):
ka, kb = kakb[i]
idx = np.where(igroup==i)[0]
for kc in range(nkpts):
kd = kconserv[ka, kc, kb]
Wvvvv = get_Wvvvv(ka, kb, kc)
for m in idx:
ki,kj,kaa,kbb = kqrts.kqrts_ibz[m]
assert kaa==ka and kbb==kb
tau = t2[ki, kj, kc].copy()
if ki == kc and kj == kd:
tau += np.einsum('ic,jd->ijcd', t1[ki], t1[kj])
Ht2[ki, kj, ka] += einsum('abcd,ijcd->ijab', Wvvvv, tau)
_Wvvvv = None
return Ht2
[docs]
def energy(cc, t1, t2, eris):
kpts = cc.kpts
kqrts = cc.kqrts
nkpts, nocc, nvir = t1.shape
fock = eris.fock
e = 0
for ki_ibz in range(kpts.nkpts_ibz):
ki = kpts.ibz2bz[ki_ibz]
weight = kpts.weights_ibz[ki_ibz]
e += 2 * einsum('ia,ia', fock[ki,:nocc,nocc:], t1[ki]) * weight
tau = ktensor.empty_like(t2)
for kq in kqrts.kqrts_ibz:
ki, kj, ka, kb = kq
tau[ki, kj, ka] = t2[ki, kj, ka]
if ki == ka and kj == kb:
tau[ki, kj, ka] += einsum('ia,jb->ijab', t1[ki], t1[kj])
kq_weights = kqrts.weights_ibz
for k, kq in enumerate(kqrts.kqrts_ibz):
ki, kj, ka, kb = kq
weight = kq_weights[k] * nkpts**3
e += 2 * einsum('ijab,ijab', tau[ki, kj, ka], eris.oovv[ki, kj, ka]) * weight
e += -einsum('ijab,ijba', tau[ki, kj, ka], eris.oovv[ki, kj, kb]) * weight
e /= nkpts
if abs(e.imag) > 1e-4:
logger.warn(cc, 'Non-zero imaginary part found in KRCCSD energy %s', e)
return e.real
[docs]
class KsymAdaptedRCCSD(RCCSD):
_keys = set(['kqrts', 'rmat', 'ktensor_direct', 'eris_outcore'])
def __init__(self, mf, frozen=None, mo_coeff=None, mo_occ=None):
'''
Attributes:
ktensor_direct : bool
If set to True, the tensors will be stored as block-sparse,
and the symmetry related blocks are computed on-the-fly when needed.
Otherwise, the tensors will be converted to dense tensors whenever
there is enough memory. Default is False.
eris_outcore : bool
If set to True, the integrals will be always stored on the disk.
Otherwise, whether the integrals are stored on the disk or in memory
depends on the available memory size. Default is False.
'''
# NOTE self._scf is a non-symmetry object, see RCCSD.__init__
RCCSD.__init__(self, mf, frozen, mo_coeff, mo_occ)
self.kqrts = KQuartets(mf.kpts).build()
self.rmat = None
self.ktensor_direct = False
self.eris_outcore = False
[docs]
def ao2mo(self, mo_coeff=None):
eris = _PhysicistsERIs()
eris._common_init_(self, mo_coeff)
# use padded mo_coeff to construct the rotation matrix
self.rmat = MORotationMatrix(self.kpts, eris.mo_coeff,
self._scf.get_ovlp(), eris.nocc, eris.nmo)
self.rmat.build()
mem_now = lib.current_memory()[0]
mem_incore = _mem_usage(self)
if not self.eris_outcore and (
self.incore_complete or mem_incore + mem_now < self.max_memory):
eris = _make_eris_incore(self, eris, self._scf.with_df.ao2mo)
else:
eris = _make_eris_outcore(self, eris, self._scf.with_df.ao2mo)
return eris
[docs]
def init_amps(self, eris):
time0 = logger.process_clock(), logger.perf_counter()
nocc = self.nocc
nvir = self.nmo - nocc
nkpts = self.nkpts
kpts = self.kpts
kqrts = self.kqrts
rmat = self.rmat
assert rmat is not None
metadata = {'kpts': kpts, 'rmat': rmat,
'label': 'ov', 'trans': 'nc', 'incore': True}
t1 = ktensor.zeros((nocc, nvir), dtype=eris.fock.dtype,
metadata=metadata)
metadata = {'kpts': kpts, 'kqrts': kqrts, 'rmat': rmat,
'label': 'oovv', 'trans': 'nncc', 'incore': True}
t2 = ktensor.empty((nocc,nocc,nvir,nvir), dtype=eris.fock.dtype,
metadata=metadata)
mo_e_o = [eris.mo_energy[k][:nocc] for k in range(nkpts)]
mo_e_v = [eris.mo_energy[k][nocc:] for k in range(nkpts)]
# Get location of padded elements in occupied and virtual space
nonzero_opadding, nonzero_vpadding = padding_k_idx(self, kind="split")
emp2 = 0
for i, kq in enumerate(kqrts.kqrts_ibz):
ki, kj, ka, kb = kq
weight = kqrts.weights_ibz[i] * nkpts**3
eia = _get_epq([0,nocc,ki,mo_e_o,nonzero_opadding],
[0,nvir,ka,mo_e_v,nonzero_vpadding],
fac=[1.0,-1.0])
ejb = _get_epq([0,nocc,kj,mo_e_o,nonzero_opadding],
[0,nvir,kb,mo_e_v,nonzero_vpadding],
fac=[1.0,-1.0])
eijab = eia[:, None, :, None] + ejb[:, None, :]
eris_ijab = eris.oovv[ki, kj, ka]
eris_ijba = eris.oovv[ki, kj, kb]
t2[ki, kj, ka] = eris_ijab.conj() / eijab
woovv = 2 * eris_ijab - eris_ijba.transpose(0, 1, 3, 2)
emp2 += np.einsum('ijab,ijab', t2[ki, kj, ka], woovv) * weight
self.emp2 = emp2.real / nkpts
logger.info(self, 'Init t2, MP2 energy (with fock eigenvalue shift) = %.15g', self.emp2)
logger.timer(self, 'init mp2', *time0)
return self.emp2, t1, t2
[docs]
def amplitudes_to_vector(self, t1, t2):
t1_raw = np.asarray(getattr(t1, 'data', t1))
t2_raw = np.asarray(getattr(t2, 'data', t2))
return np.concatenate((t1_raw, t2_raw), axis=None)
[docs]
def vector_to_amplitudes(self, vec):
kpts = self.kpts
kqrts = self.kqrts
rmat = self.rmat
nocc = self.nocc
nvir = self.nmo - nocc
t1_size = kpts.nkpts_ibz * nocc * nvir
t1_flat = vec[:t1_size]
t2_flat = vec[t1_size:]
metadata = {'kpts': kpts, 'rmat': rmat,
'label': 'ov', 'trans': 'nc', 'incore': True}
t1 = ktensor.fromraw(t1_flat, (nocc,nvir), dtype=vec.dtype,
metadata=metadata)
metadata = {'kpts': kpts, 'kqrts': kqrts, 'rmat': rmat,
'label': 'oovv', 'trans': 'nncc', 'incore': True}
t2 = ktensor.fromraw(t2_flat, (nocc,nocc,nvir,nvir), dtype=vec.dtype,
metadata=metadata)
return t1, t2
energy = energy
update_amps = update_amps
def _make_eris_incore(cc, eris, fao2mo):
log = logger.Logger(cc.stdout, cc.verbose)
log.info('using incore ERI storage')
cput0 = (logger.process_clock(), logger.perf_counter())
mo_coeff = eris.mo_coeff
nocc = eris.nocc
nvir = eris.nvir
nmo = eris.nmo
dtype = eris.dtype
common_metadata = {'kpts' : cc.kpts,
'kqrts' : cc.kqrts,
'rmat' : cc.rmat,
'trans' : 'ccnn',
'incore': True}
eris.oooo = ktensor.empty([nocc,nocc,nocc,nocc], dtype=dtype,
metadata={**common_metadata, 'label': 'oooo'})
eris.ooov = ktensor.empty([nocc,nocc,nocc,nvir], dtype=dtype,
metadata={**common_metadata, 'label': 'ooov'})
eris.oovv = ktensor.empty([nocc,nocc,nvir,nvir], dtype=dtype,
metadata={**common_metadata, 'label': 'oovv'})
eris.ovov = ktensor.empty([nocc,nvir,nocc,nvir], dtype=dtype,
metadata={**common_metadata, 'label': 'ovov'})
eris.voov = ktensor.empty([nvir,nocc,nocc,nvir], dtype=dtype,
metadata={**common_metadata, 'label': 'voov'})
eris.vovv = ktensor.empty([nvir,nocc,nvir,nvir], dtype=dtype,
metadata={**common_metadata, 'label': 'vovv'})
eris.vvvv = ktensor.empty([nvir,nvir,nvir,nvir], dtype=dtype,
metadata={**common_metadata, 'label': 'vvvv'})
kpts = cc.kpts.kpts
nkpts = len(kpts)
kqrts = cc.kqrts
khelper = cc.khelper
kconserv = khelper.kconserv
kptlist = kqrts.kqrts_ibz[:,:3][:,[0,2,1]] #chemists' notation
khelper.build_symm_map(kptlist=kptlist)
for (iki, ika, ikj) in khelper.symm_map.keys():
ikb = kconserv[iki, ika, ikj]
eri_kpt = fao2mo((mo_coeff[iki],mo_coeff[ika],mo_coeff[ikj],mo_coeff[ikb]),
(kpts[iki],kpts[ika],kpts[ikj],kpts[ikb]), compact=False)
if np.issubdtype(dtype, np.floating):
eri_kpt = eri_kpt.real
eri_kpt = eri_kpt.reshape(nmo, nmo, nmo, nmo)
for (ki, ka, kj) in khelper.symm_map[(iki, ika, ikj)]:
eri_kpt_symm = khelper.transform_symm(eri_kpt, ki, ka, kj).transpose(0, 2, 1, 3)
eris.oooo[ki, kj, ka] = eri_kpt_symm[:nocc, :nocc, :nocc, :nocc] / nkpts
eris.ooov[ki, kj, ka] = eri_kpt_symm[:nocc, :nocc, :nocc, nocc:] / nkpts
eris.oovv[ki, kj, ka] = eri_kpt_symm[:nocc, :nocc, nocc:, nocc:] / nkpts
eris.ovov[ki, kj, ka] = eri_kpt_symm[:nocc, nocc:, :nocc, nocc:] / nkpts
eris.voov[ki, kj, ka] = eri_kpt_symm[nocc:, :nocc, :nocc, nocc:] / nkpts
eris.vovv[ki, kj, ka] = eri_kpt_symm[nocc:, :nocc, nocc:, nocc:] / nkpts
eris.vvvv[ki, kj, ka] = eri_kpt_symm[nocc:, nocc:, nocc:, nocc:] / nkpts
if not cc.ktensor_direct:
eris.oooo = eris.oooo.todense()
eris.ooov = eris.ooov.todense()
eris.oovv = eris.oovv.todense()
eris.ovov = eris.ovov.todense()
eris.voov = eris.voov.todense()
eris.vovv = eris.vovv.todense()
eris.vvvv = eris.vvvv.todense()
log.timer('CCSD integral transformation', *cput0)
return eris
def _make_eris_outcore(cc, eris, fao2mo):
log = logger.Logger(cc.stdout, cc.verbose)
log.info('using outcore ERI storage')
cput0 = (logger.process_clock(), logger.perf_counter())
mo_coeff = eris.mo_coeff
nocc = eris.nocc
nvir = eris.nvir
nmo = eris.nmo
dtype = eris.dtype
common_metadata = {'kpts' : cc.kpts,
'kqrts' : cc.kqrts,
'rmat' : cc.rmat,
'trans' : 'ccnn',
'incore': False}
eris.oooo = ktensor.empty([nocc,nocc,nocc,nocc], dtype=dtype,
metadata={**common_metadata, 'label': 'oooo'})
eris.ooov = ktensor.empty([nocc,nocc,nocc,nvir], dtype=dtype,
metadata={**common_metadata, 'label': 'ooov'})
eris.oovv = ktensor.empty([nocc,nocc,nvir,nvir], dtype=dtype,
metadata={**common_metadata, 'label': 'oovv'})
eris.ovov = ktensor.empty([nocc,nvir,nocc,nvir], dtype=dtype,
metadata={**common_metadata, 'label': 'ovov'})
eris.voov = ktensor.empty([nvir,nocc,nocc,nvir], dtype=dtype,
metadata={**common_metadata, 'label': 'voov'})
eris.vovv = ktensor.empty([nvir,nocc,nvir,nvir], dtype=dtype,
metadata={**common_metadata, 'label': 'vovv'})
vvvv_required = (not cc.direct
or not isinstance(cc._scf.with_df, (GDF, RSGDF))
or cc._scf.cell.dimension == 2)
if vvvv_required:
eris.vvvv = ktensor.empty([nvir,nvir,nvir,nvir], dtype=dtype,
metadata={**common_metadata, 'label': 'vvvv'})
kpts = cc.kpts.kpts
nkpts = len(kpts)
kqrts = cc.kqrts
# <ij|pq> = (ip|jq)
cput1 = logger.process_clock(), logger.perf_counter()
for kprqs in kqrts.kqrts_ibz:
kp, kr, kq, ks = kprqs
orbo_p = mo_coeff[kp][:, :nocc]
orbo_r = mo_coeff[kr][:, :nocc]
buf_kpt = fao2mo((orbo_p, mo_coeff[kq], orbo_r, mo_coeff[ks]),
(kpts[kp], kpts[kq], kpts[kr], kpts[ks]), compact=False)
if np.issubdtype(dtype, np.floating):
buf_kpt = buf_kpt.real
buf_kpt = buf_kpt.reshape(nocc, nmo, nocc, nmo).transpose(0, 2, 1, 3)
eris.oooo[kp, kr, kq] = buf_kpt[:, :, :nocc, :nocc] / nkpts
eris.ooov[kp, kr, kq] = buf_kpt[:, :, :nocc, nocc:] / nkpts
eris.oovv[kp, kr, kq] = buf_kpt[:, :, nocc:, nocc:] / nkpts
cput1 = log.timer_debug1('transforming oopq', *cput1)
# <ia|jb> = (ij|ab)
for kiajb in kqrts.kqrts_ibz:
ki, ka, kj, kb = kiajb
orb_i = mo_coeff[ki][:, :nocc]
orb_a = mo_coeff[ka][:, nocc:]
orb_j = mo_coeff[kj][:, :nocc]
orb_b = mo_coeff[kb][:, nocc:]
buf_kpt = fao2mo((orb_i, orb_j, orb_a, orb_b),
(kpts[ki], kpts[kj], kpts[ka], kpts[kb]), compact=False)
if np.issubdtype(dtype, np.floating):
buf_kpt = buf_kpt.real
buf_kpt = buf_kpt.reshape(nocc, nocc, nvir, nvir).transpose(0, 2, 1, 3)
eris.ovov[ki, ka, kj] = buf_kpt / nkpts
cput1 = log.timer_debug1('transforming ovov', *cput1)
# <ai|pq> = (ap|iq)
for kaipq in kqrts.kqrts_ibz:
ka, ki, kp, kq = kaipq
orb_a = mo_coeff[ka][:, nocc:]
orb_i = mo_coeff[ki][:, :nocc]
buf_kpt = fao2mo((orb_a, mo_coeff[kp], orb_i, mo_coeff[kq]),
(kpts[ka], kpts[kp], kpts[ki], kpts[kq]), compact=False)
if np.issubdtype(dtype, np.floating):
buf_kpt = buf_kpt.real
buf_kpt = buf_kpt.reshape(nvir, nmo, nocc, nmo).transpose(0, 2, 1, 3)
# TODO: compute vovv on the fly
eris.vovv[ka, ki, kp] = buf_kpt[:, :, nocc:, nocc:] / nkpts
eris.voov[ka, ki, kp] = buf_kpt[:, :, :nocc, nocc:] / nkpts
cput1 = log.timer_debug1('transforming vopq', *cput1)
mem_now = lib.current_memory()[0]
if not vvvv_required:
_init_df_eris(cc, eris)
elif nvir ** 4 * 16 / 1e6 + mem_now < cc.max_memory:
khelper = cc.khelper
kconserv = khelper.kconserv
kptlist = kqrts.kqrts_ibz[:,:3][:,[0,2,1]] #chemists' notation
khelper.build_symm_map(kptlist=kptlist)
for (ikp, ikq, ikr) in khelper.symm_map.keys():
iks = kconserv[ikp, ikq, ikr]
orbv_p = mo_coeff[ikp][:, nocc:]
orbv_q = mo_coeff[ikq][:, nocc:]
orbv_r = mo_coeff[ikr][:, nocc:]
orbv_s = mo_coeff[iks][:, nocc:]
buf_kpt = fao2mo((orbv_p,orbv_q,orbv_r,orbv_s),
kpts[[ikp,ikq,ikr,iks]], compact=False)
if np.issubdtype(dtype, np.floating):
buf_kpt = buf_kpt.real
buf_kpt = buf_kpt.reshape([nvir,]*4)
for (kp, kq, kr) in khelper.symm_map[(ikp, ikq, ikr)]:
buf_kpt_symm = khelper.transform_symm(buf_kpt, kp, kq, kr).transpose(0, 2, 1, 3)
eris.vvvv[kp, kr, kq] = buf_kpt_symm / nkpts
else:
raise MemoryError(f'Minimal memory requirements '
f'{mem_now + nvir ** 4 / 1e6 * 16 * 2} MB')
cput1 = log.timer_debug1('transforming vvvv', *cput1)
log.timer('CCSD integral transformation', *cput0)
return eris
class _PhysicistsERIs():
def __init__(self, cell=None):
self.cell = cell
self.kpts = None
self.mo_coeff = None
self.nocc = None
self.nmo = None
self.nvir = None
self.fock = None
self.dtype = None
self.oooo = None
self.ooov = None
self.oovv = None
self.ovov = None
self.voov = None
self.vovv = None
self.vvvv = None
self.Lpv = None
def _common_init_(self, cc, mo_coeff=None):
from pyscf.pbc import tools
from pyscf.pbc.cc.ccsd import _adjust_occ
mf = cc._scf
self.cell = cell = mf.cell
self.kpts = kpts = cc.kpts
self.nocc = nocc = cc.nocc
self.nmo = nmo = cc.nmo
self.nvir = nmo - nocc
if mo_coeff is None:
mo_coeff = cc.mo_coeff
self.dtype = mo_coeff[-1].dtype
# Re-make our fock MO matrix elements from density and fock AO
# FIXME what if mo_coeff is not consistent with cc.mo_occ?
dm = mf.make_rdm1(mo_coeff, cc.mo_occ)
exxdiv = mf.exxdiv if cc.keep_exxdiv else None
with lib.temporary_env(mf, exxdiv=exxdiv):
# _scf.exxdiv affects eris.fock. HF exchange correction should be
# excluded from the Fock matrix.
vhf = mf.get_veff(cell, dm)
fockao = mf.get_hcore() + vhf
self.mo_coeff = mo_coeff = padded_mo_coeff(cc, mo_coeff)
fock = np.asarray([reduce(np.dot, (mo.T.conj(), fockao[k], mo))
for k, mo in enumerate(mo_coeff)])
self.fock = fock
mo_energy = [fock[k].diagonal().real for k in range(len(fock))]
if not cc.keep_exxdiv:
# Add HFX correction in the self.mo_energy to improve convergence in
# CCSD iteration. It is useful for the 2D systems since their occupied and
# the virtual orbital energies may overlap which may lead to numerical
# issue in the CCSD iterations.
# FIXME: Whether to add this correction for other exxdiv treatments?
# Without the correction, MP2 energy may be largely off the correct value.
madelung = tools.madelung(cell, kpts.kpts)
mo_energy = [_adjust_occ(mo_e, nocc, -madelung)
for k, mo_e in enumerate(mo_energy)]
# Get location of padded elements in occupied and virtual space.
nocc_per_kpt = get_nocc(cc, per_kpoint=True)
nonzero_padding = padding_k_idx(cc, kind="joint")
# Check direct and indirect gaps for possible issues with CCSD convergence.
mo_e = [mo_energy[kp][nonzero_padding[kp]] for kp in range(len(mo_energy))]
mo_e = np.sort([y for x in mo_e for y in x]) # Sort de-nested array
gap = mo_e[np.sum(nocc_per_kpt)] - mo_e[np.sum(nocc_per_kpt)-1]
if gap < 1e-5:
logger.warn(cc, 'HOMO-LUMO gap %s too small for KCCSD. '
'May cause issues in convergence.', gap)
self.mo_energy = mo_energy
def _mem_usage(cc):
nocc = cc.nocc
nmo = cc.nmo
nvir = nmo - nocc
nkpts = cc.kpts.nkpts
#Wvvvv, eris.vvvv
incore = _memory_4d(cc, [nvir,]*4) * 2
#eris.vovv
incore += _memory_4d(cc, [nvir,nvir,nvir,nocc])
#Wvoov, Wvovo, eris.oovv, eris.ovov, eris.voov, t2new, t2
incore += _memory_4d(cc, [nvir,nvir,nocc,nocc]) * 7
#eris.ooov
incore += _memory_4d(cc, [nvir,nocc,nocc,nocc])
#Woooo, eris.oooo
incore += _memory_4d(cc, [nocc,]*4) * 2
if not cc.ktensor_direct:
incore += _memory_4d(cc, [nvir,]*4, False) * 2
incore += _memory_4d(cc, [nvir,nvir,nvir,nocc], False)
incore += _memory_4d(cc, [nvir,nvir,nocc,nocc], False) * 5
incore += _memory_4d(cc, [nvir,nocc,nocc,nocc], False)
incore += _memory_4d(cc, [nocc,]*4, False) * 2
#temp
incore += nkpts * nmo**4 * 16 / 1e6
#1e, fock, t1
incore += nkpts * nmo**2 * 16 / 1e6 * 7
#t2_bz
incore += _memory_4d(cc, [nvir,nvir,nocc,nocc], False)
logger.info(cc, f"Incore memory estimation: {incore} MB")
return incore
def _memory_4d(cc, shape, ibz=True):
if ibz:
nk = len(cc.kqrts.kqrts_ibz)
else:
nk = cc.kpts.nkpts**3
return nk * np.prod(shape) * 16 / 1e6