Repository: pyscf/pyscf — pull request #2870 (python)
Title: Enable hermi=0 for Breit and Gaunt JK integrals

fix #2867 

Resolved issues:
title: [UNEXPECTED BEHAVIOUR] No Gaunt or Breit possible in TDDFT
body: **Gaunt and Breit Interactions don't work in TDDFT**
When using the later mentioned code, the regular `pyscf.dft.dks().DKS()` works, when using the options `with_gaunt = True` or `with_breit = True`. When subsequently performing a TDDFT calculation, an unimplemented Error occurs. This is not the case when `with_gaunt` and `with_breit` are turned to `False`.

**Minimal Working Example**
A minimal example, where the error occurs:
```python
from pyscf import gto, dft

mol = gto.Mole()

mol.build(
    atom = '''H 0 0 0
    F 0 0 1.5''',
    basis = '631g',
    verbose = 4,
    symmetry = True,
    )

mf = dft.dks.DKS(mol)
mf.xc = 'pbe0'
mf.with_gaunt = True
mf.with_breit = True
mf.kernel()

mytd = mf.TDDFT()

mytd.kernel()
```

**Error Message**
```bash
Traceback (most recent call last):
  File "/home/kevin/Documents/Spin-Orbit-Coupling/pyscf/HF/test.py", line 22, in <module>
    mytd.kernel()
  File "/home/kevin/.local/lib/python3.12/site-packages/pyscf/tdscf/dhf.py", line 567, in kernel
    self.converged, w, x1 = lr_eig(
                            ^^^^^^^
  File "/home/kevin/.local/lib/python3.12/site-packages/pyscf/tdscf/_lr_eig.py", line 340, in eig
    axt = aop(xt)
          ^^^^^^^
  File "/home/kevin/.local/lib/python3.12/site-packages/pyscf/tdscf/dhf.py", line 504, in vind
    v1ao = vresp(dms) # = <mj||nb> Xjb + <mb||nj> Yjb
           ^^^^^^^^^^
  File "/home/kevin/.local/lib/python3.12/site-packages/pyscf/scf/_response_functions.py", line 297, in vind
    vj, vk = mf.get_jk(mol, dm1, hermi, with_j=with_j)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/kevin/.local/lib/python3.12/site-packages/pyscf/scf/dhf.py", line 679, in get_jk
    vj1, vk1 = _call_veff_gaunt_breit(mol, dm, hermi, opt_gaunt, False)
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/kevin/.local/lib/python3.12/site-packages/pyscf/scf/dhf.py", line 1038, in _call_veff_gaunt_breit
    raise NotImplementedError
NotImplementedError
```


**Version and hardware:**
- PySCF version: 2.9.0
- PySCF install type: built from source and pip, both result in the same error
- Operating system: Ubuntu 24.04


Hints:

Fix patch:
diff --git a/pyscf/scf/dhf.py b/pyscf/scf/dhf.py
index 7a9ced7eed..06395fcef7 100644
--- a/pyscf/scf/dhf.py
+++ b/pyscf/scf/dhf.py
@@ -1010,32 +1010,57 @@ def _call_veff_gaunt_breit(mol, dm, hermi=1, mf_opt=None, with_breit=False):
     vj = numpy.zeros((n_dm,n2c*2,n2c*2), dtype=numpy.complex128)
     vk = numpy.zeros((n_dm,n2c*2,n2c*2), dtype=numpy.complex128)
 
-    jks = ('lk->s1ij', 'jk->s1il')
-    vj_ls, vk_ls = _vhf.rdirect_mapdm(intor_prefix+'ssp1ssp2_spinor', 's1', jks, dms[:n_dm], 1,
-                            mol._atm, mol._bas, mol._env, opt_gaunt_lsls)
-    vj[:,:n2c,n2c:] = vj_ls
-    vk[:,:n2c,n2c:] = vk_ls
-    t0 = log.timer_debug1('LSLS contribution', *t0)
-
-    jks = ('lk->s1ij',) * n_dm \
-        + ('li->s1kj',) * n_dm \
-        + ('jk->s1il',) * n_dm
-    vx = _vhf.rdirect_bindm(intor_prefix+'ssp1sps2_spinor', 's1', jks, dms[n_dm:], 1,
-                            mol._atm, mol._bas, mol._env, opt_gaunt_lssl)
-
-    t0 = log.timer_debug1('LSSL contribution', *t0)
-    vj[:,:n2c,n2c:]+= vx[      :n_dm  ,:,:]
-    vk[:,n2c:,n2c:] = vx[n_dm  :n_dm*2,:,:]
-    vk[:,:n2c,:n2c] = vx[n_dm*2:      ,:,:]
-
-    if hermi == 1:
-        vj[:,n2c:,:n2c] = vj[:,:n2c,n2c:].transpose(0,2,1).conj()
-        vk[:,n2c:,:n2c] = vk[:,:n2c,n2c:].transpose(0,2,1).conj()
-    elif hermi == 2:
-        vj[:,n2c:,:n2c] = -vj[:,:n2c,n2c:].transpose(0,2,1).conj()
-        vk[:,n2c:,:n2c] = -vk[:,:n2c,n2c:].transpose(0,2,1).conj()
+    if hermi != 0:
+        jks = ('lk->s1ij', 'jk->s1il')
+        vj_ls, vk_ls = _vhf.rdirect_mapdm(
+            intor_prefix+'ssp1ssp2_spinor', 's1', jks, dms[:n_dm], 1,
+            mol._atm, mol._bas, mol._env, opt_gaunt_lsls)
+        vj[:,:n2c,n2c:] = vj_ls
+        vk[:,:n2c,n2c:] = vk_ls
+        t0 = log.timer_debug1('LSLS contribution', *t0)
+
+        jks = (('lk->s1ij',) * n_dm +
+               ('li->s1kj',) * n_dm +
+               ('jk->s1il',) * n_dm)
+        vx = _vhf.rdirect_bindm(
+            intor_prefix+'ssp1sps2_spinor', 's1', jks, dms[n_dm:], 1,
+            mol._atm, mol._bas, mol._env, opt_gaunt_lssl)
+        vj[:,:n2c,n2c:]+= vx[      :n_dm  ,:,:]
+        vk[:,n2c:,n2c:] = vx[n_dm  :n_dm*2,:,:]
+        vk[:,:n2c,:n2c] = vx[n_dm*2:      ,:,:]
+        t0 = log.timer_debug1('LSSL contribution', *t0)
+
+        if hermi == 1:
+            vj[:,n2c:,:n2c] = vj[:,:n2c,n2c:].transpose(0,2,1).conj()
+            vk[:,n2c:,:n2c] = vk[:,:n2c,n2c:].transpose(0,2,1).conj()
+        elif hermi == 2:
+            vj[:,n2c:,:n2c] = -vj[:,:n2c,n2c:].transpose(0,2,1).conj()
+            vk[:,n2c:,:n2c] = -vk[:,:n2c,n2c:].transpose(0,2,1).conj()
     else:
-        raise NotImplementedError
+        jks = ('lk->s1ij', 'jk->s1il')
+        subj, subk = _vhf.rdirect_mapdm(
+            intor_prefix+'ssp1ssp2_spinor', 's1', jks,
+            (dms[:n_dm] + [x.conj().T for x in dms[n_dm:n_dm*2]]), 1,
+            mol._atm, mol._bas, mol._env, opt_gaunt_lsls)
+        vj[:,:n2c,n2c:] = subj[      :n_dm  ,:,:]
+        vj[:,n2c:,:n2c] = subj[n_dm  :n_dm*2,:,:].conj().transpose(0,2,1)
+        vk[:,:n2c,n2c:] = subk[      :n_dm  ,:,:]
+        vk[:,n2c:,:n2c] = subk[n_dm  :n_dm*2,:,:].conj().transpose(0,2,1)
+        t0 = log.timer_debug1('LSLS contribution', *t0)
+
+        jks = (('ji->s1kl',) * n_dm +
+               ('lk->s1ij',) * n_dm +
+               ('li->s1kj',) * n_dm +
+               ('jk->s1il',) * n_dm)
+        vx = _vhf.rdirect_bindm(
+            intor_prefix+'ssp1sps2_spinor', 's1', jks, dms, 1,
+            mol._atm, mol._bas, mol._env, opt_gaunt_lssl)
+        vj[:,n2c:,:n2c]+= vx[      :n_dm  ,:,:]
+        vj[:,:n2c,n2c:]+= vx[n_dm  :n_dm*2,:,:]
+        vk[:,n2c:,n2c:] = vx[n_dm*2:n_dm*3,:,:]
+        vk[:,:n2c,:n2c] = vx[n_dm*3:n_dm*4,:,:]
+        t0 = log.timer_debug1('LSSL contribution', *t0)
+
     vj = vj.reshape(dm.shape)
     vk = vk.reshape(dm.shape)
     c1 = .5 / lib.param.LIGHT_SPEED


Test patch:
diff --git a/pyscf/scf/test/test_dhf.py b/pyscf/scf/test/test_dhf.py
index 07ee57229f..2db8b48017 100644
--- a/pyscf/scf/test/test_dhf.py
+++ b/pyscf/scf/test/test_dhf.py
@@ -250,11 +250,17 @@ def test_gaunt(self):
         n4c = erig.shape[0]
         numpy.random.seed(1)
         dm = numpy.random.random((2,n4c,n4c))+numpy.random.random((2,n4c,n4c))*1j
-        dm = dm + dm.transpose(0,2,1).conj()
         c1 = .5 / lib.param.LIGHT_SPEED
         vj0 = -numpy.einsum('ijkl,xlk->xij', erig, dm) * c1**2
         vk0 = -numpy.einsum('ijkl,xjk->xil', erig, dm) * c1**2
+        vj1, vk1 = scf.dhf._call_veff_gaunt_breit(h4, dm, hermi=0)
+        self.assertAlmostEqual(abs(vj0 - vj1).max(), 0, 12)
+        self.assertAlmostEqual(abs(vk0 - vk1).max(), 0, 12)
 
+        dm = dm + dm.transpose(0,2,1).conj()
+        c1 = .5 / lib.param.LIGHT_SPEED
+        vj0 = -numpy.einsum('ijkl,xlk->xij', erig, dm) * c1**2
+        vk0 = -numpy.einsum('ijkl,xjk->xil', erig, dm) * c1**2
         vj1, vk1 = scf.dhf._call_veff_gaunt_breit(h4, dm)
         self.assertTrue(numpy.allclose(vj0, vj1))
         self.assertTrue(numpy.allclose(vk0, vk1))
@@ -279,11 +285,17 @@ def test_breit_high_cost(self):
         n4c = erig.shape[0]
         numpy.random.seed(1)
         dm = numpy.random.random((n4c,n4c))+numpy.random.random((n4c,n4c))*1j
+        c1 = .5 / lib.param.LIGHT_SPEED
+        vj0 = -numpy.einsum('ijkl,xlk->xij', erig, dm) * c1**2
+        vk0 = -numpy.einsum('ijkl,xjk->xil', erig, dm) * c1**2
+        vj1, vk1 = scf.dhf._call_veff_gaunt_breit(h4, dm, hermi=0)
+        self.assertAlmostEqual(abs(vj0 - vj1).max(), 0, 12)
+        self.assertAlmostEqual(abs(vk0 - vk1).max(), 0, 12)
+
         dm = dm + dm.T.conj()
         c1 = .5 / lib.param.LIGHT_SPEED
         vj0 = numpy.einsum('ijkl,lk->ij', erig, dm) * c1**2
         vk0 = numpy.einsum('ijkl,jk->il', erig, dm) * c1**2
-
         vj1, vk1 = scf.dhf._call_veff_gaunt_breit(h4, dm, with_breit=True)
         self.assertTrue(numpy.allclose(vj0, vj1))
         self.assertTrue(numpy.allclose(vk0, vk1))
@@ -325,23 +337,22 @@ def _fill_gaunt(mol, erig):
     eri0 = numpy.zeros((n4c,n4c,n4c,n4c), dtype=numpy.complex128)
     eri0[:n2c,n2c:,:n2c,n2c:] = erig # ssp1ssp2
 
-    eri2 = erig.take(idx,axis=0).take(idx,axis=1) # sps1ssp2
+    eri2 = erig[idx[:,None],idx] # sps1ssp2
     eri2[sign_mask,:] *= -1
     eri2[:,sign_mask] *= -1
     eri2 = -eri2.transpose(1,0,2,3)
     eri0[n2c:,:n2c,:n2c,n2c:] = eri2
 
-    eri2 = erig.take(idx,axis=2).take(idx,axis=3) # ssp1sps2
+    eri2 = erig[:,:,idx[:,None],idx] # ssp1sps2
     eri2[:,:,sign_mask,:] *= -1
     eri2[:,:,:,sign_mask] *= -1
     eri2 = -eri2.transpose(0,1,3,2)
     #self.assertTrue(numpy.allclose(eri0, eri2))
     eri0[:n2c,n2c:,n2c:,:n2c] = eri2
 
-    eri2 = erig.take(idx,axis=0).take(idx,axis=1)
-    eri2 = eri2.take(idx,axis=2).take(idx,axis=3) # sps1sps2
-    eri2 = eri2.transpose(1,0,2,3)
-    eri2 = eri2.transpose(0,1,3,2)
+    eri2 = erig[idx[:,None],idx]
+    eri2 = eri2[:,:,idx[:,None],idx] # sps1sps2
+    eri2 = eri2.transpose(1,0,3,2)
     eri2[sign_mask,:] *= -1
     eri2[:,sign_mask] *= -1
     eri2[:,:,sign_mask,:] *= -1