Call graph algorithms as SkeletonGraph methods, not networkx module functions

[IanConvy]

Converts this repo's call sites to the SkeletonGraph API that AllenInstitute/arborist now exposes: graph
algorithms as methods rather than networkx module functions, and node attribute allocation through the
registry rather than by hand. Two commits, rebased onto current main (8fe9593).

Important

Requires AllenInstitute/arborist#19 to merge first. That PR drops nx.Graph as SkeletonGraph's
parent in favour of a CSR array backing. FragmentsGraph subclasses SkeletonGraph, so nx.f(graph, ...)
stops working the moment it lands — the 15 call sites here are the breaking surface. This PR is a no-op
against today's arborist except for the registry cleanup, and it must not merge before arborist's.

Why this repo is affected at all

ProposalGraph(FragmentsGraph) → FragmentsGraph(SkeletonGraph) → SkeletonGraph(nx.Graph). Every
nx.connected_components(self) in this repo works only because of that last link. Replacing it with
SkeletonGraph(CSRGraph) breaks them, and the fix is one mechanical pattern.

The motivation is memory: nx.Graph costs ~554 B/node, and ~258 B/node of that is charged even with zero
edges. Against the MICrONS excitatory cohort (49,803 cells / 94.2M nodes) that is 52 GB on a 62 GB machine
versus 4.0 GB CSR-backed. Full numbers in the arborist PR.

What changed

refactor: use SkeletonGraph node attribute registry — FragmentsGraph.load allocated
node_component_id / node_radius / node_xyz by hand with their dtypes spelled out, and load_somas grew
them through a local resize_node_attr helper. Both now go through the registry:

• self.init_node_attrs(num_nodes) replaces the three np.zeros calls
• self.grow_node_attrs(num_nodes + num_somas) replaces three resize_node_attr calls, and the helper is
  deleted
• ProposalGraph.from_fragments_graph / to_fragments_graph replace three explicit .copy() lines with
  graph.copy_node_attrs(pg)

This is what makes the dtypes single-sourced: previously radius was float16 here and in arborist's loader
independently, and adding a fourth attribute meant editing three sites in two repos. The registry also
carries the SWC type column, so load_somas now sets self.node_type[node_id] = 1 — soma nodes were
previously written out as type 2 (axon), because the loader read type to flag somas and then discarded it.

refactor: call graph algorithms as methods, not nx module functions — 16 sites across 9 files:

call	count	sites
connected_components	5	visualization.py, geometry_util.py, groundtruth_generation.py, merge_datamodules.py, search_datasets.py
number_connected_components	3	fragments_graph.py ×2, geometry_util.py
shortest_path	2	groundtruth_generation.py, merge_inference.py
has_path	2	fragments_graph.py, split_inference.py
dfs_edges	2	search_datasets.py ×2
node_connected_component	1	search_datasets.py
nx.Graph() → SkeletonGraph()	1	visualization.py:51, the default empty gt_graph

The one non-mechanical change is ProposalGraph's two conversion paths, which built a graph with
cls.__new__(cls) + nx.Graph.__init__(pg) + pg.update(graph). update is an nx dict merge and has no CSR
equivalent, so both become pg._init_structure() + pg.copy_structure_from(graph) — an explicit structure
copy that shares indptr/indices rather than re-inserting every edge. copy_structure_from is added by the
fourth commit of the arborist PR specifically for these two sites.

networkx stays a dependency

Three files still import it, all correctly:

• merge_inference.py:318 and groundtruth_generation.py:352 catch nx.NetworkXNoPath by name —
  CSRGraph.shortest_path raises it deliberately, so the except clauses are untouched
• proposal_graph.py:466 — ProposalComputationGraph(nx.Graph) is a separate small graph, not a skeleton,
  and is left alone

fragments_graph.py and visualization.py lose their networkx imports because nothing in them uses it any
more.

Testing

tests/ held only __init__.py, so this adds tests/test_merge_integration.py — 5 tests, and it targets
the specific thing CSR is structurally bad at. merge_proposal calls add_edge and then immediately has
update_component_ids BFS the merged component through neighbors(), so the workload interleaves mutation
and adjacency reads one edge at a time. Unit coverage of the primitives (25 tests, in the arborist PR) does
not exercise that interleaving.

It builds its own 24-fragment graph rather than loading data, so it has no data dependency, and runs the real
ProposalGenerator at search_radius=25. The tests check that the accept loop is self-consistent, that the
has_path guard keeps the result a forest, that component_ids track the merges, and that the
FragmentsGraph ↔ ProposalGraph round trip preserves structure and node attributes.

5 tests pass in 4.9 s.

The full-scale check needs the real cohort and so lives outside this repo, in the capsule: 20 Patch-seq cells
with synthetic soma links suppressed to produce a genuinely fragmented graph — 159,614 nodes, 635 components,
3,662 proposals — cross-checked against a recorded pre-swap nx baseline on 11 invariants including the
proposal hash, accept/block counts, final component count and edge count, and a component_id checksum.
All 11 match. That baseline was re-earned on this exact base rather than carried over, because #760/#761/
#762 changed merge_proofreading itself after it was first recorded.

Incidentally measured there: has_path is 493× faster CSR-backed (17.27 s → 0.031 s over 2,000 calls),
which is the two has_path sites above.

One thing worth deciding separately

arborist is not declared in pyproject.toml — not in the 26 dependencies, and there is no
requirements.txt or setup.py mentioning it. The import resolves only because deployments happen to
pip install -e a local clone, which means nothing can refuse an incompatible arborist at install time. That
is exactly the failure this PR's ordering constraint is about. Adding arborist @ git+https://github.com/ AllenInstitute/arborist@<sha> would make the coupling checkable, but it is a packaging decision with its own
consequences for how this repo is deployed, so it is not bundled in here.

---

🤖 Generated with Claude Code

[IanConvy]

CI note

ci is red for a reason unrelated to this branch: the install step fails before any test runs.

ERROR: Ignored the following versions that require a different python version:
       0.1.71 Requires-Python >=3.10 ...
ERROR: No matching distribution found for tensorstore==0.1.71

pyproject.toml pins tensorstore==0.1.71, which requires Python ≥ 3.10, while .github/workflows runs the matrix at 3.8, 3.9 and 3.10. The 3.9 leg fails to resolve and fail-fast cancels 3.8 and 3.10, so no leg reports. This branch does not touch pyproject.toml or the workflow.

Locally, on Python 3.10: 5 tests pass in 4.9 s — the new tests/test_merge_integration.py, which is self-contained and has no data dependency.

---

🤖 Generated with Claude Code