Add feature to increase iPEPS bond dim

Author: JanLuca

varipeps/contractions/definitions.py

@@ -691,6 +691,66 @@ def _prepare_defs(cls):
         ],
     }
 
+    unitcell_bond_dim_change_left: Definition = {
+        "tensors": [["tensor", "tensor_conj", "C1", "T1", "T3", "C4", "T4"]],
+        "network": [
+            [
+                (5, 9, 11, -4, 4),  # tensor
+                (7, 10, 11, -3, 6),  # tensor_conj
+                (1, 3),  # C1
+                (3, 4, 6, -1),  # T1
+                (2, -2, 10, 9),  # T3
+                (2, 8),  # C4
+                (8, 7, 5, 1),  # T4
+            ]
+        ],
+    }
+
+    unitcell_bond_dim_change_right: Definition = {
+        "tensors": [["tensor", "tensor_conj", "T1", "C2", "T2", "T3", "C3"]],
+        "network": [
+            [
+                (-1, 8, 11, 5, 4),  # tensor
+                (-4, 9, 11, 7, 6),  # tensor_conj
+                (-2, 4, 6, 3),  # T1
+                (3, 1),  # C2
+                (5, 7, 10, 1),  # T2
+                (-3, 2, 9, 8),  # T3
+                (2, 10),  # C3
+            ]
+        ],
+    }
+
+    unitcell_bond_dim_change_top: Definition = {
+        "tensors": [["tensor", "tensor_conj", "C1", "T1", "C2", "T2", "T4"]],
+        "network": [
+            [
+                (8, -4, 11, 4, 5),  # tensor
+                (9, -3, 11, 6, 7),  # tensor_conj
+                (2, 10),  # C1
+                (10, 5, 7, 1),  # T1
+                (1, 3),  # C2
+                (4, 6, -2, 3),  # T2
+                (-1, 9, 8, 2),  # T4
+            ]
+        ],
+    }
+
+    unitcell_bond_dim_change_bottom: Definition = {
+        "tensors": [["tensor", "tensor_conj", "T2", "C3", "T3", "C4", "T4"]],
+        "network": [
+            [
+                (4, 5, 11, 8, -1),  # tensor
+                (6, 7, 11, 9, -4),  # tensor_conj
+                (8, 9, 2, -3),  # T2
+                (10, 2),  # C3
+                (1, 10, 7, 5),  # T3
+                (1, 3),  # C4
+                (3, 6, 4, -2),  # T4
+            ]
+        ],
+    }
+
     kagome_pess3_single_site_mapping: Definition = {
         "tensors": ["up_simplex", "down_simplex", "site1", "site2", "site3"],
         "network": [

varipeps/ctmrg/absorption.py

@@ -10,12 +10,12 @@
 from varipeps.contractions import apply_contraction, apply_contraction_jitted
 from varipeps.utils.svd import gauge_fixed_svd
 from varipeps.utils.periodic_indices import calculate_periodic_indices
+from varipeps.utils.projector_dict import Projector_Dict
 from .projectors import (
     calc_left_projectors,
     calc_right_projectors,
     calc_top_projectors,
     calc_bottom_projectors,
-    T_Projector,
 )
 from varipeps.expectation.one_site import calc_one_site_single_gate_obj
 from varipeps.config import PEPS_AD_Config
@@ -27,40 +27,6 @@
 CTMRG_Orientation = Literal["top-left", "top-right", "bottom-left", "bottom-right"]
 
 
-@dataclass
-class _Projector_Dict(collections.abc.MutableMapping):
-    max_x: int
-    max_y: int
-    projector_dict: Dict[Tuple[int, int], T_Projector] = field(default_factory=dict)  # type: ignore
-
-    def __getitem__(self, key: Tuple[int, int]) -> T_Projector:
-        return self.projector_dict[key]
-
-    def __setitem__(self, key: Tuple[int, int], value: T_Projector) -> None:
-        self.projector_dict[key] = value
-
-    def __delitem__(self, key: Tuple[int, int]) -> None:
-        self.projector_dict.__delitem__(key)
-
-    def __iter__(self):
-        return self.projector_dict.__iter__()
-
-    def __len__(self):
-        return self.projector_dict.__len__()
-
-    def get_projector(
-        self,
-        current_x: int,
-        current_y: int,
-        relative_x: int,
-        relative_y: int,
-    ) -> T_Projector:
-        select_x = (current_x + relative_x) % self.max_x
-        select_y = (current_y + relative_y) % self.max_y
-
-        return self.projector_dict[(select_x, select_y)]
-
-
 def _tensor_list_from_indices(
     peps_tensors: Sequence[jnp.ndarray], indices: Sequence[Sequence[int]]
 ) -> List[List[jnp.ndarray]]:
@@ -150,7 +116,7 @@ def do_left_absorption(
         all elements of the unitcell.
     """
     max_x, max_y = unitcell.get_size()
-    left_projectors = _Projector_Dict(max_x=max_x, max_y=max_y)
+    left_projectors = Projector_Dict(max_x=max_x, max_y=max_y)
 
     working_unitcell = unitcell.copy()
 
@@ -242,7 +208,7 @@ def do_right_absorption(
         all elements of the unitcell.
     """
     max_x, max_y = unitcell.get_size()
-    right_projectors = _Projector_Dict(max_x=max_x, max_y=max_y)
+    right_projectors = Projector_Dict(max_x=max_x, max_y=max_y)
 
     working_unitcell = unitcell.copy()
 
@@ -338,7 +304,7 @@ def do_top_absorption(
         all elements of the unitcell.
     """
     max_x, max_y = unitcell.get_size()
-    top_projectors = _Projector_Dict(max_x=max_x, max_y=max_y)
+    top_projectors = Projector_Dict(max_x=max_x, max_y=max_y)
 
     working_unitcell = unitcell.copy()
 
@@ -430,7 +396,7 @@ def do_bottom_absorption(
         all elements of the unitcell.
     """
     max_x, max_y = unitcell.get_size()
-    bottom_projectors = _Projector_Dict(max_x=max_x, max_y=max_y)
+    bottom_projectors = Projector_Dict(max_x=max_x, max_y=max_y)
 
     working_unitcell = unitcell.copy()
 

varipeps/ctmrg/projectors.py

@@ -1,4 +1,3 @@
-from collections import namedtuple
 import enum
 from functools import partial
 
@@ -10,22 +9,18 @@
 from varipeps import varipeps_config
 from varipeps.utils.func_cache import Checkpointing_Cache
 from varipeps.utils.svd import gauge_fixed_svd
+from varipeps.utils.projector_dict import (
+    Left_Projectors,
+    Right_Projectors,
+    Top_Projectors,
+    Bottom_Projectors,
+)
 from varipeps.config import Projector_Method, PEPS_AD_Config
 from varipeps.global_state import PEPS_AD_Global_State
 
 from typing import Sequence, Tuple, TypeVar
 
 
-Left_Projectors = namedtuple("Left_Projectors", ("top", "bottom"))
-Right_Projectors = namedtuple("Right_Projectors", ("top", "bottom"))
-Top_Projectors = namedtuple("Top_Projectors", ("left", "right"))
-Bottom_Projectors = namedtuple("Bottom_Projectors", ("left", "right"))
-
-T_Projector = TypeVar(
-    "T_Projector", Left_Projectors, Right_Projectors, Top_Projectors, Bottom_Projectors
-)
-
-
 class _Projectors_Func_Cache:
     _left = None
     _right = None

varipeps/peps/tensor.py

@@ -327,6 +327,7 @@ def replace_tensor(
         new_tensor: Tensor,
         *,
         reinitialize_env_as_identities: bool = True,
+        new_D: Optional[Tuple[int, int, int, int]] = None,
     ) -> T_PEPS_Tensor:
         """
         Replace the PEPS tensor and returns new object of the class.
@@ -337,31 +338,38 @@ def replace_tensor(
         Keyword args:
           reinitialize_env_as_identities (:obj:`bool`):
             Reinitialize the CTM tensors as identities.
+          new_D (:obj:`tuple` of four :obj:`int`, optional):
+            Tuple of new iPEPS bond dimensions if tensor has changed dimensions
         Returns:
           :obj:`~varipeps.peps.PEPS_Tensor`:
             New instance of the class with the tensor replaced.
         """
+        if new_D is None:
+            new_D = self.D
+        elif not isinstance(new_D, tuple) and len(new_D) != 4:
+            raise ValueError("Invalid argument for parameter new_D")
+
         if reinitialize_env_as_identities:
             return type(self)(
                 tensor=new_tensor,
                 C1=jnp.ones((1, 1), dtype=self.C1.dtype),
                 C2=jnp.ones((1, 1), dtype=self.C2.dtype),
                 C3=jnp.ones((1, 1), dtype=self.C3.dtype),
                 C4=jnp.ones((1, 1), dtype=self.C4.dtype),
-                T1=jnp.eye(self.D[3], dtype=self.T1.dtype).reshape(
-                    1, self.D[3], self.D[3], 1
+                T1=jnp.eye(new_D[3], dtype=self.T1.dtype).reshape(
+                    1, new_D[3], new_D[3], 1
                 ),
-                T2=jnp.eye(self.D[2], dtype=self.T2.dtype).reshape(
-                    self.D[2], self.D[2], 1, 1
+                T2=jnp.eye(new_D[2], dtype=self.T2.dtype).reshape(
+                    new_D[2], new_D[2], 1, 1
                 ),
-                T3=jnp.eye(self.D[1], dtype=self.T3.dtype).reshape(
-                    1, 1, self.D[1], self.D[1]
+                T3=jnp.eye(new_D[1], dtype=self.T3.dtype).reshape(
+                    1, 1, new_D[1], new_D[1]
                 ),
-                T4=jnp.eye(self.D[0], dtype=self.T4.dtype).reshape(
-                    1, self.D[0], self.D[0], 1
+                T4=jnp.eye(new_D[0], dtype=self.T4.dtype).reshape(
+                    1, new_D[0], new_D[0], 1
                 ),
                 d=self.d,
-                D=self.D,
+                D=new_D,
                 chi=self.chi,
                 max_chi=self.max_chi,
             )
@@ -377,7 +385,7 @@ def replace_tensor(
                 T3=self.T3,
                 T4=self.T4,
                 d=self.d,
-                D=self.D,
+                D=new_D,
                 chi=self.chi,
                 max_chi=self.max_chi,
             )