From 18b093226e5ec3882345a824416066f74f677a65 Mon Sep 17 00:00:00 2001
From: Matthew R Hermes <mrhermes@uchicago.edu>
Date: Mon, 9 Dec 2024 10:56:56 -0600
Subject: [PATCH] test min_2frag_tps

---
 my_pyscf/lassi/excitations.py   |  4 ++--
 my_pyscf/lassi/min_2frag_tps.py | 25 +++++++++++++++++++++++--
 2 files changed, 25 insertions(+), 4 deletions(-)

diff --git a/my_pyscf/lassi/excitations.py b/my_pyscf/lassi/excitations.py
index eb50aea8..79e31ff1 100644
--- a/my_pyscf/lassi/excitations.py
+++ b/my_pyscf/lassi/excitations.py
@@ -626,11 +626,11 @@ def truncrot_ham_pq (self, ham_pq, u, vh):
         p = np.prod (lroots)
         nstates = ham_pq.shape[1]
         h_pr = ham_pq[:p,:].reshape (lroots[1],lroots[0],nstates)
-        h_pr = np.dot (uh, np.dot (vh, h_pr)).reshape (nroots*nroots,nstates)
+        h_pr = np.dot (uh, np.dot (vh, h_pr)).reshape (u.shape[1]*v.shape[1],nstates)
         ham_pq = np.append (h_pr, ham_pq[p:,:], axis=0)
         nstates = ham_pq.shape[0]
         h_rp = ham_pq[:,:p].reshape (nstates,lroots[1],lroots[0])
-        h_rp = lib.einsum ('rij,ia,jb->rab', h_rp, u, v).reshape (nstates,nroots*nroots)
+        h_rp = lib.einsum ('rij,ia,jb->rab', h_rp, u, v).reshape (nstates,u.shape[1]*v.shape[1])
         ham_pq = np.append (h_rp, ham_pq[:,p:], axis=1)
         return ham_pq
 
diff --git a/my_pyscf/lassi/min_2frag_tps.py b/my_pyscf/lassi/min_2frag_tps.py
index 1462ec4e..71ad71e1 100644
--- a/my_pyscf/lassi/min_2frag_tps.py
+++ b/my_pyscf/lassi/min_2frag_tps.py
@@ -153,14 +153,35 @@ def quadratic_step (ham_pq, ci1, si_p, si_q):
     kappa_u[nroots:,:nroots] = x[:nelu].reshape (
             lroots[1]-nroots, nroots)
     kappa_u -= kappa_u.T
-    u = linalg.expm (kappa_u)[:,:nroots]
+    u = linalg.expm (-kappa_u/2)[:,:nroots]
     kappa_v = np.zeros ((lroots[0],lroots[0]), dtype=x.dtype)
     kappa_v[nroots:,:nroots] = x[:nelv].reshape (
             lroots[0]-nroots, nroots)
     kappa_v -= kappa_v.T
-    vh = linalg.expm (kappa_v)[:,:nroots].conj ().T
+    vh = linalg.expm (-kappa_v/2)[:,:nroots].conj ().T
     return u, vh
 
+def idx_down_ham_pq (ham_pq, lroots, nroots):
+    p = np.prod (lroots)
+    q = ham_pq.shape[0] - p
+    h_pr = ham_pq[:p,:].reshape (lroots[1],lroots[0],p+q)
+    h_pr = h_pr[:nroots,:nroots,:].reshape (nroots*nroots,p+q)
+    ham_pq = np.append (h_pr, ham_pq[p:,:], axis=0)
+    h_rp = ham_pq[:,:p].reshape (nroots*nroots+q,lroots[1],lroots[0])
+    h_rp = h_rp[:,:nroots,:nroots].reshape (nroots*nroots+q,nroots*nroots)
+    ham_pq = np.append (h_rp, ham_pq[:,p:], axis=1)
+    return ham_pq
+
+def subspace_eig (ham_pq, lroots, nroots):
+    p = np.prod (lroots)
+    q = ham_pq.shape[0] - p
+    p0 = nroots*nroots
+    evals, evecs = linalg.eigh (idx_down_ham_pq (ham_pq, lroots, nroots))
+    si_p = np.zeros ((lroots[1],lroots[0],p0+q), dtype=evecs.dtype)
+    si_p[:nroots,:nroots,:] = evecs[:p0,:].reshape (nroots,nroots,p0+q)
+    si = np.append (si_p.reshape (p,p0+q), evecs[p0:,:], axis=0)
+    return evals, si
+