Implement edge list-colouring

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "vizing"
-version = "0.1.0.9000"
+version = "0.1.0.9001"
 description = "Constraint-based list-colouring in Python."
 authors = ["Matthew Henderson <[email protected]>"]
 readme = "README.md"

src/vizing/__init__.py

@@ -27,3 +27,29 @@ def node_list_colouring_solution(G):
     G.nodes[node]['colour'] = S[node]
   return(G)
 
+def edge_list_colouring_problem(G):
+  """
+  Return a constraint problem representing an edge list-colouring instance.
+
+  :param G: A graph with lists of permissible colours assigned to edges.
+  :return An edge list-colouring constraint problem.
+  """
+  P = constraint.Problem()
+  for edge in G.edges():
+    P.addVariable(edge, G.edges[edge]['permissible'])
+  for node in G.nodes():
+    P.addConstraint(constraint.AllDifferentConstraint(), [tuple(sorted(x)) for x in G.edges(node)])
+  return(P)
+
+def edge_list_colouring_solution(G):
+  """
+  Return an edge-coloured graph.
+
+  :param G: A graph with lists of permissible colours assigned to edges.
+  :return A properly edge list-coloured graph.
+  """
+  P = edge_list_colouring_problem(G)
+  S = P.getSolution()
+  for edge in G.edges():
+    G.edges()[edge]['colour'] = S[edge]
+  return(G)

tests/test_edge-list-colouring.py

@@ -0,0 +1,20 @@
+import networkx as nx
+import vizing as vz
+
+def test_edge_list_colouring_problem():
+  G = nx.complete_graph(3)
+  permissible_dict_edge = {(0, 1): [0, 1], (0, 2): [1, 2], (1, 2): [2, 4]}
+  nx.set_edge_attributes(G, permissible_dict_edge, "permissible")
+  nx.set_edge_attributes(G, None, "colour")
+  P = vz.edge_list_colouring_problem(G)
+  assert P._variables == {(0, 1): [0, 1], (0, 2): [1, 2], (1, 2): [2, 4]}
+
+def test_edge_list_colouring_solution():
+  G = nx.complete_graph(3)
+  permissible_dict_edge = {(0, 1): [0, 1], (0, 2): [1, 2], (1, 2): [2, 4]}
+  nx.set_edge_attributes(G, permissible_dict_edge, "permissible")
+  nx.set_edge_attributes(G, None, "colour")
+  G = vz.edge_list_colouring_solution(G)
+  assert G.edges()[(0, 1)]['colour'] == 1
+  assert G.edges()[(0, 2)]['colour'] == 2
+  assert G.edges()[(1, 2)]['colour'] == 4