Task: Bugs in libxc.define_xc_ (fix #2551) (pyscf/pyscf#2574)

org: pyscf
repo: pyscf
pr_number: 2574
issue_title: Bugs in libxc.define_xc_ (fix #2551)
issue_body: Fix `dft.libxc.define_xc_` and update `examples/dft/24-define_xc_functional.py` (fix #2551)
resolved_issues: number: 2551
title: pyscf/examples/dft/24-define_xc_functional.py is not working
body: I'm trying to run examples of pyscf, and I figured out that dft/24-define_xc_functional.py is not working

Below is my error message.

```python
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[87], line 63
     61 mf = mf.define_xc_(eval_xc, 'GGA', hyb=hybrid_coeff)
     62 mf.verbose = 4
---> 63 mf.kernel()

File <string>:2, in kernel(self, dm0, **kwargs)

File ~/miniconda3/lib/python3.11/site-packages/pyscf/scf/hf.py:1793, in SCF.scf(self, dm0, **kwargs)
   1788 self.build(self.mol)
   1790 if self.max_cycle > 0 or self.mo_coeff is None:
   1791     self.converged, self.e_tot, \
   1792             self.mo_energy, self.mo_coeff, self.mo_occ = \
-> 1793             kernel(self, self.conv_tol, self.conv_tol_grad,
   1794                    dm0=dm0, callback=self.callback,
   1795                    conv_check=self.conv_check, **kwargs)
   1796 else:
   1797     # Avoid to update SCF orbitals in the non-SCF initialization
   1798     # (issue #495).  But run regular SCF for initial guess if SCF was
   1799     # not initialized.
   1800     self.e_tot = kernel(self, self.conv_tol, self.conv_tol_grad,
   1801                         dm0=dm0, callback=self.callback,
   1802                         conv_check=self.conv_check, **kwargs)[1]

File ~/miniconda3/lib/python3.11/site-packages/pyscf/scf/hf.py:126, in kernel(mf, conv_tol, conv_tol_grad, dump_chk, dm0, callback, conv_check, **kwargs)
    123     dm = dm0
    125 h1e = mf.get_hcore(mol)
--> 126 vhf = mf.get_veff(mol, dm)
    127 e_tot = mf.energy_tot(dm, h1e, vhf)
    128 logger.info(mf, 'init E= %.15g', e_tot)

File ~/miniconda3/lib/python3.11/site-packages/pyscf/dft/rks.py:80, in get_veff(ks, mol, dm, dm_last, vhf_last, hermi)
     78 else:
     79     max_memory = ks.max_memory - lib.current_memory()[0]
---> 80     n, exc, vxc = ni.nr_rks(mol, ks.grids, ks.xc, dm, max_memory=max_memory)
     81     logger.debug(ks, 'nelec by numeric integration = %s', n)
     82     if ks.do_nlc():

File ~/miniconda3/lib/python3.11/site-packages/pyscf/dft/numint.py:1158, in nr_rks(ni, mol, grids, xc_code, dms, relativity, hermi, max_memory, verbose)
   1156 elif xctype == 'GGA':
   1157     ao_deriv = 1
-> 1158     for i, ao, mask, wv in block_loop(ao_deriv):
   1159         wv[0] *= .5  # *.5 because vmat + vmat.T at the end
   1160         aow = _scale_ao_sparse(ao[:4], wv[:4], mask, ao_loc, out=aow)

File ~/miniconda3/lib/python3.11/site-packages/pyscf/dft/numint.py:1138, in nr_rks.<locals>.block_loop(ao_deriv)
   1136 for i in range(nset):
   1137     rho = make_rho(i, ao, mask, xctype)
-> 1138     exc, vxc = ni.eval_xc_eff(xc_code, rho, deriv=1, xctype=xctype)[:2]
   1139     if xctype == 'LDA':
   1140         den = rho * weight

File ~/miniconda3/lib/python3.11/site-packages/pyscf/dft/numint.py:2696, in LibXCMixin.eval_xc_eff(self, xc_code, rho, deriv, omega, xctype, verbose)
   2693 else:
   2694     spin = 0
-> 2696 out = self.libxc.eval_xc1(xc_code, rho, spin, deriv, omega)
   2697 evfk = [out[0]]
   2698 for order in range(1, deriv+1):

File ~/miniconda3/lib/python3.11/site-packages/pyscf/dft/libxc.py:1610, in eval_xc1(xc_code, rho, spin, deriv, omega)
   1606 def eval_xc1(xc_code, rho, spin=0, deriv=1, omega=None):
   1607     '''Similar to eval_xc.
   1608     Returns an array with the order of derivatives following xcfun convention.
   1609     '''
-> 1610     out = _eval_xc(xc_code, rho, spin, deriv=deriv, omega=omega)
   1611     xctype = xc_type(xc_code)
   1612     if deriv <= 1:

File ~/miniconda3/lib/python3.11/site-packages/pyscf/dft/libxc.py:1665, in _eval_xc(xc_code, rho, spin, deriv, omega)
   1663 nvar, xlen = xc_deriv._XC_NVAR[xctype, spin]
   1664 ngrids = rho.shape[-1]
-> 1665 rho = rho.reshape(spin+1,nvar,ngrids)
   1666 outlen = lib.comb(xlen+deriv, deriv)
   1667 out = numpy.zeros((outlen,ngrids))

ValueError: cannot reshape array of size 120224 into shape (1,1,30056)
``` 

And I also trying to run libxc.py's example of "define_xc_" [function](https://github.com/pyscf/pyscf/blob/c58968907adc96bc748591c9d72f1bc2c8ee66dc/pyscf/dft/libxc.py#L1123)

```python
def eval_xc(xc_code, rho, *args, **kwargs):
    exc = 0.01 * rho**2
    vrho = 0.01 * 3 * rho**2
    vxc = (vrho, None, None, None)
    fxc = None  # 2nd order functional derivative
    kxc = None  # 3rd order functional derivative
    return exc, vxc, fxc, kxc
define_xc_(mf._numint, eval_xc, xctype='LDA')
mf.kernel()

>>converged SCF energy = -75.8511902095329
``` 

but the code's results are same with 'SVWN' functional (default of pyscf).

```python
mf = dft.RKS(mol)
mf.xc = 'SVWN'
mf.kernel()

>>converged SCF energy = -75.8511902095329
``` 
I think there are some problems in running scf with self-made xc functional.


### Minimal API referenced by tests (typed signatures)
_Added only when tests call brand‑new symbols introduced by the fix; types reflect annotations present in the repo or `Unknown` if not annotated._

- Function in `examples/dft/24-define_xc_functional.py`: `def eval_mgga_xc(xc_code, rho, spin, relativity, deriv, omega, verbose) -> numpy.ndarray: ...`
  **Inputs:** xc_code: str, rho: numpy.ndarray, spin: int, relativity: int, deriv: int, omega: float, verbose: bool
  **Output:** numpy.ndarray

```python
def eval_mgga_xc(xc_code, rho, spin, relativity, deriv, omega, verbose) -> numpy.ndarray: ...
```

Reference solution:
fix_patch: diff --git a/examples/dft/24-define_xc_functional.py b/examples/dft/24-define_xc_functional.py
index 72c7f25cfd..d5e2986f53 100644
--- a/examples/dft/24-define_xc_functional.py
+++ b/examples/dft/24-define_xc_functional.py
@@ -12,6 +12,7 @@
   functionals provided by Libxc or XcFun library.
 '''
 
+import numpy as np
 from pyscf import gto
 from pyscf import dft
 
@@ -25,28 +26,19 @@
 # half-half exact exchange and GGA functional
 hybrid_coeff = 0.5
 
-def eval_xc(xc_code, rho, spin=0, relativity=0, deriv=1, omega=None, verbose=None):
+def eval_gga_xc(xc_code, rho, spin=0, relativity=0, deriv=1, omega=None, verbose=None):
     # A fictitious XC functional to demonstrate the usage
-    rho0, dx, dy, dz = rho[:4]
+    rho0, dx, dy, dz = rho
     gamma = (dx**2 + dy**2 + dz**2)
     exc = .01 * rho0**2 + .02 * (gamma+.001)**.5
     vrho = .01 * 3 * rho0**2 + .02 * (gamma+.001)**.5
     vgamma = .02 * .5 * (gamma+.001)**(-.5)
-    vlapl = None
-    vtau = None
-    vxc = (vrho, vgamma, vlapl, vtau)
+    vxc = (vrho, vgamma)
     v2rho2 = 0.01 * 6 * rho0
-    v2rhosigma = .02 * .5 * (gamma+.001)**(-.5)
+    v2rhosigma = np.zeros(gamma.shape)
     v2sigma2 = 0.02 * .5 * -.5 * (gamma+.001)**(-1.5)
-    v2lapl2 = None
-    vtau2 = None
-    v2rholapl = None
-    v2rhotau = None
-    v2lapltau = None
-    v2sigmalapl = None
-    v2sigmatau = None
     # 2nd order functional derivative
-    fxc = (v2rho2, v2rhosigma, v2sigma2, v2lapl2, vtau2, v2rholapl, v2rhotau, v2lapltau, v2sigmalapl, v2sigmatau)
+    fxc = (v2rho2, v2rhosigma, v2sigma2)
     kxc = None  # 3rd order functional derivative
 
     # Mix with existing functionals
@@ -58,15 +50,42 @@ def eval_xc(xc_code, rho, spin=0, relativity=0, deriv=1, omega=None, verbose=Non
     return exc, vxc, fxc, kxc
 
 mf = dft.RKS(mol)
-mf = mf.define_xc_(eval_xc, 'GGA', hyb=hybrid_coeff)
+mf = mf.define_xc_(eval_gga_xc, 'GGA', hyb=hybrid_coeff)
 mf.verbose = 4
 mf.kernel()
 
+def eval_mgga_xc(xc_code, rho, spin=0, relativity=0, deriv=1, omega=None, verbose=None):
+    # A fictitious XC functional to demonstrate the usage
+    rho0, dx, dy, dz, tau = rho
+    gamma = (dx**2 + dy**2 + dz**2)
+    exc = .01 * rho0**2 + .02 * (gamma+.001)**.5 +  0.01 * tau**2
+    vrho = .01 * 3 * rho0**2 + .02 * (gamma+.001)**.5
+    vgamma = .02 * .5 * (gamma+.001)**(-.5)
+    vtau = 0.02 * tau
+    vxc = (vrho, vgamma, vtau)
+    v2rho2 = 0.01 * 6 * rho0
+    v2rhosigma = np.zeros(gamma.shape)
+    v2sigma2 = 0.02 * .5 * -.5 * (gamma+.001)**(-1.5)
+    v2tau2 = np.full(tau.shape, 0.02)
+    v2rhotau = np.zeros(tau.shape)
+    v2sigmatau = np.zeros(tau.shape)
+    # 2nd order functional derivative
+    fxc = (v2rho2, v2rhosigma, v2sigma2, v2tau2, v2rhotau, v2sigmatau)
+    kxc = None  # 3rd order functional derivative
+
+    # Mix with existing functionals
+    pbe_xc = dft.libxc.eval_xc('pbe,pbe', rho[:4], 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
+
 # half exact exchange in which 40% of the exchange is computed with short
 # range part of the range-separation Coulomb operator (omega = 0.8)
 beta = 0.2
 rsh_coeff = (0.8, hybrid_coeff-beta, beta)
 mf = dft.RKS(mol)
-mf = mf.define_xc_(eval_xc, 'GGA', rsh=rsh_coeff)
+mf = mf.define_xc_(eval_mgga_xc, 'MGGA', rsh=rsh_coeff)
 mf.verbose = 4
 mf.kernel()
diff --git a/pyscf/dft/libxc.py b/pyscf/dft/libxc.py
index 2d703479d3..4f67f3c652 100644
--- a/pyscf/dft/libxc.py
+++ b/pyscf/dft/libxc.py
@@ -770,7 +770,7 @@ def eval_xc(xc_code, rho, spin=0, relativity=0, deriv=1, omega=None, verbose=Non
           | vtau[:,2]   = (u, d)
 
         * fxc for restricted case:
-          (v2rho2, v2rhosigma, v2sigma2, v2lapl2, vtau2, v2rholapl, v2rhotau, v2lapltau, v2sigmalapl, v2sigmatau)
+          (v2rho2, v2rhosigma, v2sigma2, v2lapl2, v2tau2, v2rholapl, v2rhotau, v2lapltau, v2sigmalapl, v2sigmatau)
 
         * fxc for unrestricted case:
           | v2rho2[:,3]     = (u_u, u_d, d_d)
@@ -1186,16 +1186,33 @@ def _eval_xc1(xc_code, rho, spin=0, deriv=1, omega=None):
         outlen = lib.comb(xlen+deriv, deriv)
         exc, vxc, fxc, kxc = libxc_out[:4]
         out = [exc]
-        if vxc is not None:
+        if deriv > 0:
+            assert 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:
+        if deriv > 1:
+            assert fxc is not None
+            if xctype == 'GGA':
+                assert len(fxc) == 3, 'fxc for GGA should be arranged as (v2rho2, v2rhosigma, v2sigma2)'
+            elif xctype == 'MGGA':
+                if len(fxc) == 10:
+                    fxc = [fxc[i] for i in [0, 1, 2, 6, 4, 9]]
+                else:
+                    assert len(fxc) == 6, (
+                        'fxc for MGGA should be arranged as\n'
+                        '(v2rho2, v2rhosigma, v2sigma2, v2tau2, v2rhotau, v2sigmatau)\nor\n'
+                        '(v2rho2, v2rhosigma, v2sigma2, v2lapl2, v2tau2, '
+                        'v2rholapl, v2rhotau, v2lapltau, v2sigmalapl, v2sigmatau)')
+            assert all(x is not None for x in fxc)
+            out.extend(fxc)
+        if deriv > 2:
+            assert 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)