Implement triangular CTMRG for maple leaf lattice

Author: JanLuca

varipeps/expectation/spiral_helpers.py

@@ -56,6 +56,9 @@ def apply_unitary(
     if isinstance(delta_r, jnp.ndarray):
         delta_r = (delta_r,) * len(apply_to_index)
 
+    if isinstance(q, jnp.ndarray):
+        q = (q,) * len(apply_to_index)
+
     if len(q) != len(apply_to_index) or len(q) != len(delta_r):
         raise ValueError("Length mismatch!")
 

varipeps/expectation/triangular_one_site.py

@@ -39,6 +39,10 @@ def calc_triangular_one_site(
     peps_tensor_objs,
     gates,
 ):
+    if isinstance(peps_tensors, jnp.ndarray):
+        peps_tensors = (peps_tensors,)
+        peps_tensor_objs = (peps_tensor_objs,)
+
     real_result = all(jnp.allclose(g, g.T.conj()) for g in gates)
 
     return _one_site_workhorse(

varipeps/expectation/triangular_two_sites.py

@@ -16,31 +16,16 @@
 from typing import Sequence, List, Tuple, Union, Optional
 
 
-@partial(jit, static_argnums=(3,))
-def calc_triangular_two_sites_horizontal(
-    peps_tensors,
-    peps_tensor_objs,
+def calc_triangular_two_sites_workhorse(
+    density_matrix_left,
+    density_matrix_right,
     gates,
     real_result=False,
 ):
-    density_matrix_left = apply_contraction_jitted(
-        "triangular_ctmrg_two_site_expectation_horizontal_left",
-        [peps_tensors[0]],
-        [peps_tensor_objs[0]],
-        [],
-    )
-
     density_matrix_left = density_matrix_left.reshape(
         density_matrix_left.shape[0], density_matrix_left.shape[1], -1
     )
 
-    density_matrix_right = apply_contraction_jitted(
-        "triangular_ctmrg_two_site_expectation_horizontal_right",
-        [peps_tensors[1]],
-        [peps_tensor_objs[1]],
-        [],
-    )
-
     density_matrix_right = density_matrix_right.reshape(
         -1, density_matrix_right.shape[-2], density_matrix_right.shape[-1]
     )
@@ -69,55 +54,55 @@ def calc_triangular_two_sites_horizontal(
 
 
 @partial(jit, static_argnums=(3,))
-def calc_triangular_two_sites_vertical(
+def calc_triangular_two_sites_horizontal(
     peps_tensors,
     peps_tensor_objs,
     gates,
     real_result=False,
 ):
-    density_matrix_top = apply_contraction_jitted(
-        "triangular_ctmrg_two_site_expectation_vertical_top",
+    density_matrix_left = apply_contraction_jitted(
+        "triangular_ctmrg_two_site_expectation_horizontal_left",
         [peps_tensors[0]],
         [peps_tensor_objs[0]],
         [],
     )
 
-    density_matrix_top = density_matrix_top.reshape(
-        density_matrix_top.shape[0], density_matrix_top.shape[1], -1
-    )
-
-    density_matrix_bottom = apply_contraction_jitted(
-        "triangular_ctmrg_two_site_expectation_vertical_bottom",
+    density_matrix_right = apply_contraction_jitted(
+        "triangular_ctmrg_two_site_expectation_horizontal_right",
         [peps_tensors[1]],
         [peps_tensor_objs[1]],
         [],
     )
 
-    density_matrix_bottom = density_matrix_bottom.reshape(
-        -1, density_matrix_bottom.shape[-2], density_matrix_bottom.shape[-1]
+    return calc_triangular_two_sites_workhorse(
+        density_matrix_left, density_matrix_right, gates, real_result
     )
 
-    density_matrix = jnp.tensordot(
-        density_matrix_top, density_matrix_bottom, ((2,), (0,))
-    )
 
-    density_matrix = density_matrix.transpose(0, 2, 1, 3)
-    density_matrix = density_matrix.reshape(
-        density_matrix.shape[0] * density_matrix.shape[1],
-        density_matrix.shape[2] * density_matrix.shape[3],
+@partial(jit, static_argnums=(3,))
+def calc_triangular_two_sites_vertical(
+    peps_tensors,
+    peps_tensor_objs,
+    gates,
+    real_result=False,
+):
+    density_matrix_top = apply_contraction_jitted(
+        "triangular_ctmrg_two_site_expectation_vertical_top",
+        [peps_tensors[0]],
+        [peps_tensor_objs[0]],
+        [],
     )
 
-    norm = jnp.trace(density_matrix)
+    density_matrix_bottom = apply_contraction_jitted(
+        "triangular_ctmrg_two_site_expectation_vertical_bottom",
+        [peps_tensors[1]],
+        [peps_tensor_objs[1]],
+        [],
+    )
 
-    if real_result:
-        return [
-            jnp.real(jnp.tensordot(density_matrix, g, ((0, 1), (0, 1))) / norm)
-            for g in gates
-        ]
-    else:
-        return [
-            jnp.tensordot(density_matrix, g, ((0, 1), (0, 1))) / norm for g in gates
-        ]
+    return calc_triangular_two_sites_workhorse(
+        density_matrix_top, density_matrix_bottom, gates, real_result
+    )
 
 
 @partial(jit, static_argnums=(3,))
@@ -134,43 +119,17 @@ def calc_triangular_two_sites_diagonal(
         [],
     )
 
-    density_matrix_top = density_matrix_top.reshape(
-        density_matrix_top.shape[0], density_matrix_top.shape[1], -1
-    )
-
     density_matrix_bottom = apply_contraction_jitted(
         "triangular_ctmrg_two_site_expectation_diagonal_bottom",
         [peps_tensors[1]],
         [peps_tensor_objs[1]],
         [],
     )
 
-    density_matrix_bottom = density_matrix_bottom.reshape(
-        -1, density_matrix_bottom.shape[-2], density_matrix_bottom.shape[-1]
-    )
-
-    density_matrix = jnp.tensordot(
-        density_matrix_top, density_matrix_bottom, ((2,), (0,))
+    return calc_triangular_two_sites_workhorse(
+        density_matrix_top, density_matrix_bottom, gates, real_result
     )
 
-    density_matrix = density_matrix.transpose(0, 2, 1, 3)
-    density_matrix = density_matrix.reshape(
-        density_matrix.shape[0] * density_matrix.shape[1],
-        density_matrix.shape[2] * density_matrix.shape[3],
-    )
-
-    norm = jnp.trace(density_matrix)
-
-    if real_result:
-        return [
-            jnp.real(jnp.tensordot(density_matrix, g, ((0, 1), (0, 1))) / norm)
-            for g in gates
-        ]
-    else:
-        return [
-            jnp.tensordot(density_matrix, g, ((0, 1), (0, 1))) / norm for g in gates
-        ]
-
 
 @dataclass
 class Triangular_Two_Sites_Expectation_Value(Expectation_Model):