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