Add feature for overlaps to check if only the exp value converges

Author: JanLuca

varipeps/overlap/overlap.py

@@ -1,6 +1,7 @@
 import jax.numpy as jnp
 
-from varipeps.ctmrg import calc_ctmrg_env
+from varipeps import varipeps_config
+from varipeps.ctmrg import calc_ctmrg_env, CTMRGNotConvergedError
 from varipeps.peps import PEPS_Unit_Cell
 
 from . import overlap_single_site
@@ -12,7 +13,9 @@
 }
 
 
-def calculate_overlap(unitcell_A, unitcell_B, chi, max_chi):
+def calculate_overlap(
+    unitcell_A, unitcell_B, chi, max_chi, *, test_automatic_overlap_conv=True
+):
     structure_A = tuple(tuple(i) for i in unitcell_A.data.structure)
     structure_B = tuple(tuple(i) for i in unitcell_B.data.structure)
 
@@ -43,10 +46,40 @@ def calculate_overlap(unitcell_A, unitcell_B, chi, max_chi):
 
     overlap_unitcell = PEPS_Unit_Cell.from_tensor_list(overlap_tensors, structure_A)
 
-    overlap_unitcell, _ = calc_ctmrg_env(
-        [i.tensor for i in overlap_tensors], overlap_unitcell
-    )
+    if test_automatic_overlap_conv:
+        tmp_max_steps = varipeps_config.ctmrg_max_steps
+        tmp_fail = varipeps_config.ctmrg_fail_if_not_converged
 
-    overlap_AB = overlap_func(overlap_unitcell)
+        varipeps_config.ctmrg_fail_if_not_converged = False
+        varipeps_config.ctmrg_max_steps = 10
+
+        overlap_unitcell, _ = calc_ctmrg_env(
+            [i.tensor for i in overlap_tensors], overlap_unitcell
+        )
+        overlap_AB = overlap_func(overlap_unitcell)
+
+        for count in range(tmp_max_steps):
+            old_overlap_AB = overlap_AB
+
+            overlap_unitcell, _ = calc_ctmrg_env(
+                [i.tensor for i in overlap_tensors], overlap_unitcell
+            )
+            overlap_AB = overlap_func(overlap_unitcell)
+
+            if (
+                jnp.abs(overlap_AB - old_overlap_AB)
+                <= varipeps_config.ctmrg_convergence_eps
+            ):
+                varipeps_config.ctmrg_fail_if_not_converged = tmp_fail
+                varipeps_config.ctmrg_max_steps = tmp_max_steps
+                break
+        if count == (tmp_max_steps - 1):
+            raise CTMRGNotConvergedError
+    else:
+        overlap_unitcell, _ = calc_ctmrg_env(
+            [i.tensor for i in overlap_tensors], overlap_unitcell
+        )
+
+        overlap_AB = overlap_func(overlap_unitcell)
 
     return overlap_AB