Merge pull request #129 from shreya-chem/lasscf_spin_sep_3rdm

LASSCF spin separated 3-RDMs

exploratory/citools/lasci_ominus1.py

@@ -33,7 +33,7 @@ def _n_m_s (dm1s, dm2s, _print_fn=print):
     _print_fn ('<N>,<Sz>,<S^2> = %f, %f, %f', n, m, ss)
 
 def kernel (fci, h1, h2, norb, nelec, norb_f=None, ci0_f=None,
-            tol=1e-8, gtol=1e-4, max_cycle=15000, 
+            tol=1e-8, gtol=1e-4, max_cycle=15000,
             orbsym=None, wfnsym=None, ecore=0, **kwargs):
 
     if norb_f is None: norb_f = getattr (fci, 'norb_f', [norb])
@@ -183,7 +183,7 @@ def contract_ss (fci, ci, norb, nelec):
     return fcivec
 
 def transform_ci_for_orbital_rotation (fci, ci, norb, nelec, umat):
-    fcivec = np.zeros_like (ci) 
+    fcivec = np.zeros_like (ci)
     for ne in product (range (norb+1), repeat=2):
         c = np.squeeze (fockspace.fock2hilbert (ci, norb, ne))
         c = fci_addons.transform_ci_for_orbital_rotation (c, norb, ne, umat)
@@ -257,7 +257,7 @@ def unpack (self, x_):
         x[self.var_mask] = x_[:]
         xconstr, x = x[:self.nconstr], x[self.nconstr:]
 
-        xcc = x[:self.uop.ngen_uniq] 
+        xcc = x[:self.uop.ngen_uniq]
         x = x[self.uop.ngen_uniq:]
 
         xci = []
@@ -332,7 +332,7 @@ def hc_x (self, x, h):
         huc = self.contract_h2 (h, uc)
         uhuc = self.uop (huc, transpose=True)
         return c, uc, huc, uhuc, c_f
-        
+
     def contract_h2 (self, h, ci, norb=None):
         if norb is None: norb = self.norb
         hci = h[0] * ci
@@ -343,7 +343,7 @@ def contract_h2 (self, h, ci, norb=None):
             hc = direct_spin1.contract_2e (h2eff, ci_h, norb, nelec)
             hci += np.squeeze (fockspace.hilbert2fock (hc, norb, nelec))
         return hci
-            
+
     def dp_ci (self, ci_f):
         norb, norb_f = self.norb, self.norb_f
         ci = np.ones ([1,1], dtype=ci_f[0].dtype)
@@ -357,7 +357,7 @@ def constr_h (self, xconstr, h):
         x = xconstr[0]
         norb, nelec = self.norb, self.nelec
         h = [h[0] - (x*nelec), h[1] + (x*np.eye (self.norb)), h[2]]
-        return h 
+        return h
 
     def rotate_ci0 (self, xci_f):
         ci0, norb = self.ci_f, self.norb
@@ -408,12 +408,78 @@ def get_jac_t1 (self, x, h, c=None, huc=None, uhuc=None):
         xconstr, xcc, xci_f = self.unpack (x)
         self.uop.set_uniq_amps_(xcc)
         if (c is None) or (uhuc is None):
-            c, _, _, uhuc = self.hc_x (x, h)[:4] 
+            c, _, _, uhuc = self.hc_x (x, h)[:4]
         for duc, uhuc_i in zip (self.uop.gen_deriv1 (c, _full=False), self.uop.gen_partial (uhuc)):
             g.append (2*duc.ravel ().dot (uhuc_i.ravel ()))
         g = self.uop.product_rule_pack (g)
         return np.asarray (g)
 
+
+    def get_grad_t1(self, x, h, c=None, huc=None, uhuc=None, epsilon=0.0):
+        """
+        Compute the gradients and relevant indices based on input values.
+
+        Parameters:
+        - x: array-like
+            Input array to unpack and set unique amplitudes.
+
+        - h: array-like
+            Some form of input data to be used for computation.
+
+        - c (optional): array-like
+            Precomputed value; if not provided, it will be computed using hc_x.
+
+        - huc (optional): array-like
+            Precomputed value; if not provided, it will be computed using hc_x.
+
+        - uhuc (optional): array-like
+            Precomputed value; if not provided, it will be computed using hc_x.
+
+        - epsilon (optional): float, default=0.0
+            Threshold value for considering a gradient. If epsilon is 0, all gradients are considered.
+
+        Returns:
+        - tuple
+            all_g: list of all computed gradients
+            g: list of gradients above the epsilon threshold
+            gen_indices: list of indices representing a_idx and i_idx
+            a_idxs_lst: list of a_idx values
+            i_idxs_lst: list of i_idx values
+            len(a_idxs_lst): length of a_idx list
+            len(i_idxs_lst): length of i_idx list
+        """
+
+        g = []
+        all_g = []
+
+        # Unpack and set unique amplitudes
+        xconstr, xcc, xci_f = self.unpack(x)
+        self.uop.set_uniq_amps_(xcc)
+
+        # Compute 'c' and 'uhuc' if not provided
+        if (c is None) or (uhuc is None):
+            c, _, _, uhuc = self.hc_x(x, h)[:4]
+
+        gen_indices = []
+        a_idxs_lst = []
+        i_idxs_lst = []
+        # print("self.uop.init_a_idxs[i]",self.uop.init_a_idxs)
+        for i, (duc, uhuc_i) in enumerate(zip(self.uop.gen_deriv1(c, _full=False), self.uop.gen_partial(uhuc))):
+            gradient = 2 * duc.ravel().dot(uhuc_i.ravel())
+            all_g.append((gradient, i))
+
+            # Allow all gradients if epsilon is 0, else use the abs gradient condition
+            if epsilon == 0.0 or abs(gradient) > epsilon:
+                g.append((gradient, i))
+                a_idx = self.uop.a_idxs[i]
+                i_idx = self.uop.i_idxs[i]
+
+                gen_indices.append((a_idx, i_idx))
+                a_idxs_lst.append(a_idx)
+                i_idxs_lst.append(i_idx)
+
+        return all_g, g, gen_indices, a_idxs_lst, i_idxs_lst, len(a_idxs_lst), len(i_idxs_lst)
+
     def get_jac_ci (self, x, h, uhuc=None, uci_f=None):
         # "uhuc": e^-T1 H e^T1 U|ci0>
         # "jacci": Jacobian elements for the CI degrees of freedom
@@ -542,7 +608,7 @@ def print_x (self, x, h, _print_fn=print, ci_maxlines=10, jac=None):
     def finalize_(self):
         ''' Update self.ci_f and set the corresponding part of self.x to zero '''
         xconstr, xcc, xci_f = self.unpack (self.x)
-        self.ci_f = self.rotate_ci0 (xci_f) 
+        self.ci_f = self.rotate_ci0 (xci_f)
         for xc in xci_f: xc[:] = 0.0
         self.x = self.pack (xconstr, xcc, xci_f)
         return self.x
@@ -564,7 +630,7 @@ class FCISolver (direct_spin1.FCISolver):
     def build_psi (self, *args, **kwargs):
         return LASUCCTrialState (self, *args, **kwargs)
     def save_psi (self, fname, psi):
-        psi_raw = np.concatenate ([psi.x,] 
+        psi_raw = np.concatenate ([psi.x,]
             + [c.ravel () for c in psi.ci_f])
         np.save (fname, psi_raw)
     def load_psi (self, fname, norb, nelec, norb_f=None, **kwargs):
@@ -584,5 +650,3 @@ def get_uop (self, norb, norb_f):
         for i,j in zip (np.cumsum (norb_f)-norb_f, np.cumsum(norb_f)):
             freeze_mask[i:j,i:j] = True
         return get_uccs_op (norb, freeze_mask=freeze_mask)
-
-

my_pyscf/grad/lasscf.py

@@ -72,11 +72,7 @@ def grad_elec(las_grad, mo_coeff=None, ci=None, atmlst=None, verbose=None):
     mo_occ = mo_coeff[:,:nocc]
     mo_core = mo_coeff[:,:ncore]
     mo_cas = mo_coeff[:,ncore:nocc]
-    #print ("SV nocc, ncas, ncore = ", nocc, ncas, ncore)
-    #print ("mo_cas shape =", mo_cas.shape)
-    #print ("SV entering grad_elec in lasscf")
     lasdm1 = las.make_casdm1()
-    #print ("SV lasdm1 shape = ",lasdm1.shape)
     lasdm2 = las.make_casdm2()
     #casdm1, casdm2 = mc.fcisolver.make_rdm12(ci, ncas, nelecas)