#!/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.
#
# Authors: James D. McClain
# Mario Motta
# Yang Gao
# Qiming Sun <osirpt.sun@gmail.com>
# Jason Yu
#
import itertools
import numpy as np
from pyscf import lib
from pyscf.lib import logger
from pyscf.pbc.cc import eom_kccsd_ghf as eom_kgccsd
from pyscf.pbc.cc import kccsd
from pyscf.pbc.lib import kpts_helper
from pyscf.lib.parameters import LOOSE_ZERO_TOL, LARGE_DENOM # noqa
from pyscf.pbc.cc import kintermediates_uhf
from pyscf.pbc.mp.kump2 import (get_frozen_mask, get_nocc, get_nmo,
padded_mo_coeff, padding_k_idx) # noqa
einsum = lib.einsum
########################################
# EOM-IP-CCSD
########################################
[docs]
def amplitudes_to_vector_ip(r1, r2, kshift, kconserv):
r1a, r1b = r1
r2aaa, r2baa, r2abb, r2bbb = r2
nkpts = r2aaa.shape[0]
nocca, noccb = r1a.shape[0], r1b.shape[0]
nvira, nvirb = r2aaa.shape[-1], r2bbb.shape[-1]
# From symmetry for aaa and bbb terms, only store lower
# triangular part (ki,i) < (kj,j)
idxa, idya = np.tril_indices(nkpts*nocca, -1)
idxb, idyb = np.tril_indices(nkpts*noccb, -1)
r2aaa = r2aaa.transpose(0,2,1,3,4).reshape(nkpts*nocca,nkpts*nocca,nvira)
r2bbb = r2bbb.transpose(0,2,1,3,4).reshape(nkpts*noccb,nkpts*noccb,nvirb)
return np.hstack((r1a, r1b, r2aaa[idxa,idya].ravel(),
r2baa.ravel(), r2abb.ravel(),
r2bbb[idxb,idyb].ravel()))
[docs]
def vector_to_amplitudes_ip(vector, kshift, nkpts, nmo, nocc, kconserv):
nocca, noccb = nocc
nmoa, nmob = nmo
nvira, nvirb = nmoa-nocca, nmob-noccb
sizes = (nocca, noccb, (nkpts*nocca)*(nkpts*nocca-1)*nvira//2,
nkpts**2*noccb*nocca*nvira, nkpts**2*nocca*noccb*nvirb,
nkpts*noccb*(nkpts*noccb-1)*nvirb//2)
sections = np.cumsum(sizes[:-1])
r1a, r1b, r2a, r2baa, r2abb, r2b = np.split(vector, sections)
r2a = r2a.reshape(nkpts*nocca*(nkpts*nocca-1)//2,nvira)
r2b = r2b.reshape(nkpts*noccb*(nkpts*noccb-1)//2,nvirb)
idxa, idya = np.tril_indices(nkpts*nocca, -1)
idxb, idyb = np.tril_indices(nkpts*noccb, -1)
r2aaa = np.zeros((nkpts*nocca,nkpts*nocca,nvira), dtype=r2a.dtype)
r2aaa[idxa,idya] = r2a.copy()
r2aaa[idya,idxa] = -r2a.copy() # Fill in value : kj, j < ki, i
r2aaa = r2aaa.reshape(nkpts,nocca,nkpts,nocca,nvira)
r2aaa = r2aaa.transpose(0,2,1,3,4)
r2baa = r2baa.reshape(nkpts,nkpts,noccb,nocca,nvira).copy()
r2abb = r2abb.reshape(nkpts,nkpts,nocca,noccb,nvirb).copy()
r2bbb = np.zeros((nkpts*noccb,nkpts*noccb,nvirb), dtype=r2b.dtype)
r2bbb[idxb,idyb] = r2b.copy()
r2bbb[idyb,idxb] = -r2b.copy() # Fill in value : kj, j < ki, i
r2bbb = r2bbb.reshape(nkpts,noccb,nkpts,noccb,nvirb)
r2bbb = r2bbb.transpose(0,2,1,3,4)
r1 = (r1a.copy(), r1b.copy())
r2 = (r2aaa, r2baa, r2abb, r2bbb)
return r1, r2
[docs]
def ipccsd_matvec(eom, vector, kshift, imds=None, diag=None):
'''2ph operators are of the form s_{ij}^{ b}, i.e. 'jb' indices are coupled'''
if imds is None: imds = eom.make_imds()
t1, t2= imds.t1, imds.t2
t1a, t1b = t1
t2aa, t2ab, t2bb = t2
nocca, noccb, nvira, nvirb = t2ab.shape[3:]
nmoa, nmob = nocca + nvira, noccb + nvirb
kconserv = imds.kconserv
nkpts = eom.nkpts
r1, r2 = eom.vector_to_amplitudes(vector, kshift, nkpts, (nmoa, nmob), (nocca, noccb), kconserv)
#nocc = eom.nocc
#nmo = eom.nmo
#nvir = (nmo[0]-nocc[0], nmo[1]-nocc[1])
#nocca, noccb = nocc
#nvira, nvirb = nvir
#nkpts = eom.nkpts
#r1, r2 = eom.vector_to_amplitudes(vector, nkpts, nmo[0]+nmo[1], nocc[0]+nocc[1]) # spin
#spatial_r1, spatial_r2 = eom_kgccsd.spin2spatial_ip_doublet(r1, r2, kconserv, kshift, orbspin)
#imds = imds._imds
#t2aa, t2ab, t2bb = t2
# k-point spin orbital version of ipccsd
#Hr1 = -0.0*np.einsum('mi,m->i', imds.Foo[kshift], r1)
#Hr2 = np.zeros_like(r2)
r1a, r1b = r1
r2aaa, r2baa, r2abb, r2bbb = r2
#Foo term
# -\sum_{kk,k} U_{kk,k,ki,i} s_{kk,k}
Hr1a = -np.einsum('mi,m->i', imds.Foo[kshift], r1a)
Hr1b = -np.einsum('MI,M->I', imds.FOO[kshift], r1b)
#Fov term
# \sum_{kL,kD,L,D} U_{kL,kD,L,D} S_{ki,i,kL,L}^{kD,D} + \sum_{kl,kd,l,d} U_{kl,kd,l,d} S_{ki,i,kl,l}^{kd,d}
for km in range(nkpts):
Hr1a += einsum('me,mie->i', imds.Fov[km], r2aaa[km,kshift])
Hr1a -= einsum('ME,iME->i', imds.FOV[km], r2abb[kshift,km])
Hr1b += einsum('ME,MIE->I', imds.FOV[km], r2bbb[km,kshift])
Hr1b -= einsum('me,Ime->I', imds.Fov[km], r2baa[kshift,km])
#Wooov
# \sum_{kk,kl,kd,k,l,d} W_{kk,ki,kl,kd,k,i,l,d} s_{kl,kk,l,k}^{kd,d}
# \sum_{kk,kL,kD,k,L,D} W_{kk,ki,kL,kD,k,i,L,D} s_{kL,kk,L,k}^{kD,D}
for km in range(nkpts):
for kn in range(nkpts):
Hr1a += -0.5 * einsum('nime,mne->i', imds.Wooov[kn,kshift,km], r2aaa[km,kn])
Hr1b += einsum('NIme,Nme->I', imds.WOOov[kn,kshift,km], r2baa[kn,km])
Hr1b += -0.5 * einsum('NIME,MNE->I', imds.WOOOV[kn,kshift,km], r2bbb[km,kn])
Hr1a += einsum('niME,nME->i', imds.WooOV[kn,kshift,km], r2abb[kn,km])
dtype = np.result_type(Hr1a, *r2)
Hr2aaa = np.zeros((nkpts, nkpts, nocca, nocca, nvira), dtype=dtype)
Hr2baa = np.zeros((nkpts, nkpts, noccb, nocca, nvira), dtype=dtype)
Hr2abb = np.zeros((nkpts, nkpts, nocca, noccb, nvirb), dtype=dtype)
Hr2bbb = np.zeros((nkpts, nkpts, noccb, noccb, nvirb), dtype=dtype)
# Fvv term
# \sum_{kd,d} U_{kb,kd,b,d} S_{ki,kj,i,j}^{kd,d} = (\bar{H}S)_{ki,kj,i,j}^{kb,b}
# \sum_{kD,D} S_{ki,kJ,i,J}^{kD,D} U_{kB,kD,B,D} = (\bar{H}S)_{ki,kJ,i,J}^{kB,B}
for kb, ki in itertools.product(range(nkpts),repeat=2):
kj = kconserv[kshift,ki,kb]
Hr2aaa[ki,kj] += lib.einsum('be,ije->ijb', imds.Fvv[kb], r2aaa[ki,kj])
Hr2abb[ki,kj] += lib.einsum('BE,iJE->iJB', imds.FVV[kb], r2abb[ki,kj])
Hr2bbb[ki,kj] += lib.einsum('BE,IJE->IJB', imds.FVV[kb], r2bbb[ki,kj])
Hr2baa[ki,kj] += lib.einsum('be,Ije->Ijb', imds.Fvv[kb], r2baa[ki,kj])
# Foo term
# \sum_{kl,l} U_{kl,ki,l,i} s_{kl,kj,l,j}^{kb,b} = (\bar{H}S)_{ki,kj,i,j}^{kb,b}
# \sum_{kl,l} U_{kl,kj,l,j} S_{ki,kl,i,l}^{kb,b} = (\bar{H}S)_{ki,kj,i,j}^{kb,b}
# \sum_{kl,l} S_{kl,kJ,l,J}^{kB,B} U_{kl,ki,l,i} = (\bar{H}S)_{ki,kJ,i,J}^{kB,B}
# \sum_{KL,L} S_{ki,kL,i,L}^{kB,B} U_{kL,kJ,L,J} = (\bar{H}S)_{ki,kJ,i,J}^{kB,B}
for ki, kj in itertools.product(range(nkpts), repeat=2):
tmpa = lib.einsum('mi,mjb->ijb', imds.Foo[ki], r2aaa[ki,kj])
tmpb = lib.einsum('mj,mib->ijb', imds.Foo[kj], r2aaa[kj,ki])
Hr2aaa[ki,kj] -= tmpa - tmpb
Hr2abb[ki,kj] -= lib.einsum('mi,mJB->iJB', imds.Foo[ki], r2abb[ki,kj])
Hr2abb[ki,kj] -= lib.einsum('MJ,iMB->iJB', imds.FOO[kj], r2abb[ki,kj])
Hr2baa[ki,kj] -= lib.einsum('MI,Mjb->Ijb', imds.FOO[ki], r2baa[ki,kj])
Hr2baa[ki,kj] -= lib.einsum('mj,Imb->Ijb', imds.Foo[kj], r2baa[ki,kj])
tmpb = lib.einsum('MI,MJB->IJB', imds.FOO[ki], r2bbb[ki,kj])
tmpa = lib.einsum('MJ,MIB->IJB', imds.FOO[kj], r2bbb[kj,ki])
Hr2bbb[ki,kj] -= tmpb - tmpa
# Wovoo term
# \sum_{kk,k} W_{kk,kb,kj,ki,k,b,j,i} s_{kk,k} = (\bar{H}S)_{ki,kj,i,j}^{kb,b}
# \sum_{kk,k} W_{kk,kB,ki,kJ,k,B,i,J} S_{kk,k} = (\bar{H}S)_{ki,kJ,i,J}^{kB,B}
for ki, kj in itertools.product(range(nkpts), repeat=2):
kb = kconserv[ki, kshift, kj]
Hr2aaa[ki,kj] -= einsum('mjbi,m->ijb', imds.Woovo[kshift,kj,kb], r1a)
Hr2abb[ki,kj] += einsum('miBJ,m->iJB', imds.WooVO[kshift,ki,kb], r1a)
Hr2baa[ki,kj] += einsum('MIbj,M->Ijb', imds.WOOvo[kshift,ki,kb], r1b)
Hr2bbb[ki,kj] -= einsum('MJBI,M->IJB', imds.WOOVO[kshift,kj,kb], r1b)
# Woooo term
# \sum_{kk,kl,k,l} W_{kk,ki,kl,kj,k,i,l,j} S_{kk,kl,k,l}^{kb,b} = (\bar{H}S)_{ki,kj,i,j}^{kb,b}
# \sum_{kk,kL,k,L} W_{kk,kL,ki,kJ,k,L,i,J} S_{kk,kl,k,L}^{kB,B} = (\bar{H}S)_{ki,kJ,i,J}^{kB,B}
for ki, kj in itertools.product(range(nkpts), repeat=2):
kb = kconserv[ki, kshift, kj]
for kn in range(nkpts):
km = kconserv[kj, kn, ki]
Hr2aaa[ki, kj] += .5 * lib.einsum('minj,mnb->ijb', imds.Woooo[km, ki, kn], r2aaa[km, kn])
Hr2abb[ki, kj] += lib.einsum('miNJ,mNB->iJB', imds.WooOO[km, ki, kn], r2abb[km, kn])
Hr2bbb[ki, kj] += .5 * lib.einsum('MINJ,MNB->IJB', imds.WOOOO[km, ki, kn], r2bbb[km, kn])
Hr2baa[ki, kj] += lib.einsum('njMI,Mnb->Ijb', imds.WooOO[kn, kj, km], r2baa[km, kn])
# T2 term
# - \sum_{kc,c} t_{kj,ki,j,i}^{kb,kc,b,c} [ \sum_{kk,kL,kD,k,L,D} W_{kL,kk,kD,kc,L,k,D,c} S_{kk,kL,k,L}^{kD,D}
# + \sum{kk,kl,kd,k,l,d} W_{kl,kk,kd,kc,l,k,d,c} S_{kk,kl,k,l}^{kd,d} ] = (\bar{H}S)_{ki,kj,i,j}^{kb,b}
#
# - \sum_{kc,c} t_{ki,kJ,i,J}^{kc,kB,c,B} [ \sum_{kk,kL,kD,k,L,D} W_{kL,kk,kD,kc,L,k,D,c} S_{Kk,kL,k,L}^{kD,D}
# + \sum{kk,kl,kd,k,l,d} W_{kl,kk,kd,kc,l,k,d,c} S_{kk,kl,k,l}^{kd,d} ] = (\bar{H}S)_{ki,kJ,i,J}^{kB,B}
tmp_aaa = lib.einsum('xymenf,xymnf->e', imds.Wovov[:,kshift,:], r2aaa)
tmp_bbb = lib.einsum('xyMENF,xyMNF->E', imds.WOVOV[:,kshift,:], r2bbb)
tmp_abb = lib.einsum('xymeNF,xymNF->e', imds.WovOV[:,kshift,:], r2abb)
tmp_baa = np.zeros(tmp_bbb.shape, dtype=tmp_bbb.dtype)
for km, kn in itertools.product(range(nkpts), repeat=2):
kf = kconserv[kn, kshift, km]
tmp_baa += lib.einsum('nfME, Mnf->E', imds.WovOV[kn, kf, km], r2baa[km, kn])
for ki, kj in itertools.product(range(nkpts), repeat=2):
kb = kconserv[ki, kshift, kj]
Hr2aaa[ki,kj] -= 0.5 * lib.einsum('e,jibe->ijb', tmp_aaa, t2aa[kj,ki,kb])
Hr2aaa[ki,kj] -= lib.einsum('e,jibe->ijb', tmp_abb, t2aa[kj,ki,kb])
Hr2abb[ki,kj] -= 0.5 * lib.einsum('e,iJeB->iJB', tmp_aaa, t2ab[ki,kj,kshift])
Hr2abb[ki,kj] -= lib.einsum('e,iJeB->iJB', tmp_abb, t2ab[ki,kj,kshift])
Hr2baa[ki,kj] -= 0.5 * lib.einsum('E,jIbE->Ijb', tmp_bbb, t2ab[kj,ki,kb])
Hr2baa[ki,kj] -= lib.einsum('E,jIbE->Ijb', tmp_baa, t2ab[kj,ki,kb])
Hr2bbb[ki,kj] -= 0.5 * lib.einsum('E,JIBE->IJB', tmp_bbb, t2bb[kj,ki,kb])
Hr2bbb[ki,kj] -= lib.einsum('E,JIBE->IJB', tmp_baa, t2bb[kj,ki,kb])
#idxoa = [np.where(orbspin[k][:nocca+noccb] == 0)[0] for k in range(nkpts)]
#idxva = [np.where(orbspin[k][nocca+noccb:] == 0)[0] for k in range(nkpts)]
#idxob = [np.where(orbspin[k][:nocca+noccb] == 1)[0] for k in range(nkpts)]
#idxvb = [np.where(orbspin[k][nocca+noccb:] == 1)[0] for k in range(nkpts)]
# j \/ b | i
# --- |
# /\ |
# m \/ e|
# -------
for ki, kj in itertools.product(range(nkpts), repeat=2):
kb = kconserv[ki, kshift, kj]
for km in range(nkpts):
ke = kconserv[km, kshift, ki]
# \sum_{kL,kD,L,D} W_{kL,kD,kb,kj,L,D,b,j} S_{ki,kL,i,L}^{kb,b}
# \sum_{kl,kd,l,d} W_{kl,kd,kb,kj,l,d,b,j} S_{ki,kl,i,l}^{kb,b}
Hr2aaa[ki, kj] += lib.einsum('mebj,ime->ijb', imds.Wovvo[km, ke, kb],
r2aaa[ki, km])
Hr2aaa[ki, kj] += lib.einsum('MEbj,iME->ijb', imds.WOVvo[km, ke, kb],
r2abb[ki, km])
# P(ij)
ke = kconserv[km, kshift, kj]
Hr2aaa[ki, kj] -= lib.einsum('mebi,jme->ijb', imds.Wovvo[km, ke, kb],
r2aaa[kj, km])
Hr2aaa[ki, kj] -= lib.einsum('MEbi,jME->ijb', imds.WOVvo[km, ke, kb],
r2abb[kj, km])
# \sum_{kL,kD,L,D} W_{kL,kD,kb,kJ,L,D,b,J} S_{ki,kL,i,L}^{kD,D}
# \sum_{kl,kd,l,d} W_{kl,kd,kB,kJ,l,d,B,J} S_{ki,kl,i,l}^{kd,d}
ke = kconserv[km, kshift, ki]
Hr2abb[ki, kj] += lib.einsum('meBJ,ime->iJB', imds.WovVO[km, ke, kb],
r2aaa[ki, km])
Hr2abb[ki, kj] += lib.einsum('MEBJ,iME->iJB', imds.WOVVO[km, ke, kb],
r2abb[ki, km])
ke = kconserv[km, kshift, kj]
Hr2abb[ki, kj] -= lib.einsum('miBE,mJE->iJB', imds.WooVV[km, ki, kb],
r2abb[km, kj])
ke = kconserv[km, kshift, ki]
Hr2baa[ki, kj] += lib.einsum('MEbj,IME->Ijb', imds.WOVvo[km, ke, kb],
r2bbb[ki, km])
Hr2baa[ki, kj] += lib.einsum('mebj,Ime->Ijb', imds.Wovvo[km, ke, kb],
r2baa[ki, km])
ke = kconserv[km, kshift, kj]
Hr2baa[ki, kj] -= lib.einsum('MIbe,Mje->Ijb', imds.WOOvv[km, ki, kb],
r2baa[km, kj])
ke = kconserv[km, kshift, ki]
Hr2bbb[ki, kj] += lib.einsum('MEBJ,IME->IJB', imds.WOVVO[km, ke, kb],
r2bbb[ki, km])
Hr2bbb[ki, kj] += lib.einsum('meBJ,Ime->IJB', imds.WovVO[km, ke, kb],
r2baa[ki, km])
# P(ij)
ke = kconserv[km, kshift, kj]
Hr2bbb[ki, kj] -= lib.einsum('MEBI,JME->IJB', imds.WOVVO[km, ke, kb],
r2bbb[kj, km])
Hr2bbb[ki, kj] -= lib.einsum('meBI,Jme->IJB', imds.WovVO[km, ke, kb],
r2baa[kj, km])
#spatial_Hr1 = [Hr1a, Hr1b]
#spatial_Hr2 = [Hr2aaa, Hr2baa, Hr2abb, Hr2bbb]
#spin_Hr1, spin_Hr2 = eom_kgccsd.spatial2spin_ip_doublet(spatial_Hr1, spatial_Hr2,
# kconserv, kshift, orbspin)
#Hr1 += spin_Hr1
#Hr2 += spin_Hr2
#vector = eom.amplitudes_to_vector(Hr1, Hr2)
vector = amplitudes_to_vector_ip([Hr1a, Hr1b], [Hr2aaa, Hr2baa, Hr2abb, Hr2bbb], kshift, kconserv)
return vector
[docs]
def ipccsd_diag(eom, kshift, imds=None):
if imds is None: imds = eom.make_imds()
t1, t2 = imds.t1, imds.t2
t1a, t1b = t1
t2aa, t2ab, t2bb = t2
nkpts, nocc_a, nvir_a = t1a.shape
nkpts, nocc_b, nvir_b = t1b.shape
kconserv = imds.kconserv
Hr1a = -np.diag(imds.Foo[kshift])
Hr1b = -np.diag(imds.FOO[kshift])
Hr2aaa = np.zeros((nkpts,nkpts,nocc_a,nocc_a,nvir_a), dtype=t1[0].dtype)
Hr2bbb = np.zeros((nkpts,nkpts,nocc_b,nocc_b,nvir_b), dtype=t1[0].dtype)
Hr2abb = np.zeros((nkpts,nkpts,nocc_a,nocc_b,nvir_b), dtype=t1[0].dtype)
Hr2baa = np.zeros((nkpts,nkpts,nocc_b,nocc_a,nvir_a), dtype=t1[0].dtype)
if eom.partition == 'mp':
raise Exception("MP diag is not tested") # remove this to use untested code
#foo = eris.fock[0][:,:nocc_a,:nocc_a]
#fOO = eris.fock[1][:,:nocc_b,:nocc_b]
#fvv = eris.fock[0][:,:nvir_a,:nvir_a]
#fVV = eris.fock[1][:,:nvir_b,:nvir_b]
for ki in range(nkpts):
for kj in range(nkpts):
ka = kconserv[ki,kshift,kj]
Hr2aaa[ki,kj] = imds.Fvv[ka].diagonal()
Hr2aaa[ki,kj] -= imds.Foo[ki].diagonal()[:,None,None]
Hr2aaa[ki,kj] -= imds.Foo[kj].diagonal()[None,:,None]
Hr2bbb[ki,kj] = imds.FVV[ka].diagonal()
Hr2bbb[ki,kj] -= imds.FOO[ki].diagonal()[:,None,None]
Hr2bbb[ki,kj] -= imds.FOO[kj].diagonal()[None,:,None]
Hr2aba[ki,kj] = imds.Fvv[ka].diagonal()
Hr2aba[ki,kj] -= imds.Foo[ki].diagonal()[:,None,None]
Hr2aba[ki,kj] -= imds.FOO[kj].diagonal()[None,:,None]
Hr2bab[ki,kj] = imds.FVV[ka].diagonal()
Hr2bab[ki,kj] -= imds.FOO[ki].diagonal()[:,None,None]
Hr2bab[ki,kj] -= imds.Foo[kj].diagonal()[None,:,None]
else:
for ka in range(nkpts):
for ki in range(nkpts):
kj = kconserv[kshift,ki,ka]
Hr2aaa[ki,kj] += imds.Fvv[ka].diagonal()
Hr2abb[ki,kj] += imds.FVV[ka].diagonal()
Hr2bbb[ki,kj] += imds.FVV[ka].diagonal()
Hr2baa[ki,kj] += imds.Fvv[ka].diagonal()
Hr2aaa[ki,kj] -= imds.Foo[ki].diagonal()[:,None,None]
Hr2aaa[ki,kj] -= imds.Foo[kj].diagonal()[None,:,None]
Hr2abb[ki,kj] -= imds.Foo[ki].diagonal()[:,None,None]
Hr2abb[ki,kj] -= imds.FOO[kj].diagonal()[None,:,None]
Hr2baa[ki,kj] -= imds.FOO[ki].diagonal()[:,None,None]
Hr2baa[ki,kj] -= imds.Foo[kj].diagonal()[None,:,None]
Hr2bbb[ki,kj] -= imds.FOO[ki].diagonal()[:,None,None]
Hr2bbb[ki,kj] -= imds.FOO[kj].diagonal()[None,:,None]
for ki, kj in itertools.product(range(nkpts), repeat=2):
Hr2aaa[ki, kj] += lib.einsum('iijj->ij', imds.Woooo[ki, ki, kj])[:,:,None]
Hr2abb[ki, kj] += lib.einsum('iiJJ->iJ', imds.WooOO[ki, ki, kj])[:,:,None]
Hr2bbb[ki, kj] += lib.einsum('IIJJ->IJ', imds.WOOOO[ki, ki, kj])[:,:,None]
Hr2baa[ki, kj] += lib.einsum('jjII->Ij', imds.WooOO[kj, kj, ki])[:,:,None]
kb = kconserv[ki, kshift, kj]
Hr2aaa[ki,kj] -= lib.einsum('iejb,jibe->ijb', imds.Wovov[ki,kshift,kj], t2aa[kj,ki,kb])
Hr2abb[ki,kj] -= lib.einsum('ieJB,iJeB->iJB', imds.WovOV[ki,kshift,kj], t2ab[ki,kj,kshift])
Hr2baa[ki,kj] -= lib.einsum('jbIE,jIbE->Ijb', imds.WovOV[kj,kb,ki], t2ab[kj,ki,kb])
Hr2bbb[ki,kj] -= lib.einsum('IEJB,JIBE->IJB', imds.WOVOV[ki,kshift,kj], t2bb[kj,ki,kb])
Hr2aaa[ki, kj] += lib.einsum('ibbi->ib', imds.Wovvo[ki, kb, kb])[:,None,:]
Hr2aaa[ki, kj] += lib.einsum('jbbj->jb', imds.Wovvo[kj, kb, kb])[None,:,:]
Hr2baa[ki, kj] += lib.einsum('jbbj->jb', imds.Wovvo[kj, kb, kb])[None,:,:]
Hr2baa[ki, kj] -= lib.einsum('IIbb->Ib', imds.WOOvv[ki, ki, kb])[:,None,:]
Hr2abb[ki, kj] += lib.einsum('JBBJ->JB', imds.WOVVO[kj, kb, kb])[None,:,:]
Hr2abb[ki, kj] -= lib.einsum('iiBB->iB', imds.WooVV[ki, ki, kb])[:,None,:]
Hr2bbb[ki, kj] += lib.einsum('IBBI->IB', imds.WOVVO[ki, kb, kb])[:,None,:]
Hr2bbb[ki, kj] += lib.einsum('JBBJ->JB', imds.WOVVO[kj, kb, kb])[None,:,:]
vector = amplitudes_to_vector_ip((Hr1a,Hr1b), (Hr2aaa,Hr2baa,Hr2abb,Hr2bbb), kshift, kconserv)
return vector
[docs]
def mask_frozen_ip(eom, vector, kshift, const=LARGE_DENOM):
'''Replaces all frozen orbital indices of `vector` with the value `const`.'''
nkpts = eom.nkpts
nocca, noccb = eom.nocc
nmoa, nmob = eom.nmo
kconserv = eom.kconserv
r1, r2 = eom.vector_to_amplitudes(vector, kshift, nkpts, (nmoa, nmob), (nocca, noccb), kconserv)
r1a, r1b = r1
r2aaa, r2baa, r2abb, r2bbb = r2
# Get location of padded elements in occupied and virtual space
nonzero_opadding, nonzero_vpadding = eom.nonzero_opadding, eom.nonzero_vpadding
nonzero_opadding_a, nonzero_opadding_b = nonzero_opadding
nonzero_vpadding_a, nonzero_vpadding_b = nonzero_vpadding
new_r1a = const * np.ones_like(r1a)
new_r1b = const * np.ones_like(r1b)
new_r2aaa = const * np.ones_like(r2aaa)
new_r2baa = const * np.ones_like(r2baa)
new_r2abb = const * np.ones_like(r2abb)
new_r2bbb = const * np.ones_like(r2bbb)
# r1a/b case
new_r1a[nonzero_opadding_a[kshift]] = r1a[nonzero_opadding_a[kshift]]
new_r1b[nonzero_opadding_b[kshift]] = r1b[nonzero_opadding_b[kshift]]
# r2aaa case
for ki in range(nkpts):
for kj in range(nkpts):
kb = kconserv[ki, kshift, kj]
idx = np.ix_([ki], [kj], nonzero_opadding_a[ki], nonzero_opadding_a[kj], nonzero_vpadding_a[kb])
new_r2aaa[idx] = r2aaa[idx]
# r2baa case
for ki in range(nkpts):
for kj in range(nkpts):
kb = kconserv[ki, kshift, kj]
idx = np.ix_([ki], [kj], nonzero_opadding_b[ki], nonzero_opadding_a[kj], nonzero_vpadding_a[kb])
new_r2baa[idx] = r2baa[idx]
# r2abb case
for ki in range(nkpts):
for kj in range(nkpts):
kb = kconserv[ki, kshift, kj]
idx = np.ix_([ki], [kj], nonzero_opadding_a[ki], nonzero_opadding_b[kj], nonzero_vpadding_b[kb])
new_r2abb[idx] = r2abb[idx]
# r2bbb case
for ki in range(nkpts):
for kj in range(nkpts):
kb = kconserv[ki, kshift, kj]
idx = np.ix_([ki], [kj], nonzero_opadding_b[ki], nonzero_opadding_b[kj], nonzero_vpadding_b[kb])
new_r2bbb[idx] = r2bbb[idx]
return eom.amplitudes_to_vector((new_r1a,new_r1b), (new_r2aaa,new_r2baa,new_r2abb,new_r2bbb), kshift, kconserv)
[docs]
def get_padding_k_idx(eom, cc):
# Get location of padded elements in occupied and virtual space
nonzero_padding_alpha, nonzero_padding_beta = padding_k_idx(cc, kind="split")
nonzero_opadding_alpha, nonzero_vpadding_alpha = nonzero_padding_alpha
nonzero_opadding_beta, nonzero_vpadding_beta = nonzero_padding_beta
return ((nonzero_opadding_alpha, nonzero_opadding_beta),
(nonzero_vpadding_alpha, nonzero_vpadding_beta))
[docs]
class EOMIP(eom_kgccsd.EOMIP):
def __init__(self, cc):
#if not isinstance(cc, kccsd.GCCSD):
# raise TypeError
self.kpts = cc.kpts
eom_kgccsd.EOMIP.__init__(self, cc)
get_diag = ipccsd_diag
matvec = ipccsd_matvec
get_padding_k_idx = get_padding_k_idx
mask_frozen = mask_frozen_ip
[docs]
def get_init_guess(self, kshift, nroots=1, koopmans=False, diag=None):
size = self.vector_size()
dtype = getattr(diag, 'dtype', np.complex128)
nroots = min(nroots, size)
nocca, noccb = self.nocc
guess = []
if koopmans:
idx = np.zeros(nroots, dtype=int)
tmp_oalpha, tmp_obeta = self.nonzero_opadding[kshift]
tmp_oalpha = list(tmp_oalpha)
tmp_obeta = list(tmp_obeta)
if len(tmp_obeta) + len(tmp_oalpha) < nroots:
raise ValueError("Max number of roots for k-point (idx=%3d) for koopmans "
"is %3d.\nRequested %3d." %
(kshift, len(tmp_obeta)+len(tmp_oalpha), nroots))
total_count = 0
while (total_count < nroots):
if total_count % 2 == 0 and len(tmp_oalpha) > 0:
idx[total_count] = tmp_oalpha.pop()
else:
# Careful! index depends on how we create vector
# (here the first elements are r1a, then r1b)
idx[total_count] = nocca + tmp_obeta.pop()
total_count += 1
else:
idx = diag.argsort()
for i in idx[:nroots]:
g = np.zeros(size, dtype)
g[i] = 1.0
g = self.mask_frozen(g, kshift, const=0.0)
guess.append(g)
return guess
[docs]
def gen_matvec(self, kshift, imds=None, left=False, **kwargs):
if imds is None: imds = self.make_imds()
diag = self.get_diag(kshift, imds)
if left:
raise NotImplementedError
matvec = lambda xs: [self.l_matvec(x, kshift, imds, diag) for x in xs]
else:
matvec = lambda xs: [self.matvec(x, kshift, imds, diag) for x in xs]
return matvec, diag
[docs]
def vector_to_amplitudes(self, vector, kshift, nkpts=None, nmo=None, nocc=None, kconserv=None):
if nmo is None: nmo = self.nmo
if nocc is None: nocc = self.nocc
if nkpts is None: nkpts = self.nkpts
if kconserv is None: kconserv = self.kconserv
return vector_to_amplitudes_ip(vector, kshift, nkpts, nmo, nocc, kconserv)
[docs]
def amplitudes_to_vector(self, r1, r2, kshift, kconserv=None):
if kconserv is None: kconserv = self.kconserv
return amplitudes_to_vector_ip(r1, r2, kshift, kconserv)
[docs]
def vector_size(self):
nocca, noccb = self.nocc
nmoa, nmob = self.nmo
nvira, nvirb = nmoa - nocca, nmob - noccb
nkpts = self.nkpts
return (nocca + noccb +
nkpts*nocca*(nkpts*nocca-1)*nvira//2 +
nkpts**2*noccb*nocca*nvira +
nkpts**2*nocca*noccb*nvirb +
nkpts*noccb*(nkpts*noccb-1)*nvirb//2)
[docs]
def make_imds(self, eris=None, t1=None, t2=None):
imds = _IMDS(self._cc, eris, t1, t2)
imds.make_ip()
return imds
########################################
# EOM-EA-CCSD
########################################
[docs]
def amplitudes_to_vector_ea(r1, r2, kshift, kconserv):
r1a, r1b = r1
r2a, r2aba, r2bab, r2b = r2
nkpts = r2a.shape[0]
nocca, noccb = r2a.shape[2], r2b.shape[2]
nvira, nvirb = r2a.shape[3], r2b.shape[3]
# From symmetry for aaa and bbb terms, only store lower
# triangular part (ka,a) < (kb,b)
r2aaa = np.zeros((nocca*nkpts*nvira*(nkpts*nvira-1))//2, dtype=r2a.dtype)
r2bbb = np.zeros((noccb*nkpts*nvirb*(nkpts*nvirb-1))//2, dtype=r2b.dtype)
index = 0
for kj, ka in itertools.product(range(nkpts), repeat=2):
kb = kconserv[kshift,ka,kj]
if ka < kb: # Take diagonal part
idxa, idya = np.tril_indices(nvira, 0)
else: # Don't take diagonal (equal to zero)
idxa, idya = np.tril_indices(nvira, -1)
r2aaa[index:index + nocca*len(idya)] = r2a[kj,ka,:,idxa,idya].reshape(-1)
index = index + nocca*len(idya)
index = 0
for kj, ka in itertools.product(range(nkpts), repeat=2):
kb = kconserv[kshift,ka,kj]
if ka < kb: # Take diagonal part
idxb, idyb = np.tril_indices(nvirb, 0)
else:
idxb, idyb = np.tril_indices(nvirb, -1)
r2bbb[index:index + noccb*len(idyb)] = r2b[kj,ka,:,idxb,idyb].reshape(-1)
index = index + noccb*len(idyb)
return np.hstack((r1a, r1b, r2aaa.ravel(),
r2aba.ravel(), r2bab.ravel(),
r2bbb.ravel()))
[docs]
def vector_to_amplitudes_ea(vector, kshift, nkpts, nmo, nocc, kconserv):
nocca, noccb = nocc
nmoa, nmob = nmo
nvira, nvirb = nmoa-nocca, nmob-noccb
sizes = (nvira, nvirb, nkpts*nocca*(nkpts*nvira-1)*nvira//2,
nkpts**2*nocca*nvirb*nvira, nkpts**2*noccb*nvira*nvirb,
nkpts*noccb*(nkpts*nvirb-1)*nvirb//2)
sections = np.cumsum(sizes[:-1])
r1a, r1b, r2a, r2aba, r2bab, r2b = np.split(vector, sections)
r2aaa = np.zeros((nkpts,nkpts,nocca,nvira,nvira), dtype=r2a.dtype)
r2aba = r2aba.reshape(nkpts,nkpts,nocca,nvirb,nvira).copy()
r2bab = r2bab.reshape(nkpts,nkpts,noccb,nvira,nvirb).copy()
r2bbb = np.zeros((nkpts,nkpts,noccb,nvirb,nvirb), dtype=r2b.dtype)
index = 0
for kj, ka in itertools.product(range(nkpts), repeat=2):
kb = kconserv[kshift,ka,kj]
if ka < kb: # Take diagonal part
idxa, idya = np.tril_indices(nvira, 0)
else:
idxa, idya = np.tril_indices(nvira, -1)
tmp = r2a[index:index + nocca*len(idya)].reshape(-1,nocca)
r2aaa[kj,ka,:,idxa,idya] = tmp
r2aaa[kj,kb,:,idya,idxa] = -tmp
index = index + nocca*len(idya)
index = 0
for kj, ka in itertools.product(range(nkpts), repeat=2):
kb = kconserv[kshift,ka,kj]
if ka < kb: # Take diagonal part
idxb, idyb = np.tril_indices(nvirb, 0)
else:
idxb, idyb = np.tril_indices(nvirb, -1)
tmp = r2b[index:index + noccb*len(idyb)].reshape(-1,noccb)
r2bbb[kj,ka,:,idxb,idyb] = tmp
r2bbb[kj,kb,:,idyb,idxb] = -tmp
index = index + noccb*len(idyb)
r1 = (r1a.copy(), r1b.copy())
r2 = (r2aaa, r2aba, r2bab, r2bbb)
return r1, r2
[docs]
def eaccsd_matvec(eom, vector, kshift, imds=None, diag=None):
'''2ph operators are of the form s_{ j}^{ab}, i.e. 'jb' indices are coupled'''
if imds is None: imds = eom.make_imds()
t1, t2= imds.t1, imds.t2
t1a, t1b = t1
t2aa, t2ab, t2bb = t2
nocca, noccb, nvira, nvirb = t2ab.shape[3:]
nmoa, nmob = nocca + nvira, noccb + nvirb
kconserv = imds.kconserv
nkpts = eom.nkpts
r1, r2 = eom.vector_to_amplitudes(vector, kshift, nkpts, (nmoa, nmob), (nocca, noccb), kconserv)
r1a, r1b = r1
r2aaa, r2aba, r2bab, r2bbb = r2
# BEGINNING OF MATVEC CONTRACTIONS: ref - Nooijen 1995 EOM-CC for EA
# Fvv terms
# (\bar{H}S)^a = \sum_{kc,c} U_{ac} s^c
Hr1a = einsum('ac,c->a', imds.Fvv[kshift], r1a)
Hr1b = einsum('AC,C->A', imds.FVV[kshift], r1b)
# Fov terms
# (\bar{H}S)^a = \sum_{kL,kD, L, D} U_{kL,kD,L,D} s^{a,kD,D}_{kL,L} + \sum_{kl,kd,l,d} U_{kl, d}^{a,kd,d}_{kl,l}
for kl in range(nkpts):
Hr1a += einsum('ld,lad->a', imds.Fov[kl], r2aaa[kl,kshift])
Hr1a += einsum('LD,LaD->a', imds.FOV[kl], r2bab[kl,kshift])
Hr1b += einsum('ld,lAd->A', imds.Fov[kl], r2aba[kl,kshift])
Hr1b += einsum('LD,LAD->A', imds.FOV[kl], r2bbb[kl,kshift])
# Wvovv
# (\bar{H}S)^a = \sum_{kc,kL,kD,c,L,D} W_{kL,kc,kD,a,l,c,D} s_{kL,L}^{kc,kD,c,D}
# + \sum_{kc,kd,kl,c,d,l} W_{ka,kl,kc,kd,a,l,c,d} s_{kl,l}^{kc,kd,c,d}
for kc, kl in itertools.product(range(nkpts), repeat=2):
Hr1a += 0.5*lib.einsum('acld,lcd->a', imds.Wvvov[kshift,kc,kl], r2aaa[kl,kc])
Hr1a += lib.einsum('acLD,LcD->a', imds.WvvOV[kshift,kc,kl], r2bab[kl,kc])
Hr1b += 0.5*lib.einsum('ACLD,LCD->A', imds.WVVOV[kshift,kc,kl], r2bbb[kl,kc])
Hr1b += lib.einsum('ACld,lCd->A', imds.WVVov[kshift,kc,kl], r2aba[kl,kc])
dtype = np.result_type(Hr1a, *r2)
Hr2aaa = np.zeros((nkpts, nkpts, nocca, nvira, nvira), dtype=dtype)
Hr2aba = np.zeros((nkpts, nkpts, nocca, nvirb, nvira), dtype=dtype)
Hr2bab = np.zeros((nkpts, nkpts, noccb, nvira, nvirb), dtype=dtype)
Hr2bbb = np.zeros((nkpts, nkpts, noccb, nvirb, nvirb), dtype=dtype)
# Wvvvv
# \sum_{kc,kd,c,d} W_{ka,kb,kc,kd,a,b,c,d} s_{kj,j}^{kc,kd,c,d} = (\bar{H}S)^{kb, a, b}_{kj,j}
# \sum_{kc,kD,c,D} W{ka,kB,kc,kD,a,B,c,D} s_{kJ,kc,kD,J,c,D} = (\bar{H}S)^{kB, a, B}_{kJ,J}
for kj, ka in itertools.product(range(nkpts), repeat=2):
kb = kconserv[kshift,ka,kj]
for kc in range(nkpts):
kd = kconserv[ka, kc, kb]
Wvvvv, WvvVV, WVVVV = imds.get_Wvvvv(ka, kb, kc)
Hr2aaa[kj,ka] += .5 * lib.einsum('acbd,jcd->jab', Wvvvv, r2aaa[kj,kc])
Hr2aba[kj,kb] += lib.einsum('bcad,jdc->jab', WvvVV, r2aba[kj,kd])
Hr2bab[kj,ka] += lib.einsum('acbd,jcd->jab', WvvVV, r2bab[kj,kc])
Hr2bbb[kj,ka] += .5 * lib.einsum('acbd,jcd->jab', WVVVV, r2bbb[kj,kc])
#Wvvvo
# \sum_{kc,ka,kj,c,a,j} W_{kb,kc,kj,a,b,c,j} s^{kc,c} = (\bar{H}S)^{kb, a, b}_{kj,j}
# \sum_{kc,ka,kJ,c,a,J} W_{kB,kc,kJ,a,B,c,J} s^{kc,c} = (\bar{H}S)^{kB, a, B}_{kJ,J}
for ka, kj, in itertools.product(range(nkpts),repeat=2):
kb = kconserv[kshift,ka,kj]
kc = kshift
Hr2aaa[kj,ka] += einsum('acbj,c->jab', imds.Wvvvo[ka,kc,kb], r1a)
Hr2bbb[kj,ka] += einsum('ACBJ,C->JAB', imds.WVVVO[ka,kc,kb], r1b)
Hr2bab[kj,ka] += einsum('acBJ,c->JaB', imds.WvvVO[ka,kc,kb], r1a)
Hr2aba[kj,ka] += einsum('ACbj,C->jAb', imds.WVVvo[ka,kc,kb], r1b)
#Fvv Terms
# sum_{kc,ka,kj,c,a,j} s_{kj,j}^{kc,kb,c,b} U_{ka,kc,a,c} = (\bar{H}S)^{kb, a, b}_{kj,j}
# sum_{kd,ka,kj,d,b,j} s_{kj,j}^{ka,kd,a,d} U_{kb,kd,b,d} = (\bar{H}S)^{kb, a, b}_{kj,j}
# sum_{kc,ka,kJ,c,a,J} U_{ka,kc,a,c} s_{kJ,J}^{kc,kB,c,B} = (\bar{H}S)^{kB, a, B}_{kJ,J}
# sum_{kD,ka,kj,D,a,j} U_{kb,kd,b,d} s_{kj,j}^{ka,kd,a,d} = (\bar{H}S)^{kB, a, B}_{kJ,J}
for ka, kj in itertools.product(range(nkpts), repeat=2):
# kb = kshift - ka + kj
kb = kconserv[kshift, ka, kj]
tmpa = lib.einsum('ac,jcb->jab', imds.Fvv[ka], r2aaa[kj,ka])
tmpb = lib.einsum('bc,jca->jab', imds.Fvv[kb], r2aaa[kj,kb])
Hr2aaa[kj,ka] += tmpa - tmpb
Hr2aba[kj,ka] += lib.einsum('AC,jCb->jAb', imds.FVV[ka], r2aba[kj,ka])
Hr2bab[kj,ka] += lib.einsum('ac,JcB->JaB', imds.Fvv[ka], r2bab[kj,ka])
Hr2aba[kj,ka] += lib.einsum('bc, jAc -> jAb', imds.Fvv[kb], r2aba[kj,ka])
Hr2bab[kj,ka] += lib.einsum('BC, JaC -> JaB', imds.FVV[kb], r2bab[kj,ka])
tmpb = lib.einsum('AC,JCB->JAB', imds.FVV[ka], r2bbb[kj,ka])
tmpa = lib.einsum('BC,JCA->JAB', imds.FVV[kb], r2bbb[kj,kb])
Hr2bbb[kj,ka] += tmpb - tmpa
#Foo Term
# \sum_{ka,kl,l} U_{kl,l,kj,j} s^{ka,a,kb,b}^{kl,l} = (\bar{H}S)^{kb, a, b}_{kj,j}
# \sum_{ka,kL,L} s^{ka,a,kB,B}_{kL,L} U_{kL,L,kJ,J} = (\bar{H}S)^{kB, a, B}_{kJ,J}
for kl, ka in itertools.product(range(nkpts), repeat=2):
Hr2aaa[kl,ka] -= lib.einsum('lj,lab->jab', imds.Foo[kl], r2aaa[kl,ka])
Hr2bbb[kl,ka] -= lib.einsum('LJ,LAB->JAB', imds.FOO[kl], r2bbb[kl,ka])
Hr2bab[kl,ka] -= lib.einsum('LJ,LaB->JaB', imds.FOO[kl], r2bab[kl,ka])
Hr2aba[kl,ka] -= lib.einsum('lj,lAb->jAb', imds.Foo[kl], r2aba[kl,ka])
# Woovv term
# - \sum{kk,k} t_{kk,kj,k,j}^{ka,kb,a,b} [\sum_{kc,kD,kL,c,D,L} W_{kL,kk,kD,kc,L,k,D,c} s_{kL,L}^{kc,kD,c,D}
# + \sum{kc,kd,kl,c,d,l} W_{kk,kl,kc,kd,k,l,c,d} s_{kl,l}^{kc,kd,c,d} ] = (\bar{H}S)^{kb, a, b}_{kj,j}
#
# - \sum_{kk,k} t_{kk,kJ,k,J}^{ka,kB,a,B} [ \sum{kc,kD,kL,c,D,L} W_{kk,kL,kc,kD,k,L,c,D} s_{kL,L}^{kc,kD,c,D}
# + \sum_{kc,kd,kl,c,d,l} W_{kk,kl,kc,kd,k,l,c,d} s_{kl,l}^{kc,kd,c,d} ] = (\bar{H}S)^{kB, a, B}_{kJ,J}
tmp_aaa = lib.einsum('xykcld, yxlcd->k', imds.Wovov[kshift,:,:], r2aaa)
tmp_bbb = lib.einsum('xyKCLD, yxLCD->K', imds.WOVOV[kshift,:,:], r2bbb)
tmp_bab = lib.einsum('xykcLD, yxLcD->k', imds.WovOV[kshift], r2bab)
tmp_aba = np.zeros(tmp_bbb.shape, dtype = tmp_bbb.dtype)
for kl, kc in itertools.product(range(nkpts), repeat=2):
kd = kconserv[kl,kc,kshift]
tmp_aba += lib.einsum('ldKC, lCd->K', imds.WovOV[kl,kd,kshift], r2aba[kl,kc])
Hr2aaa -= 0.5 * lib.einsum('k, xykjab->xyjab', tmp_aaa, t2aa[kshift])
Hr2bab -= 0.5 * lib.einsum('k, xykJaB->xyJaB', tmp_aaa, t2ab[kshift])
Hr2aaa -= lib.einsum('k, xykjab->xyjab', tmp_bab, t2aa[kshift])
Hr2bbb -= 0.5 * lib.einsum('K, xyKJAB->xyJAB', tmp_bbb, t2bb[kshift])
Hr2bbb -= lib.einsum('K, xyKJAB->xyJAB', tmp_aba, t2bb[kshift])
Hr2bab -= lib.einsum('k, xykJaB->xyJaB', tmp_bab, t2ab[kshift])
for kj, ka in itertools.product(range(nkpts), repeat=2):
kb = kconserv[kshift, ka, kj]
Hr2aba[kj, ka] -= lib.einsum('K, jKbA->jAb', tmp_aba, t2ab[kj, kshift, kb])
Hr2aba[kj, ka] -= 0.5 * einsum('K, jKbA->jAb', tmp_bbb, t2ab[kj, kshift, kb])
# j \/ b | a
# --- |
# /\ |
# l \/ d|
# -------
for kj, ka in itertools.product(range(nkpts), repeat=2):
kb = kconserv[kshift, ka, kj]
for kd in range(nkpts):
kl = kconserv[ka, kshift, kd]
# \sum_{kL,kD,L,D} W_{kL,kb,kD,kj,L,b,D,j} s_{kL,L}^{ka,kD,a,D} = (\bar{H}S)^{kb, a, b}_{kj,j}
# \sum_{kl,kd,l,d} W_{kl,kb,kd,kj,l,b,d,j} s_{kl,l}^{ka,kd,a,d} = (\bar{H}S)^{kb, a, b}_{kj,j}
Hr2aaa[kj, ka] += lib.einsum('ldbj,lad->jab', imds.Wovvo[kl,kd,kb],
r2aaa[kl,ka])
Hr2aaa[kj, ka] += lib.einsum('LDbj,LaD->jab', imds.WOVvo[kl,kd,kb],
r2bab[kl,ka])
# P(ab)
kl = kconserv[kb, kshift, kd]
Hr2aaa[kj, ka] -= lib.einsum('ldaj,lbd->jab', imds.Wovvo[kl,kd,ka],
r2aaa[kl,kb])
Hr2aaa[kj, ka] -= lib.einsum('LDaj,LbD->jab', imds.WOVvo[kl,kd,ka],
r2bab[kl,kb])
kl = kconserv[ka, kshift, kd]
# \sum_{kL,kD,L,D} W_{kL,kB,kD,kJ,L,B,D,J} s_{kL,L}^{ka,kD,a,D} = (\bar{H}S)^{kB, a, B}_{kJ,J}
# \sum_{kl,kd,l,d} W_{kl,kB,kd,kJ,l,B,d,J} s_{kl,l}^{ka,kd,a,d} = (\bar{H}S)^{kB, a, B}_{kJ,J}
# - \sum_{kc,kL,c,L} W_{ka,kL,kc,kJ,a,L,c,J} s_{kL,L}^{kc,kB,c,B} = (\bar{H}S)^{kB, a, B}_{kJ,J}
Hr2bab[kj, ka] += lib.einsum('ldBJ,lad->JaB', imds.WovVO[kl,kd,kb],
r2aaa[kl,ka])
Hr2bab[kj, ka] += lib.einsum('LDBJ,LaD->JaB', imds.WOVVO[kl,kd,kb],
r2bab[kl,ka])
kl = kconserv[kb, kshift, kd]
Hr2bab[kj, ka] -= lib.einsum('LJad,LdB->JaB', imds.WOOvv[kl,kj,ka],
r2bab[kl,kd])
kl = kconserv[ka, kshift, kd]
Hr2aba[kj, ka] += lib.einsum('LDbj,LAD->jAb', imds.WOVvo[kl,kd,kb],
r2bbb[kl,ka])
Hr2aba[kj, ka] += lib.einsum('ldbj,lAd->jAb', imds.Wovvo[kl,kd,kb],
r2aba[kl,ka])
kl = kconserv[kb, kshift, kd]
Hr2aba[kj, ka] -= lib.einsum('ljAD,lDb->jAb', imds.WooVV[kl,kj,ka],
r2aba[kl,kd])
kl = kconserv[ka, kshift, kd]
Hr2bbb[kj, ka] += lib.einsum('LDBJ,LAD->JAB', imds.WOVVO[kl,kd,kb],
r2bbb[kl,ka])
Hr2bbb[kj, ka] += lib.einsum('ldBJ,lAd->JAB', imds.WovVO[kl,kd,kb],
r2aba[kl,ka])
# P(ab)
kl = kconserv[kb, kshift, kd]
Hr2bbb[kj, ka] -= lib.einsum('LDAJ,LBD->JAB', imds.WOVVO[kl,kd,ka],
r2bbb[kl,kb])
Hr2bbb[kj, ka] -= lib.einsum('ldAJ,lBd->JAB', imds.WovVO[kl,kd,ka],
r2aba[kl,kb])
vector = amplitudes_to_vector_ea([Hr1a, Hr1b], [Hr2aaa, Hr2aba, Hr2bab, Hr2bbb], kshift, kconserv)
return vector
[docs]
def eaccsd_diag(eom, kshift, imds=None):
if imds is None: imds = eom.make_imds()
t1, t2 = imds.t1, imds.t2
t1a, t1b = t1
t2aa, t2ab, t2bb = t2
nkpts, nocca, nvira = t1a.shape
nkpts, noccb, nvirb = t1b.shape
kconserv = imds.kconserv
#Hr1a = np.zeros((nvira), dtype=t1a.dtype)
#Hr1b = np.zeros((nvirb), dtype=t1b.dtype)
Hr2aaa = np.zeros((nkpts,nkpts,nocca,nvira,nvira), dtype=t1a.dtype)
Hr2aba = np.zeros((nkpts,nkpts,nocca,nvirb,nvira), dtype=t1a.dtype)
Hr2bab = np.zeros((nkpts,nkpts,noccb,nvira,nvirb), dtype=t1a.dtype)
Hr2bbb = np.zeros((nkpts,nkpts,noccb,nvirb,nvirb), dtype=t1b.dtype)
Hr1a = np.diag(imds.Fvv[kshift])
Hr1b = np.diag(imds.FVV[kshift])
if eom.partition == 'mp':
raise Exception("MP diag is not tested") # remove this to use untested code
for kj, ka in itertools.product(range(nkpts), repeat=2):
kb = kconserv[kshift, ka, kj]
Hr2aaa[kj,ka] -= imds.Foo[kj,:,None,None]
Hr2aaa[kj,ka] += imds.Fvv[ka,None,:,None]
Hr2aaa[kj,ka] += imds.Fvv[kb,None,None,:]
Hr2aba[kj,ka] -= imds.Foo[kj,:,None,None]
Hr2aba[kj,ka] += imds.FVV[ka,None,:,None]
Hr2aba[kj,ka] += imds.Fvv[kb,None,None,:]
Hr2bab[kj,ka] -= imds.FOO[kj,:,None,None]
Hr2bab[kj,ka] += imds.Fvv[ka,None,:,None]
Hr2bab[kj,ka] += imds.FVV[kb,None,None,:]
Hr2bbb[kj,ka] -= imds.FOO[kj,:,None,None]
Hr2bbb[kj,ka] += imds.FVV[ka,None,:,None]
Hr2bbb[kj,ka] += imds.FVV[kb,None,None,:]
else:
for kj, ka in itertools.product(range(nkpts), repeat=2):
kb = kconserv[kshift, ka, kj]
# Fvv
Hr2aaa[kj,ka] += imds.Fvv[ka].diagonal()[None,:,None]
Hr2aaa[kj,ka] += imds.Fvv[kb].diagonal()[None,None,:]
Hr2aba[kj,ka] += imds.FVV[ka].diagonal()[None,:,None]
Hr2aba[kj,ka] += imds.Fvv[kb].diagonal()[None,None,:]
Hr2bab[kj,ka] += imds.Fvv[ka].diagonal()[None,:,None]
Hr2bab[kj,ka] += imds.FVV[kb].diagonal()[None,None,:]
Hr2bbb[kj,ka] += imds.FVV[ka].diagonal()[None,:,None]
Hr2bbb[kj,ka] += imds.FVV[kb].diagonal()[None,None,:]
# Foo
Hr2aaa[kj,ka] -= imds.Foo[kj].diagonal()[:,None,None]
Hr2bbb[kj,ka] -= imds.FOO[kj].diagonal()[:,None,None]
Hr2bab[kj,ka] -= imds.FOO[kj].diagonal()[:,None,None]
Hr2aba[kj,ka] -= imds.Foo[kj].diagonal()[:,None,None]
# Wvvvv
Wvvvv, WvvVV, WVVVV = imds.get_Wvvvv(ka, kb, ka)
# FIXME: Do Wvvvv and WVVVV have a factor 0.5?
Hr2aaa[kj,ka] += lib.einsum('aabb->ab', Wvvvv)[None,:,:]
Hr2aba[kj,kb] += lib.einsum('bbAA->Ab', WvvVV)[None,:,:]
Hr2bab[kj,ka] += lib.einsum('aaBB->aB', WvvVV)[None,:,:]
Hr2bbb[kj,ka] += lib.einsum('AABB->AB', WVVVV)[None,:,:]
# Wovov term (physicist's Woovv)
Hr2aaa[kj,ka] -= lib.einsum('kajb, kjab->jab', imds.Wovov[kshift,ka,kj], t2aa[kshift,kj,ka])
Hr2aba[kj,ka] -= lib.einsum('jbKA, jKbA->jAb', imds.WovOV[kj,kb,kshift], t2ab[kj,kshift,kb])
Hr2bab[kj,ka] -= lib.einsum('kaJB, kJaB->JaB', imds.WovOV[kshift,ka,kj], t2ab[kshift,kj,ka])
Hr2bbb[kj,ka] -= lib.einsum('kajb, kjab->jab', imds.WOVOV[kshift,ka,kj], t2bb[kshift,kj,ka])
# Wovvo term
Hr2aaa[kj, ka] += lib.einsum('jbbj->jb', imds.Wovvo[kj,kb,kb])[:,None,:]
Hr2aaa[kj, ka] += lib.einsum('jaaj->ja', imds.Wovvo[kj,ka,ka])[:,:,None]
Hr2aba[kj, ka] += lib.einsum('jbbj->jb', imds.Wovvo[kj,kb,kb])[:,None,:]
Hr2aba[kj, ka] -= lib.einsum('jjAA->jA', imds.WooVV[kj,kj,ka])[:,:,None]
Hr2bab[kj, ka] += lib.einsum('JBBJ->JB', imds.WOVVO[kj,kb,kb])[:,None,:]
Hr2bab[kj, ka] -= lib.einsum('JJaa->Ja', imds.WOOvv[kj,kj,ka])[:,:,None]
Hr2bbb[kj, ka] += lib.einsum('JBBJ->JB', imds.WOVVO[kj,kb,kb])[:,None,:]
Hr2bbb[kj, ka] += lib.einsum('JAAJ->JA', imds.WOVVO[kj,ka,ka])[:,:,None]
vector = amplitudes_to_vector_ea([Hr1a,Hr1b], [Hr2aaa,Hr2aba,Hr2bab,Hr2bbb], kshift, kconserv)
return vector
[docs]
def mask_frozen_ea(eom, vector, kshift, const=LARGE_DENOM):
'''Replaces all frozen orbital indices of `vector` with the value `const`.'''
nkpts = eom.nkpts
nocca, noccb = eom.nocc
nmoa, nmob = eom.nmo
kconserv = eom.kconserv
r1, r2 = eom.vector_to_amplitudes(vector, kshift, nkpts, (nmoa, nmob), (nocca, noccb), kconserv)
r1a, r1b = r1
r2aaa, r2aba, r2bab, r2bbb = r2
# Get location of padded elements in occupied and virtual space
nonzero_opadding, nonzero_vpadding = eom.nonzero_opadding, eom.nonzero_vpadding
nonzero_opadding_a, nonzero_opadding_b = nonzero_opadding
nonzero_vpadding_a, nonzero_vpadding_b = nonzero_vpadding
new_r1a = const * np.ones_like(r1a)
new_r1b = const * np.ones_like(r1b)
new_r2aaa = const * np.ones_like(r2aaa)
new_r2aba = const * np.ones_like(r2aba)
new_r2bab = const * np.ones_like(r2bab)
new_r2bbb = const * np.ones_like(r2bbb)
# r1a/b case
new_r1a[nonzero_vpadding_a[kshift]] = r1a[nonzero_vpadding_a[kshift]]
new_r1b[nonzero_vpadding_b[kshift]] = r1b[nonzero_vpadding_b[kshift]]
# r2aaa case
for kj in range(nkpts):
for ka in range(nkpts):
kb = kconserv[kshift, ka, kj]
idx = np.ix_([kj], [ka], nonzero_opadding_a[kj], nonzero_vpadding_a[ka], nonzero_vpadding_a[kb])
new_r2aaa[idx] = r2aaa[idx]
# r2aba case
for kj in range(nkpts):
for ka in range(nkpts):
kb = kconserv[kshift, ka, kj]
idx = np.ix_([kj], [ka], nonzero_opadding_a[kj], nonzero_vpadding_b[ka], nonzero_vpadding_a[kb])
new_r2aba[idx] = r2aba[idx]
# r2bab case
for kj in range(nkpts):
for ka in range(nkpts):
kb = kconserv[kshift, ka, kj]
idx = np.ix_([kj], [ka], nonzero_opadding_b[kj], nonzero_vpadding_a[ka], nonzero_vpadding_b[kb])
new_r2bab[idx] = r2bab[idx]
# r2bbb case
for kj in range(nkpts):
for ka in range(nkpts):
kb = kconserv[kshift, ka, kj]
idx = np.ix_([kj], [ka], nonzero_opadding_b[kj], nonzero_vpadding_b[ka], nonzero_vpadding_b[kb])
new_r2bbb[idx] = r2bbb[idx]
return eom.amplitudes_to_vector((new_r1a,new_r1b), (new_r2aaa,new_r2aba,new_r2bab,new_r2bbb), kshift)
[docs]
class EOMEA(eom_kgccsd.EOMEA):
def __init__(self, cc):
#if not isinstance(cc, kccsd.GCCSD):
# raise TypeError
self.kpts = cc.kpts
eom_kgccsd.EOMEA.__init__(self, cc)
get_diag = eaccsd_diag
matvec = eaccsd_matvec
get_padding_k_idx = get_padding_k_idx
mask_frozen = mask_frozen_ea
[docs]
def get_init_guess(self, kshift, nroots=1, koopmans=False, diag=None):
size = self.vector_size()
dtype = getattr(diag, 'dtype', np.complex128)
nroots = min(nroots, size)
nocca, noccb = self.nocc
nmoa, nmob = self.nmo
nvira = nmoa-nocca
guess = []
if koopmans:
idx = np.zeros(nroots, dtype=int)
tmp_valpha, tmp_vbeta = self.nonzero_vpadding[kshift]
tmp_valpha = list(tmp_valpha)
tmp_vbeta = list(tmp_vbeta)
if len(tmp_vbeta) + len(tmp_valpha) < nroots:
raise ValueError("Max number of roots for k-point (idx=%3d) for koopmans "
"is %3d.\nRequested %3d." %
(kshift, len(tmp_vbeta)+len(tmp_valpha), nroots))
total_count = 0
while (total_count < nroots):
if total_count % 2 == 0 and len(tmp_valpha) > 0:
idx[total_count] = tmp_valpha.pop(0)
else:
# Careful! index depends on how we create vector
# (here the first elements are r1a, then r1b)
idx[total_count] = nvira + tmp_vbeta.pop(0)
total_count += 1
else:
idx = diag.argsort()
for i in idx[:nroots]:
g = np.zeros(size, dtype)
g[i] = 1.0
g = self.mask_frozen(g, kshift, const=0.0)
guess.append(g)
return guess
[docs]
def vector_to_amplitudes(self, vector, kshift, nkpts=None, nmo=None, nocc=None, kconserv=None):
if nmo is None: nmo = self.nmo
if nocc is None: nocc = self.nocc
if nkpts is None: nkpts = self.nkpts
if kconserv is None: kconserv = self.kconserv
return vector_to_amplitudes_ea(vector, kshift, nkpts, nmo, nocc, kconserv)
[docs]
def amplitudes_to_vector(self, r1, r2, kshift, kconserv=None):
if kconserv is None: kconserv = self.kconserv
return amplitudes_to_vector_ea(r1, r2, kshift, kconserv)
[docs]
def vector_size(self):
nocca, noccb = self.nocc
nmoa, nmob = self.nmo
nvira, nvirb = nmoa - nocca, nmob - noccb
nkpts = self.nkpts
#return (nvira + nvirb +
# nocca*nkpts*nvira*nkpts*nvira +
# nkpts**2*nocca*nvirb*nvira +
# nkpts**2*noccb*nvira*nvirb +
# noccb*nkpts*nvirb*nkpts*nvirb)
return (nvira + nvirb +
nocca*nkpts*nvira*(nkpts*nvira-1)//2 +
nkpts**2*nocca*nvirb*nvira +
nkpts**2*noccb*nvira*nvirb +
noccb*nkpts*nvirb*(nkpts*nvirb-1)//2)
[docs]
def make_imds(self, eris=None, t1=None, t2=None):
imds = _IMDS(self._cc, eris, t1, t2)
imds.make_ea()
return imds
class _IMDS:
def __init__(self, cc, eris=None, t1=None, t2=None):
self._cc = cc
self.verbose = cc.verbose
self.kconserv = kpts_helper.get_kconserv(cc._scf.cell, cc.kpts)
self.stdout = cc.stdout
if t1 is None:
t1 = cc.t1
self.t1 = t1
if t2 is None:
t2 = cc.t2
self.t2 = t2
if eris is None:
eris = cc.ao2mo()
self.eris = eris
self._made_shared = False
self.made_ip_imds = False
self.made_ea_imds = False
self.made_ee_imds = False
def _make_shared(self):
cput0 = (logger.process_clock(), logger.perf_counter())
t1, t2, eris = self.t1, self.t2, self.eris
self.Foo, self.FOO = kintermediates_uhf.Foo(self._cc, t1, t2, eris)
self.Fvv, self.FVV = kintermediates_uhf.Fvv(self._cc, t1, t2, eris)
self.Fov, self.FOV = kintermediates_uhf.Fov(self._cc, t1, t2, eris)
# 2 virtuals
self.Wovvo, self.WovVO, self.WOVvo, self.WOVVO = kintermediates_uhf.Wovvo(self._cc, t1, t2, eris)
self.Woovv, self.WooVV, self.WOOvv, self.WOOVV = kintermediates_uhf.Woovv(self._cc, t1, t2, eris)
self.Wovov = eris.ovov - np.asarray(eris.ovov).transpose(2,1,0,5,4,3,6)
self.WOVOV = eris.OVOV - np.asarray(eris.OVOV).transpose(2,1,0,5,4,3,6)
self.WovOV = eris.ovOV
self.WOVov = None
self._made_shared = True
logger.timer_debug1(self, 'EOM-KCCSD shared intermediates', *cput0)
return self
def make_ip(self):
if not self._made_shared:
self._make_shared()
kconserv = self.kconserv
cput0 = (logger.process_clock(), logger.perf_counter())
t1, t2, eris = self.t1, self.t2, self.eris
# 0 or 1 virtuals
self.Woooo, self.WooOO, _ , self.WOOOO = kintermediates_uhf.Woooo(self._cc, t1, t2, eris) # noqa: E501
self.Wooov, self.WooOV, self.WOOov, self.WOOOV = kintermediates_uhf.Wooov(self._cc, t1, t2, eris, kconserv) # noqa: E501
self.Woovo, self.WooVO, self.WOOvo, self.WOOVO = kintermediates_uhf.Woovo(self._cc, t1, t2, eris) # noqa: E501
self.made_ip_imds = True
logger.timer_debug1(self, 'EOM-KUCCSD IP intermediates', *cput0)
return self
def make_ea(self):
if not self._made_shared:
self._make_shared()
cput0 = (logger.process_clock(), logger.perf_counter())
t1, t2, eris = self.t1, self.t2, self.eris
# 3 or 4 virtuals
#self.Wvovv, self.WvoVV, self.WVOvv, self.WVOVV = kintermediates_uhf.Wvovv(self._cc, t1, t2, eris)
self.Wvvov, self.WvvOV, self.WVVov, self.WVVOV = kintermediates_uhf.Wvvov(self._cc, t1, t2, eris)
if eris.vvvv is not None:
self.Wvvvv, self.WvvVV, self.WVVVV = kintermediates_uhf.Wvvvv(self._cc, t1, t2, eris)
else:
self.Wvvvv = self.WvvVV = self.WVVVV = None
self.Wvvvo, self.WvvVO, self.WVVvo, self.WVVVO = kintermediates_uhf.Wvvvo(self._cc, t1, t2, eris)
self.made_ea_imds = True
logger.timer_debug1(self, 'EOM-KUCCSD EA intermediates', *cput0)
return self
def make_ee(self):
raise NotImplementedError
def get_Wvvvv(self, ka, kb, kc):
if not self.made_ea_imds:
self.make_ea()
if self.Wvvvv is None:
return kintermediates_uhf.get_Wvvvv(self._cc, self.t1, self.t2, self.eris,
ka, kb, kc)
else:
return self.Wvvvv[ka,kc,kb], self.WvvVV[ka,kc,kb], self.WVVVV[ka,kc,kb]
if __name__ == '__main__':
from pyscf.pbc import gto
from pyscf.pbc import scf
from pyscf import lo
cell = gto.Cell()
cell.atom='''
He 0.000000000000 0.000000000000 0.000000000000
He 1.685068664391 1.685068664391 1.685068664391
'''
#cell.basis = [[0, (1., 1.)], [1, (.5, 1.)]]
cell.basis = [[0, (1., 1.)], [0, (.5, 1.)]]
cell.a = '''
0.000000000, 3.370137329, 3.370137329
3.370137329, 0.000000000, 3.370137329
3.370137329, 3.370137329, 0.000000000'''
cell.unit = 'B'
cell.mesh = [5, 5, 5]
cell.build()
np.random.seed(1)
# Running HF and CCSD with 1x1x2 Monkhorst-Pack k-point mesh
kmf = scf.KUHF(cell, kpts=cell.make_kpts([1,1,3]), exxdiv=None)
nmo = cell.nao_nr()
kmf.mo_occ = np.zeros((2,3,nmo))
kmf.mo_occ[0,:,:3] = 1
kmf.mo_occ[1,:,:1] = 1
kmf.mo_energy = np.arange(nmo) + np.random.random((2,3,nmo)) * .3
kmf.mo_energy[kmf.mo_occ == 0] += 2
mo = (np.random.random((2,3,nmo,nmo)) +
np.random.random((2,3,nmo,nmo))*1j - .5-.5j)
s = kmf.get_ovlp()
kmf.mo_coeff = np.empty_like(mo)
nkpts = len(kmf.kpts)
for k in range(nkpts):
kmf.mo_coeff[0,k] = lo.orth.vec_lowdin(mo[0,k], s[k])
kmf.mo_coeff[1,k] = lo.orth.vec_lowdin(mo[1,k], s[k])
def rand_t1_t2(mycc):
nkpts = mycc.nkpts
nocca, noccb = mycc.nocc
nmoa, nmob = mycc.nmo
nvira, nvirb = nmoa - nocca, nmob - noccb
np.random.seed(1)
t1a = (np.random.random((nkpts,nocca,nvira)) +
np.random.random((nkpts,nocca,nvira))*1j - .5-.5j)
t1b = (np.random.random((nkpts,noccb,nvirb)) +
np.random.random((nkpts,noccb,nvirb))*1j - .5-.5j)
t2aa = (np.random.random((nkpts,nkpts,nkpts,nocca,nocca,nvira,nvira)) +
np.random.random((nkpts,nkpts,nkpts,nocca,nocca,nvira,nvira))*1j - .5-.5j)
kconserv = kpts_helper.get_kconserv(kmf.cell, kmf.kpts)
t2aa = t2aa - t2aa.transpose(1,0,2,4,3,5,6)
tmp = t2aa.copy()
for ki, kj, kk in kpts_helper.loop_kkk(nkpts):
kl = kconserv[ki, kk, kj]
t2aa[ki,kj,kk] = t2aa[ki,kj,kk] - tmp[ki,kj,kl].transpose(0,1,3,2)
t2ab = (np.random.random((nkpts,nkpts,nkpts,nocca,noccb,nvira,nvirb)) +
np.random.random((nkpts,nkpts,nkpts,nocca,noccb,nvira,nvirb))*1j - .5-.5j)
t2bb = (np.random.random((nkpts,nkpts,nkpts,noccb,noccb,nvirb,nvirb)) +
np.random.random((nkpts,nkpts,nkpts,noccb,noccb,nvirb,nvirb))*1j - .5-.5j)
t2bb = t2bb - t2bb.transpose(1,0,2,4,3,5,6)
tmp = t2bb.copy()
for ki, kj, kk in kpts_helper.loop_kkk(nkpts):
kl = kconserv[ki, kk, kj]
t2bb[ki,kj,kk] = t2bb[ki,kj,kk] - tmp[ki,kj,kl].transpose(0,1,3,2)
t1 = (t1a, t1b)
t2 = (t2aa, t2ab, t2bb)
return t1, t2
from pyscf.pbc.cc import kccsd_uhf
mycc = kccsd_uhf.KUCCSD(kmf)
eris = mycc.ao2mo()
t1, t2 = rand_t1_t2(mycc)
mycc.t1 = t1
mycc.t2 = t2
kconserv = kpts_helper.get_kconserv(kmf.cell, kmf.kpts)
kgcc = kccsd.GCCSD(scf.addons.convert_to_ghf(kmf))
kccsd_eris = kccsd._make_eris_incore(kgcc, kgcc._scf.mo_coeff)
orbspin = kccsd_eris.orbspin
nkpts = mycc.nkpts
nocca, noccb = mycc.nocc
nmoa, nmob = mycc.nmo
nvira, nvirb = nmoa - nocca, nmob - noccb
kshift = 0 # excitation out of 0th k-point
nmo = nmoa + nmob
nocc = nocca + noccb
nvir = nmo - nocc
np.random.seed(0)
# IP version
myeom = EOMIP(mycc)
imds = myeom.make_imds()
imds.make_ip()
spin_r1_ip = (np.random.rand(nocc)*1j +
np.random.rand(nocc) - 0.5 - 0.5*1j)
spin_r2_ip = (np.random.rand(nkpts**2 * nocc**2 * nvir) +
np.random.rand(nkpts**2 * nocc**2 * nvir)*1j - 0.5 - 0.5*1j)
spin_r2_ip = spin_r2_ip.reshape(nkpts, nkpts, nocc, nocc, nvir)
spin_r2_ip = eom_kgccsd.enforce_2p_spin_ip_doublet(spin_r2_ip, kconserv, kshift, orbspin)
r1, r2 = eom_kgccsd.spin2spatial_ip_doublet(spin_r1_ip, spin_r2_ip, kconserv, kshift, orbspin)
vector = myeom.amplitudes_to_vector(r1, r2, kshift)
vector = myeom.matvec(vector, kshift=kshift, imds=imds)
Hr1, Hr2 = myeom.vector_to_amplitudes(vector, 0, nkpts, (nmoa, nmob), (nocca, noccb))
Hr1a, Hr1b = Hr1
Hr2aaa, Hr2baa, Hr2abb, Hr2bbb = Hr2
print('ip Hr1a', abs(lib.fp(Hr1a) - (-0.34462696543560045-1.6104596956729178j)))
print('ip Hr1b', abs(lib.fp(Hr1b) - (-0.055793611517250929+0.22169994342782473j)))
print('ip Hr2aaa', abs(lib.fp(Hr2aaa) - (0.692705827672665420-1.958639508839846943j)))
print('ip Hr2baa', abs(lib.fp(Hr2baa) - (2.892194153603884654+2.039530776282815872j)))
print('ip Hr2abb', abs(lib.fp(Hr2abb) - (1.618257685489421727-5.489218743953674817j)))
print('ip Hr2bbb', abs(lib.fp(Hr2bbb) - (0.479835513829048044+0.108406393138471210j)))
# EA version
myeom = EOMEA(mycc)
imds = myeom.make_imds()
imds.make_ea()
spin_r1_ea = (np.random.rand(nvir)*1j +
np.random.rand(nvir) - 0.5 - 0.5*1j)
spin_r2_ea = (np.random.rand(nkpts**2 * nocc * nvir**2) +
np.random.rand(nkpts**2 * nocc * nvir**2)*1j - 0.5 - 0.5*1j)
spin_r2_ea = spin_r2_ea.reshape(nkpts, nkpts, nocc, nvir, nvir)
spin_r2_ea = eom_kgccsd.enforce_2p_spin_ea_doublet(spin_r2_ea, kconserv, kshift, orbspin)
r1, r2 = eom_kgccsd.spin2spatial_ea_doublet(spin_r1_ea, spin_r2_ea, kconserv, kshift, orbspin)
vector = myeom.amplitudes_to_vector(r1, r2, kshift)
vector = myeom.matvec(vector, kshift=kshift, imds=imds)
Hr1, Hr2 = myeom.vector_to_amplitudes(vector, 0, nkpts, (nmoa, nmob), (nocca, noccb))
Hr1a, Hr1b = Hr1
Hr2aaa, Hr2aba, Hr2bab, Hr2bbb = Hr2
print('ea Hr1a', abs(lib.fp(Hr1a) - (-0.081373075311041126-0.51422895644026023j)))
print('ea Hr1b', abs(lib.fp(Hr1b) - (-0.39518588661294807-1.3063424820239824j)) )
print('ea Hr2aaa',abs(lib.fp(Hr2aaa) - (-2.6502079691200251-0.61302655915003545j)) )
print('ea Hr2aba',abs(lib.fp(Hr2aba) - (5.5723208649566036-5.4202659143496286j)) )
print('ea Hr2bab',abs(lib.fp(Hr2bab) - (-1.2822293707887937+0.3026476580141586j)) )
print('ea Hr2bbb',abs(lib.fp(Hr2bbb) - (-4.0202809577487253-0.46985725132191702j)) )
