More tests and cleanup

bqskit/qis/graph.py

@@ -141,16 +141,6 @@ def __init__(
                 'specified in the graph input.',
             )
 
-        calc_num_qudits = 0
-        for q1, q2 in self._edges:
-            calc_num_qudits = max(calc_num_qudits, max(q1, q2))
-        calc_num_qudits += 1
-
-        if num_qudits is not None and calc_num_qudits > num_qudits:
-            raise ValueError(
-                'Edges between invalid qudits or num_qudits too small.',
-            )
-
         if isinstance(graph, CouplingGraph):
             self.num_qudits: int = graph.num_qudits
             self._edges: set[tuple[int, int]] = graph._edges
@@ -161,6 +151,16 @@ def __init__(
             self.default_remote_weight: float = graph.default_remote_weight
             return
 
+        calc_num_qudits = 0
+        for q1, q2 in graph:
+            calc_num_qudits = max(calc_num_qudits, max(q1, q2))
+        calc_num_qudits += 1
+
+        if num_qudits is not None and calc_num_qudits > num_qudits:
+            raise ValueError(
+                'Edges between invalid qudits or num_qudits too small.',
+            )
+
         self.num_qudits = calc_num_qudits if num_qudits is None else num_qudits
         self._edges = {g if g[0] <= g[1] else (g[1], g[0]) for g in graph}
         self._remote_edges = {
@@ -191,6 +191,30 @@ def __init__(
             self._mat[q1][q2] = weight
             self._mat[q2][q1] = weight
 
+    def get_qpu_to_qudit_map(self) -> list[list[int]]:
+        """Return a mapping of QPU indices to qudit indices."""
+        if not hasattr(self, '_qpu_to_qudit'):
+            seen = set()
+            self._qpu_to_qudit = []
+            for qudit in range(self.num_qudits):
+                if qudit in seen:
+                    continue
+                qpu = []
+                frontier = {qudit}
+                while len(frontier) > 0:
+                    node = frontier.pop()
+                    qpu.append(node)
+                    seen.add(node)
+                    for neighbor in self._adj[node]:
+                        if (node, neighbor) in self._remote_edges:
+                            continue
+                        if (neighbor, node) in self._remote_edges:
+                            continue
+                        if neighbor not in seen:
+                            frontier.add(neighbor)
+                self._qpu_to_qudit.append(qpu)
+        return self._qpu_to_qudit
+
     def is_distributed(self) -> bool:
         """Return true if the graph represents multiple connected QPUs."""
         return len(self._remote_edges) > 0
@@ -207,45 +231,14 @@ def get_individual_qpu_graphs(self) -> list[CouplingGraph]:
         qpu_to_qudit = self.get_qpu_to_qudit_map()
         return [self.get_subgraph(qpu) for qpu in qpu_to_qudit]
 
-    def get_qpu_to_qudit_map(self) -> list[list[int]]:
-        """Return a mapping of QPU indices to qudit indices."""
-        # TODO: Cache this?
-        seen = set()
-        qpus = []
-        for qudit in range(self.num_qudits):
-            if qudit in seen:
-                continue
-            qpu = []
-            frontier = {qudit}
-            while len(frontier) > 0:
-                node = frontier.pop()
-                qpu.append(node)
-                seen.add(node)
-                for neighbor in self._adj[node]:
-                    if (node, neighbor) in self._remote_edges:
-                        continue
-                    if (neighbor, node) in self._remote_edges:
-                        continue
-                    if neighbor not in seen:
-                        frontier.add(neighbor)
-            qpus.append(qpu)
-
-        if len(seen) != self.num_qudits:
-            raise RuntimeError(
-                'Graph is not fully connected and pathological'
-                ' for distributed subroutines.',
-            )
-
-        return qpus
-
-    def get_qudit_to_qpu_map(self) -> dict[int, int]:
+    def get_qudit_to_qpu_map(self) -> list[int]:
         """Return a mapping of qudit indices to QPU indices."""
         qpu_to_qudit = self.get_qpu_to_qudit_map()
         qudit_to_qpu = {}
         for qpu, qudits in enumerate(qpu_to_qudit):
             for qudit in qudits:
                 qudit_to_qpu[qudit] = qpu
-        return qudit_to_qpu
+        return list(qudit_to_qpu.values())
 
     def get_qpu_connectivity(self) -> list[set[int]]:
         """Return the adjacency list of the QPUs."""

tests/qis/test_graph.py

@@ -1,9 +1,217 @@
 """This module tests the CouplingGraph class."""
 from __future__ import annotations
 
+from typing import Any
+
 import pytest
 
 from bqskit.qis.graph import CouplingGraph
+from bqskit.qis.graph import CouplingGraphLike
+
+
+def test_coupling_graph_init_valid() -> None:
+    # Test with valid inputs
+    graph = {(0, 1), (1, 2), (2, 3)}
+    num_qudits = 4
+    remote_edges = [(1, 2)]
+    default_weight = 1.0
+    default_remote_weight = 10.0
+    edge_weights_overrides = {(1, 2): 0.5}
+
+    coupling_graph = CouplingGraph(
+        graph,
+        num_qudits,
+        remote_edges,
+        default_weight,
+        default_remote_weight,
+        edge_weights_overrides,
+    )
+
+    assert coupling_graph.num_qudits == num_qudits
+    assert coupling_graph._edges == graph
+    assert coupling_graph._remote_edges == set(remote_edges)
+    assert coupling_graph.default_weight == default_weight
+    assert coupling_graph.default_remote_weight == default_remote_weight
+    assert all(
+        coupling_graph._mat[q1][q2] == weight
+        for (q1, q2), weight in edge_weights_overrides.items()
+    )
+
+
+@pytest.mark.parametrize(
+    'graph, num_qudits, remote_edges, default_weight, default_remote_weight,'
+    ' edge_weights_overrides, expected_exception',
+    [
+        # Invalid graph
+        (None, 4, [], 1.0, 100.0, {}, TypeError),
+        # num_qudits is not an integer
+        ({(0, 1)}, '4', [], 1.0, 100.0, {}, TypeError),
+        # num_qudits is negative
+        ({(0, 1)}, -1, [], 1.0, 100.0, {}, ValueError),
+        # Invalid remote_edges
+        ({(0, 1)}, 4, None, 1.0, 100.0, {}, TypeError),
+        # Remote edge not in graph
+        ({(0, 1)}, 4, [(1, 2)], 1.0, 100.0, {}, ValueError),
+        # Invalid default_weight
+        ({(0, 1)}, 4, [], '1.0', 100.0, {}, TypeError),
+        # Invalid default_remote_weight
+        ({(0, 1)}, 4, [], 1.0, '100.0', {}, TypeError),
+        # Invalid edge_weights_overrides
+        ({(0, 1)}, 4, [], 1.0, 100.0, None, TypeError),
+        # Non-integer value in edge_weights_overrides
+        ({(0, 1)}, 4, [], 1.0, 100.0, {(0, 1): '0.5'}, TypeError),
+        # Edge in edge_weights_overrides not in graph
+        ({(0, 1)}, 4, [], 1.0, 100.0, {(1, 2): 0.5}, ValueError),
+    ],
+)
+def test_coupling_graph_init_invalid(
+    graph: CouplingGraphLike,
+    num_qudits: Any,
+    remote_edges: Any,
+    default_weight: Any,
+    default_remote_weight: Any,
+    edge_weights_overrides: Any,
+    expected_exception: Exception,
+) -> None:
+    with pytest.raises(expected_exception):
+        CouplingGraph(
+            graph,
+            num_qudits,
+            remote_edges,
+            default_weight,
+            default_remote_weight,
+            edge_weights_overrides,
+        )
+
+
+def test_get_qpu_to_qudit_map_single_qpu() -> None:
+    graph = CouplingGraph([(0, 1), (1, 2), (2, 3)])
+    expected_map = [[0, 1, 2, 3]]
+    assert graph.get_qpu_to_qudit_map() == expected_map
+
+
+def test_get_qpu_to_qudit_map_multiple_qpus() -> None:
+    graph = CouplingGraph([(0, 1), (1, 2), (2, 3)], remote_edges=[(1, 2)])
+    expected_map = [[0, 1], [2, 3]]
+    assert graph.get_qpu_to_qudit_map() == expected_map
+
+
+def test_get_qpu_to_qudit_map_disconnected() -> None:
+    graph = CouplingGraph([(0, 1), (1, 2), (3, 4)], remote_edges=[(1, 2)])
+    expected_map = [[0, 1], [2], [3, 4]]
+    assert graph.get_qpu_to_qudit_map() == expected_map
+
+
+def test_get_qpu_to_qudit_map_empty_graph() -> None:
+    graph = CouplingGraph([])
+    expected_map = [[0]]
+    assert graph.get_qpu_to_qudit_map() == expected_map
+
+
+def test_get_qpu_to_qudit_map_complex_topology() -> None:
+    graph = CouplingGraph(
+        [(0, 1), (1, 2), (0, 2), (2, 5), (3, 4), (4, 5), (3, 5)],
+        remote_edges=[(2, 5)],
+    )
+    expected_map = [[0, 1, 2], [3, 4, 5]]
+    assert graph.get_qpu_to_qudit_map() == expected_map
+
+
+def test_get_qudit_to_qpu_map_three_qpu() -> None:
+    graph = CouplingGraph(
+        [(0, 1), (1, 2), (2, 3), (3, 4), (4, 5), (5, 6), (6, 7)],
+        remote_edges=[(2, 3), (5, 6)],
+    )
+    expected_map = [[0, 1, 2], [3, 4, 5], [6, 7]]
+    assert graph.get_qpu_to_qudit_map() == expected_map
+
+
+def test_is_distributed() -> None:
+    graph = CouplingGraph([(0, 1), (1, 2), (2, 3)])
+    assert not graph.is_distributed()
+
+    graph = CouplingGraph([(0, 1), (1, 2), (2, 3)], remote_edges=[(1, 2)])
+    assert graph.is_distributed()
+
+    graph = CouplingGraph([(0, 1), (1, 2), (2, 3)], remote_edges=[(1, 2)])
+    assert graph.is_distributed()
+
+    graph = CouplingGraph(
+        [(0, 1), (1, 2), (2, 3)],
+        remote_edges=[(1, 2), (2, 3)],
+    )
+    assert graph.is_distributed()
+
+
+def test_qpu_count() -> None:
+    graph = CouplingGraph([(0, 1), (1, 2), (2, 3)])
+    assert graph.qpu_count() == 1
+
+    graph = CouplingGraph([(0, 1), (1, 2), (2, 3)], remote_edges=[(1, 2)])
+    assert graph.qpu_count() == 2
+
+    graph = CouplingGraph(
+        [(0, 1), (1, 2), (2, 3)],
+        remote_edges=[(1, 2), (2, 3)],
+    )
+    assert graph.qpu_count() == 3
+
+    graph = CouplingGraph([])
+    assert graph.qpu_count() == 1
+
+
+def test_get_individual_qpu_graphs() -> None:
+    graph = CouplingGraph([(0, 1), (1, 2), (2, 3)])
+    qpus = graph.get_individual_qpu_graphs()
+    assert len(qpus) == 1
+    assert qpus[0] == graph
+
+    graph = CouplingGraph([(0, 1), (1, 2), (2, 3)], remote_edges=[(1, 2)])
+    qpus = graph.get_individual_qpu_graphs()
+    assert len(qpus) == 2
+    assert qpus[0] == CouplingGraph([(0, 1)])
+    assert qpus[1] == CouplingGraph([(0, 1)])
+
+    graph = CouplingGraph(
+        [(0, 1), (1, 2), (2, 3)],
+        remote_edges=[(1, 2), (2, 3)],
+    )
+    qpus = graph.get_individual_qpu_graphs()
+    assert len(qpus) == 3
+    assert qpus[0] == CouplingGraph([(0, 1)])
+    assert qpus[1] == CouplingGraph([])
+    assert qpus[2] == CouplingGraph([])
+
+
+def test_get_qudit_to_qpu_map() -> None:
+    graph = CouplingGraph([(0, 1), (1, 2), (2, 3)])
+    assert graph.get_qudit_to_qpu_map() == [0, 0, 0, 0]
+
+    graph = CouplingGraph([(0, 1), (1, 2), (2, 3)], remote_edges=[(1, 2)])
+    assert graph.get_qudit_to_qpu_map() == [0, 0, 1, 1]
+
+    graph = CouplingGraph(
+        [(0, 1), (1, 2), (2, 3)],
+        remote_edges=[(1, 2), (2, 3)],
+    )
+    assert graph.get_qudit_to_qpu_map() == [0, 0, 1, 2]
+
+    graph = CouplingGraph([])
+    assert graph.get_qudit_to_qpu_map() == [0]
+
+
+def test_get_qpu_connectivity() -> None:
+    graph = CouplingGraph([(0, 1), (1, 2), (2, 3)])
+    assert graph.get_qpu_connectivity() == [set()]
+
+    graph = CouplingGraph([(0, 1), (1, 2), (2, 3)], remote_edges=[(1, 2)])
+    assert graph.get_qpu_connectivity() == [{1}, {0}]
+
+    graph = CouplingGraph(
+        [(0, 1), (1, 2), (2, 3)],
+        remote_edges=[(1, 2), (2, 3)],
+    )
+    assert graph.get_qpu_connectivity() == [{1}, {0, 2}, {1}]
 
 
 class TestGraphGetSubgraphsOfSize: