Add split-transfer awareness to optimizer

JanLuca · JanLuca · commit a2e6ff37db3f · 2025-02-27T15:59:14.000+01:00
diff --git a/varipeps/ctmrg/routine.py b/varipeps/ctmrg/routine.py
@@ -69,8 +69,6 @@ def _calc_corner_svds(
             C1_svd, indices_are_sorted=True, unique_indices=True
         )
 
-        # debug_print("C1: {}", C1_svd)
-
         C2_svd = jnp.linalg.svd(t.C2, full_matrices=False, compute_uv=False)
         step_corner_svd = step_corner_svd.at[ti, 1, : C2_svd.shape[0]].set(
             C2_svd, indices_are_sorted=True, unique_indices=True
@@ -382,7 +380,7 @@ def _ctmrg_body_func(carry):
         config,
     ) = carry
 
-    if state.ctmrg_split_transfer:
+    if w_unitcell_last_step.is_split_transfer():
         w_unitcell, norm_smallest_S = do_absorption_step_split_transfer(
             w_tensors, w_unitcell_last_step, config, state
         )
@@ -397,7 +395,7 @@ def elementwise_func(old, new, old_corner, conv_eps, config):
             new,
             conv_eps,
             verbose=config.ctmrg_verbose_output,
-            split_transfer=state.ctmrg_split_transfer,
+            split_transfer=w_unitcell.is_split_transfer(),
         )
         return converged, measure, verbose_data, old_corner
 
@@ -535,10 +533,6 @@ def calc_ctmrg_env(
     norm_smallest_S = jnp.nan
     already_tried_chi = {working_unitcell[0, 0][0][0].chi}
 
-    varipeps_global_state.ctmrg_split_transfer = isinstance(
-        unitcell.get_unique_tensors()[0], PEPS_Tensor_Split_Transfer
-    )
-
     while True:
         tmp_count = 0
         corner_singular_vals = None
@@ -776,6 +770,7 @@ def _ctmrg_rev_while_body(carry):
         bar_fixed_point.get_unique_tensors(),
         config.ad_custom_convergence_eps,
         verbose=config.ad_custom_verbose_output,
+        split_transfer=bar_fixed_point.is_split_transfer(),
     )
 
     count += 1
@@ -796,15 +791,31 @@ def _ctmrg_rev_while_body(carry):
 
 @jit
 def _ctmrg_rev_workhorse(peps_tensors, new_unitcell, new_unitcell_bar, config, state):
-    _, vjp_peps_tensors = vjp(
-        lambda t: do_absorption_step(t, new_unitcell, config, state), peps_tensors
-    )
+    if new_unitcell.is_split_transfer():
+        _, vjp_peps_tensors = vjp(
+            lambda t: do_absorption_step_split_transfer(t, new_unitcell, config, state),
+            peps_tensors,
+        )
 
-    vjp_env = tree_util.Partial(
-        vjp(lambda u: do_absorption_step(peps_tensors, u, config, state), new_unitcell)[
-            1
-        ]
-    )
+        vjp_env = tree_util.Partial(
+            vjp(
+                lambda u: do_absorption_step_split_transfer(
+                    peps_tensors, u, config, state
+                ),
+                new_unitcell,
+            )[1]
+        )
+    else:
+        _, vjp_peps_tensors = vjp(
+            lambda t: do_absorption_step(t, new_unitcell, config, state), peps_tensors
+        )
+
+        vjp_env = tree_util.Partial(
+            vjp(
+                lambda u: do_absorption_step(peps_tensors, u, config, state),
+                new_unitcell,
+            )[1]
+        )
 
     def cond_func(carry):
         _, _, _, converged, count, config, state = carry
diff --git a/varipeps/global_state.py b/varipeps/global_state.py
@@ -21,7 +21,6 @@ class should not be modified by users.
 
     ctmrg_effective_truncation_eps: Optional[float] = None
     ctmrg_projector_method: Optional[Projector_Method] = None
-    ctmrg_split_transfer: bool = False
     basinhopping_disable_half_projector: Optional[bool] = None
 
     def tree_flatten(self) -> Tuple[Tuple[Any, ...], Tuple[Any, ...]]:
diff --git a/varipeps/optimization/inner_function.py b/varipeps/optimization/inner_function.py
@@ -94,6 +94,8 @@ def calc_ctmrg_expectation(
       :obj:`tuple`\ (:obj:`jax.numpy.ndarray`, :obj:`~varipeps.peps.PEPS_Unit_Cell`):
         Tuple consisting of the calculated expectation value and the new unitcell.
     """
+    state_split_transfer = unitcell.is_split_transfer()
+
     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(
@@ -122,20 +124,25 @@ def calc_ctmrg_expectation(
             input_tensors, unitcell, convert_to_unitcell_func, False
         )
 
+    if state_split_transfer != unitcell.is_split_transfer():
+        raise ValueError("Map function is not split transfer aware. Please fix that!")
+
     new_unitcell, max_trunc_error = calc_ctmrg_env(
         peps_tensors,
         unitcell,
         enforce_elementwise_convergence=enforce_elementwise_convergence,
     )
 
+    exp_unitcell = new_unitcell.convert_to_full_transfer()
+
     if expectation_func.is_spiral_peps:
         return cast(
-            jnp.ndarray, expectation_func(peps_tensors, new_unitcell, spiral_vectors)
+            jnp.ndarray, expectation_func(peps_tensors, exp_unitcell, spiral_vectors)
         ), (
             new_unitcell,
             max_trunc_error,
         )
-    return cast(jnp.ndarray, expectation_func(peps_tensors, new_unitcell)), (
+    return cast(jnp.ndarray, expectation_func(peps_tensors, exp_unitcell)), (
         new_unitcell,
         max_trunc_error,
     )
@@ -188,6 +195,8 @@ def calc_preconverged_ctmrg_value_and_grad(
         unitcell.
         2. The calculated gradient.
     """
+    state_split_transfer = unitcell.is_split_transfer()
+
     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(
@@ -216,6 +225,9 @@ def calc_preconverged_ctmrg_value_and_grad(
             input_tensors, unitcell, convert_to_unitcell_func, False
         )
 
+    if state_split_transfer != unitcell.is_split_transfer():
+        raise ValueError("Map function is not split transfer aware. Please fix that!")
+
     if calc_preconverged:
         preconverged_unitcell, _ = calc_ctmrg_env(
             peps_tensors,
@@ -265,6 +277,8 @@ def calc_ctmrg_expectation_custom(
       :obj:`tuple`\ (:obj:`jax.numpy.ndarray`, :obj:`~varipeps.peps.PEPS_Unit_Cell`):
         Tuple consisting of the calculated expectation value and the new unitcell.
     """
+    state_split_transfer = unitcell.is_split_transfer()
+
     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(
@@ -293,16 +307,21 @@ def calc_ctmrg_expectation_custom(
             input_tensors, unitcell, convert_to_unitcell_func, False
         )
 
+    if state_split_transfer != unitcell.is_split_transfer():
+        raise ValueError("Map function is not split transfer aware. Please fix that!")
+
     new_unitcell, max_trunc_error = calc_ctmrg_env_custom_rule(peps_tensors, unitcell)
 
+    exp_unitcell = new_unitcell.convert_to_full_transfer()
+
     if expectation_func.is_spiral_peps:
         return cast(
-            jnp.ndarray, expectation_func(peps_tensors, new_unitcell, spiral_vectors)
+            jnp.ndarray, expectation_func(peps_tensors, exp_unitcell, spiral_vectors)
         ), (
             new_unitcell,
             max_trunc_error,
         )
-    return cast(jnp.ndarray, expectation_func(peps_tensors, new_unitcell)), (
+    return cast(jnp.ndarray, expectation_func(peps_tensors, exp_unitcell)), (
         new_unitcell,
         max_trunc_error,
     )
diff --git a/varipeps/peps/tensor.py b/varipeps/peps/tensor.py
@@ -1105,7 +1105,16 @@ def load_from_group(cls: Type[T_PEPS_Tensor], grp: h5py.Group) -> T_PEPS_Tensor:
             max_chi=max_chi,
         )
 
-    def convert_to_split_transfer(self: T_PEPS_Tensor) -> T_PEPS_Tensor_Split_Transfer:
+    @property
+    def is_split_transfer(self: T_PEPS_Tensor) -> bool:
+        return False
+
+    def convert_to_split_transfer(
+        self: T_PEPS_Tensor, interlayer_chi: Optional[int] = None
+    ) -> T_PEPS_Tensor_Split_Transfer:
+        if interlayer_chi is None:
+            interlayer_chi = self.chi
+
         return PEPS_Tensor_Split_Transfer(
             tensor=self.tensor,
             C1=self.C1,
@@ -1120,9 +1129,12 @@ def convert_to_split_transfer(self: T_PEPS_Tensor) -> T_PEPS_Tensor_Split_Transf
             D=self.D,
             chi=self.chi,
             max_chi=self.max_chi,
-            interlayer_chi=self.chi,
+            interlayer_chi=interlayer_chi,
         )
 
+    def convert_to_full_transfer(self: T_PEPS_Tensor) -> T_PEPS_Tensor:
+        return self
+
     def tree_flatten(self) -> Tuple[Tuple[Any, ...], Tuple[Any, ...]]:
         data = (
             self.tensor,
@@ -2467,7 +2479,7 @@ def __add__(
                 "Both PEPS tensors must have the same tensor, d, D and chi values."
             )
 
-        return PEPS_Tensor(
+        return type(self)(
             tensor=self.tensor,
             C1=self.C1 + other.C1,
             C2=self.C2 + other.C2,
@@ -2621,6 +2633,10 @@ def load_from_group(
             interlayer_chi=interlayer_chi,
         )
 
+    @property
+    def is_split_transfer(self: T_PEPS_Tensor_Split_Transfer) -> bool:
+        return True
+
     def convert_to_split_transfer(
         self: T_PEPS_Tensor_Split_Transfer,
     ) -> T_PEPS_Tensor_Split_Transfer:
diff --git a/varipeps/peps/unitcell.py b/varipeps/peps/unitcell.py
@@ -450,6 +450,66 @@ def replace_unique_tensors(
             sanity_checks=False,
         )
 
+    def is_split_transfer(self: T_PEPS_Unit_Cell) -> bool:
+        return all(t.is_split_transfer for t in self.data.peps_tensors)
+
+    def convert_to_split_transfer(
+        self: T_PEPS_Unit_Cell, interlayer_chi: Optional[int] = None
+    ) -> T_PEPS_Unit_Cell:
+        """
+        Convert the list of unique tensors to the split transfer ansatz.
+
+        Args:
+          interlayer_chi (:obj:`int`, optional):
+            Bond dimension for the interlayer index in the split transfer
+            ansatz. If set to None, the same value as for the enviroment
+            bond dimension is used.
+        Returns:
+          PEPS_Unit_Cell:
+            New instance of PEPS unit cell with the new unique tensor list.
+        """
+        if self.is_split_transfer():
+            return self
+
+        new_unique_tensors = type(self.data.peps_tensors)(
+            t.convert_to_split_transfer(interlayer_chi) for t in self.data.peps_tensors
+        )
+
+        new_data = self.data.replace_peps_tensors(new_unique_tensors)
+
+        return type(self)(
+            data=new_data,
+            real_ix=self.real_ix,
+            real_iy=self.real_iy,
+            sanity_checks=False,
+        )
+
+    def convert_to_full_transfer(
+        self: T_PEPS_Unit_Cell, interlayer_chi: Optional[int] = None
+    ) -> T_PEPS_Unit_Cell:
+        """
+        Convert the list of unique tensors to the full transfer ansatz.
+
+        Returns:
+          PEPS_Unit_Cell:
+            New instance of PEPS unit cell with the new unique tensor list.
+        """
+        if not self.is_split_transfer():
+            return self
+
+        new_unique_tensors = type(self.data.peps_tensors)(
+            t.convert_to_full_transfer() for t in self.data.peps_tensors
+        )
+
+        new_data = self.data.replace_peps_tensors(new_unique_tensors)
+
+        return type(self)(
+            data=new_data,
+            real_ix=self.real_ix,
+            real_iy=self.real_iy,
+            sanity_checks=False,
+        )
+
     def change_chi(
         self: T_PEPS_Unit_Cell,
         new_chi: int,
