perf(mobile): index completed PRM roadmaps and benchmark queries

examples/benchmark_roadmap.py

@@ -0,0 +1,183 @@
+#!/usr/bin/env python
+"""
+Compare PRM and lattice queries with and without pgraph's spatial index.
+
+Run from a checkout with ``python examples/benchmark_roadmap.py``. Roadmaps
+are built by the real planners once; planning is excluded from query timings.
+The default cases use 1000 PRM vertices and 12 lattice iterations. Increase
+``--npoints`` and ``--iterations`` to compare larger maps. PRM has a central
+obstacle and a fixed sampling seed; lattice queries start at the root so that
+their goals are reachable. Both use a fixed batch of distinct destinations.
+
+Each repetition measures both lookup modes, alternating their order. The tree
+is invalidated through the public coordinate setter before each measurement.
+Index construction and two endpoint lookups are timed separately from complete
+queries (which include A*). All paths, and lattice costs, must agree. Timings
+are diagnostic, not test assertions, and depend on the dependency versions.
+
+With pgraph 1.0.0, lattice angles just below zero can wrap to exactly 2*pi,
+which scipy rejects as outside the periodic box. The JSON reports this as
+``index_error`` and leaves indexed timings unset; it does not time a fallback
+scan as an indexed query. Five lattice iterations already reproduce it.
+"""
+
+import argparse
+import json
+import platform
+from importlib.metadata import version
+from statistics import median
+from time import perf_counter
+from typing import Any
+
+import numpy as np
+
+from roboticstoolbox import LatticePlanner, PRMPlanner
+from roboticstoolbox.tools.types import NDArray
+
+
+def measure(
+    planner: PRMPlanner | LatticePlanner,
+    pairs: list[tuple[NDArray, NDArray]],
+    repeats: int,
+) -> dict[str, Any]:
+    """
+    Measure index construction, endpoint lookups and complete path queries.
+
+    :param planner: planner with a completed roadmap
+    :param pairs: reachable start and goal pairs
+    :param repeats: repetitions of each mode
+    :returns: timings in seconds, graph size and result comparison
+    """
+    graph = planner.graph
+    vertex = next(iter(graph))
+    vertex.coord = vertex.coord.copy()
+    reference = [planner.query(start, goal) for start, goal in pairs]
+    samples = {
+        mode: {key: [] for key in ("build", "endpoints", "query")}
+        for mode in ("scan", "kdtree")
+    }
+    index_error = None
+
+    for repetition in range(repeats):
+        modes = ("scan", "kdtree") if repetition % 2 == 0 else ("kdtree", "scan")
+        for mode in modes:
+            # Assignment invalidates the index without changing graph geometry.
+            vertex.coord = vertex.coord.copy()
+            start_time = perf_counter()
+            if mode == "kdtree":
+                try:
+                    graph.build_kdtree()
+                except ValueError as error:
+                    # Report an unsupported graph instead of timing a fallback
+                    # scan and presenting it as an indexed query.
+                    index_error = str(error)
+                    continue
+            samples[mode]["build"].append(perf_counter() - start_time)
+
+            # Warm each lookup path before measuring steady-state queries.
+            for start, goal in pairs:
+                graph.closest(start)
+                graph.closest(goal)
+            start_time = perf_counter()
+            for start, goal in pairs:
+                graph.closest(start)
+                graph.closest(goal)
+            samples[mode]["endpoints"].append(
+                (perf_counter() - start_time) / len(pairs)
+            )
+
+            start_time = perf_counter()
+            results = [planner.query(start, goal) for start, goal in pairs]
+            samples[mode]["query"].append((perf_counter() - start_time) / len(pairs))
+            for expected, actual in zip(reference, results):
+                if isinstance(planner, LatticePlanner):
+                    np.testing.assert_array_equal(actual[0], expected[0])
+                    np.testing.assert_equal(actual[1].cost, expected[1].cost)
+                else:
+                    np.testing.assert_array_equal(actual, expected)
+
+    medians = {
+        mode: {
+            key: median(values) if values else None for key, values in timings.items()
+        }
+        for mode, timings in samples.items()
+    }
+    saving = (
+        medians["scan"]["query"] - medians["kdtree"]["query"]
+        if index_error is None
+        else 0
+    )
+    return {
+        "vertices": graph.n,
+        "edges": graph.ne,
+        "queries_per_repeat": len(pairs),
+        "paths_equal": True if index_error is None else None,
+        "index_error": index_error,
+        "seconds": samples,
+        "median_seconds": medians,
+        "query_speedup": (
+            medians["scan"]["query"] / medians["kdtree"]["query"]
+            if index_error is None
+            else None
+        ),
+        "break_even_queries": (
+            int(np.ceil(medians["kdtree"]["build"] / saving)) if saving > 0 else None
+        ),
+    }
+
+
+def main() -> None:
+    """Build reproducible roadmaps and print the measurements as JSON."""
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--npoints", type=int, default=1000)
+    parser.add_argument("--iterations", type=int, default=12)
+    parser.add_argument("--queries", type=int, default=20)
+    parser.add_argument("--repeats", type=int, default=5)
+    parser.add_argument("--seed", type=int, default=0)
+    args = parser.parse_args()
+    if min(args.npoints, args.iterations, args.queries, args.repeats) < 1:
+        parser.error("npoints, iterations, queries and repeats must be positive")
+
+    grid = np.zeros((100, 100))
+    grid[40:60, 40:60] = 1
+    prm = PRMPlanner(grid, seed=args.seed, progress=False, dist_thresh=15)
+    prm.plan(npoints=args.npoints)
+    random = np.random.default_rng(args.seed)
+    goals = [vertex.coord.copy() for vertex in prm.graph]
+    indices = random.choice(
+        len(goals), size=min(args.queries, len(goals)), replace=False
+    )
+    prm_pairs = [(np.array([5.0, 5.0]), goals[i]) for i in indices]
+
+    lattice = LatticePlanner(progress=False)
+    lattice.plan(iterations=args.iterations)
+    goals = [vertex.coord.copy() for vertex in lattice.graph]
+    indices = random.choice(
+        len(goals), size=min(args.queries, len(goals)), replace=False
+    )
+    lattice_pairs = [(np.zeros(3), goals[i]) for i in indices]
+    print(
+        json.dumps(
+            {
+                "python": platform.python_version(),
+                "platform": platform.platform(),
+                "dependencies": {
+                    name: version(name)
+                    for name in (
+                        "numpy",
+                        "scipy",
+                        "spatialmath-python",
+                        "pgraph-python",
+                    )
+                },
+                "seed": args.seed,
+                "prm": measure(prm, prm_pairs, args.repeats),
+                "lattice": measure(lattice, lattice_pairs, args.repeats),
+            },
+            indent=2,
+        )
+    )
+
+
+if __name__ == "__main__":
+    main()

