TDDFT analytical gradient with large basis set not matching numerical gradient

When trying to varify TDDFT gradient with large basis set (d-aug-cc-pvdz), I notice the error between analytical and numerical gradient doesn't show a quadratic decay vs dx.

There's no problem with TDHF. There's no problem with a smaller basis like aug-cc-pvdz. When switching from PBE0 to PBE, the error is reduced by two order of magnitude, but the error is still constant wrt dx. When the grid size is increased (doubled radial grid), the error is reduced by two order of magnitude, but the error is still near constant wrt dx.

Here's a reproducible case, where the problem is most obvious:

import pyscf
import gpu4pyscf
import numpy as np

gpu4pyscf.scf.hf.remove_overlap_zero_eigenvalue = True
gpu4pyscf.scf.hf.overlap_zero_eigenvalue_threshold = 1e-10 # No basis removed

mol = pyscf.M(
    atom = """
        H      0.675261    0.254864    1.141395
        C      0.000000    0.000000    0.000000
        H      0.205100    0.824000   -0.678600
        H      0.334500   -0.931400   -0.449600
        H     -1.121420   -0.177825   -0.199081
    """,
    basis = "d-aug-cc-pvdz",
    charge = 0,
    verbose = 0,
)


mf = mol.RKS(xc = "PBE0").density_fit(auxbasis = "def2-universal-jkfit").to_gpu()
mf.grids.atom_grid = (99,590)
# mf.grids.prune = None
# mf.grids.radii_adjust = None
mf.small_rho_cutoff = 1e-30
mf.conv_tol = 1e-10
test_scf_energy = mf.kernel()
assert mf.converged

td = mf.TDDFT().set(nstates = 3)
assert td.device == 'gpu'
td.conv_tol = 1e-6
td.singlet = True
test_singlet_energies, _ = td.kernel()
assert np.all(td.converged)
test_singlet_energies += test_scf_energy

gobj = td.nuc_grad_method()
gobj.state = 1 # CIS_STATE_DERIV 3
test_singlet_1_gradient = gobj.kernel()


def numerical_tddft_gradient(mol, get_energy, dx = 1e-4):
    numerical_gradient = np.zeros([mol.natm, 3])

    mol_copy = mol.copy()
    for i_atom in range(mol.natm):
        for i_xyz in range(3):
            xyz_p = mol.atom_coords()
            xyz_p[i_atom, i_xyz] += dx
            mol_copy.set_geom_(xyz_p, unit='Bohr')
            mol_copy.build()
            e_p = get_energy(mol_copy)

            xyz_m = mol.atom_coords()
            xyz_m[i_atom, i_xyz] -= dx
            mol_copy.set_geom_(xyz_m, unit='Bohr')
            mol_copy.build()
            e_m = get_energy(mol_copy)

            numerical_gradient[i_atom, i_xyz] = (e_p - e_m) / (2 * dx)

    np.set_printoptions(linewidth = np.iinfo(np.int32).max, threshold = np.iinfo(np.int32).max, precision = 16, suppress = True)
    print(repr(numerical_gradient))
    return numerical_gradient


def get_singlet_1_energy(mol):
    mf = mol.RKS(xc = "PBE0").density_fit(auxbasis = "def2-universal-jkfit").to_gpu()
    mf.grids.atom_grid = (99,590)
    # mf.grids.prune = None
    # mf.grids.radii_adjust = None
    mf.small_rho_cutoff = 1e-30
    mf.conv_tol = 1e-10
    test_scf_energy = mf.kernel()
    assert mf.converged

    td = mf.TDDFT().set(nstates = 3)
    assert td.device == 'gpu'
    td.conv_tol = 1e-6
    td.singlet = True
    test_singlet_energies, _ = td.kernel()
    assert np.all(td.converged)
    test_singlet_energies += test_scf_energy
    return test_singlet_energies[0]


for dx in [4e-3, 2e-3, 1e-3, 5e-4, 2.5e-4, 1.25e-4]:
    ref_singlet_1_gradient = numerical_tddft_gradient(mol, get_singlet_1_energy, dx)
    diff = np.max(np.abs(test_singlet_1_gradient - ref_singlet_1_gradient))
    print(f"dx = {dx}, diff = {diff}")

And the output is:

Start by running the reproducible Python case through gpu4pyscf.scf.hf, RKS, TDDFT, and td.nuc_grad_method(). Compare the analytical gradient with numerical_tddft_gradient across the listed dx values and inspect the TDDFT nuclear-gradient implementation. Done means the large-basis PBE0 analytical and numerical gradients show the expected convergence as dx decreases.