from wavefunction_analysis import sys, np
from wavefunction_analysis.utils import print_matrix, convert_units
from wavefunction_analysis.utils.pyscf_parser import build_molecule
import pyscf
from pyscf import lib, ao2mo
from pyscf.lib import logger
from pyscf.dft import rks, gks
from pyscf.scf.hf import get_jk # for integral
[docs]
def pick_orbital_from_mom(mo0, ovlp, mo, imethod=0):
s = np.einsum('pi,pq,qj->j', mo0, ovlp, mo)
#print_matrix('s:', s)
return np.argsort(-np.abs(s))
[docs]
def get_occ(mf, mo_energy=None, mo_coeff=None):
mo_coeff0 = mf.previous_mo_coeff # reference
e_idx = pick_orbital_from_mom(mo_coeff0, mf.get_ovlp(), mo_coeff)
#print('e_idx:', e_idx)
print_matrix('mo_coeff:', mo_coeff, 10)
e_sort = mo_energy[e_idx]
nmo = mo_energy.size
mo_occ = np.zeros_like(mo_energy)
nocc = mf.mol.nelectron
mo_occ[e_idx[:nocc]] = 1
#print('mo_occ:', mo_occ)
if mf.verbose >= logger.INFO and nocc < nmo:
if e_sort[nocc-1]+1e-3 > e_sort[nocc]:
logger.warn(mf, 'HOMO %.15g == LUMO %.15g',
e_sort[nocc-1], e_sort[nocc])
else:
logger.info(mf, ' HOMO = %.15g LUMO = %.15g',
e_sort[nocc-1], e_sort[nocc])
if mf.verbose >= logger.DEBUG:
np.set_printoptions(threshold=nmo)
logger.debug(mf, ' mo_energy =\n%s', mo_energy)
np.set_printoptions(threshold=1000)
#if mo_coeff is not None and mf.verbose >= logger.DEBUG:
# ss, s = mf.spin_square(mo_coeff[:,mo_occ>0], mf.get_ovlp())
# logger.debug(mf, 'multiplicity <S^2> = %.8g 2S+1 = %.8g', ss, s)
mf.orbo = np.copy(mo_coeff[:,mo_occ==1])
if getattr(mf, 'imom', 0) < 1:
mf.previous_mo_coeff = np.copy(mo_coeff[:,mo_occ==1])
print_matrix('mf.previous_mo_coeff:', mf.previous_mo_coeff, 10)
mocc = mf.previous_mo_coeff
dm = (mocc*mo_occ[mo_occ>0]).dot(mocc.conj().T)
print_matrix('dm:', dm, 10)
return mo_occ
"""
restricted open shell scf for excited states
"""
[docs]
class Delta_UKS(gks.GKS):
get_occ = get_occ
"""
def get_jk(self, mol=None, dm=None, hermi=0, with_j=True, with_k=True,
omega=None):
if mol is None: mol = self.mol
if dm is None: dm = self.make_rdm1()
vj, vk = super().get_jk(mol, dm, hermi, with_j, with_k, omega)
ispin = self.spin_image
if ispin:
nao = mol.nao
kaa, kbb = np.copy(vk[:nao,:nao]), np.copy(vk[nao:,nao:])
print_matrix('jaa:', vj[:nao,:nao], 10)
print_matrix('kaa:', kaa, 10)
print_matrix('jbb:', vj[nao:,nao:], 10)
print_matrix('kbb:', kbb, 10)
vk[:nao,:nao] += ispin * kbb
vk[nao:,nao:] += ispin * kaa
return vj, vk
"""
[docs]
def get_veff(self, mol=None, dm=None, dm_last=0, vhf_last=0, hermi=1):
vxc = super().get_veff(mol, dm, dm_last, vhf_last, hermi)
ispin = self.spin_image
if ispin:
ecoul, exc, vj, vk = vxc.ecoul, vxc.exc, vxc.vj, vxc.vk
nao, nocc = self.orbo.shape
nao, nocc = nao//2, nocc//2
idx = np.sum(self.orbo[:nao]**2, axis=0)
idx = np.argsort(-idx).reshape(2, -1)
idx[0] = np.sort(idx[0])
idx[1] = np.sort(idx[1])
print('idx:', idx)
moa, mob = self.orbo[:nao,idx[0]], self.orbo[nao:,idx[1]]
print_matrix('moa:', moa, 10)
print_matrix('mob:', mob, 10)
ovlp = self.get_ovlp()
zeros = np.zeros((nao, nao))
ovlp = ovlp[:nao,:nao]
#print_matrix('ovlp0:', np.einsum('pi,pq,qj->ij', moa, ovlp, mob), 10)
u, s, vt = np.linalg.svd(np.einsum('pi,pq,qj->ij', moa, ovlp, mob))
print_matrix('singular values:', s, 10)
s = - np.prod(s[:-1])
orb_u = np.einsum('pi,i->p', moa, u[:,-1]).reshape(-1,1)
orb_v = np.einsum('pi,i->p', mob, vt[:,-1]).reshape(-1,1)
eri_mo = ao2mo.general(mol, [orb_u,orb_v,orb_v,orb_u], compact=False)
eri_mo *= s
print_matrix('eri_mo energy:', eri_mo)
dm1 = np.einsum('pi,qi->pq', orb_u, orb_v) # alpha-beta
vj = s * get_jk(mol, dm1, hermi, with_k=False)[0]
eri_mo = np.einsum('pq,qp->', dm1, vj)
print_matrix('eri_mo energy check:', np.array([eri_mo]))
vjab = (.5*ispin) * vj
#ovlp = np.block([[zeros,ovlp], [ovlp,zeros]])
#ovlp = np.einsum('pi,pq,qj->ij', self.orbo, ovlp, self.orbo)
#print_matrix('ovlp:', ovlp, 10)
#u, s, vt = np.linalg.svd(ovlp)
#print_matrix('s:', s, 10)
#orb_u = np.einsum('pi,ij->pj', self.orbo, u)
#orb_v = np.einsum('pi,ij->pj', self.orbo, vt)
#orb_t = np.block([orb_u[:,-2:], orb_v[:,-2:]])
#print_matrix('orb_t:', orb_t)
#eri_mo = ao2mo.general(mol, [orb_t,orb_t,orb_t,orb_t], compact=False)
#eri_mo = eri_mo.reshape(4,4,4,4)
#print_matrix('eri_mo 2:', eri_mo)
print_matrix('vjab:', vjab, 10)
#
# vj[:nao,nao:] += vjab
# vj[nao:,:nao] += vjab
# vxc[:nao,nao:] += vjab
# vxc[nao:,:nao] += vjab
#vj[:nao,:nao] += vjab[1]
#vj[nao:,nao:] += vjab[0]
#vxc[:nao,:nao] += vjab[1]
#vxc[nao:,nao:] += vjab[0]
vxc = lib.tag_array(vxc, ecoul=ecoul, exc=exc, vj=vj, vk=vk)
return vxc
[docs]
def require_spin_image(self):
ispin = getattr(self, 'spin_image', 0)
if type(ispin) is str:
if ispin.lower() == 'singlet':
ispin = 1
elif ispin.lower() == 'triplet':
ispin = -1
self.spin_image = int(ispin)
[docs]
def get_init_guess(self, mol=None, key='minao', **kwargs):
self.require_spin_image()
if not isinstance(key, str):
return key
key = key.lower()
if mol is None:
mol = self.mol
if key == 'aufbau':
return self.init_guess_by_aufbau(mol, **kwargs)
else:
raise NotImplementedError(f'{key} initial guess not implemented')
[docs]
def init_guess_by_aufbau(self, mol=None, **kwargs):
if mol is None:
mol = self.mol
mo_coeff0 = kwargs.get('mo_coeff')
nocc = mol.nelectron // 2
nbas, norb = mo_coeff0.shape
mo_occ = np.zeros(norb)
mo_occ[:nocc-1] = 1
mo_occ[nocc+1] = 1
moa, mob = np.sqrt(.5)*mo_coeff0[:,mo_occ>0], np.sqrt(.5)*mo_coeff0[:,:nocc]
zeros = np.zeros_like(moa)
self.previous_mo_coeff = np.block([[moa, zeros], [zeros, mob]])
#self.previous_mo_coeff = np.block([[moa, mob], [moa, mob]])
print_matrix('previous_mo_coeff:'+str(self.previous_mo_coeff.shape), self.previous_mo_coeff, 10)
self.orbo = np.copy(self.previous_mo_coeff)
return self.make_rdm1(self.previous_mo_coeff, np.ones(nocc*2))
if __name__ == '__main__':
atom = """
C 0. 0. 0.
O 0. 0. 1.05
"""
basis = 'sto-3g'
functional = 'hf'
charge = 0
spin = 0 # nalpha-nbeta
verbose = 0
mol = build_molecule(atom, basis, charge, spin, verbose=verbose)
mf = rks.RKS(mol)
mf.xc = functional
e_tot = mf.kernel()
print('scf energy:', e_tot)
mo_coeff = mf.mo_coeff
mo_energy = mf.mo_energy
from pyscf.tdscf import rks as tdrks
td = tdrks.TDA(mf)
td.nroots = 5
td.kernel()
td.analyze(verbose=5)
#gks = gks.GKS(mol)
#gks.xc = functional
#gks.kernel()
#print('mo_occ:', gks.mo_occ)
#print_matrix('mo_coeff:', gks.mo_coeff, 10)
#print_matrix('mo_energy:', gks.mo_energy, 10)
print('start delta scf')
dscf = Delta_UKS(mol)
dscf.xc = functional
dscf.init_guess = 'aufbau'
dscf.imom = 1
dscf.spin_image = 1
e_tot1 = dscf.kernel(mo_coeff=mo_coeff)
ee = e_tot1 - e_tot
print('excitation energy:', ee, convert_units(ee, 'eh', 'ev'))
print_matrix('mo_energy:', dscf.mo_energy, 10)