feat: TDDFT-ris store Tia Tij Tab in RAM (#461)

* store Tia Tij Tab in RAM

* ris class API change: GS -> gram_schmidt

* merge get_ab into new ris.py structure

* self.X -> self.xy; merge zhichen's chkfile attr

* print -> log.info

* log info more attri; fancy indexing lower_inv_eri2c

* converge TDA calculation in 5 iterations, enjoy it

---------

Co-authored-by: zhouzehao.dbsv <zhouzehao.dbsv@bytedance.com>

Author: John-zzh

examples/42_ris_preconditioned_TDA.py

@@ -0,0 +1,48 @@
+# Copyright 2021-2025 The PySCF Developers. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import gpu4pyscf.tdscf.ris as ris
+import pyscf
+
+from gpu4pyscf import dft
+from gpu4pyscf.tdscf import _krylov_tools
+from pyscf.data.nist import HARTREE2EV
+
+# this example shows how to use the TDA-rid (up to d orbitals) preconditioned TDA solver
+# converged the TDA calculation in 5 iterations
+
+mol = pyscf.M(atom='Vitamin_C.xyz', basis='def2-svp', verbose=3)
+
+mf = dft.RKS(mol, xc='pbe0')
+
+mf=mf.to_gpu()
+mf.kernel()
+
+ 
+
+td = mf.TDA()
+# construct the TDA matrix vector product function and the hdiag
+vind, hdiag = td.gen_vind(td._scf)
+
+#construct the preconditioner,  the TDA-ris matrix vector product function (and the hdiag)gi
+tda_ris = ris.TDA(mf, J_fit='spd')  
+ris_mvp, _hdiag = tda_ris.gen_vind()
+
+
+# use the nested krylov solver, instead of the default krylov solver
+energies, X = _krylov_tools.nested_krylov_solver(matrix_vector_product=vind,hdiag=hdiag, 
+                                                problem_type='eigenvalue', n_states=5,
+                                                init_mvp=ris_mvp, precond_mvp=ris_mvp)
+
+print('energies', energies*HARTREE2EV) 

gpu4pyscf/tdscf/_lr_eig.py

@@ -23,7 +23,7 @@
 import scipy.linalg
 import cupyx.scipy.linalg
 from gpu4pyscf.tdscf import math_helper
-import time
+from functools import partial
 from pyscf.lib.parameters import MAX_MEMORY
 from gpu4pyscf.lib import logger
 from pyscf.lib.linalg_helper import _sort_elast, _outprod_to_subspace
@@ -1119,7 +1119,7 @@ def Davidson(matrix_vector_product,
                     N_states=20,
                     conv_tol=1e-5,
                     max_iter=25,
-                    GS=True,
+                    gram_schmidt=True,
                     single=False,
                     verbose=logger.INFO):
     '''
@@ -1144,7 +1144,7 @@ def Davidson(matrix_vector_product,
              convergence tolerance
         max_iter: int
              maximum iterations
-        GS: bool
+        gram_schmidt: bool
              use Gram-Schmidt orthogonalization
         single: bool
              use single precision
@@ -1196,9 +1196,9 @@ def Davidson(matrix_vector_product,
     '''
     V_holder = math_helper.TDA_diag_initial_guess(V_holder=V_holder, N_states=size_new, hdiag=hdiag)
 
-    if GS:
+    if gram_schmidt:
         log.info('Using Gram-Schmidt orthogonalization')
-        fill_holder = math_helper.Gram_Schmidt_fill_holder
+        fill_holder = partial(math_helper.Gram_Schmidt_fill_holder, double=True)
     else:
         log.info('Using non-orthogonalized Krylov subspace (nKs) method.')
 
@@ -1247,7 +1247,7 @@ def Davidson(matrix_vector_product,
         omega[:N_states] are smallest N_states eigenvalues
         '''
         t0 = log.init_timer()
-        if GS:
+        if gram_schmidt:
             omega, x = cp.linalg.eigh(sub_A)
         else:
             s_holder = math_helper.gen_VW(s_holder, V_holder, V_holder, size_old, size_new, symmetry=False)
@@ -1320,7 +1320,7 @@ def Davidson_Casida(matrix_vector_product,
                         N_states=20,
                         conv_tol=1e-5,
                         max_iter=25,
-                        GS=True,
+                        gram_schmidt=True,
                         single=False,
                         verbose=logger.NOTE):
     '''
@@ -1345,7 +1345,7 @@ def Davidson_Casida(matrix_vector_product,
             convergence tolerance
         max_iter: int
             maximum number of iterations
-        GS: bool
+        gram_schmidt: bool
             use Gram-Schmidt orthogonalization
         single: bool
             use single precision
@@ -1425,7 +1425,7 @@ def Davidson_Casida(matrix_vector_product,
                                                 N_states=size_new,
                                                 hdiag=hdiag)
 
-    if GS:
+    if gram_schmidt:
         log.info('Using Gram-Schmidt orthogonalization')
         fill_holder = math_helper.VW_Gram_Schmidt_fill_holder
     else:
@@ -1527,7 +1527,7 @@ def Davidson_Casida(matrix_vector_product,
                                                             hdiag=hdiag)
 
         '''
-        GS and symmetric orthonormalization
+        gram_schmidt and symmetric orthonormalization
         '''
         t0 = log.init_timer()
         size_old = size_new