pyscf 2.6.2__py3-none-macosx_11_0_arm64.whl → 2.7.0__py3-none-macosx_11_0_arm64.whl

pyscf/pbc/tdscf/uhf.py

@@ -16,14 +16,212 @@
 # Author: Qiming Sun <osirpt.sun@gmail.com>
 #
 
+import numpy as np
 from pyscf import lib
 from pyscf.tdscf import uhf
 from pyscf.pbc.tdscf import rhf as td_rhf
 from pyscf.pbc.tdscf.rhf import TDBase
 
 
+def get_ab(mf):
+    r'''A and B matrices for TDDFT response function.
+
+    A[i,a,j,b] = \delta_{ab}\delta_{ij}(E_a - E_i) + (ia||bj)
+    B[i,a,j,b] = (ia||jb)
+
+    Spin symmetry is considered in the returned A, B lists.  List A has three
+    items: (A_aaaa, A_aabb, A_bbbb). A_bbaa = A_aabb.transpose(2,3,0,1).
+    B has three items: (B_aaaa, B_aabb, B_bbbb).
+    B_bbaa = B_aabb.transpose(2,3,0,1).
+    '''
+    cell = mf.cell
+    nao = cell.nao_nr()
+    mo_energy = scf.addons.mo_energy_with_exxdiv_none(mf)
+    mo = np.asarray(mf.mo_coeff)
+    mo_occ = np.asarray(mf.mo_occ)
+    kpt = mf.kpt
+
+    occidx_a = np.where(mo_occ[0]==1)[0]
+    viridx_a = np.where(mo_occ[0]==0)[0]
+    occidx_b = np.where(mo_occ[1]==1)[0]
+    viridx_b = np.where(mo_occ[1]==0)[0]
+    orbo_a = mo[0][:,occidx_a]
+    orbv_a = mo[0][:,viridx_a]
+    orbo_b = mo[1][:,occidx_b]
+    orbv_b = mo[1][:,viridx_b]
+    nocc_a = orbo_a.shape[1]
+    nvir_a = orbv_a.shape[1]
+    nocc_b = orbo_b.shape[1]
+    nvir_b = orbv_b.shape[1]
+    mo_a = np.hstack((orbo_a,orbv_a))
+    mo_b = np.hstack((orbo_b,orbv_b))
+    nmo_a = nocc_a + nvir_a
+    nmo_b = nocc_b + nvir_b
+
+    e_ia_a = (mo_energy[0][viridx_a,None] - mo_energy[0][occidx_a]).T
+    e_ia_b = (mo_energy[1][viridx_b,None] - mo_energy[1][occidx_b]).T
+    a_aa = np.diag(e_ia_a.ravel()).reshape(nocc_a,nvir_a,nocc_a,nvir_a)
+    a_bb = np.diag(e_ia_b.ravel()).reshape(nocc_b,nvir_b,nocc_b,nvir_b)
+    a_ab = np.zeros((nocc_a,nvir_a,nocc_b,nvir_b))
+    b_aa = np.zeros_like(a_aa)
+    b_ab = np.zeros_like(a_ab)
+    b_bb = np.zeros_like(a_bb)
+    a = (a_aa, a_ab, a_bb)
+    b = (b_aa, b_ab, b_bb)
+
+    def add_hf_(a, b, hyb=1):
+        eri_aa = mf.with_df.ao2mo([orbo_a,mo_a,mo_a,mo_a], kpt, compact=False)
+        eri_ab = mf.with_df.ao2mo([orbo_a,mo_a,mo_b,mo_b], kpt, compact=False)
+        eri_bb = mf.with_df.ao2mo([orbo_b,mo_b,mo_b,mo_b], kpt, compact=False)
+        eri_aa = eri_aa.reshape(nocc_a,nmo_a,nmo_a,nmo_a)
+        eri_ab = eri_ab.reshape(nocc_a,nmo_a,nmo_b,nmo_b)
+        eri_bb = eri_bb.reshape(nocc_b,nmo_b,nmo_b,nmo_b)
+        a_aa, a_ab, a_bb = a
+        b_aa, b_ab, b_bb = b
+
+        a_aa += np.einsum('iabj->iajb', eri_aa[:nocc_a,nocc_a:,nocc_a:,:nocc_a])
+        a_aa -= np.einsum('ijba->iajb', eri_aa[:nocc_a,:nocc_a,nocc_a:,nocc_a:]) * hyb
+        b_aa += np.einsum('iajb->iajb', eri_aa[:nocc_a,nocc_a:,:nocc_a,nocc_a:])
+        b_aa -= np.einsum('jaib->iajb', eri_aa[:nocc_a,nocc_a:,:nocc_a,nocc_a:]) * hyb
+
+        a_bb += np.einsum('iabj->iajb', eri_bb[:nocc_b,nocc_b:,nocc_b:,:nocc_b])
+        a_bb -= np.einsum('ijba->iajb', eri_bb[:nocc_b,:nocc_b,nocc_b:,nocc_b:]) * hyb
+        b_bb += np.einsum('iajb->iajb', eri_bb[:nocc_b,nocc_b:,:nocc_b,nocc_b:])
+        b_bb -= np.einsum('jaib->iajb', eri_bb[:nocc_b,nocc_b:,:nocc_b,nocc_b:]) * hyb
+
+        a_ab += np.einsum('iabj->iajb', eri_ab[:nocc_a,nocc_a:,nocc_b:,:nocc_b])
+        b_ab += np.einsum('iajb->iajb', eri_ab[:nocc_a,nocc_a:,:nocc_b,nocc_b:])
+
+    if isinstance(mf, scf.hf.KohnShamDFT):
+        ni = mf._numint
+        omega, alpha, hyb = ni.rsh_and_hybrid_coeff(mf.xc, cell.spin)
+
+        add_hf_(a, b, hyb)
+        if omega != 0:  # For RSH
+            raise NotImplementedError
+
+        xctype = ni._xc_type(mf.xc)
+        dm0 = mf.make_rdm1(mo, mo_occ)
+        make_rho = ni._gen_rho_evaluator(cell, dm0, hermi=1, with_lapl=False)[0]
+        mem_now = lib.current_memory()[0]
+        max_memory = max(2000, mf.max_memory*.8-mem_now)
+
+        if xctype == 'LDA':
+            ao_deriv = 0
+            for ao, _, mask, weight, coords \
+                    in ni.block_loop(cell, mf.grids, nao, ao_deriv, kpt, None, max_memory):
+                rho0a = make_rho(0, ao, mask, xctype)
+                rho0b = make_rho(1, ao, mask, xctype)
+                rho = (rho0a, rho0b)
+                fxc = ni.eval_xc_eff(mf.xc, rho, deriv=2, xctype=xctype)[2]
+                wfxc = fxc[:,0,:,0] * weight
+
+                rho_o_a = lib.einsum('rp,pi->ri', ao, orbo_a)
+                rho_v_a = lib.einsum('rp,pi->ri', ao, orbv_a)
+                rho_o_b = lib.einsum('rp,pi->ri', ao, orbo_b)
+                rho_v_b = lib.einsum('rp,pi->ri', ao, orbv_b)
+                rho_ov_a = np.einsum('ri,ra->ria', rho_o_a, rho_v_a)
+                rho_ov_b = np.einsum('ri,ra->ria', rho_o_b, rho_v_b)
+                rho_vo_a = rho_ov_a.conj()
+                rho_vo_b = rho_ov_b.conj()
+                w_ov_aa = np.einsum('ria,r->ria', rho_ov_a, wfxc[0,0])
+                w_ov_ab = np.einsum('ria,r->ria', rho_ov_a, wfxc[0,1])
+                w_ov_bb = np.einsum('ria,r->ria', rho_ov_b, wfxc[1,1])
+
+                a_aa += lib.einsum('ria,rjb->iajb', w_ov_aa, rho_vo_a)
+                b_aa += lib.einsum('ria,rjb->iajb', w_ov_aa, rho_ov_a)
+
+                a_ab += lib.einsum('ria,rjb->iajb', w_ov_ab, rho_vo_b)
+                b_ab += lib.einsum('ria,rjb->iajb', w_ov_ab, rho_ov_b)
+
+                a_bb += lib.einsum('ria,rjb->iajb', w_ov_bb, rho_vo_b)
+                b_bb += lib.einsum('ria,rjb->iajb', w_ov_bb, rho_ov_b)
+
+        elif xctype == 'GGA':
+            ao_deriv = 1
+            for ao, _, mask, weight, coords \
+                    in ni.block_loop(cell, mf.grids, nao, ao_deriv, kpt, None, max_memory):
+                rho0a = make_rho(0, ao, mask, xctype)
+                rho0b = make_rho(1, ao, mask, xctype)
+                rho = (rho0a, rho0b)
+                fxc = ni.eval_xc_eff(mf.xc, rho, deriv=2, xctype=xctype)[2]
+                wfxc = fxc * weight
+                rho_o_a = lib.einsum('xrp,pi->xri', ao, orbo_a)
+                rho_v_a = lib.einsum('xrp,pi->xri', ao, orbv_a)
+                rho_o_b = lib.einsum('xrp,pi->xri', ao, orbo_b)
+                rho_v_b = lib.einsum('xrp,pi->xri', ao, orbv_b)
+                rho_ov_a = np.einsum('xri,ra->xria', rho_o_a, rho_v_a[0])
+                rho_ov_b = np.einsum('xri,ra->xria', rho_o_b, rho_v_b[0])
+                rho_ov_a[1:4] += np.einsum('ri,xra->xria', rho_o_a[0], rho_v_a[1:4])
+                rho_ov_b[1:4] += np.einsum('ri,xra->xria', rho_o_b[0], rho_v_b[1:4])
+                rho_vo_a = rho_ov_a.conj()
+                rho_vo_b = rho_ov_b.conj()
+                w_ov_aa = np.einsum('xyr,xria->yria', wfxc[0,:,0], rho_ov_a)
+                w_ov_ab = np.einsum('xyr,xria->yria', wfxc[0,:,1], rho_ov_a)
+                w_ov_bb = np.einsum('xyr,xria->yria', wfxc[1,:,1], rho_ov_b)
+
+                a_aa += lib.einsum('xria,xrjb->iajb', w_ov_aa, rho_vo_a)
+                b_aa += lib.einsum('xria,xrjb->iajb', w_ov_aa, rho_ov_a)
+
+                a_ab += lib.einsum('xria,xrjb->iajb', w_ov_ab, rho_vo_b)
+                b_ab += lib.einsum('xria,xrjb->iajb', w_ov_ab, rho_ov_b)
+
+                a_bb += lib.einsum('xria,xrjb->iajb', w_ov_bb, rho_vo_b)
+                b_bb += lib.einsum('xria,xrjb->iajb', w_ov_bb, rho_ov_b)
+
+        elif xctype == 'HF':
+            pass
+
+        elif xctype == 'NLC':
+            raise NotImplementedError('NLC')
+
+        elif xctype == 'MGGA':
+            ao_deriv = 1
+            for ao, _, mask, weight, coords \
+                    in ni.block_loop(cell, mf.grids, nao, ao_deriv, kpt, None, max_memory):
+                rho0a = make_rho(0, ao, mask, xctype)
+                rho0b = make_rho(1, ao, mask, xctype)
+                rho = (rho0a, rho0b)
+                fxc = ni.eval_xc_eff(mf.xc, rho, deriv=2, xctype=xctype)[2]
+                wfxc = fxc * weight
+                rho_oa = lib.einsum('xrp,pi->xri', ao, orbo_a)
+                rho_ob = lib.einsum('xrp,pi->xri', ao, orbo_b)
+                rho_va = lib.einsum('xrp,pi->xri', ao, orbv_a)
+                rho_vb = lib.einsum('xrp,pi->xri', ao, orbv_b)
+                rho_ov_a = np.einsum('xri,ra->xria', rho_oa, rho_va[0])
+                rho_ov_b = np.einsum('xri,ra->xria', rho_ob, rho_vb[0])
+                rho_ov_a[1:4] += np.einsum('ri,xra->xria', rho_oa[0], rho_va[1:4])
+                rho_ov_b[1:4] += np.einsum('ri,xra->xria', rho_ob[0], rho_vb[1:4])
+                tau_ov_a = np.einsum('xri,xra->ria', rho_oa[1:4], rho_va[1:4]) * .5
+                tau_ov_b = np.einsum('xri,xra->ria', rho_ob[1:4], rho_vb[1:4]) * .5
+                rho_ov_a = np.vstack([rho_ov_a, tau_ov_a[np.newaxis]])
+                rho_ov_b = np.vstack([rho_ov_b, tau_ov_b[np.newaxis]])
+                rho_vo_a = rho_ov_a.conj()
+                rho_vo_b = rho_ov_b.conj()
+                w_ov_aa = np.einsum('xyr,xria->yria', wfxc[0,:,0], rho_ov_a)
+                w_ov_ab = np.einsum('xyr,xria->yria', wfxc[0,:,1], rho_ov_a)
+                w_ov_bb = np.einsum('xyr,xria->yria', wfxc[1,:,1], rho_ov_b)
+
+                a_aa += lib.einsum('xria,xrjb->iajb', w_ov_aa, rho_vo_a)
+                b_aa += lib.einsum('xria,xrjb->iajb', w_ov_aa, rho_ov_a)
+
+                a_ab += lib.einsum('xria,xrjb->iajb', w_ov_ab, rho_vo_b)
+                b_ab += lib.einsum('xria,xrjb->iajb', w_ov_ab, rho_ov_b)
+
+                a_bb += lib.einsum('xria,xrjb->iajb', w_ov_bb, rho_vo_b)
+                b_bb += lib.einsum('xria,xrjb->iajb', w_ov_bb, rho_ov_b)
+
+    else:
+        add_hf_(a, b)
+
+    return a, b
+
 class TDA(TDBase):
 
+    def get_ab(self, mf=None):
+        if mf is None: mf = self._scf
+        return get_ab(mf)
+
     singlet = None
 
     init_guess = uhf.TDA.init_guess
@@ -34,7 +232,11 @@ class TDA(TDBase):
 CIS = TDA
 
 
-class TDHF(TDA):
+class TDHF(TDBase):
+
+    def get_ab(self, mf=None):
+        if mf is None: mf = self._scf
+        return get_ab(mf)
 
     singlet = None
 

pyscf/pbc/tdscf/uks.py

@@ -25,7 +25,8 @@ from pyscf.pbc.tdscf.uhf import TDHF as TDDFT
 RPA = TDUKS = TDDFT
 
 
-class CasidaTDDFT(TDA):
+class CasidaTDDFT(TDDFT):
+    init_guess = TDA.init_guess
     _gen_vind = uks.TDDFTNoHybrid.gen_vind
     gen_vind = TDA.gen_vind
     kernel = uks.TDDFTNoHybrid.kernel

pyscf/pbc/tools/k2gamma.py

@@ -257,15 +257,18 @@ def k2gamma(kmf, kmesh=None):
     return mf
 
 
-def to_supercell_ao_integrals(cell, kpts, ao_ints):
+def to_supercell_ao_integrals(cell, kpts, ao_ints, kmesh=None, force_real=True):
     '''Transform from the unitcell k-point AO integrals to the supercell
     gamma-point AO integrals.
     '''
-    scell, phase = get_phase(cell, kpts)
+    scell, phase = get_phase(cell, kpts, kmesh=kmesh)
     NR, Nk = phase.shape
     nao = cell.nao
     scell_ints = np.einsum('Rk,kij,Sk->RiSj', phase, ao_ints, phase.conj())
-    return scell_ints.reshape(NR*nao,NR*nao).real
+    if force_real:
+        return scell_ints.reshape(NR*nao,NR*nao).real
+    else:
+        return scell_ints.reshape(NR*nao,NR*nao)
 
 
 def to_supercell_mo_integrals(kmf, mo_ints):
@@ -321,7 +324,7 @@ if __name__ == '__main__':
     mf = k2gamma(kmf, kmesh)
     c_g_ao = mf.mo_coeff
 
-    # The following is to check whether the MO is correctly coverted:
+    # The following is to check whether the MO is correctly converted:
 
     print("Supercell gamma MO in AO basis from conversion:")
     scell = tools.super_cell(cell, kmesh)

pyscf/pbc/tools/pbc.py

@@ -16,6 +16,7 @@
 import warnings
 import ctypes
 import numpy as np
+import scipy
 import scipy.linalg
 from pyscf import lib
 from pyscf.lib import logger
@@ -23,39 +24,55 @@ from pyscf.gto import ATM_SLOTS, BAS_SLOTS, ATOM_OF, PTR_COORD
 from pyscf.pbc.lib.kpts_helper import get_kconserv, get_kconserv3  # noqa
 from pyscf import __config__
 
