Idea: remove `numba` as a dependency

[scottstanie]

Meta-issue to track the removal of Numba.
There's only a few places we are using Numba instead of JAX for a compiled speedup. Having two accelerators seems excessive for the dependency list.

• Similarity
• interpolation
• _ps_filling.get_max_idxs
• ks SHP test

[s-sasaki-earthsea-wizard]

Hi @scottstanie — I'd like to take a crack at this checklist, starting with similarity.py, if no one else is already working on it. I have a working JAX replacement with equivalence and benchmark evidence; the performance picture has a real trade-off, so I'd like your read before opening a PR.

Equivalence (vs. the current numba implementation at c2f7c24): identical NaN patterns and max |diff| ≈ 1e-7 (float32 rounding) across 36 cases on both CPU and GPU backends — multiple radii, random/pathological masks, non-bool masks (the numba loop treats any nonzero as valid; kept), all-NaN and all-zero invalid pixels, non-complex phase inputs, images smaller than the search radius, and real OPERA CSLC blocks including one straddling the burst valid-data boundary. Comparison tolerance rtol=1e-5 / atol=5e-6 (the numba loop accumulates in float64; the JAX version computes in float32; a float64 phase input is likewise computed in complex64 rather than complex128).

Performance (16-core CPU, RTX 5080, jax/jaxlib 0.10.1; median_similarity, n_ifg=20, 512×512 blocks; "real" blocks are tiled from OPERA CSLC granules and include NaN borders). Values are median wall-time per block in seconds (3 passes; lower is better), with the ratio to the numba baseline in parentheses. I benchmarked the two configurations create_similarities is actually called with in-repo, not just its defaults:

caller config	dataset	numba	JAX CPU	JAX GPU
single.py (num_threads=1, r=7)	synthetic	2.17 s	3.26 s (1.5× slower)	0.42 s (5.1× faster)
single.py (num_threads=1, r=7)	real	2.01 s	2.99 s (1.5× slower)	0.37 s (5.4× faster)
sequential.py (num_threads=2, r=11)	synthetic	1.93 s	4.37 s (2.3× slower)	0.29 s (6.7× faster)
sequential.py (num_threads=2, r=11)	real	1.95 s	4.24 s (2.2× slower)	0.31 s (6.3× faster)

The honest summary:

• CPU-only users would regress in the current caller shapes: 1.5× slower for single.py, 2.2–2.3× for sequential.py. The cause is XLA's CPU sort/top_k being effectively single-threaded (the median is the bottleneck; I already use lax.top_k partial selection, which computes the same order statistics as np.nanmedian and halves the cost of a full sort). The r=11 case is worse because top_k cost grows superlinearly in the neighbor count.
• GPU is 5.1–6.7× faster than numba and radius-insensitive (0.29–0.42 s/block across both configs).
• The JAX version scales with block-level threading where numba self-saturates. At the function's default num_threads=5 (which no in-repo caller currently uses), JAX CPU is 1.5–1.7× faster than numba at r=7 — and 1.4–1.6× against a NUMBA_NUM_THREADS-capped control, so it's not just numba threading overhead. That does not extend to r=11: even at 5 threads the JAX version is still 1.3–1.4× slower there (measured) — the top_k cost at K=348 neighbors outweighs the threading gain.

Full harness and JSON results (equivalence suite, caller-shaped matrix, CPU-sort breakdown): https://github.com/s-sasaki-earthsea-wizard/dolphin-benchmark/tree/4a203e5e9844a5dbb724563c62482b6dc1f69d50/benchmarks/similarity_jax

Where I land, honestly: the numbers say the JAX kernel should run where it wins — on a GPU — and CPU-only users shouldn't pay for cleanup. That points to keeping the numba path on CPU and dispatching to the JAX kernel when a GPU is available. But that retains numba as a dependency and adds a dual code path, which cuts against the stated goal of this issue ("two accelerators seems excessive") — so I don't think it's my call to make. Depending on which you weight higher, I see three ways through:

1. dependency cleanup first: take the JAX version as-is — one less compiled backend, JAX-CUDA users gain 5.1–6.7×, CPU users pay 1.5–2.3× in the current caller shapes (documented above);
2. cleanup + recovering what threading can recover: pair it with a caller-side num_threads bump — that flips the r=7 single.py path in JAX's favor (1.5–1.7× faster at num_threads=5) but only narrows the r=11 sequential.py gap (still 1.3–1.4× slower at 5 threads, measured) — and it's a workflow change I'd validate separately for memory/I-O;
3. no regression anywhere: the CPU/GPU dispatch above, accepting that numba stays for now, or simply deferring similarity.py and starting from a checklist item without this trade-off (e.g. _ps_filling.get_max_idxs is pure-numpy-able).

The branch for (1) is ready; (2) would additionally need the caller-side change, which I'd validate separately. If you'd rather go (3), I'm happy to start there instead. Which trade-off would you prefer?

Two scoping notes on the checklist:

• Beyond the 4 items listed, numba is also imported in utils.py (both gpu_is_available, which probes numba.cuda, and the thread-cap helper that calls numba.set_num_threads) and in shp/_common.py — fully dropping the dependency would need those too.
• I'd hold off on shp/_ks.py until Add Gaussian multilooking option for phase linking #734 settles, since both touch the shp/ dispatcher.

---

Disclosure: this investigation and the patch were developed with assistance from AI coding tools (Claude, Codex, and Gemini). All measurements were executed on real hardware, and the evidence linked above (the equivalence suite, benchmark JSONs, and the CPU-sort breakdown) was generated and reviewed by the author.