small change

ankit76 · Oct 13, 2024 · 433bb18 · 433bb18
1 parent 3c5d497
commit 433bb18
Showing 1 changed file with 10 additions and 14 deletions.
diff --git a/ad_afqmc/propagation.py b/ad_afqmc/propagation.py
@@ -312,13 +312,11 @@ def _build_propagation_intermediates(
             optimize="optimal",
         )
         h1_mod = (ham_data["h1"][0] + ham_data["h1"][1]) / 2.0 - v0
-        h1_mod = h1_mod - jnp.real(
-            1.0j
-            * jnp.einsum(
-                "g,gik->ik",
-                ham_data["mf_shifts"],
-                ham_data["chol"].reshape(-1, trial.norb, trial.norb),
-            )
+        mf_shifts_r = (1.0j * ham_data["mf_shifts"]).real
+        h1_mod = h1_mod - jnp.einsum(
+            "g,gik->ik",
+            mf_shifts_r,
+            ham_data["chol"].reshape(-1, trial.norb, trial.norb),
         )
         ham_data["exp_h1"] = jsp.linalg.expm(-self.dt * h1_mod / 2.0)
         return ham_data
@@ -490,13 +488,11 @@ def _build_propagation_intermediates(
             ham_data["chol"].reshape(-1, trial.norb, trial.norb),
             optimize="optimal",
         )
-        v1 = jnp.real(
-            1.0j
-            * jnp.einsum(
-                "g,gik->ik",
-                ham_data["mf_shifts"],
-                ham_data["chol"].reshape(-1, trial.norb, trial.norb),
-            )
+        mf_shifts_r = (1.0j * ham_data["mf_shifts"]).real
+        v1 = jnp.einsum(
+            "g,gik->ik",
+            mf_shifts_r,
+            ham_data["chol"].reshape(-1, trial.norb, trial.norb),
         )
         h1_mod = ham_data["h1"] - jnp.array([v0 + v1, v0 + v1])
         ham_data["exp_h1"] = jnp.array(