PEPS: support bond_dim=1 cyclic constructors

jcmgray · Oct 22, 2024 · b08a113 · b08a113
1 parent 970c064
commit b08a113
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 2 deletions.
diff --git a/quimb/tensor/tensor_2d.py b/quimb/tensor/tensor_2d.py
@@ -4804,8 +4804,20 @@ def __init__(
         self._Lx = len(arrays)
         self._Ly = len(arrays[0])
 
-        cyclicx = sum(d > 1 for d in ar.shape(arrays[0][1])) == 5
-        cyclicy = sum(d > 1 for d in ar.shape(arrays[1][0])) == 5
+        cyclicx = (
+            sum(d > 1 for d in ar.shape(arrays[0][1])) == 5
+        ) or (
+            # handle D=1 PBC case
+            (ar.ndim(arrays[0][1]) == 5) and
+            (sum(d == 1 for d in ar.shape(arrays[0][1])) == 4)
+        )
+        cyclicy = (
+            sum(d > 1 for d in ar.shape(arrays[1][0])) == 5
+        ) or (
+            # handle D=1 PBC case
+            (ar.ndim(arrays[1][0]) == 5) and
+            (sum(d == 1 for d in ar.shape(arrays[1][0])) == 4)
+        )
 
         # cache for both creating and retrieving indices
         ix = defaultdict(rand_uuid)

diff --git a/tests/test_tensor/test_tensor_2d.py b/tests/test_tensor/test_tensor_2d.py
@@ -47,6 +47,12 @@ def test_basic_rand(self, Lx, Ly):
         assert f"Lx={Lx}" in psi.__str__()
         assert f"Lx={Lx}" in psi.__repr__()
 
+    def test_cyclic_edge_cases(self):
+        peps = qtn.PEPS.rand(3, 3, bond_dim=1, cyclic=True)
+        assert peps.is_cyclic_x()
+        assert peps.is_cyclic_y()
+        assert peps.num_indices == peps.num_tensors * 3
+
     def test_flatten(self):
         psi = qtn.PEPS.rand(3, 5, 3, seed=42)
         norm = psi.H & psi