Implement next-nearest interaction for triangular lattice

Author: JanLuca

docs/source/api/expectation/four_sites.rst

@@ -0,0 +1,12 @@
+.. _varipeps_expectation_four_sites:
+
+.. currentmodule:: varipeps.expectation.four_sites
+
+Calculation of four sites expectation values
+============================================
+
+.. automodule:: varipeps.expectation.four_sites
+   :members:
+   :undoc-members:
+   :show-inheritance:
+   :special-members: __call__

docs/source/api/expectation/index.rst

@@ -14,6 +14,7 @@ Calculation of expectation values (:mod:`varipeps.expectation`)
    one_site
    two_sites
    three_sites
+   four_sites
    spiral_helpers
 
 .. automodule:: varipeps.expectation

varipeps/contractions/definitions.py

@@ -503,6 +503,67 @@ def _prepare_defs(cls):
         ],
     }
 
+    density_matrix_two_sites_horizontal_rectangle: Definition = {
+        "tensors": [
+            ["tensor", "tensor_conj", "T1"],  # top_middle
+            ["tensor", "tensor_conj", "T3"],  # bottom_middle
+            "top_left",
+            "top_right",
+            "bottom_left",
+            "bottom_right",
+        ],
+        "network": [
+            [
+                (5, 11, 9, 19, 6),  # tensor
+                (7, 13, 9, 20, 8),  # tensor_conj
+                (4, 6, 8, 18),  # T1
+            ],
+            [
+                (10, 16, 14, 25, 11),  # tensor
+                (12, 17, 14, 26, 13),  # tensor_conj
+                (15, 24, 17, 16),  # T3
+            ],
+            (-1, -3, 1, 2, 3, 4, 5, 7),  # top_left
+            (18, 19, 20, 21, 22, 23),  # top_right
+            (
+                15,
+                12,
+                10,
+                1,
+                2,
+                3,
+            ),  # bottom_left
+            (-2, -4, 21, 22, 23, 24, 26, 25),  # bottom_right
+        ],
+    }
+
+    density_matrix_two_sites_vertical_rectangle: Definition = {
+        "tensors": [
+            ["tensor", "tensor_conj", "T4"],
+            ["tensor", "tensor_conj", "T2"],
+            "top_left",
+            "top_right",
+            "bottom_left",
+            "bottom_right",
+        ],
+        "network": [
+            [
+                (6, 19, 9, 11, 5),  # tensor
+                (8, 20, 9, 13, 7),  # tensor_conj
+                (18, 8, 6, 4),  # T4
+            ],
+            [
+                (11, 22, 14, 16, 10),  # tensor
+                (13, 23, 14, 17, 12),  # tensor_conj
+                (16, 17, 21, 15),  # T2
+            ],
+            (-1, -3, 4, 5, 7, 1, 2, 3),  # top_left
+            (1, 2, 3, 15, 12, 10),  # top_right
+            (24, 25, 26, 18, 19, 20),  # bottom_left
+            (-2, -4, 21, 23, 22, 24, 25, 26),  # bottom_right
+        ],
+    }
+
     density_matrix_four_sites_top_left: Definition = {
         "tensors": [["tensor", "tensor_conj", "T4", "C1", "T1"]],
         "network": [

varipeps/expectation/__init__.py

@@ -4,6 +4,7 @@
 from . import one_site
 from . import two_sites
 from . import three_sites
+from . import four_sites
 
 from .model import Expectation_Model
 from .one_site import One_Site_Expectation_Value

varipeps/expectation/four_sites.py

@@ -0,0 +1,119 @@
+from functools import partial
+
+import jax.numpy as jnp
+from jax import jit
+
+from varipeps.peps import PEPS_Tensor
+from varipeps.contractions import apply_contraction_jitted
+
+from typing import Sequence, List, Tuple, Literal
+
+Corner_Literal = Literal["top-left", "top-right", "bottom-left", "bottom-right"]
+
+
+@partial(jit, static_argnums=(5,))
+def _four_sites_quadrat_workhorse(
+    top_left: jnp.ndarray,
+    top_right: jnp.ndarray,
+    bottom_left: jnp.ndarray,
+    bottom_right: jnp.ndarray,
+    gates: Tuple[jnp.ndarray, ...],
+    real_result: bool = False,
+) -> List[jnp.ndarray]:
+    density_matrix = jnp.tensordot(top_left, top_right, ((5, 6, 7), (2, 3, 4)))
+    density_matrix = jnp.tensordot(density_matrix, bottom_left, ((2, 3, 4), (5, 6, 7)))
+    density_matrix = jnp.tensordot(
+        density_matrix, bottom_right, ((4, 5, 6, 9, 10, 11), (2, 3, 4, 5, 6, 7))
+    )
+
+    density_matrix = density_matrix.transpose(0, 2, 4, 6, 1, 3, 5, 7)
+    density_matrix = density_matrix.reshape(
+        density_matrix.shape[0]
+        * density_matrix.shape[1]
+        * density_matrix.shape[2]
+        * density_matrix.shape[3],
+        density_matrix.shape[4]
+        * density_matrix.shape[5]
+        * density_matrix.shape[6]
+        * density_matrix.shape[7],
+    )
+
+    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
+        ]
+
+
+def calc_four_sites_quadrat_multiple_gates(
+    peps_tensors: Sequence[jnp.ndarray],
+    peps_tensor_objs: Sequence[PEPS_Tensor],
+    gates: Sequence[jnp.ndarray],
+) -> List[jnp.ndarray]:
+    """
+    Calculate the four site expectation values for three as quadrat ordered
+    PEPS tensor and their environment.
+
+    The order of the PEPS sequence have to be
+    [top-left, top-right, bottom-left, bottom-right].
+
+    The gate is applied in the order [top-left, top-right, bottom-left, bottom-right].
+
+    Args:
+      peps_tensors (:term:`sequence` of :obj:`jax.numpy.ndarray`):
+        The PEPS tensor arrays. Have to be the same objects as the tensor
+        attribute of the `peps_tensor_obj` argument.
+      peps_tensor_objs (:term:`sequence` of :obj:`~varipeps.peps.PEPS_Tensor`):
+        PEPS tensor objects.
+      gates (:term:`sequence` of :obj:`jax.numpy.ndarray`):
+        Sequence with the gates which should be applied to the PEPS tensors.
+        Gates are expected to be a matrix with first axis corresponding to
+        the Hilbert space and the second axis corresponding to the dual room.
+    Returns:
+      :obj:`list` of :obj:`jax.numpy.ndarray`:
+        List with the calculated expectation values of each gate.
+    """
+    density_matrix_top_left = apply_contraction_jitted(
+        "density_matrix_four_sites_top_left",
+        [peps_tensors[0]],
+        [peps_tensor_objs[0]],
+        [],
+    )
+
+    density_matrix_top_right = apply_contraction_jitted(
+        "density_matrix_four_sites_top_right",
+        [peps_tensors[1]],
+        [peps_tensor_objs[1]],
+        [],
+    )
+
+    density_matrix_bottom_left = apply_contraction_jitted(
+        "density_matrix_four_sites_bottom_left",
+        [peps_tensors[2]],
+        [peps_tensor_objs[2]],
+        [],
+    )
+
+    density_matrix_bottom_right = apply_contraction_jitted(
+        "density_matrix_four_sites_bottom_right",
+        [peps_tensors[3]],
+        [peps_tensor_objs[3]],
+        [],
+    )
+
+    real_result = all(jnp.allclose(g, g.T.conj()) for g in gates)
+
+    return _four_sites_quadrat_workhorse(
+        density_matrix_top_left,
+        density_matrix_top_right,
+        density_matrix_bottom_left,
+        density_matrix_bottom_right,
+        tuple(gates),
+        real_result,
+    )

varipeps/expectation/spiral_helpers.py

@@ -29,7 +29,8 @@ def apply_unitary(
       gate (:obj:`jax.numpy.ndarray`):
         The gate which should be updated with the unitary operator.
       delta_r (:obj:`jax.numpy.ndarray`):
-        Vector for the spatial difference.
+        Vector for the spatial difference. Can be a sequence if the spatial
+        difference are different for the single indices.
       q (:term:`sequence` of :obj:`jax.numpy.ndarray`):
         Sequence with the relevant wavevector for the different indices of
         the gate.
@@ -52,13 +53,17 @@ def apply_unitary(
       :obj:`jax.numpy.ndarray`:
         The updated gate with the unitary applied.
     """
-    if len(q) != len(apply_to_index):
+    if isinstance(delta_r, jnp.ndarray):
+        delta_r = (delta_r,) * len(apply_to_index)
+
+    if len(q) != len(apply_to_index) or len(q) != len(delta_r):
         raise ValueError("Length mismatch!")
 
     working_gate = gate.reshape((phys_d,) * 2 * number_sites)
 
     for index, i in enumerate(apply_to_index):
         w_q = q[index]
+        w_r = delta_r[index]
 
         if w_q.ndim == 0:
             w_q = jnp.array((w_q, w_q))
@@ -72,8 +77,8 @@ def apply_unitary(
         else:
             raise ValueError("Unknown wavevector type!")
 
-        # U = jsp.linalg.expm(1j * jnp.pi * jnp.dot(w_q, delta_r) * unitary_operator)
-        U = jnp.exp(1j * jnp.pi * jnp.dot(w_q, delta_r) * unitary_operator_D)
+        # U = jsp.linalg.expm(1j * jnp.pi * jnp.dot(w_q, w_r) * unitary_operator)
+        U = jnp.exp(1j * jnp.pi * jnp.dot(w_q, w_r) * unitary_operator_D)
         U = jnp.dot(
             unitary_operator_sigma * U[jnp.newaxis, :], unitary_operator_sigma.T.conj()
         )