-FFT_ENGINE = getattr(__config__, 'pbc_tools_pbc_fft_engine', 'BLAS')
+FFT_ENGINE = getattr(__config__, 'pbc_tools_pbc_fft_engine', 'NUMPY+BLAS')
 
 def _fftn_blas(f, mesh):
-    Gx = np.fft.fftfreq(mesh[0])
-    Gy = np.fft.fftfreq(mesh[1])
-    Gz = np.fft.fftfreq(mesh[2])
-    expRGx = np.exp(np.einsum('x,k->xk', -2j*np.pi*np.arange(mesh[0]), Gx))
-    expRGy = np.exp(np.einsum('x,k->xk', -2j*np.pi*np.arange(mesh[1]), Gy))
-    expRGz = np.exp(np.einsum('x,k->xk', -2j*np.pi*np.arange(mesh[2]), Gz))
-    out = np.empty(f.shape, dtype=np.complex128)
-    buf = np.empty(mesh, dtype=np.complex128)
-    for i, fi in enumerate(f):
-        buf[:] = fi.reshape(mesh)
-        g = lib.dot(buf.reshape(mesh[0],-1).T, expRGx, c=out[i].reshape(-1,mesh[0]))
-        g = lib.dot(g.reshape(mesh[1],-1).T, expRGy, c=buf.reshape(-1,mesh[1]))
-        g = lib.dot(g.reshape(mesh[2],-1).T, expRGz, c=out[i].reshape(-1,mesh[2]))
-    return out.reshape(-1, *mesh)
+    assert f.ndim == 4
+    mx, my, mz = mesh
+    expRGx = np.exp(-2j*np.pi*np.arange(mx)[:,None] * np.fft.fftfreq(mx))
+    expRGy = np.exp(-2j*np.pi*np.arange(my)[:,None] * np.fft.fftfreq(my))
+    expRGz = np.exp(-2j*np.pi*np.arange(mz)[:,None] * np.fft.fftfreq(mz))
+    blksize = max(int(1e5 / (mx * my * mz)), 8) * 4
+    n = f.shape[0]
+    out = np.empty((n, mx*my*mz), dtype=np.complex128)
+    buf = np.empty((blksize, mx*my*mz), dtype=np.complex128)
+    for i0, i1 in lib.prange(0, n, blksize):
+        ni = i1 - i0
+        buf1 = buf[:ni]
+        out1 = out[i0:i1]
+        g = lib.transpose(f[i0:i1].reshape(ni,-1), out=buf1.reshape(-1,ni))
+        g = lib.dot(g.reshape(mx,-1).T, expRGx, c=out1.reshape(-1,mx))
+        g = lib.dot(g.reshape(my,-1).T, expRGy, c=buf1.reshape(-1,my))
+        g = lib.dot(g.reshape(mz,-1).T, expRGz, c=out1.reshape(-1,mz))
+    return out.reshape(n, *mesh)
 
 def _ifftn_blas(g, mesh):
