test_beamsearch.py 1.1 KB

1234567891011121314151617181920212223242526272829303132
  1. """Unit tests for the beam search functions."""
  2. import networkx as nx
  3. def identity(x):
  4. return x
  5. class TestBeamSearch:
  6. """Unit tests for the beam search function."""
  7. def test_narrow(self):
  8. """Tests that a narrow beam width may cause an incomplete search."""
  9. # In this search, we enqueue only the neighbor 3 at the first
  10. # step, then only the neighbor 2 at the second step. Once at
  11. # node 2, the search chooses node 3, since it has a higher value
  12. # that node 1, but node 3 has already been visited, so the
  13. # search terminates.
  14. G = nx.cycle_graph(4)
  15. edges = nx.bfs_beam_edges(G, 0, identity, width=1)
  16. assert list(edges) == [(0, 3), (3, 2)]
  17. def test_wide(self):
  18. G = nx.cycle_graph(4)
  19. edges = nx.bfs_beam_edges(G, 0, identity, width=2)
  20. assert list(edges) == [(0, 3), (0, 1), (3, 2)]
  21. def test_width_none(self):
  22. G = nx.cycle_graph(4)
  23. edges = nx.bfs_beam_edges(G, 0, identity, width=None)
  24. assert list(edges) == [(0, 3), (0, 1), (3, 2)]