JAX evaluator backend with analytic Jacobians (v0.12.0)

.claude/skills/benchmark/SKILL.md

@@ -1,6 +1,6 @@
 ---
 name: benchmark
-description: 'Benchmark GIR compiled evaluator vs interpreter. Args: example number (default: 2), `--fit` for full fit, `-n N` for repetitions, `--nfev` for residual-evaluation counts, `--plan-time` for planning vs fit wall time, `--par-variability` for fit robustness vs perturbed initial guesses (`--starts N`).'
+description: 'Benchmark GIR compiled evaluator vs interpreter (and the JAX backend when the [jax] extra is installed). Args: example number (default: 2), `--fit` for full fit, `-n N` for repetitions, `--nfev` for residual-evaluation counts, `--plan-time` for planning vs fit wall time, `--par-variability` for fit robustness vs perturbed initial guesses (`--starts N`).'
 argument-hint: '[example_num] [--fit] [-n N] [--nfev] [--plan-time] [--par-variability] [--starts N]'
 disable-model-invocation: true
 allowed-tools: Bash(.venv/bin/python .claude/skills/benchmark/benchmark_gir.py *)

.claude/skills/benchmark/benchmark_gir.py

@@ -289,6 +289,31 @@ def bench_per_call(file, *, n_warmup=5, n_calls=200):
     slow = spectra.fit_model_mcp(energy, par, True, model, 2)
     max_diff = np.max(np.abs(fast - slow))
 
+    # --- JAX path (optional [jax] extra) ---
+    jax_per_call = None
+    jax_compile = None
+    jax_diff = None
+    jax_note = None
+    try:
+        from trspecfit.eval_jax import make_evaluator_2d_jax
+
+        theta = np.asarray(par, dtype=np.float64)[theta_indices]
+        t0 = time.perf_counter()
+        evaluate_jax = make_evaluator_2d_jax(plan)
+        jax_result = evaluate_jax(theta)  # first call includes XLA compile
+        jax_compile = time.perf_counter() - t0
+        jax_diff = np.max(np.abs(jax_result - fast))
+        for _ in range(n_warmup):
+            evaluate_jax(theta)
+        t0 = time.perf_counter()
+        for _ in range(n_calls):
+            evaluate_jax(theta)
+        jax_per_call = (time.perf_counter() - t0) / n_calls
+    except ImportError:
+        jax_note = "jax not installed (pip install -e '.[jax]')"
+    except ValueError as exc:
+        jax_note = f"plan not JAX-supported ({exc})"
+
     print("=" * 60)
     print("PER-CALL EVALUATION BENCHMARK")
     print(f"  Grid:       {len(file.time)} time x {len(file.energy)} energy")
@@ -297,6 +322,14 @@ def bench_per_call(file, *, n_warmup=5, n_calls=200):
     print(f"  Interpreter:{mcp_per_call * 1e3:8.2f} ms/call")
     print(f"  Speedup:    {mcp_per_call / gir_per_call:8.2f}x")
     print(f"  Max |diff|: {max_diff:.2e}")
+    if jax_per_call is not None:
+        print(f"  JAX:        {jax_per_call * 1e3:8.2f} ms/call")
+        print(f"    vs GIR:   {gir_per_call / jax_per_call:8.2f}x")
+        print(f"    vs MCP:   {mcp_per_call / jax_per_call:8.2f}x")
+        print(f"    compile:  {jax_compile * 1e3:8.0f} ms (once per plan)")
+        print(f"    max |diff| vs GIR: {jax_diff:.2e}")
+    else:
+        print(f"  JAX:        skipped -- {jax_note}")
     print("=" * 60)
 
     return gir_per_call, mcp_per_call
@@ -608,16 +641,28 @@ def capture_par_variability(example_num, *, n_starts=4):
 
 #
 def bench_fit(example_num, dynamics_calls, *, n_reps=3):
-    """Time a complete fit_2d for both paths."""
+    """Time a complete fit_2d for both paths (+ JAX when installed).
 
-    gir_times = []
-    mcp_times = []
+    The JAX column runs ``spec_fun_str="fit_model_jax"`` -- the jitted
+    evaluator plus the analytic Jacobian for the leastsq stage. Each
+    rep rebuilds the file, so every JAX rep pays its own XLA compile
+    (the realistic single-fit cost, not a warm-cache best case).
+    """
+
+    try:
+        import jax  # noqa: F401
+
+        have_jax = True
+    except ImportError:
+        have_jax = False
+
+    backends = [("fit_model_gir", "GIR"), ("fit_model_mcp", "MCP")]
+    if have_jax:
+        backends.append(("fit_model_jax", "JAX"))
+    times: dict[str, list] = {label: [] for _, label in backends}
 
     for i in range(n_reps):
