Optimization of stitching, inference, added (optional) low-precision compute

[cophus]

Summary

Speeds up the experimental 4D-STEM reconstruction pipeline ~2.5× end-to-end on an NVIDIA L4 (233 s → 93 s over the 4-dataset benchmark; 49 s → 19 s per single reconstruction) while keeping results essentially identical. The gains come from three independent areas — position-corrected stitching, network-inference data movement, and optional low-precision compute. The default numerical path is unchanged; every new behavior is opt-in.

Results (NVIDIA L4, 4 experimental datasets)

stage	baseline	optimized (fp32, numpy stitch)	fp16 + GPU stitch	speedup
load	2.48	2.41	2.41	1.0×
inference	111.85	95.41	37.49	3.0×
rigid stitch	0.83	0.78	0.78	1.1×
position stitch	118.28	45.41	51.89	2.3×
total	233.44	144.05	92.56	2.5×

The middle column is the bitwise-safe tier (stitching rewrite + inference data movement); the right column adds opt-in float16 inference. (For single reconstructions, fp16 + the numpy stitch backend is marginally faster still — it avoids the one-time GPU-stitch compile — for ~2.7×.)

What changed

1. Position-corrected stitching (learn_stitch): 118 s → 45 s (2.6×)

Replaced three per-item Python loops with vectorized equivalents:

• Patch-position solve: the 2500-iteration Adam optimization is now a single sparse linear solve — the objective is a convex weighted graph-Laplacian least-squares, so we compute the exact minimizer directly. (Also removes the undeclared torch dependency.)
• Subpixel shift refinement: the per-pair upsampled-DFT loop → one batched einsum (factored kernel → shared GEMM).
• Patch shift + accumulate: the per-patch scipy.ndimage.shift loop → batched bilinear shift + scatter.
• Peak finding via a single maximum_filter; cross-correlation FFTs multithreaded (workers=-1).

Output differs from the old Adam result only by a smooth sub-0.1% global gauge term (exact optimum vs. under-converged iterate); the output array size can differ by ≤1 px at the masked border.

2. Inference data movement (bitwise identical)

• On-device gather: move the unique diffraction patterns to the device once and gather each batch's 25-neighbor stack there, instead of rebuilding and re-transferring the ~25×-redundant input every batch (~25 GB → ~1 GB of host→device traffic per scan).
• Flexible batch size (any value; ragged final batch handled).

3. Optional mixed-precision inference (PTYRAN_PRECISION)

Convolutions can run on the tensor cores in float16/bfloat16 (batch-norm stays float32; weights stay float32). Default float32 is bitwise-unchanged. On the L4: float16 = 2.3× faster inference at 0.23% max output difference; bfloat16 = 2.4× at 2.3%. float16 is the recommended fast path; bfloat16 is the fallback if a dataset overflows fp16's range.

precision	inference (s)	speedup	max rel. err vs fp32
float32	18.62	1.00×	0
tf32	18.47	1.01×	6.0e-4
bfloat16	7.85	2.37×	2.3e-2
float16	8.07	2.31×	2.3e-3

4. Optional GPU stitching backend (PTYRAN_STITCH_BACKEND=jax)

Runs the stitching FFTs / peak-finding / refinement / accumulate on the GPU (graph-prune + position solve stay on CPU); matches the numpy result to ~3e-7 (identical shapes and NaN masks). On the L4 this is roughly break-even with the already-vectorized CPU stitching for these sizes, so the numpy backend remains the default; the GPU backend mainly helps CPU-constrained runtimes.

Accuracy & compatibility

• The default path (float32, numpy stitch, and the original host gather all still available) is unchanged. New behavior is entirely opt-in.
• Env vars: PTYRAN_PRECISION = float32|bfloat16|float16; PTYRAN_STITCH_BACKEND = numpy|jax.
• New kwargs: learn_stitch(..., backend=), loop_batch_exp(..., on_device_gather=), Autoencoder(..., compute_dtype=).
• Removed the (undeclared) torch dependency from stitching; added scipy.sparse. Also fixed a degenerate-graph crash (empty shift-vector directions).

Colab testing notebooks

• Baseline: https://drive.google.com/file/d/11Zrg_WM0hMAIAL-_HNWX1QOvlWZYc8p6/view
• Optimized (fp32): https://drive.google.com/file/d/1RufdzkTPWwvcIEhSuY1xJBf8CZfVfBQu/view
• Fast path (fp16 + GPU stitch): https://drive.google.com/file/d/1ab02IeeAVykJ7-8WHFfhcKOoHER3r2fL/view