Repository: netket/netket

Issue:
Make graph.edge a property and not a method
I'd also change nodes to be a property, too.

@femtobit Will you open a PR to do this if you have some time?


Discussion:
> Will you open a PR to do this if you have some time?

Yes, I'll do that.
Though on a second thought, maybe we can keep the current interface, and maybe just change how the color kwarg works
I don't know. We are 1) not happy with the current interface and 2) the current interface is different from  NetKet v2, while the proposed change brings us back to the original interface. Seems like a good choice to me.
Ah, for a moment I thought it was a change.
That’s why I thought about not changing it 

Though arguably most people are playing around with one of the development, never released versions, I think... 



Resolving patch:
diff --git a/netket/graph/abstract_graph.py b/netket/graph/abstract_graph.py
--- a/netket/graph/abstract_graph.py
+++ b/netket/graph/abstract_graph.py
@@ -13,7 +13,7 @@
 # limitations under the License.
 
 import abc
-from typing import List, Iterator, Sequence, Tuple, Union
+from typing import List, Iterator, Sequence, Tuple, Union, Optional
 
 
 Edge = Tuple[int, int]
@@ -35,8 +35,27 @@ def is_bipartite(self) -> bool:
         raise NotImplementedError
 
     @abc.abstractmethod
-    def edges(self, color: Union[bool, int] = False) -> EdgeSequence:
-        r"""Iterator over the edges of the graph. Optionally filter by edge color."""
+    def edges(
+        self, return_color: bool = False, filter_color: Optional[int] = None
+    ) -> EdgeSequence:
+        r"""Returns the sequence of edges of the graph.
+
+        Arguments:
+            return_color: If :code:`True`, return edges with added color information.
+            filter_color: If set, return only edges of the given color.
+
+        Returns:
+            A sequence of edges as tuples :code:`(i, j)` or, if `return_color` is
+            passed, a sequence of tuples :code:`(i, j, c)` with :code:`c` indicating
+            the color of the respective edge.
+        """
+        raise NotImplementedError
+
+    @property
+    @abc.abstractmethod
+    def edge_colors(self) -> Sequence[int]:
+        r"""Sequence of edge colors, in the order of the edges returned by
+        :code:`self.edges`."""
         raise NotImplementedError
 
     @abc.abstractmethod
diff --git a/netket/graph/graph.py b/netket/graph/graph.py
--- a/netket/graph/graph.py
+++ b/netket/graph/graph.py
@@ -17,6 +17,7 @@
 import numpy as np
 import igraph
 
+from netket.utils.deprecation import warn_deprecation
 from netket.utils.group import Permutation, PermutationGroup
 from .abstract_graph import AbstractGraph, Edge, ColoredEdge, EdgeSequence
 
@@ -66,7 +67,7 @@ def _clean_edges(edges):
     # Conversion
     # ------------------------------------------------------------------------
     @classmethod
-    def from_igraph(cls, graph: igraph.Graph):
+    def from_igraph(cls, graph: igraph.Graph) -> "Graph":
         """
         Creates a new Graph instance from an igraph.Graph instance.
         """
@@ -79,7 +80,7 @@ def from_igraph(cls, graph: igraph.Graph):
         if "color" not in self._igraph.edge_attributes():
             self._igraph.es.set_attribute_values("color", [0] * self._igraph.ecount())
         else:
-            if not all(isinstance(c, int) for (_, _, c) in self.edges(color=True)):
+            if not all(isinstance(c, int) for c in self.edge_colors):
                 raise ValueError(
                     "graph has 'color' edge attributes, but not all colors are integers."
                 )
@@ -87,7 +88,7 @@ def from_igraph(cls, graph: igraph.Graph):
         return self
 
     @classmethod
-    def from_networkx(cls, graph):
+    def from_networkx(cls, graph) -> "Graph":
         """
         Creates a new Graph instance from a networkx graph.
         """
@@ -129,19 +130,42 @@ def n_edges(self) -> int:
     def nodes(self) -> Sequence[int]:
         return range(self._igraph.vcount())
 
-    def edges(self, color: Union[bool, int] = False) -> EdgeSequence:
-        edges = self._igraph.get_edgelist()
-        if color is False:
-            return edges
-        colors = self._edge_colors
-        if color is True:
-            return [(*e, c) for e, c in zip(edges, colors)]
+    def edges(
+        self,
+        color=None,
+        *,
+        return_color: bool = False,
+        filter_color: Optional[int] = None,
+    ) -> EdgeSequence:
+        if color is not None:
+            warn_deprecation(
+                "The color option has been split into return_color and filter_color."
+            )
+            # need to check for bool first, because bool is a subclass of int
+            if isinstance(color, bool):
+                return_color = color
+            elif isinstance(color, int):
+                filter_color = color
+            else:
+                raise TypeError("Incorrect type for 'color'")
+
+        if not return_color and filter_color is None:
+            return self._igraph.get_edgelist()
+
+        edges_with_color = zip(self._igraph.get_edgelist(), self.edge_colors)
+        if filter_color is not None:
+            edges_with_color = filter(lambda x: x[1] == filter_color, edges_with_color)
+        if return_color:
+            return [(*e, c) for (e, c) in edges_with_color]
         else:
-            return [e for e, c in zip(edges, colors) if c == color]
+            return [e for (e, _) in edges_with_color]
 
     @property
-    def _edge_colors(self):
-        return self._igraph.es.get_attribute_values("color")
+    def edge_colors(self):
+        if self.n_edges > 0:
+            return self._igraph.es.get_attribute_values("color")
+        else:
+            return []
 
     # Graph algorithms
     # ------------------------------------------------------------------------
@@ -152,7 +176,7 @@ def _compute_automorphisms(self):
         """
         Compute the graph autmorphisms of this graph.
         """
-        colors = self._edge_colors
+        colors = self.edge_colors
         result = self._igraph.get_isomorphisms_vf2(
             edge_color1=colors, edge_color2=colors
         )
diff --git a/netket/operator/_graph_operator.py b/netket/operator/_graph_operator.py
--- a/netket/operator/_graph_operator.py
+++ b/netket/operator/_graph_operator.py
@@ -106,7 +106,7 @@ def __init__(
 
             if len(bond_ops) > 0:
                 #  Use edge_colors to populate operators
-                for (u, v, color) in graph.edges(color=True):
+                for (u, v, color) in graph.edges(return_color=True):
                     edge = u, v
                     for c, bond_color in enumerate(bond_ops_colors):
                         if bond_color == color:


Test patch:
diff --git a/test/graph/test_graph.py b/test/graph/test_graph.py
--- a/test/graph/test_graph.py
+++ b/test/graph/test_graph.py
@@ -401,6 +401,7 @@ def test_computes_distances():
 
 def test_lattice_is_bipartite():
     for graph in graphs:
+        print(graph)
         g = nx.Graph()
         for edge in graph.edges():
             g.add_edge(edge[0], edge[1])
@@ -441,17 +442,17 @@ def test_adjacency_list():
 #    count = lambda it: sum(1 for _ in it)
 #
 #    g = Grid([4, 4], pbc=True, color_edges=True)
-#    assert count(g.edges(color=0)) == 16
-#    assert count(g.edges(color=1)) == 16
+#    assert count(g.edges(filter_color=0)) == 16
+#    assert count(g.edges(filter_color=1)) == 16
 #    assert g.n_edges == 32
 #
 #    g = Grid([4, 2], pbc=True, color_edges=True)
-#    assert count(g.edges(color=0)) == 8
-#    assert count(g.edges(color=1)) == 4
+#    assert count(g.edges(filter_color=0)) == 8
+#    assert count(g.edges(filter_color=1)) == 4
 #
 #    g = Grid([4, 2], pbc=False, color_edges=True)
-#    assert count(g.edges(color=0)) == 6
-#    assert count(g.edges(color=1)) == 4
+#    assert count(g.edges(filter_color=0)) == 6
+#    assert count(g.edges(filter_color=1)) == 4
 #
 #    with pytest.raises(ValueError, match="Directions with length <= 2 cannot have PBC"):
 #        g = Grid([2, 4], pbc=[True, True])
@@ -651,11 +652,11 @@ def test_edge_color_accessor():
     edges = [(0, 1, 0), (0, 3, 1), (1, 2, 1), (2, 3, 0)]
     g = Graph(edges=edges)
 
-    assert edges == sorted(g.edges(color=True))
+    assert edges == sorted(g.edges(return_color=True))
 
     g = Hypercube(4, 1)
 
-    assert [(i, j, 0) for (i, j, _) in edges] == sorted(g.edges(color=True))
+    assert [(i, j, 0) for (i, j, _) in edges] == sorted(g.edges(return_color=True))
 
 
 def test_union():
@@ -673,13 +674,13 @@ def test_graph_conversions():
     g = Graph.from_igraph(igraph)
     assert g.n_nodes == igraph.vcount()
     assert g.edges() == igraph.get_edgelist()
-    assert all(c == 0 for c in g._edge_colors)
+    assert all(c == 0 for c in g.edge_colors)
 
     nxg = nx.star_graph(5)
     g = Graph.from_networkx(nxg)
     assert g.n_nodes == igraph.vcount()
     assert g.edges() == igraph.get_edgelist()
-    assert all(c == 0 for c in g._edge_colors)
+    assert all(c == 0 for c in g.edge_colors)
 
     igraph = ig.Graph()
     igraph.add_vertices(3)
@@ -701,5 +702,17 @@ def test_graph_conversions():
         },
     )
     g = Graph.from_igraph(igraph)
-    assert g.edges(color=0) == [(0, 1)]
-    assert g.edges(color=1) == [(1, 2)]
+    assert g.edges(filter_color=0) == [(0, 1)]
+    assert g.edges(filter_color=1) == [(1, 2)]
+    assert g.edges(filter_color=0, return_color=True) == [(0, 1, 0)]
+    assert g.edges(filter_color=1, return_color=True) == [(1, 2, 1)]
+
+    with pytest.warns(FutureWarning):
+        assert g.edges(color=0) == g.edges(filter_color=0)
+    with pytest.warns(FutureWarning):
+        assert g.edges(color=True) == g.edges(return_color=True)
+
+
+def test_edge_colors():
+    for g in graphs:
+        assert all(isinstance(c, int) for c in g.edge_colors)