import unittest
import numpy
import retworkx
class TestFloydWarshall(unittest.TestCase):
def test_floyd_warshall(self):
dag = retworkx.PyDAG()
qr_0 = dag.add_node('qr[0]')
qr_1 = dag.add_node('qr[1]')
qr_2 = dag.add_node('qr[2]')
cr_0 = dag.add_node('cr[0]')
cr_1 = dag.add_node('cr[1]')
cx_1 = dag.add_node('cx_1')
dag.add_edge(qr_0, cx_1, 'qr[0]')
dag.add_edge(qr_1, cx_1, 'qr[1]')
h_1 = dag.add_node('h_1')
dag.add_edge(cx_1, h_1, 'qr[0]')
cx_2 = dag.add_node('cx_2')
dag.add_edge(cx_1, cx_2, 'qr[1]')
dag.add_edge(qr_2, cx_2, 'qr[2]')
cx_3 = dag.add_node('cx_3')
dag.add_edge(h_1, cx_3, 'qr[0]')
dag.add_edge(cx_2, cx_3, 'qr[2]')
h_2 = dag.add_node('h_2')
dag.add_edge(cx_3, h_2, 'qr[2]')
qr_0_out = dag.add_node('qr[0]_out')
dag.add_edge(cx_3, qr_0_out, 'qr[0]')
qr_1_out = dag.add_node('qr[1]_out')
dag.add_edge(cx_2, qr_1_out, 'qr[1]')
qr_2_out = dag.add_node('qr[2]_out')
dag.add_edge(h_2, qr_2_out, 'qr[2]')
cr_0_out = dag.add_node('cr[0]_out')
dag.add_edge(cr_0, cr_0_out, 'qr[2]')
cr_1_out = dag.add_node('cr[1]_out')
dag.add_edge(cr_1, cr_1_out, 'cr[1]')
result = retworkx.floyd_warshall(dag)
expected = {
0: {0: 0, 5: 1, 6: 2, 7: 2, 8: 3, 9: 4, 10: 4, 11: 3, 12: 5},
1: {1: 0, 5: 1, 6: 2, 7: 2, 8: 3, 9: 4, 10: 4, 11: 3, 12: 5},
2: {2: 0, 7: 1, 8: 2, 9: 3, 10: 3, 11: 2, 12: 4},
3: {3: 0, 13: 1},
4: {4: 0, 14: 1},
5: {5: 0, 6: 1, 7: 1, 8: 2, 9: 3, 10: 3, 11: 2, 12: 4},
6: {6: 0, 8: 1, 9: 2, 10: 2, 12: 3},
7: {7: 0, 8: 1, 9: 2, 10: 2, 11: 1, 12: 3},
8: {8: 0, 9: 1, 10: 1, 12: 2},
9: {9: 0, 12: 1},
10: {10: 0},
11: {11: 0},
12: {12: 0},
13: {13: 0},
14: {14: 0},
}
self.assertDictEqual(result, expected)
def test_floyd_warshall_numpy_three_edges(self):
graph = retworkx.PyGraph()
graph.add_nodes_from(list(range(6)))
weights = [2, 12, 1, 5, 1]
graph.add_edges_from([(i, i + 1, weights[i]) for i in range(5)])
graph.add_edge(5, 0, 10)
dist = retworkx.graph_floyd_warshall_numpy(graph, lambda x: x)
self.assertEqual(dist[0, 3], 15)
self.assertEqual(dist[3, 0], 15)
def test_weighted_numpy_two_edges(self):
graph = retworkx.PyGraph()
graph.add_nodes_from(list(range(8)))
graph.add_edges_from([
(0, 1, 2),
(1, 2, 2),
(2, 3, 1),
(3, 4, 1),
(4, 5, 1),
(5, 6, 1),
(6, 7, 1),
(7, 0, 1),
])
dist = retworkx.graph_floyd_warshall_numpy(graph, lambda x: x)
self.assertEqual(dist[0, 2], 4)
self.assertEqual(dist[2, 0], 4)
def test_weighted_numpy_negative_cycle(self):
graph = retworkx.PyGraph()
graph.add_nodes_from(list(range(4)))
graph.add_edges_from([
(0, 1, 1),
(1, 2, -1),
(2, 3, -1),
(3, 0, -1),
])
dist = retworkx.graph_floyd_warshall_numpy(graph, lambda x: x)
expected = numpy.array(
[[-6, -7, -8, -13],
[-7, -8, -9, -14],
[-8, -9, -10, -15],
[-13, -14, -15, -20]], dtype=numpy.float64)
self.assertTrue(numpy.array_equal(dist, expected))
def test_floyd_warshall_numpy_cycle(self):
graph = retworkx.PyGraph()
graph.add_nodes_from(list(range(7)))
graph.add_edges_from_no_data(
[(0, 1), (0, 6), (1, 2), (2, 3), (3, 4), (4, 5), (5, 6)])
dist = retworkx.graph_floyd_warshall_numpy(graph, lambda x: 1)
self.assertEqual(dist[0, 3], 3)
self.assertEqual(dist[0, 4], 3)
def test_directed_floyd_warshall_numpy_cycle_as_undirected(self):
graph = retworkx.PyDiGraph()
graph.add_nodes_from(list(range(7)))
graph.add_edges_from_no_data(
[(0, 1), (0, 6), (1, 2), (2, 3), (3, 4), (4, 5), (5, 6)])
dist = retworkx.digraph_floyd_warshall_numpy(graph, lambda x: 1,
as_undirected=True)
expected = numpy.array([[0., 1., 2., 3., 3., 2., 1.],
[1., 0., 1., 2., 3., 3., 2.],
[2., 1., 0., 1., 2., 3., 3.],
[3., 2., 1., 0., 1., 2., 3.],
[3., 3., 2., 1., 0., 1., 2.],
[2., 3., 3., 2., 1., 0., 1.],
[1., 2., 3., 3., 2., 1., 0.]])
self.assertTrue(numpy.array_equal(dist, expected))
def test_floyd_warshall_numpy_digraph_three_edges(self):
graph = retworkx.PyDiGraph()
graph.add_nodes_from(list(range(6)))
weights = [2, 12, 1, 5, 1]
graph.add_edges_from([(i, i + 1, weights[i]) for i in range(5)])
graph.add_edge(5, 0, 10)
dist = retworkx.digraph_floyd_warshall_numpy(graph, lambda x: x)
self.assertEqual(dist[0, 3], 15)
self.assertEqual(dist[3, 0], 16)
def test_weighted_numpy_digraph_two_edges(self):
graph = retworkx.PyDiGraph()
graph.add_nodes_from(list(range(8)))
graph.add_edges_from([
(0, 1, 2),
(1, 2, 2),
(2, 3, 1),
(3, 4, 1),
(4, 5, 1),
(5, 6, 1),
(6, 7, 1),
(7, 0, 1),
])
dist = retworkx.digraph_floyd_warshall_numpy(graph, lambda x: x)
self.assertEqual(dist[0, 2], 4)
self.assertEqual(dist[2, 0], 6)
def test_floyd_warshall_numpy_digraph_cycle(self):
graph = retworkx.PyDiGraph()
graph.add_nodes_from(list(range(7)))
graph.add_edges_from_no_data(
[(0, 1), (0, 6), (1, 2), (2, 3), (3, 4), (4, 5), (5, 6)])
dist = retworkx.digraph_floyd_warshall_numpy(graph, lambda x: 1)
self.assertEqual(dist[0, 3], 3)
self.assertEqual(dist[0, 4], 4)
def test_weighted_numpy_directed_negative_cycle(self):
graph = retworkx.PyDiGraph()
graph.add_nodes_from(list(range(4)))
graph.add_edges_from([
(0, 1, 1),
(1, 2, -1),
(2, 3, -1),
(3, 0, -1),
])
dist = retworkx.digraph_floyd_warshall_numpy(graph, lambda x: x)
expected = numpy.array(
[[-2, -1, -2, -3],
[-3, -2, -3, -4],
[-2, -1, -2, -3],
[-3, -2, -3, -4]], dtype=numpy.float64)
self.assertTrue(numpy.array_equal(dist, expected))
def test_numpy_directed_no_edges(self):
graph = retworkx.PyDiGraph()
graph.add_nodes_from(list(range(4)))
dist = retworkx.digraph_floyd_warshall_numpy(graph, lambda x: x)
expected = numpy.full((4, 4), numpy.inf)
numpy.fill_diagonal(expected, 0)
self.assertTrue(numpy.array_equal(dist, expected))
def test_numpy_no_edges(self):
graph = retworkx.PyGraph()
graph.add_nodes_from(list(range(4)))
dist = retworkx.graph_floyd_warshall_numpy(graph, lambda x: x)
expected = numpy.full((4, 4), numpy.inf)
numpy.fill_diagonal(expected, 0)
self.assertTrue(numpy.array_equal(dist, expected))
def test_floyd_warshall_numpy_digraph_cycle_with_removals(self):
graph = retworkx.PyDiGraph()
graph.add_nodes_from(list(range(8)))
graph.remove_node(0)
graph.add_edges_from_no_data(
[(1, 2), (1, 7), (2, 3), (3, 4), (4, 5), (5, 6), (6, 7)])
dist = retworkx.digraph_floyd_warshall_numpy(graph, lambda x: 1)
self.assertEqual(dist[0, 3], 3)
self.assertEqual(dist[0, 4], 4)
def test_floyd_warshall_numpy_graph_cycle_with_removals(self):
graph = retworkx.PyGraph()
graph.add_nodes_from(list(range(8)))
graph.remove_node(0)
graph.add_edges_from_no_data(
[(1, 2), (1, 7), (2, 3), (3, 4), (4, 5), (5, 6), (6, 7)])
dist = retworkx.graph_floyd_warshall_numpy(graph, lambda x: 1)
self.assertEqual(dist[0, 3], 3)
self.assertEqual(dist[0, 4], 3)
def test_floyd_warshall_numpy_digraph_cycle_no_weight_fn(self):
graph = retworkx.PyDiGraph()
graph.add_nodes_from(list(range(8)))
graph.remove_node(0)
graph.add_edges_from_no_data(
[(1, 2), (1, 7), (2, 3), (3, 4), (4, 5), (5, 6), (6, 7)])
dist = retworkx.digraph_floyd_warshall_numpy(graph)
self.assertEqual(dist[0, 3], 3)
self.assertEqual(dist[0, 4], 4)
def test_floyd_warshall_numpy_graph_cycle_no_weight_fn(self):
graph = retworkx.PyGraph()
graph.add_nodes_from(list(range(8)))
graph.remove_node(0)
graph.add_edges_from_no_data(
[(1, 2), (1, 7), (2, 3), (3, 4), (4, 5), (5, 6), (6, 7)])
dist = retworkx.graph_floyd_warshall_numpy(graph)
self.assertEqual(dist[0, 3], 3)
self.assertEqual(dist[0, 4], 3)
def test_floyd_warshall_numpy_digraph_cycle_default_weight(self):
graph = retworkx.PyDiGraph()
graph.add_nodes_from(list(range(8)))
graph.remove_node(0)
graph.add_edges_from_no_data(
[(1, 2), (1, 7), (2, 3), (3, 4), (4, 5), (5, 6), (6, 7)])
dist = retworkx.digraph_floyd_warshall_numpy(graph, default_weight=2)
self.assertEqual(dist[0, 3], 6)
self.assertEqual(dist[0, 4], 8)
def test_floyd_warshall_numpy_graph_cycle_default_weight(self):
graph = retworkx.PyGraph()
graph.add_nodes_from(list(range(8)))
graph.remove_node(0)
graph.add_edges_from_no_data(
[(1, 2), (1, 7), (2, 3), (3, 4), (4, 5), (5, 6), (6, 7)])
dist = retworkx.graph_floyd_warshall_numpy(graph, default_weight=2)
self.assertEqual(dist[0, 3], 6)
self.assertEqual(dist[0, 4], 6)