Support parameters in other PEPS_Tensor classes as well

Author: JanLuca

varipeps/peps/tensor.py

@@ -1390,6 +1390,7 @@ def change_chi(
         new_chi: int,
         *,
         reinitialize_env_as_identities: bool = True,
+        reset_max_chi: bool = False,
     ) -> T_PEPS_Tensor:
         """
         Change the environment bond dimension and returns new object of the class.
@@ -1400,6 +1401,8 @@ def change_chi(
         Keyword args:
           reinitialize_env_as_identities (:obj:`bool`):
             Reinitialize the CTM tensors as identities if decreasing the dimension.
+          reset_max_chi (:obj:`bool`):
+            Set maximal bond dimension to the same new value.
         Returns:
           :obj:`~varipeps.peps.PEPS_Tensor`:
             New instance of the class with the increased value.
@@ -1409,6 +1412,8 @@ def change_chi(
                 "Increase above the max value for environment bond dimension."
             )
 
+        new_max_chi = new_chi if reset_max_chi else self.max_chi
+
         if new_chi < self.chi and reinitialize_env_as_identities:
             return type(self)(
                 tensor=self.tensor,
@@ -1447,7 +1452,7 @@ def change_chi(
                 d=self.d,
                 D=self.D,
                 chi=new_chi,
-                max_chi=self.max_chi,
+                max_chi=new_max_chi,
             )
         else:
             return type(self)(
@@ -1471,7 +1476,7 @@ def change_chi(
                 d=self.d,
                 D=self.D,
                 chi=new_chi,
-                max_chi=self.max_chi,
+                max_chi=new_max_chi,
             )
 
     def increase_max_chi(
@@ -2167,6 +2172,8 @@ def change_chi(
         new_chi: int,
         *,
         reinitialize_env_as_identities: bool = True,
+        reset_max_chi: bool = False,
+        reset_interlayer_chi: bool = True,
     ) -> T_PEPS_Tensor_Split_Transfer:
         """
         Change the environment bond dimension and returns new object of the class.
@@ -2177,6 +2184,10 @@ def change_chi(
         Keyword args:
           reinitialize_env_as_identities (:obj:`bool`):
             Reinitialize the CTM tensors as identities if decreasing the dimension.
+          reset_max_chi (:obj:`bool`):
+            Set maximal bond dimension to the same new value.
+          reset_interlayer_chi (:obj:`bool`):
+            Set interlayer bond dimension to the same new value.
         Returns:
           :obj:`~peps_ad.peps.PEPS_Tensor_Split_Transfer`:
             New instance of the class with the increased value.
@@ -2186,6 +2197,10 @@ def change_chi(
                 "Increase above the max value for environment bond dimension."
             )
 
+        new_max_chi = new_chi if reset_max_chi else self.max_chi
+
+        new_interlayer_chi = new_chi if reset_interlayer_chi else self.interlayer_chi
+
         if new_chi < self.chi and reinitialize_env_as_identities:
             return type(self)(
                 tensor=self.tensor,
@@ -2208,8 +2223,8 @@ def change_chi(
                 d=self.d,
                 D=self.D,
                 chi=new_chi,
-                max_chi=self.max_chi,
-                interlayer_chi=self.interlayer_chi,
+                max_chi=new_max_chi,
+                interlayer_chi=new_interlayer_chi,
             )
         else:
             return type(self)(
@@ -2229,8 +2244,8 @@ def change_chi(
                 d=self.d,
                 D=self.D,
                 chi=new_chi,
-                max_chi=self.max_chi,
-                interlayer_chi=self.interlayer_chi,
+                max_chi=new_max_chi,
+                interlayer_chi=new_interlayer_chi,
             )
 
     def increase_max_chi(

varipeps/peps/unitcell.py

@@ -496,7 +496,7 @@ def convert_to_split_transfer(
         )
 
     def convert_to_full_transfer(
-        self: T_PEPS_Unit_Cell, interlayer_chi: Optional[int] = None
+        self: T_PEPS_Unit_Cell
     ) -> T_PEPS_Unit_Cell:
         """
         Convert the list of unique tensors to the full transfer ansatz.