Merge branch 'master' of https://github.com/block-hczhai/block2

block-hczhai · Dec 18, 2021 · 67ea82a · 67ea82a
2 parents d71877c + 4e19a32
commit 67ea82a
Showing 1 changed file with 8 additions and 2 deletions.
diff --git a/pyblock2/driver/block2main b/pyblock2/driver/block2main
@@ -1919,10 +1919,16 @@ if not pre_run:
             dmr.deallocate()
             dm = dm.reshape((dm.shape[0] // 2, 2, dm.shape[1] // 2, 2))
             dm = np.transpose(dm, (0, 2, 1, 3))
-            dm = np.concatenate(
-                [dm[None, :, :, 0, 0], dm[None, :, :, 1, 1]], axis=0)
+            if "use_general_spin" in dic:
+                dm = dm[None]
+            else:
+                dm = np.concatenate(
+                    [dm[None, :, :, 0, 0], dm[None, :, :, 1, 1]], axis=0)
             if orb_idx is not None:
                 rev_idx = np.argsort(orb_idx)
+                if "trans_integral_to_spin_orbital" in dic:
+                    rev_idx = np.array(list(zip(rev_idx*2, rev_idx*2+1))).ravel()
+                assert dm.shape[-1] == len(rev_idx)
                 dm[:, :, :] = dm[:, rev_idx, :][:, :, rev_idx]
             return dm
         else: