Task: Fix initial guess dimension problems when solving for multiple roots (pyscf/pyscf#2306)

org: pyscf
repo: pyscf
pr_number: 2306
issue_title: Fix initial guess dimension problems when solving for multiple roots
issue_body: fix #2292 
resolved_issues: number: 2292
title: PySCF >=v2.2.1 fci.direct_spin1 davidson solver issue
body: Hello

I have the following minimal example using FCI to solve a given Hamiltonian which produces an error when using a PySCF version greater than or equal to v2.2.1.
``` python
from pyscf import fci
import numpy as np

norb=4
nelec=(2, 2)
nroots=36 # (4 choose 2)*(4 choose 2)

e = -0.5
t = -0.1
h1 = np.array([[e, t, 0, 0],
               [t, e, t, 0],
               [0, t, e, t],
               [0, 0, t, e],
              ])

U = 0.1
h2 = np.zeros((norb, norb, norb, norb))
for i in range(norb):
    h2[i,i,i,i] = U

e, fcivec = fci.direct_spin1.kernel(h1, h2, norb, nelec, nroots=nroots, max_space=30, max_cycle=100,
                                    davidson_only=True,
                                   )

print("Energies:", e)
```

The error message:
``` bash
  File "/lib/python3.11/site-packages/pyscf/fci/direct_spin1.py", line 257, in kernel
    return _kfactory(FCISolver, h1e, eri, norb, nelec, ci0, level_shift,
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/lib/python3.11/site-packages/pyscf/fci/direct_spin1.py", line 281, in _kfactory
    e, c = cis.kernel(h1e, eri, norb, nelec, ci0, ecore=ecore, **unknown)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/lib/python3.11/site-packages/pyscf/fci/direct_spin1.py", line 842, in kernel
    e, c = kernel_ms1(self, h1e, eri, norb, nelec, ci0, link_index,
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/lib/python3.11/site-packages/pyscf/fci/direct_spin1.py", line 612, in kernel_ms1
    e, c = fci.eig(hop, ci0, precond, tol=tol, lindep=lindep,
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/lib/python3.11/site-packages/pyscf/fci/direct_spin1.py", line 810, in eig
    lib.davidson1(lambda xs: [op(x) for x in xs],
  File "/lib/python3.11/site-packages/pyscf/lib/linalg_helper.py", line 388, in davidson1
    x0 = x0()
         ^^^^
  File "/lib/python3.11/site-packages/pyscf/fci/direct_spin1.py", line 580, in init_guess
    return fci.get_init_guess(norb, nelec, nroots, hdiag)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/lib/python3.11/site-packages/pyscf/fci/direct_spin1.py", line 829, in get_init_guess
    return get_init_guess(norb, nelec, nroots, hdiag)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/lib/python3.11/site-packages/pyscf/fci/direct_spin1.py", line 469, in get_init_guess
    return _get_init_guess(na, nb, nroots, hdiag, nelec)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/lib/python3.11/site-packages/pyscf/fci/direct_spin1.py", line 444, in _get_init_guess
    addrs = numpy.argpartition(lib.pack_tril(hdiag), nroots-1)[:nroots]
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/lib/python3.11/site-packages/numpy/_core/fromnumeric.py", line 908, in argpartition
    return _wrapfunc(a, 'argpartition', kth, axis=axis, kind=kind, order=order)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/lib/python3.11/site-packages/numpy/_core/fromnumeric.py", line 57, in _wrapfunc
    return bound(*args, **kwds)
           ^^^^^^^^^^^^^^^^^^^^
ValueError: kth(=35) out of bounds (21)
```

I have only encountered this error with versions of PySCF >=v2.2.1, the versions before that run with the correct output. I believe the issue is related to `davidson_only=True` as setting `davidson_only=False` for >=v2.2.1 also does not produce an error.

Is this a bug with the davidson solver?

Thanks!

Reference solution:
fix_patch: diff --git a/pyscf/fci/direct_spin0_symm.py b/pyscf/fci/direct_spin0_symm.py
index ed6feb92f9..5e1e8f1f82 100644
--- a/pyscf/fci/direct_spin0_symm.py
+++ b/pyscf/fci/direct_spin0_symm.py
@@ -74,15 +74,20 @@ def get_init_guess(norb, nelec, nroots, hdiag, orbsym, wfnsym=0):
     strsa = cistring.gen_strings4orblist(range(norb), neleca)
     na = len(strsa)
     airreps = direct_spin1_symm._gen_strs_irrep(strsa, orbsym)
-    hdiag = hdiag.reshape(na,na)
 
     sym_allowed = (airreps[:,None] ^ airreps) == wfnsym
     idx = numpy.arange(na)
     sym_allowed[idx[:,None] < idx] = False
     idx_a, idx_b = numpy.where(sym_allowed)
 
+    hdiag = hdiag.reshape(na,na)[idx_a,idx_b]
+    if hdiag.size <= nroots:
+        hdiag_indices = numpy.arange(hdiag.size)
+    else:
+        hdiag_indices = numpy.argpartition(hdiag, nroots-1)[:nroots]
+
     ci0 = []
-    for k in numpy.argpartition(hdiag[idx_a,idx_b], nroots-1)[:nroots]:
+    for k in hdiag_indices:
         addra, addrb = idx_a[k], idx_b[k]
         x = numpy.zeros((na, na))
         if addra == addrb:
diff --git a/pyscf/fci/direct_spin1.py b/pyscf/fci/direct_spin1.py
index 7f8d04f3d3..710acc872c 100644
--- a/pyscf/fci/direct_spin1.py
+++ b/pyscf/fci/direct_spin1.py
@@ -440,8 +440,11 @@ def _get_init_guess(na, nb, nroots, hdiag, nelec):
     ci0 = []
     neleca, nelecb = _unpack_nelec(nelec)
     if neleca == nelecb and na == nb:
-        hdiag = hdiag.reshape(na, na)
-        addrs = numpy.argpartition(lib.pack_tril(hdiag), nroots-1)[:nroots]
+        hdiag = lib.pack_tril(hdiag.reshape(na, na))
+        if hdiag.size <= nroots:
+            addrs = numpy.arange(hdiag.size)
+        else:
+            addrs = numpy.argpartition(hdiag, nroots-1)[:nroots]
         for addr in addrs:
             addra = (int)((2*addr+.25)**.5 - .5 + 1e-7)
             addrb = addr - addra*(addra+1)//2
@@ -449,7 +452,10 @@ def _get_init_guess(na, nb, nroots, hdiag, nelec):
             x[addra,addrb] = 1
             ci0.append(x.ravel().view(FCIvector))
     else:
-        addrs = numpy.argpartition(hdiag, nroots-1)[:nroots]
+        if hdiag.size <= nroots:
+            addrs = numpy.arange(hdiag.size)
+        else:
+            addrs = numpy.argpartition(hdiag, nroots-1)[:nroots]
         for addr in addrs:
             x = numpy.zeros((na*nb))
             x[addr] = 1
diff --git a/pyscf/fci/direct_spin1_symm.py b/pyscf/fci/direct_spin1_symm.py
index 017505836a..ee10487b8f 100644
--- a/pyscf/fci/direct_spin1_symm.py
+++ b/pyscf/fci/direct_spin1_symm.py
@@ -217,15 +217,20 @@ def _get_init_guess(airreps, birreps, nroots, hdiag, nelec, orbsym, wfnsym=0):
     neleca, nelecb = _unpack_nelec(nelec)
     na = len(airreps)
     nb = len(birreps)
-    hdiag = hdiag.reshape(na,nb)
     sym_allowed = airreps[:,None] == wfnsym ^ birreps
     if neleca == nelecb and na == nb:
         idx = np.arange(na)
         sym_allowed[idx[:,None] < idx] = False
     idx_a, idx_b = np.where(sym_allowed)
 
+    hdiag = hdiag.reshape(na,nb)[idx_a,idx_b]
+    if hdiag.size <= nroots:
+        hdiag_indices = np.arange(hdiag.size)
+    else:
+        hdiag_indices = np.argpartition(hdiag, nroots-1)[:nroots]
+
     ci0 = []
-    for k in np.argpartition(hdiag[idx_a,idx_b], nroots-1)[:nroots]:
+    for k in hdiag_indices:
         addra, addrb = idx_a[k], idx_b[k]
         x = np.zeros((na, nb))
         x[addra,addrb] = 1
@@ -248,7 +253,11 @@ def get_init_guess(norb, nelec, nroots, hdiag, orbsym, wfnsym=0):
         return _get_init_guess(airreps, birreps, nroots, hdiag, nelec, orbsym, wfnsym)
 
     ci0 = []
-    for k in np.argpartition(hdiag, nroots-1)[:nroots]:
+    if hdiag.size <= nroots:
+        hdiag_indices = np.arange(hdiag.size)
+    else:
+        hdiag_indices = np.argpartition(hdiag, nroots-1)[:nroots]
+    for k in hdiag_indices:
         x = np.zeros_like(hdiag)
         x[k] = 1.
         ci0.append(x.ravel().view(direct_spin1.FCIvector))
diff --git a/pyscf/fci/direct_uhf.py b/pyscf/fci/direct_uhf.py
index e5c7cda67a..6aaea941b3 100644
--- a/pyscf/fci/direct_uhf.py
+++ b/pyscf/fci/direct_uhf.py
@@ -198,7 +198,7 @@ def pspace(h1e, eri, norb, nelec, hdiag=None, np=400):
     na = cistring.num_strings(norb, neleca)
     nb = cistring.num_strings(norb, nelecb)
     assert hdiag.size == na * nb
-    if hdiag.size < np:
+    if hdiag.size <= np:
         addr = numpy.arange(hdiag.size)
     else:
         try:
diff --git a/pyscf/scf/dispersion.py b/pyscf/scf/dispersion.py
index d63dd47cd6..4e038420ff 100644
--- a/pyscf/scf/dispersion.py
+++ b/pyscf/scf/dispersion.py
@@ -23,6 +23,7 @@
 import warnings
 from functools import lru_cache
 from pyscf.lib import logger
+from pyscf import scf
 
 # supported dispersion corrections
 DISP_VERSIONS = ['d3bj', 'd3zero', 'd3bjm', 'd3zerom', 'd3op', 'd4']
@@ -115,8 +116,13 @@ def check_disp(mf, disp=None):
     if disp == 0: # disp = False
         return False
 
-    # The disp method for both HF and MCSCF is set to 'hf'
-    method = getattr(mf, 'xc', 'hf')
+    # To prevent mf.do_disp() triggering the SCF.__getattr__ method, do not use
+    # method = getattr(mf, 'xc', 'hf').
+    if isinstance(mf, scf.hf.KohnShamDFT):
+        method = mf.xc
+    else:
+        # Set the disp method for both HF and MCSCF to 'hf'
+        method = 'hf'
     disp_version = parse_disp(method)[1]
 
     if disp is None: # Using the disp version decoded from the mf.xc attribute
@@ -172,6 +178,5 @@ def get_dispersion(mf, disp=None, with_3body=None, verbose=None):
         raise RuntimeError(f'dipersion correction: {disp_version} is not supported.')
 
 # Inject to SCF class
-from pyscf import scf
 scf.hf.SCF.do_disp = check_disp
 scf.hf.SCF.get_dispersion = get_dispersion