Add GPU density-fitted CASCI and diagonal-Hessian CASSCF

examples/post_HF/21-df-casscf.py

@@ -0,0 +1,40 @@
+#!/usr/bin/env python
+# Copyright 2026 The PySCF Developers. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+########################################
+#  Example of GPU DF-CASSCF
+########################################
+
+import pyscf
+from pyscf import mcscf as cpu_mcscf
+
+from gpu4pyscf import mcscf as gpu_mcscf
+
+
+mol = pyscf.M(
+    atom='N 0 0 -0.7; N 0 0 0.7',
+    basis='6-31g*',
+)
+mf = mol.RHF().to_gpu().density_fit().run()
+
+mc_cpu = cpu_mcscf.DFCASSCF(mf.to_cpu(), 6, 6)
+e_cpu = mc_cpu.kernel()[0]
+
+mc_gpu = gpu_mcscf.DFCASSCF(mf.to_gpu(), 6, 6)
+e_gpu = mc_gpu.kernel()[0]
+
+print(f'CPU DF-CASSCF energy: {e_cpu:.12f}')
+print(f'GPU DF-CASSCF energy: {e_gpu:.12f}')
+print(f'CPU-GPU difference:   {e_gpu - e_cpu:.3e}')

gpu4pyscf/df/df.py

@@ -20,9 +20,11 @@
 import cupy as cp
 from cupyx.scipy.linalg import solve_triangular
 from pyscf import lib
+from pyscf.ao2mo.incore import iden_coeffs
 from pyscf.df import addons, incore
 from gpu4pyscf.lib.cupy_helper import (
-    cholesky, get_avail_mem, fill_symmetric, asarray, empty_mapped, ndarray)
+    cholesky, contract, get_avail_mem, fill_symmetric, asarray, empty_mapped,
+    ndarray)
 from gpu4pyscf.lib import logger
 from gpu4pyscf.lib import utils
 from gpu4pyscf.lib import multi_gpu
@@ -282,6 +284,61 @@ def loop(self, blksize=None, unpack=True):
                                          out=work[:,:,:p1-p0])
                     yield out.transpose(2,0,1), cderi_blk
 
+    def ao2mo(self, mo_coeffs, compact=True):
+        '''Transform density-fitting integrals to MO integrals on GPU.'''
+        if isinstance(mo_coeffs, (np.ndarray, cupy.ndarray)):
+            if mo_coeffs.ndim != 2:
+                raise ValueError('mo_coeffs must be a 2D array or four 2D arrays')
+            mo_coeffs = (mo_coeffs,) * 4
+        elif len(mo_coeffs) != 4:
+            raise ValueError('mo_coeffs must contain four coefficient matrices')
+        else:
+            mo_coeffs = tuple(mo_coeffs)
+
+        if any(mo.ndim != 2 or mo.shape[0] != self.mol.nao
+               for mo in mo_coeffs):
+            raise ValueError('MO coefficients must have shape (mol.nao, nmo)')
+        if any(cupy.iscomplexobj(mo) for mo in mo_coeffs):
+            raise NotImplementedError('GPU DF AO2MO does not support complex orbitals')
+
+        ij_compact = compact and iden_coeffs(mo_coeffs[0], mo_coeffs[1])
+        kl_compact = compact and iden_coeffs(mo_coeffs[2], mo_coeffs[3])
+        same_pairs = (iden_coeffs(mo_coeffs[0], mo_coeffs[2]) and
+                      iden_coeffs(mo_coeffs[1], mo_coeffs[3]))
+
+        if self._cderi is None:
+            self.build()
+        mo_coeffs = [cupy.asarray(mo, dtype=cupy.float64)
+                     for mo in mo_coeffs]
+        ni, nj, nk, nl = [mo.shape[1] for mo in mo_coeffs]
+        nij = ni * (ni + 1) // 2 if ij_compact else ni * nj
+        nkl = nk * (nk + 1) // 2 if kl_compact else nk * nl
+
+        def transform():
+            coeffs = [asarray(mo) for mo in mo_coeffs]
+            eri = cupy.zeros((nij, nkl))
+
+            def transform_pair(cderi, left, right, compact):
+                buf = contract('Luv,up->Lpv', cderi, left)
+                out = contract('Lpv,vq->Lpq', buf, right)
+                if compact:
+                    idx = cp.tril_indices(left.shape[1])
+                    return out[:, idx[0], idx[1]]
+                return out.reshape(len(cderi), -1)
+
+            for cderi, _ in self.loop(unpack=True):
+                lij = transform_pair(
+                    cderi, coeffs[0], coeffs[1], ij_compact)
+                lkl = lij if same_pairs else transform_pair(
+                    cderi, coeffs[2], coeffs[3], kl_compact)
+                contract('Li,Lj->ij', lij, lkl, beta=1, out=eri)
+            return eri
+
+        eri = multi_gpu.run(transform, non_blocking=True)
+        return multi_gpu.array_reduce(eri, inplace=True)
+
+    get_mo_eri = ao2mo
+
     def reset(self, mol=None):
         '''Reset mol and clean up relevant attributes for scanner mode'''
         if mol is not None:
@@ -336,7 +393,6 @@ def range_coulomb(self, omega):
                 auxmol.omega = auxmol_omega
 
     get_ao_eri = get_eri = NotImplemented
-    get_mo_eri = ao2mo = NotImplemented
 
 def cholesky_eri_gpu(intopt, mol, auxmol, cd_low,
                      omega=None, sr_only=False, use_gpu_memory=True):

gpu4pyscf/df/tests/test_df_ao2mo.py

@@ -0,0 +1,63 @@
+# Copyright 2026 The PySCF Developers. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import pyscf
+from pyscf import df as cpu_df
+from pyscf import scf
+
+from gpu4pyscf import df
+
+
+class KnownValues(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.mol = pyscf.M(
+            atom='O 0 0 0; H 0 -0.757 0.587; H 0 0.757 0.587',
+            basis='cc-pvdz', verbose=0, output='/dev/null')
+        cls.mo = scf.RHF(cls.mol).run(conv_tol=1e-12).mo_coeff
+        cls.cpu_df = cpu_df.DF(cls.mol, auxbasis='weigend').build()
+        cls.gpu_df = df.DF(cls.mol, auxbasis='weigend').build()
+
+    @classmethod
+    def tearDownClass(cls):
+        cls.mol.stdout.close()
+
+    def test_compact(self):
+        ref = self.cpu_df.ao2mo(self.mo)
+        out = self.gpu_df.ao2mo(self.mo).get()
+        self.assertEqual(out.shape, ref.shape)
+        self.assertLess(abs(out - ref).max(), 1e-8)
+
+    def test_general(self):
+        coeffs = (self.mo[:, :3], self.mo[:, 1:5],
+                  self.mo[:, 2:6], self.mo[:, :2])
+        ref = self.cpu_df.ao2mo(coeffs, compact=False)
+        out = self.gpu_df.ao2mo(coeffs, compact=False).get()
+        self.assertEqual(out.shape, ref.shape)
+        self.assertLess(abs(out - ref).max(), 1e-8)
+
+    def test_host_cderi(self):
+        with_df = df.DF(self.mol, auxbasis='weigend')
+        with_df.use_gpu_memory = False
+        with_df.build()
+        ref = self.cpu_df.ao2mo(self.mo[:, 2:9])
+        out = with_df.ao2mo(self.mo[:, 2:9]).get()
+        self.assertLess(abs(out - ref).max(), 1e-8)
+
+
+if __name__ == '__main__':
+    print('Full tests for GPU DF AO2MO')
+    unittest.main()

gpu4pyscf/fci/__init__.py

@@ -0,0 +1,24 @@
+# Copyright 2026 The PySCF Developers. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from gpu4pyscf.fci.direct_spin1 import FCI
+from gpu4pyscf.fci.direct_spin1 import FCISolver
+
+
+def solver(mol=None, singlet=False, symm=None):
+    if singlet:
+        raise NotImplementedError('GPU spin-adapted FCI is not implemented')
+    if symm or (symm is None and mol is not None and mol.symmetry):
+        raise NotImplementedError('GPU symmetry-adapted FCI is not implemented')
+    return FCISolver(mol)