def test_shls_slice1(self):
mol = gto.M(atom='H 0 -.5 0; H 0 .5 0', basis='cc-pvdz')
nao = mol.nao_nr()
dm = numpy.random.random((nao,nao))
mol1 = gto.M(atom='He 2 0 0', basis='6-31g')
nao1 = mol1.nao_nr()
dm1 = numpy.random.random((nao1,nao1))
eri0 = gto.conc_mol(mol, mol1).intor('int2e_sph').reshape([nao+nao1]*4)
j1part = jk.get_jk((mol1,mol1,mol,mol), dm1[:1,:1], scripts='ijkl,ji->kl', intor='int2e',
shls_slice=(0,1,0,1,0,mol.nbas,0,mol.nbas))
j1ref = numpy.einsum('ijkl,ji->kl', eri0[nao:nao+1,nao:nao+1,:nao,:nao], dm1[:1,:1])
self.assertAlmostEqual(abs(j1part - j1ref).max(), 0, 12)
k1part = jk.get_jk((mol1,mol,mol,mol1), dm1[:,:1], scripts='ijkl,li->kj', intor='int2e',
shls_slice=(0,1,0,1,0,mol.nbas,0,mol1.nbas))
k1ref = numpy.einsum('ijkl,li->kj', eri0[nao:nao+1,:1,:nao,nao:], dm1[:,:1])
self.assertAlmostEqual(abs(k1part - k1ref).max(), 0, 12)
j1part = jk.get_jk(mol, dm[:1,1:2], scripts='ijkl,ji->kl', intor='int2e',
shls_slice=(1,2,0,1,0,mol.nbas,0,mol.nbas))
j1ref = numpy.einsum('ijkl,ji->kl', eri0[1:2,:1,:nao,:nao], dm[:1,1:2])
self.assertAlmostEqual(abs(j1part - j1ref).max(), 0, 12)
k1part = jk.get_jk(mol, dm[:,1:2], scripts='ijkl,li->kj', intor='int2e',
shls_slice=(1,2,0,1,0,mol.nbas,0,mol.nbas))
k1ref = numpy.einsum('ijkl,li->kj', eri0[:1,1:2,:nao,:nao], dm[:,1:2])
self.assertAlmostEqual(abs(k1part - k1ref).max(), 0, 12)
def dia(magobj, gauge_orig=None):
mol = magobj.mol
mf = magobj._scf
mo_energy = magobj._scf.mo_energy
mo_coeff = magobj._scf.mo_coeff
mo_occ = magobj._scf.mo_occ
orboa = mo_coeff[0][:,mo_occ[0] > 0]
orbob = mo_coeff[1][:,mo_occ[1] > 0]
dm0a = numpy.dot(orboa, orboa.T)
dm0b = numpy.dot(orbob, orbob.T)
dm0 = dm0a + dm0b
dme0a = numpy.dot(orboa * mo_energy[0][mo_occ[0] > 0], orboa.T)
dme0b = numpy.dot(orbob * mo_energy[1][mo_occ[1] > 0], orbob.T)
dme0 = dme0a + dme0b
e2 = rhf_mag._get_dia_1e(magobj, gauge_orig, dm0, dme0).ravel()
if gauge_orig is None:
vs = jk.get_jk(mol, [dm0, dm0a, dm0a, dm0b, dm0b],
['ijkl,ji->s2kl',
'ijkl,jk->s1il', 'ijkl,li->s1kj',
'ijkl,jk->s1il', 'ijkl,li->s1kj'],
'int2e_gg1', 's4', 9, hermi=1)
e2 += numpy.einsum('xpq,qp->x', vs[0], dm0)
e2 -= numpy.einsum('xpq,qp->x', vs[1], dm0a) * .5
e2 -= numpy.einsum('xpq,qp->x', vs[2], dm0a) * .5
e2 -= numpy.einsum('xpq,qp->x', vs[3], dm0b) * .5
e2 -= numpy.einsum('xpq,qp->x', vs[4], dm0b) * .5
vk = jk.get_jk(mol, [dm0a, dm0b], ['ijkl,jk->s1il', 'ijkl,jk->s1il'],
'int2e_g1g2', 'aa4', 9, hermi=0)
e2 -= numpy.einsum('xpq,qp->x', vk[0], dm0a)
e2 -= numpy.einsum('xpq,qp->x', vk[1], dm0b)
return -e2.reshape(3, 3)
def test_shls_slice1(self):
mol = gto.M(atom='H 0 -.5 0; H 0 .5 0', basis='cc-pvdz')
nao = mol.nao_nr()
dm = numpy.random.random((nao,nao))
mol1 = gto.M(atom='He 2 0 0', basis='6-31g')
nao1 = mol1.nao_nr()
dm1 = numpy.random.random((nao1,nao1))
eri0 = gto.conc_mol(mol, mol1).intor('int2e_sph').reshape([nao+nao1]*4)
j1part = jk.get_jk((mol1,mol1,mol,mol), dm1[:1,:1], scripts='ijkl,ji->kl', intor='int2e',
shls_slice=(0,1,0,1,0,mol.nbas,0,mol.nbas))
j1ref = numpy.einsum('ijkl,ji->kl', eri0[nao:nao+1,nao:nao+1,:nao,:nao], dm1[:1,:1])
self.assertAlmostEqual(abs(j1part - j1ref).max(), 0, 12)
k1part = jk.get_jk((mol1,mol,mol,mol1), dm1[:,:1], scripts='ijkl,li->kj', intor='int2e',
shls_slice=(0,1,0,1,0,mol.nbas,0,mol1.nbas))
k1ref = numpy.einsum('ijkl,li->kj', eri0[nao:nao+1,:1,:nao,nao:], dm1[:,:1])
self.assertAlmostEqual(abs(k1part - k1ref).max(), 0, 12)
j1part = jk.get_jk(mol, dm[:1,1:2], scripts='ijkl,ji->kl', intor='int2e',
shls_slice=(1,2,0,1,0,mol.nbas,0,mol.nbas))
j1ref = numpy.einsum('ijkl,ji->kl', eri0[1:2,:1,:nao,:nao], dm[:1,1:2])
self.assertAlmostEqual(abs(j1part - j1ref).max(), 0, 12)
k1part = jk.get_jk(mol, dm[:,1:2], scripts='ijkl,li->kj', intor='int2e',
shls_slice=(1,2,0,1,0,mol.nbas,0,mol.nbas))
k1ref = numpy.einsum('ijkl,li->kj', eri0[:1,1:2,:nao,:nao], dm[:,1:2])
self.assertAlmostEqual(abs(k1part - k1ref).max(), 0, 12)
def dia(magobj, gauge_orig=None):
'''Diamagnetic term of magnetizability.
See also J. Olsen et al., Theor. Chem. Acc., 90, 421 (1995)
'''
mol = magobj.mol
mf = magobj._scf
mo_energy = mf.mo_energy
mo_coeff = mf.mo_coeff
mo_occ = mf.mo_occ
orbo = mo_coeff[:, mo_occ > 0]
dm0 = numpy.dot(orbo, orbo.T) * 2
# Energy weighted density matrix
dme0 = numpy.dot(orbo * mo_energy[mo_occ > 0], orbo.T) * 2
e2 = _get_dia_1e(magobj, gauge_orig, dm0, dme0)
# Add the contributions from two-electron interactions
if gauge_orig is None:
e2 = e2.ravel()
# + 1/2 Tr[(J^{[20]} - K^{[20]}), DM]
# Symmetry between 'ijkl,ji->s2kl' and 'ijkl,lk->s2ij' can be used.
#vs = jk.get_jk(mol, [dm0]*4, ['ijkl,ji->s2kl',
# 'ijkl,lk->s2ij',
# 'ijkl,jk->s1il',
# 'ijkl,li->s1kj'],
# 'int2e_gg1', 's4', 9, hermi=1)
#e2 += numpy.einsum('xpq,qp->x', vs[0], dm0) * .5
#e2 += numpy.einsum('xpq,qp->x', vs[1], dm0) * .5
#e2 -= numpy.einsum('xpq,qp->x', vs[2], dm0) * .25
#e2 -= numpy.einsum('xpq,qp->x', vs[3], dm0) * .25
vs = jk.get_jk(mol, [dm0] * 3,
['ijkl,ji->s2kl', 'ijkl,jk->s1il', 'ijkl,li->s1kj'],
'int2e_gg1',
's4',
9,
hermi=1)
e2 += numpy.einsum('xpq,qp->x', vs[0], dm0)
e2 -= numpy.einsum('xpq,qp->x', vs[1], dm0) * .25
e2 -= numpy.einsum('xpq,qp->x', vs[2], dm0) * .25
# J does not have contribution because integrals are anti-symmetric
#vs = jk.get_jk(mol, [dm0]*2, ['ijkl,ji->s2kl',
# 'ijkl,jk->s1il'],
# 'int2e_g1g2', 'aa4', 9, hermi=0)
#e2 += numpy.einsum('xpq,qp->x', vs[0], dm0)
#e2 -= numpy.einsum('xpq,qp->x', vs[1], dm0) * .5
vk = jk.get_jk(mol,
dm0,
'ijkl,jk->s1il',
'int2e_g1g2',
'aa4',
9,
hermi=0)
e2 -= numpy.einsum('xpq,qp->x', vk, dm0) * .5
# Note the sign of magnetizability (ksi) in the Talyor expansion
# E(B) = E0 - m * B - 1/2 B ksi B + ...
# Magnetic susceptibility chi = mu0 * ksi
return -e2.reshape(3, 3)
def get_vk_s4(mol, dm):
ao_loc = mol.ao_loc_nr()
nao = ao_loc[-1]
vk = numpy.zeros((nao,nao))
bas_groups = list(lib.prange(0, mol.nbas, 3))
for ip, (ish0, ish1) in enumerate(bas_groups):
for jp, (jsh0, jsh1) in enumerate(bas_groups[:ip]):
for kp, (ksh0, ksh1) in enumerate(bas_groups):
for lp, (lsh0, lsh1) in enumerate(bas_groups[:kp]):
shls_slice = (ish0, ish1, jsh0, jsh1, ksh0, ksh1, lsh0, lsh1)
i0, i1, j0, j1, k0, k1, l0, l1 = [ao_loc[x] for x in shls_slice]
dms = [dm[j0:j1,k0:k1],
dm[i0:i1,k0:k1],
dm[j0:j1,l0:l1],
dm[i0:i1,l0:l1]]
scripts = ['ijkl,jk->il',
'ijkl,ik->jl',
'ijkl,jl->ik',
'ijkl,il->jk']
kparts = jk.get_jk(mol, dms, scripts, shls_slice=shls_slice)
vk[i0:i1,l0:l1] += kparts[0]
vk[j0:j1,l0:l1] += kparts[1]
vk[i0:i1,k0:k1] += kparts[2]
vk[j0:j1,k0:k1] += kparts[3]
lsh0, lsh1 = ksh0, ksh1
shls_slice = (ish0, ish1, jsh0, jsh1, ksh0, ksh1, lsh0, lsh1)
i0, i1, j0, j1, k0, k1, l0, l1 = [ao_loc[x] for x in shls_slice]
kparts = jk.get_jk(mol,
[dm[j0:j1,k0:k1], dm[i0:i1,k0:k1]],
scripts=['ijkl,jk->il', 'ijkl,ik->jl'],
shls_slice=shls_slice)
vk[i0:i1,l0:l1] += kparts[0]
vk[j0:j1,l0:l1] += kparts[1]
jsh0, jsh1 = ish0, ish1
for kp, (ksh0, ksh1) in enumerate(bas_groups):
for lp, (lsh0, lsh1) in enumerate(bas_groups[:kp]):
shls_slice = (ish0, ish1, jsh0, jsh1, ksh0, ksh1, lsh0, lsh1)
i0, i1, j0, j1, k0, k1, l0, l1 = [ao_loc[x] for x in shls_slice]
kparts = jk.get_jk(mol,
[dm[j0:j1,k0:k1], dm[j0:j1,l0:l1]],
scripts=['ijkl,jk->il', 'ijkl,jl->ik'],
shls_slice=shls_slice)
vk[i0:i1,l0:l1] += kparts[0]
vk[j0:j1,k0:k1] += kparts[1]
lsh0, lsh1 = ksh0, ksh1
shls_slice = (ish0, ish1, jsh0, jsh1, ksh0, ksh1, lsh0, lsh1)
i0, i1, j0, j1, k0, k1, l0, l1 = [ao_loc[x] for x in shls_slice]
kparts = jk.get_jk(mol,
[dm[j0:j1,k0:k1]],
scripts=['ijkl,jk->il'],
shls_slice=shls_slice)
vk[i0:i1,l0:l1] += kparts[0]
return vk
def get_vk_s4(mol, dm):
ao_loc = mol.ao_loc_nr()
nao = ao_loc[-1]
vk = numpy.zeros((nao,nao))
bas_groups = list(lib.prange(0, mol.nbas, 3))
for ip, (ish0, ish1) in enumerate(bas_groups):
for jp, (jsh0, jsh1) in enumerate(bas_groups[:ip]):
for kp, (ksh0, ksh1) in enumerate(bas_groups):
for lp, (lsh0, lsh1) in enumerate(bas_groups[:kp]):
shls_slice = (ish0, ish1, jsh0, jsh1, ksh0, ksh1, lsh0, lsh1)
i0, i1, j0, j1, k0, k1, l0, l1 = [ao_loc[x] for x in shls_slice]
dms = [dm[j0:j1,k0:k1],
dm[i0:i1,k0:k1],
dm[j0:j1,l0:l1],
dm[i0:i1,l0:l1]]
scripts = ['ijkl,jk->il',
'ijkl,ik->jl',
'ijkl,jl->ik',
'ijkl,il->jk']
kparts = jk.get_jk(mol, dms, scripts, shls_slice=shls_slice)
vk[i0:i1,l0:l1] += kparts[0]
vk[j0:j1,l0:l1] += kparts[1]
vk[i0:i1,k0:k1] += kparts[2]
vk[j0:j1,k0:k1] += kparts[3]
lsh0, lsh1 = ksh0, ksh1
shls_slice = (ish0, ish1, jsh0, jsh1, ksh0, ksh1, lsh0, lsh1)
i0, i1, j0, j1, k0, k1, l0, l1 = [ao_loc[x] for x in shls_slice]
kparts = jk.get_jk(mol,
[dm[j0:j1,k0:k1], dm[i0:i1,k0:k1]],
scripts=['ijkl,jk->il', 'ijkl,ik->jl'],
shls_slice=shls_slice)
vk[i0:i1,l0:l1] += kparts[0]
vk[j0:j1,l0:l1] += kparts[1]
jsh0, jsh1 = ish0, ish1
for kp, (ksh0, ksh1) in enumerate(bas_groups):
for lp, (lsh0, lsh1) in enumerate(bas_groups[:kp]):
shls_slice = (ish0, ish1, jsh0, jsh1, ksh0, ksh1, lsh0, lsh1)
i0, i1, j0, j1, k0, k1, l0, l1 = [ao_loc[x] for x in shls_slice]
kparts = jk.get_jk(mol,
[dm[j0:j1,k0:k1], dm[j0:j1,l0:l1]],
scripts=['ijkl,jk->il', 'ijkl,jl->ik'],
shls_slice=shls_slice)
vk[i0:i1,l0:l1] += kparts[0]
vk[j0:j1,k0:k1] += kparts[1]
lsh0, lsh1 = ksh0, ksh1
shls_slice = (ish0, ish1, jsh0, jsh1, ksh0, ksh1, lsh0, lsh1)
i0, i1, j0, j1, k0, k1, l0, l1 = [ao_loc[x] for x in shls_slice]
kparts = jk.get_jk(mol,
[dm[j0:j1,k0:k1]],
scripts=['ijkl,jk->il'],
shls_slice=shls_slice)
vk[i0:i1,l0:l1] += kparts[0]
return vk
def dia(magobj, gauge_orig=None):
'''Diamagnetic term of magnetizability.
See also J. Olsen et al., Theor. Chem. Acc., 90, 421 (1995)
'''
mol = magobj.mol
mf = magobj._scf
mo_energy = mf.mo_energy
mo_coeff = mf.mo_coeff
mo_occ = mf.mo_occ
orbo = mo_coeff[:,mo_occ > 0]
dm0 = numpy.dot(orbo, orbo.T) * 2
# Energy weighted density matrix
dme0 = numpy.dot(orbo * mo_energy[mo_occ > 0], orbo.T) * 2
e2 = _get_dia_1e(magobj, gauge_orig, dm0, dme0)
# Add the contributions from two-electron interactions
if gauge_orig is None:
e2 = e2.ravel()
# + 1/2 Tr[(J^{[20]} - K^{[20]}), DM]
# Symmetry between 'ijkl,ji->s2kl' and 'ijkl,lk->s2ij' can be used.
#vs = jk.get_jk(mol, [dm0]*4, ['ijkl,ji->s2kl',
# 'ijkl,lk->s2ij',
# 'ijkl,jk->s1il',
# 'ijkl,li->s1kj'],
# 'int2e_gg1', 's4', 9, hermi=1)
#e2 += numpy.einsum('xpq,qp->x', vs[0], dm0) * .5
#e2 += numpy.einsum('xpq,qp->x', vs[1], dm0) * .5
#e2 -= numpy.einsum('xpq,qp->x', vs[2], dm0) * .25
#e2 -= numpy.einsum('xpq,qp->x', vs[3], dm0) * .25
vs = jk.get_jk(mol, [dm0]*3, ['ijkl,ji->s2kl',
'ijkl,jk->s1il',
'ijkl,li->s1kj'],
'int2e_gg1', 's4', 9, hermi=1)
e2 += numpy.einsum('xpq,qp->x', vs[0], dm0)
e2 -= numpy.einsum('xpq,qp->x', vs[1], dm0) * .25
e2 -= numpy.einsum('xpq,qp->x', vs[2], dm0) * .25
# J does not have contribution because integrals are anti-symmetric
#vs = jk.get_jk(mol, [dm0]*2, ['ijkl,ji->s2kl',
# 'ijkl,jk->s1il'],
# 'int2e_g1g2', 'aa4', 9, hermi=0)
#e2 += numpy.einsum('xpq,qp->x', vs[0], dm0)
#e2 -= numpy.einsum('xpq,qp->x', vs[1], dm0) * .5
vk = jk.get_jk(mol, dm0, 'ijkl,jk->s1il',
'int2e_g1g2', 'aa4', 9, hermi=0)
e2 -= numpy.einsum('xpq,qp->x', vk, dm0) * .5
# Note the sign of magnetizability (ksi) in the Talyor expansion
# E(B) = E0 - m * B - 1/2 B ksi B + ...
# Magnetic susceptibility chi = mu0 * ksi
return -e2.reshape(3, 3)
def test_range_separated_Coulomb(self):
'''test range-separated Coulomb'''
with mol.with_range_coulomb(0.2):
dm = mf.make_rdm1()
vk0 = jk.get_jk(mol, dm, 'ijkl,jk->il', hermi=0)
vk1 = jk.get_jk(mol, (dm,dm), ['ijkl,jk->il','ijkl,li->kj'], hermi=1)
self.assertAlmostEqual(abs(vk1[0]-vk0).max(), 0, 9)
self.assertAlmostEqual(abs(vk1[1]-vk0).max(), 0, 9)
self.assertAlmostEqual(lib.finger(vk0), 0.87325708945599279, 9)
vk = scf.hf.get_jk(mol, dm)[1]
self.assertAlmostEqual(abs(vk-vk0).max(), 0, 12)
vk = scf.hf.get_jk(mol, dm)[1]
self.assertTrue(abs(vk-vk0).max() > 0.1)
def test_range_separated_Coulomb(self):
'''test range-separated Coulomb'''
with mol.with_range_coulomb(0.2):
dm = mf.make_rdm1()
vk0 = jk.get_jk(mol, dm, 'ijkl,jk->il', hermi=0)
vk1 = jk.get_jk(mol, (dm,dm), ['ijkl,jk->il','ijkl,li->kj'], hermi=1)
self.assertAlmostEqual(abs(vk1[0]-vk0).max(), 0, 9)
self.assertAlmostEqual(abs(vk1[1]-vk0).max(), 0, 9)
self.assertAlmostEqual(lib.finger(vk0), 0.87325708945599279, 9)
vk = scf.hf.get_jk(mol, dm)[1]
self.assertAlmostEqual(abs(vk-vk0).max(), 0, 12)
vk = scf.hf.get_jk(mol, dm)[1]
self.assertTrue(abs(vk-vk0).max() > 0.1)
def _get_k_lr(mol, dm, omega=0, hermi=0, vhfopt=None):
import sys
sys.stderr.write('This function is deprecated. '
'It is replaced by mol.get_k(mol, dm, omege=omega)')
dm = numpy.asarray(dm)
# Note, ks object caches the ERIs for small systems. The cached eris are
# computed with regular Coulomb operator. ks.get_jk or ks.get_k do not evalute
# the K matrix with the range separated Coulomb operator. Here jk.get_jk
# function computes the K matrix with the modified Coulomb operator.
nao = dm.shape[-1]
dms = dm.reshape(-1, nao, nao)
with mol.with_range_coulomb(omega):
# Compute the long range part of ERIs temporarily with omega. Restore
# the original omega when the block ends
if vhfopt is None:
contents = lambda: None # just a place_holder
else:
contents = vhfopt._this.contents
with lib.temporary_env(
contents, fprescreen=_vhf._fpointer('CVHFnrs8_vk_prescreen')):
intor = mol._add_suffix('int2e')
vklr = jk.get_jk(mol,
dms, ['ijkl,jk->il'] * len(dms),
intor=intor,
vhfopt=vhfopt)
return numpy.asarray(vklr).reshape(dm.shape)
def test_vk_s8(self):
mol = gto.M(atom='H 0 -.5 0; H 0 .5 0; H 1.1 0.2 0.2; H 0.6 0.5 0.4',
basis='cc-pvdz')
ao_loc = mol.ao_loc_nr()
eri0 = mol.intor('int2e')
nao = mol.nao
numpy.random.seed(1)
dm = numpy.random.random((nao,nao))
vk0 = numpy.einsum('ijkl,jk->il', eri0, dm)
self.assertAlmostEqual(abs(vk0-get_vk_s4(mol, dm)).max(), 0, 12)
self.assertAlmostEqual(abs(vk0-get_vk_s8(mol, dm)).max(), 0, 12)
shls_slice = (2,4,0,2,0,8,0,5)
i0,i1,j0,j1,k0,k1,l0,l1 = [ao_loc[x] for x in shls_slice]
vk0 = numpy.einsum('ijkl,jk->il', eri0[i0:i1,j0:j1,k0:k1,l0:l1], dm[j0:j1,k0:k1])
vk1 = numpy.einsum('ijkl,jl->ik', eri0[i0:i1,j0:j1,k0:k1,l0:l1], dm[j0:j1,l0:l1])
vk2 = numpy.einsum('ijkl,ik->jl', eri0[i0:i1,j0:j1,k0:k1,l0:l1], dm[i0:i1,k0:k1])
vk3 = numpy.einsum('ijkl,il->jk', eri0[i0:i1,j0:j1,k0:k1,l0:l1], dm[i0:i1,l0:l1])
vk = jk.get_jk(mol,
[dm[j0:j1,k0:k1], dm[j0:j1,l0:l1], dm[i0:i1,k0:k1], dm[i0:i1,l0:l1]],
scripts=['ijkl,jk->il', 'ijkl,jl->ik', 'ijkl,ik->jl', 'ijkl,il->jk'],
shls_slice=shls_slice)
self.assertAlmostEqual(abs(vk0-vk[0]).max(), 0, 12)
self.assertAlmostEqual(abs(vk1-vk[1]).max(), 0, 12)
self.assertAlmostEqual(abs(vk2-vk[2]).max(), 0, 12)
self.assertAlmostEqual(abs(vk3-vk[3]).max(), 0, 12)
def test_vk_s8(self):
mol = gto.M(atom='H 0 -.5 0; H 0 .5 0; H 1.1 0.2 0.2; H 0.6 0.5 0.4',
basis='cc-pvdz')
ao_loc = mol.ao_loc_nr()
eri0 = mol.intor('int2e')
nao = mol.nao
numpy.random.seed(1)
dm = numpy.random.random((nao,nao))
vk0 = numpy.einsum('ijkl,jk->il', eri0, dm)
self.assertAlmostEqual(abs(vk0-get_vk_s4(mol, dm)).max(), 0, 12)
self.assertAlmostEqual(abs(vk0-get_vk_s8(mol, dm)).max(), 0, 12)
shls_slice = (2,4,0,2,0,8,0,5)
i0,i1,j0,j1,k0,k1,l0,l1 = [ao_loc[x] for x in shls_slice]
vk0 = numpy.einsum('ijkl,jk->il', eri0[i0:i1,j0:j1,k0:k1,l0:l1], dm[j0:j1,k0:k1])
vk1 = numpy.einsum('ijkl,jl->ik', eri0[i0:i1,j0:j1,k0:k1,l0:l1], dm[j0:j1,l0:l1])
vk2 = numpy.einsum('ijkl,ik->jl', eri0[i0:i1,j0:j1,k0:k1,l0:l1], dm[i0:i1,k0:k1])
vk3 = numpy.einsum('ijkl,il->jk', eri0[i0:i1,j0:j1,k0:k1,l0:l1], dm[i0:i1,l0:l1])
vk = jk.get_jk(mol,
[dm[j0:j1,k0:k1], dm[j0:j1,l0:l1], dm[i0:i1,k0:k1], dm[i0:i1,l0:l1]],
scripts=['ijkl,jk->il', 'ijkl,jl->ik', 'ijkl,ik->jl', 'ijkl,il->jk'],
shls_slice=shls_slice)
self.assertAlmostEqual(abs(vk0-vk[0]).max(), 0, 12)
self.assertAlmostEqual(abs(vk1-vk[1]).max(), 0, 12)
self.assertAlmostEqual(abs(vk2-vk[2]).max(), 0, 12)
self.assertAlmostEqual(abs(vk3-vk[3]).max(), 0, 12)
def test_mols(self):
pmol = copy.copy(mol)
mols = (mol, pmol, pmol, mol)
dm = mf.make_rdm1()
vj0 = jk.get_jk(mols, dm, 'ijkl,lk->ij')
vj1 = scf.hf.get_jk(mol, dm)[0]
self.assertAlmostEqual(abs(vj1-vj0).max(), 0, 9)
self.assertAlmostEqual(lib.finger(vj0), 28.36214139459754, 9)
def jk_ints(molA, molB, dm_ia, dm_jb):
from pyscf.scf import jk, _vhf
naoA = molA.nao
naoB = molB.nao
assert (dm_ia.shape == (naoA, naoA))
assert (dm_jb.shape == (naoB, naoB))
molAB = molA + molB
vhfopt = _vhf.VHFOpt(molAB, 'int2e', 'CVHFnrs8_prescreen',
'CVHFsetnr_direct_scf', 'CVHFsetnr_direct_scf_dm')
dmAB = scipy.linalg.block_diag(dm_ia, dm_jb)
# The prescreen function CVHFnrs8_prescreen indexes q_cond and dm_cond
# over the entire basis. "set_dm" in function jk.get_jk/direct_bindm only
# creates a subblock of dm_cond which is not compatible with
# CVHFnrs8_prescreen.
vhfopt.set_dm(dmAB, molAB._atm, molAB._bas, molAB._env)
# Then skip the "set_dm" initialization in function jk.get_jk/direct_bindm.
vhfopt._dmcondname = None
with lib.temporary_env(vhfopt._this.contents,
fprescreen=_vhf._fpointer('CVHFnrs8_vj_prescreen')):
shls_slice = (0, molA.nbas, 0, molA.nbas, molA.nbas, molAB.nbas,
molA.nbas, molAB.nbas) # AABB
vJ = jk.get_jk(molAB,
dm_jb,
'ijkl,lk->s2ij',
shls_slice=shls_slice,
vhfopt=vhfopt,
aosym='s4',
hermi=1)
cJ = np.einsum('ia,ia->', vJ, dm_ia)
with lib.temporary_env(vhfopt._this.contents,
fprescreen=_vhf._fpointer('CVHFnrs8_vk_prescreen')):
shls_slice = (0, molA.nbas, molA.nbas, molAB.nbas, molA.nbas,
molAB.nbas, 0, molA.nbas) # ABBA
vK = jk.get_jk(molAB,
dm_jb,
'ijkl,jk->il',
shls_slice=shls_slice,
vhfopt=vhfopt,
aosym='s1',
hermi=0)
cK = np.einsum('ia,ia->', vK, dm_ia)
return cJ, cK
def test_mols(self):
pmol = copy.copy(mol)
mols = (mol, pmol, pmol, mol)
dm = mf.make_rdm1()
vj0 = jk.get_jk(mols, dm, 'ijkl,lk->ij')
vj1 = scf.hf.get_jk(mol, dm)[0]
self.assertAlmostEqual(abs(vj1-vj0).max(), 0, 9)
self.assertAlmostEqual(lib.finger(vj0), 28.36214139459754, 9)
def test_shls_slice(self):
dm = mf.make_rdm1()
ao_loc = mol.ao_loc_nr()
shls_slice = [0, 2, 1, 4, 2, 5, 0, 4]
locs = [ao_loc[i] for i in shls_slice]
i0, i1, j0, j1, k0, k1, l0, l1 = locs
vs = jk.get_jk(mol, (dm[j0:j1,k0:k1], dm[l0:l1,k0:k1]),
['ijkl,jk->il', 'ijkl,lk->ij'], hermi=0,
intor='int2e_ip1', shls_slice=shls_slice)
self.assertEqual(vs[0].shape, (3,2,6))
self.assertEqual(vs[1].shape, (3,2,5))
def test_shls_slice(self):
dm = mf.make_rdm1()
ao_loc = mol.ao_loc_nr()
shls_slice = [0, 2, 1, 4, 2, 5, 0, 4]
locs = [ao_loc[i] for i in shls_slice]
i0, i1, j0, j1, k0, k1, l0, l1 = locs
vs = jk.get_jk(mol, (dm[j0:j1,k0:k1], dm[l0:l1,k0:k1]),
['ijkl,jk->il', 'ijkl,lk->ij'], hermi=0,
intor='int2e_ip1', shls_slice=shls_slice)
self.assertEqual(vs[0].shape, (3,2,6))
self.assertEqual(vs[1].shape, (3,2,5))
def _get_k_lr(mol, dm, omega=0, hermi=0):
dm = numpy.asarray(dm)
# Note, ks object caches the ERIs for small systems. The cached eris are
# computed with regular Coulomb operator. ks.get_jk or ks.get_k do not evalute
# the K matrix with the range separated Coulomb operator. Here jk.get_jk
# function computes the K matrix with the modified Coulomb operator.
nao = dm.shape[-1]
dms = dm.reshape(-1, nao, nao)
with mol.with_range_coulomb(omega):
# Compute the long range part of ERIs temporarily with omega. Restore
# the original omega when the block ends
vklr = jk.get_jk(mol, dms, ['ijkl,jk->il'] * len(dms))
return numpy.asarray(vklr).reshape(dm.shape)
def _get_k_lr(mol, dm, omega=0, hermi=0):
dm = numpy.asarray(dm)
# Note, ks object caches the ERIs for small systems. The cached eris are
# computed with regular Coulomb operator. ks.get_jk or ks.get_k do not evalute
# the K matrix with the range separated Coulomb operator. Here jk.get_jk
# function computes the K matrix with the modified Coulomb operator.
nao = dm.shape[-1]
dms = dm.reshape(-1,nao,nao)
with mol.with_range_coulomb(omega):
# Compute the long range part of ERIs temporarily with omega. Restore
# the original omega when the block ends
vklr = jk.get_jk(mol, dms, ['ijkl,jk->il']*len(dms))
return numpy.asarray(vklr).reshape(dm.shape)
def _get_k_lr(mol, dm, omega=0, hermi=0):
omega_bak = mol._env[gto.PTR_RANGE_OMEGA]
mol.set_range_coulomb(omega)
dm = numpy.asarray(dm)
# Note, ks object caches the ERIs for small systems. The cached eris are
# computed with regular Coulomb operator. ks.get_jk or ks.get_k do not evalute
# the K matrix with the range separated Coulomb operator. Here jk.get_jk
# function computes the K matrix with the modified Coulomb operator.
nao = dm.shape[-1]
dms = dm.reshape(-1, nao, nao)
vklr = jk.get_jk(mol, dms, ['ijkl,jk->il'] * len(dms))
mol.set_range_coulomb(omega_bak)
return numpy.asarray(vklr).reshape(dm.shape)
def dia(magobj, gauge_orig=None):
mol = magobj.mol
mf = magobj._scf
mo_energy = mf.mo_energy
mo_coeff = mf.mo_coeff
mo_occ = mf.mo_occ
orbo = mo_coeff[:, mo_occ > 0]
dm0 = numpy.dot(orbo, orbo.T) * 2
dm0 = lib.tag_array(dm0, mo_coeff=mo_coeff, mo_occ=mo_occ)
dme0 = numpy.dot(orbo * mo_energy[mo_occ > 0], orbo.T) * 2
e2 = rhf_mag._get_dia_1e(magobj, gauge_orig, dm0, dme0)
if gauge_orig is not None:
return -e2
# Computing the 2nd order Vxc integrals from GIAO
grids = mf.grids
ni = mf._numint
xc_code = mf.xc
xctype = ni._xc_type(xc_code)
omega, alpha, hyb = ni.rsh_and_hybrid_coeff(xc_code, mol.spin)
make_rho, nset, nao = ni._gen_rho_evaluator(mol, dm0, hermi=1)
ngrids = len(grids.weights)
mem_now = lib.current_memory()[0]
max_memory = max(2000, mf.max_memory * .9 - mem_now)
BLKSIZE = numint.BLKSIZE
blksize = min(
int(max_memory / 12 * 1e6 / 8 / nao / BLKSIZE) * BLKSIZE, ngrids)
vmat = numpy.zeros((3, 3, nao, nao))
if xctype == 'LDA':
ao_deriv = 0
for ao, mask, weight, coords \
in ni.block_loop(mol, grids, nao, ao_deriv, max_memory,
blksize=blksize):
rho = make_rho(0, ao, mask, 'LDA')
vxc = ni.eval_xc(xc_code, rho, 0, deriv=1)[1]
vrho = vxc[0]
r_ao = numpy.einsum('pi,px->pxi', ao, coords)
aow = numpy.einsum('pxi,p,p->pxi', r_ao, weight, vrho)
vmat += lib.einsum('pxi,pyj->xyij', r_ao, aow)
rho = vxc = vrho = aow = None
elif xctype == 'GGA':
ao_deriv = 1
for ao, mask, weight, coords \
in ni.block_loop(mol, grids, nao, ao_deriv, max_memory,
blksize=blksize):
rho = make_rho(0, ao, mask, 'GGA')
vxc = ni.eval_xc(xc_code, rho, 0, deriv=1)[1]
wv = numint._rks_gga_wv0(rho, vxc, weight)
# Computing \nabla (r * AO) = r * \nabla AO + [\nabla,r]_- * AO
r_ao = numpy.einsum('npi,px->npxi', ao, coords)
r_ao[1, :, 0] += ao[0]
r_ao[2, :, 1] += ao[0]
r_ao[3, :, 2] += ao[0]
aow = numpy.einsum('npxi,np->pxi', r_ao, wv)
vmat += lib.einsum('pxi,pyj->xyij', r_ao[0], aow)
rho = vxc = vrho = vsigma = wv = aow = None
vmat = vmat + vmat.transpose(0, 1, 3, 2)
elif xctype == 'MGGA':
raise NotImplementedError('meta-GGA')
vmat = _add_giao_phase(mol, vmat)
e2 += numpy.einsum('qp,xypq->xy', dm0, vmat)
vmat = None
e2 = e2.ravel()
# Handle the hybrid functional and the range-separated functional
if abs(hyb) > 1e-10:
vs = jk.get_jk(mol, [dm0] * 3,
['ijkl,ji->s2kl', 'ijkl,jk->s1il', 'ijkl,li->s1kj'],
'int2e_gg1',
's4',
9,
hermi=1)
e2 += numpy.einsum('xpq,qp->x', vs[0], dm0)
e2 -= numpy.einsum('xpq,qp->x', vs[1], dm0) * .25 * hyb
e2 -= numpy.einsum('xpq,qp->x', vs[2], dm0) * .25 * hyb
vk = jk.get_jk(mol,
dm0,
'ijkl,jk->s1il',
'int2e_g1g2',
'aa4',
9,
hermi=0)
e2 -= numpy.einsum('xpq,qp->x', vk, dm0) * .5 * hyb
if abs(omega) > 1e-10:
with mol.with_range_coulomb(omega):
vs = jk.get_jk(mol, [dm0] * 2,
['ijkl,jk->s1il', 'ijkl,li->s1kj'],
'int2e_gg1',
's4',
9,
hermi=1)
e2 -= numpy.einsum('xpq,qp->x', vs[0],
dm0) * .25 * (alpha - hyb)
e2 -= numpy.einsum('xpq,qp->x', vs[1],
dm0) * .25 * (alpha - hyb)
vk = jk.get_jk(mol,
dm0,
'ijkl,jk->s1il',
'int2e_g1g2',
'aa4',
9,
hermi=0)
e2 -= numpy.einsum('xpq,qp->x', vk, dm0) * .5 * (alpha - hyb)
else:
vj = jk.get_jk(mol,
dm0,
'ijkl,ji->s2kl',
'int2e_gg1',
's4',
9,
hermi=1)
e2 += numpy.einsum('xpq,qp->x', vj, dm0)
return -e2.reshape(3, 3)
def get_vk_s8(mol, dm):
ao_loc = mol.ao_loc_nr()
nao = ao_loc[-1]
vk = numpy.zeros((nao,nao))
bas_groups = list(lib.prange(0, mol.nbas, 3))
for ip, (ish0, ish1) in enumerate(bas_groups):
for jp, (jsh0, jsh1) in enumerate(bas_groups[:ip]):
for kp, (ksh0, ksh1) in enumerate(bas_groups[:ip]):
for lp, (lsh0, lsh1) in enumerate(bas_groups[:kp]):
shls_slice = (ish0, ish1, jsh0, jsh1, ksh0, ksh1, lsh0, lsh1)
i0, i1, j0, j1, k0, k1, l0, l1 = [ao_loc[x] for x in shls_slice]
dms = [dm[j0:j1,k0:k1],
dm[j0:j1,l0:l1],
dm[i0:i1,k0:k1],
dm[i0:i1,l0:l1],
dm[l0:l1,i0:i1],
dm[l0:l1,j0:j1],
dm[k0:k1,i0:i1],
dm[k0:k1,j0:j1]]
scripts = ['ijkl,jk->il',
'ijkl,jl->ik',
'ijkl,ik->jl',
'ijkl,il->jk',
'ijkl,li->kj',
'ijkl,lj->ki',
'ijkl,ki->lj',
'ijkl,kj->li']
kparts = jk.get_jk(mol, dms, scripts, shls_slice=shls_slice)
vk[i0:i1,l0:l1] += kparts[0]
vk[i0:i1,k0:k1] += kparts[1]
vk[j0:j1,l0:l1] += kparts[2]
vk[j0:j1,k0:k1] += kparts[3]
vk[k0:k1,j0:j1] += kparts[4]
vk[k0:k1,i0:i1] += kparts[5]
vk[l0:l1,j0:j1] += kparts[6]
vk[l0:l1,i0:i1] += kparts[7]
# ip > jp, ip > kp, kp == lp
lsh0, lsh1 = ksh0, ksh1
shls_slice = (ish0, ish1, jsh0, jsh1, ksh0, ksh1, lsh0, lsh1)
i0, i1, j0, j1, k0, k1, l0, l1 = [ao_loc[x] for x in shls_slice]
dms = [dm[j0:j1,k0:k1],
dm[i0:i1,k0:k1],
dm[k0:k1,j0:j1],
dm[k0:k1,i0:i1]]
scripts = ['ijkl,jk->il',
'ijkl,ik->jl',
'ijkl,kj->li',
'ijkl,ki->lj']
kparts = jk.get_jk(mol, dms, scripts, shls_slice=shls_slice)
vk[i0:i1,l0:l1] += kparts[0]
vk[j0:j1,l0:l1] += kparts[1]
vk[l0:l1,i0:i1] += kparts[2]
vk[l0:l1,j0:j1] += kparts[3]
# ip == kp and ip > jp and kp > lp
kp, ksh0, ksh1 = ip, ish0, ish1
for lp, (lsh0, lsh1) in enumerate(bas_groups[:kp]):
shls_slice = (ish0, ish1, jsh0, jsh1, ksh0, ksh1, lsh0, lsh1)
i0, i1, j0, j1, k0, k1, l0, l1 = [ao_loc[x] for x in shls_slice]
dms = [dm[j0:j1,k0:k1],
dm[i0:i1,k0:k1],
dm[j0:j1,l0:l1],
dm[i0:i1,l0:l1]]
scripts = ['ijkl,jk->il',
'ijkl,ik->jl',
'ijkl,jl->ik',
'ijkl,il->jk']
kparts = jk.get_jk(mol, dms, scripts, shls_slice=shls_slice)
vk[i0:i1,l0:l1] += kparts[0]
vk[j0:j1,l0:l1] += kparts[1]
vk[i0:i1,k0:k1] += kparts[2]
vk[j0:j1,k0:k1] += kparts[3]
# ip == jp and ip >= kp
jsh0, jsh1 = ish0, ish1
for kp, (ksh0, ksh1) in enumerate(bas_groups[:ip+1]):
for lp, (lsh0, lsh1) in enumerate(bas_groups[:kp]):
shls_slice = (ish0, ish1, jsh0, jsh1, ksh0, ksh1, lsh0, lsh1)
i0, i1, j0, j1, k0, k1, l0, l1 = [ao_loc[x] for x in shls_slice]
dms = [dm[j0:j1,k0:k1],
dm[j0:j1,l0:l1],
dm[k0:k1,j0:j1],
dm[l0:l1,j0:j1]]
scripts = ['ijkl,jk->il',
'ijkl,jl->ik',
'ijkl,kj->li',
'ijkl,lj->ki']
kparts = jk.get_jk(mol, dms, scripts, shls_slice=shls_slice)
vk[i0:i1,l0:l1] += kparts[0]
vk[i0:i1,k0:k1] += kparts[1]
vk[l0:l1,i0:i1] += kparts[2]
vk[k0:k1,i0:i1] += kparts[3]
# ip == jp and ip > kp and kp == lp
for kp, (ksh0, ksh1) in enumerate(bas_groups[:ip]):
lsh0, lsh1 = ksh0, ksh1
shls_slice = (ish0, ish1, jsh0, jsh1, ksh0, ksh1, lsh0, lsh1)
i0, i1, j0, j1, k0, k1, l0, l1 = [ao_loc[x] for x in shls_slice]
dms = [dm[j0:j1,k0:k1],
dm[l0:l1,i0:i1]]
scripts = ['ijkl,jk->il',
'ijkl,li->kj']
kparts = jk.get_jk(mol, dms, scripts, shls_slice=shls_slice)
vk[i0:i1,l0:l1] += kparts[0]
vk[k0:k1,j0:j1] += kparts[1]
# ip == jp == kp == lp
kp, ksh0, ksh1 = ip, ish0, ish1
lsh0, lsh1 = ksh0, ksh1
shls_slice = (ish0, ish1, jsh0, jsh1, ksh0, ksh1, lsh0, lsh1)
i0, i1, j0, j1, k0, k1, l0, l1 = [ao_loc[x] for x in shls_slice]
kparts = jk.get_jk(mol,
[dm[j0:j1,k0:k1]],
scripts=['ijkl,jk->il'],
shls_slice=shls_slice)
vk[i0:i1,l0:l1] += kparts[0]
return vk
def get_vk_s8(mol, dm):
ao_loc = mol.ao_loc_nr()
nao = ao_loc[-1]
vk = numpy.zeros((nao,nao))
bas_groups = list(lib.prange(0, mol.nbas, 3))
for ip, (ish0, ish1) in enumerate(bas_groups):
for jp, (jsh0, jsh1) in enumerate(bas_groups[:ip]):
for kp, (ksh0, ksh1) in enumerate(bas_groups[:ip]):
for lp, (lsh0, lsh1) in enumerate(bas_groups[:kp]):
shls_slice = (ish0, ish1, jsh0, jsh1, ksh0, ksh1, lsh0, lsh1)
i0, i1, j0, j1, k0, k1, l0, l1 = [ao_loc[x] for x in shls_slice]
dms = [dm[j0:j1,k0:k1],
dm[j0:j1,l0:l1],
dm[i0:i1,k0:k1],
dm[i0:i1,l0:l1],
dm[l0:l1,i0:i1],
dm[l0:l1,j0:j1],
dm[k0:k1,i0:i1],
dm[k0:k1,j0:j1]]
scripts = ['ijkl,jk->il',
'ijkl,jl->ik',
'ijkl,ik->jl',
'ijkl,il->jk',
'ijkl,li->kj',
'ijkl,lj->ki',
'ijkl,ki->lj',
'ijkl,kj->li']
kparts = jk.get_jk(mol, dms, scripts, shls_slice=shls_slice)
vk[i0:i1,l0:l1] += kparts[0]
vk[i0:i1,k0:k1] += kparts[1]
vk[j0:j1,l0:l1] += kparts[2]
vk[j0:j1,k0:k1] += kparts[3]
vk[k0:k1,j0:j1] += kparts[4]
vk[k0:k1,i0:i1] += kparts[5]
vk[l0:l1,j0:j1] += kparts[6]
vk[l0:l1,i0:i1] += kparts[7]
# ip > jp, ip > kp, kp == lp
lsh0, lsh1 = ksh0, ksh1
shls_slice = (ish0, ish1, jsh0, jsh1, ksh0, ksh1, lsh0, lsh1)
i0, i1, j0, j1, k0, k1, l0, l1 = [ao_loc[x] for x in shls_slice]
dms = [dm[j0:j1,k0:k1],
dm[i0:i1,k0:k1],
dm[k0:k1,j0:j1],
dm[k0:k1,i0:i1]]
scripts = ['ijkl,jk->il',
'ijkl,ik->jl',
'ijkl,kj->li',
'ijkl,ki->lj']
kparts = jk.get_jk(mol, dms, scripts, shls_slice=shls_slice)
vk[i0:i1,l0:l1] += kparts[0]
vk[j0:j1,l0:l1] += kparts[1]
vk[l0:l1,i0:i1] += kparts[2]
vk[l0:l1,j0:j1] += kparts[3]
# ip == kp and ip > jp and kp > lp
kp, ksh0, ksh1 = ip, ish0, ish1
for lp, (lsh0, lsh1) in enumerate(bas_groups[:kp]):
shls_slice = (ish0, ish1, jsh0, jsh1, ksh0, ksh1, lsh0, lsh1)
i0, i1, j0, j1, k0, k1, l0, l1 = [ao_loc[x] for x in shls_slice]
dms = [dm[j0:j1,k0:k1],
dm[i0:i1,k0:k1],
dm[j0:j1,l0:l1],
dm[i0:i1,l0:l1]]
scripts = ['ijkl,jk->il',
'ijkl,ik->jl',
'ijkl,jl->ik',
'ijkl,il->jk']
kparts = jk.get_jk(mol, dms, scripts, shls_slice=shls_slice)
vk[i0:i1,l0:l1] += kparts[0]
vk[j0:j1,l0:l1] += kparts[1]
vk[i0:i1,k0:k1] += kparts[2]
vk[j0:j1,k0:k1] += kparts[3]
# ip == jp and ip >= kp
jsh0, jsh1 = ish0, ish1
for kp, (ksh0, ksh1) in enumerate(bas_groups[:ip+1]):
for lp, (lsh0, lsh1) in enumerate(bas_groups[:kp]):
shls_slice = (ish0, ish1, jsh0, jsh1, ksh0, ksh1, lsh0, lsh1)
i0, i1, j0, j1, k0, k1, l0, l1 = [ao_loc[x] for x in shls_slice]
dms = [dm[j0:j1,k0:k1],
dm[j0:j1,l0:l1],
dm[k0:k1,j0:j1],
dm[l0:l1,j0:j1]]
scripts = ['ijkl,jk->il',
'ijkl,jl->ik',
'ijkl,kj->li',
'ijkl,lj->ki']
kparts = jk.get_jk(mol, dms, scripts, shls_slice=shls_slice)
vk[i0:i1,l0:l1] += kparts[0]
vk[i0:i1,k0:k1] += kparts[1]
vk[l0:l1,i0:i1] += kparts[2]
vk[k0:k1,i0:i1] += kparts[3]
# ip == jp and ip > kp and kp == lp
for kp, (ksh0, ksh1) in enumerate(bas_groups[:ip]):
lsh0, lsh1 = ksh0, ksh1
shls_slice = (ish0, ish1, jsh0, jsh1, ksh0, ksh1, lsh0, lsh1)
i0, i1, j0, j1, k0, k1, l0, l1 = [ao_loc[x] for x in shls_slice]
dms = [dm[j0:j1,k0:k1],
dm[l0:l1,i0:i1]]
scripts = ['ijkl,jk->il',
'ijkl,li->kj']
kparts = jk.get_jk(mol, dms, scripts, shls_slice=shls_slice)
vk[i0:i1,l0:l1] += kparts[0]
vk[k0:k1,j0:j1] += kparts[1]
# ip == jp == kp == lp
kp, ksh0, ksh1 = ip, ish0, ish1
lsh0, lsh1 = ksh0, ksh1
shls_slice = (ish0, ish1, jsh0, jsh1, ksh0, ksh1, lsh0, lsh1)
i0, i1, j0, j1, k0, k1, l0, l1 = [ao_loc[x] for x in shls_slice]
kparts = jk.get_jk(mol,
[dm[j0:j1,k0:k1]],
scripts=['ijkl,jk->il'],
shls_slice=shls_slice)
vk[i0:i1,l0:l1] += kparts[0]
return vk
def get_j(dfobj, dm, hermi=1, direct_scf_tol=1e-13):
from pyscf.scf import _vhf
from pyscf.scf import jk
from pyscf.df import addons
t0 = t1 = (time.clock(), time.time())
mol = dfobj.mol
if dfobj._vjopt is None:
dfobj.auxmol = auxmol = addons.make_auxmol(mol, dfobj.auxbasis)
opt = _vhf.VHFOpt(mol, 'int3c2e', 'CVHFnr3c2e_schwarz_cond')
opt.direct_scf_tol = direct_scf_tol
# q_cond part 1: the regular int2e (ij|ij) for mol's basis
opt.init_cvhf_direct(mol, 'int2e', 'CVHFsetnr_direct_scf')
mol_q_cond = lib.frompointer(opt._this.contents.q_cond, mol.nbas**2)
# Update q_cond to include the 2e-integrals (auxmol|auxmol)
j2c = auxmol.intor('int2c2e', hermi=1)
j2c_diag = numpy.sqrt(abs(j2c.diagonal()))
aux_loc = auxmol.ao_loc
aux_q_cond = [
j2c_diag[i0:i1].max() for i0, i1 in zip(aux_loc[:-1], aux_loc[1:])
]
q_cond = numpy.hstack((mol_q_cond, aux_q_cond))
fsetqcond = _vhf.libcvhf.CVHFset_q_cond
fsetqcond(opt._this, q_cond.ctypes.data_as(ctypes.c_void_p),
ctypes.c_int(q_cond.size))
try:
opt.j2c = j2c = scipy.linalg.cho_factor(j2c, lower=True)
opt.j2c_type = 'cd'
except scipy.linalg.LinAlgError:
opt.j2c = j2c
opt.j2c_type = 'regular'
# jk.get_jk function supports 4-index integrals. Use bas_placeholder
# (l=0, nctr=1, 1 function) to hold the last index.
bas_placeholder = numpy.array([0, 0, 1, 1, 0, 0, 0, 0],
dtype=numpy.int32)
fakemol = mol + auxmol
fakemol._bas = numpy.vstack((fakemol._bas, bas_placeholder))
opt.fakemol = fakemol
dfobj._vjopt = opt
t1 = logger.timer_debug1(dfobj, 'df-vj init_direct_scf', *t1)
opt = dfobj._vjopt
fakemol = opt.fakemol
dm = numpy.asarray(dm, order='C')
dm_shape = dm.shape
nao = dm_shape[-1]
dm = dm.reshape(-1, nao, nao)
n_dm = dm.shape[0]
# First compute the density in auxiliary basis
# j3c = fauxe2(mol, auxmol)
# jaux = numpy.einsum('ijk,ji->k', j3c, dm)
# rho = numpy.linalg.solve(auxmol.intor('int2c2e'), jaux)
nbas = mol.nbas
nbas1 = mol.nbas + dfobj.auxmol.nbas
shls_slice = (0, nbas, 0, nbas, nbas, nbas1, nbas1, nbas1 + 1)
with lib.temporary_env(opt,
prescreen='CVHFnr3c2e_vj_pass1_prescreen',
_dmcondname='CVHFsetnr_direct_scf_dm'):
jaux = jk.get_jk(fakemol,
dm, ['ijkl,ji->kl'] * n_dm,
'int3c2e',
aosym='s2ij',
hermi=0,
shls_slice=shls_slice,
vhfopt=opt)
# remove the index corresponding to bas_placeholder
jaux = numpy.array(jaux)[:, :, 0]
t1 = logger.timer_debug1(dfobj, 'df-vj pass 1', *t1)
if opt.j2c_type == 'cd':
rho = scipy.linalg.cho_solve(opt.j2c, jaux.T)
else:
rho = scipy.linalg.solve(opt.j2c, jaux.T)
# transform rho to shape (:,1,naux), to adapt to 3c2e integrals (ij|k)
rho = rho.T[:, numpy.newaxis, :]
t1 = logger.timer_debug1(dfobj, 'df-vj solve ', *t1)
# Next compute the Coulomb matrix
# j3c = fauxe2(mol, auxmol)
# vj = numpy.einsum('ijk,k->ij', j3c, rho)
with lib.temporary_env(opt,
prescreen='CVHFnr3c2e_vj_pass2_prescreen',
_dmcondname=None):
# CVHFnr3c2e_vj_pass2_prescreen requires custom dm_cond
aux_loc = dfobj.auxmol.ao_loc
dm_cond = [
abs(rho[:, :, i0:i1]).max()
for i0, i1 in zip(aux_loc[:-1], aux_loc[1:])
]
dm_cond = numpy.array(dm_cond)
fsetcond = _vhf.libcvhf.CVHFset_dm_cond
fsetcond(opt._this, dm_cond.ctypes.data_as(ctypes.c_void_p),
ctypes.c_int(dm_cond.size))
vj = jk.get_jk(fakemol,
rho, ['ijkl,lk->ij'] * n_dm,
'int3c2e',
aosym='s2ij',
hermi=1,
shls_slice=shls_slice,
vhfopt=opt)
t1 = logger.timer_debug1(dfobj, 'df-vj pass 2', *t1)
logger.timer(dfobj, 'df-vj', *t0)
return numpy.asarray(vj).reshape(dm_shape)
def _eval_jk(mf, dm, hermi, gen_jobs):
cpu0 = (logger.process_clock(), logger.perf_counter())
mol = mf.mol
ao_loc = mol.ao_loc_nr()
nao = ao_loc[-1]
bas_groups = _partition_bas(mol)
jobs = gen_jobs(len(bas_groups), hermi)
njobs = len(jobs)
logger.debug1(mf, 'njobs %d', njobs)
# Each job has multiple recipes.
n_recipes = len(jobs[0][1:])
dm = numpy.asarray(dm).reshape(-1, nao, nao)
n_dm = dm.shape[0]
vk = numpy.zeros((n_recipes, n_dm, nao, nao))
if mf.opt is None:
vhfopt = mf.init_direct_scf(mol)
else:
vhfopt = mf.opt
# Assign the entire dm_cond to vhfopt.
# The prescreen function CVHFnrs8_prescreen will index q_cond and dm_cond
# over the entire basis. "set_dm" in function jk.get_jk/direct_bindm only
# creates a subblock of dm_cond which is not compatible with
# CVHFnrs8_prescreen.
vhfopt.set_dm(dm, mol._atm, mol._bas, mol._env)
# Then skip the "set_dm" initialization in function jk.get_jk/direct_bindm.
vhfopt._dmcondname = None
logger.timer_debug1(mf, 'get_jk initialization', *cpu0)
for job_id in mpi.work_stealing_partition(range(njobs)):
group_ids = jobs[job_id][0]
recipes = jobs[job_id][1:]
shls_slice = lib.flatten([bas_groups[i] for i in group_ids])
loc = ao_loc[shls_slice].reshape(4, 2)
dm_blks = []
for i_dm in range(n_dm):
for ir, recipe in enumerate(recipes):
for i, rec in enumerate(recipe):
p0, p1 = loc[rec[0]]
q0, q1 = loc[rec[1]]
dm_blks.append(dm[i_dm, p0:p1, q0:q1])
scripts = [
'ijkl,%s%s->%s%s' % tuple(['ijkl'[x] for x in rec])
for recipe in recipes for rec in recipe
] * n_dm
kparts = jk.get_jk(mol,
dm_blks,
scripts,
shls_slice=shls_slice,
vhfopt=vhfopt)
for i_dm in range(n_dm):
for ir, recipe in enumerate(recipes):
for i, rec in enumerate(recipe):
p0, p1 = loc[rec[2]]
q0, q1 = loc[rec[3]]
vk[ir, i_dm, p0:p1, q0:q1] += kparts[i]
# Pop the results of one recipe
kparts = kparts[i + 1:]
vk = mpi.reduce(vk)
if rank == 0:
if hermi:
for i in range(n_recipes):
for j in range(n_dm):
lib.hermi_triu(vk[i, j], hermi, inplace=True)
else:
# Zero out vk on workers. If reduce(get_jk()) is called twice,
# non-zero vk on workers can cause error.
vk[:] = 0
logger.timer(mf, 'get_jk', *cpu0)
return vk
def dia(magobj, gauge_orig=None):
mol = magobj.mol
mf = magobj._scf
mo_energy = magobj._scf.mo_energy
mo_coeff = magobj._scf.mo_coeff
mo_occ = magobj._scf.mo_occ
orboa = mo_coeff[0][:,mo_occ[0] > 0]
orbob = mo_coeff[1][:,mo_occ[1] > 0]
dm0a = lib.tag_array(orboa.dot(orboa.T), mo_coeff=mo_coeff[0], mo_occ=mo_occ[0])
dm0b = lib.tag_array(orbob.dot(orbob.T), mo_coeff=mo_coeff[1], mo_occ=mo_occ[1])
dm0 = dm0a + dm0b
dme0a = numpy.dot(orboa * mo_energy[0][mo_occ[0] > 0], orboa.T)
dme0b = numpy.dot(orbob * mo_energy[1][mo_occ[1] > 0], orbob.T)
dme0 = dme0a + dme0b
e2 = rhf_mag._get_dia_1e(magobj, gauge_orig, dm0, dme0)
if gauge_orig is not None:
return -e2
# Computing the 2nd order Vxc integrals from GIAO
grids = mf.grids
ni = mf._numint
xc_code = mf.xc
xctype = ni._xc_type(xc_code)
omega, alpha, hyb = ni.rsh_and_hybrid_coeff(xc_code, mol.spin)
make_rhoa, nset, nao = ni._gen_rho_evaluator(mol, dm0a, hermi=1)
make_rhob = ni._gen_rho_evaluator(mol, dm0b, hermi=1)[0]
ngrids = len(grids.weights)
mem_now = lib.current_memory()[0]
max_memory = max(2000, mf.max_memory*.9-mem_now)
BLKSIZE = numint.BLKSIZE
blksize = min(int(max_memory/12*1e6/8/nao/BLKSIZE)*BLKSIZE, ngrids)
vmata = numpy.zeros((3,3,nao,nao))
vmatb = numpy.zeros((3,3,nao,nao))
if xctype == 'LDA':
ao_deriv = 0
for ao, mask, weight, coords \
in ni.block_loop(mol, grids, nao, ao_deriv, max_memory,
blksize=blksize):
rho = (make_rhoa(0, ao, mask, 'LDA'),
make_rhob(0, ao, mask, 'LDA'))
vxc = ni.eval_xc(xc_code, rho, spin=1, deriv=1)[1]
vrho = vxc[0]
r_ao = numpy.einsum('pi,px->pxi', ao, coords)
aow = numpy.einsum('pxi,p,p->pxi', r_ao, weight, vrho[:,0])
vmata += lib.einsum('pxi,pyj->xyij', r_ao, aow)
aow = numpy.einsum('pxi,p,p->pxi', r_ao, weight, vrho[:,1])
vmatb += lib.einsum('pxi,pyj->xyij', r_ao, aow)
rho = vxc = vrho = aow = None
elif xctype == 'GGA':
ao_deriv = 1
for ao, mask, weight, coords \
in ni.block_loop(mol, grids, nao, ao_deriv, max_memory,
blksize=blksize):
rho = (make_rhoa(0, ao, mask, 'GGA'),
make_rhob(0, ao, mask, 'GGA'))
vxc = ni.eval_xc(xc_code, rho, spin=1, deriv=1)[1]
wva, wvb = numint._uks_gga_wv0(rho, vxc, weight)
# Computing \nabla (r * AO) = r * \nabla AO + [\nabla,r]_- * AO
r_ao = numpy.einsum('npi,px->npxi', ao, coords)
r_ao[1,:,0] += ao[0]
r_ao[2,:,1] += ao[0]
r_ao[3,:,2] += ao[0]
aow = numpy.einsum('npxi,np->pxi', r_ao, wva)
vmata += lib.einsum('pxi,pyj->xyij', r_ao[0], aow)
aow = numpy.einsum('npxi,np->pxi', r_ao, wvb)
vmata += lib.einsum('pxi,pyj->xyij', r_ao[0], aow)
rho = vxc = vrho = vsigma = wv = aow = None
vmata = vmata + vmata.transpose(0,1,3,2)
vmatb = vmatb + vmatb.transpose(0,1,3,2)
elif xctype == 'MGGA':
raise NotImplementedError('meta-GGA')
vmata = rks_mag._add_giao_phase(mol, vmata)
vmatb = rks_mag._add_giao_phase(mol, vmatb)
e2 += numpy.einsum('qp,xypq->xy', dm0a, vmata)
e2 += numpy.einsum('qp,xypq->xy', dm0b, vmatb)
vmata = vmatb = None
e2 = e2.ravel()
# Handle the hybrid functional and the range-separated functional
if abs(hyb) > 1e-10:
vs = jk.get_jk(mol, [dm0, dm0a, dm0a, dm0b, dm0b],
['ijkl,ji->s2kl',
'ijkl,jk->s1il', 'ijkl,li->s1kj',
'ijkl,jk->s1il', 'ijkl,li->s1kj'],
'int2e_gg1', 's4', 9, hermi=1)
e2 += numpy.einsum('xpq,qp->x', vs[0], dm0)
e2 -= numpy.einsum('xpq,qp->x', vs[1], dm0a) * .5 * hyb
e2 -= numpy.einsum('xpq,qp->x', vs[2], dm0a) * .5 * hyb
e2 -= numpy.einsum('xpq,qp->x', vs[3], dm0b) * .5 * hyb
e2 -= numpy.einsum('xpq,qp->x', vs[4], dm0b) * .5 * hyb
vk = jk.get_jk(mol, [dm0a, dm0b], ['ijkl,jk->s1il', 'ijkl,jk->s1il'],
'int2e_g1g2', 'aa4', 9, hermi=0)
e2 -= numpy.einsum('xpq,qp->x', vk[0], dm0a) * hyb
e2 -= numpy.einsum('xpq,qp->x', vk[1], dm0b) * hyb
if abs(omega) > 1e-10:
with mol.with_range_coulomb(omega):
vs = jk.get_jk(mol, [dm0a, dm0a, dm0b, dm0b],
['ijkl,jk->s1il', 'ijkl,li->s1kj',
'ijkl,jk->s1il', 'ijkl,li->s1kj'],
'int2e_gg1', 's4', 9, hermi=1)
e2 -= numpy.einsum('xpq,qp->x', vs[0], dm0a) * .5 * (alpha-hyb)
e2 -= numpy.einsum('xpq,qp->x', vs[1], dm0a) * .5 * (alpha-hyb)
e2 -= numpy.einsum('xpq,qp->x', vs[2], dm0b) * .5 * (alpha-hyb)
e2 -= numpy.einsum('xpq,qp->x', vs[3], dm0b) * .5 * (alpha-hyb)
vk = jk.get_jk(mol, [dm0a, dm0b], ['ijkl,jk->s1il', 'ijkl,jk->s1il'],
'int2e_g1g2', 'aa4', 9, hermi=0)
e2 -= numpy.einsum('xpq,qp->x', vk[0], dm0a) * (alpha-hyb)
e2 -= numpy.einsum('xpq,qp->x', vk[1], dm0b) * (alpha-hyb)
else:
vj = jk.get_jk(mol, dm0, 'ijkl,ji->s2kl',
'int2e_gg1', 's4', 9, hermi=1)
e2 += numpy.einsum('xpq,qp->x', vj, dm0)
return -e2.reshape(3, 3)