Reuse CTMRG env tensors if gradient in optimization is already small

Author: JanLuca

peps_ad/config.py

@@ -137,6 +137,9 @@ class PEPS_AD_Config:
         Maximal retries for optimization with random noise.
       optimizer_random_noise_relative_amplitude (:obj:`float`):
         Relative amplitude used for random noise.
+      optimizer_reuse_env_eps (:obj:`float`):
+        Reuse CTMRG environment of previous step if norm of gradient is below
+        this threshold.
       line_search_method (:obj:`Line_Search_Methods`):
         Method used for the line search routine.
       line_search_initial_step_size (:obj:`float`):
@@ -208,6 +211,7 @@ class PEPS_AD_Config:
     optimizer_random_noise_eps: float = 1e-4
     optimizer_random_noise_max_retries: int = 5
     optimizer_random_noise_relative_amplitude: float = 1e-1
+    optimizer_reuse_env_eps: float = 1e-3
 
     # Line search
     line_search_method: Line_Search_Methods = Line_Search_Methods.WOLFE

peps_ad/optimization/line_search.py

@@ -147,6 +147,7 @@ def line_search(
     generate_unitcell: bool = False,
     spiral_indices: Optional[Sequence[int]] = None,
     additional_input: Dict[str, jnp.ndarray] = {},
+    reinitialize_env_as_identities: bool = True,
 ) -> Tuple[
     List[jnp.ndarray],
     PEPS_Unit_Cell,
@@ -243,7 +244,8 @@ def line_search(
             new_unitcell = unitcell.replace_unique_tensors(
                 [
                     unitcell_tensors[i].replace_tensor(
-                        new_tensors[i], reinitialize_env_as_identities=True
+                        new_tensors[i],
+                        reinitialize_env_as_identities=reinitialize_env_as_identities,
                     )
                     for i in range(unitcell.get_len_unique_tensors())
                 ]

peps_ad/optimization/optimizer.py

@@ -435,6 +435,7 @@ def random_noise(a):
                     generate_unitcell,
                     spiral_indices,
                     additional_input,
+                    conv > peps_ad_config.optimizer_reuse_env_eps,
                 )
             except NoSuitableStepSizeError:
                 if peps_ad_config.optimizer_fail_if_no_step_size_found: