Refactor and optimize varous df modules and bugfixes (pyscf#1168)

* Refactor and optimize various df modules and bugfixes

    * Refactor rsjk and gdf, mdf, rsdf, ft_ao, aft;
    * Fix bug for large kmesh;
    * Fix memory usage issue in pbc.df.incore;
    * Add new vhf driver nr_direct_ex_drv;
    * Fix r_direct;
    * Fix bug in GDF when using Cartesian basis;
    * Fix issue pyscf#1154 for mdf;
    * Tune guess_omega;
    * Improve lattice Ls;
    * Tune k2gamma.kpts_to_kmesh;
    * Improve function to decontract basis;

* Fix condense function

* Fix h5py compatibility issue and rsdf_builder memory footprint

* Bugfix for gdf_builder if exclude_dd_block=True

* Update tests

* Optimize memory footprint

* Bugfix for pbc.df.rsdf prescreen

* Update tests

* Custom mesh for pbc get_pp

* Update test_pbc_fill_int

* Update multigrid tests

* Fix mol.decontract_basis for partial decontraction

* Update tests and fix pvxp in pbc-x2c hcore

* Improve numerical stability

* updates due to changes in group_by_conj_pairs

* Update kmp2_stagger against the latest GDF implemantions

* Solve conflicts to master branch

* Numerical instability in tests

* Fix stdout for pbc KPoints class

* kpts wraparound convention updated for kpts_helper

pyscf/__init__.py

@@ -1,4 +1,4 @@
-# Copyright 2014-2020 The PySCF Developers. All Rights Reserved.
+# Copyright 2014-2022 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.
@@ -86,7 +86,7 @@
         sys.stderr.write('pyscf plugins found in \n%s\n'
                          'When PYTHONPATH is set, it is recommended to load '
                          'these plugins through the environment variable '
-                         'PYSCF_EXT_PATH' % '\n'.join(__path__[1:]))
+                         'PYSCF_EXT_PATH\n' % '\n'.join(__path__[1:]))
 
 from distutils.version import LooseVersion
 import numpy

pyscf/df/df.py

@@ -62,8 +62,8 @@ class DF(lib.StreamObject):
             'j3c'.
             The DF integral tensor :math:`V_{x,ij}` should be a 2D array in C
             (row-major) convention, where x corresponds to index of auxiliary
-            basis, and the combined index ij is the orbital pair index. The
-            hermitian symmetry is assumed for the combined ij index, ie
+            basis, and the composite index ij is the orbital pair index. The
+            hermitian symmetry is assumed for the composite ij index, ie
             the elements of :math:`V_{x,i,j}` with :math:`i\geq j` are existed
             in the DF integral tensor.  Thus the shape of DF integral tensor
             is (M,N*(N+1)/2), where M is the number of auxbasis functions and

pyscf/df/outcore.py

@@ -261,14 +261,14 @@ def load(row_slice):
     log.timer('AO->MO CD eri transformation', *time0)
     return erifile
 
-def _guess_shell_ranges(mol, buflen, aosym):
+def _guess_shell_ranges(mol, buflen, aosym, start=0, stop=None):
     from pyscf.ao2mo.outcore import balance_partition
     ao_loc = mol.ao_loc_nr()
     if 's2' in aosym:
-        return balance_partition(ao_loc*(ao_loc+1)//2, buflen)
+        return balance_partition(ao_loc*(ao_loc+1)//2, buflen, start, stop)
     else:
         nao = ao_loc[-1]
-        return balance_partition(ao_loc*nao, buflen)
+        return balance_partition(ao_loc*nao, buflen, start, stop)
 
 def _create_h5file(erifile, dataname):
     if h5py.is_hdf5(erifile):

pyscf/gto/ft_ao.py

@@ -1,5 +1,5 @@
 #!/usr/bin/env python
-# Copyright 2014-2020 The PySCF Developers. All Rights Reserved.
+# 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.
@@ -17,7 +17,7 @@
 #
 
 '''
-Analytic Fourier transformation for AO and AO-pair value
+Analytical Fourier transformation for AO and AO-pair value
 '''
 
 import ctypes
@@ -27,8 +27,6 @@
 from pyscf import gto
 from pyscf.gto.moleintor import libcgto
 
-# TODO: in C code, store complex data in two vectors for real and imag part
-
 #
 # \int mu*nu*exp(-ik*r) dr
 #
@@ -66,38 +64,40 @@ def ft_aopair(mol, Gv, shls_slice=None, aosym='s1', b=numpy.eye(3),
         p_gs = (ctypes.c_int*3)(*[len(x) for x in Gvbase])
 
     ao_loc = gto.moleintor.make_loc(mol._bas, intor)
-    ni = ao_loc[shls_slice[1]] - ao_loc[shls_slice[0]]
-    nj = ao_loc[shls_slice[3]] - ao_loc[shls_slice[2]]
 
     if aosym == 's1':
         if (shls_slice[:2] == shls_slice[2:4] and
             intor.startswith('GTO_ft_ovlp')):
             fill = getattr(libcgto, 'GTO_ft_fill_s1hermi')
         else:
             fill = getattr(libcgto, 'GTO_ft_fill_s1')
-        shape = (nGv,ni,nj,comp)
+        ni = ao_loc[shls_slice[1]] - ao_loc[shls_slice[0]]
+        nj = ao_loc[shls_slice[3]] - ao_loc[shls_slice[2]]
+        shape = (comp, ni, nj, nGv)
     else:
         fill = getattr(libcgto, 'GTO_ft_fill_s2')
         i0 = ao_loc[shls_slice[0]]
         i1 = ao_loc[shls_slice[1]]
         nij = i1*(i1+1)//2 - i0*(i0+1)//2
-        shape = (nGv,nij,comp)
-    mat = numpy.ndarray(shape, order='F', dtype=numpy.complex128, buffer=buf)
+        shape = (comp, nij, nGv)
+    mat = numpy.ndarray(shape, order='C', dtype=numpy.complex128, buffer=buf)
+    mat[:] = 0
 
     fn = libcgto.GTO_ft_fill_drv
     intor = getattr(libcgto, intor)
     eval_gz = getattr(libcgto, eval_gz)
+    phase = 0
 
     fn(intor, eval_gz, fill, mat.ctypes.data_as(ctypes.c_void_p),
        ctypes.c_int(comp), (ctypes.c_int*4)(*shls_slice),
-       ao_loc.ctypes.data_as(ctypes.c_void_p), ctypes.c_double(0),
+       ao_loc.ctypes.data_as(ctypes.c_void_p), ctypes.c_double(phase),
        GvT.ctypes.data_as(ctypes.c_void_p),
        p_b, p_gxyzT, p_gs, ctypes.c_int(nGv),
        mol._atm.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(mol.natm),
        mol._bas.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(mol.nbas),
        mol._env.ctypes.data_as(ctypes.c_void_p))
 
-    mat = numpy.rollaxis(mat, -1, 0)
+    mat = numpy.rollaxis(mat, -1, 1)
     if comp == 1:
         mat = mat[0]
     return mat
@@ -106,7 +106,10 @@ def ft_aopair(mol, Gv, shls_slice=None, aosym='s1', b=numpy.eye(3),
 # gxyz is the index for Gvbase
 def ft_ao(mol, Gv, shls_slice=None, b=numpy.eye(3),
           gxyz=None, Gvbase=None, verbose=None):
-    ''' FT transform AO
+    r'''Analytical FT transform AO
+    \int mu(r) exp(-ikr) dr^3
+
+    The output tensor has the shape [nGv, nao]
     '''
     if shls_slice is None:
         shls_slice = (0, mol.nbas)
@@ -149,42 +152,17 @@ def ft_ao(mol, Gv, shls_slice=None, b=numpy.eye(3),
     nao = ao_loc[mol.nbas]
     ao_loc = numpy.asarray(numpy.hstack((ao_loc, [nao+1])), dtype=numpy.int32)
     ni = ao_loc[shls_slice[1]] - ao_loc[shls_slice[0]]
-    mat = numpy.zeros((nGv,ni), order='F', dtype=numpy.complex128)
+    shape = (ni, nGv)
+    mat = numpy.zeros(shape, order='C', dtype=numpy.complex128)
+    phase = 0
 
     shls_slice = shls_slice + (mol.nbas, mol.nbas+1)
     fn(intor, eval_gz, fill, mat.ctypes.data_as(ctypes.c_void_p),
        ctypes.c_int(1), (ctypes.c_int*4)(*shls_slice),
-       ao_loc.ctypes.data_as(ctypes.c_void_p),
-       ctypes.c_double(0),
+       ao_loc.ctypes.data_as(ctypes.c_void_p), ctypes.c_double(phase),
        GvT.ctypes.data_as(ctypes.c_void_p),
        p_b, p_gxyzT, p_gs, ctypes.c_int(nGv),
        atm.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(len(atm)),
        bas.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(len(bas)),
        env.ctypes.data_as(ctypes.c_void_p))
-    return mat
-
-
-if __name__ == '__main__':
-    mol = gto.Mole()
-    mol.atom = '''C    1.3    .2       .3
-                  C     .1    .1      1.1
-                  '''
-    mol.basis = 'ccpvdz'
-    #mol.basis = {'C': [[0, (2.4, .1, .6), (1.0,.8, .4)], [1, (1.1, 1)]]}
-    #mol.basis = {'C': [[0, (2.4, 1)]]}
-    mol.unit = 'B'
-    mol.build(0,0)
-
-    L = 5.
-    n = 20
-    a = numpy.diag([L,L,L])
-    b = scipy.linalg.inv(a)
-    gs = [n,n,n]
-    gxrange = range(gs[0]+1)+range(-gs[0],0)
-    gyrange = range(gs[1]+1)+range(-gs[1],0)
-    gzrange = range(gs[2]+1)+range(-gs[2],0)
-    gxyz = lib.cartesian_prod((gxrange, gyrange, gzrange))
-    Gv = 2*numpy.pi * numpy.dot(gxyz, b)
-
-    print(numpy.linalg.norm(ft_aopair(mol, Gv, None, 's1', b, gxyz, gs)) - 63.0239113778)
-    print(numpy.linalg.norm(ft_ao(mol, Gv, None, b, gxyz, gs))-56.8273147065)
+    return mat.T

pyscf/gto/mole.py

@@ -1,5 +1,5 @@
 #!/usr/bin/env python
-# Copyright 2014-2020 The PySCF Developers. All Rights Reserved.
+# 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.
@@ -153,7 +153,7 @@ def gto_norm(l, expnt):
         #return math.sqrt(f)
         return 1/numpy.sqrt(gaussian_int(l*2+2, 2*expnt))
     else:
-        raise ValueError('l should be > 0')
+        raise ValueError('l should be >= 0')
 
 def cart2sph(l, c_tensor=None, normalized=None):
     '''
@@ -527,9 +527,9 @@ def uncontracted_basis(_basis):
 
 def to_uncontracted_cartesian_basis(mol):
     '''Decontract the basis of a Mole or a Cell.  Returns a Mole (Cell) object
-    with the uncontracted basis environment and a list of coefficients that
-    transform the uncontracted cartesian basis to the original basis.  Each
-    element in the list corresponds to one shell of the original Mole (Cell).
+    with uncontracted Cartesian basis and a list of coefficients that
+    transform the uncontracted basis to the original basis. Each element in
+    the coefficients list corresponds to one shell of the original Mole (Cell).
 
     Examples:
 
@@ -540,37 +540,109 @@ def to_uncontracted_cartesian_basis(mol):
     >>> abs(s-mol.intor('int1e_ovlp')).max()
     0.0
     '''
+    return decontract_basis(mol, to_cart=True)
+
+def decontract_basis(mol, atoms=None, to_cart=False):
+    '''Decontract the basis of a Mole or a Cell.  Returns a Mole (Cell) object
+    with the uncontracted basis environment and a list of coefficients that
+    transform the uncontracted basis to the original basis. Each element in
+    the coefficients list corresponds to one shell of the original Mole (Cell).
+
+    Kwargs:
+        atoms: list or str
+            Atoms on which the basis to be decontracted. By default, all basis
+            are decontracted
+        to_cart: bool
+            Decontract basis and transfer to Cartesian basis
+
+    Examples:
+
+    >>> mol = gto.M(atom='Ne', basis='ccpvdz')
+    >>> pmol, ctr_coeff = mol.decontract_basis()
+    >>> c = scipy.linalg.block_diag(*ctr_coeff)
+    >>> s = reduce(numpy.dot, (c.T, pmol.intor('int1e_ovlp'), c))
+    >>> abs(s-mol.intor('int1e_ovlp')).max()
+    0.0
+    '''
     import copy
+    pmol = copy.copy(mol)
+
+    # Some input basis may be segmented basis from a general contracted set.
+    # This may lead to duplicated pGTOs. First contract all basis to remove
+    # duplicated primitive functions.
+    bas_exps = mol.bas_exps()
+    def _to_full_contraction(mol, bas_idx):
+        es = numpy.hstack([bas_exps[i] for i in bas_idx])
+        cs = scipy.linalg.block_diag(*[mol._libcint_ctr_coeff(ib) for ib in bas_idx])
+
+        es, e_idx, rev_idx = numpy.unique(es.round(9), True, True)
+        cs_new = numpy.zeros((es.size, cs.shape[1]))
+        for i, j in enumerate(rev_idx):
+            cs_new[j] += cs[i]
+        return es[::-1], cs_new[::-1]
+
+    _bas = []
+    _env = mol._env.copy()
+    contr_coeff = []
+
     lmax = mol._bas[:,ANG_OF].max()
     if mol.cart:
         c2s = [numpy.eye((l+1)*(l+2)//2) for l in range(lmax+1)]
-    else:
+    elif to_cart:
         c2s = [cart2sph(l, normalized='sp') for l in range(lmax+1)]
+        pmol.cart = True
+    else:
+        c2s = [numpy.eye(l*2+1) for l in range(lmax+1)]
+
+    aoslices = mol.aoslice_by_atom()
+    for ia, (ib0, ib1) in enumerate(aoslices[:,:2]):
+        if atoms is not None:
+            if isinstance(atoms, str):
+                to_apply = ((atoms == mol.atom_pure_symbol(ia)) or
+                            (atoms == mol.atom_symbol(ia)))
+            elif isinstance(atoms, (tuple, list)):
+                to_apply = ((mol.atom_pure_symbol(ia) in atoms) or
+                            (mol.atom_symbol(ia) in atoms) or
+                            (ia in atoms))
+            else:
+                to_apply = True
+            if not to_apply:
+                for ib in range(ib0, ib1):
+                    l = mol.bas_angular(ib)
+                    nc = mol.bas_nctr(ib)
+                    c = numpy.einsum('pi,xm->pxim', numpy.eye(nc), c2s[l])
+                    contr_coeff.append(c.reshape(nc * c2s[l].shape[0], -1))
+                _bas.append(mol._bas[ib0:ib1])
+                continue
 
-    pmol = copy.copy(mol)
-    pmol.cart = True
-    _bas = []
-    _env = mol._env.copy()
-    contr_coeff = []
-    for ib in range(mol.nbas):
-        l = mol._bas[ib,ANG_OF]
-        ncart = (l+1)*(l+2)//2
-        es = mol.bas_exp(ib)
-        cs = mol._libcint_ctr_coeff(ib)
-        nprim = cs.shape[0]
-        norm = gto_norm(l, es)
-        c = numpy.einsum('pi,p,xm->pxim', cs, 1./norm, c2s[l])
-        contr_coeff.append(c.reshape(nprim*ncart,-1))
-
-        pexp = mol._bas[ib,PTR_EXP]
-        pc = mol._bas[ib,PTR_COEFF]
-        bs = numpy.empty((nprim,8), dtype=numpy.int32)
-        bs[:] = mol._bas[ib]
-        bs[:,NCTR_OF] = bs[:,NPRIM_OF] = 1
-        bs[:,PTR_EXP] = numpy.arange(pexp, pexp+nprim)
-        bs[:,PTR_COEFF] = numpy.arange(pc, pc+nprim)
-        _env[pc:pc+nprim] = norm
-        _bas.append(bs)
+        lmax = mol._bas[ib0:ib1,ANG_OF].max()
+        pexp = mol._bas[ib0,PTR_EXP]
+        for l in range(lmax+1):
+            bas_idx = ib0 + numpy.where(mol._bas[ib0:ib1,ANG_OF] == l)[0]
+            if len(bas_idx) == 0:
+                continue
+
+            mol_exps, b_coeff = _to_full_contraction(mol, bas_idx)
+            nprim, nc = b_coeff.shape
+            bs = numpy.zeros((nprim, BAS_SLOTS), dtype=numpy.int32)
+            bs[:,ATOM_OF] = ia
+            bs[:,ANG_OF ] = l
+            bs[:,NCTR_OF] = bs[:,NPRIM_OF] = 1
+            norm = gto_norm(l, mol_exps)
+            if atoms is None:
+                bs[:,PTR_EXP] = pexp + numpy.arange(nprim)
+                bs[:,PTR_COEFF] = pexp + numpy.arange(nprim, nprim*2)
+                _env[pexp:pexp+nprim] = mol_exps
+                _env[pexp+nprim:pexp+nprim*2] = norm
+                pexp += nprim * 2
+            else:
+                bs[:,PTR_EXP] = _env.size + numpy.arange(nprim)
+                bs[:,PTR_COEFF] = _env.size + numpy.arange(nprim, nprim*2)
+                _env = np.hstack([_env, mol_exps, norm])
+            _bas.append(bs)
+
+            c = numpy.einsum('pi,p,xm->pxim', b_coeff, 1./norm, c2s[l])
+            contr_coeff.append(c.reshape(nprim * c2s[l].shape[0], -1))
 
     pmol._bas = numpy.asarray(numpy.vstack(_bas), dtype=numpy.int32)
     pmol._env = _env
@@ -3226,6 +3298,15 @@ def bas_exp(self, bas_id):
         ptr = self._bas[bas_id,PTR_EXP]
         return self._env[ptr:ptr+nprim].copy()
 
+    def bas_exps(self):
+        '''exponents of all basis
+        return [mol.bas_exp(i) for i in range(self.nbas)]
+        '''
+        nprims = self._bas[:,NPRIM_OF]
+        pexps = self._bas[:,PTR_EXP]
+        exps = [self._env[i0:i1] for i0, i1 in zip(pexps, pexps + nprims)]
+        return exps
+
     def _libcint_ctr_coeff(self, bas_id):
         nprim = self.bas_nprim(bas_id)
         nctr = self.bas_nctr(bas_id)
@@ -3465,6 +3546,14 @@ def intor_by_shell(self, intor, shells, comp=None, grids=None):
     def get_enuc(self):
         return self.energy_nuc()
 
+    def get_ao_indices(self, bas_list, ao_loc=None):
+        '''
+        Generate (dis-continued) AO indices for basis specified in bas_list
+        '''
+        if ao_loc is None:
+            ao_loc = self.ao_loc
+        return lib.locs_to_indices(ao_loc, bas_list)
+
     sph_labels = spheric_labels = sph_labels
     cart_labels = cart_labels
     ao_labels = ao_labels
@@ -3486,6 +3575,7 @@ def search_shell_id(self, atm_id, l):
     get_overlap_cond = get_overlap_cond
 
     to_uncontracted_cartesian_basis = to_uncontracted_cartesian_basis
+    decontract_basis = decontract_basis
 
     __add__ = conc_mol