pyscf 2.4.0__py3-none-macosx_11_0_arm64.whl → 2.6.0__py3-none-macosx_11_0_arm64.whl

pyscf/cc/ccsd.py

@@ -923,7 +923,7 @@ class CCSDBase(lib.StreamObject):
     cc2 = getattr(__config__, 'cc_ccsd_CCSD_cc2', False)
     callback = None
 
-    _keys = set((
+    _keys = {
         'max_cycle', 'conv_tol', 'iterative_damping',
         'conv_tol_normt', 'diis', 'diis_space', 'diis_file',
         'diis_start_cycle', 'diis_start_energy_diff', 'direct',
@@ -931,7 +931,7 @@ class CCSDBase(lib.StreamObject):
         'mol', 'verbose', 'stdout', 'frozen', 'level_shift',
         'mo_coeff', 'mo_occ', 'converged', 'converged_lambda', 'emp2', 'e_hf',
         'e_corr', 't1', 't2', 'l1', 'l2', 'chkfile',
-    ))
+    }
 
     def __init__(self, mf, frozen=None, mo_coeff=None, mo_occ=None):
         from pyscf.scf import hf
@@ -1240,6 +1240,15 @@ class CCSDBase(lib.StreamObject):
     def nuc_grad_method(self):
         raise NotImplementedError
 
+    # to_gpu can be reused only when __init__ still takes mf
+    def to_gpu(self):
+        mf = self.base.to_gpu()
+        from importlib import import_module
+        mod = import_module(self.__module__.replace('pyscf', 'gpu4pyscf'))
+        cls = getattr(mod, self.__class__.__name__)
+        obj = cls(mf)
+        return obj
+
 class CCSD(CCSDBase):
     __doc__ = CCSDBase.__doc__
 
@@ -1496,6 +1505,8 @@ def _make_eris_incore(mycc, mo_coeff=None):
     return eris
 
 def _make_eris_outcore(mycc, mo_coeff=None):
+    from pyscf.scf.hf import RHF
+    assert isinstance(mycc._scf, RHF)
     cput0 = (logger.process_clock(), logger.perf_counter())
     log = logger.Logger(mycc.stdout, mycc.verbose)
     eris = _ChemistsERIs()

pyscf/cc/ccsd_t.py

@@ -176,7 +176,7 @@ def _sort_eri(mycc, eris, nocc, nvir, vvop, log):
 
 def _sort_t2_vooo_(mycc, orbsym, t1, t2, eris):
     assert (t2.flags.c_contiguous)
-    vooo = numpy.asarray(eris.ovoo).transpose(1,0,3,2).conj().copy()
+    vooo = numpy.asarray(eris.ovoo).transpose(1,0,2,3).conj().copy()
     nocc, nvir = t1.shape
     if mycc.mol.symmetry:
         orbsym = numpy.asarray(orbsym, dtype=numpy.int32)

pyscf/cc/ccsd_t_slow.py

@@ -45,7 +45,7 @@ def kernel(mycc, eris, t1=None, t2=None, verbose=logger.NOTE):
     eijk = lib.direct_sum('i,j,k->ijk', e_occ, e_occ, e_occ)
 
     eris_vvov = eris.get_ovvv().conj().transpose(1,3,0,2)
-    eris_vooo = numpy.asarray(eris.ovoo).conj().transpose(1,0,3,2)
+    eris_vooo = numpy.asarray(eris.ovoo).conj().transpose(1,0,2,3)
     eris_vvoo = numpy.asarray(eris.ovov).conj().transpose(1,3,0,2)
     fvo = eris.fock[nocc:,:nocc]
     def get_w(a, b, c):

pyscf/cc/dfccsd.py

@@ -27,7 +27,7 @@ from pyscf import __config__
 MEMORYMIN = getattr(__config__, 'cc_ccsd_memorymin', 2000)
 
 class RCCSD(ccsd.CCSD):
