Add config setting for the grad fixed point method

Author: JanLuca

varipeps/config.py

@@ -40,6 +40,17 @@ class Wavevector_Type(IntEnum):
     TWO_PI_SYMMETRIC = auto()  #: Use interval [-2pi, 2pi) for q vectors
 
 
+@unique
+class Grad_Fixed_Point_Method(IntEnum):
+    ITERATIVE = auto()  #: Use iterative method to calculate gradient of CTMRG routine
+    LINEAR_SOLVER = (
+        auto()
+    )  #: Use linear solver method to calculate gradient of CTMRG routine
+    EIGEN_SOLVER = (
+        auto()
+    )  #: Use eigen solver method to calculate gradient of CTMRG routine
+
+
 @unique
 class Slurm_Restart_Mode(IntEnum):
     DISABLED = (
@@ -73,6 +84,9 @@ class VariPEPS_Config:
       ad_custom_max_steps (:obj:`int`):
         Maximal number of steps for fix-pointer iteration of the custom VJP
         function.
+      ad_custom_fixed_point_method (:obj:`~varipeps.config.Grad_Fixed_Point_Method`):
+        Select method how the gradient of the CTMRG fixed point routine is
+        calculated.
       checkpointing_ncon (:obj:`bool`):
         Enable AD checkpointing for the ncon calls.
       checkpointing_projectors (:obj:`bool`):
@@ -242,6 +256,9 @@ class VariPEPS_Config:
     ad_custom_verbose_output: bool = False
     ad_custom_convergence_eps: float = 1e-7
     ad_custom_max_steps: int = 75
+    ad_custom_fixed_point_method: Grad_Fixed_Point_Method = (
+        Grad_Fixed_Point_Method.LINEAR_SOLVER
+    )
     checkpointing_ncon: bool = False
     checkpointing_projectors: bool = False
 
@@ -406,6 +423,7 @@ class ConfigModuleWrapper:
         "Line_Search_Methods",
         "Projector_Method",
         "Wavevector_Type",
+        "Grad_Fixed_Point_Method",
         "Slurm_Restart_Mode",
         "VariPEPS_Config",
         "config",

varipeps/ctmrg/routine.py

@@ -10,6 +10,7 @@
 import jax.debug as jdebug
 
 from varipeps import varipeps_config, varipeps_global_state
+from varipeps.config import Grad_Fixed_Point_Method
 from varipeps.peps import PEPS_Tensor, PEPS_Tensor_Split_Transfer, PEPS_Unit_Cell
 from varipeps.utils.debug_print import debug_print
 from .absorption import do_absorption_step, do_absorption_step_split_transfer
@@ -1033,9 +1034,7 @@ def _ctmrg_rev_workhorse(peps_tensors, new_unitcell, new_unitcell_bar, config, s
             )[1]
         )
 
-    old_method = False
-
-    if old_method:
+    if config.ad_custom_fixed_point_method is Grad_Fixed_Point_Method.ITERATIVE:
 
         def cond_func(carry):
             _, _, _, converged, count, config, state = carry
@@ -1048,18 +1047,21 @@ def cond_func(carry):
             (vjp_env, new_unitcell_bar, new_unitcell_bar, False, 0, config, state),
         )
     else:
-        env_fixed_point, arnoldi_worked = jax.pure_callback(
-            _ctmrg_rev_arnoldi,
-            jax.eval_shape(lambda x: (x, True), new_unitcell_bar),
-            vjp(
-                lambda u: do_absorption_step(peps_tensors, u, config, state),
-                new_unitcell,
-            )[1],
-            new_unitcell_bar,
-        )
-        end_count = 0
+        if config.ad_custom_fixed_point_method is Grad_Fixed_Point_Method.EIGEN_SOLVER:
+            env_fixed_point, arnoldi_worked = jax.pure_callback(
+                _ctmrg_rev_arnoldi,
+                jax.eval_shape(lambda x: (x, True), new_unitcell_bar),
+                vjp(
+                    lambda u: do_absorption_step(peps_tensors, u, config, state),
+                    new_unitcell,
+                )[1],
+                new_unitcell_bar,
+            )
+        else:
+            env_fixed_point = new_unitcell_bar
+            arnoldi_worked = False
 
-        debug_print("Arnoldi: {}", arnoldi_worked)
+        end_count = 0
 
         def run_gmres(v, e):
             def f_gmres(w):