#!/usr/bin/env python
# Copyright 2017-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.
import numpy
from pyscf import lib
import itertools
#einsum = numpy.einsum
einsum = lib.einsum
#TODO: optimize memory use
def _gamma1_intermediates(cc, t1, t2, l1=None, l2=None):
if l1 is None:
l1 = [amp.conj() for amp in t1]
if l2 is None:
l2 = [amp.conj() for amp in t2]
t1a, t1b = t1
t2aa, t2ab, t2bb = t2
l1a, l1b = l1
l2aa, l2ab, l2bb = l2
nkpts, nocca, nvira = t1a.shape
_, noccb, nvirb = t1b.shape
kconserv = cc.khelper.kconserv
dooa = -einsum('xie,xje->xij', l1a, t1a)
dooa -= einsum('xyzimef,xyzjmef->xij', l2ab, t2ab)
dooa -= einsum('xyzimef,xyzjmef->xij', l2aa, t2aa) * .5
doob = -einsum('xie,xje->xij', l1b, t1b)
doob -= einsum('yxzmief,yxzmjef->xij', l2ab, t2ab)
doob -= einsum('xyzimef,xyzjmef->xij', l2bb, t2bb) * .5
dvva = einsum('xma,xmb->xab', t1a, l1a)
dvva += einsum('xyzmnae,xyzmnbe->zab', t2ab, l2ab)
dvva += einsum('xyzmnae,xyzmnbe->zab', t2aa, l2aa) * .5
dvvb = einsum('xma,xmb->xab', t1b, l1b)
for km, kn, ke in itertools.product(range(nkpts),repeat=3):
ka = kconserv[km,ke,kn]
dvvb[ka] += einsum('mnea,mneb->ab', t2ab[km,kn,ke], l2ab[km,kn,ke])
dvvb += einsum('xyzmnae,xyzmnbe->zab', t2bb, l2bb) * .5
xt1a = einsum('xyzmnef,xyzinef->xmi', l2aa, t2aa) * .5
xt1a += einsum('xyzmnef,xyzinef->xmi', l2ab, t2ab)
xt2a = einsum('xyzmnaf,xyzmnef->zae', t2aa, l2aa) * .5
xt2a += einsum('xyzmnaf,xyzmnef->zae', t2ab, l2ab)
xt2a += einsum('xma,xme->xae', t1a, l1a)
dvoa = t1a.copy().transpose(0,2,1)
for ka, km in itertools.product(range(nkpts),repeat=2):
dvoa[ka] += einsum('imae, me->ai', t2aa[ka,km,ka], l1a[km])
dvoa[ka] += einsum('imae, me->ai', t2ab[ka,km,ka], l1b[km])
dvoa -= einsum('xmi,xma->xai', xt1a, t1a)
dvoa -= einsum('xie,xae->xai', t1a, xt2a)
xt1b = einsum('xyzmnef,xyzinef->xmi', l2bb, t2bb) * .5
xt1b += einsum('xyznmef,xyznief->ymi', l2ab, t2ab)
xt2b = einsum('xyzmnaf,xyzmnef->zae', t2bb, l2bb) * .5
for km, kn, kf in itertools.product(range(nkpts),repeat=3):
ka = kconserv[km,kf,kn]
xt2b[ka] += einsum('mnfa,mnfe->ae', t2ab[km,kn,kf], l2ab[km,kn,kf])
xt2b += einsum('xma,xme->xae', t1b, l1b)
dvob = t1b.copy().transpose(0,2,1)
for ka, km in itertools.product(range(nkpts),repeat=2):
dvob[ka] += einsum('imae,me->ai', t2bb[ka,km,ka], l1b[km])
dvob[ka] += einsum('miea,me->ai', t2ab[km,ka,km], l1a[km])
dvob -= einsum('xmi, xma->xai',xt1b, t1b)
dvob -= einsum('xie,xae->xai', t1b, xt2b)
dova = l1a
dovb = l1b
return ((dooa, doob), (dova, dovb), (dvoa, dvob), (dvva, dvvb))
[docs]
def make_rdm1(mycc, t1, t2, l1=None, l2=None, ao_repr=False):
r'''
One-particle spin density matrices dm1a, dm1b in MO basis (the
occupied-virtual blocks due to the orbital response contribution are not
included).
dm1a[p,q] = <q_alpha^\dagger p_alpha>
dm1b[p,q] = <q_beta^\dagger p_beta>
The convention of 1-pdm is based on McWeeney's book, Eq (5.4.20).
'''
d1 = _gamma1_intermediates(mycc, t1, t2, l1, l2)
return _make_rdm1(mycc, d1, with_frozen=True, ao_repr=ao_repr)
def _make_rdm1(mycc, d1, with_frozen=True, ao_repr=False):
doo, dOO = d1[0]
dov, dOV = d1[1]
dvo, dVO = d1[2]
dvv, dVV = d1[3]
nkpts, nocca, nvira = dov.shape
_, noccb, nvirb = dOV.shape
nmoa = nocca + nvira
nmob = noccb + nvirb
dtype = numpy.result_type(doo, dOO, dov, dOV, dvo, dVO, dvv, dVV)
dm1a = numpy.empty((nkpts,nmoa,nmoa), dtype=dtype)
dm1a[:,:nocca,:nocca] = doo + doo.conj().transpose(0,2,1)
dm1a[:,:nocca,nocca:] = dov + dvo.conj().transpose(0,2,1)
dm1a[:,nocca:,:nocca] = dm1a[:,:nocca,nocca:].conj().transpose(0,2,1)
dm1a[:,nocca:,nocca:] = dvv + dvv.conj().transpose(0,2,1)
dm1a *= .5
for k in range(nkpts):
dm1a[k][numpy.diag_indices(nocca)] +=1
dm1b = numpy.empty((nkpts,nmob,nmob), dtype=dtype)
dm1b[:,:noccb,:noccb] = dOO + dOO.conj().transpose(0,2,1)
dm1b[:,:noccb,noccb:] = dOV + dVO.conj().transpose(0,2,1)
dm1b[:,noccb:,:noccb] = dm1b[:,:noccb,noccb:].conj().transpose(0,2,1)
dm1b[:,noccb:,noccb:] = dVV + dVV.conj().transpose(0,2,1)
dm1b *= .5
for k in range(nkpts):
dm1b[k][numpy.diag_indices(noccb)] +=1
if with_frozen and mycc.frozen is not None:
raise NotImplementedError
_, nmoa = mycc.mo_occ[0].size
_, nmob = mycc.mo_occ[1].size
nocca = numpy.count_nonzero(mycc.mo_occ[0] > 0)
noccb = numpy.count_nonzero(mycc.mo_occ[1] > 0)
rdm1a = numpy.zeros((nkpts,nmoa,nmoa), dtype=dm1a.dtype)
rdm1b = numpy.zeros((nkpts,nmob,nmob), dtype=dm1b.dtype)
rdm1a[:,numpy.diag_indices(nocca)] = 1
rdm1b[:,numpy.diag_indices(noccb)] = 1
moidx = mycc.get_frozen_mask()
moidxa = numpy.where(moidx[0])[0]
moidxb = numpy.where(moidx[1])[0]
rdm1a[moidxa[:,None],moidxa] = dm1a
rdm1b[moidxb[:,None],moidxb] = dm1b
dm1a = rdm1a
dm1b = rdm1b
if ao_repr:
mo_a, mo_b = mycc.mo_coeff
dm1a = lib.einsum('xpi,xij,xqj->xpq', mo_a, dm1a, mo_a.conj())
dm1b = lib.einsum('xpi,xij,xqj->xpq', mo_b, dm1b, mo_b.conj())
return dm1a, dm1b