-    _keys = set(['with_df'])
+    _keys = {'with_df'}
 
     def __init__(self, mf, frozen=None, mo_coeff=None, mo_occ=None):
         ccsd.CCSD.__init__(self, mf, frozen, mo_coeff, mo_occ)
@@ -48,6 +48,8 @@ class RCCSD(ccsd.CCSD):
         assert (not self.direct)
         return ccsd.CCSD._add_vvvv(self, t1, t2, eris, out, with_ovvv, t2sym)
 
+    to_gpu = lib.to_gpu
+
 
 def _contract_vvvv_t2(mycc, mol, vvL, t2, out=None, verbose=None):
     '''Ht2 = numpy.einsum('ijcd,acdb->ijab', t2, vvvv)
@@ -121,6 +123,7 @@ class _ChemistsERIs(ccsd._ChemistsERIs):
         return _contract_vvvv_t2(mycc, self.mol, self.vvL, t2, out, verbose)
 
 def _make_df_eris(cc, mo_coeff=None):
+    assert cc._scf.istype('RHF')
     eris = _ChemistsERIs()
     eris._common_init_(cc, mo_coeff)
     nocc = eris.nocc

pyscf/cc/eom_rccsd.py

@@ -96,10 +96,10 @@ def kernel(eom, nroots=1, koopmans=False, guess=None, left=False,
 
 
 class EOM(lib.StreamObject):
-    _keys = set((
+    _keys = {
         'mol', 'max_space', 'max_cycle', 'conv_tol', 'partition',
         'e', 'v', 'nocc', 'nmo',
-    ))
+    }
 
     def __init__(self, cc):
         self.mol = cc.mol

pyscf/cc/gccsd.py

@@ -117,7 +117,6 @@ class GCCSD(ccsd.CCSDBase):
     conv_tol_normt = getattr(__config__, 'cc_gccsd_GCCSD_conv_tol_normt', 1e-6)
 
     def __init__(self, mf, frozen=None, mo_coeff=None, mo_occ=None):
-        assert (isinstance(mf, scf.ghf.GHF))
         ccsd.CCSDBase.__init__(self, mf, frozen, mo_coeff, mo_occ)
 
     def init_amps(self, eris=None):
@@ -290,6 +289,8 @@ class GCCSD(ccsd.CCSDBase):
                 orbspin = orbspin[self.get_frozen_mask()]
         return spin2spatial(tx, orbspin)
 
+    to_gpu = lib.to_gpu
+
 CCSD = GCCSD
 
 
@@ -363,6 +364,7 @@ def _make_eris_incore(mycc, mo_coeff=None, ao2mofn=None):
             mo = mo_a + mo_b
             eri = ao2mo.kernel(mycc._scf._eri, mo)
             if eri.size == nmo**4:  # if mycc._scf._eri is a complex array
+                eri = eri.reshape(nmo**2, nmo**2)
                 sym_forbid = (orbspin[:,None] != orbspin).ravel()
             else:  # 4-fold symmetry
                 sym_forbid = (orbspin[:,None] != orbspin)[np.tril_indices(nmo)]
@@ -385,6 +387,8 @@ def _make_eris_incore(mycc, mo_coeff=None, ao2mofn=None):
     return eris
 
 def _make_eris_outcore(mycc, mo_coeff=None):
+    from pyscf.scf.ghf import GHF
+    assert isinstance(mycc._scf, GHF)
     cput0 = (logger.process_clock(), logger.perf_counter())
     log = logger.Logger(mycc.stdout, mycc.verbose)
 

pyscf/cc/gccsd_t.py

@@ -85,13 +85,13 @@ if __name__ == '__main__':
     mycc = cc.CCSD(mf).set(conv_tol=1e-11).run()
     et = mycc.ccsd_t()
 
-    mycc = cc.GCCSD(scf.addons.convert_to_ghf(mf)).set(conv_tol=1e-11).run()
+    mycc = cc.GCCSD(mf.to_ghf()).set(conv_tol=1e-11).run()
     eris = mycc.ao2mo()
     print(kernel(mycc, eris) - et)
 
     numpy.random.seed(1)
     mf.mo_coeff = numpy.random.random(mf.mo_coeff.shape) - .9
-    mycc = cc.GCCSD(scf.addons.convert_to_ghf(mf))
+    mycc = cc.GCCSD(mf.to_ghf())
     eris = mycc.ao2mo()
     nocc = 10
     nvir = mol.nao_nr() * 2 - nocc

pyscf/cc/gccsd_t_rdm.py

@@ -120,7 +120,7 @@ if __name__ == '__main__':
     mol.basis = '631g'
     mol.build()
     mf0 = mf = scf.RHF(mol).run(conv_tol=1.)
-    mf = scf.addons.convert_to_ghf(mf)
+    mf = mf.to_ghf()
 
     from pyscf.cc import ccsd_t_lambda_slow as ccsd_t_lambda
     from pyscf.cc import ccsd_t_rdm_slow as ccsd_t_rdm

pyscf/cc/momgfccsd.py

@@ -607,10 +607,10 @@ class MomGFCCSD(lib.StreamObject):
             particle Green's function
     """
 
-    _keys = set((
+    _keys = {
         'verbose', 'stdout', 'niter', 'weight_tol',
         'hermi_moments', 'hermi_solver', 'eh', 'ep', 'vh', 'vp', 'chkfile',
-    ))
+    }
 
     def __init__(self, mycc, niter=(2, 2)):
         self._cc = mycc

pyscf/cc/qcisd_t_slow.py

@@ -47,7 +47,7 @@ def kernel(mycc, eris, t1=None, t2=None, verbose=logger.NOTE):
     eijk = lib.direct_sum('i,j,k->ijk', e_occ, e_occ, e_occ)
 
     eris_vvov = eris.get_ovvv().conj().transpose(1,3,0,2)
-    eris_vooo = numpy.asarray(eris.ovoo).conj().transpose(1,0,3,2)
+    eris_vooo = numpy.asarray(eris.ovoo).conj().transpose(1,0,2,3)
     eris_vvoo = numpy.asarray(eris.ovov).conj().transpose(1,3,0,2)
     fvo = eris.fock[nocc:,:nocc]
     def get_w(a, b, c):

pyscf/cc/rccsd.py

@@ -259,6 +259,8 @@ def _make_eris_incore(mycc, mo_coeff=None, ao2mofn=None):
     return eris
 
 def _make_eris_outcore(mycc, mo_coeff=None):
+    from pyscf.scf.hf import RHF
+    assert isinstance(mycc._scf, RHF)
     cput0 = (logger.process_clock(), logger.perf_counter())
     log = logger.Logger(mycc.stdout, mycc.verbose)
     eris = _ChemistsERIs()

pyscf/cc/rccsd_slow.py

@@ -148,7 +148,7 @@ def energy(cc, t1, t2, eris):
 
 
 class RCCSD(ccsd.CCSD):
-    _keys = set(['max_space'])
+    _keys = {'max_space'}
 
     def __init__(self, mf, frozen=None, mo_coeff=None, mo_occ=None):
         ccsd.CCSD.__init__(self, mf, frozen, mo_coeff, mo_occ)

pyscf/cc/uccsd.py

@@ -547,7 +547,6 @@ class UCCSD(ccsd.CCSDBase):
 # * A pair of list : First list is the orbital indices to be frozen for alpha
 #       orbitals, second list is for beta orbitals
     def __init__(self, mf, frozen=None, mo_coeff=None, mo_occ=None):
-        assert isinstance(mf, scf.uhf.UHF)
         ccsd.CCSDBase.__init__(self, mf, frozen, mo_coeff, mo_occ)
 
     get_nocc = get_nocc
@@ -759,6 +758,8 @@ class UCCSD(ccsd.CCSDBase):
     def amplitudes_from_rccsd(self, t1, t2):
         return amplitudes_from_rccsd(t1, t2)
 
+    to_gpu = lib.to_gpu
+
 CCSD = UCCSD
 
 
@@ -942,6 +943,7 @@ def _make_eris_incore(mycc, mo_coeff=None, ao2mofn=None):
     return eris
 
 def _make_df_eris_outcore(mycc, mo_coeff=None):
+    assert mycc._scf.istype('UHF')
     cput0 = (logger.process_clock(), logger.perf_counter())
     log = logger.Logger(mycc.stdout, mycc.verbose)
     eris = _ChemistsERIs()
@@ -1059,6 +1061,8 @@ def _make_df_eris_outcore(mycc, mo_coeff=None):
     return eris
 
 def _make_eris_outcore(mycc, mo_coeff=None):
+    from pyscf.scf.uhf import UHF
+    assert isinstance(mycc._scf, UHF)
     eris = _ChemistsERIs()
     eris._common_init_(mycc, mo_coeff)
 

pyscf/cc/uccsd_slow.py

@@ -610,7 +610,7 @@ def uspatial2spin(cc, moidx, mo_coeff):
 
     dm = cc._scf.make_rdm1(cc.mo_coeff, cc.mo_occ)
     fockao = cc._scf.get_hcore() + cc._scf.get_veff(cc.mol, dm)
-    fockab = list()
+    fockab = []
     for a in range(2):
         fockab.append( reduce(numpy.dot, (mo_coeff[a].T, fockao[a], mo_coeff[a])) )
 

pyscf/ci/__init__.py

@@ -13,7 +13,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from pyscf import lib
 from pyscf import scf
 from pyscf.ci import cisd
 from pyscf.ci import ucisd
@@ -21,11 +20,10 @@ from pyscf.ci import gcisd
 from pyscf.cc import qcisd
 
 def CISD(mf, frozen=None, mo_coeff=None, mo_occ=None):
-    from pyscf.soscf import newton_ah
-
-    if isinstance(mf, scf.uhf.UHF):
+    if mf.istype('UHF'):
         return UCISD(mf, frozen, mo_coeff, mo_occ)
-    elif isinstance(mf, scf.rohf.ROHF):
+    elif mf.istype('ROHF'):
+        from pyscf import lib
         lib.logger.warn(mf, 'RCISD method does not support ROHF method. ROHF object '
                         'is converted to UHF object and UCISD method is called.')
         return UCISD(mf, frozen, mo_coeff, mo_occ)
@@ -34,24 +32,32 @@ def CISD(mf, frozen=None, mo_coeff=None, mo_occ=None):
 CISD.__doc__ = cisd.CISD.__doc__
 
 def RCISD(mf, frozen=None, mo_coeff=None, mo_occ=None):
-    from pyscf.soscf import newton_ah
+    from pyscf.df.df_jk import _DFHF
 
-    if isinstance(mf, newton_ah._CIAH_SOSCF) or not isinstance(mf, scf.hf.RHF):
-        mf = scf.addons.convert_to_rhf(mf)
+    mf = mf.remove_soscf()
+    if not mf.istype('RHF'):
+        mf = mf.to_rhf()
 
-    if getattr(mf, 'with_df', None):
+    if isinstance(mf, _DFHF) and mf.with_df:
+        from pyscf import lib
+        lib.logger.warn(mf, f'DF-RCISD for DFHF method {mf} is not available. '
+                        'Normal RCISD method is called.')
         return cisd.RCISD(mf, frozen, mo_coeff, mo_occ)
     else:
         return cisd.RCISD(mf, frozen, mo_coeff, mo_occ)
 RCISD.__doc__ = cisd.RCISD.__doc__
 
 def UCISD(mf, frozen=None, mo_coeff=None, mo_occ=None):
-    from pyscf.soscf import newton_ah
+    from pyscf.df.df_jk import _DFHF
 
-    if isinstance(mf, newton_ah._CIAH_SOSCF) or not isinstance(mf, scf.uhf.UHF):
-        mf = scf.addons.convert_to_uhf(mf)
+    mf = mf.remove_soscf()
+    if not mf.istype('UHF'):
+        mf = mf.to_uhf()
 
-    if getattr(mf, 'with_df', None):
+    if isinstance(mf, _DFHF) and mf.with_df:
+        from pyscf import lib
+        lib.logger.warn(mf, f'DF-UCISD for DFHF method {mf} is not available. '
+                        'Normal UCISD method is called.')
         return ucisd.UCISD(mf, frozen, mo_coeff, mo_occ)
     else:
         return ucisd.UCISD(mf, frozen, mo_coeff, mo_occ)
@@ -59,12 +65,13 @@ UCISD.__doc__ = ucisd.UCISD.__doc__
 
 
 def GCISD(mf, frozen=None, mo_coeff=None, mo_occ=None):
-    from pyscf.soscf import newton_ah
+    from pyscf.df.df_jk import _DFHF
 
-    if isinstance(mf, newton_ah._CIAH_SOSCF) or not isinstance(mf, scf.ghf.GHF):
-        mf = scf.addons.convert_to_ghf(mf)
+    mf = mf.remove_soscf()
+    if not mf.istype('GHF'):
+        mf = mf.to_ghf()
 
-    if getattr(mf, 'with_df', None):
+    if isinstance(mf, _DFHF) and mf.with_df:
         raise NotImplementedError('DF-GCISD')
     else:
         return gcisd.GCISD(mf, frozen, mo_coeff, mo_occ)
@@ -72,9 +79,9 @@ GCISD.__doc__ = gcisd.GCISD.__doc__
 
 
 def QCISD(mf, frozen=None, mo_coeff=None, mo_occ=None):
-    if isinstance(mf, scf.uhf.UHF):
+    if mf.istype('UHF'):
         raise NotImplementedError
-    elif isinstance(mf, scf.ghf.GHF):
+    elif mf.istype('GHF'):
         raise NotImplementedError
     else:
         return RQCISD(mf, frozen, mo_coeff, mo_occ)
@@ -85,18 +92,17 @@ scf.hf.SCF.QCISD = QCISD
 def RQCISD(mf, frozen=None, mo_coeff=None, mo_occ=None):
     import numpy
     from pyscf import lib
-    from pyscf.soscf import newton_ah
 
-    if isinstance(mf, scf.uhf.UHF):
+    if mf.istype('UHF'):
         raise RuntimeError('RQCISD cannot be used with UHF method.')
-    elif isinstance(mf, scf.rohf.ROHF):
+    elif mf.istype('ROHF'):
         lib.logger.warn(mf, 'RQCISD method does not support ROHF method. ROHF object '
                         'is converted to UHF object and UQCISD method is called.')
-        mf = scf.addons.convert_to_uhf(mf)
         raise NotImplementedError
 
-    if isinstance(mf, newton_ah._CIAH_SOSCF) or not isinstance(mf, scf.hf.RHF):
-        mf = scf.addons.convert_to_rhf(mf)
+    mf = mf.remove_soscf()
+    if not mf.istype('RHF'):
+        mf = mf.to_rhf()
 
     elif numpy.iscomplexobj(mo_coeff) or numpy.iscomplexobj(mf.mo_coeff):
         raise NotImplementedError

pyscf/ci/addons.py

@@ -17,14 +17,13 @@
 #
 
 import numpy as np
-from pyscf import scf
 
 def convert_to_gcisd(myci):
     from pyscf.ci import gcisd
     if isinstance(myci, gcisd.GCISD):
         return myci
 
-    mf = scf.addons.convert_to_ghf(myci._scf)
+    mf = myci._scf.to_ghf()
     gci = gcisd.GCISD(mf)
     assert (myci._nocc is None)
     assert (myci._nmo is None)

pyscf/ci/cisd.py

@@ -860,12 +860,12 @@ class CISD(lib.StreamObject):
     direct = getattr(__config__, 'ci_cisd_CISD_direct', False)
     async_io = getattr(__config__, 'ci_cisd_CISD_async_io', True)
 
-    _keys = set((
+    _keys = {
         'conv_tol', 'max_cycle', 'max_space', 'lindep',
         'level_shift', 'direct', 'async_io', 'mol', 'max_memory',
         'nroots', 'frozen', 'chkfile', 'converged', 'mo_coeff', 'mo_occ',
         'e_hf', 'e_corr', 'emp2', 'ci',
-    ))
+    }
 
     def __init__(self, mf, frozen=None, mo_coeff=None, mo_occ=None):
         from pyscf.scf import hf
@@ -1131,6 +1131,8 @@ class CISD(lib.StreamObject):
         from pyscf.grad import cisd
         return cisd.Gradients(self)
 
+    to_gpu = lib.to_gpu
+
 class RCISD(CISD):
     pass
 

pyscf/data/elements.py

@@ -30,7 +30,7 @@ ELEMENTS = [
     'Md', 'No', 'Lr', 'Rf', 'Db', 'Sg', 'Bh', 'Hs', 'Mt', 'Ds',
     'Rg', 'Cn', 'Nh', 'Fl', 'Mc', 'Lv', 'Ts', 'Og',
 ]
-NUC = dict(((x,i) for i,x in enumerate(ELEMENTS)))
+NUC = {x: i for i,x in enumerate(ELEMENTS)}
 NUC.update((x.upper(),i) for i,x in enumerate(ELEMENTS))
 NUC['GHOST'] = 0
 ELEMENTS_PROTON = NUC
@@ -1130,7 +1130,7 @@ def _rm_digit(symb):
     else:
         return ''.join([i for i in symb if i.isalpha()])
 
-_ELEMENTS_UPPER = dict((x.upper(),x) for x in ELEMENTS)
+_ELEMENTS_UPPER = {x.upper(): x for x in ELEMENTS}
 _ELEMENTS_UPPER['GHOST'] = 'Ghost'
 
 def charge(symb_or_chg):

pyscf/df/addons.py

@@ -79,7 +79,7 @@ def aug_etb_for_dfbasis(mol, dfbasis=DFBASIS, beta=ETB_BETA,
     exps = alpha*beta^i for i = 1..N
     '''
     nuc_start = gto.charge(start_at)
