From 2fcc63b25f63dc21a416bebd98fa86cfc8d2fb75 Mon Sep 17 00:00:00 2001
From: Matthew R Hermes <mrhermes@uchicago.edu>
Date: Fri, 13 Dec 2024 14:01:25 -0600
Subject: [PATCH] Record/report disc svals in lassis

---
 my_pyscf/lassi/excitations.py   | 10 +++++-----
 my_pyscf/lassi/lassis.py        | 26 +++++++++++++++-----------
 tests/lassi/test_excitations.py |  6 +++---
 3 files changed, 23 insertions(+), 19 deletions(-)

diff --git a/my_pyscf/lassi/excitations.py b/my_pyscf/lassi/excitations.py
index 79e31ff1..39326612 100644
--- a/my_pyscf/lassi/excitations.py
+++ b/my_pyscf/lassi/excitations.py
@@ -194,7 +194,7 @@ def kernel (self, h1, h2, ecore=0, ci0=None,
         tdm1s_f = self.get_tdm1s_f (ci1, ci1, norb_f, nelec_f)
         e = 0
         grad_max = conv_tol_grad * 100
-        disc_sval_sum = sum (disc_svals)
+        disc_sval_max = max (list(disc_svals)+[0.0,])
         converged = False
         log.info ('Entering product-state fixed-point CI iteration')
         for it in range (max_cycle_macro):
@@ -204,8 +204,8 @@ def kernel (self, h1, h2, ecore=0, ci0=None,
             e, si = self.eig1 (ham_pq, ci0)
             _, u, si_p, si_q, vh = self.schmidt_trunc (si, ci0, nroots=nroots)
             ci1 = self.truncrot_ci (ci0, u, vh)
-            log.info ('Cycle %d: max grad = %e ; e = %e, |delta| = %e, ||discarded|| = %e',
-                      it, grad_max, e, e - e_last, disc_sval_sum)
+            log.info ('Cycle %d: max grad = %e ; e = %e, |delta| = %e, max (discarded) = %e',
+                      it, grad_max, e, e - e_last, disc_sval_max)
             if ((grad_max < conv_tol_grad) and (abs (e-e_last) < conv_tol_self)):
                 converged = True
                 break
@@ -247,7 +247,7 @@ def kernel (self, h1, h2, ecore=0, ci0=None,
             hci_pspace_diag = self.truncrot_hci_pspace_diag (hci_pspace_diag, u, vh)
             tdm1s_f = self.truncrot_tdm1s_f (tdm1s_f, u, vh)
             log.debug ('Discarded singular values: {}'.format (disc_svals))
-            disc_sval_sum = sum (disc_svals)
+            disc_sval_max = max (list (disc_svals) + [0.0,])
         conv_str = ['NOT converged','converged'][int (converged)]
         log.info (('Product_state fixed-point CI iteration {} after {} '
                    'cycles').format (conv_str, it))
@@ -261,7 +261,7 @@ def kernel (self, h1, h2, ecore=0, ci0=None,
         for ifrag, c in zip (self.excited_frags, ci1_active):
             ci1[ifrag] = np.asarray (c)
         t1 = self.log.timer ('ExcitationPSFCISolver kernel', *t0)
-        return converged, energy_elec, ci1
+        return converged, energy_elec, ci1, disc_sval_max
 
     def get_nq (self):
         lroots = get_lroots ([self.ci_ref[ifrag] for ifrag in self.excited_frags])
diff --git a/my_pyscf/lassi/lassis.py b/my_pyscf/lassi/lassis.py
index 5b35ebaa..3cf86097 100644
--- a/my_pyscf/lassi/lassis.py
+++ b/my_pyscf/lassi/lassis.py
@@ -96,13 +96,13 @@ def prepare_fbf (lsi, ci_ref, ci_sf, ci_ch, ncharge=1, nspin=0, sa_heff=True,
     # 3. Charge excitations
     if ncharge:
         las2 = all_single_excitations (las1)
-        conv_ch, ci_ch = single_excitations_ci (
+        conv_ch, ci_ch, max_disc_sval = single_excitations_ci (
             lsi, las2, las1, ci_ch, ncharge=ncharge, sa_heff=sa_heff,
             deactivate_vrv=deactivate_vrv, spin_flips=spin_flips, crash_locmin=crash_locmin,
             ham_2q=ham_2q
         )
     log.timer ("LASSIS fragment basis functions preparation", *t0)
-    return conv_sf and conv_ch, ci_sf, ci_ch
+    return conv_sf and conv_ch, ci_sf, ci_ch, max_disc_sval
 
 def filter_single_excitation_spin_shuffles (lsi, spaces, nroots_ref=1):
     spaces_ref = spaces[:nroots_ref]
@@ -202,6 +202,7 @@ def single_excitations_ci (lsi, las2, las1, ci_ch, ncharge=1, sa_heff=True, deac
     converged = True
     spaces = filter_single_excitation_spin_shuffles (lsi, spaces, nroots_ref=las1.nroots)
     keys = set ()
+    max_max_disc = 0
     for i in range (las1.nroots, len (spaces)):
         # compute lroots
         psref_ix = [j for j, space in enumerate (spaces[:las1.nroots])
@@ -264,10 +265,10 @@ def single_excitations_ci (lsi, las2, las1, ci_ch, ncharge=1, sa_heff=True, deac
             ci0[1] = mdown (ci0[1], norb_a, nelec_a, smult_a)
             if lroots[afrag,i] == 1 and ci0[1].ndim==3: ci0[1] = ci0[1][0]
         ci0 = [ci0[int (afrag<ifrag)], ci0[int (ifrag<afrag)]]
-        conv, e_roots[i], ci1 = psexc.kernel (h1, h2, ecore=h0, ci0=ci0,
-                                              max_cycle_macro=lsi.max_cycle_macro,
-                                              conv_tol_self=lsi.conv_tol_self,
-                                              nroots=ncharge_i)
+        conv, e_roots[i], ci1, disc_svals_max = psexc.kernel (
+            h1, h2, ecore=h0, ci0=ci0, max_cycle_macro=lsi.max_cycle_macro,
+            conv_tol_self=lsi.conv_tol_self, nroots=ncharge_i
+        )
         ci_ch_ias[0] = mup (ci1[ifrag], norb_i, nelec_i, smult_i)
         if lroots[ifrag,i]==1 and ci_ch_ias[0].ndim == 2:
             ci_ch_ias[0] = ci_ch_ias[0][None,:,:]
@@ -279,8 +280,10 @@ def single_excitations_ci (lsi, las2, las1, ci_ch, ncharge=1, sa_heff=True, deac
         spaces[i].ci = ci1
         if not conv: log.warn ("CI vectors for charge-separated rootspace %s not converged",keystr)
         converged = converged and conv
+        log.info ('Electron hop {} max disc sval: {}'.format (keystr, disc_svals_max))
+        max_max_disc = max (max_max_disc, disc_svals_max)
         t0 = log.timer ("Electron hop {}".format (keystr), *t0)
-    return converged, ci_ch
+    return converged, ci_ch, max_max_disc
 
 class SpinFlips (object):
     '''For a single fragment, bundle the ci vectors of various spin-flipped states with their
@@ -665,10 +668,10 @@ def prepare_states_(self, ncharge=None, nspin=None, sa_heff=None, deactivate_vrv
         ci_ref = self.get_ci_ref ()
         ci_sf = self.ci_spin_flips
         ci_ch = self.ci_charge_hops
-        self.converged, ci_sf, ci_ch = self.prepare_fbf (ci_ref, ci_sf, ci_ch, ncharge=ncharge,
-                                                         nspin=nspin, sa_heff=sa_heff,
-                                                         deactivate_vrv=deactivate_vrv,
-                                                         crash_locmin=crash_locmin)
+        self.converged, ci_sf, ci_ch, self.max_disc_sval = self.prepare_fbf (
+            ci_ref, ci_sf, ci_ch, ncharge=ncharge, nspin=nspin, sa_heff=sa_heff,
+            deactivate_vrv=deactivate_vrv, crash_locmin=crash_locmin
+        )
         self.ci_spin_flips = ci_sf
         self.ci_charge_hops = ci_ch
 
@@ -682,6 +685,7 @@ def prepare_states_(self, ncharge=None, nspin=None, sa_heff=None, deactivate_vrv
         self.e_states = las.e_states
         log.info ('LASSIS model state summary: %d rootspaces; %d model states; converged? %s',
                   self.nroots, self.get_lroots ().prod (0).sum (), str (self.converged))
+        log.info ('LASSIS overall max disc sval: %e', self.max_disc_sval)
         return self.converged
 
     eig = LASSI.kernel
diff --git a/tests/lassi/test_excitations.py b/tests/lassi/test_excitations.py
index 93936f4a..90cf8100 100644
--- a/tests/lassi/test_excitations.py
+++ b/tests/lassi/test_excitations.py
@@ -133,7 +133,7 @@ def lassi_ref (ci1, iroot):
                                                 smults[iroot,i])
                 nroots = np.amin (lroots[:,iroot])
                 conv, energy_tot, ci1 = psexc.kernel (h1, h2, ecore=h0, davidson_only=True,
-                                                      nroots=nroots)
+                                                      nroots=nroots)[:3]
                 if (opt==0):
                     for i in range (2):
                         c = np.asarray (ci1[i+1]).reshape (nroots,-1)
@@ -237,7 +237,7 @@ def lassi_ref (ci1, iroot):
                     weights = np.ones (lroots[i,iroot]) / lroots[i,iroot]
                     psexc.set_excited_fragment_(1+i, (neleca[iroot,i], nelecb[iroot,i]),
                                                 smults[iroot,i], weights=weights)
-                conv, energy_tot, ci1 = psexc.kernel (h1, h2, ecore=h0)
+                conv, energy_tot, ci1 = psexc.kernel (h1, h2, ecore=h0)[:3]
                 with self.subTest (rootspace=iroot):
                     self.assertTrue (conv)
                     e_roots1, si1 = lassi_ref (ci1, iroot)
@@ -245,7 +245,7 @@ def lassi_ref (ci1, iroot):
                     self.assertAlmostEqual (energy_tot, e_roots1[idx_match], 6)
                     self.assertEqual (idx_match, 0) # local minimum problems
                 conv, energy_tot, ci1 = psexc.kernel (h1, h2, ecore=h0,
-                                                      davidson_only=True)
+                                                      davidson_only=True)[:3]
                 with self.subTest ('davidson only', rootspace=iroot):
                     self.assertTrue (conv)
                     e_roots1, si1 = lassi_ref (ci1, iroot)