def test_get_jk_high_cost(self):
kpts = cell1.make_kpts([3, 1, 1])
np.random.seed(1)
dm = (np.random.rand(len(kpts), cell1.nao, cell1.nao) +
np.random.rand(len(kpts), cell1.nao, cell1.nao) * 1j)
dm = dm + dm.transpose(0, 2, 1).conj()
kmesh = k2gamma.kpts_to_kmesh(cell1, kpts)
phase = k2gamma.get_phase(cell1, kpts, kmesh)[1]
dm = np.einsum('Rk,kuv,Sk->RSuv', phase.conj().T, dm, phase.T)
dm = np.einsum('Rk,RSuv,Sk->kuv', phase, dm.real, phase.conj())
mf = cell1.KRHF(kpts=kpts)
jref, kref = mf.get_jk(cell1, dm, kpts=kpts)
ej = np.einsum('kij,kji->', jref, dm)
ek = np.einsum('kij,kji->', kref, dm) * .5
jk_builder = rsjk.RangeSeparationJKBuilder(cell1, kpts)
jk_builder.omega = 0.5
vj, vk = jk_builder.get_jk(dm, kpts=kpts, exxdiv=mf.exxdiv)
self.assertAlmostEqual(abs(vj - jref).max(), 0, 7)
self.assertAlmostEqual(abs(vk - kref).max(), 0, 7)
vj, vk = jk_builder.get_jk(dm,
kpts=kpts,
exxdiv=mf.exxdiv,
with_k=False)
self.assertAlmostEqual(abs(vj - jref).max(), 0, 7)
vj, vk = jk_builder.get_jk(dm,
kpts=kpts,
exxdiv=mf.exxdiv,
with_j=False)
self.assertAlmostEqual(abs(vk - kref).max(), 0, 7)
vj, vk = jk_builder.get_jk(dm, hermi=0, kpts=kpts, exxdiv=mf.exxdiv)
self.assertAlmostEqual(abs(vj - jref).max(), 0, 7)
self.assertAlmostEqual(abs(vk - kref).max(), 0, 7)
vj, vk = jk_builder.get_jk(dm,
hermi=0,
kpts=kpts,
exxdiv=mf.exxdiv,
with_k=False)
self.assertAlmostEqual(abs(vj - jref).max(), 0, 7)
vj, vk = jk_builder.get_jk(dm,
hermi=0,
kpts=kpts,
exxdiv=mf.exxdiv,
with_j=False)
self.assertAlmostEqual(abs(vk - kref).max(), 0, 7)
def build(self, omega=None, direct_scf_tol=None):
cpu0 = (time.clock(), time.time())
cell = self.cell
kpts = self.kpts
k_scaled = cell.get_scaled_kpts(kpts).sum(axis=0)
k_mod_to_half = k_scaled * 2 - (k_scaled * 2).round(0)
if abs(k_mod_to_half).sum() > 1e-5:
raise NotImplementedError('k-points must be symmetryic')
if omega is not None:
self.omega = omega
if self.omega is None:
# Search a proper range-separation parameter omega that can balance the
# computational cost between the real space integrals and moment space
# integrals
self.omega, self.mesh, self.ke_cutoff = _guess_omega(
cell, kpts, self.mesh)
else:
self.ke_cutoff = aft.estimate_ke_cutoff_for_omega(cell, self.omega)
self.mesh = pbctools.cutoff_to_mesh(cell.lattice_vectors(),
self.ke_cutoff)
logger.info(self, 'omega = %.15g ke_cutoff = %s mesh = %s',
self.omega, self.ke_cutoff, self.mesh)
if direct_scf_tol is None:
direct_scf_tol = cell.precision**1.5
logger.debug(self, 'Set direct_scf_tol %g', direct_scf_tol)
self.cell_rs = cell_rs = _re_contract_cell(cell, self.ke_cutoff)
self.bvk_kmesh = kmesh = k2gamma.kpts_to_kmesh(cell_rs, kpts)
bvkcell, phase = k2gamma.get_phase(cell_rs, kpts, kmesh)
self.bvkmesh_Ls = Ks = k2gamma.translation_vectors_for_kmesh(
cell_rs, kmesh)
self.bvkcell = bvkcell
self.phase = phase
# Given ke_cutoff, eta corresponds to the most steep Gaussian basis
# of which the Coulomb integrals can be accurately computed in moment
# space.
eta = aft.estimate_eta_for_ke_cutoff(cell,
self.ke_cutoff,
precision=cell.precision)
# * Assuming the most steep function in smooth basis has exponent eta,
# with attenuation parameter omega, rcut_sr is the distance of which
# the value of attenuated Coulomb integrals of four shells |eta> is
# smaller than the required precision.
# * The attenuated coulomb integrals between four s-type Gaussians
# (2*a/pi)^{3/4}exp(-a*r^2) is
# (erfc(omega*a^0.5/(omega^2+a)^0.5*R) - erfc(a^0.5*R)) / R
# if two Gaussians on one center and the other two on another center
# and the distance between the two centers are R.
# * The attenuated coulomb integrals between two spherical charge
# distributions is
# ~(pi/eta)^3/2 (erfc(tau*(eta/2)^0.5*R) - erfc((eta/2)^0.5*R)) / R
# tau = omega/sqrt(omega^2 + eta/2)
# if the spherical charge distribution is the product of above s-type
# Gaussian with exponent eta and a very smooth function.
# When R is large, the attenuated Coulomb integral is
# ~= (pi/eta)^3/2 erfc(tau*(eta/2)^0.5*R) / R
# ~= pi/(tau*eta^2*R^2) exp(-tau^2*eta*R^2/2)
tau = self.omega / (self.omega**2 + eta / 2)**.5
rcut_sr = 10 # initial guess
rcut_sr = (-np.log(direct_scf_tol * tau * (eta * rcut_sr)**2 / np.pi) /
(tau**2 * eta / 2))**.5
logger.debug(self, 'eta = %g rcut_sr = %g', eta, rcut_sr)
# Ls is the translation vectors to mimic periodicity of a cell
Ls = bvkcell.get_lattice_Ls(rcut=cell.rcut + rcut_sr)
self.supmol_Ls = Ls = Ls[np.linalg.norm(Ls, axis=1).argsort()]
supmol = _make_extended_mole(cell_rs, Ls, Ks, self.omega,
direct_scf_tol)
self.supmol = supmol
nkpts = len(self.bvkmesh_Ls)
nbas = cell_rs.nbas
n_steep, n_local, n_diffused = cell_rs._nbas_each_set
n_compact = n_steep + n_local
bas_mask = supmol._bas_mask
self.bvk_bas_mask = bvk_bas_mask = bas_mask.any(axis=2)
# Some basis in bvk-cell are not presented in the supmol. They can be
# skipped when computing SR integrals
self.bvkcell._bas = bvkcell._bas[bvk_bas_mask.ravel()]
# Record the mapping between the dense bvkcell basis and the
# original sparse bvkcell basis
bvk_cell_idx = np.repeat(np.arange(nkpts)[:, None], nbas, axis=1)
self.bvk_cell_id = bvk_cell_idx[bvk_bas_mask].astype(np.int32)
cell0_shl_idx = np.repeat(np.arange(nbas)[None, :], nkpts, axis=0)
self.cell0_shl_id = cell0_shl_idx[bvk_bas_mask].astype(np.int32)
logger.timer_debug1(self, 'initializing supmol', *cpu0)
logger.info(self, 'sup-mol nbas = %d cGTO = %d pGTO = %d', supmol.nbas,
supmol.nao, supmol.npgto_nr())
supmol.omega = -self.omega # Set short range coulomb
with supmol.with_integral_screen(direct_scf_tol**2):
vhfopt = _vhf.VHFOpt(supmol,
'int2e_sph',
qcondname=libpbc.PBCVHFsetnr_direct_scf)
vhfopt.direct_scf_tol = direct_scf_tol
self.vhfopt = vhfopt
logger.timer(self, 'initializing vhfopt', *cpu0)
q_cond = vhfopt.get_q_cond((supmol.nbas, supmol.nbas))
idx = supmol._images_loc
bvk_q_cond = lib.condense('NP_absmax', q_cond, idx, idx)
ovlp_mask = bvk_q_cond > direct_scf_tol
# Remove diffused-diffused block
if n_diffused > 0:
diffused_mask = np.zeros_like(bvk_bas_mask)
diffused_mask[:, n_compact:] = True
diffused_mask = diffused_mask[bvk_bas_mask]
ovlp_mask[diffused_mask[:, None] & diffused_mask] = False
self.ovlp_mask = ovlp_mask.astype(np.int8)
# mute rcut_threshold, divide basis into two sets only
cell_lr_aft = _re_contract_cell(cell, self.ke_cutoff, -1, verbose=0)
self.lr_aft = lr_aft = _LongRangeAFT(cell_lr_aft, kpts, self.omega,
self.bvk_kmesh)
lr_aft.ke_cutoff = self.ke_cutoff
lr_aft.mesh = self.mesh
lr_aft.eta = eta
return self
cell.atom = '''
H 0. 0. 0.
H 0.5 0.3 0.4
'''
cell.basis = 'gth-dzvp'
cell.pseudo = 'gth-pade'
cell.a = np.eye(3) * 4.
cell.unit = 'B'
cell.build()
kmesh = [2, 2, 1]
kpts = cell.make_kpts(kmesh)
print("Transform k-point integrals to supercell integral")
scell, phase = k2gamma.get_phase(cell, kpts)
NR, Nk = phase.shape
nao = cell.nao
s_k = cell.pbc_intor('int1e_ovlp', kpts=kpts)
s = scell.pbc_intor('int1e_ovlp')
s1 = np.einsum('Rk,kij,Sk->RiSj', phase, s_k, phase.conj())
print(abs(s - s1.reshape(s.shape)).max())
s = scell.pbc_intor('int1e_ovlp').reshape(NR, nao, NR, nao)
s1 = np.einsum('Rk,RiSj,Sk->kij', phase.conj(), s, phase)
print(abs(s1 - s_k).max())
kmf = dft.KRKS(cell, kpts)
ekpt = kmf.run()
mf = k2gamma.k2gamma(kmf, kmesh)
0.4037767149 -0.4712295093
0.1187877657 -0.4058039291
''')
cell.verbose = 6
cells.append(cell)
for cell in cells:
kpts = cell.make_kpts([3, 1, 1])
mf = cell.KRHF(kpts=kpts) #.run()
#dm = mf.make_rdm1()
np.random.seed(1)
dm = (np.random.rand(len(kpts), cell.nao, cell.nao) +
np.random.rand(len(kpts), cell.nao, cell.nao) * 1j)
dm = dm + dm.transpose(0, 2, 1).conj()
kmesh = k2gamma.kpts_to_kmesh(cell, kpts)
phase = k2gamma.get_phase(cell, kpts, kmesh)[1]
dm = lib.einsum('Rk,kuv,Sk->RSuv', phase.conj().T, dm, phase.T)
dm = lib.einsum('Rk,RSuv,Sk->kuv', phase, dm.real, phase.conj())
jref, kref = mf.get_jk(cell, dm, kpts=kpts)
ej = np.einsum('kij,kji->', jref, dm)
ek = np.einsum('kij,kji->', kref, dm) * .5
jk_builder = RangeSeparationJKBuilder(cell, kpts)
jk_builder.build(omega=0.5)
#jk_builder.mesh = [6,6,6]
#print(jk_builder.omega, jk_builder.mesh)
vj, vk = jk_builder.get_jk(dm, kpts=kpts, exxdiv=mf.exxdiv)
print(abs(vj - jref).max())
print(abs(vk - kref).max())
print('ej_ref', ej, 'ek_ref', ek)
cell.atom = '''
H 0. 0. 0.
H 0.5 0.3 0.4
'''
cell.basis = 'gth-dzvp'
cell.pseudo = 'gth-pade'
cell.a = np.eye(3) * 4.
cell.unit='B'
cell.build()
kmesh = [2, 2, 1]
kpts = cell.make_kpts(kmesh)
print("Transform k-point integrals to supercell integral")
scell, phase = k2gamma.get_phase(cell, kpts)
NR, Nk = phase.shape
nao = cell.nao
s_k = cell.pbc_intor('int1e_ovlp', kpts=kpts)
s = scell.pbc_intor('int1e_ovlp')
s1 = np.einsum('Rk,kij,Sk->RiSj', phase, s_k, phase.conj())
print(abs(s-s1.reshape(s.shape)).max())
s = scell.pbc_intor('int1e_ovlp').reshape(NR,nao,NR,nao)
s1 = np.einsum('Rk,RiSj,Sk->kij', phase.conj(), s, phase)
print(abs(s1-s_k).max())
kmf = dft.KRKS(cell, kpts)
ekpt = kmf.run()
mf = k2gamma.k2gamma(kmf, kmesh)