Design Note: Project-Level Shared Fits on a Compiled Backend

Forward-looking note (2026-07-11). Captures the design direction for the “Project-level fit backend” TODO item, written just after the JAX backend landed (see jax-planning.md and the eval_jax.py section of repo_architecture.md). Implemented 2026-07-13 as designed — see the implementation status and measured results at the end of this note.

Current state

Project.fit_2d() already supports Project/File/Static vary levels, but evaluates through fit_project_mcp(): per residual call it distributes the combined parameter vector to each file’s model via name/dict lookups, runs the full interpreter (create_value_2d) per file, applies per-file e_lim/t_lim windows, and concatenates. This is the slowest evaluation path left in the codebase — none of the lowered-evaluator work applies to it yet.

The original TODO question was: lower the multi-file residual to GIR, or keep project-managed per-file loops?

Recommended direction: plans + JAX + joint analytic Jacobian

These are not alternatives — the JAX evaluator consumes ScheduledPlans, so per-file lowering (graphs, plans, and packing the combined-theta -> per-file mapping into index arrays instead of name dicts) is the prerequisite step either way. The question is only what evaluates the plans, and the shared-fit workload is where JAX’s relative edge over a NumPy per-file loop is largest:

1. The Jacobian argument scales with file count. A shared fit’s combined theta is large (~n_files x per-file params, minus shared). Numeric differencing costs one full multi-file evaluation per theta entry per iteration — 10 files x ~6 params is 40+ multi-file evaluations per Jacobian estimate. jax.jacfwd shares the forward computation across all columns inside one XLA program, and the shared parameters’ columns (which cut across every file) come out exactly rather than via noisy differencing. The leastsq nfev collapse measured on single files (8 -> 2, 22 -> 4) multiplies by the per-call cost of evaluating all files.

2. The multi-file residual is one fusable program. Concatenating N independent per-file evaluations is exactly what XLA fuses and internally parallelizes — no Python loop between files. The common shared-fit scenario (same physical model across a measurement series) is the ideal vmap case: identical plan structure, stacked data, one batched call. Unlike batched slice-by-slice fitting (ui.md), none of this needs an optimizer replacement — it is one joint lmfit leastsq with a fused residual and an analytic Dfun, i.e. Phase-D machinery, not Phase E.

3. The baseline is the interpreter, not GIR. Headroom per call is interpreter-vs-JAX (10-70x measured on single files), on top of the Jacobian effect.

Implementation sketch

• Per file: build_graph + schedule_2d (unchanged machinery).

• Pack the vary-level mapping once at fit setup: combined theta -> per-plan theta scatter as index arrays (replacing the per-call name/dict distribution in fit_project_mcp).

• A JAX factory over the list of plans: evaluate each plan’s traced function, apply the per-file windows (static slices — they trace cleanly), flatten and concatenate; jax.jit the whole residual and jax.jacfwd it for the joint Jacobian.

• Wire into the existing flow like the single-file path: a project variant of fit_model_jax plus fitlib.jacobian_fun-style Dfun column reordering against the combined lmfit parameter set.

Caveats and open questions

• Compile time grows with total op count across files; a 10-file program may take a few seconds to compile — amortized over a joint fit, but worth measuring.

• Heterogeneous files (different grids or model structures) forgo vmap batching; unrolled fusion in one program still applies.

• Mixed lowerability: one non-JAX-lowerable file currently implies falling back for the whole project. Mixed-backend execution is explicitly next-track in jax-planning.md; the first implementation should fall back whole-project (JAX -> NumPy plans -> MCP) rather than mix.

• Same constraints as the single-file JAX path: closures do not pickle (no parallel-worker MCMC), and the untraceable runtime value checks are absent.

• Weighted residuals (sigma_type expansion, tracked in TODO.md) should be designed in from the start if it lands first.

Implementation status and measured results (2026-07-13)

Landed exactly per the sketch above (v0.13.0):

• spectra.pack_project_theta packs the combined-parameter mapping into gather index arrays at fit setup.

• eval_jax.make_project_evaluator_2d_jax / make_project_jacobian_2d_jax build one fused jitted program over all per-file plans (windowed static slices, flatten, concatenate; jacfwd for the joint Jacobian).

• spectra.fit_project_jax + fitlib.jacobian_fun_project follow the single-file dispatch conventions; Project.fit_2d gates via Project._build_project_jax_args (every file must pass can_lower_2d + can_lower_jax_2d, jax importable) and falls back whole-project to fit_project_mcp otherwise — no mixed backends, as recommended.

• Heterogeneous grids fuse without vmap (unrolled per-file traces); vmap batching for homogeneous series remains a TODO follow-on.

Benchmark (GLP + mono-exponential x0 dynamics, 50x60 grid per file, shared tau, 2-stage fit with analytic Dfun on the leastsq stage, noiseless synthetic data, CPU):

files	opt params	eval jit	jac jit	jac call	fit JAX	fit MCP	speedup
2	7	0.09 s	0.12 s	0.17 ms	0.25 s	0.65 s	2.6x
4	13	0.10 s	0.16 s	0.38 ms	0.63 s	4.8 s	7.6x
8	25	0.14 s	0.25 s	0.88 ms	3.7 s	63 s	17x

This resolves the compile-time open question: XLA compile grows sub-linearly with file count and stays well under a second at 8 files — negligible against the fit itself. The speedup grows with file count because the interpreter path pays numeric differencing (one full multi-file evaluation per combined theta entry) while the fused Jacobian shares the forward pass across all columns, exactly the scaling argument above.