gpu4pyscf 0.7.2__tar.gz → 0.7.4__tar.gz

{gpu4pyscf-0.7.2/gpu4pyscf.egg-info → gpu4pyscf-0.7.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: gpu4pyscf
-Version: 0.7.2
+Version: 0.7.4
 Summary: GPU extensions for PySCF
 Home-page: UNKNOWN
 Author: Qiming Sun

{gpu4pyscf-0.7.2 → gpu4pyscf-0.7.4}/README.md

@@ -51,6 +51,7 @@ Features
 - Nonlocal functional correction (vv10) for SCF and gradient;
 - ECP is supported and calculated on CPU;
 - PCM models, SMD model, their analytical gradients, and semi-analytical Hessian matrix;
+- Unrestricted Hartree-Fock and Unrestricted DFT, gradient, and Hessian
 
 Limitations
 --------
@@ -87,7 +88,25 @@ h = mf.Hessian()
 h_dft = h.kernel()   # compute analytical Hessian
 
 ```
-Find more examples in gpu4pyscf/examples
+
+`to_gpu` is supported since PySCF 2.5.0
+```python
+import pyscf
+from pyscf.dft import rks
+
+atom ='''
+O       0.0000000000    -0.0000000000     0.1174000000
+H      -0.7570000000    -0.0000000000    -0.4696000000
+H       0.7570000000     0.0000000000    -0.4696000000
+'''
+
+mol = pyscf.M(atom=atom, basis='def2-tzvpp')
+mf = rks.RKS(mol, xc='LDA').density_fit().to_gpu()  # move PySCF object to GPU4PySCF object
+e_dft = mf.kernel()  # compute total energy
+
+```
+
+Find more examples in [gpu4pyscf/examples](https://github.com/pyscf/gpu4pyscf/tree/master/examples)
 
 Benchmarks
 --------
@@ -108,5 +127,5 @@ Speedup with GPU4PySCF v0.6.0 on A100-80G over Q-Chem 6.1 on 32-cores CPU (Desit
 | 095_Azadirachtin  |     95 |   5.58 |   7.72 |   24.18 |  26.84 |     25.21 |
 | 113_Taxol         |    113 |   5.44 |   6.81 |   24.58 |  29.14 |    nan    |
 
-Find more benchmarks in gpu4pyscf/benchmarks
+Find more benchmarks in [gpu4pyscf/benchmarks](https://github.com/pyscf/gpu4pyscf/tree/master/benchmarks)
 

{gpu4pyscf-0.7.2 → gpu4pyscf-0.7.4}/gpu4pyscf/__init__.py

@@ -1,6 +1,6 @@
 from . import lib, grad, hessian, solvent, scf, dft
 
-__version__ = '0.7.2'
+__version__ = '0.7.4'
 
 # monkey patch libxc reference due to a bug in nvcc
 from pyscf.dft import libxc

gpu4pyscf-0.7.4/gpu4pyscf/cc/__init__.py

@@ -0,0 +1 @@
+from .cc import ccsd_incore

gpu4pyscf-0.7.4/gpu4pyscf/cc/ccsd_incore.py

@@ -0,0 +1,542 @@
+# Copyright 2024 The GPU4PySCF Authors. All Rights Reserved.
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program.  If not, see <http://www.gnu.org/licenses/>.
+
+
+'''
+Rewrite the pyscf/cc/ccsd.py using cupy, and GPU for ERIs.
+This implementation requires that the GPU memory is large enough to hold at
+least two t2 tensors.
+'''
+
+import time
+import ctypes
+import cupy
+import numpy as np
+from pyscf import gto
+from pyscf import lib
+from pyscf.ao2mo.outcore import balance_partition
+from pyscf.ao2mo import _ao2mo
+from pyscf.cc import ccsd
+from pyscf.cc import _ccsd
+from pyscf import __config__
+from gpu4pyscf.scf import hf as gpu_hf
+from gpu4pyscf.lib.cupy_helper import load_library
+from gpu4pyscf.lib import logger
+
+FREE_CUPY_CACHE = True
+
+BLKMIN = getattr(__config__, 'cc_ccsd_blkmin', 4)
+MEMORYMIN = getattr(__config__, 'cc_ccsd_memorymin', 2000)
+
+libgint = load_library('libgint')
+libgint.GINTfill_int2e.restype = ctypes.c_int
+
+def update_amps(mycc, t1, t2, eris):
+    time0 = logger.process_clock(), logger.perf_counter()
+    log = logger.Logger(mycc.stdout, mycc.verbose)
+    nocc, nvir = t1.shape
+    fock = eris.fock
+    mo_e_o = eris.mo_energy[:nocc]
+    mo_e_v = eris.mo_energy[nocc:] + mycc.level_shift
+    orbo = eris.mo_coeff[:,:nocc]
+    orbv = eris.mo_coeff[:,nocc:]
+
+    wpq, t1new, t2new, wVOov, wVooV = _direct_ovvv_vvvv(mycc, t1, t2)
+    t2new *= .5  # *.5 because t2+t2.transpose(1,0,3,2) at the end
+
+    _einsum = cupy.einsum
+
+    fov = fock[:nocc,nocc:].copy()
+    t1new += fock[:nocc,nocc:]
+
+    foo = fock[:nocc,:nocc] - np.diag(mo_e_o)
+    foo += .5 * np.einsum('ia,ja->ij', fock[:nocc,nocc:], t1)
+
+    fvv = lib.einsum('pa,qp,qb->ab', orbv, wpq, orbv)
+    t1new -= lib.einsum('ab,ib->ia', fvv, t1)
+
+    fvv += fock[nocc:,nocc:] - np.diag(mo_e_v)
+    fvv -= .5 * np.einsum('ia,ib->ab', t1, fock[:nocc,nocc:])
+
+    foo += lib.einsum('pi,qp,qj->ij', orbo, wpq, orbo)
+    fov += lib.einsum('pi,qp,qa->ia', orbo, wpq, orbv)
+
+    t1, t1_cpu = cupy.asarray(t1), t1
+    t2, t2_cpu = cupy.asarray(t2), t2
+    tau = _einsum('ia,jb->ijab', t1, t1)
+    tau += t2
+    woooo = _einsum('ijab,kabl->ijkl', tau, eris.ovvo)
+    woooo += cupy.asarray(eris.oooo).transpose(0,2,1,3)
+    tmp = _einsum('la,jaik->lkji', t1, eris.ovoo)
+    woooo += tmp
+    woooo += tmp.transpose(1,0,3,2)
+    t2new += .5 * _einsum('ijkl,klab->ijab', woooo, tau).get()
+    woooo = tau = None
+
+    wVOov -= lib.einsum('jbik,ka->bjia', eris.ovoo, t1_cpu)
+    t2new += wVOov.transpose(1,2,0,3)
+
+    wVooV += lib.einsum('kbij,ka->bija', eris.ovoo, t1_cpu)
+    wVooV -= eris.oovv.transpose(2,0,1,3)
+    wVOov += wVooV*.5  #: bjia + bija*.5
+    wVOov += eris.ovvo.transpose(2,3,0,1)
+
+    t2new += (eris.ovvo*0.5).transpose(0,3,1,2)
+    t1new += lib.einsum('pi,pq,qa->ia', orbo, wpq, orbv)
+
+    tmp  = lib.einsum('ic,kjbc->ikjb', t1_cpu, eris.oovv)
+    tmp += lib.einsum('jbck,ic->jkib', eris.ovvo, t1_cpu)
+    t2new -= lib.einsum('ka,jkib->jiba', t1_cpu, tmp)
+    tmp = None
+
+    tau  = t2 * .5
+    tau += _einsum('ia,jb->ijab', t1, t1)
+    wVooV += _einsum('kbci,jkca->bija', eris.ovvo, tau).get()
+    tau = None
+
+    tmp = _einsum('jkca,ckib->jaib', t2, wVooV).get()
+    t2new += tmp.transpose(2,0,1,3)
+    tmp *= .5
+    t2new += tmp.transpose(0,2,1,3)
+    tmp = None
+
+    tau  = np.einsum('ia,jb->iajb', t1_cpu*.5, t1_cpu)
+    tau += t2_cpu.transpose(0,2,1,3)
+    eris_ovOV = eris.ovvo.transpose(0,1,3,2) * 2
+    eris_ovOV -= eris.ovvo.transpose(3,1,0,2)
+    fvv -= lib.einsum('jcia,jcib->ab', tau, eris_ovOV)
+    foo += lib.einsum('iakb,jakb->ij', eris_ovOV, tau)
+
+    theta  = t2.transpose(0,2,1,3) * 2
+    theta -= t2.transpose(1,2,0,3)
+    tau = theta * .25
+    tau -= _einsum('ia,jb->jaib', t1*.5, t1)
+    wVOov += _einsum('kcia,kcjb->aijb', eris_ovOV, tau).get()
+    eris_ovOV = tau = None
+
+    t2new += _einsum('kcia,ckjb->ijab', theta, wVOov).get()
+    theta = wVOov = wVooV = None
+
+    t1new += np.einsum('jb,ijab->ia', fov, t2_cpu) * 2
+    t1new -= np.einsum('jb,ijba->ia', fov, t2_cpu)
+    ovoo = eris.ovoo * 2
+    ovoo -= eris.ovoo.transpose(2,1,0,3)
+    t1new -= lib.einsum('jbki,jkba->ia', ovoo, t2_cpu)
+    ovoo = None
+
+    ft_ij = foo + np.einsum('ja,ia->ij', .5*t1_cpu, fov)
+    ft_ab = fvv - np.einsum('ia,ib->ab', .5*t1_cpu, fov)
+    t2new += lib.einsum('ijac,bc->ijab', t2_cpu, ft_ab)
+    t2new -= lib.einsum('ki,kjab->ijab', ft_ij, t2_cpu)
+
+    eia = mo_e_o[:,None] - mo_e_v
+    t1new += lib.einsum('ib,ab->ia', t1_cpu, fvv)
+    t1new -= lib.einsum('ja,ji->ia', t1_cpu, foo)
+    t1new /= eia
+
+    t2new = t2new + t2new.transpose(1,0,3,2)
+    t2new /= eia[:,None,:,None] + eia[:,None,:]
+
+    time0 = log.timer_debug1('update t1 t2', *time0)
+    return t1new, t2new
+
+# Corresponds to the _add_vvvv_tril function in pyscf.cc.ccsd
+def _direct_ovvv_vvvv(mycc, t1, t2):
+    nocc, nvir = t1.shape
+    nocc2 = nocc*(nocc+1)//2
+    nao_cart = mycc.mol.nao_nr(cart=True)
+    max_memory = max(MEMORYMIN, mycc.max_memory - lib.current_memory()[0])
+    blksize = ((max_memory*.9e6-t2.size*4*8)/8/nao_cart**2/3.5)**.5
+    Ht2_mem = nocc2*nao_cart**2 * 8 * 2  # x2 and Ht2
+    mem_avail = int(cupy.cuda.runtime.memGetInfo()[0] * .75)
+    if mem_avail * .9 < Ht2_mem:
+        raise RuntimeError(
+            f'Not enough GPU memory. Available {mem_avail*1e-6} MB, required {Ht2_mem/.9e-6} MB')
+    # Reserve some memory for ERIs?
+    cupy.get_default_memory_pool().set_limit(mem_avail)
+
+    blksize = max(BLKMIN, int(min((nao_cart+3)/4, blksize,
+                                  ((mem_avail-Ht2_mem)*.5/8/nao_cart**2)**.5)))
+    logger.debug1(mycc, 'blksize %d nao %d', blksize, nao_cart)
+
+    vhfopt = gpu_hf._VHFOpt(mycc.mol, 'int2e')
+    vhfopt.build(group_size=blksize, diag_block_with_triu=True)
+    mol = vhfopt.mol
+
+    _einsum = cupy.einsum
+
+    mo = vhfopt.coeff.dot(cupy.asarray(mycc.mo_coeff))
+    orbo = cupy.asarray(mo[:,:nocc])
+    orbv = cupy.asarray(mo[:,nocc:])
+    t1po = orbv.dot(cupy.asarray(t1).T)
+    tau = make_tau_tril(t1, t2)
+    x2 = _einsum('xab,pa->xpb', tau, orbv)
+    x2 = _einsum('xpb,qb->xpq', x2, orbv)
+    tau = None
+
+    nao, nmo = mo.shape
+    ao_loc = mol.ao_loc
+    nao2 = nao * nao
+
+    x2 = cupy.asarray(x2, order='C')
+    Ht2ao = cupy.zeros_like(x2)
+    _dgemm = cupy.cuda.cublas.dgemm
+    handle = cupy.cuda.device.get_cublas_handle()
+    N = cupy.cuda.cublas.CUBLAS_OP_N
+    T = cupy.cuda.cublas.CUBLAS_OP_T
+    one = np.ones(1)
+    one_ptr = one.ctypes.data
+    x2_ptr = x2.data.ptr
+    Ht2ao_ptr = Ht2ao.data.ptr
+    def contract_vvvv_(eri, i0, i1, j0, j1):
+        ic = i1 - i0
+        jc = j1 - j0
+        eri = eri.reshape(-1,jc*nao)
+        #:Ht2[:,j0:j1] += np.einsum('xef,efab->xab', x2[:,i0:i1], eri)
+        _dgemm(handle, N, N, jc*nao, nocc2, ic*nao,
+               one_ptr, eri.data.ptr, jc*nao, x2_ptr+i0*nao*8, nao2,
+               one_ptr, Ht2ao_ptr+j0*nao*8, nao2)
+
+        if i0 > j0:
+            #:Ht2[:,i0:i1] += np.einsum('xef,abef->xab', x2[:,j0:j1], eri)
+            _dgemm(handle, T, N, ic*nao, nocc2, jc*nao,
+                   one_ptr, eri.data.ptr, jc*nao, x2_ptr+j0*nao*8, nao2,
+                   one_ptr, Ht2ao_ptr+i0*nao*8, nao2)
+
+    l_ctr_offsets = vhfopt.l_ctr_offsets
+    log_qs = vhfopt.log_qs
+    cp_idx, cp_jdx = np.tril_indices(len(vhfopt.uniq_l_ctr))
+
+    if vhfopt.uniq_l_ctr[:,0].max() <= gpu_hf.LMAX_ON_GPU:
+        # Computing ERIs on GPU
+        idx, idy = cupy.tril_indices(nao)
+        #eribuf = cupy.empty(blksize**2*nao**2)
+        def fint(ish0, ish1, jsh0, jsh1, group_id):
+            i0, i1 = ao_loc[ish0], ao_loc[ish1]
+            j0, j1 = ao_loc[jsh0], ao_loc[jsh1]
+            #eri = cupy.ndarray((i1-i0, nao, j1-j0, nao), memptr=eribuf.data)
+            #eri.fill(0.)
+            eri = cupy.zeros([i1-i0,nao,j1-j0,nao])
+
+            # strides to ensure data order consistent with eri(k1-k0,nao,l1-l0,nao)
+            strides = [1, (j1-j0)*nao, (j1-j0)*nao**2, nao]
+            ao_offsets = [0, 0, i0, j0]
+            _fill_eri_block(eri, strides, ao_offsets, vhfopt, group_id)
+            # Fill lower triangular part
+            eri[:,idx,:,idy] = eri[:,idy,:,idx]
+            return eri
+    else:
+        intor = mol._add_suffix('int2e')
+        ao2mopt = _ao2mo.AO2MOpt(mol, intor, 'CVHFnr_schwarz_cond',
+                                 'CVHFsetnr_direct_scf')
+        eribuf = np.empty((blksize,blksize,nao,nao))
+        loadbuf = np.empty((blksize,blksize,nao,nao))
+        def fint(ish0, ish1, jsh0, jsh1, group_id):
+            if ish0 != jsh0:
+                i0, i1 = ao_loc[ish0], ao_loc[ish1]
+                j0, j1 = ao_loc[jsh0], ao_loc[jsh1]
+                eri = gto.moleintor.getints4c(
+                    intor, mol._atm, mol._bas, mol._env,
+                    shls_slice=(ish0,ish1,jsh0,jsh1), aosym='s2kl',
+                    ao_loc=ao_loc, cintopt=ao2mopt._cintopt, out=eribuf)
+                aoblk = np.ndarray((i1-i0,nao,j1-j0,nao), buffer=loadbuf)
+                _ccsd.libcc.CCload_eri(aoblk.ctypes.data_as(ctypes.c_void_p),
+                                       eri.ctypes.data_as(ctypes.c_void_p),
+                                       (ctypes.c_int*4)(i0, i1, j0, j1),
+                                       ctypes.c_int(nao))
+            else:
+                i0, i1 = ao_loc[ish0], ao_loc[ish1]
+                eri = gto.moleintor.getints4c(
+                    intor, mol._atm, mol._bas, mol._env,
+                    shls_slice=(ish0,ish1,ish0,ish1), aosym='s4',
+                    ao_loc=ao_loc, cintopt=ao2mopt._cintopt, out=eribuf)
+                eri = lib.unpack_tril(eri, axis=0)
+                aoblk = np.ndarray((i1-i0,nao,i1-i0,nao), buffer=loadbuf)
+                _ccsd.libcc.CCload_eri(aoblk.ctypes.data_as(ctypes.c_void_p),
+                                       eri.ctypes.data_as(ctypes.c_void_p),
+                                       (ctypes.c_int*4)(i0, i1, i0, i1),
+                                       ctypes.c_int(nao))
+            return cupy.asarray(aoblk)
+
+    wVVoo = np.zeros((nao,nao,nocc,nocc))
+    wVvoO = np.zeros((nao,nao,nocc,nocc))
+
+    #mempool = cupy.get_default_memory_pool()
+    for cp_ij_id, log_q_ij in enumerate(log_qs):
+        cpi = cp_idx[cp_ij_id]
+        cpj = cp_jdx[cp_ij_id]
+        li = vhfopt.uniq_l_ctr[cpi,0]
+        lj = vhfopt.uniq_l_ctr[cpj,0]
+        if li > gpu_hf.LMAX_ON_GPU or lj > gpu_hf.LMAX_ON_GPU or log_q_ij.size == 0:
+            continue
+
+        ish0 = l_ctr_offsets[cpi]
+        jsh0 = l_ctr_offsets[cpj]
+        ish1 = l_ctr_offsets[cpi+1]
+        jsh1 = l_ctr_offsets[cpj+1]
+        aoblk = fint(ish0, ish1, jsh0, jsh1, cp_ij_id)
+
+        i0, i1 = ao_loc[ish0], ao_loc[ish1]
+        j0, j1 = ao_loc[jsh0], ao_loc[jsh1]
+        contract_vvvv_(aoblk, i0, i1, j0, j1)
+
+        #:fvv += 2*np.einsum('kc,kcab->ab', t1, eris_ovvv)
+        #:fvv -= np.einsum('kc,kbca->ab', t1, eris_ovvv)
+        pppo = _einsum('prqs,si->prqi', aoblk, orbo)
+        wVvoO[j0:j1] += _einsum('prqi,pj->qrij', pppo, t1po[i0:i1]).get()
+        wVVoo[i0:i1,j0:j1] = _einsum('prqi,rj->pqij', pppo, t1po).get()
+        pppo = None
+
+        if ish0 != jsh0:
+            wVVoo[j0:j1,i0:i1] = wVVoo[i0:i1,j0:j1].transpose(1,0,2,3)
+            #mempool.free_all_blocks()
+            wVvoO[i0:i1] += _einsum('prqs,ri,qj->psij',
+                                    aoblk, orbo, t1po[j0:j1]).get()
+        aoblk = None
+    eribuf = loadbuf = x2 = None
+
+    #:t1new += 2*lib.einsum('edac,ikcd->ikea', eris_ovvv, t2)
+    #:t1new +=  -lib.einsum('edac,ikdc->ikea', eris_ovvv, t2)
+    Ht2full = _unpack_t2_tril(Ht2ao, nocc, nao)
+    t1tmp  = _einsum('ijpq,qj->ip', Ht2full, orbo) * 2
+    t1tmp -= _einsum('ijqp,qj->ip', Ht2full, orbo)
+    t1new = t1tmp.dot(orbv).get()
+
+    # vvvv-t2 contractions back to MO repr.
+    Ht2tril = _einsum('xpq,pa->xaq', Ht2ao, orbv)
+    Ht2tril = _einsum('xaq,qb->xab', Ht2tril, orbv)
+
+    # part of ovvv-t2 contractions back to MO repr.
+    #: tmp = np.einsum('ijcd,ka,kdcb->ijba', tau, t1, eris.ovvv)
+    #: t2new -= tmp + tmp.transpose(1,0,3,2)
+    t1pv = orbo.dot(cupy.asarray(t1))
+    tmp = _einsum('xpq,pa->xaq', Ht2ao, orbv)
+    Ht2tril -= _einsum('xaq,qb->xab', tmp, t1pv)
+
+    tmp = _einsum('xpq,pa->xaq', Ht2ao, t1pv)
+    Ht2tril -= _einsum('xaq,qb->xab', tmp, orbv)#_einsum('xpq,pa,qb->xab', Ht2ao, t1pv, orbv)
+
+    t2new = _unpack_t2_tril(Ht2tril, nocc, nvir).get()
+    Ht2ao = Ht2full = None
+
+    c = vhfopt.coeff.get()
+    wpq = 2 * lib.einsum('pqkk,pi,qj->ij', wVVoo, c, c)
+    wpq -= lib.einsum('pqkk,pi,qj->ji', wVvoO, c, c)
+
+    tmp = _einsum('pqji,qb->pbji', cupy.asarray(wVvoO), orbv)
+    wVOov = _einsum('pbji,pa->bjia', tmp, orbv).get()
+    #wVOov = _einsum('pqji,qb,pa->bjia', cupy.asarray(wVvoO), orbv, orbv).get()
+
+    tmp = _einsum('pqji,pa->aqji', cupy.asarray(wVVoo), -orbv)
+    wVooV = _einsum('aqji,qb->bjia', tmp, orbv).get()
+    #wVooV = _einsum('pqji,pa,qb->bjia', cupy.asarray(wVVoo),-orbv, orbv).get()
+    wVVoo = None
+
+    if FREE_CUPY_CACHE:
+        cupy.get_default_memory_pool().free_all_blocks()
+    return wpq, t1new, t2new, wVOov, wVooV
+
+def make_tau_tril(t1, t2):
+    nocc, nvir = t1.shape
+    t1 = cupy.asarray(t1)
+    tau = cupy.einsum('ia,jb->ijab', t1, t1)
+    tau += cupy.asarray(t2)
+    return tau[cupy.tril_indices(nocc)]
+
+def _unpack_t2_tril(t2tril, nocc, nvir):
+    t2 = cupy.empty((nocc,nocc,nvir,nvir))
+    idx,idy = cupy.tril_indices(nocc)
+    t2[idy,idx] = t2tril.transpose(0,2,1)
+    t2[idx,idy] = t2tril
+    return t2
+
+def _fill_eri_block(eri, strides, ao_offsets, vhfopt, group_id):
+    log_qs = vhfopt.log_qs
+    cp_kl_id = group_id
+    log_q_kl = log_qs[cp_kl_id]
+    if log_q_kl.size == 0:
+        return eri
+
+    cp_idx, cp_jdx = np.tril_indices(len(vhfopt.uniq_l_ctr))
+    cpk = cp_idx[cp_kl_id]
+    cpl = cp_jdx[cp_kl_id]
+    lk = vhfopt.uniq_l_ctr[cpk,0]
+    ll = vhfopt.uniq_l_ctr[cpl,0]
+    if lk > gpu_hf.LMAX_ON_GPU or ll > gpu_hf.LMAX_ON_GPU:
+        raise NotImplementedError
+
+    stream = cupy.cuda.get_current_stream()
+    log_cutoff = np.log(vhfopt.direct_scf_tol)
+    omega = 0.
+
+    l_symb = lib.param.ANGULAR
+    nao = vhfopt.coeff.shape[0]
+    bins_locs_kl = vhfopt.bins[cp_kl_id]
+    bins_floor_kl = vhfopt.bins_floor[cp_kl_id]
+    nbins_kl = len(bins_locs_kl) - 1
+
+    fn = libgint.GINTfill_int2e
+    for cp_ij_id, log_q_ij in enumerate(log_qs):
+        cpi = cp_idx[cp_ij_id]
+        cpj = cp_jdx[cp_ij_id]
+        li = vhfopt.uniq_l_ctr[cpi,0]
+        lj = vhfopt.uniq_l_ctr[cpj,0]
+        if li > gpu_hf.LMAX_ON_GPU or lj > gpu_hf.LMAX_ON_GPU or log_q_ij.size == 0:
+            continue
+
+        t0 = time.perf_counter()
+        bins_locs_ij = vhfopt.bins[cp_ij_id]
+        bins_floor_ij = vhfopt.bins_floor[cp_ij_id]
+        nbins_ij = len(bins_locs_ij) - 1
+
+        err = fn(ctypes.cast(stream.ptr, ctypes.c_void_p), vhfopt.bpcache,
+                 ctypes.cast(eri.data.ptr, ctypes.c_void_p), ctypes.c_int(nao),
+                 (ctypes.c_int*4)(*strides), (ctypes.c_int*4)(*ao_offsets),
+                 bins_locs_ij.ctypes.data_as(ctypes.c_void_p),
+                 bins_locs_kl.ctypes.data_as(ctypes.c_void_p),
+                 bins_floor_ij.ctypes.data_as(ctypes.c_void_p),
+                 bins_floor_kl.ctypes.data_as(ctypes.c_void_p),
+                 ctypes.c_int(nbins_ij), ctypes.c_int(nbins_kl),
+                 ctypes.c_int(cp_ij_id), ctypes.c_int(cp_kl_id),
+                 ctypes.c_double(log_cutoff),
+                 ctypes.c_double(omega))
+        if err != 0:
+            detail = f'CUDA Error for ({l_symb[li]}{l_symb[lj]}|{l_symb[lk]}{l_symb[ll]})'
+            raise RuntimeError(detail)
+        logger.debug1(vhfopt.mol, '(%s%s|%s%s) on GPU %.3fs',
+                      l_symb[li], l_symb[lj], l_symb[lk], l_symb[ll],
+                      time.perf_counter() - t0)
+    return eri
+
+def _make_eris_incore(mycc, mo_coeff=None):
+    cput0 = (logger.process_clock(), logger.perf_counter())
+    log = logger.Logger(mycc.stdout, mycc.verbose)
+    eris = ccsd._ChemistsERIs()
+    eris._common_init_(mycc, mo_coeff)
+
+    # Cupy memory buffer may be created in previous SCF calculations.
+    if FREE_CUPY_CACHE:
+        cupy.get_default_memory_pool().free_all_blocks()
+
+    mol = mycc.mol
+    mo_coeff = cupy.asarray(eris.mo_coeff, order='F')
+    nocc = eris.nocc
+    nmo = mo_coeff.shape[1]
+    nvir = nmo - nocc
+
+    nao_cart = mycc.mol.nao_nr(cart=True)
+    max_memory = max(MEMORYMIN, mycc.max_memory - lib.current_memory()[0])
+    blksize = ((max_memory*.9e6-nocc**2*nao_cart**2*2*8)/8/nao_cart**2/2.5)**.5
+    mem_avail = int(cupy.cuda.runtime.memGetInfo()[0] * .75)
+    cupy.get_default_memory_pool().set_limit(mem_avail)
+    blksize = max(BLKMIN, int(min((nao_cart+3)/4, blksize,
+                                  (mem_avail*.5/8/nao_cart**2)**.5)))
+    logger.debug1(mycc, 'blksize %d nao %d', blksize, nao_cart)
+
+    vhfopt = gpu_hf._VHFOpt(mycc.mol, 'int2e')
+    vhfopt.build(group_size=blksize, diag_block_with_triu=True)
+    mol = vhfopt.mol
+    mo = vhfopt.coeff.dot(mo_coeff)
+    orbo = cupy.asarray(mo[:,:nocc])
+    orbv = cupy.asarray(mo[:,nocc:])
+    ao_loc = mol.ao_loc
+    nao = mo.shape[0]
+
+    l_ctr_offsets = vhfopt.l_ctr_offsets
+    log_qs = vhfopt.log_qs
+    cp_idx, cp_jdx = np.tril_indices(len(vhfopt.uniq_l_ctr))
+
+    ppOO = np.empty((nao,nao,nocc,nocc))
+    pPoO = np.zeros((nao,nao,nocc,nocc))
+    eribuf = cupy.empty(blksize**2*nao**2)
+    #mempool = cupy.get_default_memory_pool()
+    idx, idy = cupy.tril_indices(nao)
+
+    for cp_ij_id, log_q_ij in enumerate(log_qs):
+        cpi = cp_idx[cp_ij_id]
+        cpj = cp_jdx[cp_ij_id]
+        li = vhfopt.uniq_l_ctr[cpi,0]
+        lj = vhfopt.uniq_l_ctr[cpj,0]
+        if li > gpu_hf.LMAX_ON_GPU or lj > gpu_hf.LMAX_ON_GPU or log_q_ij.size == 0:
+            continue
+
+        ish0 = l_ctr_offsets[cpi]
+        jsh0 = l_ctr_offsets[cpj]
+        ish1 = l_ctr_offsets[cpi+1]
+        jsh1 = l_ctr_offsets[cpj+1]
+        i0, i1 = ao_loc[ish0], ao_loc[ish1]
+        j0, j1 = ao_loc[jsh0], ao_loc[jsh1]
+        eri = cupy.ndarray((nao, i1-i0, j1-j0, nao), memptr=eribuf.data)
+        eri.fill(0.)
+        # strides to ensure data order consistent with eri(nao,k1-k0,l1-l0,nao)
+        strides = [1, (i1-i0)*(j1-j0)*nao, (j1-j0)*nao, nao]
+        ao_offsets = [0, 0, i0, j0]
+        _fill_eri_block(eri, strides, ao_offsets, vhfopt, cp_ij_id)
+        # Fill lower triangular part
+        eri[idx,:,:,idy] = eri[idy,:,:,idx]
+
+        pijo = cupy.dot(eri.reshape(-1,nao), orbo)
+        ijoo = cupy.dot(pijo.reshape(nao,-1).T, orbo)
+        ppOO[i0:i1,j0:j1] = ijoo.get().reshape(i1-i0,j1-j0,nocc,nocc)
+        ijoo = None
+
+        jopi = cupy.asarray(pijo.reshape(nao*(i1-i0),(j1-j0)*nocc).T, order='C')
+        jopo = cupy.dot(jopi.reshape(-1,i1-i0), orbo[i0:i1])
+        pPoO[j0:j1] += jopo.get().reshape(j1-j0,nocc,nao,nocc).transpose(0,2,1,3)
+        pijo = jopo = None
+
+        if ish0 != jsh0:
+            ppOO[j0:j1,i0:i1] = ppOO[i0:i1,j0:j1].transpose(1,0,2,3)
+            opio = cupy.dot(jopi.reshape(j1-j0,-1).T, orbo[j0:j1])
+            pPoO[i0:i1] += opio.get().reshape(nocc,nao,i1-i0,nocc).transpose(2,1,0,3)
+            jopi = opio = None
+
+    ppOO = cupy.asarray(ppOO)
+    pooo = cupy.dot(ppOO.reshape(nao,-1).T, orbo)
+    oooo = cupy.dot(pooo.reshape(nao,-1).T, orbo).reshape(nocc,nocc,nocc,nocc)
+    ooov = cupy.dot(pooo.reshape(nao,-1).T, orbv).reshape(nocc,nocc,nocc,nvir)
+    eris.oooo = oooo.get()
+    eris.ovoo = lib.transpose(ooov.get().reshape(nocc*nocc,nocc*nvir)).reshape(nocc,nvir,nocc,nocc)
+    pooo = oooo = ooov = None
+
+    poov = cupy.dot(ppOO.reshape(nao,-1).T, orbv)
+    oovv = cupy.dot(poov.reshape(nao,-1).T, orbv).reshape(nocc,nocc,nvir,nvir)
+    eris.oovv = oovv.get()
+    ppOO = poov = oovv = None
+
+    pPoO = cupy.asarray(pPoO)
+    poov = cupy.dot(pPoO.reshape(nao,-1).T, orbv)
+    voov = cupy.dot(orbv.T, poov.reshape(nao,-1))
+    eris.ovvo = lib.transpose(voov.get().reshape(nvir*nocc,nocc*nvir)).reshape(nocc,nvir,nvir,nocc)
+    eris.ovov = eris.ovvo.transpose(0,1,3,2)
+    pPoO = poov = voov = None
+    log.timer('CCSD integral transformation', *cput0)
+
+    if FREE_CUPY_CACHE:
+        cupy.get_default_memory_pool().free_all_blocks()
+    return eris
+
+class CCSD(ccsd.CCSD):
+    from gpu4pyscf.lib.utils import to_cpu, to_gpu, device
+
+    def __init__(self, mf, *args, **kwargs):
+        if hasattr(mf, 'to_cpu'):
+            mf = mf.to_cpu()
+        if hasattr(mf, 'with_df') and mf.with_df:
+            lib.logger.warn(mf.mol, 'DF-CCSD not available. Run the standard CCSD.')
+        ccsd.CCSD.__init__(self, mf, *args, **kwargs)
+
+    update_amps = update_amps
+    ao2mo = _make_eris_incore

{gpu4pyscf-0.7.2 → gpu4pyscf-0.7.4}/gpu4pyscf/df/grad/uhf.py

@@ -254,6 +254,9 @@ def get_jk(mf_grad, mol=None, dm0=None, hermi=0, with_j=True, with_k=True, omega
 
 class Gradients(uhf_grad.Gradients):
     '''Unrestricted density-fitting Hartree-Fock gradients'''
+
+    _keys = {'with_df', 'auxbasis_response'}
+
     def __init__(self, mf):
         # Whether to include the response of DF auxiliary basis when computing
         # nuclear gradients of J/K matrices
@@ -262,7 +265,7 @@ class Gradients(uhf_grad.Gradients):
         uhf_grad.Gradients.__init__(self, mf)
 
     get_jk = get_jk
-    
+
     # TODO: finish these two functions
     def get_j(self, mol=None, dm=None, hermi=0, mo_coeff=None, mo_occ=None, dm2 = None):
         vj, _, vjaux, _ =  self.get_jk(mol, dm, with_k=False, mo_coeff=mo_coeff, mo_occ=mo_occ, dm2=dm2)
@@ -271,7 +274,7 @@ class Gradients(uhf_grad.Gradients):
     def get_k(self, mol=None, dm=None, hermi=0, mo_coeff=None, mo_occ=None, dm2 = None):
         _, vk, _, vkaux = self.get_jk(mol, dm, with_j=False, mo_coeff=mo_coeff, mo_occ=mo_occ, dm2=dm2)
         return vk, vkaux
-        
+
 
     def get_veff(self, mol=None, dm=None):
         vj0, vk0, vjaux0, vkaux0 = self.get_jk(mol, dm[0], mo_coeff=self.base.mo_coeff[0], mo_occ=self.base.mo_occ[0])

{gpu4pyscf-0.7.2 → gpu4pyscf-0.7.4}/gpu4pyscf/df/grad/uks.py

@@ -100,7 +100,7 @@ def get_veff(ks_grad, mol=None, dm=None):
         vxc += vj0 + vj1 + vj0_m1 + vj1_m0
     else:
         omega, alpha, hyb = ni.rsh_and_hybrid_coeff(mf.xc, spin=mol.spin)
-        
+
         vj0, vk0, vjaux0, vkaux0 = ks_grad.get_jk(mol, dm[0], mo_coeff=ks_grad.base.mo_coeff[0], mo_occ=ks_grad.base.mo_occ[0])
         vj1, vk1, vjaux1, vkaux1 = ks_grad.get_jk(mol, dm[1], mo_coeff=ks_grad.base.mo_coeff[1], mo_occ=ks_grad.base.mo_occ[1])
         vj0_m1, vjaux0_m1 = ks_grad.get_j(mol, dm[0], mo_coeff=ks_grad.base.mo_coeff[0], mo_occ=ks_grad.base.mo_occ[0], dm2=dm[1])
@@ -110,29 +110,32 @@ def get_veff(ks_grad, mol=None, dm=None):
         if ks_grad.auxbasis_response:
             vj_aux = vjaux0 + vjaux1 + vjaux0_m1 + vjaux1_m0
             vk_aux = (vkaux0+vkaux1) * hyb
-        
+
         if omega != 0:
             vk_lr0, vkaux_lr0 = ks_grad.get_k(mol, dm[0], mo_coeff=ks_grad.base.mo_coeff[0], mo_occ=ks_grad.base.mo_occ[0], omega=omega)
-            vk_lr1, vkaux_lr1 = ks_grad.get_k(mol, dm[1], mo_coeff=ks_grad.base.mo_coeff[1], mo_occ=ks_grad.base.mo_occ[1], omega=omega) 
+            vk_lr1, vkaux_lr1 = ks_grad.get_k(mol, dm[1], mo_coeff=ks_grad.base.mo_coeff[1], mo_occ=ks_grad.base.mo_occ[1], omega=omega)
             vk += (vk_lr0 + vk_lr1) * (alpha-hyb)
             if ks_grad.auxbasis_response:
                 vk_aux += (vkaux_lr0 + vkaux_lr1) * (alpha-hyb)
-                
+
         vxc += vj - vk
         if ks_grad.auxbasis_response:
             e1_aux = vj_aux - vk_aux
-            
+
     if ks_grad.auxbasis_response:
         logger.debug1(ks_grad, 'sum(auxbasis response) %s', e1_aux.sum(axis=0))
     else:
-        e1_aux = None   
-             
+        e1_aux = None
+
     vxc = tag_array(vxc, aux=e1_aux)
-    
+
     return vxc
 
 
 class Gradients(uks_grad.Gradients):
+
+    _keys = {'with_df', 'auxbasis_response'}
+
     def __init__(self, mf):
         # Whether to include the response of DF auxiliary basis when computing
         # nuclear gradients of J/K matrices
@@ -140,7 +143,7 @@ class Gradients(uks_grad.Gradients):
         uks_grad.Gradients.__init__(self, mf)
 
     get_jk = get_jk
-    
+
     def get_j(self, mol=None, dm=None, hermi=0, mo_coeff=None, mo_occ=None, dm2 = None, omega=None):
         vj, _, vjaux, _ = self.get_jk(mol, dm, with_k=False, mo_coeff=mo_coeff, mo_occ=mo_occ, dm2=dm2, omega=omega)
         return vj, vjaux
@@ -148,7 +151,7 @@ class Gradients(uks_grad.Gradients):
     def get_k(self, mol=None, dm=None, hermi=0, mo_coeff=None, mo_occ=None, dm2 = None, omega=None):
         _, vk, _, vkaux = self.get_jk(mol, dm, with_j=False, mo_coeff=mo_coeff, mo_occ=mo_occ, dm2=dm2, omega=omega)
         return vk, vkaux
-    
+
     get_veff = get_veff
 
     def extra_force(self, atom_id, envs):

{gpu4pyscf-0.7.2 → gpu4pyscf-0.7.4}/gpu4pyscf/scf/uhf.py

@@ -113,7 +113,7 @@ def get_fock(mf, h1e=None, s1e=None, vhf=None, dm=None, cycle=-1, diis=None,
 class UHF(uhf.UHF):
     from gpu4pyscf.lib.utils import to_cpu, to_gpu, device
 
-    _keys = {'e_disp', 'screen_tol', 'conv_tol_cpscf'}
+    _keys = {'e_disp', 'screen_tol', 'conv_tol_cpscf', 'h1e', 's1e'}
 
     screen_tol = 1e-14
     conv_tol_cpscf = 1e-3

{gpu4pyscf-0.7.2 → gpu4pyscf-0.7.4}/gpu4pyscf/solvent/grad/pcm.py

@@ -344,7 +344,8 @@ class WithSolventGrad:
         dm = kwargs.pop('dm', None)
         if dm is None:
             dm = self.base.make_rdm1(ao_repr=True)
-
+        if dm.ndim == 3:
+            dm = dm[0] + dm[1]
         self.de_solute = super().kernel(*args, **kwargs)
         self.de_solvent = grad_qv(self.base.with_solvent, dm)
         self.de_solvent+= grad_solver(self.base.with_solvent, dm)

{gpu4pyscf-0.7.2 → gpu4pyscf-0.7.4}/gpu4pyscf/solvent/grad/smd.py

@@ -255,14 +255,15 @@ class WithSolventGrad:
         dm = kwargs.pop('dm', None)
         if dm is None:
             dm = self.base.make_rdm1(ao_repr=True)
-
-        self.de_solute = super().kernel(*args, **kwargs)
+        if dm.ndim == 3:
+            dm = dm[0] + dm[1]
+        self.de_solute  = super().kernel(*args, **kwargs)
         self.de_solvent = pcm_grad.grad_qv(self.base.with_solvent, dm)
         self.de_solvent+= pcm_grad.grad_solver(self.base.with_solvent, dm)
         self.de_solvent+= pcm_grad.grad_nuc(self.base.with_solvent, dm)
+        self.de_cds     = get_cds(self.base.with_solvent)
+        self.de = self.de_solute + self.de_solvent + self.de_cds
 
-        self.de = self.de_solute + self.de_solvent
-        self.de += get_cds(self.base.with_solvent)
         if self.verbose >= logger.NOTE:
             logger.note(self, '--------------- %s (+%s) gradients ---------------',
                         self.base.__class__.__name__,

{gpu4pyscf-0.7.2 → gpu4pyscf-0.7.4}/gpu4pyscf/solvent/hessian/pcm.py

@@ -21,8 +21,7 @@ Gradient of PCM family solvent model
 import numpy
 import cupy
 from cupyx import scipy
-from pyscf import lib
-from pyscf import gto, df
+from pyscf import lib, scf, gto, df
 from pyscf.grad import rhf as rhf_grad
 from gpu4pyscf.solvent.pcm import PI, switch_h
 from gpu4pyscf.solvent.grad.pcm import grad_switch_h, get_dF_dA, get_dD_dS, grad_qv, grad_solver, grad_nuc
@@ -155,7 +154,7 @@ def hess_elec(pcmobj, dm, verbose=None):
     pcmobj.reset(pmol)
     return de
 
-def fd_grad_vmat(pcmobj, mo_coeff, mo_occ, atmlst=None, verbose=None):
+def fd_grad_vmat(pcmobj, dm, mo_coeff, mo_occ, atmlst=None, verbose=None):
     '''
     dv_solv / da
     slow version with finite difference
@@ -169,7 +168,6 @@ def fd_grad_vmat(pcmobj, mo_coeff, mo_occ, atmlst=None, verbose=None):
     nao, nmo = mo_coeff.shape
     mocc = mo_coeff[:,mo_occ>0]
     nocc = mocc.shape[1]
-    dm = cupy.dot(mocc, mocc.T) * 2
     coords = mol.atom_coords(unit='Bohr')
     def pcm_vmat_scanner(mol):
         pcmobj.reset(mol)
@@ -250,6 +248,8 @@ class WithSolventHess:
         dm = kwargs.pop('dm', None)
         if dm is None:
             dm = self.base.make_rdm1(ao_repr=True)
+        if dm.ndim == 3:
+            dm = dm[0] + dm[1]
         is_equilibrium = self.base.with_solvent.equilibrium_solvation
         self.base.with_solvent.equilibrium_solvation = True
         self.de_solvent = hess_elec(self.base.with_solvent, dm, verbose=self.verbose)
@@ -263,11 +263,25 @@ class WithSolventHess:
         if atmlst is None:
             atmlst = range(self.mol.natm)
         h1ao = super().make_h1(mo_coeff, mo_occ, atmlst=atmlst, verbose=verbose)
-        dv = fd_grad_vmat(self.base.with_solvent, mo_coeff, mo_occ, atmlst=atmlst, verbose=verbose)
-        for i0, ia in enumerate(atmlst):
-            h1ao[i0] += dv[i0]
-        return h1ao
-
+        if isinstance(self.base, scf.hf.RHF):
+            dm = self.base.make_rdm1(ao_repr=True)
+            dv = fd_grad_vmat(self.base.with_solvent, dm, mo_coeff, mo_occ, atmlst=atmlst, verbose=verbose)
+            for i0, ia in enumerate(atmlst):
+                h1ao[i0] += dv[i0]
+            return h1ao
+        elif isinstance(self.base, scf.uhf.UHF):
+            h1aoa, h1aob = h1ao
+            solvent = self.base.with_solvent
+            dm = self.base.make_rdm1(ao_repr=True)
+            dm = dm[0] + dm[1]
+            dva = fd_grad_vmat(solvent, dm, mo_coeff[0], mo_occ[0], atmlst=atmlst, verbose=verbose)
+            dvb = fd_grad_vmat(solvent, dm, mo_coeff[1], mo_occ[1], atmlst=atmlst, verbose=verbose)
+            for i0, ia in enumerate(atmlst):
+                h1aoa[i0] += dva[i0]
+                h1aob[i0] += dvb[i0]
+            return h1aoa, h1aob
+        else:
+            raise NotImplementedError('Base object is not supported')
     def _finalize(self):
         # disable _finalize. It is called in grad_method.kernel method
         # where self.de was not yet initialized.

{gpu4pyscf-0.7.2 → gpu4pyscf-0.7.4}/gpu4pyscf/solvent/hessian/smd.py

@@ -20,9 +20,7 @@ Gradient of PCM family solvent model
 
 import numpy as np
 import cupy
-#from cupyx import scipy, jit
-from pyscf import lib
-from pyscf import gto, df
+from pyscf import gto, df, scf, lib
 from pyscf.grad import rhf as rhf_grad
 from pyscf.data import radii
 from pyscf.solvent import ddcosmo_grad
@@ -278,13 +276,15 @@ class WithSolventHess:
         dm = kwargs.pop('dm', None)
         if dm is None:
             dm = self.base.make_rdm1(ao_repr=True)
+        if dm.ndim == 3:
+            dm = dm[0] + dm[1]
         is_equilibrium = self.base.with_solvent.equilibrium_solvation
         self.base.with_solvent.equilibrium_solvation = True
         self.de_solvent = pcm_hess.hess_elec(self.base.with_solvent, dm, verbose=self.verbose)
         #self.de_solvent+= hess_nuc(self.base.with_solvent)
         self.de_solute = super().kernel(*args, **kwargs)
-        self.de = self.de_solute + self.de_solvent
-        self.de += get_cds(self.base.with_solvent)
+        self.de_cds = get_cds(self.base.with_solvent)
+        self.de = self.de_solute + self.de_solvent + self.de_cds
         self.base.with_solvent.equilibrium_solvation = is_equilibrium
         return self.de
 
@@ -292,11 +292,25 @@ class WithSolventHess:
         if atmlst is None:
             atmlst = range(self.mol.natm)
         h1ao = super().make_h1(mo_coeff, mo_occ, atmlst=atmlst, verbose=verbose)
-        dv = pcm_hess.fd_grad_vmat(self.base.with_solvent, mo_coeff, mo_occ, atmlst=atmlst, verbose=verbose)
-        for i0, ia in enumerate(atmlst):
-            h1ao[i0] += dv[i0]
-        return h1ao
-
+        if isinstance(self.base, scf.hf.RHF):
+            dm = self.base.make_rdm1(ao_repr=True)
+            dv = pcm_hess.fd_grad_vmat(self.base.with_solvent, dm, mo_coeff, mo_occ, atmlst=atmlst, verbose=verbose)
+            for i0, ia in enumerate(atmlst):
+                h1ao[i0] += dv[i0]
+            return h1ao
+        elif isinstance(self.base, scf.uhf.UHF):
+            h1aoa, h1aob = h1ao
+            solvent = self.base.with_solvent
+            dm = self.base.make_rdm1(ao_repr=True)
+            dm = dm[0] + dm[1]
+            dva = pcm_hess.fd_grad_vmat(solvent, dm, mo_coeff[0], mo_occ[0], atmlst=atmlst, verbose=verbose)
+            dvb = pcm_hess.fd_grad_vmat(solvent, dm, mo_coeff[1], mo_occ[1], atmlst=atmlst, verbose=verbose)
+            for i0, ia in enumerate(atmlst):
+                h1aoa[i0] += dva[i0]
+                h1aob[i0] += dvb[i0]
+            return h1aoa, h1aob
+        else:
+            raise NotImplementedError('Base object is not supported')
     def _finalize(self):
         # disable _finalize. It is called in grad_method.kernel method
         # where self.de was not yet initialized.

{gpu4pyscf-0.7.2 → gpu4pyscf-0.7.4}/gpu4pyscf/solvent/pcm.py

@@ -310,7 +310,6 @@ class PCM(ddcosmo.DDCOSMO):
     def _get_vind(self, dms):
         if not self._intermediates:
             self.build()
-
         nao = dms.shape[-1]
         dms = dms.reshape(-1,nao,nao)
         if dms.shape[0] == 2:
@@ -335,7 +334,6 @@ class PCM(ddcosmo.DDCOSMO):
         self._intermediates['q'] = q[0]
         self._intermediates['q_sym'] = q_sym[0]
         self._intermediates['v_grids'] = v_grids[0]
-
         return epcm, vmat[0]
 
     def _get_v(self, dms):
@@ -363,7 +361,7 @@ class PCM(ddcosmo.DDCOSMO):
         if self.frozen:
             raise RuntimeError('Frozen solvent model is not supported')
         from gpu4pyscf import scf
-        if isinstance(grad_method.base, scf.hf.RHF):
+        if isinstance(grad_method.base, (scf.hf.RHF, scf.uhf.UHF)):
             return pcm_grad.make_grad_object(grad_method)
         else:
             raise RuntimeError('Only SCF gradient is supported')
@@ -373,7 +371,7 @@ class PCM(ddcosmo.DDCOSMO):
         if self.frozen:
             raise RuntimeError('Frozen solvent model is not supported')
         from gpu4pyscf import scf
-        if isinstance(hess_method.base, scf.hf.RHF):
+        if isinstance(hess_method.base, (scf.hf.RHF, scf.uhf.UHF)):
             return pcm_hess.make_hess_object(hess_method)
         else:
             raise RuntimeError('Only SCF gradient is supported')
@@ -387,7 +385,7 @@ class PCM(ddcosmo.DDCOSMO):
     def _B_dot_x(self, dms):
         if not self._intermediates:
             self.build()
-
+        out_shape = dms.shape
         nao = dms.shape[-1]
         dms = dms.reshape(-1,nao,nao)
 
@@ -403,6 +401,5 @@ class PCM(ddcosmo.DDCOSMO):
         q_sym = (q + qt)/2.0
 
         vmat = self._get_vmat(q_sym)
-
-        return vmat
+        return vmat.reshape(out_shape)
 

{gpu4pyscf-0.7.2 → gpu4pyscf-0.7.4}/gpu4pyscf/solvent/smd.py

@@ -572,7 +572,7 @@ class SMD(pcm.PCM):
         if self.frozen:
             raise RuntimeError('Frozen solvent model is not supported')
         from gpu4pyscf import scf
-        if isinstance(grad_method.base, scf.hf.RHF):
+        if isinstance(grad_method.base, (scf.hf.RHF, scf.uhf.UHF)):
             return smd_grad.make_grad_object(grad_method)
         else:
             raise RuntimeError('Only SCF gradient is supported')
@@ -582,7 +582,7 @@ class SMD(pcm.PCM):
         if self.frozen:
             raise RuntimeError('Frozen solvent model is not supported')
         from gpu4pyscf import scf
-        if isinstance(hess_method.base, scf.hf.RHF):
+        if isinstance(hess_method.base, (scf.hf.RHF, scf.uhf.UHF)):
             return smd_hess.make_hess_object(hess_method)
         else:
             raise RuntimeError('Only SCF gradient is supported')

{gpu4pyscf-0.7.2 → gpu4pyscf-0.7.4/gpu4pyscf.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: gpu4pyscf
-Version: 0.7.2
+Version: 0.7.4
 Summary: GPU extensions for PySCF
 Home-page: UNKNOWN
 Author: Qiming Sun

{gpu4pyscf-0.7.2 → gpu4pyscf-0.7.4}/gpu4pyscf.egg-info/SOURCES.txt

@@ -10,6 +10,8 @@ gpu4pyscf.egg-info/SOURCES.txt
 gpu4pyscf.egg-info/dependency_links.txt
 gpu4pyscf.egg-info/requires.txt
 gpu4pyscf.egg-info/top_level.txt
+gpu4pyscf/cc/__init__.py
+gpu4pyscf/cc/ccsd_incore.py
 gpu4pyscf/df/__init__.py
 gpu4pyscf/df/cderi.py
 gpu4pyscf/df/df.py