Task: Fix define_xc example (pyscf/pyscf#2364)

org: pyscf
repo: pyscf
pr_number: 2364
issue_title: Fix define_xc example
issue_body: Update examples: 24-define_xc_functional.py and 31-xc_value_on_grid.py . (Fix issue #2092 , #2011)
resolved_issues: number: 2011
title: Obtaining Hartree-Fock exchange energy density on grids
body: Hello pySCF team, I hope the new year is treating you well. 

I'm trying to obtain Hartree-Fock exchange energy densities on grids (as the title would suggest). While there was an issue ( #1485 ) posted and seemingly resolved a little over a year ago, there was no input from any pySCF devs and the issue was closed the same day that it was posted with no commentary from anyone other than the author. I've tried executing this method to no success, as it doesn't seem to return values that get us the energy when combined with rho and the grid weights (as you would be able to with exc in a Kohn-Sham calculation).

In short, I'm looking for an analogue to exc from KS, but in HF (an ex, if you will).

If there is any insight on how this may be accomplished, please let me know and I look forward to your responses!
number: 2092
title: Unable to use custom xc functional
body: Hi PySCF developers,

I'm new to using PySCF. I wanted to plugin my own XC functional using the define_xc_() interface in the dft module. To begin with, I attempted to run the example that demonstrates the use of custom XC functional (https://github.com/pyscf/pyscf/blob/master/examples/dft/24-define_xc_functional.py). However, running it throws the following error:

```
Traceback (most recent call last):
  File "custom_xc_test.py", line 63, in <module>
    mf.kernel()
  File "<string>", line 2, in kernel
  File "/home/biku/miniconda3/lib/python3.8/site-packages/pyscf/scf/hf.py", line 1779, in scf
    kernel(self, self.conv_tol, self.conv_tol_grad,
  File "/home/biku/miniconda3/lib/python3.8/site-packages/pyscf/scf/hf.py", line 124, in kernel
    vhf = mf.get_veff(mol, dm)
  File "/home/biku/miniconda3/lib/python3.8/site-packages/pyscf/dft/rks.py", line 80, in get_veff
    n, exc, vxc = ni.nr_rks(mol, ks.grids, ks.xc, dm, max_memory=max_memory)
  File "/home/biku/miniconda3/lib/python3.8/site-packages/pyscf/dft/numint.py", line 1158, in nr_rks
    for i, ao, mask, wv in block_loop(ao_deriv):
  File "/home/biku/miniconda3/lib/python3.8/site-packages/pyscf/dft/numint.py", line 1138, in block_loop
    exc, vxc = ni.eval_xc_eff(xc_code, rho, deriv=1, xctype=xctype)[:2]
  File "/home/biku/miniconda3/lib/python3.8/site-packages/pyscf/dft/numint.py", line 2696, in eval_xc_eff
    out = self.libxc.eval_xc1(xc_code, rho, spin, deriv, omega)
  File "/home/biku/miniconda3/lib/python3.8/site-packages/pyscf/dft/libxc.py", line 1599, in eval_xc1
    out = _eval_xc(xc_code, rho, spin, deriv=deriv, omega=omega)
  File "/home/biku/miniconda3/lib/python3.8/site-packages/pyscf/dft/libxc.py", line 1654, in _eval_xc
    rho = rho.reshape(spin+1,nvar,ngrids)
ValueError: cannot reshape array of size 120224 into shape (1,1,30056)
```

The custom_xc_test.py is just a copy of the example link provided above. 

I used `pip install --prefer-binary pyscf` to install my PySCF.

I would really appreciate it if you can help me figure out how to use a custom XC functional. 

Regards,
Bikash

Reference solution:
fix_patch: diff --git a/examples/dft/24-define_xc_functional.py b/examples/dft/24-define_xc_functional.py
index e4efae3089..72c7f25cfd 100644
--- a/examples/dft/24-define_xc_functional.py
+++ b/examples/dft/24-define_xc_functional.py
@@ -50,11 +50,11 @@ def eval_xc(xc_code, rho, spin=0, relativity=0, deriv=1, omega=None, verbose=Non
     kxc = None  # 3rd order functional derivative
 
     # Mix with existing functionals
-    pbe_xc = dft.libxc.eval_xc('pbe,pbe', rho, spin, relativity, deriv,
-                               verbose)
+    pbe_xc = dft.libxc.eval_xc('pbe,pbe', rho, spin, relativity, deriv, verbose)
     exc += pbe_xc[0] * 0.5
     vrho += pbe_xc[1][0] * 0.5
     vgamma += pbe_xc[1][1] * 0.5
+    # The output follows the libxc.eval_xc API convention
     return exc, vxc, fxc, kxc
 
 mf = dft.RKS(mol)
@@ -70,4 +70,3 @@ def eval_xc(xc_code, rho, spin=0, relativity=0, deriv=1, omega=None, verbose=Non
 mf = mf.define_xc_(eval_xc, 'GGA', rsh=rsh_coeff)
 mf.verbose = 4
 mf.kernel()
-
diff --git a/examples/dft/31-xc_value_on_grid.py b/examples/dft/31-xc_value_on_grid.py
index c78592334e..d0e7a980c9 100644
--- a/examples/dft/31-xc_value_on_grid.py
+++ b/examples/dft/31-xc_value_on_grid.py
@@ -5,7 +5,7 @@
 potential on given grid coordinates.
 '''
 
-import numpy
+import numpy as np
 from pyscf import gto, dft, lib
 from pyscf.dft import numint
 from pyscf.dft import r_numint
@@ -37,19 +37,19 @@
 #
 ex, vx = dft.libxc.eval_xc('B88', rho)[:2]
 ec, vc = dft.libxc.eval_xc(',P86', rho)[:2]
-print('Exc = %.12f' % numpy.einsum('i,i,i->', ex+ec, rho[0], weights))
+print('Exc = %.12f' % np.einsum('i,i,i->', ex+ec, rho[0], weights))
 
 #
 # Evaluate XC functional together
 #
 exc, vxc = dft.libxc.eval_xc('B88,P86', rho)[:2]
-print('Exc = %.12f' % numpy.einsum('i,i,i->', exc, rho[0], weights))
+print('Exc = %.12f' % np.einsum('i,i,i->', exc, rho[0], weights))
 
 #
 # Evaluate XC functional for user specified functional
 #
 exc, vxc = dft.libxc.eval_xc('.2*HF + .08*SLATER + .72*B88, .81*LYP + .19*VWN', rho)[:2]
-print('Exc = %.12f  ref = -7.520014202688' % numpy.einsum('i,i,i->', exc, rho[0], weights))
+print('Exc = %.12f  ref = -7.520014202688' % np.einsum('i,i,i->', exc, rho[0], weights))
 
 #
 # For density matrix computed with GKS in spin-orbital basis, plain density in
@@ -82,3 +82,20 @@
 rhoS, mS = r_numint.eval_rho(mol, aoS_value, dmSS)
 rho = rhoL + rhoS
 mx, my, mz = mL + mS
+
+#
+# The values of exact exchange
+# Exx(r) = \sum_{ij} 2 (ij(r)|ji)
+#
+mf = mol.RHF().run()
+# Compute analytically v_ij(r') = \int \chi_i(r) \chi_j(r) / |r-r'| dr
+rpq = mol.intor('int1e_grids', grids=coords)
+ao = mol.eval_gto('GTOval', coords)
+nocc = np.count_nonzero(mf.mo_occ)
+occ_orb = mf.mo_coeff[:,:nocc]
+mo = ao.dot(occ_orb)
+exx = 2 * np.einsum('rpq,ri,rj,pj,qi->r', rpq, mo, mo, occ_orb, occ_orb)
+print('Semi-numerical HFX:', exx.dot(weights))
+dm = mf.make_rdm1()
+vj, vk = mf.get_jk(dm)
+print('Analytical HFX:', .5 * np.einsum('pq,qp->', vk, dm))
diff --git a/examples/fci/03-cylindrical_symmetry.py b/examples/fci/03-cylindrical_symmetry.py
index 78ad41eca7..3b3c382f88 100644
--- a/examples/fci/03-cylindrical_symmetry.py
+++ b/examples/fci/03-cylindrical_symmetry.py
@@ -26,5 +26,5 @@
 solver = direct_spin1_cyl_sym.FCI(mol)
 for sym in ['A1g', 'A1u', 'E1ux', 'E1uy', 'E1gx', 'E1gy', 'E2ux', 'E2uy', 'E2gx', 'E2gy', 'E3ux', 'E3uy', 'E3gx', 'E3gy']:
     e, v = solver.kernel(h1e, eri, c.shape[1], mol.nelec, orbsym=orbsym,
-                         wfnsym=sym, nroots=5, verbose=0)
+                         wfnsym=sym, nroots=5, ecore=mol.energy_nuc(), verbose=0)
     print(f'{sym}: {e}')
diff --git a/pyscf/dft/libxc.py b/pyscf/dft/libxc.py
index 0f248fb505..26aef161bf 100644
--- a/pyscf/dft/libxc.py
+++ b/pyscf/dft/libxc.py
@@ -1609,13 +1609,17 @@ def eval_xc1(xc_code, rho, spin=0, deriv=1, omega=None):
     '''
     out = _eval_xc(xc_code, rho, spin, deriv=deriv, omega=omega)
     xctype = xc_type(xc_code)
+    idx = _libxc_to_xcfun_indices(xctype, spin, deriv)
+    return out[idx]
+
+def _libxc_to_xcfun_indices(xctype, spin=0, deriv=1):
     if deriv <= 1:
-        return out
+        return slice(None)
     elif xctype == 'LDA' or xctype == 'HF':
-        return out
+        return slice(None)
     elif xctype == 'GGA':
         if spin == 0:
-            return out
+            return slice(None)
         else:
             idx = [numpy.arange(6)] # up to deriv=1
             for i in range(2, deriv+1):
@@ -1627,7 +1631,7 @@ def eval_xc1(xc_code, rho, spin=0, deriv=1, omega=None):
             idx = [numpy.arange(8)] # up to deriv=1
         for i in range(2, deriv+1):
             idx.append(_MGGA_SORT[(spin, i)])
-    return out[numpy.hstack(idx)]
+    return numpy.hstack(idx)
 
 def _eval_xc(xc_code, rho, spin=0, deriv=1, omega=None):
     assert deriv <= max_deriv_order(xc_code)
@@ -1762,19 +1766,43 @@ def define_xc_(ni, description, xctype='LDA', hyb=0, rsh=(0,0,0)):
     48.8525211046668
     '''
     if isinstance(description, str):
-        ni.eval_xc = lambda xc_code, rho, *args, **kwargs: \
-                eval_xc(description, rho, *args, **kwargs)
+        def _eval_xc(xc_code, rho, *args, **kwargs):
+            return eval_xc(description, rho, *args, **kwargs)
+        ni.eval_xc = _eval_xc
         ni.hybrid_coeff = lambda *args, **kwargs: hybrid_coeff(description)
         ni.rsh_coeff = lambda *args: rsh_coeff(description)
         ni._xc_type = lambda *args: xc_type(description)
 
     elif callable(description):
-        ni.eval_xc = description
+        ni.eval_xc = _eval_xc = description
         ni.hybrid_coeff = lambda *args, **kwargs: hyb
         ni.rsh_coeff = lambda *args, **kwargs: rsh
         ni._xc_type = lambda *args: xctype
+
     else:
         raise ValueError('Unknown description %s' % description)
+
+    def _eval_xc1(xc_code, rho, spin=0, deriv=1, omega=None):
+        libxc_out = _eval_xc(xc_code, rho, spin, deriv=deriv, omega=omega)
+        nvar, xlen = xc_deriv._XC_NVAR[xctype, spin]
+        outlen = lib.comb(xlen+deriv, deriv)
+        exc, vxc, fxc, kxc = libxc_out[:4]
+        out = [exc]
+        if vxc is not None:
+            out.extend([x for x in vxc if x is not None])
+        if fxc is not None:
+            out.extend([fxc[i] for i in [0, 1, 2, 6, 4, 9]])
+        if kxc is not None:
+            out.extend([x for x in kxc if x is not None])
+        if spin == 1:
+            # Returns of eval_xc are structured as [grid_id,deriv_component]
+            # for each term in libxc_out. Change the shape to [deriv_comp, grid_id]
+            out = [x.T for x in out]
+        out = numpy.vstack(out)[:outlen]
+        assert len(out) == outlen
+        idx = _libxc_to_xcfun_indices(xctype, spin, deriv)
+        return out[idx]
+    ni.eval_xc1 = _eval_xc1
     return ni
 
 def define_xc(ni, description, xctype='LDA', hyb=0, rsh=(0,0,0)):
diff --git a/pyscf/dft/numint.py b/pyscf/dft/numint.py
index 7eef303d2c..0f9dae11e5 100644
--- a/pyscf/dft/numint.py
+++ b/pyscf/dft/numint.py
@@ -2654,6 +2654,10 @@ def eval_xc(self, xc_code, rho, spin=0, relativity=0, deriv=1, omega=None,
         return self.libxc.eval_xc(xc_code, rho, spin, relativity, deriv,
                                   omega, verbose)
 
+    def eval_xc1(self, xc_code, rho, spin=0, deriv=1, omega=None):
+        if omega is None: omega = self.omega
+        return self.libxc.eval_xc1(xc_code, rho, spin, deriv, omega)
+
     def eval_xc_eff(self, xc_code, rho, deriv=1, omega=None, xctype=None,
                     verbose=None):
         r'''Returns the derivative tensor against the density parameters
@@ -2693,12 +2697,12 @@ def eval_xc_eff(self, xc_code, rho, deriv=1, omega=None, xctype=None,
         else:
             spin = 0
 
-        out = self.libxc.eval_xc1(xc_code, rho, spin, deriv, omega)
+        out = self.eval_xc1(xc_code, rho, spin, deriv, omega)
         evfk = [out[0]]
         for order in range(1, deriv+1):
             evfk.append(xc_deriv.transform_xc(rho, out, xctype, spin, order))
         if deriv < 3:
-            # The return has at least [e, v, f, k] terms
+            # Returns at least [e, v, f, k] terms
             evfk.extend([None] * (3 - deriv))
         return evfk