def test_condense(self):
locs = numpy.arange(5)
a = numpy.random.random((locs[-1],locs[-1])) - .5
self.assertTrue(numpy.allclose(a, lib.condense('sum', a, locs)))
self.assertTrue(numpy.allclose(a, lib.condense('max', a, locs)))
self.assertTrue(numpy.allclose(a, lib.condense('min', a, locs)))
self.assertTrue(numpy.allclose(abs(a), lib.condense('abssum', a, locs)))
self.assertTrue(numpy.allclose(abs(a), lib.condense('absmax', a, locs)))
self.assertTrue(numpy.allclose(abs(a), lib.condense('absmin', a, locs)))
self.assertTrue(numpy.allclose(abs(a), lib.condense('norm', a, locs)))
def test_condense(self):
locs = numpy.arange(5)
a = numpy.random.random((locs[-1], locs[-1])) - .5
self.assertTrue(numpy.allclose(a, lib.condense('sum', a, locs)))
self.assertTrue(numpy.allclose(a, lib.condense('max', a, locs)))
self.assertTrue(numpy.allclose(a, lib.condense('min', a, locs)))
self.assertTrue(numpy.allclose(abs(a), lib.condense('abssum', a,
locs)))
self.assertTrue(numpy.allclose(abs(a), lib.condense('absmax', a,
locs)))
self.assertTrue(numpy.allclose(abs(a), lib.condense('absmin', a,
locs)))
self.assertTrue(numpy.allclose(abs(a), lib.condense('norm', a, locs)))
def get_jk(self,
dm_kpts,
hermi=1,
kpts=None,
kpts_band=None,
with_j=True,
with_k=True,
omega=None,
exxdiv=None):
if omega is not None: # J/K for RSH functionals
# TODO: call AFTDF.get_jk function
raise NotImplementedError
# Does not support to specify arbitrary kpts
if kpts is not None and abs(kpts - self.kpts).max() > 1e-7:
raise RuntimeError('kpts error')
kpts = self.kpts
if kpts_band is not None:
raise NotImplementedError
cpu0 = (time.clock(), time.time())
if self.supmol is None:
self.build()
nkpts = kpts.shape[0]
vhfopt = self.vhfopt
supmol = self.supmol
bvkcell = self.bvkcell
phase = self.phase
cell = self.cell_rs
nao = cell.nao
orig_nao = self.cell.nao
# * dense_bvk_ao_loc are the AOs which appear in supmol (some basis
# are removed)
# * sparse_ao_loc has dimension (Nk,nbas), corresponding to the
# bvkcell with all basis
dense_bvk_ao_loc = bvkcell.ao_loc
sparse_ao_loc = nao * np.arange(nkpts)[:, None] + cell.ao_loc[:-1]
sparse_ao_loc = np.append(sparse_ao_loc.ravel(), nao * nkpts)
nbands = nkpts
if dm_kpts.ndim != 4:
dm = dm_kpts.reshape(-1, nkpts, orig_nao, orig_nao)
else:
dm = dm_kpts
n_dm = dm.shape[0]
rs_c_coeff = cell._contr_coeff
sc_dm = lib.einsum('nkij,pi,qj->nkpq', dm, rs_c_coeff, rs_c_coeff)
# Utilized symmetry sc_dm[R,S] = sc_dm[S-R] = sc_dm[(S-R)%N]
#:sc_dm = lib.einsum('Rk,nkuv,Sk->nRuSv', phase, sc_dm, phase.conj())
sc_dm = lib.einsum('k,Sk,nkuv->nSuv', phase[0], phase.conj(), sc_dm)
dm_translation = k2gamma.double_translation_indices(
self.bvk_kmesh).astype(np.int32)
dm_imag_max = abs(sc_dm.imag).max()
is_complex_dm = dm_imag_max > 1e-6
if is_complex_dm:
if dm_imag_max < 1e-2:
logger.warn(
self, 'DM in (BvK) cell has small imaginary part. '
'It may be a signal of symmetry broken in k-point symmetry'
)
sc_dm = np.vstack([sc_dm.real, sc_dm.imag])
else:
sc_dm = sc_dm.real
sc_dm = np.asarray(sc_dm.reshape(-1, nkpts, nao, nao), order='C')
n_sc_dm = sc_dm.shape[0]
dm_cond = [
lib.condense('NP_absmax', d, sparse_ao_loc,
sparse_ao_loc[:cell.nbas + 1]) for d in sc_dm
]
dm_cond = np.asarray(np.max(dm_cond, axis=0), order='C')
libpbc.CVHFset_dm_cond(vhfopt._this,
dm_cond.ctypes.data_as(ctypes.c_void_p),
dm_cond.size)
dm_cond = None
bvk_nbas = bvkcell.nbas
shls_slice = (0, cell.nbas, 0, bvk_nbas, 0, bvk_nbas, 0, bvk_nbas)
if hermi:
fdot_suffix = 's2kl'
else:
fdot_suffix = 's1'
if with_j and with_k:
fdot = 'PBCVHF_contract_jk_' + fdot_suffix
vs = np.zeros((2, n_sc_dm, nao, nkpts, nao))
elif with_j:
fdot = 'PBCVHF_contract_j_' + fdot_suffix
vs = np.zeros((1, n_sc_dm, nao, nkpts, nao))
else: # with_k
fdot = 'PBCVHF_contract_k_' + fdot_suffix
vs = np.zeros((1, n_sc_dm, nao, nkpts, nao))
drv = libpbc.PBCVHF_direct_drv
drv(getattr(libpbc, fdot), vs.ctypes.data_as(ctypes.c_void_p),
sc_dm.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(n_dm),
ctypes.c_int(nkpts), ctypes.c_int(nbands), ctypes.c_int(cell.nbas),
self.ovlp_mask.ctypes.data_as(ctypes.c_void_p),
self.bvk_cell_id.ctypes.data_as(ctypes.c_void_p),
self.cell0_shl_id.ctypes.data_as(ctypes.c_void_p),
supmol._images_loc.ctypes.data_as(ctypes.c_void_p),
(ctypes.c_int * 8)(*shls_slice),
dense_bvk_ao_loc.ctypes.data_as(ctypes.c_void_p),
dm_translation.ctypes.data_as(ctypes.c_void_p), vhfopt._cintopt,
vhfopt._this, supmol._atm.ctypes.data_as(ctypes.c_void_p),
ctypes.c_int(supmol.natm),
supmol._bas.ctypes.data_as(ctypes.c_void_p),
ctypes.c_int(supmol.nbas),
supmol._env.ctypes.data_as(ctypes.c_void_p))
if is_complex_dm:
vs = vs[:, :n_dm] + vs[:, n_dm:] * 1j
if with_j and with_k:
vj, vk = vs
elif with_j:
vj, vk = vs[0], None
else:
vj, vk = None, vs[0]
cpu1 = logger.timer(self, 'short range part vj and vk', *cpu0)
lr_c_coeff = self.lr_aft.cell._contr_coeff
lr_dm = lib.einsum('nkij,pi,qj->nkpq', dm, lr_c_coeff, lr_c_coeff)
# For rho product other than diffused-diffused block, construct LR
# parts in terms of full ERIs and SR ERIs
vj1, vk1 = self.lr_aft.get_jk(lr_dm,
hermi,
kpts,
kpts_band,
with_j,
with_k,
exxdiv=exxdiv)
cpu1 = logger.timer(self, 'AFT-vj and AFT-vk', *cpu1)
# expRk is almost the same to phase, except a normalization factor
expRk = np.exp(1j * np.dot(self.bvkmesh_Ls, kpts.T))
if with_j:
vj = lib.einsum('npRq,pi,qj,Rk->nkij', vj, rs_c_coeff, rs_c_coeff,
expRk)
vj += lib.einsum('nkpq,pi,qj->nkij', vj1, lr_c_coeff, lr_c_coeff)
if self.purify and kpts_band is None:
vj = _purify(vj, phase)
if gamma_point(kpts) and dm_kpts.dtype == np.double:
vj = vj.real
if hermi:
vj = (vj + vj.conj().transpose(0, 1, 3, 2)) * .5
vj = vj.reshape(dm_kpts.shape)
if with_k:
vk = lib.einsum('npRq,pi,qj,Rk->nkij', vk, rs_c_coeff, rs_c_coeff,
expRk)
vk += lib.einsum('nkpq,pi,qj->nkij', vk1, lr_c_coeff, lr_c_coeff)
if self.purify and kpts_band is None:
vk = _purify(vk, phase)
if gamma_point(kpts) and dm_kpts.dtype == np.double:
vk = vk.real
if hermi:
vk = (vk + vk.conj().transpose(0, 1, 3, 2)) * .5
vk = vk.reshape(dm_kpts.shape)
return vj, vk
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
