Optimize dijkstra implementation

Author: constantinpape

MIGRATION_GUIDE.md

@@ -2395,6 +2395,9 @@ TEASAR therefore uses Dijkstra where it must reconstruct and join paths.
 
 See `development/distance/benchmark_dijkstra.py` for a standalone benchmark of
 the Dijkstra field, predecessor field, weighted field, and early-stopping path.
+The implementation uses reusable fixed neighbor metadata and specialized heap
+strategies internally, but these do not change dtype, tie-breaking, or public
+results.
 
 ## scikit-fmm
 
@@ -2545,15 +2548,20 @@ Important differences and current scope:
   cross-section metadata, or postprocessing heuristics. These differences can
   change branch positions and vertex counts, so output is not expected to be
   vertex-for-vertex identical to kimimaro.
-- The C++ core is dependency-free and single-threaded. It deliberately reuses
-  the same public grid-Dijkstra implementation described above for root fields
-  and penalized rail paths.
+- The C++ core is dependency-free and single-threaded. TEASAR uses an internal
+  FP64 compact-foreground Dijkstra backend for root fields and penalized rail
+  paths. Compact IDs retain C-order tie-breaking, so it is bitwise identical to
+  the dense public grid solver on the tested TEASAR workloads while avoiding
+  full-volume shortest-path fields. The public Dijkstra functions remain dense
+  and exact FP64.
 
 Correctness tests are under `tests/skeleton/test_teasar.py`. The independent
 Dijkstra benchmark is `development/distance/benchmark_dijkstra.py`; the
 end-to-end synthetic branching-tube benchmark is
 `development/skeleton/benchmark_teasar.py` and can optionally add kimimaro with
-`--kimimaro` when it is installed.
+`--kimimaro` when it is installed. Pass `--sequential-backends` to reproduce
+the dense/compact/precision design matrix; extended density, spacing, and PDRF
+regimes are in `development/skeleton/benchmark_teasar_sequential.py`.
 
 ## I/O and Build Dependencies
 

development/distance/DIJKSTRA_PERFORMANCE_NOTES.md

@@ -233,6 +233,71 @@ intentionally change the graph. Land those separately with explicit expected
 topology/radius tests and rerun the kimimaro comparison without asserting
 vertex-for-vertex parity.
 
-Current full-suite baseline: `1126 passed`. Use the `256³` three-repeat median
+Pre-optimization full-suite baseline: `1126 passed`. Use the `256³` three-repeat median
 and peak RSS as the performance acceptance benchmark for the next optimization
 step.
+
+## Implemented compact-indexing results
+
+The sequential design matrix in `OPTIM_DIJKSTRA.md` was implemented and
+measured on 2026-07-14. Reproduce the size sweep and extended regimes with:
+
+```bash
+python development/skeleton/benchmark_teasar.py \
+    --large --sequential-backends --repeats 5
+python development/skeleton/benchmark_teasar_sequential.py --repeats 5
+```
+
+Compact IDs follow ascending full C-order indices. Compact FP64 therefore
+preserves dense FP64 heap tie-breaking and produced bitwise-identical vertices,
+edges, and radii for every size, spacing, density, and PDRF regime tested.
+
+Five-repeat end-to-end medians for the main anisotropic branching tubes:
+
+| backend | `128^3` | `192^3` | `256^3` |
+| --- | ---: | ---: | ---: |
+| optimized dense FP64 | 134.66 ms | 618.82 ms | 1.169 s |
+| compact on-the-fly FP64 | **74.04 ms** | **267.21 ms** | **686.58 ms** |
+| compact CSR FP64 | 75.19 ms | 284.57 ms | 730.72 ms |
+| compact CSR FP32 | 74.66 ms | 284.70 ms | 730.21 ms |
+
+CSR was 1.6%, 6.5%, and 6.4% slower than on-the-fly FP64 on these tiers. It
+was only 2.3% faster on the separate extremely sparse `192^3` case and was
+40% slower on the relatively dense `96^3` ball. It therefore failed the rule
+requiring a 5% win on both large tiers. The selected production backend is
+compact on-the-fly FP64: a temporary/full `uint32` lookup plus
+`compact_to_full`, with root, PDRF, heap state, predecessors, skeleton targets,
+and voxel-to-vertex data indexed only over foreground.
+
+On the `256^3` case, process peak RSS fell from 1,143,572 KiB in the original
+implementation to 266,536 KiB, a 76.7% reduction. The optimized dense backend
+measured 839,596 KiB in the same single-call process. Compact on-the-fly and
+CSR had effectively identical peaks because both must first construct the
+65.5 MiB full-to-compact lookup; CSR releases it only after its adjacency is
+built.
+
+Profile-build phase totals on `256^3` show where the improvement lands:
+
+| phase | optimized dense FP64 | selected compact FP64 |
+| --- | ---: | ---: |
+| compact-domain build | -- | 41.6 ms |
+| root Dijkstra (two fields) | 178.7 ms | 58.8 ms |
+| PDRF construction | 53.4 ms | 3.7 ms |
+| rail-path Dijkstra | 274.3 ms | 29.0 ms |
+| total TEASAR | 995.0 ms | 639.9 ms |
+
+The exact distance-to-boundary transform is now 77% of selected-backend time;
+combined root/path Dijkstra is 5.2x faster than the optimized dense backend and
+12.5x faster than the original recorded Dijkstra phases.
+
+### FP32 decision
+
+Internal FP32 was rejected. On `256^3`, CSR FP32 combined root/path Dijkstra
+was 113.9 ms versus 108.9 ms for CSR FP64, end-to-end time was slightly slower,
+and peak RSS was unchanged at about 266.5 MiB. It also failed the quality gate
+on `192^3`: the contracted degree signature changed, bidirectional 95th
+percentile distance was 2.0 voxels, Hausdorff distance was 7.07 voxels, and
+total physical length differed by 2.26%. Automatic TEASAR execution therefore
+remains FP64. Parallel shortest paths remain a separate follow-up.
+
+Final verification after selecting compact on-the-fly FP64: `1129 passed`.

