From e2f122fbd0f1ba6effce465fc86accd7fb145504 Mon Sep 17 00:00:00 2001
From: Qiming Sun <osirpt.sun@gmail.com>
Date: Thu, 25 Jan 2024 23:43:47 -0800
Subject: [PATCH] Fix the class of supercell scf object for the k2gamma
 conversion (#2030)

---
 pyscf/pbc/tools/k2gamma.py           | 27 ++++++++++++++++++++++++---
 pyscf/pbc/tools/test/test_k2gamma.py |  6 +++++-
 2 files changed, 29 insertions(+), 4 deletions(-)

diff --git a/pyscf/pbc/tools/k2gamma.py b/pyscf/pbc/tools/k2gamma.py
index 32bcd3c32a..6b0cb42cf8 100644
--- a/pyscf/pbc/tools/k2gamma.py
+++ b/pyscf/pbc/tools/k2gamma.py
@@ -205,7 +205,7 @@ def k2gamma(kmf, kmesh=None):
          C_{\nu ' n'} = C_{\vecR\mu, \veck m} = \frac{1}{\sqrt{N_{\UC}}}
          \e^{\ii \veck\cdot\vecR} C^{\veck}_{\mu  m}
     '''
-    from pyscf.pbc import scf
+    from pyscf.pbc import scf, dft
     def transform(mo_energy, mo_coeff, mo_occ):
         if isinstance(kmf.kpts, KPoints):
             kpts = kmf.kpts.kpts
@@ -228,17 +228,38 @@ def transform(mo_energy, mo_coeff, mo_occ):
 
     if isinstance(kmf, scf.khf.KRHF):
         scell, E_g, C_gamma, mo_occ = transform(mo_energy, mo_coeff, mo_occ)
-        mf = scf.RHF(scell)
     elif isinstance(kmf, scf.kuhf.KUHF):
         scell, Ea, Ca, occ_a = transform(mo_energy[0], mo_coeff[0], mo_occ[0])
         scell, Eb, Cb, occ_b = transform(mo_energy[1], mo_coeff[1], mo_occ[1])
-        mf = scf.UHF(scell)
         E_g = [Ea, Eb]
         C_gamma = [Ca, Cb]
         mo_occ = [occ_a, occ_b]
     else:
         raise NotImplementedError('SCF object %s not supported' % kmf)
 
+    known_cls = {
+        dft.kuks.KUKS  : dft.uks.UKS  ,
+        dft.kroks.KROKS: dft.roks.ROKS,
+        dft.krks.KRKS  : dft.rks.RKS  ,
+        dft.kgks.KGKS  : dft.gks.GKS  ,
+        scf.kuhf.KUHF  : scf.uhf.UHF  ,
+        scf.krohf.KROHF: scf.rohf.ROHF,
+        scf.khf.KRHF   : scf.hf.RHF   ,
+        scf.kghf.KGHF  : scf.ghf.GHF  ,
+    }
+    if kmf.__class__ in known_cls:
+        mf = known_cls[kmf.__class__](scell)
+        mf.exxdiv = kmf.exxdiv
+        if isinstance(mf, KohnShamDFT):
+            mf.xc = kmf.xc
+    else:
+        kmf.warn(f'Unknown SCF object {kmf}. '
+                 'The supercell SCF object is initialized to an HF object.')
+        if isinstance(kmf, scf.khf.KRHF):
+            mf = scf.RHF(scell)
+        else:
+            mf = scf.UHF(scell)
+
     mf.mo_coeff = C_gamma
     mf.mo_energy = E_g
     mf.mo_occ = mo_occ
diff --git a/pyscf/pbc/tools/test/test_k2gamma.py b/pyscf/pbc/tools/test/test_k2gamma.py
index ae70bc3982..fa4f04c98c 100644
--- a/pyscf/pbc/tools/test/test_k2gamma.py
+++ b/pyscf/pbc/tools/test/test_k2gamma.py
@@ -88,10 +88,14 @@ def test_k2gamma_ksymm(self):
         kpts = cell.make_kpts(kmesh, space_group_symmetry=True)
         kmf = scf.KRKS(cell, kpts).density_fit()
         kmf.kernel()
-        c_g_ao = k2gamma.k2gamma(kmf).mo_coeff
+        mf = k2gamma.k2gamma(kmf)
+        c_g_ao = mf.mo_coeff
 
         scell = tools.super_cell(cell, kmesh)
         mf_sc = scf.RKS(scell).density_fit()
+        self.assertEqual(mf.__class__, mf_sc.__class__)
+        self.assertEqual(mf.xc, kmf.xc)
+
         s = mf_sc.get_ovlp()
         mf_sc.run()
         sc_mo = mf_sc.mo_coeff