Source code for pyscf.pbc.df.incore

#!/usr/bin/env python
# Copyright 2014-2021 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#

import ctypes
import numpy as np
from pyscf import lib
from pyscf import gto
from pyscf.lib import logger
import pyscf.df
from pyscf.scf import _vhf
from pyscf.pbc.df import ft_ao
from pyscf.pbc import gto as pbcgto
from pyscf.pbc.gto.pseudo import pp_int
from pyscf.pbc.lib.kpts_helper import is_zero, unique
from pyscf.pbc.tools import k2gamma
from pyscf.pbc.tools import pbc as pbctools
from pyscf import __config__

RCUT_THRESHOLD = getattr(__config__, 'pbc_scf_rsjk_rcut_threshold', 2.5)
KECUT_THRESHOLD = getattr(__config__, 'pbc_scf_rsjk_kecut_threshold', 10.0)
LOG_ADJUST = 32

libpbc = lib.load_library('libpbc')

[docs] def make_auxcell(cell, auxbasis=None): ''' See pyscf.df.addons.make_auxmol ''' auxcell = pyscf.df.addons.make_auxmol(cell, auxbasis) auxcell.rcut = pbcgto.estimate_rcut(auxcell) ke_cutoff = pbcgto.estimate_ke_cutoff(auxcell) a = cell.lattice_vectors() mesh = pbctools.cutoff_to_mesh(a, ke_cutoff) dimension = cell.dimension if dimension < 2 or (dimension == 2 and cell.low_dim_ft_type == 'inf_vacuum'): mesh[dimension:] = cell.mesh[dimension:] auxcell.mesh = mesh return auxcell
make_auxmol = make_auxcell
[docs] def format_aux_basis(cell, auxbasis='weigend+etb'): '''For backward compatibility''' return make_auxcell(cell, auxbasis)
[docs] def aux_e2(cell, auxcell_or_auxbasis, intor='int3c2e', aosym='s1', comp=None, kptij_lst=np.zeros((1,2,3)), shls_slice=None, **kwargs): r'''3-center AO integrals (ij|L) with double lattice sum: \sum_{lm} (i[l]j[m]|L[0]), where L is the auxiliary basis. Returns: (nao_pair, naux) array ''' if isinstance(auxcell_or_auxbasis, gto.MoleBase): auxcell = auxcell_or_auxbasis else: assert isinstance(auxcell_or_auxbasis, str) auxcell = make_auxcell(cell, auxcell_or_auxbasis) # For some unkown reasons, the pre-decontracted basis 'is slower than ## Slighly decontract basis. The decontracted basis has better locality. ## The locality can be used in the lattice sum to reduce cost. # if shls_slice is None and cell.nao_nr() < 200: # cell, contr_coeff = pbcgto.cell._split_basis(cell) # else: # contr_coeff = None int3c = wrap_int3c(cell, auxcell, intor, aosym, comp, kptij_lst, **kwargs) out = int3c(shls_slice) # if contr_coeff is not None: # if aosym == 's2': # tmp = out.reshape(nkptij,comp,ni,ni,naux) # idx, idy = np.tril_indices(ni) # tmp[:,:,idy,idx] = out.conj() # tmp[:,:,idx,idy] = out # out, tmp = tmp, None # out = lib.einsum('kcpql,pi->kciql', out, contr_coeff) # out = lib.einsum('kciql,qj->kcijl', out, contr_coeff) # idx, idy = np.tril_indices(contr_coeff.shape[1]) # out = out[:,:,idx,idy] # else: # out = out.reshape(nkptij,comp,ni,nj,naux) # out = lib.einsum('kcpql,pi->kciql', out, contr_coeff) # out = lib.einsum('kciql,qj->kcijl', out, contr_coeff) # out = out.reshape(nkptij,comp,-1,naux) if len(kptij_lst) == 1: out = out[0] return out
[docs] def wrap_int3c(cell, auxcell, intor='int3c2e', aosym='s1', comp=1, kptij_lst=np.zeros((1,2,3)), cintopt=None, pbcopt=None): '''Generate a 3-center integral kernel which can be called repeatedly in the incore or outcore driver. The kernel function has a simple function signature f(shls_slice) ''' kpti = kptij_lst[:,0] kptj = kptij_lst[:,1] j_only = is_zero(kpti - kptj) if j_only: kpts = kpti nkpts = len(kpts) kptij_idx = np.arange(nkpts) else: kpts, _, kptij_idx = unique(np.vstack([kpti,kptj])) wherei = kptij_idx[:len(kpti)] wherej = kptij_idx[len(kpti):] nkpts = len(kpts) kptij_idx = wherei * nkpts + wherej dfbuilder = Int3cBuilder(cell, auxcell, kpts).build() int3c = dfbuilder.gen_int3c_kernel(intor, aosym, comp, j_only, kptij_idx, return_complex=True) return int3c
[docs] def fill_2c2e(cell, auxcell_or_auxbasis, intor='int2c2e', hermi=0, kpt=np.zeros(3)): '''2-center 2-electron AO integrals (L|ij), where L is the auxiliary basis. ''' if isinstance(auxcell_or_auxbasis, gto.MoleBase): auxcell = auxcell_or_auxbasis else: auxcell = make_auxcell(cell, auxcell_or_auxbasis) if hermi != 0: hermi = pyscf.lib.HERMITIAN # pbcopt use the value of AO-pair to prescreening PBC integrals in the lattice # summation. Pass NULL pointer to pbcopt to prevent the prescreening return auxcell.pbc_intor(intor, 1, hermi, kpt, pbcopt=lib.c_null_ptr())
[docs] class Int3cBuilder(lib.StreamObject): '''helper functions to compute 3-center integral tensor with double-lattice sum ''' _keys = set(( 'cell', 'auxcell', 'kpts', 'rs_cell', 'bvk_kmesh', 'supmol', 'ke_cutoff', 'direct_scf_tol', )) def __init__(self, cell, auxcell, kpts=None): self.cell = cell self.auxcell = auxcell self.stdout = cell.stdout self.verbose = cell.verbose self.max_memory = cell.max_memory if kpts is None: self.kpts = np.zeros((1, 3)) else: self.kpts = np.reshape(kpts, (-1, 3)) self.rs_cell = None # mesh to generate Born-von Karman supercell self.bvk_kmesh = None self.supmol = None self.ke_cutoff = None self.direct_scf_tol = None
[docs] def reset(self, cell=None): if cell is not None: self.cell = cell self.rs_cell = None self.supmol = None self.direct_scf_tol = None return self
[docs] def build(self): log = logger.new_logger(self) cell = self.cell kpts = self.kpts self.bvk_kmesh = kmesh = k2gamma.kpts_to_kmesh(cell, kpts) log.debug('kmesh for bvk-cell = %s', kmesh) self.rs_cell = rs_cell = ft_ao._RangeSeparatedCell.from_cell( cell, self.ke_cutoff, RCUT_THRESHOLD, verbose=log) # TODO: strip supmol basis based on the intor type. rcut = estimate_rcut(cell, self.auxcell) supmol = ft_ao.ExtendedMole.from_cell(rs_cell, kmesh, rcut.max(), log) self.supmol = supmol.strip_basis(rcut) log.debug('sup-mol nbas = %d cGTO = %d pGTO = %d', supmol.nbas, supmol.nao, supmol.npgto_nr()) return self
[docs] def gen_int3c_kernel(self, intor='int3c2e', aosym='s2', comp=None, j_only=False, reindex_k=None, rs_auxcell=None, auxcell=None, supmol=None, return_complex=False): '''Generate function to compute int3c2e with double lattice-sum rs_auxcell: range-separated auxcell for gdf/rsdf module reindex_k: an index array to sort the order of k-points in output ''' log = logger.new_logger(self) cput0 = logger.process_clock(), logger.perf_counter() if self.rs_cell is None: self.build() if auxcell is None: auxcell = self.auxcell if rs_auxcell is None: rs_auxcell = ft_ao._RangeSeparatedCell.from_cell(auxcell) elif not isinstance(rs_auxcell, ft_ao._RangeSeparatedCell): rs_auxcell = ft_ao._RangeSeparatedCell.from_cell(rs_auxcell) if supmol is None: supmol = self.supmol cell = self.cell kpts = self.kpts nkpts = len(kpts) bvk_ncells, rs_nbas, nimgs = supmol.bas_mask.shape intor, comp = gto.moleintor._get_intor_and_comp(cell._add_suffix(intor), comp) nbasp = cell.nbas # The number of shells in the primitive cell if self.direct_scf_tol is None: omega = supmol.omega aux_exp = np.hstack(auxcell.bas_exps()).min() cell_exp = np.hstack(cell.bas_exps()).min() if omega == 0: theta = 1./(1./cell_exp + 1./aux_exp) else: theta = 1./(1./cell_exp + 1./aux_exp + omega**-2) lattice_sum_factor = max(2*np.pi*cell.rcut/(cell.vol*theta), 1) cutoff = cell.precision / lattice_sum_factor**2 * .1 log.debug1('int3c_kernel integral omega=%g theta=%g cutoff=%g', omega, theta, cutoff) else: cutoff = self.direct_scf_tol log_cutoff = int(np.log(cutoff) * LOG_ADJUST) atm, bas, env = gto.conc_env(supmol._atm, supmol._bas, supmol._env, rs_auxcell._atm, rs_auxcell._bas, rs_auxcell._env) cell0_ao_loc = _conc_locs(cell.ao_loc, auxcell.ao_loc) seg_loc = _conc_locs(supmol.seg_loc, rs_auxcell.sh_loc) # mimic the lattice sum with range 1 aux_seg2sh = np.arange(rs_auxcell.nbas + 1) seg2sh = _conc_locs(supmol.seg2sh, aux_seg2sh) if 'ECP' in intor: # rs_auxcell is a placeholder only to represent the ecpbas. # Ensure the ECPBAS_OFFSET be consistent with the treatment in pbc.gto.ecp env[gto.AS_ECPBAS_OFFSET] = len(bas) bas = np.asarray(np.vstack([bas, cell._ecpbas]), dtype=np.int32) cintopt = _vhf.make_cintopt(atm, bas, env, intor) # sindex may not be accurate enough to screen ECP integral. # Add penalty 1e-2 to reduce the screening error log_cutoff = int(np.log(cutoff*1e-2) * LOG_ADJUST) else: cintopt = _vhf.make_cintopt(supmol._atm, supmol._bas, supmol._env, intor) sindex = self.get_q_cond(supmol) ovlp_mask = sindex > log_cutoff bvk_ovlp_mask = lib.condense('np.any', ovlp_mask, supmol.sh_loc) cell0_ovlp_mask = bvk_ovlp_mask.reshape( bvk_ncells, nbasp, bvk_ncells, nbasp).any(axis=2).any(axis=0) cell0_ovlp_mask = cell0_ovlp_mask.astype(np.int8) ovlp_mask = None # Estimate the buffer size required by PBCfill_nr3c functions cache_size = max(_get_cache_size(cell, intor), _get_cache_size(rs_auxcell, intor)) cell0_dims = cell0_ao_loc[1:] - cell0_ao_loc[:-1] dijk = cell0_dims[:nbasp].max()**2 * cell0_dims[nbasp:].max() * comp aosym = aosym[:2] gamma_point_only = is_zero(kpts) if gamma_point_only: assert nkpts == 1 fill = f'PBCfill_nr3c_g{aosym}' nkpts_ij = 1 cache_size += dijk elif nkpts == 1: fill = f'PBCfill_nr3c_nk1{aosym}' nkpts_ij = 1 cache_size += dijk * 3 elif j_only: fill = f'PBCfill_nr3c_k{aosym}' # sort kpts then reindex_k in sort_k can be skipped if reindex_k is not None: kpts = kpts[reindex_k] nkpts = len(kpts) nkpts_ij = nkpts cache_size = (dijk * bvk_ncells + dijk * nkpts * 2 + max(dijk * nkpts * 2, cache_size)) else: assert nkpts < 45000 fill = f'PBCfill_nr3c_kk{aosym}' nkpts_ij = nkpts * nkpts cache_size = (max(dijk * bvk_ncells**2 + cache_size, dijk * nkpts**2 * 2) + dijk * bvk_ncells * nkpts * 2) if aosym == 's2' and reindex_k is not None: kk_mask = np.zeros((nkpts*nkpts), dtype=bool) kk_mask[reindex_k] = True kk_mask = kk_mask.reshape(nkpts, nkpts) if not np.all(kk_mask == kk_mask.T): log.warn('aosym=s2 not found in required kpts pairs') expLk = np.exp(1j*np.dot(supmol.bvkmesh_Ls, kpts.T)) expLkR = np.asarray(expLk.real, order='C') expLkI = np.asarray(expLk.imag, order='C') expLk = None if reindex_k is None: reindex_k = np.arange(nkpts_ij, dtype=np.int32) else: reindex_k = np.asarray(reindex_k, dtype=np.int32) nkpts_ij = reindex_k.size drv = libpbc.PBCfill_nr3c_drv # is_pbcintor controls whether to use memory efficient functions # Only suppots int3c2e_sph, int3c2e_cart in current C library is_pbcintor = intor in ('int3c2e_sph', 'int3c2e_cart') or intor[:3] == 'ECP' if is_pbcintor and not intor.startswith('PBC'): intor = 'PBC' + intor log.debug1('is_pbcintor = %d, intor = %s', is_pbcintor, intor) log.timer('int3c kernel initialization', *cput0) def int3c(shls_slice=None, outR=None, outI=None): cput0 = logger.process_clock(), logger.perf_counter() if shls_slice is None: shls_slice = [0, nbasp, 0, nbasp, nbasp, len(cell0_dims)] else: ksh0 = nbasp + shls_slice[4] ksh1 = nbasp + shls_slice[5] shls_slice = list(shls_slice[:4]) + [ksh0, ksh1] i0, i1, j0, j1, k0, k1 = cell0_ao_loc[shls_slice] if aosym == 's1': nrow = (i1-i0)*(j1-j0) else: nrow = i1*(i1+1)//2 - i0*(i0+1)//2 if comp == 1: shape = (nkpts_ij, nrow, k1-k0) else: shape = (nkpts_ij, comp, nrow, k1-k0) # output has to be filled with zero first because certain integrals # may be skipped by fill_ints driver outR = np.ndarray(shape, buffer=outR) outR[:] = 0 if gamma_point_only: outI = np.zeros(0) else: outI = np.ndarray(shape, buffer=outI) outI[:] = 0 if sindex is None: sindex_ptr = lib.c_null_ptr() else: sindex_ptr = sindex.ctypes.data_as(ctypes.c_void_p) drv(getattr(libpbc, intor), getattr(libpbc, fill), ctypes.c_int(is_pbcintor), outR.ctypes.data_as(ctypes.c_void_p), outI.ctypes.data_as(ctypes.c_void_p), expLkR.ctypes.data_as(ctypes.c_void_p), expLkI.ctypes.data_as(ctypes.c_void_p), reindex_k.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(nkpts_ij), ctypes.c_int(bvk_ncells), ctypes.c_int(nimgs), ctypes.c_int(nkpts), ctypes.c_int(nbasp), ctypes.c_int(comp), seg_loc.ctypes.data_as(ctypes.c_void_p), seg2sh.ctypes.data_as(ctypes.c_void_p), cell0_ao_loc.ctypes.data_as(ctypes.c_void_p), (ctypes.c_int*6)(*shls_slice), cell0_ovlp_mask.ctypes.data_as(ctypes.c_void_p), sindex_ptr, ctypes.c_int(log_cutoff), cintopt, ctypes.c_int(cache_size), atm.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(supmol.natm), bas.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(supmol.nbas), env.ctypes.data_as(ctypes.c_void_p)) log.timer_debug1(f'pbc integral {intor}', *cput0) if return_complex: if gamma_point_only: return outR else: return outR + outI * 1j else: if gamma_point_only: return outR, None else: return outR, outI return int3c
[docs] def get_q_cond(self, supmol=None): '''Integral screening condition max(sqrt((ij|ij))) inside the supmol''' if supmol is None: supmol = self.supmol nbas = supmol.nbas sindex = np.empty((nbas,nbas), dtype=np.int16) libpbc.PBCVHFsetnr_sindex( sindex.ctypes.data_as(ctypes.c_void_p), 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)) return sindex
libpbc.GTOmax_cache_size.restype = ctypes.c_int def _get_cache_size(cell, intor): '''Cache size for libcint integrals. Cache size cannot be accurately estimated in function PBC_ft_bvk_drv ''' cache_size = libpbc.GTOmax_cache_size( getattr(libpbc, intor), (ctypes.c_int*2)(0, cell.nbas), ctypes.c_int(1), cell._atm.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(cell.natm), cell._bas.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(cell.nbas), cell._env.ctypes.data_as(ctypes.c_void_p)) return cache_size
[docs] def estimate_rcut(cell, auxcell, precision=None): '''Estimate rcut for 3c2e integrals''' if precision is None: precision = cell.precision cell_exps = np.array([e.min() for e in cell.bas_exps()]) aux_exps = np.array([e.min() for e in auxcell.bas_exps()]) if cell_exps.size == 0 or aux_exps.size == 0: return np.zeros(1) ls = cell._bas[:,gto.ANG_OF] cs = gto.gto_norm(ls, cell_exps) ai_idx = cell_exps.argmin() ak_idx = aux_exps.argmin() ai = cell_exps[ai_idx] aj = cell_exps ak = aux_exps[ak_idx] li = cell._bas[ai_idx,gto.ANG_OF] lj = ls lk = auxcell._bas[ak_idx,gto.ANG_OF] ci = cs[ai_idx] cj = cs # Note ck normalizes the auxiliary basis \int \chi_k dr to 1 ck = 1./(4*np.pi) / gto.gaussian_int(lk+2, ak) aij = ai + aj lij = li + lj l3 = lij + lk theta = 1./(1./aij + 1./ak) norm_ang = ((2*li+1)*(2*lj+1))**.5/(4*np.pi) c1 = ci * cj * ck * norm_ang sfac = aij*aj/(aij*aj + ai*theta) fl = 2 fac = 2**(li+1)*np.pi**3.5*c1 * theta**(l3-1.5) / aij**(lij+1.5) / ak**(lk+1.5) fac *= (1 + ai/aj)**lj * fl / precision r0 = cell.rcut r0 = (np.log(fac * r0 * (sfac*r0)**(l3-2) + 1.) / (sfac*theta))**.5 r0 = (np.log(fac * r0 * (sfac*r0)**(l3-2) + 1.) / (sfac*theta))**.5 return r0
def _conc_locs(ao_loc1, ao_loc2): '''auxiliary basis was appended to regular AO basis when calling int3c2e integrals. Composite loc combines locs from regular AO basis and auxiliary basis accordingly.''' comp_loc = np.append(ao_loc1[:-1], ao_loc1[-1] + ao_loc2) return np.asarray(comp_loc, dtype=np.int32)