development/skeleton/PERFORMANCE_NOTES.md

@@ -10,6 +10,7 @@
 --------
 python development/skeleton/benchmark_teasar.py --small --repeats 3
 python development/skeleton/benchmark_teasar.py --repeats 5
+python development/skeleton/benchmark_teasar.py --large --sequential-backends
 python development/skeleton/benchmark_teasar.py --large --kimimaro --repeats 3 \
     --json /tmp/teasar.json
 """
@@ -27,6 +28,7 @@
 import numpy as np
 
 import bioimage_cpp as bic
+from bioimage_cpp import _core
 
 
 def draw_ball(mask: np.ndarray, center: np.ndarray, radius: int) -> None:
@@ -85,6 +87,43 @@ def count_bic(result) -> tuple[int, int]:
     return len(vertices), len(edges)
 
 
+def compare_skeletons(reference, candidate, spacing):
+    """Return the FP32 acceptance metrics used by OPTIM_DIJKSTRA.md."""
+    ref_vertices, ref_edges, _ = reference
+    got_vertices, got_edges, _ = candidate
+
+    def degree_signature(edges, n_vertices):
+        degrees = np.bincount(edges.ravel().astype(np.int64), minlength=n_vertices)
+        return tuple(sorted(int(degree) for degree in degrees if degree != 2))
+
+    ref_voxels = ref_vertices / np.asarray(spacing, dtype=np.float64)
+    got_voxels = got_vertices / np.asarray(spacing, dtype=np.float64)
+    pairwise = np.linalg.norm(
+        ref_voxels[:, None, :] - got_voxels[None, :, :], axis=2
+    )
+    ref_to_got = pairwise.min(axis=1)
+    got_to_ref = pairwise.min(axis=0)
+
+    def total_length(vertices, edges):
+        if len(edges) == 0:
+            return 0.0
+        return float(
+            np.linalg.norm(vertices[edges[:, 0]] - vertices[edges[:, 1]], axis=1).sum()
+        )
+
+    ref_length = total_length(ref_vertices, ref_edges)
+    got_length = total_length(got_vertices, got_edges)
+    return {
+        "topology_match": degree_signature(ref_edges, len(ref_vertices))
+        == degree_signature(got_edges, len(got_vertices)),
+        "bidirectional_p95_voxels": float(
+            max(np.percentile(ref_to_got, 95), np.percentile(got_to_ref, 95))
+        ),
+        "hausdorff_voxels": float(max(ref_to_got.max(), got_to_ref.max())),
+        "relative_length_error": abs(got_length - ref_length) / max(ref_length, 1e-300),
+    }
+
+
 def kimimaro_call(mask, spacing, scale, constant, pdrf_scale, pdrf_exponent):
     import kimimaro
 
@@ -112,6 +151,19 @@ def kimimaro_call(mask, spacing, scale, constant, pdrf_scale, pdrf_exponent):
     return skeletons[1]
 
 
+def bic_backend_call(mask, spacing, parameters, backend):
+    """Call a development-only C++ backend without changing the public API."""
+    return _core._teasar_uint8_backend(
+        mask,
+        spacing,
+        parameters["scale"],
+        parameters["constant"],
+        parameters["pdrf_scale"],
+        parameters["pdrf_exponent"],
+        backend,
+    )
+
+
 def main() -> int:
     parser = argparse.ArgumentParser(description=__doc__)
     sizes = parser.add_mutually_exclusive_group()
@@ -120,6 +172,11 @@ def main() -> int:
     parser.add_argument("--repeats", type=int, default=5)
     parser.add_argument("--warmup", type=int, default=1)
     parser.add_argument("--kimimaro", action="store_true")
+    parser.add_argument(
+        "--sequential-backends",
+        action="store_true",
+        help="compare dense FP64, compact on-the-fly/CSR FP64, and CSR FP32",
+    )
     parser.add_argument("--json", default="", help="optional JSON result path")
     args = parser.parse_args()
     if args.repeats < 1 or args.warmup < 0:
@@ -142,13 +199,30 @@ def main() -> int:
         "pdrf_scale": 100000.0,
         "pdrf_exponent": 4.0,
     }
-    backends = [
-        (
-            "bioimage-cpp",
-            lambda mask: bic.skeleton.teasar(mask, spacing=spacing, **parameters),
-            count_bic,
-        )
-    ]
+    if args.sequential_backends:
+        backends = [
+            (
+                backend,
+                lambda mask, backend=backend: bic_backend_call(
+                    mask, spacing, parameters, backend
+                ),
+                count_bic,
+            )
+            for backend in (
+                "dense-fp64",
+                "compact-on-the-fly-fp64",
+                "compact-csr-fp64",
+                "compact-csr-fp32",
+            )
+        ]
+    else:
+        backends = [
+            (
+                "bioimage-cpp",
+                lambda mask: bic.skeleton.teasar(mask, spacing=spacing, **parameters),
+                count_bic,
+            )
+        ]
     if args.kimimaro:
         backends.append(
             (
@@ -162,6 +236,7 @@ def main() -> int:
         )
 
     rows = []
+    quality_rows = []
     header = f"{'backend':>14} {'shape':>14} {'foreground':>11} {'vertices':>9} {'median ms':>11} {'min ms':>10}"
     print(header)
     print("-" * len(header))
@@ -193,7 +268,34 @@ def main() -> int:
                 f"{name:>14} {str(mask.shape):>14} {foreground:11d} {vertices:9d} "
                 f"{median_s * 1e3:11.2f} {min(samples) * 1e3:10.2f}"
             )
-        if len(case_rows) == 2:
+        if args.sequential_backends:
+            dense_result = results_by_backend["dense-fp64"]
+            for exact_backend in (
+                "compact-on-the-fly-fp64", "compact-csr-fp64"
+            ):
+                exact = all(
+                    np.array_equal(got, expected)
+                    for got, expected in zip(
+                        results_by_backend[exact_backend], dense_result
+                    )
+                )
+                print(f"  {exact_backend} exact dense parity: {exact}")
+            quality = compare_skeletons(
+                results_by_backend["compact-csr-fp64"],
+                results_by_backend["compact-csr-fp32"],
+                spacing,
+            )
+            quality_rows.append({"shape": list(mask.shape), **quality})
+            print(
+                "  FP32 quality: "
+                f"topology={quality['topology_match']} "
+                f"p95={quality['bidirectional_p95_voxels']:.3f} vox "
+                f"hausdorff={quality['hausdorff_voxels']:.3f} vox "
+                f"length_error={quality['relative_length_error']:.3%}"
+            )
+        if len(case_rows) == 2 and {row["backend"] for row in case_rows} == {
+            "bioimage-cpp", "kimimaro"
+        }:
             by_name = {row["backend"]: row for row in case_rows}
             ratio = (
                 by_name["bioimage-cpp"]["median_s"]
@@ -203,7 +305,11 @@ def main() -> int:
 
     if args.json:
         with open(args.json, "w", encoding="utf-8") as file:
-            json.dump({"repeats": args.repeats, "results": rows}, file, indent=2)
+            json.dump(
+                {"repeats": args.repeats, "results": rows, "quality": quality_rows},
+                file,
+                indent=2,
+            )
         print(f"wrote {args.json}", file=sys.stderr)
     return 0