efficiency of transpose for derivative of MO cderi

ajz34 · May 22, 2021 · 7943e50 · 7943e50
1 parent 8ae4412
commit 7943e50
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 3 deletions.
diff --git a/pyscf/dh/dhutil.py b/pyscf/dh/dhutil.py
@@ -212,3 +212,12 @@ def restricted_biorthogonalize(t_ijab, cc, c_os, c_ss):
         res *= coef_1
         res += coef_0 * t_ijab
         return res
+
+
+def hermi_sum_last2dim_inplace(tsr, hermi=1):
+    # shameless call lib.hermi_sum, just for a tensor wrapper
+    tsr_shape = tsr.shape
+    tsr.shape = (-1, tsr.shape[-1], tsr.shape[-2])
+    res = lib.hermi_sum(tsr, axes=(0, 2, 1), hermi=hermi, inplace=True)
+    res.shape = tsr_shape
+    return res
diff --git a/pyscf/dh/polar/rdfdh.py b/pyscf/dh/polar/rdfdh.py
@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from dh import RDFDH
-from dh.dhutil import gen_batch, get_rho_from_dm_gga, restricted_biorthogonalize
+from dh.dhutil import gen_batch, get_rho_from_dm_gga, restricted_biorthogonalize, hermi_sum_last2dim_inplace
 from pyscf import gto, lib, dft
 import numpy as np
 
@@ -265,12 +265,14 @@ def get_SCR3(self):
             pdA_G_blk = np.asarray(pdA_G_ia_ri[:, saux])
             # pdA_Y_ij part
             pdA_Y_blk = einsum("Ami, Pmj -> APij", U_1[:, :, so], Y_blk[:, :, so])
-            pdA_Y_blk += pdA_Y_blk.swapaxes(-1, -2)
+            # pdA_Y_blk += pdA_Y_blk.swapaxes(-1, -2)
+            hermi_sum_last2dim_inplace(pdA_Y_blk)
             SCR3 -= 4 * einsum("APja, Pij -> Aai", pdA_G_blk, Y_blk[:, so, so])
             SCR3 -= 4 * einsum("Pja, APij -> Aai", G_blk, pdA_Y_blk)
             # pdA_Y_ab part
             pdA_Y_blk = einsum("Ama, Pmb -> APab", U_1[:, :, sv], Y_blk[:, :, sv])
-            pdA_Y_blk += pdA_Y_blk.swapaxes(-1, -2)
+            # pdA_Y_blk += pdA_Y_blk.swapaxes(-1, -2)
+            hermi_sum_last2dim_inplace(pdA_Y_blk)
             SCR3 += 4 * einsum("APib, Pab -> Aai", pdA_G_blk, Y_blk[:, sv, sv])
             SCR3 += 4 * einsum("Pib, APab -> Aai", G_blk, pdA_Y_blk)
         if self.xc_n: