diff --git a/.github/workflows/release.yaml b/.github/workflows/release.yaml index b20b561..6d185ed 100644 --- a/.github/workflows/release.yaml +++ b/.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 diff --git a/CHANGELOG.md b/CHANGELOG.md index cdd888f..3276468 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/TODO.md b/TODO.md index 88894a5..c78de8c 100644 --- a/TODO.md +++ b/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. diff --git a/docs/design/project-level-fits.md b/docs/design/project-level-fits.md index d79b547..ef98e77 100644 --- a/docs/design/project-level-fits.md +++ b/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. diff --git a/docs/design/repo_architecture.md b/docs/design/repo_architecture.md index 1f5c8d3..bc2e23a 100644 --- a/docs/design/repo_architecture.md +++ b/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 diff --git a/docs/design/roundtrip_test_matrix.md b/docs/design/roundtrip_test_matrix.md index a20d1ed..c20c699 100644 --- a/docs/design/roundtrip_test_matrix.md +++ b/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 diff --git a/docs/design/supported_models.md b/docs/design/supported_models.md index ff23cc9..129d2af 100644 --- a/docs/design/supported_models.md +++ b/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 diff --git a/pyproject.toml b/pyproject.toml index a4ab547..369a2dc 100644 --- a/pyproject.toml +++ b/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"}, ] diff --git a/src/trspecfit/eval_jax.py b/src/trspecfit/eval_jax.py index 286ef06..8751bc5 100644 --- a/src/trspecfit/eval_jax.py +++ b/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) diff --git a/src/trspecfit/fitlib.py b/src/trspecfit/fitlib.py index d20e1c2..6024668 100644 --- a/src/trspecfit/fitlib.py +++ b/src/trspecfit/fitlib.py @@ -352,6 +352,65 @@ def jacobian_fun( return d_res[:, columns] +# +def jacobian_fun_project( + par: Any, + x: ArrayLike, + data: np.ndarray, + fit_fun_str: str, + unpack: int = 0, + e_lim: list[int] | None = None, + t_lim: list[int] | None = None, + res_type: str = "lmfit", + args: Sequence[Any] | None = None, +) -> np.ndarray: + """Analytic joint Jacobian for project-level fits (lmfit ``Dfun``). + + Project counterpart of :func:`jacobian_fun`. Requires the + ``fit_project_jax`` dispatch convention: + ``args = (evaluator, jacobian, theta_c_indices, var_names, dim)`` + with *jacobian* from ``eval_jax.make_project_jacobian_2d_jax``. + Per-file fit windows are applied inside the fused jacobian, so no + window slicing happens here (``e_lim``/``t_lim`` are empty for + project fits). + + Returns + ------- + ndarray + ``d(residual)/d(varying params)``, shape + ``(n_residuals_total, n_varys)``, columns in lmfit + varying-parameter order (``col_deriv=0``). Residual is + ``data - fit``, so this is the negated fused model Jacobian. + """ + + if args is None or not callable(args[1]): + raise ValueError( + "jacobian_fun_project requires the fit_project_jax dispatch " + "args (evaluator, jacobian, theta_c_indices, var_names, dim)." + ) + jacobian = args[1] + theta_c_indices: np.ndarray = args[2] + theta_names: list[str] = list(args[3]) + + par_values = np.asarray( + ulmfit.par_extract(par, return_type="list"), dtype=np.float64 + ) + # (n_residuals_total, n_opt), columns in theta_c order + jac = np.asarray(jacobian(par_values[theta_c_indices]), dtype=np.float64) + d_res = -jac + + # Column order: theta_c order -> lmfit varying-parameter order. + var_names = [name for name in par if par[name].vary] + if sorted(theta_names) != sorted(var_names): + raise RuntimeError( + "Project JAX Jacobian column mismatch: combined optimizer " + f"parameters {theta_names} do not match lmfit varying " + f"parameters {var_names}." + ) + columns = [theta_names.index(name) for name in var_names] + return d_res[:, columns] + + # def time_display( t_start: float, print_str: str = "", *, return_delta_seconds: bool = False diff --git a/src/trspecfit/spectra.py b/src/trspecfit/spectra.py index 791607f..0790874 100644 --- a/src/trspecfit/spectra.py +++ b/src/trspecfit/spectra.py @@ -423,3 +423,143 @@ def fit_project_mcp( slices.append(fit_2d.flatten()) return np.concatenate(slices) + + +# +def pack_project_theta( + plans: Sequence[ScheduledPlan2D], + *, + mapping: list[tuple[str, int, str]], + par_names: list[str], + var_names: list[str], +) -> tuple[np.ndarray, list[np.ndarray]]: + """Pack the combined-parameter mapping into gather index arrays. + + Converts the name-based parameter distribution of project-level + fitting (``project_fit_info["mapping"]`` from + ``Project._build_fit_params``) into integer index arrays so the + fused evaluator can scatter the combined optimizer vector to + per-file plan thetas with plain array gathers — no name lookups + per residual call. + + Parameters + ---------- + plans : sequence of ScheduledPlan2D + One compiled plan per file, in project file order. Each plan's + ``opt_param_names`` are local (per-file) parameter names. + mapping : list of tuple + ``(combined_name, file_idx, local_name)`` triples from + ``Project._build_fit_params``. + par_names : list of str + All combined parameter names, in combined-vector order. + var_names : list of str + Varying combined parameter names, ordered as in ``par_names``. + Defines the combined theta vector ``theta_c``. + + Returns + ------- + theta_c_indices : ndarray + Positions of ``var_names`` within ``par_names``: + ``theta_c = par_full[theta_c_indices]``. + plan_gathers : list of ndarray + Per file, positions within ``theta_c`` of that plan's opt + params: ``theta_f = theta_c[plan_gathers[file_idx]]`` yields + the plan's theta in ``opt_param_names`` order. + + Raises + ------ + RuntimeError + If a plan opt param has no combined counterpart, its + counterpart is not varying, or a varying combined param feeds + no plan — all indicate plans built with vary flags that do not + match the combined parameter set. + """ + + remaps: list[dict[str, str]] = [{} for _ in plans] + for combined_name, file_idx, local_name in mapping: + remaps[file_idx][local_name] = combined_name + + full_pos = {name: i for i, name in enumerate(par_names)} + var_pos = {name: i for i, name in enumerate(var_names)} + theta_c_indices = np.array([full_pos[name] for name in var_names], dtype=np.intp) + + consumed: set[str] = set() + plan_gathers: list[np.ndarray] = [] + for file_idx, plan in enumerate(plans): + remap = remaps[file_idx] + gather: list[int] = [] + for local_name in plan.opt_param_names: + mapped_name = remap.get(local_name) + if mapped_name is None: + raise RuntimeError( + f"Plan opt param '{local_name}' (file {file_idx}) has " + f"no combined-parameter mapping entry." + ) + pos = var_pos.get(mapped_name) + if pos is None: + raise RuntimeError( + f"Plan opt param '{local_name}' (file {file_idx}) maps " + f"to combined param '{mapped_name}', which is not " + f"varying." + ) + gather.append(pos) + consumed.add(mapped_name) + plan_gathers.append(np.array(gather, dtype=np.intp)) + + leftover = [name for name in var_names if name not in consumed] + if leftover: + raise RuntimeError( + f"Varying combined params feed no plan: {leftover}. " + f"Plans were likely built with stale vary flags." + ) + + return theta_c_indices, plan_gathers + + +# +def fit_project_jax( + x: Sequence[float] | np.ndarray, + par: Sequence[float] | np.ndarray, + plot_sum: bool, + *args: Any, +) -> np.ndarray: + """Generate concatenated multi-file prediction via the fused JAX evaluator. + + Project-fit counterpart of :func:`fit_model_jax`. The fused + evaluator applies per-file fit windows, flattens, and concatenates + internally, so the return value aligns element-for-element with + the concatenated data vector assembled by ``Project.fit_2d`` — no + slicing in ``fitlib.residual_fun`` (empty ``e_lim``/``t_lim``). + + Unlike ``fit_model_jax`` there is no fallback branch: ``Project`` + dispatches here only when every file's graph passed the JAX gate; + otherwise the const carries ``fit_project_mcp``. + + Parameters + ---------- + x : array-like + Unused (kept for fit-function signature compatibility). + par : array-like + Full combined parameter vector (varying + static + expr), in + ``project_fit_info["par_names"]`` order. + plot_sum : bool + Unused (the fused path always returns the sum). + *args + ``(evaluator, jacobian, theta_c_indices, var_names, dim)`` — + *evaluator*/*jacobian* from + ``eval_jax.make_project_evaluator_2d_jax`` / + ``make_project_jacobian_2d_jax``; *jacobian* and *var_names* + are carried for ``fitlib.jacobian_fun_project`` (lmfit + ``Dfun``), not used here. + + Notes + ----- + The evaluator/jacobian entries are fused closures and do not + pickle; MCMC via ``lmfit.emcee`` with ``workers > 1`` is not + supported on this path (single-worker MCMC works). + """ + + evaluator = args[0] + theta_c_indices: np.ndarray = args[2] + par_arr = np.asarray(par, dtype=np.float64) + return np.asarray(evaluator(par_arr[theta_c_indices])) diff --git a/src/trspecfit/trspecfit.py b/src/trspecfit/trspecfit.py index d2c0f7d..17f3013 100644 --- a/src/trspecfit/trspecfit.py +++ b/src/trspecfit/trspecfit.py @@ -1243,6 +1243,86 @@ def _build_fit_params( return combined_pars, project_fit_info + # + def _build_project_jax_args( + self, + *, + combined_pars: lmfit.Parameters, + project_fit_info: dict, + windows: list[tuple[slice, ...]], + ) -> tuple[Any, ...] | None: + """ + Build the fused-JAX dispatch args for a project fit, or None. + + Gates the fused fast path: every file's graph must lower to a + scheduled 2D plan inside the JAX slice, and jax must be + importable. Any miss returns None and the whole project falls + back to the interpreter path (no mixed backends). + + Parameters + ---------- + combined_pars : lmfit.Parameters + Combined optimizer parameters from ``_build_fit_params``. + project_fit_info : dict + Context dict from ``_build_fit_params``. + windows : list of tuple of slice + Per-file ``(time_slice, energy_slice)`` fit windows. + + Returns + ------- + tuple or None + ``(evaluator, jacobian, theta_c_indices, var_names, 2)`` + following the ``spectra.fit_project_jax`` args convention, + or None when the project cannot use the fused path. + """ + + from trspecfit import spectra + from trspecfit.graph_ir import ( + build_graph, + can_lower_2d, + can_lower_jax_2d, + schedule_2d, + ) + + plans = [] + for model in project_fit_info["models"]: + graph = build_graph(model) + if not (can_lower_2d(graph) and can_lower_jax_2d(graph)): + return None + plans.append(schedule_2d(graph)) + + par_names = project_fit_info["par_names"] + var_names = [name for name in par_names if combined_pars[name].vary] + theta_c_indices, plan_gathers = spectra.pack_project_theta( + plans, + mapping=project_fit_info["mapping"], + par_names=par_names, + var_names=var_names, + ) + + try: + from trspecfit.eval_jax import ( + make_project_evaluator_2d_jax, + make_project_jacobian_2d_jax, + ) + + evaluator = make_project_evaluator_2d_jax( + plans, + plan_gathers=plan_gathers, + windows=cast("list[tuple[slice, slice]]", windows), + n_theta=len(var_names), + ) + jacobian = make_project_jacobian_2d_jax( + plans, + plan_gathers=plan_gathers, + windows=cast("list[tuple[slice, slice]]", windows), + n_theta=len(var_names), + ) + except ImportError: + return None + + return (evaluator, jacobian, theta_c_indices, var_names, 2) + # def fit_2d( self, @@ -1300,32 +1380,42 @@ def fit_2d( model_name=model_name, ) - # Build concatenated data array (sliced per file) + # Build concatenated data array (sliced per file); the same + # windows feed the fused JAX evaluator so data and prediction + # vectors align element-for-element. + windows = [fitlib._fit_window_slices(2, f.e_lim, f.t_lim) for f in self.files] data_slices: list[np.ndarray] = [] - for f in self.files: + for f, window in zip(self.files, windows, strict=True): assert f.data is not None # type guard - d = f.data - if f.e_lim and f.t_lim: - d = d[f.t_lim[0] : f.t_lim[1], f.e_lim[0] : f.e_lim[1]] - elif f.e_lim: - d = d[:, f.e_lim[0] : f.e_lim[1]] - elif f.t_lim: - d = d[f.t_lim[0] : f.t_lim[1], :] - data_slices.append(d.flatten()) + data_slices.append(f.data[window].flatten()) concat_data = np.concatenate(data_slices) - # const: (x, data, package, fit_fun_str, unpack, e_lim, t_lim) - # e_lim and t_lim are empty — slicing is handled inside - # fit_project_mcp + # --- dispatch: fused JAX fast path vs interpreter --- + fit_fun_str = "fit_project_mcp" + args: tuple[Any, ...] = (project_fit_info, 2) + if self.spec_fun_str == "fit_model_jax": + jax_args = self._build_project_jax_args( + combined_pars=combined_pars, + project_fit_info=project_fit_info, + windows=windows, + ) + if jax_args is not None: + fit_fun_str = "fit_project_jax" + args = jax_args + # Analytic joint Jacobian for leastsq stages (lmfit Dfun) + fit_wrapper_kwargs.setdefault("jac_fun", fitlib.jacobian_fun_project) + + # const: (x, data, fit_fun_str, unpack, e_lim, t_lim) + # e_lim and t_lim are empty — per-file windows are applied + # inside fit_project_mcp / the fused evaluator. const: tuple[Any, ...] = ( - np.array([]), # x — unused by fit_project_mcp + np.array([]), # x — unused by the project fit functions concat_data, - "fit_project_mcp", + fit_fun_str, 0, [], # no additional e_lim slicing [], # no additional t_lim slicing ) - args = (project_fit_info, 2) par_names = project_fit_info["par_names"] @@ -1341,10 +1431,11 @@ def fit_2d( if p.vary and p.name.startswith("file") ) n_static = sum(1 for p in combined_pars.values() if not p.vary) + backend = "JAX" if fit_fun_str == "fit_project_jax" else "interpreter" print( f"Project fit: {len(self.files)} files, " f"{n_project} project-vary, {n_file} file-vary, " - f"{n_static} static params" + f"{n_static} static params ({backend} backend)" ) result = fitlib.fit_wrapper( @@ -1399,8 +1490,9 @@ def fit_2d( pd.DataFrame(), ] # const/args mirror File.fit_2d so _append_2d_slot can evaluate - # the per-file fit grid via fitlib.residual_fun. Project fits - # always go through the MCP path. + # the per-file fit grid via fitlib.residual_fun. Per-file + # re-evaluation always uses the interpreter — the fused JAX + # closures cover the whole project, not single files. model.const = (f.energy, f.data, "fit_model_mcp", 0, f.e_lim, f.t_lim) model.args = (model, 2) diff --git a/tests/test_evaluate_jax.py b/tests/test_evaluate_jax.py index d6ee2f8..14af034 100644 --- a/tests/test_evaluate_jax.py +++ b/tests/test_evaluate_jax.py @@ -32,6 +32,8 @@ from trspecfit.eval_jax import ( # noqa: E402 make_evaluator_2d_jax, make_jacobian_2d_jax, + make_project_evaluator_2d_jax, + make_project_jacobian_2d_jax, ) from trspecfit.graph_ir import ( # noqa: E402 build_graph, @@ -370,3 +372,146 @@ def test_fit_recovers_truth_with_analytic_jacobian(self): assert np.isclose(true_val, fit_val, rtol=1e-8, atol=1e-10), ( f"{name}: true={true_val:.6f}, fit={fit_val:.6f}" ) + + +# --------------------------------------------------------------------------- +# Project-level fused evaluator / Jacobian +# --------------------------------------------------------------------------- + + +# +def _make_project_pair(): + """Two identical-structure plans sharing theta_c slot 0. + + Mimics a 2-file project fit: opt param 0 is project-vary (one + combined slot feeding both plans), the rest are file-vary. File 0 + gets a proper fit window, file 1 is unwindowed. + """ + + plan_a, model_a = _make_plan(["glp_expression"], [("GLP_01_A", ["MonoExpPos"])]) + plan_b, model_b = _make_plan(["glp_expression"], [("GLP_01_A", ["MonoExpPos"])]) + theta_a = _extract_theta(plan_a, model_a) + theta_b = theta_a * 1.1 + 0.02 + n = len(theta_a) + theta_c = np.concatenate([theta_a, theta_b[1:]]) + gather_a = np.arange(n, dtype=np.intp) + gather_b = np.concatenate([[0], np.arange(n, 2 * n - 1)]).astype(np.intp) + windows = [ + (slice(2, len(plan_a.time) - 3), slice(5, len(plan_a.energy) - 7)), + (slice(None), slice(None)), + ] + return [plan_a, plan_b], [gather_a, gather_b], windows, theta_c + + +# +def _fused_reference(plans, gathers, windows, theta_c): + """NumPy reference: per-plan evaluate_2d, windowed, flattened, concat.""" + + pieces = [ + evaluate_2d(plan, theta_c[gather])[window].ravel() + for plan, gather, window in zip(plans, gathers, windows, strict=True) + ] + return np.concatenate(pieces) + + +# +# +class TestProjectFused: + """Fused multi-file evaluator and joint Jacobian.""" + + # + def test_evaluator_parity(self): + """Fused output matches per-file NumPy reference (incl. jit reuse).""" + + plans, gathers, windows, theta_c = _make_project_pair() + fused = make_project_evaluator_2d_jax( + plans, plan_gathers=gathers, windows=windows, n_theta=len(theta_c) + ) + + ref = _fused_reference(plans, gathers, windows, theta_c) + got = fused(theta_c) + assert got.ndim == 1 + np.testing.assert_allclose(got, ref, rtol=1e-12, atol=1e-12) + + theta_new = theta_c * 1.05 + 0.01 + ref_new = _fused_reference(plans, gathers, windows, theta_new) + np.testing.assert_allclose(fused(theta_new), ref_new, rtol=1e-12, atol=1e-12) + + # + def test_jacobian_matches_fd(self): + """Joint Jacobian vs central differences of the NumPy reference.""" + + plans, gathers, windows, theta_c = _make_project_pair() + jacobian = make_project_jacobian_2d_jax( + plans, plan_gathers=gathers, windows=windows, n_theta=len(theta_c) + ) + jac = jacobian(theta_c) + n_res = len(_fused_reference(plans, gathers, windows, theta_c)) + assert jac.shape == (n_res, len(theta_c)) + + f_scale = np.max(np.abs(_fused_reference(plans, gathers, windows, theta_c))) + for i in range(len(theta_c)): + h = 1e-6 * max(1.0, abs(theta_c[i])) + theta_plus = theta_c.copy() + theta_plus[i] += h + theta_minus = theta_c.copy() + theta_minus[i] -= h + fd = ( + _fused_reference(plans, gathers, windows, theta_plus) + - _fused_reference(plans, gathers, windows, theta_minus) + ) / (2 * h) + tol = 1e-5 * np.max(np.abs(fd)) + 1e-8 * f_scale + max_err = np.max(np.abs(jac[:, i] - fd)) + assert max_err <= tol, ( + f"Joint Jacobian column {i}: " + f"max |analytic - fd| = {max_err:.3e} exceeds {tol:.3e}" + ) + + # + def test_shared_column_spans_both_files(self): + """The shared slot's column is nonzero in both files' row blocks.""" + + plans, gathers, windows, theta_c = _make_project_pair() + jacobian = make_project_jacobian_2d_jax( + plans, plan_gathers=gathers, windows=windows, n_theta=len(theta_c) + ) + jac = jacobian(theta_c) + + n_rows_a = evaluate_2d(plans[0], theta_c[gathers[0]])[windows[0]].size + shared_col = jac[:, 0] + assert np.any(shared_col[:n_rows_a] != 0.0) + assert np.any(shared_col[n_rows_a:] != 0.0) + # file-vary slots touch only their own file's rows + n_opt = len(gathers[0]) + file_a_col = jac[:, 1] # slot 1: file A only + file_b_col = jac[:, n_opt] # first file-B-only slot + assert np.all(file_a_col[n_rows_a:] == 0.0) + assert np.all(file_b_col[:n_rows_a] == 0.0) + + # + def test_gather_length_mismatch_raises(self): + plans, gathers, windows, theta_c = _make_project_pair() + bad = [gathers[0][:-1], gathers[1]] + with pytest.raises(ValueError, match="expected"): + make_project_evaluator_2d_jax( + plans, plan_gathers=bad, windows=windows, n_theta=len(theta_c) + ) + + # + def test_gather_out_of_range_raises(self): + plans, gathers, windows, theta_c = _make_project_pair() + bad = [gathers[0], gathers[1].copy()] + bad[1][0] = len(theta_c) + with pytest.raises(ValueError, match="outside"): + make_project_evaluator_2d_jax( + plans, plan_gathers=bad, windows=windows, n_theta=len(theta_c) + ) + + # + def test_theta_length_mismatch_raises(self): + plans, gathers, windows, theta_c = _make_project_pair() + fused = make_project_evaluator_2d_jax( + plans, plan_gathers=gathers, windows=windows, n_theta=len(theta_c) + ) + with pytest.raises(ValueError, match="does not match"): + fused(theta_c[:-1]) diff --git a/tests/test_project_fit.py b/tests/test_project_fit.py index e04e0b9..793083e 100644 --- a/tests/test_project_fit.py +++ b/tests/test_project_fit.py @@ -17,16 +17,20 @@ # -def _make_truth_file(*, amplitude=20.0, x0_shift=3.0, tau=5.0): +def _make_truth_file( + *, amplitude=20.0, x0_shift=3.0, tau=5.0, energy=None, time_ax=None +): """Create a file with known parameters for data generation. Uses a throwaway project so truth files don't pollute the fit project. + Pass ``energy``/``time_ax`` to override the default grids (used by + heterogeneous-grid tests). """ truth_project = make_project(name="truth") - energy = np.linspace(83, 87, 30) - time_ax = np.linspace(-2, 10, 24) + energy = np.linspace(83, 87, 30) if energy is None else energy + time_ax = np.linspace(-2, 10, 24) if time_ax is None else time_ax file = File(parent_project=truth_project) file.energy = energy @@ -628,3 +632,316 @@ def test_num_fmt_and_delim_propagate_to_csv_outputs(self): assert fit_1d_lines[0].startswith("energy;sum;") energy_field = fit_1d_lines[1].split(";")[0] assert "." in energy_field and "e" not in energy_field.lower(), energy_field + + +# +# +class TestPackProjectTheta: + """Test spectra.pack_project_theta index-array assembly.""" + + # + def _stub_plan(self, opt_param_names): + import types + + return types.SimpleNamespace(opt_param_names=opt_param_names) + + # + def test_synthetic_shared_and_file_params(self): + """Shared param gathers to one theta_c slot; plan order preserved.""" + + mapping = [ + ("tau", 0, "tau"), + ("file00_A", 0, "A"), + ("file00_F", 0, "F"), + ("tau", 1, "tau"), + ("file01_A", 1, "A"), + ("file01_F", 1, "F"), + ] + par_names = ["tau", "file00_A", "file00_F", "file01_A", "file01_F"] + var_names = ["tau", "file00_A", "file01_A"] + # opt order differs between the plans on purpose + plans = [self._stub_plan(["A", "tau"]), self._stub_plan(["tau", "A"])] + + from trspecfit import spectra + + theta_c_indices, plan_gathers = spectra.pack_project_theta( + plans, + mapping=mapping, + par_names=par_names, + var_names=var_names, + ) + + assert theta_c_indices.tolist() == [0, 1, 3] + assert plan_gathers[0].tolist() == [1, 0] + assert plan_gathers[1].tolist() == [0, 2] + + # end-to-end gather: full combined vector -> per-plan theta + par_full = np.array([5.0, 20.0, 0.1, 30.0, 0.2]) + theta_c = par_full[theta_c_indices] + assert theta_c[plan_gathers[0]].tolist() == [20.0, 5.0] + assert theta_c[plan_gathers[1]].tolist() == [5.0, 30.0] + + # + def test_missing_mapping_entry_raises(self): + """Plan opt param without a mapping entry is an internal error.""" + + from trspecfit import spectra + + with pytest.raises(RuntimeError, match="no combined-parameter mapping"): + spectra.pack_project_theta( + [self._stub_plan(["GLP_01_A"])], + mapping=[], + par_names=[], + var_names=[], + ) + + # + def test_non_varying_counterpart_raises(self): + """Plan opt param mapping to a static combined param is an error.""" + + from trspecfit import spectra + + with pytest.raises(RuntimeError, match="not\\s+varying"): + spectra.pack_project_theta( + [self._stub_plan(["GLP_01_F"])], + mapping=[("file00_GLP_01_F", 0, "GLP_01_F")], + par_names=["file00_GLP_01_F"], + var_names=[], + ) + + # + def test_unconsumed_varying_param_raises(self): + """A varying combined param feeding no plan is an error.""" + + from trspecfit import spectra + + with pytest.raises(RuntimeError, match="feed no plan"): + spectra.pack_project_theta( + [self._stub_plan(["GLP_01_A"])], + mapping=[("file00_GLP_01_A", 0, "GLP_01_A")], + par_names=["file00_GLP_01_A", "orphan"], + var_names=["file00_GLP_01_A", "orphan"], + ) + + # + def test_real_models_roundtrip(self): + """Packing built from real plans reproduces name-based distribution.""" + + from trspecfit import spectra + from trspecfit.graph_ir import build_graph, can_lower_2d, schedule_2d + + project = make_project(name="project_fit") + + for i in range(2): + f = File(parent_project=project, name=f"file_{i}") + f.energy = np.linspace(83, 87, 10) + f.time = np.linspace(-2, 10, 10) + f.dim = 2 + f.load_model( + model_yaml="models/project_energy.yaml", + model_info="project_glp", + ) + f.add_time_dependence( + target_model="project_glp", + target_parameter="GLP_01_x0", + dynamics_yaml="models/project_time.yaml", + dynamics_model=["MonoExpProject"], + ) + + combined, info = project._build_fit_params(model_name="project_glp") + par_names = info["par_names"] + var_names = [n for n in par_names if combined[n].vary] + + plans = [] + for model in info["models"]: + graph = build_graph(model) + assert can_lower_2d(graph) + plans.append(schedule_2d(graph)) + + theta_c_indices, plan_gathers = spectra.pack_project_theta( + plans, + mapping=info["mapping"], + par_names=par_names, + var_names=var_names, + ) + + # gathered values must equal the name-based lookup per file + remaps = [{}, {}] + for combined_name, file_idx, local_name in info["mapping"]: + remaps[file_idx][local_name] = combined_name + par_full = np.array([combined[n].value for n in par_names]) + theta_c = par_full[theta_c_indices] + for file_idx, plan in enumerate(plans): + expected = [ + combined[remaps[file_idx][local]].value + for local in plan.opt_param_names + ] + assert theta_c[plan_gathers[file_idx]].tolist() == expected + + # shared tau lands on the same theta_c slot for both files + tau_name = "GLP_01_x0_expFun_01_tau" + tau_slot = var_names.index(tau_name) + for file_idx, plan in enumerate(plans): + local_tau = [ + local for local, comb in remaps[file_idx].items() if comb == tau_name + ] + assert len(local_tau) == 1 + opt_pos = plan.opt_param_names.index(local_tau[0]) + assert plan_gathers[file_idx][opt_pos] == tau_slot + + +# +def _make_shared_tau_project(*, spec_fun_str, grids=None, show_output=0): + """Build a ready-to-fit 2-file project with shared tau, per-file A. + + Simulates noiseless data from two truth files (differing amplitudes, + identical tau) and assembles fit files via the real workflow. + ``grids`` optionally gives per-file ``(energy, time_ax)`` pairs for + heterogeneous-grid tests. ``show_output`` is applied after setup so + baseline fits stay silent. + """ + + if grids is None: + grids = [(None, None), (None, None)] + amplitudes = [20.0, 14.0] + seeds = [42, 43] + + project = make_project( + name=f"jax_project_{spec_fun_str}", + spec_fun_str=spec_fun_str, + auto_export=False, + ) + for i, ((energy, time_ax), amplitude, seed) in enumerate( + zip(grids, amplitudes, seeds, strict=True) + ): + truth = _make_truth_file(amplitude=amplitude, energy=energy, time_ax=time_ax) + data = simulate_clean(truth.model_active, seed=seed) + _make_fit_file(project, data, truth.energy, truth.time, name=f"file_{i}") + project.show_output = show_output + return project + + +# +# +class TestProjectFitJax: + """Project.fit_2d dispatch: fused JAX backend and interpreter fallback. + + The fallback tests monkeypatch the gate/factory seams because the + JAX slice currently covers the full lowered 2D surface — no public + YAML construct builds a 2D model that fails ``can_lower_jax_2d``. + They run without jax installed; only the tests that execute the + fused path importorskip. + """ + + TRUE_TAU = 5.0 + + # + def test_jax_parity_with_interpreter(self, capsys): + """Both backends converge to the same parameters on clean data.""" + + pytest.importorskip("jax") + + results = {} + for spec_fun_str in ("fit_model_gir", "fit_model_jax"): + project = _make_shared_tau_project(spec_fun_str=spec_fun_str, show_output=1) + # exercise per-file fit windows on one file + project.files[0].e_lim = [2, 28] + project.files[0].t_lim = [1, 23] + project.fit_2d(model_name="project_glp", stages=2, try_ci=0) + + backend = "JAX" if spec_fun_str == "fit_model_jax" else "interpreter" + assert f"({backend} backend)" in capsys.readouterr().out + + per_file = [] + for f in project.files: + m = f.select_model("project_glp") + assert m is not None # type guard + per_file.append({n: m.lmfit_pars[n].value for n in m.parameter_names}) + results[spec_fun_str] = per_file + + for file_idx, (pars_gir, pars_jax) in enumerate( + zip(results["fit_model_gir"], results["fit_model_jax"], strict=True) + ): + for name, value in pars_gir.items(): + assert np.isclose(pars_jax[name], value, rtol=1e-6, atol=1e-9), ( + f"file {file_idx} par {name}: gir={value!r} jax={pars_jax[name]!r}" + ) + + tau_jax = results["fit_model_jax"][0]["GLP_01_x0_expFun_01_tau"] + assert np.isclose(tau_jax, self.TRUE_TAU, atol=0.01) + + # + def test_heterogeneous_grids_fuse(self, capsys): + """Files with different energy/time grids fit on the fused path.""" + + pytest.importorskip("jax") + + grids = [ + (np.linspace(83, 87, 30), np.linspace(-2, 10, 24)), + (np.linspace(83.2, 86.8, 37), np.linspace(-1.5, 9, 19)), + ] + project = _make_shared_tau_project( + spec_fun_str="fit_model_jax", grids=grids, show_output=1 + ) + project.fit_2d(model_name="project_glp", stages=2, try_ci=0) + assert "(JAX backend)" in capsys.readouterr().out + + m0 = project.files[0].select_model("project_glp") + m1 = project.files[1].select_model("project_glp") + assert m0 is not None # type guard + assert m1 is not None # type guard + tau_0 = m0.lmfit_pars["GLP_01_x0_expFun_01_tau"].value + tau_1 = m1.lmfit_pars["GLP_01_x0_expFun_01_tau"].value + assert tau_0 == tau_1 # project-shared + assert np.isclose(tau_0, self.TRUE_TAU, atol=0.05) + A_0 = m0.lmfit_pars["GLP_01_A"].value + A_1 = m1.lmfit_pars["GLP_01_A"].value + assert np.isclose(A_0, 20.0, atol=0.1) + assert np.isclose(A_1, 14.0, atol=0.1) + + # + def test_fallback_when_one_file_not_jax_lowerable(self, monkeypatch, capsys): + """One file failing the JAX gate sends the whole project to MCP.""" + + from trspecfit import graph_ir + + project = _make_shared_tau_project(spec_fun_str="fit_model_jax", show_output=1) + + # First file passes the real gate, second is rejected. + real_gate = graph_ir.can_lower_jax_2d + gate_calls: list[bool] = [] + + def fail_from_second_call(graph): + gate_calls.append(True) + if len(gate_calls) >= 2: + return False + return real_gate(graph) + + monkeypatch.setattr( + "trspecfit.graph_ir.can_lower_jax_2d", fail_from_second_call + ) + + project.fit_2d(model_name="project_glp", stages=1, try_ci=0) + assert "(interpreter backend)" in capsys.readouterr().out + assert len(gate_calls) >= 2 + + # The interpreter path still recovers the shared tau. + m = project.files[0].select_model("project_glp") + assert m is not None # type guard + tau_fit = m.lmfit_pars["GLP_01_x0_expFun_01_tau"].value + assert np.isclose(tau_fit, self.TRUE_TAU, atol=0.05) + + # + def test_fallback_when_jax_unavailable(self, monkeypatch, capsys): + """Factory raising ImportError (jax missing) falls back to MCP.""" + + def raise_import_error(*args, **kwargs): + raise ImportError("jax is not installed") + + project = _make_shared_tau_project(spec_fun_str="fit_model_jax", show_output=1) + monkeypatch.setattr( + "trspecfit.eval_jax.make_project_evaluator_2d_jax", raise_import_error + ) + + project.fit_2d(model_name="project_glp", stages=1, try_ci=0) + assert "(interpreter backend)" in capsys.readouterr().out