Repository: pyscf/pyscf — pull request #2574 (python)
Title: Bugs in libxc.define_xc_ (fix #2551)

Fix `dft.libxc.define_xc_` and update `examples/dft/24-define_xc_functional.py` (fix #2551)

Resolved issues:
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: ...
```

Hints:

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)


Test patch:
diff --git a/pyscf/dft/test/test_libxc.py b/pyscf/dft/test/test_libxc.py
index 6dc223789d..04b28d4ded 100644
--- a/pyscf/dft/test/test_libxc.py
+++ b/pyscf/dft/test/test_libxc.py
@@ -218,7 +218,7 @@ def test_lyp(self):
     #    rho_b = numpy.array([[4.53272893e-06, 4.18968775e-06,-2.83034672e-06, 2.61832978e-06, 5.63360737e-06, 8.97541777e-07]]).T
     #    e, v = dft.libxc.eval_xc('tpss,', (rho_a, rho_b), spin=1, deriv=1)[:2]
 
-    def test_define_xc(self):
+    def test_define_xc_gga(self):
         def eval_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]
@@ -242,11 +242,19 @@ def eval_xc(xc_code, rho, spin=0, relativity=0, deriv=1, omega=None, verbose=Non
         self.assertAlmostEqual(lib.fp(vxc[0]), 0.0065565189784811129, 9)
         self.assertAlmostEqual(lib.fp(vxc[1]), 0.0049270110162854116, 9)
 
+        with self.assertRaises(AssertionError):
+            ni.eval_xc_eff(None, rho, deriv=2)
+
+        # Ensure the xc_code in the input has no impact to the result.
         n, exc, vxc = ni.nr_rks(mol, mf.grids, 'm06x', dm)
         self.assertAlmostEqual(n, 4, 5)
         self.assertAlmostEqual(lib.fp(exc), 0.01197588220700074, 6)
         self.assertAlmostEqual(lib.fp(vxc), -0.043974912389152986, 6)
 
+        ni = dft.libxc.define_xc(ni, eval_xc, 'MGGA')
+        with self.assertRaises(AssertionError):
+            ni.eval_xc_eff(None, rho, deriv=1)
+
         mf = mf.define_xc_('0.5*B3LYP+0.5*B3LYP')
         exc0, vxc0 = mf._numint.eval_xc(None, rho, 0, deriv=1)[:2]
         exc1, vxc1 = dft.libxc.eval_xc('0.5*B3LYP+0.5*B3LYP', rho, 0, deriv=1)[:2]
@@ -263,6 +271,33 @@ def eval_xc(xc_code, rho, spin=0, relativity=0, deriv=1, omega=None, verbose=Non
         self.assertAlmostEqual(abs(vxc0[0]-vxc1[0]).max(), 0, 9)
         self.assertAlmostEqual(abs(vxc0[1]-vxc1[1]).max(), 0, 9)
 
+    def test_define_xc_mgga(self):
+        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 + .01 * tau**2
+            vrho = .01 * 2 * rho0
+            vgamma = .02 * .5 * (gamma+.001)**(-.5)
+            vtau = .02 * tau
+            vxc = (vrho, vgamma, vtau)
+            zeros = numpy.zeros_like(rho0)
+            fxc = (zeros, zeros, zeros)
+            kxc = None  # 3rd order functional derivative
+            return exc, vxc, fxc, kxc
+
+        mf = dft.RKS(mol)
+        ni = dft.libxc.define_xc(mf._numint, eval_mgga_xc, 'MGGA')
+        numpy.random.seed(1)
+        rho = numpy.random.random((5,10))
+        exc, vxc = ni.eval_xc(None, rho, 0, deriv=1)[:2]
+        self.assertAlmostEqual(lib.fp(exc), 0.016353861820705, 9)
+        self.assertAlmostEqual(lib.fp(vxc[0]), 0.0065565189784811129, 9)
+        self.assertAlmostEqual(lib.fp(vxc[1]), 0.0049270110162854116, 9)
+
+        with self.assertRaises(AssertionError):
+            ni.eval_xc_eff(None, rho, deriv=2)
+
     def test_m05x(self):
         rho =(numpy.array([1., 1., 0., 0., 0., 0.165 ]).reshape(-1,1),
               numpy.array([.8, 1., 0., 0., 0., 0.1050]).reshape(-1,1))