diff --git a/qoolqit/graphs/base_graph.py b/qoolqit/graphs/base_graph.py index 8e5f66455..022916ee5 100644 --- a/qoolqit/graphs/base_graph.py +++ b/qoolqit/graphs/base_graph.py @@ -21,9 +21,12 @@ class BaseGraph(nx.Graph): - """Base graph class, directly inheriting from the NetworkX Graph. + """Base graph class, directly inheriting from `networkx.Graph`. - On top of the standard networkx.Graph functionalities, adds alternative + Represents a simple graph (undirected and without self-loops), + with optional node coordinates and node/edge weights. + + On top of the standard networkx.Graph functionality, adds alternative constructors, node coordinates and weights as first-class attributes, distance and Rydberg-interaction calculations, unit-disk graph analysis, and plotting. @@ -40,7 +43,7 @@ class BaseGraph(nx.Graph): `max_distance`, `rescale_coords`. Unit-disk analysis: `is_ud_graph`, `ud_radius_range`, `ud_edges`, `set_ud_edges`. - Rydberg-analog interactions: `interactions`, `interaction_matrix`. + Rydberg analog model utils: `interactions`, `interaction_matrix`. Plotting: `draw`. """ @@ -275,24 +278,40 @@ def edge_weights(self) -> dict: return {(u, v): w for u, v, w in self.edges(data="weight")} @edge_weights.setter - def edge_weights(self, weights: list | dict) -> None: + def edge_weights(self, weights: dict) -> None: """Set the dictionary of edge weights. + Each edge may be keyed as (u, v) or (v, u) since the graph is undirected. + + Weights must be specified for all edges, otherwise a ValueError is raised. + To partially update edge weights, use the attribute-like access pattern: + ```python + graph.edges[0,1]["weight"] = 0.5 # Update weight of edge (0,1) + ``` + Arguments: - weights: list or dictionary of weights. + weights: a dictionary of edge weights. + + Raises: + ValueError: if the set of edges in the given dictionary does not match + the graph's edges. + + Example: + >>> graph = BaseGraph([(0, 1), (1, 2), (2, 0)]) + >>> graph.edge_weights = {(0, 1): 0.3, (1, 2): 0.4, (2, 0): 0.5} + >>> graph.edge_weights + {(0, 1): 0.3, (0, 2): 0.5, (1, 2): 0.4} """ - if isinstance(weights, list): - if len(weights) != self.number_of_edges(): - raise ValueError("Size of the weights list does not match the number of nodes.") - weights_dict = {i: w for i, w in zip(self.sorted_edges, weights)} - elif isinstance(weights, dict): - edges = set(weights.keys()) - if set(self.sorted_edges) != edges: - raise ValueError( - "Set of edges in the given dictionary does not match the graph ordered edges." - ) - weights_dict = weights - nx.set_edge_attributes(self, weights_dict, "weight") + # frozenset canonicalizes each edge regardless of orientation, e.g. + # frozenset((u, v)) == frozenset((v, u)). + given = {frozenset(e) for e in weights} + expected = {frozenset(e) for e in self.edges} + if given != expected: + raise ValueError( + "Set of edges in the given dictionary does not match the graph's edges." + ) + for (u, v), w in weights.items(): + self.edges[u, v]["weight"] = w @property def coords(self) -> dict: diff --git a/tests/test_graphs/test_base_graph.py b/tests/test_graphs/test_base_graph.py index 25db72735..e8eb27fd8 100644 --- a/tests/test_graphs/test_base_graph.py +++ b/tests/test_graphs/test_base_graph.py @@ -148,7 +148,7 @@ def test_edge_weights_update_missing_edge() -> None: graph = BaseGraph([(0, 1), (1, 2), (2, 0)]) with pytest.raises( ValueError, - match="Set of edges in the given dictionary does not match the graph ordered edges.", + match="Set of edges in the given dictionary does not match the graph's edges.", ): graph.edge_weights = {(0, 1): 0.3, (1, 2): 0.4} @@ -157,7 +157,7 @@ def test_edge_weights_update_extra_edge() -> None: graph = BaseGraph([("a", "b"), ("b", "c"), ("c", "a")]) with pytest.raises( ValueError, - match="Set of edges in the given dictionary does not match the graph ordered edges.", + match="Set of edges in the given dictionary does not match the graph's edges.", ): graph.edge_weights = { ("a", "b"): 0.3, @@ -371,6 +371,17 @@ def test_to_matrix_roundtrip(n_nodes: int, seed: int) -> None: np.testing.assert_allclose(graph.to_matrix(), matrix, atol=1e-8) +def test_edge_weights_accepts_either_orientation() -> None: + # Regression test for #447: setting edge_weights should not depend on + # knowing which of (u, v) / (v, u) the graph happens to report internally. + graph = BaseGraph() + graph.add_edges_from([(0, 1), (1, 2), (2, 0)]) + + graph.edge_weights = {(0, 1): 0.3, (1, 2): 0.4, (2, 0): 0.5} + + assert graph.edge_weights == {(0, 1): 0.3, (0, 2): 0.5, (1, 2): 0.4} + + def test_to_matrix_custom_node_labels_and_none_weights() -> None: # Ensure `to_matrix()` respects `self.nodes` ordering and handles None weights. # Also ensures it works with non-0..N-1 node labels.