From aba0c43f26be68d2094f8200090d03927fdb2750 Mon Sep 17 00:00:00 2001
From: Matthew R Hermes <mrhermes@uchicago.edu>
Date: Fri, 6 Oct 2023 12:42:34 -0500
Subject: [PATCH] LASSI chkfile (needs unittests)

---
 my_pyscf/lassi/excitations.py  |  2 ++
 my_pyscf/mcscf/chkfile.py      | 28 ++++++++++++++++++++--------
 my_pyscf/mcscf/productstate.py |  1 +
 3 files changed, 23 insertions(+), 8 deletions(-)

diff --git a/my_pyscf/lassi/excitations.py b/my_pyscf/lassi/excitations.py
index 0424d78b..0eb71e6d 100644
--- a/my_pyscf/lassi/excitations.py
+++ b/my_pyscf/lassi/excitations.py
@@ -589,6 +589,7 @@ def kernel (self, h1e, h2e, norb, nelec, ecore=0, ci0=None, orbsym=None, **kwarg
         ci1 = ci0
         self.denom_q = e0 - self.e_q
         log.debug ("Self-energy singularities in VRVSolver: {}".format (self.e_q))
+        log.debug ("Denominators in VRVSolver: {}".format (self.denom_q))
         self.test_locmin (e0, ci1, warntag='Saddle-point initial guess')
         warn_swap = False # annoying loud warning not necessary
         #print (lib.fp (ci0), self.denom_q)
@@ -607,6 +608,7 @@ def kernel (self, h1e, h2e, norb, nelec, ecore=0, ci0=None, orbsym=None, **kwarg
                 ket, e0 = ci1, e
             e0 = self.solve_e0 (ecore, h1e, h2e, norb, nelec, ket)
             self.denom_q = e0 - self.e_q
+            log.debug ("Denominators in VRVSolver: {}".format (self.denom_q))
             #hket = self.contract_2e (self.absorb_h1e (h1e, h2e, norb, nelec, 0.5), ket, norb, nelec)
             #brahket = np.dot (ket.ravel (), hket.ravel ())
             #e0_test = ecore + brahket
diff --git a/my_pyscf/mcscf/chkfile.py b/my_pyscf/mcscf/chkfile.py
index 421f9f87..d9067102 100644
--- a/my_pyscf/mcscf/chkfile.py
+++ b/my_pyscf/mcscf/chkfile.py
@@ -2,15 +2,18 @@
 from pyscf.lib.chkfile import load
 from pyscf.lib.chkfile import load_mol, save_mol
 
-keys_config = ['ncas', 'nelecas', 'ncore', 'ncas_sub', 'nelecas_sub']
-keys_saconstr = ['weights', 'charges', 'spins', 'smults', 'wfnsyms']
-keys_results = ['e_states', 'states_converged', 'e_tot', 'mo_coeff']
+KEYS_CONFIG_LASSCF = ['ncas', 'nelecas', 'ncore', 'ncas_sub', 'nelecas_sub']
+KEYS_SACONSTR_LASSCF = ['weights', 'charges', 'spins', 'smults', 'wfnsyms']
+KEYS_RESULTS_LASSCF = ['e_states', 'states_converged', 'e_tot', 'mo_coeff']
 
-def load_las_(mc, chkfile=None):
+def load_las_(mc, chkfile=None, method_key='las', 
+              keys_config=KEYS_CONFIG_LASSCF,
+              keys_saconstr=KEYS_SACONSTR_LASSCF,
+              keys_results=KEYS_RESULTS_LASSCF):
     if chkfile is None: chkfile = mc.chkfile
     if chkfile is None: raise RuntimeError ('chkfile not specified')
-    data = load (chkfile, 'las')
-    if data is None: raise KeyError ('LAS record not in chkfile')
+    data = load (chkfile, method_key)
+    if data is None: raise KeyError ('{} record not in chkfile'.format (method_key.upper()))
 
     # conditionals for backwards compatibility with older chkfiles that
     # only stored mo_coeff and ci
@@ -20,7 +23,13 @@ def load_las_(mc, chkfile=None):
     # this needs to happen before some of the results attributes
     if all ([key in data for key in keys_saconstr]):
         sakwargs = {key: data[key] for key in keys_saconstr}
-        mc.state_average_(**sakwargs)
+        try:
+            mc.state_average_(**sakwargs)
+        except AttributeError as err:
+            las = mc._las.state_average (**sakwargs)
+            mc.fciboxes = las.fciboxes
+            mc.nroots=las.nroots
+            mc.weights=las.weights
     for key in keys_results:
         if key in data:
             setattr (mc, key, data[key])
@@ -47,7 +56,10 @@ def load_las_(mc, chkfile=None):
     return mc
 
 def dump_las (mc, chkfile=None, method_key='las', mo_coeff=None, ci=None,
-              overwrite_mol=True, **kwargs):
+              overwrite_mol=True, keys_config=KEYS_CONFIG_LASSCF,
+              keys_saconstr=KEYS_SACONSTR_LASSCF,
+              keys_results=KEYS_RESULTS_LASSCF,
+              **kwargs):
     if chkfile is None: chkfile = mc.chkfile
     if not chkfile: return mc
     if mo_coeff is None: mo_coeff = mc.mo_coeff
diff --git a/my_pyscf/mcscf/productstate.py b/my_pyscf/mcscf/productstate.py
index ce87aee9..bccfe3c4 100644
--- a/my_pyscf/mcscf/productstate.py
+++ b/my_pyscf/mcscf/productstate.py
@@ -40,6 +40,7 @@ def kernel (self, h1, h2, norb_f, nelec_f, ecore=0, ci0=None, orbsym=None,
             grad_max = np.amax (np.abs (grad))
             log.info ('Cycle %d: max grad = %e ; sigma = %e', it, grad_max,
                 e_sigma)
+            log.debug ('e vector = {}'.format (e))
             solvers_converged = [np.all (np.asarray (s.converged)) for s in self.fcisolvers]
             if ((grad_max < conv_tol_grad) and (e_sigma < conv_tol_self)
                 and all ([solvers_converged])):