diff --git a/wizard-work/beam/astar.py b/wizard-work/beam/astar.py new file mode 100644 index 0000000..d1ecb6b --- /dev/null +++ b/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 diff --git a/wizard-work/beam/test_astar.py b/wizard-work/beam/test_astar.py new file mode 100644 index 0000000..73cd42f --- /dev/null +++ b/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")