src/roboticstoolbox/mobile/LatticePlanner.py

@@ -201,6 +201,14 @@ def plan(
         If an occupancy grid exists the if ``iterations`` is None the area of the
         grid will be completely filled.
 
+        .. note:: A spatial index can only be built after the last iteration:
+            adding a vertex invalidates it. Automatic indexing is deferred
+            because pgraph 1.0.0 can reject lattice angles just below zero:
+            wrapping them rounds to the upper boundary of its periodic k-d
+            tree. The unindexed nearest-vertex search remains valid. See
+            ``examples/benchmark_roadmap.py`` for a comparison on completed
+            roadmaps, including index-construction failures.
+
         .. seealso:: :meth:`query`
         """
         if iterations is None and self.occgrid is None:

src/roboticstoolbox/mobile/PRMPlanner.py

@@ -252,7 +252,14 @@ def plan(
         created between points if the distance between them is less than
         ``dist_thresh``.
 
-        The roadmap is a ``UGraph``.
+        The roadmap is a ``UGraph``. A k-d tree is built after the roadmap is
+        complete to accelerate the nearest-vertex lookups in :meth:`query`.
+        The index does not accelerate roadmap construction or A* search.
+
+        .. note:: Adding or removing a vertex, or assigning ``vertex.coord``,
+            invalidates the index. Queries then use a linear scan until
+            ``planner.graph.build_kdtree()`` is called again. Coordinate edits
+            in place are not detected: assign the whole coordinate instead.
 
         .. deprecated:: 1.5.0
             ``animate`` was accepted but had no effect.
@@ -279,6 +286,11 @@ def plan(
         self.random_init()  # reset the random number generator
         self._create_roadmap(npoints, dist_thresh)
 
+        # Vertex additions invalidate the index, so build it only once the
+        # roadmap is complete. An empty roadmap has no coordinates to index.
+        if self._graph.n:
+            self._graph.build_kdtree()
+
     def query(
         self,
         start: ArrayLike2 | None = None,
@@ -298,6 +310,10 @@ def query(
         The path is a sparse sequence of waypoints, with variable distance
         between them.
 
+        If several roadmap vertices are exactly equidistant from an endpoint,
+        the spatial index may select a different vertex than a linear scan,
+        and hence a different path.
+
         .. warning:: Waypoints 1 to N-2 are part of the roadmap, while waypoints
             0 and N-1 are the start and goal respectively.  The first and last
             motion segment is not guaranteed to be obstacle free.

tests/test_graph_planners.py

@@ -22,6 +22,7 @@
 from collections.abc import Callable
 from pathlib import Path
 from typing import Any, TypeVar
+from unittest.mock import patch
 
 import matplotlib
 
@@ -112,6 +113,45 @@ def test_different_seed_different_roadmap(self) -> None:
         self.assertEqual(a.shape, b.shape)
         self.assertFalse(np.allclose(a, b))
 
+    def test_index_built_once_on_completed_roadmap(self) -> None:
+        sizes = []
+        build_kdtree = pgraph.UGraph.build_kdtree
+
+        def record_build(graph: pgraph.UGraph) -> None:
+            sizes.append(graph.n)
+            build_kdtree(graph)
+
+        with patch.object(pgraph.UGraph, "build_kdtree", record_build):
+            self.planned(npoints=60)
+        self.assertEqual(sizes, [60])
+
+    def test_indexed_queries_match_scan(self) -> None:
+        prm = self.planned()
+        goals = [vertex.coord.copy() for vertex in list(prm.graph)[::5]]
+        indexed = [prm.query(self.START, goal) for goal in goals]
+
+        # The coordinate setter discards the index without changing geometry.
+        vertex = next(iter(prm.graph))
+        vertex.coord = vertex.coord.copy()
+        for goal, expected in zip(goals, indexed):
+            nt.assert_array_equal(prm.query(self.START, goal), expected)
+
+    def test_replan_indexes_new_vertices(self) -> None:
+        prm = self.planned()
+        old_graph = prm.graph
+        quiet(prm.plan, npoints=20)
+        self.assertIsNot(prm.graph, old_graph)
+        self.assertEqual(prm.graph.n, 20)
+        for vertex in prm.graph:
+            closest, distance = prm.graph.closest(vertex.coord)
+            self.assertIs(closest, vertex)
+            self.assertEqual(distance, 0)
+
+    def test_empty_roadmap_is_not_indexed(self) -> None:
+        prm = self.planned(npoints=0)
+        with self.assertRaisesRegex(RuntimeError, "no plan computed"):
+            prm.query(self.START, self.GOAL)
+
     def test_goal_inside_obstacle(self) -> None:
         prm = self.planned()
         with self.assertRaises(ValueError):