Added Kitaev model Hamiltonian

pennylane/spin/__init__.py

@@ -16,4 +16,4 @@
 """
 
 from .lattice import Lattice
-from .spin_hamiltonian import fermi_hubbard, heisenberg, transverse_ising
+from .spin_hamiltonian import fermi_hubbard, heisenberg, kitaev, transverse_ising

pennylane/spin/lattice.py

@@ -35,7 +35,7 @@ class Lattice:
        positions (list[list[float]]): Initial positions of spin cites. Default value is
            ``[[0.0]`` :math:`\times` ``number of dimensions]``.
 
-       boundary_condition (bool or list[bool]): Defines boundary conditions in different lattice axes,
+       boundary_condition (bool or list[bool]): Defines boundary conditions for different lattice axes,
            default is ``False`` indicating open boundary condition.
        neighbour_order (int): Specifies the interaction level for neighbors within the lattice.
            Default is 1 (nearest neighbour).
@@ -72,6 +72,7 @@ def __init__(
         positions=None,
         boundary_condition=False,
         neighbour_order=1,
+        custom_edges=None,
         distance_tol=1e-5,
     ):
 
@@ -112,10 +113,19 @@ def __init__(
         n_sl = len(self.positions)
         self.n_sites = math.prod(n_cells) * n_sl
         self.lattice_points, lattice_map = self._generate_grid(neighbour_order)
+        if custom_edges is None:
+            cutoff = (
+                neighbour_order * math.max(math.linalg.norm(self.vectors, axis=1)) + distance_tol
+            )
+            edges = self._identify_neighbours(cutoff)
+            self.edges = Lattice._generate_true_edges(edges, lattice_map, neighbour_order)
+        else:
+            if neighbour_order > 1:
+                raise ValueError(
+                    "custom_edges and neighbour_order cannot be specified at the same time"
+                )
+            self.edges = self.get_custom_edges(custom_edges, lattice_map)
 
-        cutoff = neighbour_order * math.max(math.linalg.norm(self.vectors, axis=1)) + distance_tol
-        edges = self._identify_neighbours(cutoff)
-        self.edges = Lattice._generate_true_edges(edges, lattice_map, neighbour_order)
         self.edges_indices = [(v1, v2) for (v1, v2, color) in self.edges]
 
     def _identify_neighbours(self, cutoff):
@@ -190,7 +200,61 @@ def _generate_grid(self, neighbour_order):
             lattice_points.append(point)
             lattice_map.append(node_index)
 
-        return math.array(lattice_points), math.array(lattice_map)
+        return math.array(lattice_points), lattice_map
+
+    def get_custom_edges(self, custom_edges, lattice_map):
+        """Generates the edges described in `custom_edges` for all unit cells."""
+
+        if not all([len(edge) in (1, 2) for edge in custom_edges]):
+            raise TypeError(
+                """
+                custom_edges must be a list of tuples of length 1 or 2.
+                Every tuple must contain two lattice indices to represent the edge
+                and can optionally include a list to represent the operation and coefficient for that edge.
+                """
+            )
+
+        edges = []
+        n_sl = len(self.positions)
+        nsites_axis = math.cumprod([n_sl, *self.n_cells[:0:-1]])[::-1]
+
+        for i, custom_edge in enumerate(custom_edges):
+            edge = custom_edge[0]
+
+            if edge[0] >= self.n_sites or edge[1] >= self.n_sites:
+                raise ValueError(
+                    f"The edge {edge} has vertices greater than n_sites, {self.n_sites}"
+                )
+
+            edge_operation = custom_edge[1] if len(custom_edge) == 2 else i
+            map_edge1 = lattice_map.index(edge[0])
+            map_edge2 = lattice_map.index(edge[1])
+            edge_distance = self.lattice_points[map_edge2] - self.lattice_points[map_edge1]
+            v1, v2 = math.mod(edge, n_sl)
+            translation_vector = (
+                edge_distance + self.positions[v1] - self.positions[v2]
+            ) @ math.linalg.inv(self.vectors)
+            translation_vector = math.asarray(math.rint(translation_vector), dtype=int)
+            edge_ranges = []
+            for idx, cell in enumerate(self.n_cells):
+                if self.boundary_condition[idx]:
+                    edge_ranges.append(range(0, cell))
+                else:
+                    edge_ranges.append(
+                        range(
+                            math.maximum(0, -translation_vector[idx]),
+                            cell - math.maximum(0, translation_vector[idx]),
+                        )
+                    )
+
+            for cell in itertools.product(*edge_ranges):
+                node1_idx = math.dot(math.mod(cell, self.n_cells), nsites_axis) + v1
+                node2_idx = (
+                    math.dot(math.mod(cell + translation_vector, self.n_cells), nsites_axis) + v2
+                )
+                edges.append((node1_idx, node2_idx, edge_operation))
+
+        return edges
 
     def add_edge(self, edge_indices):
         r"""Adds a specific edge based on the site index without translating it.

pennylane/spin/spin_hamiltonian.py

@@ -19,7 +19,7 @@
 from pennylane import X, Y, Z, math
 from pennylane.fermi import FermiWord
 
-from .lattice import _generate_lattice
+from .lattice import Lattice, _generate_lattice
 
 # pylint: disable=too-many-arguments
 
@@ -309,3 +309,79 @@ def fermi_hubbard(
     qubit_ham = qml.qchem.qubit_observable(hamiltonian, mapping=mapping)
 
     return qubit_ham.simplify()
+
+
+def kitaev(n_cells, coupling=None, boundary_condition=False):
+    r"""Generates the Hamiltonian for the Kitaev model on the Honeycomb lattice.
+
+    The Hamiltonian is represented as:
+
+    .. math::
+        \begin{align*}
+          \hat{H} = K_x.\sum_{\langle i,j \rangle \epsilon X}\sigma_i^x\sigma_j^x +
+          K_y.\sum_{\langle i,j \rangle \epsilon Y}\sigma_i^y\sigma_j^y +
+          K_z.\sum_{\langle i,j \rangle \epsilon Z}\sigma_i^z\sigma_j^z
+        \end{align*}
+
+    where :math:`K_x`, :math:`K_y`, :math:`K_z` are the coupling constants defined for the Hamiltonian, and ``i,j`` represent
+    the indices for neighbouring spins.
+
+    Args:
+       n_cells (List[int]): Number of cells in each direction of the grid.
+       coupling (List[math.array[float]]): Coupling between spins, it is a list of length 3.
+                            Default value is [1.0, 1.0, 1.0].
+       boundary_condition (bool or list[bool]): Defines boundary conditions for different lattice axes,
+           default is ``False`` indicating open boundary condition.
+
+    Returns:
+       ~ops.op_math.Sum: Hamiltonian for the Kitaev model.
+
+    **Example**
+
+    >>> n_cells = [2,2]
+    >>> k = [0.5, 0.6, 0.7]
+    >>> spin_ham = qml.spin.kitaev(n_cells, coupling=k)
+    >>> spin_ham
+    (
+    0.5 * (X(0) @ X(1))
+    + 0.5 * (X(2) @ X(3))
+    + 0.5 * (X(4) @ X(5))
+    + 0.5 * (X(6) @ X(7))
+    + 0.6 * (Y(1) @ Y(2))
+    + 0.6 * (Y(5) @ Y(6))
+    + 0.7 * (Z(1) @ Z(4))
+    + 0.7 * (Z(3) @ Z(6))
+    )
+
+    """
+
+    if coupling is None:
+        coupling = [1.0, 1.0, 1.0]
+
+    if len(coupling) != 3:
+        raise ValueError("The coupling parameter should be a list of length 3")
+
+    vectors = [[1, 0], [0.5, 0.75**0.5]]
+    positions = [[0, 0], [0.5, 0.5 / 3**0.5]]
+    custom_edges = [
+        [(0, 1), ("XX", coupling[0])],
+        [(1, 2), ("YY", coupling[1])],
+        [(1, n_cells[1] * 2), ("ZZ", coupling[2])],
+    ]
+
+    lattice = Lattice(
+        n_cells=n_cells[0:2],
+        vectors=vectors,
+        positions=positions,
+        boundary_condition=boundary_condition,
+        custom_edges=custom_edges,
+    )
+    opmap = {"X": X, "Y": Y, "Z": Z}
+    hamiltonian = 0.0 * qml.I(0)
+    for edge in lattice.edges:
+        v1, v2 = edge[0:2]
+        op1, op2 = edge[2][0]
+        coeff = edge[2][1]
+        hamiltonian += coeff * (opmap[op1](v1) @ opmap[op2](v2))
+
+    return hamiltonian.simplify()

tests/spin/test_lattice.py

@@ -14,6 +14,8 @@
 """
 Unit tests for functions and classes needed for construct a lattice.
 """
+import re
+
 import numpy as np
 import pytest
 
@@ -539,3 +541,72 @@ def test_shape_error():
     lattice = "Octagon"
     with pytest.raises(ValueError, match="Lattice shape, 'Octagon' is not supported."):
         _generate_lattice(lattice=lattice, n_cells=n_cells)
+
+
+def test_neighbour_order_error():
+    r"""Test that an error is raised if neighbour order is greater than 1 when custom_edges are provided."""
+
+    vectors = [[0, 1], [1, 0]]
+    n_cells = [3, 3]
+    custom_edges = [[(0, 1)], [(0, 5)], [(0, 4)]]
+    with pytest.raises(
+        ValueError, match="custom_edges and neighbour_order cannot be specified at the same time"
+    ):
+        Lattice(n_cells=n_cells, vectors=vectors, neighbour_order=2, custom_edges=custom_edges)
+
+
+def test_custom_edge_type_error():
+    r"""Test that an error is raised if custom_edges are not provided as a list of length 1 or 2."""
+
+    vectors = [[0, 1], [1, 0]]
+    n_cells = [3, 3]
+    custom_edges = [[(0, 1), 1, 3], [(0, 5)], [(0, 4)]]
+    with pytest.raises(TypeError, match="custom_edges must be a list of tuples of length 1 or 2."):
+        Lattice(n_cells=n_cells, vectors=vectors, custom_edges=custom_edges)
+
+
+def test_custom_edge_value_error():
+    r"""Test that an error is raised if the custom_edges contains an edge with site_index greater than number of sites"""
+
+    vectors = [[0, 1], [1, 0]]
+    n_cells = [3, 3]
+    custom_edges = [[(0, 1)], [(0, 5)], [(0, 12)]]
+    with pytest.raises(
+        ValueError, match=re.escape("The edge (0, 12) has vertices greater than n_sites, 9")
+    ):
+        Lattice(n_cells=n_cells, vectors=vectors, custom_edges=custom_edges)
+
+
+@pytest.mark.parametrize(
+    # expected_edges here were obtained manually
+    ("vectors", "positions", "n_cells", "custom_edges", "expected_edges"),
+    [
+        (
+            [[0, 1], [1, 0]],
+            [[0, 0]],
+            [3, 3],
+            [[(0, 1)], [(0, 5)], [(0, 4)]],
+            [(0, 1, 0), (0, 5, 1), (0, 4, 2), (1, 2, 0), (3, 4, 0), (0, 4, 2), (1, 5, 2)],
+        ),
+        (
+            [[0, 1], [1, 0]],
+            [[0, 0]],
+            [3, 4],
+            [[(0, 1)], [(1, 4)], [(1, 5)]],
+            [(0, 1, 0), (1, 2, 0), (2, 3, 0), (1, 4, 1), (2, 5, 1), (0, 4, 2), (2, 6, 2)],
+        ),
+        (
+            [[1, 0], [0.5, np.sqrt(3) / 2]],
+            [[0.5, 0.5 / 3**0.5], [1, 1 / 3**0.5]],
+            [2, 2],
+            [[(0, 1)], [(1, 2)], [(1, 5)]],
+            [(2, 3, 0), (4, 5, 0), (1, 2, 1), (5, 6, 1), (1, 5, 2), (3, 7, 2)],
+        ),
+    ],
+)
+def test_custom_edges(vectors, positions, n_cells, custom_edges, expected_edges):
+    r"""Test that the edges are added as per custom_edges provided"""
+    lattice = Lattice(
+        n_cells=n_cells, vectors=vectors, positions=positions, custom_edges=custom_edges
+    )
+    assert np.all(np.isin(expected_edges, lattice.edges))

tests/spin/test_spin_hamiltonian.py

@@ -21,7 +21,7 @@
 
 import pennylane as qml
 from pennylane import I, X, Y, Z
-from pennylane.spin import fermi_hubbard, heisenberg, transverse_ising
+from pennylane.spin import fermi_hubbard, heisenberg, kitaev, transverse_ising
 
 pytestmark = pytest.mark.usefixtures("new_opmath_only")
 
@@ -782,3 +782,96 @@ def test_fermi_hubbard_hamiltonian_matrix(shape, n_cells, t, coulomb, expected_h
     fermi_hub_ham = fermi_hubbard(lattice=shape, n_cells=n_cells, hopping=t, coulomb=coulomb)
 
     qml.assert_equal(fermi_hub_ham, expected_ham)
+
+
+def test_coupling_error_kitaev():
+    r"""Test that an error is raised when the provided coupling shape is wrong for
+    Kitaev Hamiltonian."""
+    with pytest.raises(
+        ValueError,
+        match=re.escape("The coupling parameter should be a list of length 3"),
+    ):
+        kitaev(n_cells=[3, 4], coupling=[1.0, 2.0])
+
+
+@pytest.mark.parametrize(
+    # expected_ham here was obtained manually
+    ("n_cells", "j", "boundary_condition", "expected_ham"),
+    [
+        (
+            [2, 2, 1],
+            None,
+            False,
+            1.0 * (Z(1) @ Z(4))
+            + 1.0 * (Z(3) @ Z(6))
+            + 1.0 * (X(0) @ X(1))
+            + 1.0 * (X(2) @ X(3))
+            + 1.0 * (X(4) @ X(5))
+            + 1.0 * (X(6) @ X(7))
+            + 1.0 * (Y(1) @ Y(2))
+            + 1.0 * (Y(5) @ Y(6)),
+        ),
+        (
+            [2, 2],
+            [0.5, 0.6, 0.7],
+            False,
+            0.7 * (Z(1) @ Z(4))
+            + 0.7 * (Z(3) @ Z(6))
+            + 0.5 * (X(0) @ X(1))
+            + 0.5 * (X(2) @ X(3))
+            + 0.5 * (X(4) @ X(5))
+            + 0.5 * (X(6) @ X(7))
+            + 0.6 * (Y(1) @ Y(2))
+            + 0.6 * (Y(5) @ Y(6)),
+        ),
+        (
+            [2, 3],
+            [0.1, 0.2, 0.3],
+            True,
+            0.3 * (Z(1) @ Z(6))
+            + 0.3 * (Z(3) @ Z(8))
+            + 0.3 * (Z(5) @ Z(10))
+            + 0.3 * (Z(0) @ Z(7))
+            + 0.3 * (Z(2) @ Z(9))
+            + 0.3 * (Z(4) @ Z(11))
+            + 0.1 * (X(0) @ X(1))
+            + 0.1 * (X(2) @ X(3))
+            + 0.1 * (X(4) @ X(5))
+            + 0.1 * (X(6) @ X(7))
+            + 0.1 * (X(8) @ X(9))
+            + 0.1 * (X(10) @ X(11))
+            + 0.2 * (Y(1) @ Y(2))
+            + 0.2 * (Y(3) @ Y(4))
+            + 0.2 * (Y(0) @ Y(5))
+            + 0.2 * (Y(7) @ Y(8))
+            + 0.2 * (Y(9) @ Y(10))
+            + 0.2 * (Y(11) @ Y(6)),
+        ),
+        (
+            [2, 3],
+            [0.1, 0.2, 0.3],
+            [True, False],
+            0.3 * (Z(1) @ Z(6))
+            + 0.3 * (Z(3) @ Z(8))
+            + 0.3 * (Z(5) @ Z(10))
+            + 0.3 * (Z(0) @ Z(7))
+            + 0.3 * (Z(2) @ Z(9))
+            + 0.3 * (Z(4) @ Z(11))
+            + 0.1 * (X(0) @ X(1))
+            + 0.1 * (X(2) @ X(3))
+            + 0.1 * (X(4) @ X(5))
+            + 0.1 * (X(6) @ X(7))
+            + 0.1 * (X(8) @ X(9))
+            + 0.1 * (X(10) @ X(11))
+            + 0.2 * (Y(1) @ Y(2))
+            + 0.2 * (Y(3) @ Y(4))
+            + 0.2 * (Y(7) @ Y(8))
+            + 0.2 * (Y(9) @ Y(10)),
+        ),
+    ],
+)
+def test_kitaev_hamiltonian(n_cells, j, boundary_condition, expected_ham):
+    r"""Test that the correct Hamiltonian is generated"""
+    kitaev_ham = kitaev(n_cells=n_cells, coupling=j, boundary_condition=boundary_condition)
+
+    qml.assert_equal(kitaev_ham, expected_ham)