beam: A* shortest-path with admissible heuristic + randomized cross-check vs Dijkstra

wizard-work/beam/astar.py

@@ -0,0 +1,84 @@
+"""A* shortest-path search on a weighted graph."""
+from __future__ import annotations
+import heapq
+from typing import Callable, Dict, Hashable, List, Tuple
+
+Node = Hashable
+Graph = Dict[Node, Dict[Node, float]]
+
+
+def astar(
+    graph: Graph,
+    start: Node,
+    goal: Node,
+    h: Callable[[Node], float] = None,
+) -> Tuple[float, List[Node]]:
+    if start not in graph:
+        raise ValueError(f"start node {start!r} not in graph")
+    if goal not in graph:
+        raise ValueError(f"goal node {goal!r} not in graph")
+    if h is None:
+        h = lambda _n: 0.0
+
+    g: Dict[Node, float] = {start: 0.0}
+    came: Dict[Node, Node] = {}
+    closed: set = set()
+    counter = 0
+    open_heap: list = [(h(start), counter, start)]
+    counter += 1
+
+    while open_heap:
+        _f, _c, current = heapq.heappop(open_heap)
+        if current in closed:
+            continue
+        if current == goal:
+            path = [current]
+            while current in came:
+                current = came[current]
+                path.append(current)
+            path.reverse()
+            return g[goal], path
+        closed.add(current)
+        for nbr, w in graph.get(current, {}).items():
+            if nbr in closed:
+                continue
+            if w < 0:
+                raise ValueError("negative edge weights are not supported")
+            tentative = g[current] + w
+            if tentative < g.get(nbr, float("inf")):
+                g[nbr] = tentative
+                came[nbr] = current
+                f = tentative + h(nbr)
+                heapq.heappush(open_heap, (f, counter, nbr))
+                counter += 1
+
+    return float("inf"), []
+
+
+def dijkstra(graph, start, goal):
+    if start not in graph or goal not in graph:
+        raise ValueError("node not in graph")
+    dist = {start: 0.0}
+    prev = {}
+    pq = [(0.0, start)]
+    visited = set()
+    while pq:
+        d, u = heapq.heappop(pq)
+        if u in visited:
+            continue
+        visited.add(u)
+        if u == goal:
+            break
+        for v, w in graph.get(u, {}).items():
+            nd = d + w
+            if nd < dist.get(v, float("inf")):
+                dist[v] = nd
+                prev[v] = u
+                heapq.heappush(pq, (nd, v))
+    if goal not in dist:
+        return float("inf"), []
+    path = [goal]
+    while path[-1] != start:
+        path.append(prev[path[-1]])
+    path.reverse()
+    return dist[goal], path

wizard-work/beam/test_astar.py

@@ -0,0 +1,93 @@
+import random
+from astar import astar, dijkstra
+
+
+def test_simple():
+    g = {'s': {'a': 1, 'b': 4}, 'a': {'g': 2}, 'b': {'g': 1}, 'g': {}}
+    d, p = astar(g, 's', 'g')
+    assert d == 3 and p == ['s', 'a', 'g']
+
+
+def test_no_path():
+    g = {'s': {'a': 1}, 'a': {}, 'g': {}}
+    d, p = astar(g, 's', 'g')
+    assert d == float('inf') and p == []
+
+
+def test_start_goal_same():
+    g = {'s': {}}
+    d, p = astar(g, 's', 's')
+    assert d == 0 and p == ['s']
+
+
+def test_invalid_node():
+    g = {'s': {}}
+    try:
+        astar(g, 'x', 's')
+    except ValueError:
+        pass
+    else:
+        raise AssertionError("expected ValueError")
+
+
+def test_negative_edge():
+    g = {'s': {'a': -1}, 'a': {}}
+    try:
+        astar(g, 's', 'a')
+    except ValueError:
+        pass
+    else:
+        raise AssertionError("expected ValueError")
+
+
+def test_heuristic_admissible():
+    # grid 3x3, h = manhattan distance
+    nodes = [(r, c) for r in range(3) for c in range(3)]
+    g = {n: {} for n in nodes}
+    for r, c in nodes:
+        for dr, dc in ((1, 0), (-1, 0), (0, 1), (0, -1)):
+            nr, nc = r + dr, c + dc
+            if 0 <= nr < 3 and 0 <= nc < 3:
+                g[(r, c)][(nr, nc)] = 1.0
+
+    def manhattan(n):
+        return abs(n[0] - 2) + abs(n[1] - 2)
+
+    d, p = astar(g, (0, 0), (2, 2), h=manhattan)
+    assert d == 4
+
+
+def test_matches_dijkstra_random():
+    random.seed(42)
+    for _ in range(200):
+        n = random.randint(2, 8)
+        nodes = list(range(n))
+        g = {i: {} for i in nodes}
+        for i in nodes:
+            for j in nodes:
+                if i != j and random.random() < 0.35:
+                    g[i][j] = random.randint(1, 9)
+        s, t = random.choice(nodes), random.choice(nodes)
+        # ensure both present (they are)
+        da, pa = astar(g, s, t)
+        dd, pd = dijkstra(g, s, t)
+        assert da == dd, (da, dd)
+        # path validity
+        if da != float('inf'):
+            # verify path cost
+            cost = 0.0
+            for a, b in zip(pa, pa[1:]):
+                cost += g[a][b]
+            assert abs(cost - da) < 1e-9
+            assert pa[0] == s and pa[-1] == t
+
+
+if __name__ == '__main__':
+    test_simple()
+    test_no_path()
+    test_start_goal_same()
+    test_invalid_node()
+    test_negative_edge()
+    test_heuristic_admissible()
+    test_matches_dijkstra_random()
+    print("ALL TESTS PASSED")