Merge branch 'master' into implicit-casting

doc/releases/changelog-dev.md

@@ -4,6 +4,10 @@
 
 <h3>New features since last release</h3>
 
+* Function is added for generating the spin Hamiltonian for the
+  [Kitaev](https://arxiv.org/abs/cond-mat/0506438) model on a lattice.
+  [(#6174)](https://github.com/PennyLaneAI/pennylane/pull/6174)
+
 * Functions are added for generating spin Hamiltonians for [Emery]
   (https://journals.aps.org/prl/abstract/10.1103/PhysRevLett.58.2794) and
   [Haldane](https://journals.aps.org/prl/pdf/10.1103/PhysRevLett.61.2015) models on a lattice.

pennylane/spin/__init__.py

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

pennylane/spin/lattice.py

@@ -34,11 +34,16 @@ class Lattice:
        vectors (list[list[float]]): Primitive vectors for the 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.
-           Default is ``False`` indicating open boundary condition.
+       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, indicating nearest neighbour.
+           Default is 1, indicating nearest neighbour. This cannot be greater than 1 if custom_edges is defined.
+       custom_edges (Optional[list(list(tuples))]): Specifies the edges to be added in the lattice.
+           Default value is ``None``, which adds the edges based on ``neighbour_order``.
+           Each element in the list is for a separate edge, and can contain 1 or 2 tuples.
+           First tuple contains the indices of the starting and ending vertices of the edge.
+           Second tuple is optional and contains the operator on that edge and coefficient
+           of that operator. Default value is the index of edge in custom_edges list.
        distance_tol (float): Distance below which spatial points are considered equal for the
            purpose of identifying nearest neighbours. Default value is 1e-5.
 
@@ -72,6 +77,7 @@ def __init__(
         positions=None,
         boundary_condition=False,
         neighbour_order=1,
+        custom_edges=None,
         distance_tol=1e-5,
     ):
 
@@ -112,10 +118,20 @@ 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 cannot be specified if neighbour_order argument is set to greater than 1."
+                )
+            lattice_map = dict(zip(lattice_map, self.lattice_points))
+            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):
@@ -187,7 +203,113 @@ 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.
+
+        Args:
+          custom_edges (Optional[list(list(tuples))]): Specifies the edges to be added in the lattice.
+              Default value is None, which adds the edges based on neighbour_order.
+              Each element in the list is for a separate edge, and can contain 1 or 2 tuples.
+              First tuple contains the index of the starting and ending vertex of the edge.
+              Second tuple is optional and contains the operator on that edge and coefficient
+              of that operator.
+          lattice_map (list[int]): A list to represent the node number for each lattice_point.
+
+        Returns:
+          List of edges.
+
+        **Example**
+
+        Generates a square lattice with a single diagonal and assigns a different operation
+        to horizontal, vertical, and diagonal edges.
+        >>> n_cells = [3,3]
+        >>> vectors = [[1, 0], [0,1]]
+        >>> custom_edges = [
+                [(0, 1), ("XX", 0.1)],
+                [(0, 3), ("YY", 0.2)],
+                [(0, 4), ("XY", 0.3)],
+            ]
+        >>> lattice = qml.spin.Lattice(n_cells=n_cells, vectors=vectors, custom_edges=custom_edges)
+        >>> lattice.edges
+        [(0, 1, ('XX', 0.1)),
+         (1, 2, ('XX', 0.1)),
+         (3, 4, ('XX', 0.1)),
+         (4, 5, ('XX', 0.1)),
+         (6, 7, ('XX', 0.1)),
+         (7, 8, ('XX', 0.1)),
+         (0, 3, ('YY', 0.2)),
+         (1, 4, ('YY', 0.2)),
+         (2, 5, ('YY', 0.2)),
+         (3, 6, ('YY', 0.2)),
+         (4, 7, ('YY', 0.2)),
+         (5, 8, ('YY', 0.2)),
+         (0, 4, ('XY', 0.3)),
+         (1, 5, ('XY', 0.3)),
+         (3, 7, ('XY', 0.3)),
+         (4, 8, ('XY', 0.3))
+        ]
+
+        """
+
+        for edge in custom_edges:
+            if len(edge) not in (1, 2):
+                raise TypeError(
+                    """
+                    The elements of custom_edges should be lists of length 1 or 2.
+                    Inside said lists should be a tuple that contains two lattice
+                    indices to represent the edge and, optionally, a tuple that represents
+                    the operation and coefficient for that edge.
+                    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.
+                    """
+                )
+
+            if edge[0][0] >= self.n_sites or edge[0][1] >= self.n_sites:
+                raise ValueError(
+                    f"The edge {edge[0]} has vertices greater than n_sites, {self.n_sites}"
+                )
+
+        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]
+
+            edge_operation = custom_edge[1] if len(custom_edge) == 2 else i
+
+            # Finds the coordinates of starting and ending vertices of the edge
+            # and the vector distance between the coordinates
+            vertex1 = lattice_map[edge[0]]
+            vertex2 = lattice_map[edge[1]]
+            edge_distance = vertex2 - vertex1
+
+            # Calculates the number of unit cells that a given edge spans in each direction
+            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)
+
+            # Finds the minimum and maximum range for a given edge based on boundary_conditions
+            edge_ranges = []
+            for idx, cell in enumerate(self.n_cells):
+                t_point = 0 if self.boundary_condition[idx] else translation_vector[idx]
+                edge_ranges.append(
+                    range(math.maximum(0, -t_point), cell - math.maximum(0, t_point))
+                )
+
+            # Finds the indices for starting and ending vertices of the edge
+            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, too-many-branches
 
@@ -608,3 +608,84 @@ def haldane(
     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 `Kitaev <https://arxiv.org/abs/cond-mat/0506438>`_ model Hamiltonian is represented as:
+
+    .. math::
+        \begin{align*}
+          \hat{H} = K_x.\sum_{\langle i,j \rangle \in X}\sigma_i^x\sigma_j^x +
+          \:\: K_y.\sum_{\langle i,j \rangle \in Y}\sigma_i^y\sigma_j^y +
+          \:\: K_z.\sum_{\langle i,j \rangle \in 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 :math:`X`, :math:`Y`, :math:`Z` represent the set of edges in the Honeycomb lattice between spins
+    :math:`i` and :math:`j` with real-space bond directions :math:`[0, 1], [\frac{\sqrt{3}}{2}, \frac{1}{2}],
+    \frac{\sqrt{3}}{2}, -\frac{1}{2}]`, respectively.
+
+    Args:
+       n_cells (list[int]): Number of cells in each direction of the grid.
+       coupling (Optional[list[float] or tensor_like(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.
+           The default is ``False``, indicating open boundary conditions for all.
+
+    Raises:
+       ValueError: if ``coupling`` doesn't have correct dimensions.
+
+    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
 
@@ -791,6 +793,78 @@ def test_shape_error():
         _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 cannot be specified if neighbour_order argument is set to greater than 1.",
+    ):
+        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="The elements of custom_edges should be lists 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))
+
+
 def test_dimension_error():
     r"""Test that an error is raised if wrong dimension is provided for a given lattice shape."""
     n_cells = [5, 5, 5]