-    uniq_atoms = set([a[0] for a in mol._atom])
+    uniq_atoms = {a[0] for a in mol._atom}
 
     newbasis = {}
     for symb in uniq_atoms:
@@ -143,16 +143,16 @@ def make_auxbasis(mol, mp2fit=False):
     '''Depending on the orbital basis, generating even-tempered Gaussians or
     the optimized auxiliary basis defined in DEFAULT_AUXBASIS
     '''
-    uniq_atoms = set([a[0] for a in mol._atom])
+    uniq_atoms = {a[0] for a in mol._atom}
     if isinstance(mol.basis, str):
-        _basis = dict(((a, mol.basis) for a in uniq_atoms))
+        _basis = {a: mol.basis for a in uniq_atoms}
     elif 'default' in mol.basis:
         default_basis = mol.basis['default']
-        _basis = dict(((a, default_basis) for a in uniq_atoms))
+        _basis = {a: default_basis for a in uniq_atoms}
         _basis.update(mol.basis)
         del (_basis['default'])
     else:
-        _basis = mol._basis
+        _basis = mol._basis or {}
 
     auxbasis = {}
     for k in _basis:
@@ -209,11 +209,11 @@ def make_auxmol(mol, auxbasis=None):
     pmol.basis = auxbasis
 
     if isinstance(auxbasis, (str, list, tuple)):
-        uniq_atoms = set([a[0] for a in mol._atom])
-        _basis = dict([(a, auxbasis) for a in uniq_atoms])
+        uniq_atoms = {a[0] for a in mol._atom}
+        _basis = {a: auxbasis for a in uniq_atoms}
     elif 'default' in auxbasis:
-        uniq_atoms = set([a[0] for a in mol._atom])
-        _basis = dict(((a, auxbasis['default']) for a in uniq_atoms))
+        uniq_atoms = {a[0] for a in mol._atom}
+        _basis = {a: auxbasis['default'] for a in uniq_atoms}
         _basis.update(auxbasis)
         del (_basis['default'])
     else:

pyscf/df/df.py

@@ -79,7 +79,7 @@ class DF(lib.StreamObject):
     _compatible_format = getattr(__config__, 'df_df_DF_compatible_format', False)
     _dataname = 'j3c'
 
-    _keys = set(('mol', 'auxmol'))
+    _keys = {'mol', 'auxmol'}
 
     def __init__(self, mol, auxbasis=None):
         self.mol = mol
@@ -284,7 +284,7 @@ class DF(lib.StreamObject):
         else:
             rsh_df = self._rsh_df[key] = self.copy().reset()
             if hasattr(self, '_dataname'):
-                rsh_df._dataname = f'{self._dataname}/lr/{key}'
+                rsh_df._dataname = f'{self._dataname}-lr/{key}'
             logger.info(self, 'Create RSH-DF object %s for omega=%s', rsh_df, omega)
 
         mol = self.mol
