SCR3 should be calculated before aux_ri comes

ajz34 · Jun 1, 2021 · 61d78aa · 61d78aa
1 parent ef5094c
commit 61d78aa
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 11 deletions.
diff --git a/pyscf/dh/polar/rdfdh.py b/pyscf/dh/polar/rdfdh.py
@@ -259,18 +259,18 @@ def get_SCR3(self):
         naux = self.df_ri.get_naoaux()
         nprop = self.nprop
 
-        U_1 = tensors.load("U_1")
-        G_ia_ri = tensors.load("G_ia_ri")
-        pdA_G_ia_ri = tensors.load("pdA_G_ia_ri")
-        Y_mo_ri = tensors["Y_mo_ri"]
-
         SCR3 = np.zeros((nprop, self.nvir, self.nocc))
         if self.xc_n:
             pdA_F_0_mo_n = tensors.load("pdA_F_0_mo_n")
             SCR3 += 4 * pdA_F_0_mo_n[:, sv, so]
         if not self.eval_pt2:
             return SCR3
 
+        U_1 = tensors.load("U_1")
+        G_ia_ri = tensors.load("G_ia_ri")
+        pdA_G_ia_ri = tensors.load("pdA_G_ia_ri")
+        Y_mo_ri = tensors["Y_mo_ri"]
+
         nbatch = self.calc_batch_size(10 * self.nmo**2, G_ia_ri.size + pdA_G_ia_ri.size)
         for saux in gen_batch(0, naux, nbatch):
             G_blk = G_ia_ri[saux]

diff --git a/pyscf/dh/polar/udfdh.py b/pyscf/dh/polar/udfdh.py
@@ -306,19 +306,22 @@ def get_SCR3(self):
         nocc, nvir, nmo = self.nocc, self.nvir, self.nmo
         nprop = self.nprop
 
+        SCR3 = [np.zeros((nprop, nvir[σ], nocc[σ])) for σ in (α, β)]
+
+        for σ in (α, β):
+            if self.xc_n:
+                pdA_F_0_mo_n = tensors.load("pdA_F_0_mo_n")
+                SCR3[σ] += 2 * pdA_F_0_mo_n[σ][:, sv[σ], so[σ]]
+        if not self.eval_pt2:
+            return SCR3
+
         U_1 = tensors.load("U_1")
         G_ia_ri = [tensors.load("G_ia_ri" + str(σ)) for σ in (α, β)]
         pdA_G_ia_ri = [tensors.load("pdA_G_ia_ri" + str(σ)) for σ in (α, β)]
         Y_mo_ri = [tensors["Y_mo_ri" + str(σ)] for σ in (α, β)]
 
-        SCR3 = [np.zeros((nprop, nvir[σ], nocc[σ])) for σ in (α, β)]
         nbatch = self.calc_batch_size(10 * nmo**2, tot_size(G_ia_ri, pdA_G_ia_ri, U_1))
         for σ in (α, β):
-            if self.xc_n:
-                pdA_F_0_mo_n = tensors.load("pdA_F_0_mo_n")
-                SCR3[σ] += 2 * pdA_F_0_mo_n[σ][:, sv[σ], so[σ]]
-            if not self.eval_pt2:
-                continue
             for saux in gen_batch(0, naux, nbatch):
                 G_blk = G_ia_ri[σ][saux]
                 Y_blk = np.asarray(Y_mo_ri[σ][saux])