-        for fun_str, times_list in [
-            ("fit_model_gir", gir_times),
-            ("fit_model_mcp", mcp_times),
-        ]:
+        for fun_str, label in backends:
             file, _, _ = load_example(example_num, add_dynamics=False)
             file.define_baseline(
                 time_start=0, time_stop=10, time_type="ind", show_plot=False
@@ -628,17 +673,15 @@ def bench_fit(example_num, dynamics_calls, *, n_reps=3):
             file.p.spec_fun_str = fun_str
             t0 = time.perf_counter()
             file.fit_2d(model_name="2D", stages=2, try_ci=0)
-            times_list.append(time.perf_counter() - t0)
+            times[label].append(time.perf_counter() - t0)
 
         print(
             f"  Rep {i + 1}/{n_reps}: "
-            f"GIR={gir_times[-1]:.2f}s  MCP={mcp_times[-1]:.2f}s"
+            + "  ".join(f"{label}={times[label][-1]:.2f}s" for _, label in backends)
         )
 
-    gir_med = np.median(gir_times)
-    mcp_med = np.median(mcp_times)
-    formatted_gir_times = [f"{t:.2f}" for t in gir_times]
-    formatted_mcp_times = [f"{t:.2f}" for t in mcp_times]
+    gir_med = np.median(times["GIR"])
+    mcp_med = np.median(times["MCP"])
 
     print()
     print("=" * 60)
@@ -647,8 +690,15 @@ def bench_fit(example_num, dynamics_calls, *, n_reps=3):
     print(f"  GIR:        {gir_med:8.2f} s (median)")
     print(f"  Interpreter:{mcp_med:8.2f} s (median)")
     print(f"  Speedup:    {mcp_med / gir_med:8.2f}x")
-    print(f"  GIR all:    {formatted_gir_times}")
-    print(f"  MCP all:    {formatted_mcp_times}")
+    if have_jax:
+        jax_med = np.median(times["JAX"])
+        print(f"  JAX+Dfun:   {jax_med:8.2f} s (median, incl. XLA compile)")
+        print(f"    vs GIR:   {gir_med / jax_med:8.2f}x")
+    else:
+        print("  JAX+Dfun:   skipped -- jax not installed")
+    for _, label in backends:
+        formatted = [f"{t:.2f}" for t in times[label]]
+        print(f"  {label} all:    {formatted}")
     print("=" * 60)
 
 

.github/workflows/ci.yaml

@@ -28,10 +28,14 @@ jobs:
           python-version: ${{ matrix.python-version }}
 
       - name: Install package + dev tools
+        # [jax] included so the JAX-backend parity tests run here; the
+        # min-versions job stays jax-free (jax needs numpy>=2.0, which
+        # cannot exist in the numpy==1.26 floor world) and those tests
+        # skip there via pytest.importorskip.
         run: |
           python -m venv .venv
           .venv/bin/python -m pip install --upgrade pip
-          .venv/bin/pip install -e ".[dev]"
+          .venv/bin/pip install -e ".[dev,jax]"
 
       - name: Run tests
         run: .venv/bin/pytest -q -m ""

CHANGELOG.md

@@ -7,10 +7,25 @@ and this project adheres to [Semantic Versioning](https://semver.org/).
 This file is maintained using the shared changelog workflow in
 [`docs/ai/changelog.md`](docs/ai/changelog.md).
 
-## [Unreleased]
+## [0.12.0] - 2026-07-11
+
+### Added
+
+- **Experimental JAX evaluator backend for 2D fits**: set `Project.spec_fun_str = "fit_model_jax"` (requires the optional extra: `pip install "trspecfit[jax]"`). The jitted evaluator covers the full compiled 2D model surface — dynamics, expressions, subcycle dynamics, parameter profiles, IRF convolution, and every lineshape — with measured per-call speedups of 2.4–70x over the compiled NumPy evaluator (largest for `Voigt`). Models the JAX gate rejects fall back to the compiled NumPy path automatically; 1D fits always use the NumPy path.
+- **Analytic Jacobians on the JAX backend**: `File.fit_2d` derives the model Jacobian with `jax.jacfwd` and wires it into lmfit's `leastsq` stages (`Dfun`), cutting residual evaluations roughly 4x at identical optima — a full Voigt 2D fit dropped from 25.3 s to 7.1 s end to end. Pass `fit_2d(..., jac_fun=None)` to disable.
+- `Voigt` on the JAX path evaluates the Faddeeva function via a Weideman rational approximation, matching `scipy.special.wofz` to better than 1e-12 relative over the physical parameter domain.
+- The benchmark workflow (`/benchmark`) reports a JAX column (per-call and `--fit`) whenever the `[jax]` extra is installed.
+
+### Changed
+
+- **JAX backend caveats to be aware of** (documented in `eval_jax.py`): parameter-value checks that require host access are skipped on traced values (`LinBack` ordering, convolution-kernel positivity — keep fit bounds ordered and widths positive), `lmfit.emcee` with `workers > 1` is not supported on this path, and importing the backend enables JAX 64-bit mode globally.
+- `ScheduledPlan1D`/`ScheduledPlan2D` store compiled expressions as packed CSR instruction arrays (`expr_instructions` + `expr_indptr`) instead of a list of `ExprProgram` objects (which is gone); execution plans are now fully array-native. Internal module, but breaking for code that inspected plan internals.
+
+## [0.11.0] - 2026-07-11
 
 ### Changed
 
+- **Reference lines default to black** (`refline_color`: grey `#808080` -> black `#000000`) for better contrast on 2D colormaps, and 2D fit-result limit lines now honor `refline_color`/`refline_style` from `project.yaml`/`PlotConfig` instead of being hardcoded.
 - **Breaking: IRF convolution is now a quadrature-weighted kernel-matrix operator** (`conv_matrix_operator` / `conv_matrix_apply` in `utils/arrays.py`), replacing the sampled-1D-kernel convolution (`my_conv` + per-theta kernel support) on both the mcp and compiled (GIR) evaluation paths. Consequences:
   - **Non-uniform time axes are convolved correctly.** The axis enters through the time-difference matrix and trapezoid quadrature weights, so measured fine-around-t0 + coarse-tail delay axes no longer produce silently biased fits (the old sample-index convolution deviated by up to ~7% of the trace maximum on example 21's stepped axis). Uniform axes agree with the previous path within its documented kernel-truncation tolerance.
   - **Exact edge handling for any kernel width**: kernel mass beyond the time window contributes through the edge samples (edge-value padding semantics), computed analytically per kernel via edge-mass companions (`functions/time.py::CONV_EDGE_MASS`, cancellation-safe erfc/exp/clip tail forms). No support truncation at any width or axis spacing; a kernel without a registered companion is rejected at model validation (mcp) and scheduling (GIR).

CLAUDE.md

@@ -1,7 +1,7 @@
 # General behavior
 
 - **Confidence Rule:** Do not make changes until you have 95% confidence. Understand the relevant files before editing and ask follow-up questions until you reach this threshold or when tradeoffs are non-obvious.
-- **Commits:** Never commit or rewrite history unless explicitly asked. Always show the exact commit message and wait for approval before committing. Never add AI attribution trailers (`Co-Authored-By` etc.) to commit messages.
+- **Commits:** Never commit or rewrite history unless explicitly asked. Always show the exact commit message and wait for approval before committing. Never add AI attribution trailers (`Co-Authored-By` etc.) to commit messages. When suggesting a commit message, check whether `version` in `pyproject.toml` has already been incremented relative to `main`; if not, suggest an appropriate increment alongside the message.
 - **Context Discipline:** Monitor context usage. At 60% usage (or if it starts getting tight), summarize progress and prompt me to `/compact` or `/clear`.
 - **Token Efficiency:** Be concise. Reference file paths and line numbers rather than quoting large code blocks.
 - **Subagent Protocol:** Use subagents for repo-wide scans, parallel research, or scanning large directories. Instruct them to return only concise summaries to keep the main context window lean.

TODO.md

@@ -12,8 +12,12 @@
 
 ## Performance & architecture
 
-- [ ] **Project-level fit backend**: `Project.fit_2d()` already supports `Project`/`File`/`Static` vary levels, but it currently evaluates through `fit_project_mcp()` and `Model.create_value_2d()` rather than the GIR scheduler/evaluator path. Decide whether to lower the multi-file residual to GIR or explicitly prefer project-managed per-file loops when we want maximum graph-IR speedups.
-- [ ] **JAX backend / Jacobian follow-on**: if we revisit a JAX evaluator, analytic Jacobians, or optimizer replacement, use [docs/design/jax-planning.md](docs/design/jax-planning.md) as the roadmap for scope, sequencing, and open technical constraints.
+- [ ] **Project-level fit backend**: `Project.fit_2d()` already supports `Project`/`File`/`Static` vary levels, but it currently evaluates through `fit_project_mcp()` and `Model.create_value_2d()` rather than any compiled path. Design direction settled (2026-07-11): lower per-file plans and evaluate them through a fused JAX residual with a joint analytic Jacobian — see [docs/design/project-level-fits.md](docs/design/project-level-fits.md) for the rationale, implementation sketch, and open questions (compile-time scaling, mixed lowerability, vmap batching for homogeneous series).
+- [ ] **JAX backend follow-ons**: the backend itself shipped in v0.12.0 (Phases A–D of [docs/design/jax-planning.md](docs/design/jax-planning.md); execution record in [docs/design/archive/jax-backend.md](docs/design/archive/jax-backend.md)). Remaining candidates, none scheduled:
+  - Full-parameter-vector evaluator variant for interactive use (fixed-value edits without recompile) plus session-level evaluator caching — see [docs/design/ui.md](docs/design/ui.md).
+  - vmap-batched slice-by-slice solver (the one workload where lmfit overhead plausibly dominates; would be the Phase E pilot) — see [docs/design/ui.md](docs/design/ui.md).
+  - Project-level fused shared fits with joint analytic Jacobian — see [docs/design/project-level-fits.md](docs/design/project-level-fits.md).
+  - `fit_model_compare`-style runtime JAX parity mode, or a cheaper one-shot pre-fit parity check on the JAX path.
 - [ ] **MCMC multiprocessing context**: `lmfit.emcee(workers=N)` currently inherits Python's default multiprocessing start method, which triggers a Python 3.12 `fork()` deprecation warning in multithreaded test runs. Investigate whether we can supply a `spawn`-backed worker pool or otherwise steer emcee/lmfit away from raw `fork`.
 - [ ] **Define the results-data ownership boundary**: take a look at what should live as class attributes on the `trspecfit`/`mcp` Python classes (`File`/`Model`) versus inside the `FitResults` class. Where should the line be — should fit outputs (params, `conf_ci`, MCMC payload, correlations, acceptance fraction, diagnostics) all be unified under `FitResults`, or stay split between live `model.result[...]` and persisted slots? Then update all callers and the `get_*` accessor methods to match the chosen boundary. Sub-items:
   - **Persist `correl` and `acceptance_fraction` into the slots**: 2026-06 added live-only accessors (`get_correlations`, `get_conf_intervals`, `get_mcmc`) reading `model.result` as a stopgap for notebook 12, so these are NOT yet saved. Add per-parameter correlations to the slot `params` payload and `acceptance_fraction` to the slot `mcmc` payload, with `.fit.h5` read/write support and save/load round-trip tests, so they survive persistence like the rest of the slot.

docs/ai/benchmark.md

@@ -1,10 +1,16 @@
-# Benchmark GIR vs Interpreter
+# Benchmark GIR vs Interpreter (vs JAX)
 
 Shared source of truth for benchmarking the compiled GIR evaluator against the
 interpreter (MCP) path.
 
 Run `benchmark_gir.py` to compare the compiled and interpreter evaluation paths
-on an example fitting workflow.
+on an example fitting workflow. When the optional `[jax]` extra is installed
+(`.venv/bin/pip install -e ".[jax]"`), the per-call benchmark adds a JAX
+column (jitted evaluator; reports per-call time, speedup vs both paths, the
+one-time XLA compile cost, and max |diff| vs GIR) and `--fit` adds a
+`fit_model_jax` run (jitted residuals + analytic Jacobian on the leastsq
+stage; each rep pays its own compile). Without jax both report
+"skipped".
 
 ## Available examples
 
@@ -50,9 +56,10 @@ Run:
 .venv/bin/python .claude/skills/benchmark/benchmark_gir.py --example <N> --calls 200 [--fit] [-n <N>]
 ```
 
-Report the results to the user. Highlight the speedup ratio, the
-`Max |diff|` correctness check, and note which GIR path the example exercises
-(convolution / subcycle / profile / plain).
+Report the results to the user. Highlight the speedup ratios, the
+`Max |diff|` correctness checks, and note which GIR path the example exercises
+(convolution / subcycle / profile / plain). When the JAX column is present,
+mention the compile cost separately — it is paid once per plan, not per call.
 
 ## Fit-count and planning-cost modes
 

docs/api/eval_jax.rst

@@ -0,0 +1,7 @@
+JAX Evaluator (experimental)
+============================
+
+.. automodule:: trspecfit.eval_jax
+   :members:
+   :undoc-members:
+   :show-inheritance:

docs/api/index.rst

@@ -18,3 +18,4 @@ This section contains the auto-generated API documentation.
    graph_ir
    eval_1d
    eval_2d
+   eval_jax