feat: implement GPU UCCSD(T)

cucc/README.md

@@ -3,3 +3,9 @@
 This package follows the interfaces in PySCF packages, and some codes are adapted from it to keep the same interface. The examples are in the `byteqc/cucc/test` folder.
 
 The `buffer.py` and `culib.py` modules enable automatic backend determination. The arrays in `cucc` can be stored on the GPU, CPU, or disk, all with the same interface. Users and developers can work with these arrays without having to consider the underlying backend.
+
+`cucc.CCSD` accepts PySCF UHF references and returns unrestricted amplitudes
+in the usual `(t1a, t1b)` and `(t2aa, t2ab, t2bb)` form.  The UCCSD and
+UCCSD(T) contractions run on the GPU in an antisymmetrized spin-orbital
+representation.  The initial implementation keeps transformed integrals
+in-core and checks their footprint against `gpulim * mem_ratio` before upload.

cucc/__init__.py

@@ -13,23 +13,30 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from pyscf import scf
-from byteqc.cucc import ccsd
-from byteqc.cucc import dfccsd
 from pyscf.lib import param
 
 
 def CCSD(mf, frozen=None, mo_coeff=None, mo_occ=None, gpulim=None,
          cpulim=None, pool=None, path=param.TMPDIR, mem_ratio=0.65):
-    if isinstance(mf, scf.uhf.UHF) or isinstance(mf, scf.ghf.GHF):
-        AssertionError('Not implement')
-    else:
-        return RCCSD(mf, frozen, mo_coeff, mo_occ, gpulim, cpulim,
+    if mf.istype('UHF'):
+        return UCCSD(mf, frozen, mo_coeff, mo_occ, gpulim, cpulim,
                      pool=pool, path=path, mem_ratio=mem_ratio)
+    if mf.istype('GHF'):
+        raise NotImplementedError('GHF-CCSD is not implemented')
+    return RCCSD(mf, frozen, mo_coeff, mo_occ, gpulim, cpulim,
+                 pool=pool, path=path, mem_ratio=mem_ratio)
+
+
+def UCCSD(mf, frozen=None, mo_coeff=None, mo_occ=None, gpulim=None,
+          cpulim=None, pool=None, path=param.TMPDIR, mem_ratio=0.65):
+    from byteqc.cucc import uccsd
+    return uccsd.UCCSD(mf, frozen, mo_coeff, mo_occ, gpulim, cpulim,
+                       pool=pool, path=path, mem_ratio=mem_ratio)
 
 
 def RCCSD(mf, frozen=None, mo_coeff=None, mo_occ=None, gpulim=None,
           cpulim=None, pool=None, path=param.TMPDIR, mem_ratio=0.65):
+    from byteqc.cucc import ccsd, dfccsd
     if getattr(mf, 'with_df', None):
         return dfccsd.RCCSD(mf, frozen, mo_coeff, mo_occ, gpulim,
                             cpulim, pool=pool, path=path, mem_ratio=mem_ratio)

cucc/test/test_uccsd.py

@@ -0,0 +1,60 @@
+# Copyright (c) 2024 Bytedance Ltd. and/or its affiliates
+# This file is part of ByteQC.
+#
+# 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
+#
+# https: // www.apache.org/licenses/LICENSE-2.0
+
+import unittest
+
+import cupy
+import numpy
+from byteqc import cucc
+from pyscf import cc, gto, scf
+
+
+class UCCSDTest(unittest.TestCase):
+
+    @classmethod
+    def setUpClass(cls):
+        mol = gto.M(
+            atom='Li 0 0 0; H 0 0 1.6',
+            basis='sto-3g',
+            charge=1,
+            spin=1,
+            verbose=0,
+        )
+        cls.mf = scf.UHF(mol).run(conv_tol=1e-12)
+
+    def test_uccsd_and_triples_match_pyscf(self):
+        ref = cc.UCCSD(self.mf).run(conv_tol=1e-10,
+                                    conv_tol_normt=1e-8)
+        ref_et = ref.ccsd_t()
+
+        gpu = cucc.CCSD(self.mf)
+        gpu.conv_tol = 1e-10
+        gpu.conv_tol_normt = 1e-8
+        gpu.kernel()
+        gpu_et = gpu.ccsd_t()
+
+        self.assertTrue(gpu.converged)
+        self.assertAlmostEqual(gpu.e_corr, ref.e_corr, 9)
+        self.assertAlmostEqual(gpu_et, ref_et, 9)
+        self.assertEqual(len(gpu.t1), 2)
+        self.assertEqual(len(gpu.t2), 3)
+        for actual, expected in zip(gpu.t1, ref.t1):
+            numpy.testing.assert_allclose(cupy.asnumpy(actual), expected,
+                                          atol=1e-8, rtol=1e-7)
+        for actual, expected in zip(gpu.t2, ref.t2):
+            numpy.testing.assert_allclose(cupy.asnumpy(actual), expected,
+                                          atol=1e-8, rtol=1e-7)
+
+    def test_ghf_failure_is_explicit(self):
+        with self.assertRaisesRegex(NotImplementedError, 'GHF-CCSD'):
+            cucc.CCSD(self.mf.to_ghf())
+
+
+if __name__ == '__main__':
+    unittest.main()