Project-level shared fits on the fused JAX backend (v0.13.0)

.github/workflows/release.yaml

@@ -30,10 +30,12 @@ jobs:
         with:
           python-version: "3.12"
       - name: Install package + dev tools
+        # [jax] included so the JAX-backend parity tests run and pyright
+        # can resolve the jax imports in eval_jax.py, matching ci.yaml.
         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 ""
       - name: Run ruff

CHANGELOG.md

@@ -7,6 +7,26 @@ 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).
 
+## [0.13.0] - 2026-07-13
+
+### Added
+
+- **Project-level shared fits run on the JAX backend**: with `Project.spec_fun_str = "fit_model_jax"`, `Project.fit_2d()` evaluates all files through one fused jitted program with a joint analytic Jacobian (`jax.jacfwd`) wired into the `leastsq` stages — shared (project-vary) parameter columns span all files exactly instead of via numeric differencing. Measured end-to-end speedup over the interpreter path grows with file count: 2.6x at 2 files to 17x at 8 files (50x60 grid per file), with XLA compile staying under ~0.4 s. Files with different grids fuse fine. If any file's model fails the JAX gate — or jax is not installed — the whole project falls back to the interpreter path (no mixed backends); `show_output >= 1` reports which backend ran. Same caveats as single-file JAX fits: no parallel-worker MCMC, traced-value checks skipped.
+- `Simulator.sigma_data` property: the constant per-point sigma implied by an analog-detection Gaussian simulation, aligned with the fit-side noise schema (`File.set_sigma`), and stored in saved HDF5 metadata (per config group in parameter sweeps, surfaced by `SweepDataset`).
+- A pre-1.0 stability and deprecation policy page in the docs.
+- `AGENTS.md` (repo development) and `llms.txt` (library usage) orientation docs for AI coding agents.
+- `normalize_time` warns when time samples coincide with subcycle boundaries, where the cycle assignment is ambiguous.
+
+### Changed
+
+- **Breaking: `File.create_model_path()` is now `File.model_path()`** and no longer creates directories as a side effect of computing a path — result directories are created on write, so `auto_export=False` runs leave no empty directory trees behind.
+- MCMC (`lmfit.emcee`) with `workers > 1` now runs on a spawn-based process pool instead of the deprecated fork-based default, avoiding deadlocks in multithreaded processes.
+- Simulator noise semantics (`detection` / `noise_type` / `noise_level`) are documented precisely, with per-type sigma formulas.
+
+### Fixed
+
+- JAX-backend Jacobian returned NaN for `sqrtFun` dynamics with a varying onset `t0`.
+
 ## [0.12.0] - 2026-07-11
 
 ### Added

TODO.md

@@ -10,11 +10,10 @@
 
 ## 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 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).
+  - `vmap`-batch homogeneous file series in the fused project fit (unrolled per-file fusion shipped in v0.13.0) — 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.
 - [ ] **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/design/project-level-fits.md

@@ -8,7 +8,8 @@ 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](jax-planning.md) and the
 `eval_jax.py` section of [repo_architecture.md](repo_architecture.md)).
-Nothing here is implemented.
+**Implemented 2026-07-13** as designed — see the implementation status
+and measured results at the end of this note.
 
 ## Current state
 
@@ -86,3 +87,40 @@ evaluates the plans, and the shared-fit workload is where JAX's
   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.

docs/design/repo_architecture.md

@@ -87,7 +87,11 @@ rejected graphs fall back to the compiled NumPy path, never straight to
 the interpreter). Selected via `Project.spec_fun_str = "fit_model_jax"`;
 the Jacobian reaches lmfit's leastsq through `fitlib.jacobian_fun`
 (`Dfun`). Voigt uses a Weideman rational `wofz` approximation instead
-of SciPy's.
+of SciPy's. For project-level shared fits, `make_project_evaluator_2d_jax`
+/ `make_project_jacobian_2d_jax` fuse all per-file plans into one jitted
+program (windowed, flattened, concatenated; `jacfwd` for the joint
+Jacobian) — see
+[project-level-fits.md](project-level-fits.md).
 
 ### `spectra.py` — evaluator bridge
 
@@ -97,7 +101,11 @@ Thin module that the fitting engine calls on every residual evaluation.
 falls back to `fit_model_mcp` — the mcp reference evaluator — when the
 model is not lowerable or when 1D component-wise spectra are requested
 for plotting. Users can swap in a custom spectrum function via
-`Project.spec_fun_str`.
+`Project.spec_fun_str`. Project-level shared fits evaluate through
+`fit_project_mcp` (interpreter, name-based parameter distribution) or
+`fit_project_jax` (fused jitted evaluator; `pack_project_theta`
+converts the combined-parameter mapping into gather index arrays at
+fit setup).
 
 ### `fitlib.py` — lmfit wrappers, CI, MCMC, plotting
 

docs/design/roundtrip_test_matrix.md

@@ -95,13 +95,20 @@ Scope: `SF` (single-file)
 
 Scope: `P` (`Project.fit_2d()`)
 
-Project-level fitting should be tracked separately because it is currently
-wired through `fit_project_mcp`, so the single-file GIR/MCP expectations do not
-apply cleanly yet.
+Project-level fitting should be tracked separately because its backend axis
+differs from the single-file one: with `spec_fun_str = "fit_model_jax"` (and
+every file passing the JAX gate) `Project.fit_2d()` runs the fused JAX path;
+everything else — including the default `fit_model_gir` — evaluates through
+`fit_project_mcp`. There is no project-level GIR path, so the single-file
+`G`/`C` expectations do not apply. Project-scope backend letters:
+
+- `J`: force the fused JAX path and recover truth
+- `CJ`: assert the MCP and JAX project fits agree (end-to-end final
+  parameters)
 
 | Family | 2D | Notes |
 | --- | --- | --- |
-| `PF1` Shared plain dynamics across files | `M` | Current core project roundtrip surface |
+| `PF1` Shared plain dynamics across files | `M/J/CJ` | Core project surface; `J` includes fallback-gating and heterogeneous-grid variants (`TestProjectFitJax`) |
 | `PF2` Project-level expressions | `M` | Includes file/project prefix rewriting and shared refs |
 | `PF3` Shared dynamics with IRF | `M` | Covered with `BiExpProject` + `gaussCONV` |
 | `PF4` Shared subcycle dynamics | `M` | Add once project fixtures exist |
@@ -204,10 +211,16 @@ table above.
   - focused MCMC checks for `MC1`, `MC2`, expression-sensitive `MC2`, and 2D `MC2`
   - focused `W2` coverage for plain and profile-bearing `fit_slice_by_slice()`
   - project-level `M` roundtrips for `PF1`, `PF2`, and `PF3`
+  - project-level `J`/`CJ` for `PF1` (parity, fallback gating, heterogeneous
+    grids in `TestProjectFitJax`); the fused evaluator/Jacobian factories
+    additionally have plan-level unit parity including expressions
+    (`TestProjectFused`)
 
 - Thin or missing today:
   - project-level `PF4` shared subcycle dynamics
-  - project-level `G/C` coverage, because project fitting is still MCP-only
+  - project-level `J`/`CJ` for `PF2` and `PF3` — no end-to-end project JAX
+    roundtrip with expressions or IRF yet (expressions are covered at plan
+    level only)
   - expression-heavy `W2` coverage for `fit_slice_by_slice()`
   - MCMC assertions beyond no-crash / process-boundary coverage
   - exhaustive noisy coverage, intentionally kept out of the main matrix
@@ -221,7 +234,9 @@ The original single-file matrix is implemented. Highest-value next steps:
    up beyond the existing plain/profile cases.
 3. Add lightweight recovery or constraint-preservation assertions to focused
    MCMC tests when runtime allows.
-4. Upgrade project-level cells from `M` to `M/G/C` if project-level GIR lands.
+4. Extend project-level `J`/`CJ` to `PF2` and `PF3`; `G`/`C` stay
+   inapplicable at project scope (the whole-project fallback is the
+   interpreter by design).
 
 ## Non-goals
 

docs/design/supported_models.md

@@ -45,8 +45,14 @@ The sections above describe model semantics. The graph intermediate representati
 - Expression parameters lower only when the expression is arithmetic-only.
   Function calls, attribute access, subscripts, and other non-arithmetic AST
   forms fall back to MCP.
-- Project-level fitting is still wired through ``fit_project_mcp`` even when
-  the underlying per-file models are lowerable.
+- Project-level fitting compiles to the fused JAX backend when
+  ``Project.spec_fun_str == "fit_model_jax"`` and every file's model passes
+  ``can_lower_2d`` + ``can_lower_jax_2d``; any miss — or jax not installed —
+  falls back to ``fit_project_mcp`` for the whole project (no mixed
+  backends). There is no project-level NumPy-plan (GIR) path: with the
+  default ``fit_model_gir``, project fits evaluate through
+  ``fit_project_mcp`` even when the underlying per-file models are
+  lowerable.
 
 ## Notes
 

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "trspecfit"
-version = "0.12.11"
+version = "0.13.0"
 authors = [
     {name = "Johannes Mahl", email = "infinitymonkeyatwork@gmail.com"},
 ]

src/trspecfit/eval_jax.py

@@ -24,7 +24,7 @@
 JAX is an optional dependency: ``pip install "trspecfit[jax]"``.
 """
 
-from collections.abc import Callable
+from collections.abc import Callable, Sequence
 from typing import TYPE_CHECKING
 
 import numpy as np
@@ -705,3 +705,141 @@ def make_jacobian_2d_jax(
     _check_plan_supported(plan)
     jitted = jax.jit(jax.jacfwd(_build_evaluate_2d(plan)))
     return _wrap_jitted(jitted, len(plan.opt_indices))
+
+
+#
+def _build_project_fused_2d(
+    plans: Sequence[ScheduledPlan2D],
+    *,
+    plan_gathers: Sequence[np.ndarray],
+    windows: Sequence[tuple[slice, slice]],
+    n_theta: int,
+) -> Callable:
+    """Build the fused traced function ``theta_c -> concat residual grid``.
+
+    Evaluates every plan from its gathered slice of the combined
+    optimizer vector, applies the per-file fit window as static
+    slices, flattens, and concatenates — one traceable program over
+    all files. Shared by the fused evaluator (jit) and fused Jacobian
+    (jit of jacfwd).
+    """
+
+    builders = [_build_evaluate_2d(plan) for plan in plans]
+    gathers = [np.asarray(g, dtype=np.intp) for g in plan_gathers]
+    for file_idx, (plan, gather) in enumerate(zip(plans, gathers, strict=True)):
+        if len(gather) != len(plan.opt_indices):
+            raise ValueError(
+                f"plan_gathers[{file_idx}] has length {len(gather)}, "
+                f"expected {len(plan.opt_indices)} (plan opt params)."
+            )
+        if len(gather) and (gather.min() < 0 or gather.max() >= n_theta):
+            raise ValueError(
+                f"plan_gathers[{file_idx}] indexes outside the combined "
+                f"theta vector (n_theta={n_theta})."
+            )
+
+    #
+    def _fused(theta_c):
+        pieces = []
+        for build, gather, window in zip(builders, gathers, windows, strict=True):
+            out = build(theta_c[gather])
+            pieces.append(out[window].reshape(-1))
+        return jnp.concatenate(pieces)
+
+    return _fused
+
+
+#
+def make_project_evaluator_2d_jax(
+    plans: Sequence[ScheduledPlan2D],
+    *,
+    plan_gathers: Sequence[np.ndarray],
+    windows: Sequence[tuple[slice, slice]],
+    n_theta: int,
+) -> Callable[[np.ndarray], np.ndarray]:
+    """Compile a fused jitted evaluator over all files of a project fit.
+
+    Parameters
+    ----------
+    plans : sequence of ScheduledPlan2D
+        One compiled plan per file, in project file order. Every
+        source graph must pass ``can_lower_jax_2d``.
+    plan_gathers : sequence of ndarray
+        Per file, positions within the combined optimizer vector of
+        that plan's opt params (from ``spectra.pack_project_theta``).
+    windows : sequence of tuple of slice
+        Per file ``(time_slice, energy_slice)`` fit window, applied to
+        the evaluated grid before flattening (static — no dynamic
+        shapes inside the trace).
+    n_theta : int
+        Length of the combined optimizer vector ``theta_c``.
+
+    Returns
+    -------
+    Callable[[np.ndarray], np.ndarray]
+        ``evaluate(theta_c) -> 1D array``: the windowed, flattened,
+        concatenated model prediction across all files — the fit
+        counterpart of the concatenated data vector assembled by
+        ``Project.fit_2d``.
+
+    Raises
+    ------
+    ImportError
+        If jax is not installed.
+    ValueError
+        If a plan uses features outside the JAX slice, or gathers are
+        inconsistent with the plans.
+    """
+
+    _require_jax()
+    for plan in plans:
+        _check_plan_supported(plan)
+    fused = _build_project_fused_2d(
+        plans, plan_gathers=plan_gathers, windows=windows, n_theta=n_theta
+    )
+    return _wrap_jitted(jax.jit(fused), n_theta)
+
+
+#
+def make_project_jacobian_2d_jax(
+    plans: Sequence[ScheduledPlan2D],
+    *,
+    plan_gathers: Sequence[np.ndarray],
+    windows: Sequence[tuple[slice, slice]],
+    n_theta: int,
+) -> Callable[[np.ndarray], np.ndarray]:
+    """Compile a fused jitted joint Jacobian for a project fit.
+
+    Forward-mode over the fused multi-file function: shared-parameter
+    columns (which cut across every file's rows) come out exactly; the
+    block sparsity of file-level columns is left to XLA. Derivative
+    caveats match ``make_jacobian_2d_jax``.
+
+    Parameters
+    ----------
+    plans, plan_gathers, windows, n_theta
+        As for ``make_project_evaluator_2d_jax``.
+
+    Returns
+    -------
+    Callable[[np.ndarray], np.ndarray]
+        ``jacobian(theta_c) -> (n_residuals_total, n_theta)`` —
+        derivative of the concatenated windowed prediction w.r.t. each
+        combined optimizer parameter, columns in ``theta_c`` order.
+
+    Raises
+    ------
+    ImportError
+        If jax is not installed.
+    ValueError
+        If a plan uses features outside the JAX slice, or gathers are
+        inconsistent with the plans.
+    """
+
+    _require_jax()
+    for plan in plans:
+        _check_plan_supported(plan)
+    fused = _build_project_fused_2d(
+        plans, plan_gathers=plan_gathers, windows=windows, n_theta=n_theta
+    )
+    return _wrap_jitted(jax.jit(jax.jacfwd(fused)), n_theta)