Add feature to supply additional input to the optimizer routine

Author: JanLuca

peps_ad/optimization/inner_function.py

@@ -8,7 +8,7 @@
 from peps_ad.ctmrg import calc_ctmrg_env, calc_ctmrg_env_custom_rule
 from peps_ad.mapping import Map_To_PEPS_Model
 
-from typing import Sequence, Tuple, cast, Optional, Callable
+from typing import Sequence, Tuple, cast, Optional, Callable, Dict
 
 
 def _map_tensors(
@@ -65,6 +65,7 @@ def calc_ctmrg_expectation(
     unitcell: PEPS_Unit_Cell,
     expectation_func: Expectation_Model,
     convert_to_unitcell_func: Optional[Map_To_PEPS_Model],
+    additional_input: Dict[str, jnp.ndarray] = dict(),
     *,
     enforce_elementwise_convergence: Optional[bool] = None,
 ) -> Tuple[jnp.ndarray, PEPS_Unit_Cell]:
@@ -82,6 +83,9 @@ def calc_ctmrg_expectation(
       convert_to_unitcell_func (:obj:`~peps_ad.mapping.Map_To_PEPS_Model`):
         Function to convert the `input_tensors` to a PEPS unitcell. If ommited,
         it is assumed that a PEPS unitcell is the input.
+      additional_input (:obj:`dict` of :obj:`str` to :obj:`jax.numpy.ndarray` mapping):
+        Optional dict with additional inputs which should be considered in the
+        calculation of the expectation value.
     Keyword args:
       enforce_elementwise_convergence (obj:`bool`):
         Enforce elementwise convergence of the CTM tensors instead of only
@@ -90,7 +94,8 @@ def calc_ctmrg_expectation(
       :obj:`tuple`\ (:obj:`jax.numpy.ndarray`, :obj:`~peps_ad.peps.PEPS_Unit_Cell`):
         Tuple consisting of the calculated expectation value and the new unitcell.
     """
-    if expectation_func.is_spiral_peps:
+    spiral_vectors = additional_input.get("spiral_vectors")
+    if expectation_func.is_spiral_peps and spiral_vectors is None:
         peps_tensors, unitcell, spiral_vectors = _map_tensors(
             input_tensors, unitcell, convert_to_unitcell_func, True
         )
@@ -128,6 +133,7 @@ def calc_preconverged_ctmrg_value_and_grad(
     unitcell: PEPS_Unit_Cell,
     expectation_func: Expectation_Model,
     convert_to_unitcell_func: Optional[Map_To_PEPS_Model],
+    additional_input: Dict[str, jnp.ndarray] = dict(),
     *,
     calc_preconverged: bool = True,
 ) -> Tuple[Tuple[jnp.ndarray, PEPS_Unit_Cell], Sequence[jnp.ndarray]]:
@@ -150,6 +156,9 @@ def calc_preconverged_ctmrg_value_and_grad(
       convert_to_unitcell_func (:obj:`~peps_ad.mapping.Map_To_PEPS_Model`):
         Function to convert the `input_tensors` to a PEPS unitcell. If ommited,
         it is assumed that a PEPS unitcell is the input.
+      additional_input (:obj:`dict` of :obj:`str` to :obj:`jax.numpy.ndarray` mapping):
+        Optional dict with additional inputs which should be considered in the
+        calculation of the expectation value.
     Keyword args:
       calc_preconverged (:obj:`bool`):
         Flag if the above described procedure to calculate a pre-converged
@@ -161,7 +170,8 @@ def calc_preconverged_ctmrg_value_and_grad(
         unitcell.
         2. The calculated gradient.
     """
-    if expectation_func.is_spiral_peps:
+    spiral_vectors = additional_input.get("spiral_vectors")
+    if expectation_func.is_spiral_peps and spiral_vectors is None:
         peps_tensors, unitcell, spiral_vectors = _map_tensors(
             input_tensors, unitcell, convert_to_unitcell_func, True
         )
@@ -196,6 +206,7 @@ def calc_ctmrg_expectation_custom(
     unitcell: PEPS_Unit_Cell,
     expectation_func: Expectation_Model,
     convert_to_unitcell_func: Optional[Map_To_PEPS_Model],
+    additional_input: Dict[str, jnp.ndarray] = dict(),
 ) -> Tuple[jnp.ndarray, PEPS_Unit_Cell]:
     """
     Calculate the CTMRG environment and the (energy) expectation value for a
@@ -211,11 +222,15 @@ def calc_ctmrg_expectation_custom(
       convert_to_unitcell_func (:obj:`~peps_ad.mapping.Map_To_PEPS_Model`):
         Function to convert the `input_tensors` to a PEPS unitcell. If ommited,
         it is assumed that a PEPS unitcell is the input.
+      additional_input (:obj:`dict` of :obj:`str` to :obj:`jax.numpy.ndarray` mapping):
+        Dict with additional inputs which should be considered in the
+        calculation of the expectation value.
     Returns:
       :obj:`tuple`\ (:obj:`jax.numpy.ndarray`, :obj:`~peps_ad.peps.PEPS_Unit_Cell`):
         Tuple consisting of the calculated expectation value and the new unitcell.
     """
-    if expectation_func.is_spiral_peps:
+    spiral_vectors = additional_input.get("spiral_vectors")
+    if expectation_func.is_spiral_peps and spiral_vectors is None:
         peps_tensors, unitcell, spiral_vectors = _map_tensors(
             input_tensors, unitcell, convert_to_unitcell_func, True
         )

peps_ad/optimization/line_search.py

@@ -18,7 +18,7 @@
     calc_ctmrg_expectation_custom_value_and_grad,
 )
 
-from typing import Sequence, Tuple, List, Union, Optional
+from typing import Sequence, Tuple, List, Union, Optional, Dict
 
 
 @jit
@@ -146,6 +146,7 @@ def line_search(
     convert_to_unitcell_func: Optional[Map_To_PEPS_Model] = None,
     generate_unitcell: bool = False,
     spiral_indices: Optional[Sequence[int]] = None,
+    additional_input: Dict[str, jnp.ndarray] = {},
 ) -> Tuple[
     List[jnp.ndarray],
     PEPS_Unit_Cell,
@@ -260,6 +261,7 @@ def line_search(
                     new_unitcell,
                     expectation_func,
                     convert_to_unitcell_func,
+                    additional_input,
                     enforce_elementwise_convergence=enforce_elementwise_convergence,
                 )
 
@@ -296,6 +298,7 @@ def line_search(
                                 unitcell,
                                 expectation_func,
                                 convert_to_unitcell_func,
+                                additional_input,
                             )
                         else:
                             (
@@ -306,6 +309,7 @@ def line_search(
                                 unitcell,
                                 expectation_func,
                                 convert_to_unitcell_func,
+                                additional_input,
                                 calc_preconverged=True,
                             )
                         gradient = [elem.conj() for elem in tmp_gradient_seq]
@@ -334,6 +338,7 @@ def line_search(
                         new_unitcell,
                         expectation_func,
                         convert_to_unitcell_func,
+                        additional_input,
                     )
                 else:
                     (
@@ -344,6 +349,7 @@ def line_search(
                         new_unitcell,
                         expectation_func,
                         convert_to_unitcell_func,
+                        additional_input,
                         calc_preconverged=True,
                     )
                 new_gradient = [elem.conj() for elem in new_gradient_seq]
@@ -381,6 +387,7 @@ def line_search(
                                 unitcell,
                                 expectation_func,
                                 convert_to_unitcell_func,
+                                additional_input,
                             )
                         else:
                             (
@@ -391,6 +398,7 @@ def line_search(
                                 unitcell,
                                 expectation_func,
                                 convert_to_unitcell_func,
+                                additional_input,
                                 calc_preconverged=True,
                             )
                         gradient = [elem.conj() for elem in tmp_gradient_seq]

peps_ad/optimization/optimizer.py

@@ -205,6 +205,7 @@ def optimize_peps_network(
     autosave_func: Callable[
         [PathLike, Sequence[jnp.ndarray], PEPS_Unit_Cell], None
     ] = autosave_function,
+    additional_input: Dict[str, jnp.ndarray] = {},
 ) -> Tuple[Sequence[jnp.ndarray], PEPS_Unit_Cell, Union[float, jnp.ndarray]]:
     """
     Optimize a PEPS unitcell using a variational method.
@@ -279,7 +280,11 @@ def random_noise(a):
     signal_reset_descent_dir = False
 
     spiral_indices = None
-    if hasattr(expectation_func, "is_spiral_peps") and expectation_func.is_spiral_peps:
+    if (
+        hasattr(expectation_func, "is_spiral_peps")
+        and expectation_func.is_spiral_peps
+        and additional_input.get("spiral_vectors") is None
+    ):
         if isinstance(input_tensors, collections.abc.Sequence) and isinstance(
             input_tensors[0], PEPS_Unit_Cell
         ):
@@ -321,6 +326,7 @@ def random_noise(a):
                         working_unitcell,
                         expectation_func,
                         convert_to_unitcell_func,
+                        additional_input,
                     )
                 else:
                     (
@@ -331,6 +337,7 @@ def random_noise(a):
                         working_unitcell,
                         expectation_func,
                         convert_to_unitcell_func,
+                        additional_input,
                         calc_preconverged=(count == 0),
                     )
             except (CTMRGNotConvergedError, CTMRGGradientNotConvergedError) as e:
@@ -427,6 +434,7 @@ def random_noise(a):
                     convert_to_unitcell_func,
                     generate_unitcell,
                     spiral_indices,
+                    additional_input,
                 )
             except NoSuitableStepSizeError:
                 if peps_ad_config.optimizer_fail_if_no_step_size_found:
@@ -521,6 +529,7 @@ def random_noise(a):
                     working_unitcell,
                     expectation_func,
                     convert_to_unitcell_func,
+                    additional_input,
                     enforce_elementwise_convergence=peps_ad_config.ad_use_custom_vjp,
                 )
                 descent_dir = None