diff --git a/qoolqit/graphs/data_graph.py b/qoolqit/graphs/data_graph.py index cd5e122e9..e5492e08a 100644 --- a/qoolqit/graphs/data_graph.py +++ b/qoolqit/graphs/data_graph.py @@ -10,6 +10,7 @@ import networkx as nx import numpy as np +import numpy.typing as npt from .base_graph import BaseGraph from .utils import random_coords @@ -254,7 +255,7 @@ def random_ud( return graph @classmethod - def from_matrix(cls, data: np.ndarray) -> DataGraph: + def from_matrix(cls, data: npt.NDArray[np.float64]) -> DataGraph: """Constructs a graph from a symmetric square matrix. The diagonal values are set as the node weights. For each entry (i, j) @@ -294,6 +295,34 @@ def from_matrix(cls, data: np.ndarray) -> DataGraph: graph.edge_weights = edge_weights return graph + def to_matrix(self) -> npt.NDArray[np.float64]: + """Return the adjacency matrix of this graph. + + The inverse of `from_matrix`. + Nodes are mapped to indices 0, ..., N-1 according to `self.nodes` insertion order. + - Node weights are stored in the diagonal since self-loops are not supported. + Nodes with no weight set (None) are left at 0.0 in the diagonal. + - For each edge (i, j), the entries (i,j) and (j,i) are set to its weight, + or to 1.0 if the edge has no weight set. + + Returns: + Symmetric N x N matrix of dtype float64, where N is the number of nodes. + """ + n_nodes = len(self.nodes) + index = {node: i for i, node in enumerate(self.nodes)} + matrix = np.zeros((n_nodes, n_nodes), dtype=np.float64) + + for node, weight in self.node_weights.items(): + if weight is not None: + i = index[node] + matrix[i, i] = weight + + for (u, v), weight in self.edge_weights.items(): + i, j = index[u], index[v] + matrix[i, j] = matrix[j, i] = weight if weight is not None else 1.0 + + return matrix + @classmethod def from_pyg( cls, diff --git a/tests/test_graphs/test_data_graph.py b/tests/test_graphs/test_data_graph.py index 86fa9536a..1da3b404c 100644 --- a/tests/test_graphs/test_data_graph.py +++ b/tests/test_graphs/test_data_graph.py @@ -69,7 +69,7 @@ def test_datagraph_random_er(n_nodes: int) -> None: @pytest.mark.parametrize("n_nodes", [5, 10, 50]) def test_datagraph_from_matrix(n_nodes: int) -> None: - + np.random.seed(0) data = np.random.rand(n_nodes, n_nodes) with pytest.raises(ValueError): @@ -119,6 +119,79 @@ def test_datagraph_from_matrix(n_nodes: int) -> None: np.testing.assert_allclose(edge_weights, data_edge_weights) +@pytest.mark.parametrize("n_nodes", [5, 10, 50]) +def test_datagraph_to_matrix_unweighted(n_nodes: int) -> None: + graph = DataGraph.random_er(n_nodes, p=0.5, seed=0) + assert not graph.has_node_weights + assert not graph.has_edge_weights + + matrix = graph.to_matrix() + + np.testing.assert_equal(matrix, matrix.T) + np.testing.assert_equal(np.diag(matrix), np.zeros(n_nodes)) + + for i, j in graph.sorted_edges: + assert matrix[i, j] == 1.0 + assert matrix[j, i] == 1.0 + + non_edges = graph.all_node_pairs - graph.sorted_edges + for i, j in non_edges: + assert matrix[i, j] == 0.0 + assert matrix[j, i] == 0.0 + + +@pytest.mark.parametrize("n_nodes", [5, 10, 50]) +def test_datagraph_to_matrix_weighted(n_nodes: int) -> None: + graph = DataGraph.random_er(n_nodes, p=0.5, seed=0) + graph.node_weights = {i: np.random.rand() for i in graph.nodes} + graph.edge_weights = {e: np.random.rand() for e in graph.sorted_edges} + + matrix = graph.to_matrix() + + np.testing.assert_equal(matrix, matrix.T) + np.testing.assert_allclose(np.diag(matrix), list(graph.node_weights.values())) + + for (i, j), weight in graph.edge_weights.items(): + assert matrix[i, j] == weight + assert matrix[j, i] == weight + + +@pytest.mark.parametrize("n_nodes", [3, 7, 21]) +def test_datagraph_to_matrix_roundtrip(n_nodes: int) -> None: + rng = np.random.default_rng(12345) + matrix = rng.normal(0, 1, size=(n_nodes, n_nodes)) + # zero out some elements for testing + rows, cols = rng.integers(n_nodes, size=3), rng.integers(n_nodes, size=3) + matrix[rows, cols] = 0.0 + matrix[cols, rows] = 0.0 + matrix += matrix.T + + graph = DataGraph.from_matrix(matrix) + np.testing.assert_allclose(graph.to_matrix(), matrix, atol=1e-8) + + +def test_datagraph_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. + graph = DataGraph.from_nodes(["b", "a", "c"]) # preserve explicit order + graph.add_edges_from([("b", "a"), ("a", "c")]) + + # Mix of real weights and None + graph.node_weights = {"b": 2.0, "a": None, "c": -3.0} + # TOFIX: order is not guaranteed because DataGraph maintains arbitrarily sort edges + graph.edge_weights = {("a", "b"): None, ("a", "c"): 0.25} + + matrix = graph.to_matrix() + expected = np.array( + [ + [2.0, 1.0, 0.0], + [1.0, 0.0, 0.25], + [0.0, 0.25, -3.0], + ], + ) + np.testing.assert_equal(matrix, expected) + + def test_triangular() -> None: graph = DataGraph.triangular(2, 2, spacing=2.71)