Repository: pyscf/pyscf — pull request #2457 (python)
Title: Fixes calculating exchange integrals when only_dfj=True, issue #2456

This PR fixes an assertion error in the following minimal example:
```
from pyscf import gto, dft

mol = gto.M(
    atom="""
        O  0.00000000  0.00000000  0.00000000
        H -0.75390364  0.00000000 -0.58783729
        H  0.75390364  0.00000000 -0.58783729
    """
)

mf = dft.RKS(mol).density_fit(only_dfj=True).run()
vk = mf.get_k(mol=mol)
```


Resolved issues:
title: Calculating exchange integrals fails when only_dfj=True
body: Hi everyone,
the following minimal example fails due to an assertion error:
```
from pyscf import gto, dft

mol = gto.M(
    atom="""
        O  0.00000000  0.00000000  0.00000000
        H -0.75390364  0.00000000 -0.58783729
        H  0.75390364  0.00000000 -0.58783729
    """
)

mf = dft.RKS(mol).density_fit(only_dfj=True).run()
vk = mf.get_k(mol=mol)
```
It should still be possible to get the regular exchange integrals which are also returned for `get_jk` if `with_j=True` and `with_k=True`. The issue appears to be in the following function:
https://github.com/pyscf/pyscf/blob/410d9608a73ad3a507b3008a6f68549288ebfc33/pyscf/df/df_jk.py#L133-L159

I would propose the following solution:
 ```
   def get_jk(self, mol=None, dm=None, hermi=1, with_j=True, with_k=True,
               omega=None):
        assert (with_j or with_k)
        if dm is None: dm = self.make_rdm1()
        if not self.with_df:
            return super().get_jk(mol, dm, hermi, with_j, with_k, omega)

        with_dfk = with_k and not self.only_dfj
        if with_j:
            if isinstance(self, scf.ghf.GHF):
                def jkbuild(mol, dm, hermi, with_j, with_k, omega=None):
                    vj, vk = self.with_df.get_jk(dm.real, hermi, with_j, with_k,
                                                self.direct_scf_tol, omega)
                    if dm.dtype == numpy.complex128:
                        vjI, vkI = self.with_df.get_jk(dm.imag, hermi, with_j, with_k,
                                                    self.direct_scf_tol, omega)
                        if with_j:
                            vj = vj + vjI * 1j
                        if with_k:
                            vk = vk + vkI * 1j
                    return vj, vk
                vj, vk = scf.ghf.get_jk(mol, dm, hermi, with_j, with_dfk,
                                        jkbuild, omega)
            else:
                vj, vk = self.with_df.get_jk(dm, hermi, with_j, with_dfk,
                                            self.direct_scf_tol, omega)
        else:
            vj = None
        if with_k and not with_dfk:
            vk = super().get_jk(mol, dm, hermi, False, True, omega)[1]
        return vj, vk
```
This change would still require that either the Coulomb or the exchange matrix are requested while ensuring the regular exchange matrix is returned whenever the corresponding density-fitted integrals are not available. I can open a PR if this behavior/fix is desired.

Hints:

Fix patch:
diff --git a/pyscf/df/df_jk.py b/pyscf/df/df_jk.py
index 2ac82c243c..972b7e2033 100644
--- a/pyscf/df/df_jk.py
+++ b/pyscf/df/df_jk.py
@@ -132,28 +132,31 @@ def reset(self, mol=None):
 
     def get_jk(self, mol=None, dm=None, hermi=1, with_j=True, with_k=True,
                omega=None):
+        assert (with_j or with_k)
         if dm is None: dm = self.make_rdm1()
         if not self.with_df:
             return super().get_jk(mol, dm, hermi, with_j, with_k, omega)
 
+        vj = vk = None
         with_dfk = with_k and not self.only_dfj
-        if isinstance(self, scf.ghf.GHF):
-            def jkbuild(mol, dm, hermi, with_j, with_k, omega=None):
-                vj, vk = self.with_df.get_jk(dm.real, hermi, with_j, with_k,
-                                             self.direct_scf_tol, omega)
-                if dm.dtype == numpy.complex128:
-                    vjI, vkI = self.with_df.get_jk(dm.imag, hermi, with_j, with_k,
-                                                   self.direct_scf_tol, omega)
-                    if with_j:
-                        vj = vj + vjI * 1j
-                    if with_k:
-                        vk = vk + vkI * 1j
-                return vj, vk
-            vj, vk = scf.ghf.get_jk(mol, dm, hermi, with_j, with_dfk,
-                                    jkbuild, omega)
-        else:
-            vj, vk = self.with_df.get_jk(dm, hermi, with_j, with_dfk,
-                                         self.direct_scf_tol, omega)
+        if with_j or with_dfk:
+            if isinstance(self, scf.ghf.GHF):
+                def jkbuild(mol, dm, hermi, with_j, with_k, omega=None):
+                    vj, vk = self.with_df.get_jk(dm.real, hermi, with_j, with_k,
+                                                self.direct_scf_tol, omega)
+                    if dm.dtype == numpy.complex128:
+                        vjI, vkI = self.with_df.get_jk(dm.imag, hermi, with_j, with_k,
+                                                    self.direct_scf_tol, omega)
+                        if with_j:
+                            vj = vj + vjI * 1j
+                        if with_k:
+                            vk = vk + vkI * 1j
+                    return vj, vk
+                vj, vk = scf.ghf.get_jk(mol, dm, hermi, with_j, with_dfk,
+                                        jkbuild, omega)
+            else:
+                vj, vk = self.with_df.get_jk(dm, hermi, with_j, with_dfk,
+                                            self.direct_scf_tol, omega)
         if with_k and not with_dfk:
             vk = super().get_jk(mol, dm, hermi, False, True, omega)[1]
         return vj, vk


Test patch:
diff --git a/pyscf/df/test/test_df.py b/pyscf/df/test/test_df.py
index 394d6d5a55..11b07b91be 100644
--- a/pyscf/df/test/test_df.py
+++ b/pyscf/df/test/test_df.py
@@ -142,6 +142,44 @@ def test_rsh_df_custom_storage(self):
         mf.run()
         self.assertAlmostEqual(mf.e_tot, -103.4965622991, 6)
 
+    def test_only_dfj(self):
+        mol = gto.M(atom='H 0 0 0; H 0 0 1')
+        dm = numpy.eye(mol.nao)
+        mf = mol.RHF().density_fit()
+        refj, refk = mf.get_jk(mol, dm)
+        vk = mf.get_k(mol, dm)
+        self.assertAlmostEqual(abs(vk - refk).max(), 0, 9)
+        vj = mf.get_j(mol, dm)
+        self.assertAlmostEqual(abs(vj - refj).max(), 0, 9)
+
+        mf = mol.RHF().density_fit(only_dfj=True)
+        refk = mol.RHF().get_k(mol, dm)
+        vj, vk = mf.get_jk(mol, dm)
+        self.assertAlmostEqual(abs(vk - refk).max(), 0, 9)
+        self.assertAlmostEqual(abs(vj - refj).max(), 0, 9)
+        vk = mf.get_k(mol, dm)
+        self.assertAlmostEqual(abs(vk - refk).max(), 0, 9)
+        vj = mf.get_j(mol, dm)
+        self.assertAlmostEqual(abs(vj - refj).max(), 0, 9)
+
+        dm = numpy.eye(mol.nao*2)
+        mf = mol.GHF().density_fit()
+        refj, refk = mf.get_jk(mol, dm)
+        vk = mf.get_k(mol, dm)
+        self.assertAlmostEqual(abs(vk - refk).max(), 0, 9)
+        vj = mf.get_j(mol, dm)
+        self.assertAlmostEqual(abs(vj - refj).max(), 0, 9)
+
+        mf = mol.GHF().density_fit(only_dfj=True)
+        refk = mol.GHF().get_k(mol, dm)
+        vj, vk = mf.get_jk(mol, dm)
+        self.assertAlmostEqual(abs(vk - refk).max(), 0, 9)
+        self.assertAlmostEqual(abs(vj - refj).max(), 0, 9)
+        vk = mf.get_k(mol, dm)
+        self.assertAlmostEqual(abs(vk - refk).max(), 0, 9)
+        vj = mf.get_j(mol, dm)
+        self.assertAlmostEqual(abs(vj - refj).max(), 0, 9)
+
 if __name__ == "__main__":
     print("Full Tests for df")
     unittest.main()