-    Gx = np.fft.fftfreq(mesh[0])
-    Gy = np.fft.fftfreq(mesh[1])
-    Gz = np.fft.fftfreq(mesh[2])
-    expRGx = np.exp(np.einsum('x,k->xk', 2j*np.pi*np.arange(mesh[0]), Gx))
-    expRGy = np.exp(np.einsum('x,k->xk', 2j*np.pi*np.arange(mesh[1]), Gy))
-    expRGz = np.exp(np.einsum('x,k->xk', 2j*np.pi*np.arange(mesh[2]), Gz))
-    out = np.empty(g.shape, dtype=np.complex128)
-    buf = np.empty(mesh, dtype=np.complex128)
-    for i, gi in enumerate(g):
-        buf[:] = gi.reshape(mesh)
-        f = lib.dot(buf.reshape(mesh[0],-1).T, expRGx, 1./mesh[0], c=out[i].reshape(-1,mesh[0]))
-        f = lib.dot(f.reshape(mesh[1],-1).T, expRGy, 1./mesh[1], c=buf.reshape(-1,mesh[1]))
-        f = lib.dot(f.reshape(mesh[2],-1).T, expRGz, 1./mesh[2], c=out[i].reshape(-1,mesh[2]))
-    return out.reshape(-1, *mesh)
+    assert g.ndim == 4
+    mx, my, mz = mesh
+    expRGx = np.exp(2j*np.pi*np.fft.fftfreq(mx)[:,None] * np.arange(mx))
+    expRGy = np.exp(2j*np.pi*np.fft.fftfreq(my)[:,None] * np.arange(my))
+    expRGz = np.exp(2j*np.pi*np.fft.fftfreq(mz)[:,None] * np.arange(mz))
+    blksize = max(int(1e5 / (mx * my * mz)), 8) * 4
+    n = g.shape[0]
+    out = np.empty((n, mx*my*mz), dtype=np.complex128)
+    buf = np.empty((blksize, mx*my*mz), dtype=np.complex128)
+    for i0, i1 in lib.prange(0, n, blksize):
+        ni = i1 - i0
+        buf1 = buf[:ni]
+        out1 = out[i0:i1]
+        f = lib.transpose(g[i0:i1].reshape(ni,-1), out=buf1.reshape(-1,ni))
+        f = lib.dot(f.reshape(mx,-1).T, expRGx, 1./mx, c=out1.reshape(-1,mx))
+        f = lib.dot(f.reshape(my,-1).T, expRGy, 1./my, c=buf1.reshape(-1,my))
+        f = lib.dot(f.reshape(mz,-1).T, expRGz, 1./mz, c=out1.reshape(-1,mz))
+    return out.reshape(n, *mesh)
+
+nproc = lib.num_threads()
+
+def _fftn_wrapper(a):  # noqa
+    return scipy.fft.fftn(a, axes=(1,2,3), workers=nproc)
+
+def _ifftn_wrapper(a):  # noqa
+    return scipy.fft.ifftn(a, axes=(1,2,3), workers=nproc)
 
 if FFT_ENGINE == 'FFTW':
     try:
@@ -88,60 +105,50 @@ if FFT_ENGINE == 'FFTW':
                ctypes.c_int(rank))
         return out
 
-    def _fftn_wrapper(a):
+    def _fftn_wrapper(a):  # noqa
         mesh = a.shape[1:]
         return _complex_fftn_fftw(a, mesh, 'fft')
-    def _ifftn_wrapper(a):
+    def _ifftn_wrapper(a):  # noqa
         mesh = a.shape[1:]
         return _complex_fftn_fftw(a, mesh, 'ifft')
 
 elif FFT_ENGINE == 'PYFFTW':
-    # pyfftw is slower than np.fft in most cases
+    # Note: pyfftw is likely slower than scipy.fft in multi-threading environments
     try:
         import pyfftw
+        pyfftw.config.PLANNER_EFFORT = 'FFTW_MEASURE'
         pyfftw.interfaces.cache.enable()
-        nproc = lib.num_threads()
-        def _fftn_wrapper(a):
+        def _fftn_wrapper(a):  # noqa
             return pyfftw.interfaces.numpy_fft.fftn(a, axes=(1,2,3), threads=nproc)
-        def _ifftn_wrapper(a):
+        def _ifftn_wrapper(a):  # noqa
             return pyfftw.interfaces.numpy_fft.ifftn(a, axes=(1,2,3), threads=nproc)
     except ImportError:
-        def _fftn_wrapper(a):
-            return np.fft.fftn(a, axes=(1,2,3))
-        def _ifftn_wrapper(a):
-            return np.fft.ifftn(a, axes=(1,2,3))
-
-elif FFT_ENGINE == 'NUMPY':
-    def _fftn_wrapper(a):
-        return np.fft.fftn(a, axes=(1,2,3))
-    def _ifftn_wrapper(a):
-        return np.fft.ifftn(a, axes=(1,2,3))
+        print('PyFFTW not installed. SciPy fft module will be used.')
 
 elif FFT_ENGINE == 'NUMPY+BLAS':
     _EXCLUDE = [17, 19, 23, 29, 31, 37, 41, 43, 47, 53, 59, 61, 67, 71, 73, 79,
                 83, 89, 97,101,103,107,109,113,127,131,137,139,149,151,157,163,
                 167,173,179,181,191,193,197,199,211,223,227,229,233,239,241,251,
                 257,263,269,271,277,281,283,293]
-    _EXCLUDE = set(_EXCLUDE + [n*2 for n in _EXCLUDE] + [n*3 for n in _EXCLUDE])
-    def _fftn_wrapper(a):
+    _EXCLUDE = set(_EXCLUDE + [n*2 for n in _EXCLUDE[:30]] + [n*3 for n in _EXCLUDE[:20]])
+    def _fftn_wrapper(a):  # noqa
         mesh = a.shape[1:]
         if mesh[0] in _EXCLUDE and mesh[1] in _EXCLUDE and mesh[2] in _EXCLUDE:
             return _fftn_blas(a, mesh)
         else:
-            return np.fft.fftn(a, axes=(1,2,3))
-    def _ifftn_wrapper(a):
+            return scipy.fft.fftn(a, axes=(1,2,3), workers=nproc)
+    def _ifftn_wrapper(a):  # noqa
         mesh = a.shape[1:]
         if mesh[0] in _EXCLUDE and mesh[1] in _EXCLUDE and mesh[2] in _EXCLUDE:
             return _ifftn_blas(a, mesh)
         else:
-            return np.fft.ifftn(a, axes=(1,2,3))
+            return scipy.fft.ifftn(a, axes=(1,2,3), workers=nproc)
 
-#?elif:  # 'FFTW+BLAS'
-else:  # 'BLAS'
-    def _fftn_wrapper(a):
+elif FFT_ENGINE == 'BLAS':
+    def _fftn_wrapper(a):  # noqa
         mesh = a.shape[1:]
         return _fftn_blas(a, mesh)
-    def _ifftn_wrapper(a):
+    def _ifftn_wrapper(a):  # noqa
         mesh = a.shape[1:]
         return _ifftn_blas(a, mesh)
 
@@ -384,7 +391,7 @@ def get_coulG(cell, k=np.zeros(3), exx=False, mf=None, mesh=None, Gv=None,
                 weights = 1 + Gp*Rc * scipy.special.j1(Gp*Rc) * scipy.special.k0(Gx*Rc)
                 weights -= Gx*Rc * scipy.special.j0(Gp*Rc) * scipy.special.k1(Gx*Rc)
                 coulG = 4*np.pi/absG2 * weights
-                # TODO: numerical integation
+                # TODO: numerical integration
                 # coulG[Gx==0] = -4*np.pi * (dr * r * scipy.special.j0(Gp*r) * np.log(r)).sum()
             if len(G0_idx) > 0:
                 coulG[G0_idx] = -np.pi*Rc**2 * (2*np.log(Rc) - 1)
@@ -685,7 +692,7 @@ def _build_supcell_(supcell, cell, Ls):
     _env = np.append(cell._env, coords.ravel())
     _atm = np.repeat(cell._atm[None,:,:], nimgs, axis=0)
     _atm = _atm.reshape(-1, ATM_SLOTS)
-    # Point to the corrdinates appended to _env
+    # Point to the coordinates appended to _env
     _atm[:,PTR_COORD] = cell._env.size + np.arange(nimgs * cell.natm) * 3
 
     _bas = np.repeat(cell._bas[None,:,:], nimgs, axis=0)

pyscf/pbc/tools/pyscf_ase.py

@@ -24,7 +24,6 @@ ASE package interface
 import numpy as np
 from ase.calculators.calculator import Calculator
 import ase.dft.kpoints
-from ase.lattice import bulk
 
 def pyscf_to_ase_atoms(cell):
     '''

pyscf/pbc/tools/pywannier90.py

@@ -325,7 +325,7 @@ def get_wigner_seitz_supercell(w90, ws_search_size=[2,2,2], ws_distance_tol=1e-6
     '''
     Return a grid that contains all the lattice within the Wigner-Seitz supercell
 
-    Ref: the hamiltonian_wigner_seitz(count_pts) in wannier90/src/hamittonian.F90
+    Ref: the hamiltonian_wigner_seitz(count_pts) in wannier90/src/hamiltonian.F90
     '''
 
     real_metric = w90.real_lattice_loc.T.dot(w90.real_lattice_loc)

pyscf/qmmm/mm_mole.py

@@ -97,7 +97,7 @@ def create_mm_mol(atoms_or_coords, charges=None, radii=None, unit='Angstrom'):
         charges : 1D array
             The charges of MM atoms.
         radii : 1D array
-            The Gaussian charge distribuction radii of MM atoms.
+            The Gaussian charge distribution radii of MM atoms.
         unit : string
             The unit of the input. Default is 'Angstrom'.
     '''

pyscf/scf/_response_functions.py

@@ -47,7 +47,7 @@ def _gen_rhf_response(mf, mo_coeff=None, mo_occ=None,
         ni.libxc.test_deriv_order(mf.xc, 2, raise_error=True)
         if mf.do_nlc():
             logger.warn(mf, 'NLC functional found in DFT object.  Its second '
-                        'deriviative is not available. Its contribution is '
+                        'derivative is not available. Its contribution is '
                         'not included in the response function.')
         omega, alpha, hyb = ni.rsh_and_hybrid_coeff(mf.xc, mol.spin)
         hybrid = ni.libxc.is_hybrid_xc(mf.xc)
@@ -157,7 +157,7 @@ def _gen_uhf_response(mf, mo_coeff=None, mo_occ=None,
         ni.libxc.test_deriv_order(mf.xc, 2, raise_error=True)
         if mf.do_nlc():
             logger.warn(mf, 'NLC functional found in DFT object.  Its second '
-                        'deriviative is not available. Its contribution is '
+                        'derivative is not available. Its contribution is '
                         'not included in the response function.')
         omega, alpha, hyb = ni.rsh_and_hybrid_coeff(mf.xc, mol.spin)
         hybrid = ni.libxc.is_hybrid_xc(mf.xc)

pyscf/scf/_vhf.py

@@ -79,10 +79,10 @@ class VHFOpt:
 
     @property
     def direct_scf_tol(self):
-        return self._this.contents.direct_scf_cutoff
+        return self._this.contents.direct_scf_tol
     @direct_scf_tol.setter
     def direct_scf_tol(self, v):
-        self._this.contents.direct_scf_cutoff = v
+        self._this.contents.direct_scf_tol = v
 
     @property
     def prescreen(self):
@@ -162,7 +162,7 @@ class _VHFOpt:
         self.mol = mol
         self._this = cvhfopt = _CVHFOpt()
         cvhfopt.nbas = mol.nbas
-        cvhfopt.direct_scf_cutoff = direct_scf_tol
+        cvhfopt.direct_scf_tol = direct_scf_tol
         cvhfopt.fprescreen = _fpointer(prescreen)
         cvhfopt.r_vkscreen = _fpointer('CVHFr_vknoscreen')
         self._q_cond = None
@@ -207,10 +207,10 @@ class _VHFOpt:
 
     @property
     def direct_scf_tol(self):
-        return self._this.direct_scf_cutoff
+        return self._this.direct_scf_tol
     @direct_scf_tol.setter
     def direct_scf_tol(self, v):
-        self._this.direct_scf_cutoff = v
+        self._this.direct_scf_tol = v
 
     @property
     def prescreen(self):
@@ -265,9 +265,9 @@ class _VHFOpt:
 
 class SGXOpt(_VHFOpt):
     def __init__(self, mol, intor=None, prescreen='CVHFnoscreen',
-                 qcondname=None, dmcondname=None, direct_scf_cutoff=1e-14):
+                 qcondname=None, dmcondname=None, direct_scf_tol=1e-14):
         _VHFOpt.__init__(self, mol, intor, prescreen, qcondname, dmcondname,
-                         direct_scf_cutoff)
+                         direct_scf_tol)
         self.ngrids = None
 
     def set_dm(self, dm, atm, bas, env):
@@ -287,18 +287,22 @@ class SGXOpt(_VHFOpt):
             if self._dmcondname != 'SGXnr_dm_cond':
                 raise ValueError('SGXOpt only supports SGXnr_dm_cond')
             fdmcond = getattr(libcvhf, self._dmcondname)
-        dm_cond = numpy.empty((mol.nbas, self.ngrids))
+        if self.ngrids is None:
+            ngrids = int(env[gto.NGRIDS])
+        else:
+            ngrids = self.ngrids
+        dm_cond = numpy.empty((mol.nbas, ngrids))
         fdmcond(dm_cond.ctypes, dm.ctypes, ctypes.c_int(n_dm),
                 ao_loc.ctypes, mol._atm.ctypes, ctypes.c_int(mol.natm),
                 mol._bas.ctypes, ctypes.c_int(mol.nbas), mol._env.ctypes,
-                ctypes.c_int(self.ngrids))
+                ctypes.c_int(ngrids))
         self.dm_cond = dm_cond
 
 
 class _CVHFOpt(ctypes.Structure):
     _fields_ = [('nbas', ctypes.c_int),
                 ('ngrids', ctypes.c_int),
-                ('direct_scf_cutoff', ctypes.c_double),
+                ('direct_scf_tol', ctypes.c_double),
                 ('q_cond', ctypes.c_void_p),
                 ('dm_cond', ctypes.c_void_p),
                 ('fprescreen', ctypes.c_void_p),