Minor changes to colorings and term representation; added Heisenberg …

source/pip/qsharp/magnets/models/model.py

@@ -4,7 +4,7 @@
 # pyright: reportPrivateImportUsage=false
 
 from collections.abc import Sequence
-from typing import Optional
+from typing import Iterator, Optional
 
 
 """Base Model class for quantum spin models.
@@ -54,8 +54,11 @@ def __init__(self, geometry: Hypergraph):
 
         Creates a quantum spin model on the given geometry.
 
-        The model stores operators lazily in ``_ops`` as terms are defined.
-        ``_terms`` is initialized with one empty term group.
+        The model stores operators lazily in ``_ops`` as interaction operators
+        are defined. Noncommuting collections of operators are collected in
+        ``_terms`` that stores the indices of its interaction operators. This
+        list of arrays seperate terms into parallizable groups by color. It is
+        initialized as one empty term group.
 
         Args:
             geometry: Hypergraph defining the interaction topology. The number
@@ -66,14 +69,15 @@ def __init__(self, geometry: Hypergraph):
         self._ops: list[PauliString] = []
         for edge in geometry.edges():
             self._qubits.update(edge.vertices)
-        self._terms: dict[int, list[int]] = {}
+        self._terms: dict[int, dict[int, list[int]]] = {}
 
     def add_interaction(
         self,
         edge: Hyperedge,
         pauli_string: Sequence[int | str] | str,
         coefficient: complex = 1.0,
         term: Optional[int] = None,
+        color: int = 0,
     ) -> None:
         """Add an interaction term to the model.
 
