diff --git a/dwave_networkx/algorithms/tsp.py b/dwave_networkx/algorithms/tsp.py index 32fc46de..c453c3e1 100644 --- a/dwave_networkx/algorithms/tsp.py +++ b/dwave_networkx/algorithms/tsp.py @@ -114,7 +114,7 @@ def traveling_salesperson(G, sampler=None, lagrange=None, weight='weight', traveling_salesman = traveling_salesperson -def traveling_salesperson_qubo(G, lagrange=None, weight='weight', missing_edge_penalty='None'): +def traveling_salesperson_qubo(G, lagrange=None, weight='weight', missing_edge_weight=None): """Return the QUBO with ground states corresponding to a minimum TSP route. If :math:`|G|` is the number of nodes in the graph, the resulting qubo will have: @@ -134,8 +134,10 @@ def traveling_salesperson_qubo(G, lagrange=None, weight='weight', missing_edge_p weight : optional (default 'weight') The name of the edge attribute containing the weight. - missing_edge_penalty : number, optional (default 'sum') - For bi-directional graphs, the penalty associated with the back missing edges. Default is none. Alternatives include using the sum all weights. + missing_edge_weight : number, optional (default None) + For bi-directional graphs, the weight given to missing edges. + If None is given (the default), missing edges will be set to + the sum of all weights. Returns ------- @@ -158,15 +160,14 @@ def traveling_salesperson_qubo(G, lagrange=None, weight='weight', missing_edge_p else: lagrange = 2 - missing_edge_weight="" - - # default penalty format is sum - if missing_edge_penalty == "sum": - missing_edge_weight = sum(weight for _, _, weight in G.edges.data('weight', default=0)) + # calculate default missing_edge_weight if required + if missing_edge_weight is None: + # networkx method to calculate sum of all weights + missing_edge_weight = G.size(weight=weight) # some input checking if N in (1, 2): - msg = "graph must be a complete graph" + msg = "graph must have at least 3 nodes or be empty" raise ValueError(msg) # Creating the QUBO @@ -194,10 +195,10 @@ def traveling_salesperson_qubo(G, lagrange=None, weight='weight', missing_edge_p nextpos = (pos + 1) % N # going from u -> v - Q[((u, pos), (v, nextpos))] += G[u][v].get('weight', missing_edge_weight) + Q[((u, pos), (v, nextpos))] += G[u][v].get(weight, missing_edge_weight) # going from v -> u - Q[((v, pos), (u, nextpos))] += G[u][v].get('weight', missing_edge_weight) + Q[((v, pos), (u, nextpos))] += G[u][v].get(weight, missing_edge_weight) return Q diff --git a/tests/test_tsp.py b/tests/test_tsp.py index 048acedd..c0bd6532 100644 --- a/tests/test_tsp.py +++ b/tests/test_tsp.py @@ -97,6 +97,29 @@ def test_start(self): self.assertEqual(route[0], 1) + def test_weighted_complete_digraph(self): + G = nx.DiGraph() + G.add_weighted_edges_from([ + (0, 1, 2), + (1, 0, 1), + (0, 2, 2), + (2, 0, 2), + (0, 3, 1), + (3, 0, 2), + (1, 2, 2), + (2, 1, 1), + (1, 3, 2), + (3, 1, 2), + (2, 3, 2), + (3, 2, 1), + ]) + + route = dnx.traveling_salesperson(G, dimod.ExactSolver(), start=1) + + self.assertEqual(len(route), len(G)) + self.assertEqual(nx.path_weight(G, route, 'weight'), 4) + self.assertListEqual(route, [1, 0, 3, 2]) + class TestTSPQUBO(unittest.TestCase): def test_empty(self): @@ -137,7 +160,7 @@ def test_k3(self): self.assertEqual(ground_count, len(min_routes)) def test_k3_bidirectional(self): - G = nx.Graph() + G = nx.DiGraph() G.add_weighted_edges_from([('a', 'b', 0.5), ('b', 'c', 1.0), ('a', 'c', 2.0), @@ -171,6 +194,26 @@ def test_k3_bidirectional(self): self.assertEqual(ground_count, len(min_routes)) + def test_graph_missing_edges(self): + G = nx.Graph() + G.add_weighted_edges_from([ + (1, 0, 1), + (0, 3, 1), + (3, 2, 1), + (2, 1, 1), + ]) + dnx.traveling_salesperson_qubo(G, missing_edge_weight=15) + + def test_digraph_missing_edges(self): + G = nx.DiGraph() + G.add_weighted_edges_from([ + (1, 0, 1), + (0, 3, 1), + (3, 2, 1), + (2, 1, 1), + ]) + dnx.traveling_salesperson_qubo(G, missing_edge_weight=15) + def test_k4_equal_weights(self): # k5 with all equal weights so all paths are equally good G = nx.Graph()