def test_multigrid_kroks(self):
mf = dft.KROKS(cell_he)
mf.xc = 'lda,'
nao = cell_he.nao
mo = dm_he
mo_occ = numpy.ones((2, nao))
dm1 = numpy.einsum('kpi,ki,kqi->kpq', mo, mo_occ, mo)
dm1 = lib.tag_array(numpy.array([dm1, dm1]),
mo_coeff=mo,
mo_occ=mo_occ * 2)
ref = mf.get_veff(cell_he, dm1, kpts=kpts)
out = multigrid.multigrid(mf).get_veff(cell_he, dm1, kpts=kpts)
self.assertAlmostEqual(float(abs(ref - out).max()), 0, 9)
self.assertAlmostEqual(abs(ref.exc - out.exc).max(), 0, 9)
self.assertAlmostEqual(abs(ref.ecoul - out.ecoul).max(), 0, 9)
def KROKS(cell, *args):
from pyscf.pbc import dft
return dft.KROKS(cell, *args)