@@ -88,8 +92,16 @@ def add_interaction(
         self._ops.append(s)
         if term is not None:
             if term not in self._terms:
-                self._terms[term] = []
-            self._terms[term].append(len(self._ops) - 1)
+                self._terms[term] = {}
+            if color not in self._terms[term]:
+                self._terms[term][color] = []
+            self._terms[term][color].append(len(self._ops) - 1)
+
+    def terms(self, t: int) -> Iterator[PauliString]:
+        """Get the list of PauliStrings corresponding to a term group."""
+        if t not in self._terms:
+            raise ValueError("Term group does not exist.")
+        return iter([self._ops[i] for i in self._terms[t]])
 
     @property
     def nqubits(self) -> int:
@@ -126,12 +138,67 @@ class IsingModel(Model):
 
     def __init__(self, geometry: Hypergraph, h: float, J: float):
         super().__init__(geometry)
-        self.coloring: HypergraphEdgeColoring = geometry.edge_coloring()
-        self._terms = {0: [], 1: []}
+        self.h = h
+        self.J = J
+        self._terms = {0: {}, 1: {}}
 
+        coloring: HypergraphEdgeColoring = geometry.edge_coloring()
         for edge in geometry.edges():
             vertices = edge.vertices
             if len(vertices) == 1:
-                self.add_interaction(edge, "X", -h, term=0)
+                self.add_interaction(edge, "X", -h, term=0, color=0)
             elif len(vertices) == 2:
-                self.add_interaction(edge, "ZZ", -J, term=1)
+                color = coloring.color(edge.vertices)
+                if color is None:
+                    raise ValueError("Geometry edge coloring failed to assign a color.")
+                self.add_interaction(edge, "ZZ", -J, term=1, color=color)
+
+    def __str__(self) -> str:
+        return (
+            f"Ising model with {self.nterms} terms on {self.nqubits} qubits "
+            f"(h={self.h}, J={self.J})."
+        )
+
+    def __repr__(self) -> str:
+        return (
+            f"IsingModel(nqubits={self.nqubits}, nterms={self.nterms}, "
+            f"h={self.h}, J={self.J})"
+        )
+
+
+class HeisenbergModel(Model):
+    """Translation-invariant Heisenberg model on a hypergraph geometry.
+
+    The Hamiltonian is:
+        H = -J * Σ_{<i,j>} (X_i X_j + Y_i Y_j + Z_i Z_j)
+
+    - Two-vertex edges define XX, YY, and ZZ coupling terms with coefficient ``-J``.
+    - Terms are grouped into three parts: ``0`` for XX, ``1`` for YY, and ``2`` for ZZ.
+    """
+
+    def __init__(self, geometry: Hypergraph, J: float):
+        super().__init__(geometry)
+        self.J = J
+        self.coloring: HypergraphEdgeColoring = geometry.edge_coloring()
+        self._terms = {0: {}, 1: {}, 2: {}}
+        for edge in geometry.edges():
+            vertices = edge.vertices
+            if len(vertices) == 2:
+                color = self.coloring.color(edge.vertices)
+                if color is None:
+                    raise ValueError("Geometry edge coloring failed to assign a color.")
+                self.add_interaction(edge, "XX", -J, term=0, color=color)
+                self.add_interaction(edge, "YY", -J, term=1, color=color)
+                self.add_interaction(edge, "ZZ", -J, term=2, color=color)
+
+    def __str__(self) -> str:
+        return (
+            f"Heisenberg model with {self.nterms} terms on {self.nqubits} qubits "
+            f"(J={self.J})."
+        )
+
+    def __repr__(self) -> str:
+        return (
+            f"HeisenbergModel(nqubits={self.nqubits}, nterms={self.nterms}, "
+            f"J={self.J})"
+        )

source/pip/qsharp/magnets/utilities/hypergraph.py

@@ -205,9 +205,9 @@ class HypergraphEdgeColoring:
 
     Note:
         Colors are keyed by edge vertex tuples (``edge.vertices``), not by
-        ``Hyperedge`` object identity. As a result, :meth:`color` accepts any
-        ``Hyperedge`` with matching vertices, while :meth:`add_edge` still
-        requires an edge instance that belongs to :attr:`hypergraph`.
+        ``Hyperedge`` object identity. As a result, :meth:`color` accepts edge
+        vertex tuples directly, while :meth:`add_edge` still requires an edge
+        instance that belongs to :attr:`hypergraph`.
 
     Attributes:
         hypergraph: The supporting :class:`Hypergraph` whose edges can be
@@ -226,20 +226,22 @@ def ncolors(self) -> int:
         """Return the number of distinct nonnegative colors in the coloring."""
         return len(self._used_vertices)
 
-    def color(self, edge: Hyperedge) -> Optional[int]:
-        """Return the color assigned to a specific edge.
+    def color(self, vertices: tuple[int, ...]) -> Optional[int]:
+        """Return the color assigned to edge vertices.
 
         Args:
-            edge: Hyperedge to query. Any ``Hyperedge`` with the same
-                ``vertices`` tuple resolves to the same stored color.
+            vertices: Canonical vertex tuple for the edge to query (typically
+                ``edge.vertices``).
 
         Returns:
-            The color assigned to ``edge``, or ``None`` if the edge has not
-            been added to this coloring.
+            The color assigned to ``vertices``, or ``None`` if the edge has
+            not been added to this coloring.
         """
-        if not isinstance(edge, Hyperedge):
-            raise TypeError(f"edge must be Hyperedge, got {type(edge).__name__}")
-        return self._colors.get(edge.vertices)
+        if not isinstance(vertices, tuple) or not all(
+            isinstance(vertex, int) for vertex in vertices
+        ):
+            raise TypeError("vertices must be tuple[int, ...]")
+        return self._colors.get(vertices)
 
     def colors(self) -> Iterator[int]:
         """Iterate over distinct nonnegative colors present in the coloring.

source/pip/tests/magnets/test_complete.py

@@ -221,7 +221,7 @@ def test_complete_bipartite_graph_coloring_non_overlapping():
     # Group edges by color
     colors = {}
     for edge in graph.edges():
-        color = coloring.color(edge)
+        color = coloring.color(edge.vertices)
         assert color is not None
         edge_vertices = edge.vertices
         if color not in colors:

source/pip/tests/magnets/test_hypergraph.py

@@ -173,13 +173,13 @@ def test_hypergraph_edge_coloring_rejects_equivalent_edge_not_in_hypergraph():
 
 
 def test_hypergraph_edge_coloring_color_matches_equivalent_vertices():
-    """Test color lookup uses edge vertices, not Hyperedge object identity."""
+    """Test color lookup uses edge vertex tuples as keys."""
     edge = Hyperedge([0, 1])
     graph = Hypergraph([edge])
     coloring = HypergraphEdgeColoring(graph)
 
     coloring.add_edge(edge, 3)
-    assert coloring.color(Hyperedge([1, 0])) == 3
+    assert coloring.color((0, 1)) == 3
 
 
 def test_hypergraph_edge_coloring_rejects_negative_color_for_nontrivial_edge():
@@ -225,7 +225,7 @@ def test_hypergraph_add_edge_with_color():
     coloring = HypergraphEdgeColoring(graph)
     coloring.add_edge(edge, color=1)
     assert graph.nedges == 2
-    assert coloring.color(edge) == 1
+    assert coloring.color(edge.vertices) == 1
 
 
 def test_hypergraph_color_default():
@@ -297,7 +297,7 @@ def test_greedy_edge_coloring_single_edge():
     edge = Hyperedge([0, 1])
     graph = Hypergraph([edge])
     colored = graph.edge_coloring(seed=42)
-    assert colored.color(edge) == 0
+    assert colored.color(edge.vertices) == 0
     assert colored.ncolors == 1
 
 
@@ -307,8 +307,8 @@ def test_greedy_edge_coloring_non_overlapping():
     graph = Hypergraph(edges)
     colored = graph.edge_coloring(seed=42)
     # Non-overlapping edges can be in the same color
-    assert colored.color(edges[0]) is not None
-    assert colored.color(edges[1]) is not None
+    assert colored.color(edges[0].vertices) is not None
+    assert colored.color(edges[1].vertices) is not None
     assert colored.ncolors == 1
 
 
@@ -318,8 +318,8 @@ def test_greedy_edge_coloring_overlapping():
     graph = Hypergraph(edges)
     colored = graph.edge_coloring(seed=42)
     # Overlapping edges need different colors
-    assert colored.color(edges[0]) is not None
-    assert colored.color(edges[1]) is not None
+    assert colored.color(edges[0].vertices) is not None
+    assert colored.color(edges[1].vertices) is not None
     assert colored.ncolors == 2
 
 
@@ -329,9 +329,9 @@ def test_greedy_edge_coloring_triangle():
     graph = Hypergraph(edges)
     colored = graph.edge_coloring(seed=42)
     # All edges share vertices pairwise, so need 3 colors
-    assert colored.color(edges[0]) is not None
-    assert colored.color(edges[1]) is not None
-    assert colored.color(edges[2]) is not None
+    assert colored.color(edges[0].vertices) is not None
+    assert colored.color(edges[1].vertices) is not None
+    assert colored.color(edges[2].vertices) is not None
     assert colored.ncolors == 3
 
 
@@ -350,7 +350,7 @@ def test_greedy_edge_coloring_validity():
     # Group edges by color
     colors = {}
     for edge in edges:
-        color = colored.color(edge)
+        color = colored.color(edge.vertices)
         assert color is not None
         if color not in colors:
             colors[color] = []
@@ -372,9 +372,9 @@ def test_greedy_edge_coloring_all_edges_colored():
     colored = graph.edge_coloring(seed=42)
 
     # All edges should have a color assigned
-    assert colored.color(edges[0]) is not None
-    assert colored.color(edges[1]) is not None
-    assert colored.color(edges[2]) is not None
+    assert colored.color(edges[0].vertices) is not None
+    assert colored.color(edges[1].vertices) is not None
+    assert colored.color(edges[2].vertices) is not None
 
 
 def test_greedy_edge_coloring_reproducible_with_seed():
@@ -385,8 +385,8 @@ def test_greedy_edge_coloring_reproducible_with_seed():
     colored1 = graph.edge_coloring(seed=123)
     colored2 = graph.edge_coloring(seed=123)
 
-    color_map_1 = {edge.vertices: colored1.color(edge) for edge in edges}
-    color_map_2 = {edge.vertices: colored2.color(edge) for edge in edges}
+    color_map_1 = {edge.vertices: colored1.color(edge.vertices) for edge in edges}
+    color_map_2 = {edge.vertices: colored2.color(edge.vertices) for edge in edges}
     assert color_map_1 == color_map_2
 
 
@@ -415,9 +415,9 @@ def test_greedy_edge_coloring_hyperedges():
     colored = graph.edge_coloring(seed=42)
 
     # First two share vertex 2, third is independent
-    assert colored.color(edges[0]) is not None
-    assert colored.color(edges[1]) is not None
-    assert colored.color(edges[2]) is not None
+    assert colored.color(edges[0].vertices) is not None
+    assert colored.color(edges[1].vertices) is not None
+    assert colored.color(edges[2].vertices) is not None
     assert colored.ncolors >= 2
 
 
@@ -428,7 +428,7 @@ def test_greedy_edge_coloring_self_loops():
     colored = graph.edge_coloring(seed=42)
 
     # Self-loops use the special -1 color and do not contribute to ncolors.
-    assert colored.color(edges[0]) == -1
-    assert colored.color(edges[1]) == -1
-    assert colored.color(edges[2]) == -1
+    assert colored.color(edges[0].vertices) == -1
+    assert colored.color(edges[1].vertices) == -1
+    assert colored.color(edges[2].vertices) == -1
     assert colored.ncolors == 0

source/pip/tests/magnets/test_lattice1d.py

@@ -9,7 +9,7 @@
 
 def _vertex_color_map(graph) -> dict[tuple[int, ...], int | None]:
     coloring = graph.edge_coloring()
-    return {edge.vertices: coloring.color(edge) for edge in graph.edges()}
+    return {edge.vertices: coloring.color(edge.vertices) for edge in graph.edges()}
 
 
 # Chain1D tests
@@ -101,7 +101,7 @@ def test_chain1d_coloring_non_overlapping():
     # Group edges by color
     colors = {}
     for edge in chain.edges():
-        color = coloring.color(edge)
+        color = coloring.color(edge.vertices)
         assert color is not None
         edge_vertices = edge.vertices
         if color not in colors:
@@ -211,7 +211,7 @@ def test_ring1d_coloring_non_overlapping():
     # Group edges by color
     colors = {}
     for edge in ring.edges():
-        color = coloring.color(edge)
+        color = coloring.color(edge.vertices)
         assert color is not None
         edge_vertices = edge.vertices
         if color not in colors:

source/pip/tests/magnets/test_lattice2d.py

@@ -9,7 +9,7 @@
 
 def _vertex_color_map(graph) -> dict[tuple[int, ...], int | None]:
     coloring = graph.edge_coloring()
-    return {edge.vertices: coloring.color(edge) for edge in graph.edges()}
+    return {edge.vertices: coloring.color(edge.vertices) for edge in graph.edges()}
 
 
 # Patch2D tests
@@ -117,7 +117,7 @@ def test_patch2d_coloring_non_overlapping():
     # Group edges by color
     colors = {}
     for edge in patch.edges():
-        color = coloring.color(edge)
+        color = coloring.color(edge.vertices)
         assert color is not None
         edge_vertices = edge.vertices
         if color not in colors:
@@ -243,7 +243,7 @@ def test_torus2d_coloring_non_overlapping():
     # Group edges by color
     colors = {}
     for edge in torus.edges():
-        color = coloring.color(edge)
+        color = coloring.color(edge.vertices)
         assert color is not None
         edge_vertices = edge.vertices
         if color not in colors: