fix edge weights does not reflect edges added after construction

qoolqit/graphs/base_graph.py

@@ -44,26 +44,6 @@ class BaseGraph(nx.Graph):
         Plotting: `draw`.
     """
 
-    def __init__(self, edges: Iterable = []) -> None:
-        """
-        Default constructor for the BaseGraph.
-
-        Arguments:
-            edges: set of edge tuples (i, j)
-        """
-        if edges and not isinstance(edges, Iterable):
-            raise TypeError("Input is not a valid edge list.")
-
-        super().__init__()
-        self.add_edges_from(edges)
-        self._coords = {i: None for i in self.nodes}
-        self._reset_dicts()
-
-    def _reset_dicts(self) -> None:
-        """Reset the default weight dictionaries."""
-        self._node_weights = {n: None for n in self.nodes}
-        self._edge_weights = {e: None for e in self.sorted_edges}
-
     @classmethod
     def from_nodes(cls, nodes: Iterable) -> BaseGraph:
         """Construct a base graph from a set of nodes.
@@ -73,8 +53,6 @@ def from_nodes(cls, nodes: Iterable) -> BaseGraph:
         """
         graph = cls()
         graph.add_nodes_from(nodes)
-        graph._coords = {i: None for i in graph.nodes}
-        graph._reset_dicts()
         return graph
 
     @classmethod
@@ -85,15 +63,11 @@ def from_coordinates(cls, coords: list | dict) -> BaseGraph:
             coords: list or dictionary of coordinate pairs.
         """
         if isinstance(coords, list):
-            nodes = list(range(len(coords)))
-            coords_dict = {i: pos for i, pos in enumerate(coords)}
+            coords_tuple = ((i, {"pos": pos}) for i, pos in enumerate(coords))
         elif isinstance(coords, dict):
-            nodes = list(coords.keys())
-            coords_dict = coords
-        graph = cls.from_nodes(nodes)
-        graph._coords = coords_dict
-        graph._reset_dicts()
-        return graph
+            coords_tuple = ((key, {"pos": pos}) for key, pos in coords.items())
+
+        return cls.from_nodes(coords_tuple)
 
     @classmethod
     def from_nx(cls, g: nx.Graph) -> BaseGraph:
@@ -108,9 +82,9 @@ def from_nx(cls, g: nx.Graph) -> BaseGraph:
             weight: represents the edge weight. Must be a real number.
 
         Returns an instance of the class with following attributes:
-            - _node_weights : dict[node, float or None]
-            - _edge_weights : dict[(u,v), float or None]
-            - _coords       : dict[node, (float,float) or None]
+            - node_weights : dict[node, float or None]
+            - edge_weights : dict[(u,v), float or None]
+            - coords       : dict[node, (float,float) or None]
         """
         if not isinstance(g, nx.Graph):
             raise TypeError("Input must be a networkx.Graph instance.")
@@ -165,15 +139,7 @@ def from_nx(cls, g: nx.Graph) -> BaseGraph:
                         f"got {type(weight)} instead."
                     )
 
-        # build the instance of the graph
-        graph = cls()
-        graph.add_nodes_from(g.nodes)
-        graph.add_edges_from(g.edges)
-        graph._node_weights = nx.get_node_attributes(g, "weight", default=None)
-        graph._coords = nx.get_node_attributes(g, "pos", default=None)
-        graph._edge_weights = nx.get_edge_attributes(g, "weight", default=None)
-
-        return graph
+        return cls(g)
 
     @classmethod
     def from_matrix(cls, data: npt.NDArray[np.float64]) -> BaseGraph:
@@ -259,38 +225,29 @@ def all_node_pairs(self) -> set:
 
     @property
     def has_coords(self) -> bool:
-        """Check if the graph has coordinates.
-
-        Requires all nodes to have coordinates.
-        """
-        is_any_coord_none = any(value is None for value in self._coords.values())
-        return not (is_any_coord_none or len(self._coords) == 0)
-
-    @property
-    def has_edges(self) -> bool:
-        """Check if the graph has edges."""
-        return len(self.edges) > 0
+        """Check if the graph has coordinates on all nodes."""
+        return self.number_of_nodes() > 0 and all(
+            pos is not None for _, pos in self.nodes(data="pos")
+        )
 
     @property
     def has_node_weights(self) -> bool:
-        """Check if the graph has node weights.
-
-        Requires all nodes to have a weight.
-        """
-        return not ((None in self._node_weights.values()) or len(self._node_weights) == 0)
+        """Check if the graph has node weights on all nodes."""
+        return self.number_of_nodes() > 0 and all(
+            w is not None for _, w in self.nodes(data="weight")
+        )
 
     @property
     def has_edge_weights(self) -> bool:
-        """Check if the graph has edge weights.
-
-        Requires all edges to have a weight.
-        """
-        return not ((None in self._edge_weights.values()) or len(self._edge_weights) == 0)
+        """Check if the graph has edge weights on all edges."""
+        return self.number_of_edges() > 0 and all(
+            w is not None for _, _, w in self.edges(data="weight")
+        )
 
     @property
     def node_weights(self) -> dict:
         """Return the dictionary of node weights."""
-        return self._node_weights
+        return {n: w for n, w in self.nodes(data="weight")}
 
     @node_weights.setter
     def node_weights(self, weights: list | dict) -> None:
@@ -310,12 +267,12 @@ def node_weights(self, weights: list | dict) -> None:
                     "Set of nodes in the given dictionary does not match the graph nodes."
                 )
             weights_dict = weights
-        self._node_weights = weights_dict
+        nx.set_node_attributes(self, weights_dict, "weight")
 
     @property
     def edge_weights(self) -> dict:
         """Return the dictionary of edge weights."""
-        return self._edge_weights
+        return {(u, v): w for u, v, w in self.edges(data="weight")}
 
     @edge_weights.setter
     def edge_weights(self, weights: list | dict) -> None:
@@ -335,12 +292,12 @@ def edge_weights(self, weights: list | dict) -> None:
                     "Set of edges in the given dictionary does not match the graph ordered edges."
                 )
             weights_dict = weights
-        self._edge_weights = weights_dict
+        nx.set_edge_attributes(self, weights_dict, "weight")
 
     @property
     def coords(self) -> dict:
-        """Return the dictionary of node coordinates."""
-        return self._coords
+        """Return a dictionary of node coordinates."""
+        return dict(self.nodes(data="pos"))
 
     @coords.setter
     def coords(self, coords: list | dict) -> None:
@@ -358,9 +315,7 @@ def coords(self, coords: list | dict) -> None:
                     "Set of nodes in the given dictionary does not match the graph nodes."
                 )
             coords_dict = coords
-        self._coords = coords_dict
-
-    # methods
+        nx.set_node_attributes(self, coords_dict, "pos")
 
     def distances(self, edge_list: Iterable | None = None) -> dict:
         """Returns a dictionary of distances for a given set of edges.
@@ -495,9 +450,9 @@ def rescale_coords(
             if (len(args) > 0) or (scaling is None and spacing is None):
                 raise TypeError(msg)
             if scaling is None and spacing is not None:
-                self._coords = space_coords(self._coords, spacing)
+                self.coords = space_coords(self.coords, spacing)
             elif spacing is None and scaling is not None:
-                self._coords = scale_coords(self._coords, scaling)
+                self.coords = scale_coords(self.coords, scaling)
             else:
                 raise TypeError(msg)
         else:

qoolqit/graphs/data_graph.py

@@ -36,15 +36,6 @@ class DataGraph(BaseGraph):
         PyTorch Geometric conversion: `from_pyg`, `to_pyg`.
     """
 
-    def __init__(self, edges: Iterable = []) -> None:
-        """
-        Default constructor for the BaseGraph.
-
-        Arguments:
-            edges: set of edge tuples (i, j)
-        """
-        super().__init__(edges)
-
     @classmethod
     def line(cls, n: int, spacing: float = 1.0) -> DataGraph:
         """Constructs a line graph, with the respective coordinates.
@@ -57,7 +48,6 @@ def line(cls, n: int, spacing: float = 1.0) -> DataGraph:
         graph = cls.from_coordinates(coords)
         edges = [(i, i + 1) for i in range(0, n - 1)]
         graph.add_edges_from(edges)
-        graph._reset_dicts()
         return graph
 
     @classmethod
@@ -84,12 +74,11 @@ def circle(
         edges = [(i, i + 1) for i in range(n - 1)] + [(n - 1, 0)]
         graph = cls.from_coordinates(coords)
         graph.add_edges_from(edges)
-        graph._reset_dicts()
         return graph
 
     @classmethod
     def random_er(cls, n: int, p: float, seed: int | None = None) -> DataGraph:
-        """Constructs an Erdős–Rényi random graph.
+        """Constructs an Erdős-Rényi random graph.
 
         Arguments:
             n: number of nodes.
@@ -99,7 +88,6 @@ def random_er(cls, n: int, p: float, seed: int | None = None) -> DataGraph:
         base_graph = nx.erdos_renyi_graph(n, p, seed)
         graph = DataGraph.from_nodes(list(base_graph.nodes))
         graph.add_edges_from(base_graph.edges)
-        graph._reset_dicts()
         return graph
 
     @classmethod
@@ -125,7 +113,6 @@ def triangular(
 
         graph = cls.from_coordinates(final_pos)
         graph.add_edges_from(G.edges)
-        graph._reset_dicts()
         return graph
 
     @classmethod
@@ -151,7 +138,6 @@ def hexagonal(
 
         graph = cls.from_coordinates(final_pos)
         graph.add_edges_from(G.edges)
-        graph._reset_dicts()
         return graph
 
     @classmethod
@@ -212,7 +198,6 @@ def heavy_hexagonal(
 
         graph = cls.from_coordinates(final_coords)
         graph.add_edges_from(final_edges)
-        graph._reset_dicts()
         return graph
 
     @classmethod
@@ -236,7 +221,6 @@ def square(
 
         graph = DataGraph.from_coordinates(final_coords)
         graph.add_edges_from(G.edges)
-        graph._reset_dicts()
         return graph
 
     @classmethod
@@ -264,7 +248,6 @@ def random_ud(
         graph = cls.from_coordinates(coords)
         edges = graph.ud_edges(radius)
         graph.add_edges_from(edges)
-        graph._reset_dicts()
         return graph
 
     @classmethod
@@ -283,13 +266,13 @@ def from_pyg(
 
         **Default attributes copied (if present on** ``data`` **):**
 
-        - Node: ``x``, ``pos`` (``pos`` is also stored in ``_coords``)
+        - Node: ``x``, ``pos`` (``pos`` is also stored in ``coords``)
         - Edge: ``edge_attr``
         - Graph: ``y``
 
         Use ``node_attrs``, ``edge_attrs``, ``graph_attrs`` for extras.
 
-        **QoolQit weights** (``_node_weights``, ``_edge_weights``) are not
+        **QoolQit weights** (``node_weights``, ``edge_weights``) are not
         populated automatically — use the explicit parameters:
 
         - ``node_weights_attr``: real-valued tensor of shape ``(N,)`` or
@@ -309,7 +292,7 @@ def from_pyg(
             edge_weights_attr: Data attribute to use as edge weights.
 
         Returns:
-            DataGraph with ``_coords``, ``_node_weights``, ``_edge_weights``
+            DataGraph with ``coords``, ``node_weights``, ``edge_weights``
             populated where applicable.
 
         Raises:
@@ -368,33 +351,22 @@ def from_pyg(
         )
 
         # Build the DataGraph: edges carry their data, nodes carry their data
-        graph = cls(nx_graph.edges(data=True))
-        graph.add_nodes_from(nx_graph.nodes(data=True))
-        graph.graph = nx_graph.graph
-
-        # Re-initialize QoolQit internal dicts for all nodes/edges
-        graph._coords = {n: None for n in graph.nodes}
-        graph._reset_dicts()
-
-        # pos → _coords (stored as list [x, y] by to_networkx)
-        for node, node_data in nx_graph.nodes(data=True):
-            if "pos" in node_data:
-                graph._coords[node] = tuple(node_data["pos"])  # type: ignore[assignment]
+        graph = cls(nx_graph)
 
-        # node_weights_attr → _node_weights
+        # node_weights_attr → node_weights
         if node_tensor is not None:
             for i in range(data.num_nodes):
-                graph._node_weights[i] = node_tensor[i].item()
+                graph.nodes[i]["weight"] = node_tensor[i].item()
 
-        # edge_weights_attr → _edge_weights
+        # edge_weights_attr → edge_weights
         if edge_tensor is not None:
             seen: set = set()
             for idx in range(data.edge_index.shape[1]):
                 u = int(data.edge_index[0, idx].item())
                 v = int(data.edge_index[1, idx].item())
                 key = (min(u, v), max(u, v))
                 if key not in seen:
-                    graph._edge_weights[key] = edge_tensor[idx].item()
+                    graph.edges[key]["weight"] = edge_tensor[idx].item()
                     seen.add(key)
 
         return graph
@@ -420,10 +392,10 @@ def to_pyg(
 
         **QoolQit internal dicts exported when populated:**
 
-        - ``_coords`` → ``data.pos`` (float64, shape ``(N, 2)``)
-        - ``_node_weights`` → ``data.<node_weights_attr>`` (float64, shape
+        - ``coords`` → ``data.pos`` (float64, shape ``(N, 2)``)
+        - ``node_weights`` → ``data.<node_weights_attr>`` (float64, shape
           ``(N,)``). Defaults to ``"weight"``.
-        - ``_edge_weights`` → ``data.<edge_weights_attr>`` (float64, shape
+        - ``edge_weights`` → ``data.<edge_weights_attr>`` (float64, shape
           ``(2*E,)``). Defaults to ``"edge_weight"``.
 
         Arguments:
@@ -478,23 +450,23 @@ def to_pyg(
 
         data = from_networkx(filtered_graph)
 
-        # Export _coords → pos
+        # Export coords → pos
         if self.has_coords:
-            positions = [self._coords[n] for n in sorted(self.nodes())]
+            positions = [self.coords[n] for n in sorted(self.nodes())]
             data.pos = torch.tensor(positions, dtype=torch.float64)
 
-        # Export _node_weights → node_weights_attr
+        # Export node_weights → node_weights_attr
         if self.has_node_weights:
-            weights = [self._node_weights[n] for n in sorted(self.nodes())]
+            weights = [self.node_weights[n] for n in sorted(self.nodes())]
             setattr(data, node_weights_attr, torch.tensor(weights, dtype=torch.float64))
 
-        # Export _edge_weights → edge_weights_attr (one value per directed edge in edge_index)
+        # Export edge_weights → edge_weights_attr (one value per directed edge in edge_index)
         if self.has_edge_weights:
             edge_weights: list[float] = []
             for i in range(data.edge_index.shape[1]):
                 u, v = int(data.edge_index[0, i].item()), int(data.edge_index[1, i].item())
                 edge_key = (min(u, v), max(u, v))
-                edge_weights.append(float(self._edge_weights[edge_key]))  # type: ignore[arg-type]
+                edge_weights.append(float(self.edge_weights[edge_key]))
             setattr(data, edge_weights_attr, torch.tensor(edge_weights, dtype=torch.float64))
 
         return data
@@ -567,4 +539,4 @@ def _validate_weights_attr(
     def set_ud_edges(self, radius: float) -> None:
         """Reset the set of edges to be equal to the set of unit-disk edges."""
         super().set_ud_edges(radius=radius)
-        self._edge_weights = {e: None for e in self.sorted_edges}
+        self.edge_weights = {e: None for e in self.sorted_edges}