dijkstra3d-sparse

Dijkstra shortest paths, distance fields and connected components over sparse 3D voxel sets given as an (N, 3) integer coordinate array — a sparse analogue of seung-lab/dijkstra3d, which operates on dense 3D arrays. Rust core, Python/NumPy frontend.

Why

dijkstra3d is fast because it never builds an explicit graph: it walks an implicit rectangular grid where a voxel's neighbours are generated by coordinate offset. The only thing making it "dense" — and the reason it needs memory proportional to the bounding-box volume W·H·D — is that its coordinate → payload lookup is a dense array sized to the full box.

For sparse objects (a thin structure inside a large box, N ≪ W·H·D) that is wasteful. This library keeps the implicit-grid walk and swaps that one dense component for a sparse hash coordinate → compact index [0, N). Everything else — binary heap, edge relaxation, parent tracking, path reconstruction — is unchanged. No adjacency list is ever materialized (this is not a CSR/explicit-graph Dijkstra), and all working memory is O(N), independent of the bounding box.

From benchmarks/RESULTS.md: a 1.5M-voxel helical tube in a 16,267 × 4,005 × 4,006 bounding box solves in ~0.4 s within ~470 MiB peak RSS — where a dense field over the same box would need 4 TiB. Going coordinates → distance field through scipy.sparse.csgraph.dijkstra instead takes ~3.5 s and 1.6 GiB peak on the same workload: SciPy's solver itself is fast, but it first needs the ~30M-edge CSR graph materialized — exactly the step the implicit-grid walk skips.

One caveat: if you already hold a CSR graph and only solve on it repeatedly, SciPy's solver alone is competitive (~0.14 s on this workload once the graph exists). The advantage here is going from raw coordinates to a field — the typical starting point for voxel data — without ever paying the time and memory to build an edge list.

Install

pip install dijkstra3d-sparse

Pre-built wheels cover Linux / macOS / Windows, Python 3.9+. Building from source needs a Rust toolchain (pip invokes it automatically via maturin).

Quickstart

import numpy as np
import dijkstra3d_sparse as ds

# a sparse voxel set: (N, 3) integer coordinates, any origin, unsorted OK
voxels = np.argwhere(volume > 0).astype(np.int32)   # e.g. from a dense mask
# ... or coordinates that never lived in a dense array at all

# distance + predecessor field from voxel row 0
dist, pred = ds.dijkstra_field(voxels, sources=0, connectivity=26,
                               anisotropy=(16.0, 16.0, 40.0))

# shortest path to the voxel farthest from the source
target = int(np.argmax(np.where(np.isfinite(dist), dist, -1)))
coords = ds.path(voxels, pred, target, dist=dist)    # (M, 3), source → target

# connected components over the same implicit grid
n_components, labels = ds.connected_components(voxels, connectivity=26)

# hold coordinates instead of row indices? map them first
src = ds.index_of(voxels, [[10, 4, 2], [0, 0, 0]])
dist, pred = ds.dijkstra_field(voxels, src)          # multi-source: dist to nearest

dist/pred are 1-D arrays aligned 1:1 with the rows of voxels (the key difference from dijkstra3d, whose field is a dense 3D array). Unreached voxels get dist = +inf, pred = -1; -1 matches SciPy's "no predecessor" sentinel, so (dist, pred) is a drop-in for scipy.sparse.csgraph.dijkstra(..., return_predecessors=True) on the equivalent explicit graph.

API

dijkstra_field(voxels, sources, *, node_cost=None, connectivity=26,
               anisotropy=(1.0, 1.0, 1.0), cost_mode="vertex",
               free_mask=None, free_eps=1e-6, min_only=True,
               stop_mask=None, stop_count=1,
               index_kind="hash") -> (dist, pred)

shortest_path(voxels, source, target, **kw) -> (path, cost)  # early exit

shortest_path_to_set(voxels, source, stop_mask, **kw) -> (path, hit, cost)

path(voxels, pred, target, *, dist=None) -> (M, 3) int32   # source → target

connected_components(voxels, *, group=None,
                     connectivity=26) -> (n_components, labels)

label_adjacency(voxels, labels, *, connectivity=26) -> (K, 2) int64

exposed_faces(voxels, *, index_kind="hash") -> (N,) uint8   # surface-face mask

factorize(voxels, *, return_index=False,                    # dedup coords -> labels
          index_kind="hash") -> (n_labels, labels[, reps])

index_of(voxels, coords, *, strict=True) -> int | (M,) int64

Graph(voxels, *, index_kind="hash")   # reusable handle, methods below

Reusable Graph handle

Every free function above rebuilds the coordinate → row spatial index — the one O(N) setup cost — on each call. For repeated queries over the same voxel set, build a Graph once; it holds the index and exposes the same operations as methods, minus the voxels/index_kind arguments:

g = ds.Graph(voxels, index_kind="hash")   # O(N) index build happens here, once

dist, pred = g.dijkstra_field(0, cost_mode="geometric")     # reuses the index
dist2, _   = g.dijkstra_field([3, 7], node_cost=penalty,    # different cost model,
                              cost_mode="additive")         # same handle
coords, hit, cost = g.shortest_path_to_set(q, anchors)      # grafting primitive
n_comp, labels = g.connected_components()
n_rings, rings = g.connected_components(group=level)        # components per group
ring_edges = g.label_adjacency(rings)                       # which rings touch
face_mask = g.exposed_faces()                               # surface-face mask
rows = g.index_of(coords)
g.n, g.voxels, g.index_kind                                 # introspection

Only voxels and index_kind are fixed at construction — connectivity, anisotropy, cost_mode, node_cost and the masks stay per-call, so one handle serves queries with different cost models. Results are identical to the free functions (same code runs; only where the index is built moves), duplicate coordinates are rejected at construction, and the handle keeps its own copy of the coordinates, so it is unaffected by later mutation of the input array. The payoff scales with call count — grafting loops that issue one shortest_path_to_set per path are the motivating case (see benchmarks/RESULTS.md).

Edge-cost model

Step lengths are precomputed per offset from anisotropy = (wx, wy, wz), matching dijkstra3d exactly: axis moves cost wx/wy/wz, face diagonals sqrt(wa² + wb²), corner diagonals sqrt(wa² + wb² + wc²). The cost of the directed edge cur → nbr is then:

With node_cost=None every mode reduces to the pure geometric step length. Costs must be finite and non-negative (Dijkstra invariant; validated at the boundary).

free_mask: edges into masked voxels cost free_eps (small, strictly positive) in total. This supports incremental path extraction where later paths should ride an already-selected node set for ~free before diverging.

min_only=False runs one Dijkstra per source and returns (S, N) arrays, mirroring SciPy; the default True returns a single (N,) field of distances to the nearest source.

Early termination & search-to-a-set

Dijkstra settles nodes in non-decreasing distance order, so the moment a node is popped its distance and path are final. stop_mask exploits this: the search stops as soon as stop_count masked voxels have been settled (default 1 — i.e. at the nearest member of the set), returning a partial field that is exact on everything it touched and +inf/-1 beyond. SciPy's limit distance cutoff cannot express "stop when you reach node X / this set". Two wrappers make this ergonomic:

# point → point, terminating the instant the target settles
coords, cost = ds.shortest_path(voxels, source, target)

# point → nearest member of an anchor set
coords, hit, cost = ds.shortest_path_to_set(voxels, source, anchor_mask)
# hit = row index of the anchor reached (-1 + empty path if unreachable)

This is the primitive for incremental tree construction (grafting — e.g. centerline/skeleton extraction): repeatedly connect a query voxel to a growing anchor set, where each query only explores the local catchment between the query and the nearest anchor instead of the full voxel set:

anchors = np.zeros(len(voxels), dtype=bool)
anchors[seed] = True
for query in queries:
    coords, hit, cost = ds.shortest_path_to_set(voxels, query, anchors)
    anchors[ds.index_of(voxels, coords)] = True   # graft the spur

On the benchmark tube (1.5M voxels), 60 such grafts run in ~1.4 s total, with per-query touched voxels falling from ~10% of N (sparse anchors) to ~0.3% (dense anchors) — versus 100% of N per query for repeated full fields. stop_mask composes with everything else: with multiple sources and min_only=True it means "grow a field from all sources until it first touches the anchor set". It is also the recommended replacement for free_mask-based grafting tricks — cleaner (no cost distortion) and cheaper (early exit); if both are given they stay independent (free_mask changes edge costs, stop_mask only changes termination).

Graph contraction without an edge list

Two primitives run over the same implicit-grid probe as everything else, so a caller can contract the voxel graph — collapse voxels into groups, then ask which groups touch — without ever materializing adjacencies:

• connected_components(voxels, group=values) connects two voxels only when they are connectivity-adjacent and group[u] == group[v], i.e. the components of the sub-graph induced by each group value. Grouping only constrains unions: the same value in two spatially separate places stays two components, and a voxel whose group differs from all its neighbours' becomes a singleton. group=None is the plain component labelling.
• label_adjacency(voxels, labels) returns the distinct pairs of different labels that touch, as a sorted (K, 2) array — the edges of the quotient graph. It deduplicates during the probe, so the (typically enormous) intermediate adjacency count never exists. labels need not be dense or non-negative.

Together they express level-set / Reeb-graph constructions such as wavefront skeletonization in three passes over one Graph:

g = ds.Graph(voxels)
n_comp, comp = g.connected_components()                    # one wave per component
seeds = [int(np.flatnonzero(comp == c)[0]) for c in range(n_comp)]
dist, _ = g.dijkstra_field(seeds, cost_mode="geometric")

level = np.floor(dist / step_size).astype(np.int64)        # geodesic level sets
n_rings, rings = g.connected_components(group=level)       # rings = level components
skeleton_edges = g.label_adjacency(rings)                  # contract onto rings

On the benchmark tube (1.5M voxels) that pipeline runs in 1.1 s at 514 MiB peak RSS, versus 5.7 s at 3.6 GiB for the same result via an explicit edge list plus SciPy — 30.5M adjacencies materialized to yield 2,429 distinct ring pairs (see benchmarks/RESULTS.md).

Surface faces

exposed_faces(voxels) answers, for every voxel in one pass, which of its six face-neighbours are absent from the set — the first stage of any voxel mesher / surface extraction. It returns an (N,) uint8 mask, bit k set iff the neighbour across face k is missing, in the order +x, -x, +y, -y, +z, -z:

mask = ds.exposed_faces(voxels)             # (N,) uint8; 0 = interior, 63 = isolated
right = voxels[(mask & (1 << 0)) != 0]      # voxels whose +x face is exposed

Unlike the other free functions it does not route through a Graph: it builds the spatial index, runs the probe, and frees the index inside the one call, so the surface pass leaves nothing behind for a later stage's peak to stack on. (Graph.exposed_faces() exists too, but reuses — and keeps — the handle's index, so it gives up that transient-memory win.)

Deduplicating coordinates

factorize(voxels) is the sparse np.unique(coords, axis=0, return_inverse=True): it assigns every row a dense label, equal exactly when the coordinates are equal, in one hash pass instead of a sort. It is the only primitive here that accepts duplicate coordinates — collapsing them is the whole point (the others reject repeats).

n, labels = ds.factorize(cells)                      # (E, 3) with repeats -> labels
n, labels, reps = ds.factorize(cells, return_index=True)
unique = cells[reps]                                 # one row per label...
assert np.array_equal(unique[labels], cells)         # ...and labels index back

Labels are 0 .. n-1 in order of first appearance by row; reps[k] is the first row carrying label k. This is the dedup a voxel mesher runs on its per-quad corners, and the coarse-cell assignment (fine // scale) a downsampler runs on its nodes — the same operation, so it lives here rather than as a hand-rolled argsort in each caller. It groups by exact coordinate equality, which is a different question from connected_components (spatial adjacency, unique input) despite the shared (n, labels) return.

Notes

• Multiple sources: seed them all — one pass computes distance-to-nearest-source and predecessors pointing back to each voxel's nearest source.
• Output is deterministic: heap ties break on row index, so identical inputs give identical fields across runs and platforms.
• Duplicate coordinates in voxels raise ValueError.
• Coordinates may be negative and use the full int32 range; there is no bounding-box extent limit.
• index_kind selects the spatial-index backend ("hash" FxHashMap probes, default; "sorted" binary search over sorted keys, slightly lower memory). Results are identical.

Development

uv venv && source .venv/bin/activate
uv pip install numpy scipy pytest maturin
maturin develop --release --uv   # build the Rust extension into the venv
pytest                           # Python test suite (SciPy parity + properties)
cargo test                       # Rust unit tests
python benchmarks/bench.py      # benchmark + O(N) memory gate

The test suite asserts parity with scipy.sparse.csgraph on the equivalent explicit CSR graph for all cost modes, connectivities and anisotropies, plus structural invariants (source distance 0, triangle inequality along edges, path adjacency/cost).

License

GPL-3.0-or-later, like dijkstra3d.