@@ -308,6 +308,8 @@ class DF(lib.StreamObject):
             if auxmol_omega is not None:
                 auxmol.omega = auxmol_omega
 
+    to_gpu = lib.to_gpu
+
 GDF = DF
 
 
@@ -353,7 +355,12 @@ class DF4C(DF):
         with self.range_coulomb(omega) as rsh_df:
             return df_jk.r_get_jk(rsh_df, dm, hermi, with_j, with_k)
 
+    def get_eri(self):
+        raise NotImplementedError
+    get_ao_eri = get_eri
+
     def ao2mo(self, mo_coeffs):
         raise NotImplementedError
+    get_mo_eri = ao2mo
 
 GDF4C = DF4C

pyscf/df/df_jk.py

@@ -108,7 +108,7 @@ class _DFHF:
 
     __name_mixin__ = 'DF'
 
-    _keys = set(['with_df', 'only_dfj'])
+    _keys = {'with_df', 'only_dfj'}
 
     def __init__(self, mf, df=None, only_dfj=None):
         self.__dict__.update(mf.__dict__)
@@ -123,8 +123,7 @@ class _DFHF:
     def undo_df(self):
         '''Remove the DFHF Mixin'''
         obj = lib.view(self, lib.drop_class(self.__class__, _DFHF))
-        del obj.with_df
-        del obj.only_dfj
+        del obj.with_df, obj.only_dfj
         return obj
 
     def reset(self, mol=None):
@@ -227,6 +226,10 @@ class _DFHF:
         from pyscf import mcscf
         return mcscf.DFCASSCF(self, ncas, nelecas, auxbasis, ncore, frozen)
 
+    def to_gpu(self):
+        obj = self.undo_df().to_gpu().density_fit()
+        return lib.to_gpu(self, obj)
+
 
 def get_jk(dfobj, dm, hermi=1, with_j=True, with_k=True, direct_scf_tol=1e-13):
     assert (with_j or with_k)
@@ -257,8 +260,9 @@ def get_jk(dfobj, dm, hermi=1, with_j=True, with_k=True, direct_scf_tol=1e-13):
 
     if not with_k:
         for eri1 in dfobj.loop():
-            rho = numpy.einsum('ix,px->ip', dmtril, eri1)
-            vj += numpy.einsum('ip,px->ix', rho, eri1)
+            # uses numpy.matmul
+            vj += (dmtril @ eri1.T) @ eri1
+
 
     elif getattr(dm, 'mo_coeff', None) is not None:
         #TODO: test whether dm.mo_coeff matching dm
@@ -287,9 +291,8 @@ def get_jk(dfobj, dm, hermi=1, with_j=True, with_k=True, direct_scf_tol=1e-13):
             naux, nao_pair = eri1.shape
             assert (nao_pair == nao*(nao+1)//2)
             if with_j:
-                rho = numpy.einsum('ix,px->ip', dmtril, eri1)
-                vj += numpy.einsum('ip,px->ix', rho, eri1)
-
+                # uses numpy.matmul
+                vj += (dmtril @ eri1.T) @ eri1
             for k in range(nset):
                 nocc = orbo[k].shape[1]
                 if nocc > 0:
@@ -315,8 +318,9 @@ def get_jk(dfobj, dm, hermi=1, with_j=True, with_k=True, direct_scf_tol=1e-13):
         for eri1 in dfobj.loop(blksize):
             naux, nao_pair = eri1.shape
             if with_j:
-                rho = numpy.einsum('ix,px->ip', dmtril, eri1)
-                vj += numpy.einsum('ip,px->ix', rho, eri1)
+                # uses numpy.matmul
+                vj += (dmtril @ eri1.T) @ eri1
+
 
             for k in range(nset):
                 buf1 = buf[0,:naux]

pyscf/df/grad/casdm2_util.py

@@ -476,7 +476,7 @@ def grad_elec_dferi (mc_grad, mo_cas=None, ci=None, dfcasdm2=None, casdm2=None,
 
     aoslices = mol.aoslice_by_atom ()
     dE = np.array ([dE[:,p0:p1].sum (axis=1) for p0, p1 in aoslices[:,2:]]).transpose (1,0,2)
-    return np.ascontiguousarray (dE)
+    return np.ascontiguousarray (dE)[:,atmlst,:]
 
 if __name__ == '__main__':
     from pyscf.tools import molden

pyscf/df/grad/casscf.py

@@ -224,6 +224,8 @@ class Gradients(casci_grad.Gradients):
 
     as_scanner = as_scanner
 
+    to_gpu = lib.to_gpu
+
 Grad = Gradients
 
 #from pyscf import mcscf

pyscf/df/grad/rhf.py

@@ -38,8 +38,9 @@ from functools import reduce
 from itertools import product
 from pyscf.ao2mo import _ao2mo
 from pyscf.df import df_jk
+from pyscf.df.incore import LINEAR_DEP_THR
 
-LINEAR_DEP_THRESHOLD = 1e-9
+LINEAR_DEP_THRESHOLD = LINEAR_DEP_THR
 
 def get_jk(mf_grad, mol=None, dm=None, hermi=0, with_j=True, with_k=True,
            decompose_j2c='CD', lindep=LINEAR_DEP_THRESHOLD):
@@ -342,6 +343,9 @@ def _decompose_rdm1 (mf_grad, mol, dm):
     if hasattr (dm, 'mo_coeff') and hasattr (dm, 'mo_occ'):
         mo_coeff = dm.mo_coeff
         mo_occ = dm.mo_occ
+        if getattr(mo_occ, 'ndim', None) == 1: # RHF orbitals
+            mo_coeff = [mo_coeff]
+            mo_occ = [mo_occ]
     else:
         s0 = mol.intor ('int1e_ovlp')
         mo_occ = []
@@ -352,10 +356,7 @@ def _decompose_rdm1 (mf_grad, mol, dm):
             mo_occ.append (n)
             mo_coeff.append (c)
         mo_occ = numpy.stack (mo_occ, axis=0)
-    nmo = mo_occ.shape[-1]
 
-    mo_coeff = numpy.asarray(mo_coeff).reshape(-1,nao,nmo)
-    mo_occ   = numpy.asarray(mo_occ).reshape(-1,nmo)
     orbor = []
     orbol = []
     for i in range(nset):
@@ -482,12 +483,15 @@ class Gradients(rhf_grad.Gradients):
     _keys = {'with_df', 'auxbasis_response'}
 
     def __init__(self, mf):
-        assert isinstance(mf, df.df_jk._DFHF)
-        # Whether to include the response of DF auxiliary basis when computing
-        # nuclear gradients of J/K matrices
-        self.auxbasis_response = True
         rhf_grad.Gradients.__init__(self, mf)
 
+    # Whether to include the response of DF auxiliary basis when computing
+    # nuclear gradients of J/K matrices
+    auxbasis_response = True
+
+    def check_sanity(self):
+        assert isinstance(self.base, df.df_jk._DFHF)
+
     def get_jk(self, mol=None, dm=None, hermi=0, with_j=True, with_k=True,
                omega=None):
         if omega is None: