Build hamiltonians from custom lattices (#6226)

[sc-70985]

---------

Co-authored-by: ddhawan11 <ddhawan@umich.edu>
Co-authored-by: Diksha Dhawan <40900030+ddhawan11@users.noreply.github.com>
Co-authored-by: Utkarsh <utkarshazad98@gmail.com>
Co-authored-by: soranjh <40344468+soranjh@users.noreply.github.com>

doc/releases/changelog-dev.md

@@ -8,6 +8,9 @@
   [Kitaev](https://arxiv.org/abs/cond-mat/0506438) model on a lattice.
   [(#6174)](https://github.com/PennyLaneAI/pennylane/pull/6174)
 
+* Function is added for generating the spin Hamiltonians for custom lattices.
+  [(#6226)](https://github.com/PennyLaneAI/pennylane/pull/6226)
+
 * 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.
@@ -203,6 +206,7 @@ Diksha Dhawan,
 Lillian M. A. Frederiksen,
 Pietropaolo Frisoni,
 Emiliano Godinez,
+Austin Huang,
 Christina Lee,
 William Maxwell,
 Lee J. O'Riordan,

pennylane/spin/__init__.py

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

pennylane/spin/lattice.py

@@ -24,6 +24,7 @@
 
 # pylint: disable=too-many-arguments, too-many-instance-attributes
 # pylint: disable=use-a-generator, too-few-public-methods
+# pylint: disable=too-many-branches
 
 
 class Lattice:
@@ -32,7 +33,7 @@ class Lattice:
     Args:
        n_cells (list[int]): Number of cells in each direction of the grid.
        vectors (list[list[float]]): Primitive vectors for the lattice.
-       positions (list[list[float]]): Initial positions of spin cites. Default value is
+       positions (list[list[float]]): Initial positions of spin sites. Default value is
            ``[[0.0]`` :math:`\times` ``number of dimensions]``.
        boundary_condition (bool or list[bool]): Defines boundary conditions for different lattice axes,
            default is ``False`` indicating open boundary condition.
@@ -44,6 +45,11 @@ class Lattice:
            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.
+       custom_nodes (Optional(list(list(int, tuples)))): Specifies the on-site potentials and
+           operators for nodes in the lattice. The default value is `None`, which means no on-site
+           potentials. Each element in the list is for a separate node. For each element, the first
+           value is the index of the node, and the second element is a tuple which contains the
+           operator and coefficient.
        distance_tol (float): Distance below which spatial points are considered equal for the
            purpose of identifying nearest neighbours. Default value is 1e-5.
 
@@ -54,6 +60,7 @@ class Lattice:
           if ``positions`` doesn't have a dimension of 2.
           if ``vectors`` doesn't have a dimension of 2 or the length of vectors is not equal to the number of vectors.
           if ``boundary_condition`` is not a bool or a list of bools with length equal to the number of vectors
+          if ``custom_nodes`` contains nodes with negative indices or indices greater than number of sites
 
     Returns:
        Lattice object
@@ -78,6 +85,7 @@ def __init__(
         boundary_condition=False,
         neighbour_order=1,
         custom_edges=None,
+        custom_nodes=None,
         distance_tol=1e-5,
     ):
 
@@ -134,6 +142,17 @@ def __init__(
 
         self.edges_indices = [(v1, v2) for (v1, v2, color) in self.edges]
 
+        if custom_nodes is not None:
+            for node in custom_nodes:
+                if node[0] > self.n_sites:
+                    raise ValueError(
+                        "The custom node has an index larger than the number of sites."
+                    )
+                if node[0] < 0:
+                    raise ValueError("The custom node has an index smaller than 0.")
+
+        self.nodes = custom_nodes
+
     def _identify_neighbours(self, cutoff):
         r"""Identifies the connections between lattice points and returns the unique connections
         based on the neighbour_order. This function uses KDTree to identify neighbours, which

pennylane/spin/spin_hamiltonian.py

@@ -16,7 +16,7 @@
 """
 
 import pennylane as qml
-from pennylane import X, Y, Z, math
+from pennylane import I, X, Y, Z, math
 from pennylane.fermi import FermiWord
 
 from .lattice import Lattice, _generate_lattice
@@ -689,3 +689,66 @@ def kitaev(n_cells, coupling=None, boundary_condition=False):
         hamiltonian += coeff * (opmap[op1](v1) @ opmap[op2](v2))
 
     return hamiltonian.simplify()
+
+
+def spin_hamiltonian(lattice):
+    r"""Generates a spin Hamiltonian for a custom lattice.
+
+    Args:
+        lattice (Lattice): custom lattice defined with custom_edges
+
+    Raises:
+        ValueError: if the provided Lattice does not have ``custom_edges`` defined with operators
+
+    Returns:
+        ~ops.op_math.Sum: Hamiltonian for the lattice
+
+    **Example**
+
+    .. code-block:: python
+
+        >>> lattice = qml.spin.Lattice(
+        ...     n_cells=[2, 2],
+        ...     vectors=[[1, 0], [0, 1]],
+        ...     positions=[[0, 0], [1, 5]],
+        ...     boundary_condition=False,
+        ...     custom_edges=[[(0, 1), ("XX", 0.5)], [(1, 2), ("YY", 0.6)], [(1, 4), ("ZZ", 0.7)]],
+        ...     custom_nodes=[[0, ("X", 0.5)], [1, ("Y", 0.3)]],
+        ... )
+        >>> qml.spin.spin_hamiltonian(lattice=lattice)
+        (
+            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))
+          + 0.5 * X(0)
+          + 0.3 * Y(1)
+        )
+
+    """
+    if not isinstance(lattice.edges[0][2], tuple):
+        raise ValueError(
+            "Custom edges need to be defined and should have an operator defined as a `str`"
+        )
+
+    opmap = {"I": I, "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))
+
+    if lattice.nodes is not None:
+        for node in lattice.nodes:
+            n = node[0]
+            op = node[1][0]
+            coeff = node[1][1]
+            hamiltonian += coeff * (opmap[op](n))
+
+    return hamiltonian.simplify()

tests/spin/test_lattice.py

@@ -865,6 +865,59 @@ def test_custom_edges(vectors, positions, n_cells, custom_edges, expected_edges)
     assert np.all(np.isin(expected_edges, lattice.edges))
 
 
+@pytest.mark.parametrize(
+    # expected_nodes here were obtained manually
+    ("vectors", "positions", "n_cells", "custom_nodes", "expected_nodes"),
+    [
+        (
+            [[0, 1], [1, 0]],
+            [[0, 0]],
+            [3, 3],
+            [[0, ("X", 0.3)], [2, ("Y", 0.3)]],
+            [[0, ("X", 0.3)], [2, ("Y", 0.3)]],
+        ),
+        (
+            [[1, 0], [0.5, np.sqrt(3) / 2]],
+            [[0.5, 0.5 / 3**0.5], [1, 1 / 3**0.5]],
+            [2, 2],
+            [[0, ("X", 0.3)], [2, ("Y", 0.3)], [1, ("Z", 0.9)]],
+            [[0, ("X", 0.3)], [2, ("Y", 0.3)], [1, ("Z", 0.9)]],
+        ),
+    ],
+)
+def test_custom_nodes(vectors, positions, n_cells, custom_nodes, expected_nodes):
+    r"""Test that the nodes are added as per custom_nodes provided"""
+    lattice = Lattice(
+        n_cells=n_cells, vectors=vectors, positions=positions, custom_nodes=custom_nodes
+    )
+
+    assert lattice.nodes == expected_nodes
+
+
+@pytest.mark.parametrize(
+    ("vectors", "positions", "n_cells", "custom_nodes"),
+    [
+        (
+            [[0, 1], [1, 0]],
+            [[0, 0]],
+            [3, 3],
+            [[0, ("X", 0.3)], [-202, ("Y", 0.3)]],
+        ),
+        (
+            [[1, 0], [0.5, np.sqrt(3) / 2]],
+            [[0.5, 0.5 / 3**0.5], [1, 1 / 3**0.5]],
+            [2, 2],
+            [[0, ("X", 0.3)], [204, ("Y", 0.3)], [1, ("Z", 0.9)]],
+        ),
+    ],
+)
+def test_custom_nodes_error(vectors, positions, n_cells, custom_nodes):
+    r"""Test that the incompatible `custom_nodes` raise correct error"""
+
+    with pytest.raises(ValueError, match="The custom node has"):
+        Lattice(n_cells=n_cells, vectors=vectors, positions=positions, custom_nodes=custom_nodes)
+
+
 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]

tests/spin/test_spin_hamiltonian.py

@@ -21,7 +21,16 @@
 
 import pennylane as qml
 from pennylane import I, X, Y, Z
-from pennylane.spin import emery, fermi_hubbard, haldane, heisenberg, kitaev, transverse_ising
+from pennylane.spin import (
+    Lattice,
+    emery,
+    fermi_hubbard,
+    haldane,
+    heisenberg,
+    kitaev,
+    spin_hamiltonian,
+    transverse_ising,
+)
 
 # pylint: disable=too-many-arguments
 pytestmark = pytest.mark.usefixtures("new_opmath_only")
@@ -1394,3 +1403,89 @@ def test_kitaev_hamiltonian(n_cells, j, boundary_condition, expected_ham):
     kitaev_ham = kitaev(n_cells=n_cells, coupling=j, boundary_condition=boundary_condition)
 
     qml.assert_equal(kitaev_ham, expected_ham)
+
+
+@pytest.mark.parametrize(
+    ("lattice", "expected_ham"),
+    [
+        # This is the Hamiltonian for the Kitaev model on the Honeycomb lattice
+        (
+            Lattice(
+                n_cells=[2, 2],
+                vectors=[[1, 0], [0, 1]],
+                positions=[[0, 0], [1, 5]],
+                boundary_condition=False,
+                custom_edges=[[(0, 1), ("XX", 0.5)], [(1, 2), ("YY", 0.6)], [(1, 4), ("ZZ", 0.7)]],
+            ),
+            (
+                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))
+            ),
+        ),
+        (
+            Lattice(
+                n_cells=[2, 2],
+                vectors=[[1, 0], [0, 1]],
+                positions=[[0, 0], [1, 5]],
+                boundary_condition=False,
+                custom_edges=[[(0, 1), ("XX", 0.5)], [(1, 2), ("YY", 0.6)], [(1, 4), ("ZZ", 0.7)]],
+                custom_nodes=[[0, ("X", 0.3)], [7, ("Y", 0.9)]],
+            ),
+            (
+                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))
+                + 0.3 * X(0)
+                + 0.9 * Y(7)
+            ),
+        ),
+        (
+            Lattice(
+                n_cells=[2, 2],
+                vectors=[[1, 0], [0, 1]],
+                positions=[[0, 0], [1, 5]],
+                boundary_condition=False,
+                custom_edges=[[(0, 1), ("XX", 0.5)], [(1, 2), ("YY", 0.6)], [(1, 4), ("ZZ", 0.7)]],
+                custom_nodes=[[0, ("X", 0.3)], [7, ("Y", 0.9)], [0, ("X", 0.5)]],
+            ),
+            (
+                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))
+                + 0.8 * X(0)
+                + 0.9 * Y(7)
+            ),
+        ),
+    ],
+)
+def test_spin_hamiltonian(lattice, expected_ham):
+    r"""Test that the correct Hamiltonian is generated from a given Lattice"""
+    spin_ham = spin_hamiltonian(lattice=lattice)
+
+    qml.assert_equal(spin_ham, expected_ham)
+
+
+def test_spin_hamiltonian_error():
+    r"""Test that the correct error is raised Hamiltonian with incompatible Lattice"""
+    lattice = Lattice(n_cells=[2, 2], vectors=[[1, 0], [0, 1]], positions=[[0, 0], [1, 1]])
+    with pytest.raises(
+        ValueError,
+        match="Custom edges need to be defined and should have an operator defined as a `str`",
+    ):
+        spin_hamiltonian(lattice=lattice)