Allow selection of PEPS type in unit cell random() function

Author: JanLuca

varipeps/contractions/apply.py

@@ -14,9 +14,7 @@
 from typing import Sequence, List, Tuple, Dict, Union, Optional
 
 
-@partial(
-    jax.jit, static_argnames=("name", "disable_identity_check")
-)
+@partial(jax.jit, static_argnames=("name", "disable_identity_check"))
 def apply_contraction(
     name: str,
     peps_tensors: Sequence[jnp.ndarray],

varipeps/peps/__init__.py

@@ -1,6 +1,7 @@
 from . import tensor
 from . import unitcell
 from .tensor import (
+    PEPS_Type,
     PEPS_Tensor,
     PEPS_Tensor_Structure_Factor,
     PEPS_Tensor_Split_Transfer,

varipeps/peps/tensor.py

@@ -6,6 +6,7 @@
 
 import collections
 from dataclasses import dataclass
+from enum import Enum, IntEnum, auto, unique
 
 import numpy as np
 import jax
@@ -26,6 +27,15 @@
 )
 
 
+@unique
+class PEPS_Type(IntEnum):
+    SQUARE = auto()  #: Square-lattice based iPEPS state with full transfer tensors
+    SQUARE_SPLIT = (
+        auto()
+    )  #: Square-lattice based iPEPS state with split transfer tensors
+    TRIANGULAR = auto()  #: Triangular-lattice based iPEPS state
+
+
 @dataclass
 @register_pytree_node_class
 class PEPS_Tensor:
@@ -1172,6 +1182,10 @@ def load_from_group(cls: Type[T_PEPS_Tensor], grp: h5py.Group) -> T_PEPS_Tensor:
     def is_split_transfer(self: T_PEPS_Tensor) -> bool:
         return False
 
+    @property
+    def peps_type(self) -> PEPS_Type:
+        return PEPS_Type.SQUARE
+
     def convert_to_split_transfer(
         self: T_PEPS_Tensor, interlayer_chi: Optional[int] = None
     ) -> T_PEPS_Tensor_Split_Transfer:
@@ -2916,6 +2930,10 @@ def load_from_group(
     def is_split_transfer(self: T_PEPS_Tensor_Split_Transfer) -> bool:
         return True
 
+    @property
+    def peps_type(self) -> PEPS_Type:
+        return PEPS_Type.SQUARE_SPLIT
+
     def convert_to_split_transfer(
         self: T_PEPS_Tensor_Split_Transfer,
     ) -> T_PEPS_Tensor_Split_Transfer:
@@ -4043,6 +4061,10 @@ def is_split_transfer(self) -> bool:
     def is_triangular_peps(self) -> bool:
         return True
 
+    @property
+    def peps_type(self) -> PEPS_Type:
+        return PEPS_Type.TRIANGULAR
+
     def tree_flatten(self) -> Tuple[Tuple[Any, ...], Tuple[Any, ...]]:
         data = (
             self.tensor,

varipeps/peps/unitcell.py

@@ -17,7 +17,12 @@
 import jax.numpy as jnp
 from jax.tree_util import register_pytree_node_class
 
-from .tensor import PEPS_Tensor, PEPS_Tensor_Split_Transfer, PEPS_Tensor_Triangular
+from .tensor import (
+    PEPS_Tensor,
+    PEPS_Tensor_Split_Transfer,
+    PEPS_Tensor_Triangular,
+    PEPS_Type,
+)
 import varipeps
 from varipeps.utils.random import PEPS_Random_Number_Generator
 from varipeps.utils.periodic_indices import calculate_periodic_indices
@@ -273,6 +278,7 @@ def random(
         chi: Union[int, Sequence[int]],
         dtype: Type[jnp.number],
         max_chi: Optional[int] = None,
+        peps_type: PEPS_Type = PEPS_Type.SQUARE,
         *,
         seed: Optional[int] = None,
         destroy_random_state: bool = True,
@@ -320,7 +326,10 @@ def random(
             )
 
         if isinstance(D, int):
-            D = [(D, D, D, D) for _ in range(tensors_i.size)]
+            if peps_type is PEPS_Type.TRIANGULAR:
+                D = [(D, D, D, D, D, D) for _ in range(tensors_i.size)]
+            else:
+                D = [(D, D, D, D) for _ in range(tensors_i.size)]
 
         if (
             not all(isinstance(j, int) for i in D for j in i)
@@ -344,19 +353,28 @@ def random(
 
         peps_tensors = []
 
+        if peps_type is PEPS_Type.TRIANGULAR:
+            peps_tensor_class = PEPS_Tensor_Triangular
+        else:
+            peps_tensor_class = PEPS_Tensor
+
         for i in tensors_i:
             if i > 0:
                 seed = None
 
             peps_tensors.append(
-                PEPS_Tensor.random(
+                peps_tensor_class.random(
                     d=d[i], D=D[i], chi=chi[i], dtype=dtype, seed=seed, max_chi=max_chi
                 )
             )
 
         data = cls.Unit_Cell_Data(peps_tensors=peps_tensors, structure=structure)
 
-        return cls(data=data)
+        result = cls(data=data)
+
+        if peps_type is PEPS_Type.SQUARE_SPLIT:
+            return result.convert_to_split_transfer()
+        return result
 
     def get_size(self) -> Tuple[int, int]:
         """