[BUG] Probabilistic all-zero roots in SpinFlipTDA (numerical instability in Davidson solver)

[wtpeter]

I encountered an issue where uhf.SpinFlipTDA calculations occasionally yield all-zero (or near-zero) excitation energies. This behavior appears to be probabilistic but is tied to the specific SCF object instance. It seems to be caused by numerical instability during the iterative diagonalization process, as the the full diagonalization of A matrix is correct.

import numpy as np
from pyscf import gto
from gpu4pyscf.tdscf import uhf

def diagonalize_tda(a, nroots=5):
    nocc, nvir = a.shape[:2]
    nov = nocc * nvir
    a = a.reshape(nocc*nvir, nocc*nvir)
    e, xy = np.linalg.eigh(a)
    return e[:nroots]

mol = gto.M(atom='O; O 1 1.08', charge=0, spin=2, basis='631G')
mf = mol.UKS(xc='B3LYP').to_gpu()
mf.kernel()

td = uhf.SpinFlipTDA(mf)
td.nstates = 5
td.extype = 1
td.collinear = 'mcol'
td.collinear_samples = 30
td.kernel()

a, b = td.get_ab()
e_exact = diagonalize_tda(a[1], nroots=td.nstates)

print("Solver Converged:", td.converged)
print("Solver Energies (eV):", td.e * 27.2114)
print("Exact Energies (eV): ", e_exact * 27.2114)

Example output:

Solver Converged: [ True  True  True  True  True]
Solver Energies (eV): [-7.71511441e-14 -2.45620156e-14 -2.25533817e-14  3.12359720e-12
  5.18304186e-11]
Exact Energies (eV):  [0.07848584 0.98143623 0.98143731 2.34050132 9.52798575]

Key Observations

1. Probabilistic: The issue is probabilistic regarding the initial setup but deterministic for a given SCF object (i.e., if it fails for an SCF object, it always fails for that object).
2. Geometry/Basis dependence: It occurs across various functionals and basis sets, but seems more frequent with shorter bond lengths (e.g., O-O at 0.95 Å).
3. Subspace size: Increasing nstates (e.g., from 5 to 6 in the example above) avoids the issue.

Suggested fix:
The issue seems to stem from loss of orthogonality in the Davidson solver. Modifying gpu4pyscf/tdscf/_lr_eig.py to perform the Gram-Schmidt orthogonalization twice (Double Gram-Schmidt) resolves the instability completely.

# Original
# xt -= xs.conj().dot(xt.T).T.dot(xs)
# Proposed Change (Double Gram-Schmidt)
for _ in range(2):
    xt -= xs.conj().dot(xt.T).T.dot(xs)