From 364e2b25319fefdcf394dbd155cc88fff21bb65c Mon Sep 17 00:00:00 2001
From: Qiming Sun <osirpt.sun@gmail.com>
Date: Sun, 24 Nov 2024 21:45:53 -0800
Subject: [PATCH] Fix the CCSD 2-rdm for complex orbitals

---
 pyscf/cc/ccsd_rdm.py        |  7 ++++++-
 pyscf/cc/test/test_gccsd.py | 15 +++++++++------
 pyscf/cc/uccsd_rdm.py       |  4 ++--
 3 files changed, 17 insertions(+), 9 deletions(-)

diff --git a/pyscf/cc/ccsd_rdm.py b/pyscf/cc/ccsd_rdm.py
index de1b38f2b4..9353dcf3fe 100644
--- a/pyscf/cc/ccsd_rdm.py
+++ b/pyscf/cc/ccsd_rdm.py
@@ -377,8 +377,13 @@ def _make_rdm2(mycc, d1, d2, with_dm1=True, with_frozen=True, ao_repr=False):
 
 
 def _rdm2_mo2ao(dm2, mo):
+    '''
+    Back transform the two-particle density matrices to AO representation, where
+    the dm2 is defined in accordance with the chemist's ERI notation:
+    E = einsum('pqrs,pqrs', eri, rdm2)
+    '''
     mo_C = mo.conj()
-    return lib.einsum('ijkl,pi,qj,rk,sl->pqrs', dm2, mo, mo_C, mo, mo_C)
+    return lib.einsum('ijkl,pi,qj,rk,sl->pqrs', dm2, mo_C, mo, mo_C, mo)
 
 
 if __name__ == '__main__':
diff --git a/pyscf/cc/test/test_gccsd.py b/pyscf/cc/test/test_gccsd.py
index 393e8bbced..fffbabf1b1 100644
--- a/pyscf/cc/test/test_gccsd.py
+++ b/pyscf/cc/test/test_gccsd.py
@@ -285,6 +285,7 @@ def test_rdm_complex(self):
         eri = eri + eri.transpose(1,0,3,2).conj()
         eri = eri + eri.transpose(2,3,0,1)
         eri *= .1
+        mf._eri = eri
 
         def get_jk(mol, dm, *args,**kwargs):
             vj = numpy.einsum('ijkl,lk->ij', eri, dm)
@@ -294,17 +295,19 @@ def get_veff(mol, dm, *args, **kwargs):
             vj, vk = get_jk(mol, dm)
             return vj - vk
         def ao2mofn(mos):
-            return eri
+            c = mos
+            return lib.einsum('pqrs,pi,qj,rk,sl->ijkl', eri, c.conj(), c, c.conj(), c)
 
         mf.get_jk = get_jk
         mf.get_veff = get_veff
         hcore = numpy.random.random((nmo,nmo)) * .2 + numpy.random.random((nmo,nmo))* .2j
-        hcore = hcore + hcore.T.conj() + numpy.diag(range(nmo))*2
+        hcore = hcore + hcore.T.conj() + numpy.diag(numpy.arange(nmo)*2)
         mf.get_hcore = lambda *args: hcore
         mf.get_ovlp = lambda *args: numpy.eye(nmo)
         orbspin = numpy.zeros(nmo, dtype=int)
         orbspin[1::2] = 1
-        mf.mo_coeff = lib.tag_array(numpy.eye(nmo) + 0j, orbspin=orbspin)
+        u = numpy.linalg.eigh(hcore)[1]
+        mf.mo_coeff = lib.tag_array(u, orbspin=orbspin)
         mf.mo_occ = numpy.zeros(nmo)
         mf.mo_occ[:nocc] = 1
 
@@ -313,12 +316,12 @@ def ao2mofn(mos):
         mycc.ao2mo = lambda *args, **kwargs: eris
         mycc.kernel(eris=eris)
         mycc.solve_lambda(eris=eris)
-        dm1 = mycc.make_rdm1()
-        dm2 = mycc.make_rdm2()
+        dm1 = mycc.make_rdm1(ao_repr=True)
+        dm2 = mycc.make_rdm2(ao_repr=True)
 
         e1 = numpy.einsum('ij,ji', hcore, dm1)
         e1+= numpy.einsum('ijkl,ijkl', eri, dm2) * .5
-        self.assertAlmostEqual(e1, mycc.e_tot, 6)
+        self.assertAlmostEqual(e1- mycc.e_tot, 6)
 
         self.assertAlmostEqual(abs(dm2-dm2.transpose(1,0,3,2).conj()).max(), 0, 9)
         self.assertAlmostEqual(abs(dm2-dm2.transpose(2,3,0,1)       ).max(), 0, 9)
diff --git a/pyscf/cc/uccsd_rdm.py b/pyscf/cc/uccsd_rdm.py
index 7e096973f0..8c2bada6ee 100644
--- a/pyscf/cc/uccsd_rdm.py
+++ b/pyscf/cc/uccsd_rdm.py
@@ -625,8 +625,8 @@ def _make_rdm2(mycc, d1, d2, with_dm1=True, with_frozen=True, ao_repr=False):
 
 
 def _dm2ab_mo2ao(dm2, mo_a, mo_b):
-    return lib.einsum('ijkl,pi,qj,rk,sl->pqrs', dm2, mo_a, mo_a.conj(),
-                      mo_b, mo_b.conj())
+    return lib.einsum('ijkl,pi,qj,rk,sl->pqrs', dm2, mo_a.conj(), mo_a,
+                      mo_b.conj(), mo_b)
 
 
 if __name__ == '__main__':