diff --git a/.claude/skills/benchmark/SKILL.md b/.claude/skills/benchmark/SKILL.md index 1d4ff4c..5d9a5fc 100644 --- a/.claude/skills/benchmark/SKILL.md +++ b/.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 *) diff --git a/.claude/skills/benchmark/benchmark_gir.py b/.claude/skills/benchmark/benchmark_gir.py index 3ec2749..4b47604 100644 --- a/.claude/skills/benchmark/benchmark_gir.py +++ b/.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) diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index 52dc3cf..92a7e1f 100644 --- a/.github/workflows/ci.yaml +++ b/.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 "" diff --git a/CHANGELOG.md b/CHANGELOG.md index 7c05f16..cdd888f 100644 --- a/CHANGELOG.md +++ b/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). diff --git a/CLAUDE.md b/CLAUDE.md index 63c0aa0..57789bb 100644 --- a/CLAUDE.md +++ b/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. diff --git a/TODO.md b/TODO.md index b025118..ad2aabc 100644 --- a/TODO.md +++ b/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. diff --git a/docs/ai/benchmark.md b/docs/ai/benchmark.md index f8d17be..8fcc333 100644 --- a/docs/ai/benchmark.md +++ b/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 --calls 200 [--fit] [-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 diff --git a/docs/api/eval_jax.rst b/docs/api/eval_jax.rst new file mode 100644 index 0000000..58a89c0 --- /dev/null +++ b/docs/api/eval_jax.rst @@ -0,0 +1,7 @@ +JAX Evaluator (experimental) +============================ + +.. automodule:: trspecfit.eval_jax + :members: + :undoc-members: + :show-inheritance: diff --git a/docs/api/index.rst b/docs/api/index.rst index 935c4d0..98bbda8 100644 --- a/docs/api/index.rst +++ b/docs/api/index.rst @@ -18,3 +18,4 @@ This section contains the auto-generated API documentation. graph_ir eval_1d eval_2d + eval_jax diff --git a/docs/design/archive/jax-backend.md b/docs/design/archive/jax-backend.md new file mode 100644 index 0000000..a4b467a --- /dev/null +++ b/docs/design/archive/jax-backend.md @@ -0,0 +1,129 @@ +--- +orphan: true +--- + +# JAX Backend Track (Phases A–D) + +> **Status: implemented** (2026-07-11, `jax-backend` branch, v0.12.0). +> Execution record of the JAX track planned in +> [../jax-planning.md](../jax-planning.md). The backend lives in +> `eval_jax.py`, gated by `graph_ir.can_lower_jax_2d`, selected via +> `Project.spec_fun_str = "fit_model_jax"`. Phase E (optimizer +> replacement) was deferred. Follow-on design notes: +> [../ui.md](../ui.md), [../project-level-fits.md](../project-level-fits.md). + +## Guiding decisions + +- The NumPy GIR evaluator stays the reference backend throughout; JAX + must earn default status via parity + fit-level benchmarks. +- **Backend capability gating**: mirrors the `can_lower_2d` pattern + with a JAX-specific capability check (own lowerable-function set and + node-kind set, per the "backend-specific sets" carve-out in + `graph_ir.py`). +- **Fallback chain is JAX → NumPy GIR → MCP**: a model that fails the + JAX gate but passes `can_lower_2d` runs on the compiled NumPy path — + no user-visible regression from partial JAX coverage. Enforced by + construction: `can_lower_jax_2d` composes with `can_lower_2d`. +- "JAX evaluator + analytic Jacobian (keep lmfit)" and "replace lmfit" + are separate milestones; the latter is optional and evidence-driven. + +## Phase A: backend-agnostic prep (no JAX dependency; landed alone) + +- Flattened `ScheduledPlan2D.expr_programs` into packed arrays + (`expr_instructions` + `expr_indptr`, CSR-style); `ExprProgram` + dataclass removed, plans are fully array-native. +- Same for `profile_expr_programs` (both 1D and 2D plans). +- Evaluator API unchanged: `evaluate(plan, theta) -> ndarray`. + +## Phase B: first JAX evaluator slice (file-level 2D fits) + +- Packaging: optional `[jax]` extra; `eval_jax.py` guards the import + and raises a helpful ImportError; tests skip without jax. Importing + `eval_jax` enables `jax_enable_x64` globally (float64 parity). +- Functional JAX 2D evaluator (`make_evaluator_2d_jax`): jnp kernel + mirrors of the `functions/` bodies. LinBack drops its host-side + ordering ValueError — untraceable; fit bounds must keep order. +- Dispatch strategy: trace-time unrolling — one jitted XLA program per + plan, no dispatch inside the compiled path; only theta is traced. +- Host-side checks outside jit (theta shape; plan-level feature check + in `_check_plan_supported`). +- Timing sanity (212x1131 grid, glp + 2 dynamics): JAX jit 0.78 + ms/call vs NumPy GIR 7.6 ms/call (~9.8x); compile ~230 ms once per + plan; max |diff| 1.4e-13. + +## Phase C: widened to the full lowered 2D surface + +- Profile-varying parameters: sample groups, per-sample profile + expressions (broadcast virtual rows), profiled ops vectorized over + aux + averaged (the NumPy path loops per aux point to avoid + temporaries; under XLA the fused form is simpler and still 2.6x + faster at n_aux=50). +- Subcycle-aware dynamics: free after scheduling — subcycle info is + pure data (`dyn_sub_time_axes`/`dyn_sub_masks`); gate carve-out plus + parity tests only. +- Resolved-trace convolution: kernel-matrix apply in jnp (gather, + quadrature weights, matmul, analytic edge masses via + `jax.scipy.special.erfc`). The NumPy path's runtime value checks + (kernel positivity, row sums) are untraceable and omitted. +- Voigt via Weideman (1994) rational `wofz` approximation, 64 terms, + coefficients precomputed at import (host-side FFT); accuracy vs + scipy `wofz` < 1e-12 rel over the physical domain (tested). ~69x + faster than the NumPy/scipy path at 212x1131. +- `can_lower_jax_2d` spans the full `can_lower_2d` surface; the + separate sets remain so future NumPy widening doesn't silently imply + JAX support. +- Timing (per call): profiled x0 n_aux=50 175x280: 10.6 ms vs 27.2 ms; + gaussCONV 212x1131: 1.0 ms vs 5.0 ms; Voigt+dynamics 212x1131: + 0.63 ms vs 44.0 ms. + +## Phase D: analytic Jacobian (keeping lmfit) + +- `make_jacobian_2d_jax(plan)`: `jax.jit(jax.jacfwd(...))` over the + shared traced evaluator -> `(n_time, n_energy, n_opt)`. Forward mode + (theta is short, output grid large). Caveats: boxCONV width + derivative is 0 a.e.; step-like `where` switches give subgradient + zeros at the exact switching point. +- `spec_fun_str = "fit_model_jax"`: 2D fits run residuals on the + jitted evaluator (`spectra.fit_model_jax`); JAX-rejected graphs fall + back to the compiled NumPy plan; 1D fits lower to the NumPy plan + (JAX backend is 2D-only). Evaluator/jacobian are per-plan closures — + `lmfit.emcee` with `workers > 1` will not pickle them. +- `Dfun` plumbing: `fitlib.jacobian_fun` mirrors `residual_fun`'s + signature (lmfit calls Dfun with the same fcn_args), negates the + windowed model Jacobian, and reorders columns to lmfit's varying- + parameter order; `fit_wrapper(jac_fun=...)` forwards it as + `Dfun`/`col_deriv=0` for leastsq stages only (auto-set by + `File.fit_2d` on the JAX path; pass `jac_fun=None` to disable). +- Jacobian validated three ways: central finite differences of + `evaluate_2d` (cross-backend) for dynamics/expressions, convolution, + subcycles, profiles, Voigt; lmfit's internal-variable FD at a real + fit's start point (exact match through the bounds transformation); + slow e2e `File.fit_2d` truth-recovery test. +- Fit-level benchmarks: leastsq stage nfev 22 -> 4 (simple GLP model, + same optimum); Voigt 2D fit (212x1131, stages=2, 1% noise): 25.3 s + -> 7.1 s wall (~3.6x, JAX time includes compiles), identical redchi, + leastsq nfev 8 -> 2. Benchmark-skill fit mode (example 2): JAX+Dfun + 2.10 s vs GIR 4.86 s vs interpreter 17.0 s. +- Investigated: from aggressively perturbed starts, plain single-stage + leastsq lands in different local minima with analytic vs FD + Jacobians (nonconvex objective, not a defect) — the default + two-stage workflow (Nelder first) is the intended guard, unchanged. + +## Test coverage decision + +Every evaluator-vs-interpreter comparison in `test_evaluate_2d.py` +also asserts JAX parity when jax is importable (`_assert_jax_parity`), +so the full 2D matrix including regression fixtures runs three-way; +`test_evaluate_jax.py` keeps only JAX-specific tests (gate, Jacobian, +wofz, chained conv, e2e fit). The main CI job installs `.[dev,jax]`; +the min-versions job stays jax-free (jax needs numpy>=2.0, +incompatible with the numpy==1.26 floor world) and JAX tests skip +there. + +## Phase E: JAX-native optimizer — deferred + +Deferred (2026-07-11). The Phase D benchmarks show the evaluator and +Jacobian — not lmfit — dominate fit wall time at current model sizes. +Revisit only if profiling a real workload shows lmfit overhead +dominating; the most plausible pilot is a vmap-batched slice-by-slice +solver (see [../ui.md](../ui.md)). diff --git a/docs/design/jax-planning.md b/docs/design/jax-planning.md index aeca284..7ef18ce 100644 --- a/docs/design/jax-planning.md +++ b/docs/design/jax-planning.md @@ -4,6 +4,12 @@ orphan: true # Planning Note: JAX Backend, Jacobians, and Optimizer +> **Status (2026-07-11):** Phases A–D of this plan are implemented +> (`eval_jax.py`, `can_lower_jax_2d`, `fit_model_jax` + `Dfun` +> plumbing); Phase E (optimizer replacement) is deferred for lack of +> evidence that lmfit is the bottleneck. This note is preserved as the +> planning-time rationale. + ## Summary The current NumPy GIR backend is ready to be treated as complete for the diff --git a/docs/design/lowered_evaluator.md b/docs/design/lowered_evaluator.md index d028a7f..f09e232 100644 --- a/docs/design/lowered_evaluator.md +++ b/docs/design/lowered_evaluator.md @@ -489,7 +489,7 @@ class OpKind(IntEnum): class ScheduledPlan2D: """Compiled 2D execution schedule. - No Python objects in the hot path (except ``expr_programs``). + No Python objects in the hot path. """ energy: np.ndarray # (n_energy,) @@ -522,9 +522,12 @@ class ScheduledPlan2D: dyn_sub_n_params: np.ndarray # (n_substeps,) int # --- Expression evaluation --- + # Packed RPN programs, CSR-style: program i occupies + # expr_instructions[expr_indptr[i]:expr_indptr[i + 1]]. n_expressions: int expr_target_rows: np.ndarray # (n_expressions,) int -- which row to write - expr_programs: list["ExprProgram"] # compiled RPN programs + expr_instructions: np.ndarray # (total_expr_words,) int64 + expr_indptr: np.ndarray # (n_expressions + 1,) int # --- Interleaved parameter resolution schedule --- # Dynamics groups and expressions may depend on each other @@ -532,7 +535,7 @@ class ScheduledPlan2D: # dynamics parameter). These arrays encode the correct topological # execution order: # kind=0 -> dynamics group step, index into dyn_group_* arrays - # kind=1 -> expression step, index into expr_* arrays / expr_programs + # kind=1 -> expression step, index into expr_* arrays resolution_kinds: np.ndarray # (n_dyn_groups + n_expressions,) int8 resolution_indices: np.ndarray # (n_dyn_groups + n_expressions,) int @@ -569,15 +572,14 @@ class ExprNodeKind(IntEnum): POW = 7 # pop 2, push power -@dataclass(frozen=True) -class ExprProgram: - """Compiled expression: flat int array encoding an RPN program.""" - - # Encoding: pairs of (node_kind, operand). - # CONST: operand is float bits (np.float64.view(np.int64)) - # PARAM_REF: operand is row index into trace matrix - # Operators: operand is 0 (unused) - instructions: np.ndarray # (2 * n_instructions,) int64 +# Each compiled expression is a flat int64 instruction array; all +# programs are concatenated into the plan's packed expr_instructions, +# with expr_indptr marking per-program offsets (CSR-style). +# +# Encoding: pairs of (node_kind, operand). +# CONST: operand is float bits (np.float64.view(np.int64)) +# PARAM_REF: operand is row index into trace matrix +# Operators: operand is 0 (unused) ``` All values flowing through the RPN evaluator are `(n_time,)` arrays @@ -1011,9 +1013,9 @@ Once `evaluate_2d(plan, theta) -> spectrum` exists as a pure function: **Numba:** `@njit` on the component eval dispatch loop. Most of the plan (int index arrays, CSR param maps, dense trace matrix) is -Numba-compatible. The expression programs (`list[ExprProgram]`) would -need flattening to a CSR-style encoding first -- this is noted in the -plan as a v1 simplification that can be tightened later. +Numba-compatible. The expression programs are stored as packed +CSR-style instruction arrays (`expr_instructions` + `expr_indptr`), +so no Python-object flattening step remains. **JAX:** Replace `np` with `jnp` in the evaluator. The plan's array structure maps directly to JAX arrays. Key wins: diff --git a/docs/design/project-level-fits.md b/docs/design/project-level-fits.md new file mode 100644 index 0000000..d79b547 --- /dev/null +++ b/docs/design/project-level-fits.md @@ -0,0 +1,88 @@ +--- +orphan: true +--- + +# 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](jax-planning.md) and the +`eval_jax.py` section of [repo_architecture.md](repo_architecture.md)). +Nothing here is implemented. + +## 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 +`ScheduledPlan`s, 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](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](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. diff --git a/docs/design/repo_architecture.md b/docs/design/repo_architecture.md index 5759bc0..1f5c8d3 100644 --- a/docs/design/repo_architecture.md +++ b/docs/design/repo_architecture.md @@ -77,6 +77,18 @@ dispatch tables live here (`DYNAMICS_DISPATCH`, `CONV_KERNEL_DISPATCH`). **Performance-critical — prefer array operations, avoid Python-level branching on model structure (the plan already captured it).** +### `eval_jax.py` — experimental JAX backend (optional `[jax]` extra) + +Jitted mirror of the 2D evaluator plus an analytic Jacobian +(`make_evaluator_2d_jax`, `make_jacobian_2d_jax`), compiled per plan by +trace-time unrolling of the schedule arrays. Gated by +`graph_ir.can_lower_jax_2d` (at most as wide as `can_lower_2d`, so +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. + ### `spectra.py` — evaluator bridge Thin module that the fitting engine calls on every residual evaluation. diff --git a/docs/design/ui.md b/docs/design/ui.md new file mode 100644 index 0000000..277b4b7 --- /dev/null +++ b/docs/design/ui.md @@ -0,0 +1,78 @@ +--- +orphan: true +--- + +# Design Note: Interactive UI Backend Requirements + +Forward-looking note (2026-07-11) collecting backend facts and open +design choices for a future interactive UI (live fit preview while the +user edits a model: initial guesses, fit limits, fixed values). Written +while the JAX backend landed; nothing here is implemented UI work. + +## Latency picture per interaction + +The UI's hot loop is: user tweaks something -> re-evaluate (or re-fit) +-> redraw. What each tweak costs on the compiled backends: + +| User action | JAX backend cost | Recompile? | +|---|---|---| +| Change a varying parameter's value/guess | one jitted call (~0.6-10 ms for 2D) | no — theta is the traced input | +| Change a fit limit (`e_lim`/`t_lim`) | free on the compiled side | no — windowing is host-side slicing of the full-grid output; axes are never cropped | +| Change a *fixed* parameter's value | re-schedule + re-trace (~100-500 ms hiccup) | yes — fixed values are baked into the plan and become XLA constants | +| Flip a `vary` flag | re-schedule + re-trace | yes — theta layout changes | +| Change model structure (components, dynamics, profiles) | re-schedule + re-trace | yes — unavoidable, the plan itself changes | + +## Known fixes / prerequisites for fluid interaction + +- **Full-parameter-vector evaluator variant.** Trace the evaluator over + the full parameter vector instead of theta (a small + `make_evaluator_2d_jax` variant). Fixed-value edits and vary-flips + then become plain input changes; only true structure changes + recompile. Negligible runtime cost. +- **Session-level evaluator caching.** `File.fit_2d` currently rebuilds + graph + plan + evaluator on every call. A UI must hold the compiled + evaluator/Jacobian across interactions (the factory API supports this + directly) and invalidate only on structure changes. +- **Persistent compilation cache.** `jax_compilation_cache_dir` makes + XLA compiles survive process restarts, so even a session's first + render can be warm. + +## 1D does not need JAX for interactivity + +Measured on the NumPy GIR 1D path (2001-point grid, 2026-07-11): +24 us/eval (GLP) to 188 us/eval (Voigt via scipy `wofz`). A full 1D +leastsq fit is single-digit-to-tens of milliseconds — far below +perception, with no compile hiccups ever. The interaction where the +backend genuinely matters is **2D fit preview** (seconds -> sub-second), +which the JAX backend covers. Keep 1D on the NumPy path. + +## SbS throughput: `n_workers` (today) vs `vmap` (possible future) + +Slice-by-slice fitting is the other latency-relevant workload (a UI +would want a full SbS refresh after a seed change). Two very different +parallelization models: + +- **`fit_slice_by_slice(n_workers=N)` — what exists.** Process-level + parallelism over whole *fits*: each worker owns a pickled model copy + and runs complete lmfit optimizations per slice. Backend-agnostic and + robust; scales ~min(N_cores, n_slices) minus process spawn/pickle + overhead, which is why it only pays off above ~20 slices. Keeps all + lmfit machinery (per-slice stderr from the covariance, etc.). +- **`vmap` over slices — what JAX could add.** Array-level batching of + the *evaluator*: all slices evaluated in one XLA program with no + per-process copies or pickling. But vmap alone does not parallelize + the *fits* — each slice is an independent lmfit optimization with its + own iteration trajectory, and lmfit cannot step hundreds of + optimizations in lockstep. Exploiting vmap for SbS therefore requires + a **batched least-squares solver** (a JAX-native LM stepping all + slices simultaneously until each converges) — i.e., the Phase E + "replace lmfit" decision from + [jax-planning.md](jax-planning.md), scoped to the SbS inner loop. + +Notably, SbS is the one workload where lmfit itself plausibly *is* the +bottleneck (per-eval cost is tens of microseconds, so per-iteration +Python/lmfit overhead dominates) — making a batched SbS solver the +natural Phase E pilot if throughput ever demands it. It must reproduce +the per-slice error bars that downstream SbS analysis consumes before +it can replace the default. The two models compose poorly (process +pools each re-import jax/XLA), so it would be either/or per fit. diff --git a/pyproject.toml b/pyproject.toml index 2466263..bdd490f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "trspecfit" -version = "0.11.0" +version = "0.12.0" authors = [ {name = "Johannes Mahl", email = "infinitymonkeyatwork@gmail.com"}, ] @@ -68,6 +68,13 @@ dev = [ "twine==6.2.0", ] +# Experimental JAX evaluator backend (CPU wheel; see +# docs/design/jax-planning.md). Tests that exercise it skip when +# jax is not installed. +jax = [ + "jax>=0.4.35", +] + # Optional notebook/visualization support # Also requires the Graphviz "dot" executable on your PATH. # Linux example: sudo apt-get install graphviz @@ -131,6 +138,7 @@ optional_dependencies_dev_groups = ["dev"] # confidence interval estimation needs to work with standard install # keep these as runtime deps even though used indirectly via lmfit # py-spy is a CLI tool invoked via shell (see benchmark skill), not imported +# jax is the optional [jax] extra, imported behind a TYPE_CHECKING/try guard [tool.deptry.per_rule_ignores] DEP002 = [ "numdifftools", @@ -140,6 +148,7 @@ DEP002 = [ "jupyterlab", "graphviz", "py-spy", + "jax", ] # scripts/normalize_notebooks.py is dev-only tooling (pre-commit hook), # so its dev-dep imports are intentional. deptry treats anything outside diff --git a/src/trspecfit/eval_1d.py b/src/trspecfit/eval_1d.py index e1fd39e..638d4bd 100644 --- a/src/trspecfit/eval_1d.py +++ b/src/trspecfit/eval_1d.py @@ -60,7 +60,10 @@ def evaluate_1d(plan: ScheduledPlan1D, theta: np.ndarray) -> np.ndarray: # 1c. Resolve expressions in topological order for i in range(plan.n_expressions): target = int(plan.expr_target_indices[i]) - values[target] = _eval_expr_scalar(plan.expr_programs[i], values) + values[target] = _eval_expr_scalar( + plan.expr_instructions[plan.expr_indptr[i] : plan.expr_indptr[i + 1]], + values, + ) profile_sample_values = _evaluate_profile_sample_values( plan.aux_axis, @@ -75,7 +78,8 @@ def evaluate_1d(plan: ScheduledPlan1D, theta: np.ndarray) -> np.ndarray: values, profile_sample_values, plan.n_params, - plan.profile_expr_programs, + plan.profile_expr_instructions, + plan.profile_expr_indptr, ) # 2. Component evaluation diff --git a/src/trspecfit/eval_2d.py b/src/trspecfit/eval_2d.py index 44bda7e..4890508 100644 --- a/src/trspecfit/eval_2d.py +++ b/src/trspecfit/eval_2d.py @@ -20,7 +20,6 @@ ConvKernelKind, DynFuncKind, ExprNodeKind, - ExprProgram, ParamSourceKind, ScheduledPlan2D, ) @@ -72,10 +71,10 @@ # def eval_expr_program( - program: ExprProgram, + instructions: np.ndarray, traces: np.ndarray, ) -> np.ndarray: - """Evaluate an RPN ExprProgram against the trace matrix. + """Evaluate a compiled RPN program against the trace matrix. Works for both plan initialization and hot-path evaluation. Each PARAM_REF pushes a *view* of its ``(n_time,)`` trace row and @@ -85,8 +84,9 @@ def eval_expr_program( Parameters ---------- - program - Compiled RPN instruction array. + instructions + Compiled RPN instruction array (one program's slice of the + packed ``expr_instructions``). traces ``(n_params, n_time)`` trace matrix (current state). @@ -98,7 +98,7 @@ def eval_expr_program( n_time = traces.shape[1] stack: list[np.ndarray | np.float64] = [] - instr = program.instructions + instr = instructions n_instr = len(instr) // 2 for i in range(n_instr): @@ -160,7 +160,8 @@ def resolve_param_traces( dyn_sub_time_axes: np.ndarray, dyn_sub_masks: np.ndarray, expr_target_rows: np.ndarray, - expr_programs: list[ExprProgram], + expr_instructions: np.ndarray, + expr_indptr: np.ndarray, conv_target_rows: np.ndarray, conv_func_ids: np.ndarray, conv_param_indptr: np.ndarray, @@ -204,7 +205,10 @@ def resolve_param_traces( ) elif kind == 1: # expression target = int(expr_target_rows[idx]) - traces[target, :] = eval_expr_program(expr_programs[idx], traces) + traces[target, :] = eval_expr_program( + expr_instructions[expr_indptr[idx] : expr_indptr[idx + 1]], + traces, + ) else: # kind == 2: resolved-trace convolution assert conv_operator is not None # type guard: set when steps exist target = int(conv_target_rows[idx]) @@ -285,7 +289,8 @@ def _evaluate_profile_expr_values_2d( traces: np.ndarray, profile_sample_values: np.ndarray, n_params: int, - profile_expr_programs: list[ExprProgram], + profile_expr_instructions: np.ndarray, + profile_expr_indptr: np.ndarray, ) -> np.ndarray: """Evaluate lowered per-sample profile expressions over (n_time, n_aux). @@ -293,7 +298,7 @@ def _evaluate_profile_expr_values_2d( so the standard RPN evaluator can be reused unchanged. """ - n_exprs = len(profile_expr_programs) + n_exprs = len(profile_expr_indptr) - 1 if n_exprs == 0: n_time = traces.shape[1] n_aux = profile_sample_values.shape[2] if profile_sample_values.size else 0 @@ -314,8 +319,12 @@ def _evaluate_profile_expr_values_2d( virtual[n_params:, :] = profile_sample_values.reshape(n_groups, n_cols) expr_values = np.empty((n_exprs, n_time, n_aux), dtype=np.float64) - for expr_idx, program in enumerate(profile_expr_programs): - result = eval_expr_program(program, virtual) # (n_time*n_aux,) + for expr_idx in range(n_exprs): + start = int(profile_expr_indptr[expr_idx]) + end = int(profile_expr_indptr[expr_idx + 1]) + result = eval_expr_program( + profile_expr_instructions[start:end], virtual + ) # (n_time*n_aux,) expr_values[expr_idx] = result.reshape(n_time, n_aux) return expr_values @@ -436,7 +445,8 @@ def evaluate_2d(plan: ScheduledPlan2D, theta: np.ndarray) -> np.ndarray: plan.dyn_sub_time_axes, plan.dyn_sub_masks, plan.expr_target_rows, - plan.expr_programs, + plan.expr_instructions, + plan.expr_indptr, plan.conv_target_rows, plan.conv_func_ids, plan.conv_param_indptr, @@ -458,7 +468,8 @@ def evaluate_2d(plan: ScheduledPlan2D, theta: np.ndarray) -> np.ndarray: traces, profile_sample_values, plan.n_params, - plan.profile_expr_programs, + plan.profile_expr_instructions, + plan.profile_expr_indptr, ) # 2. Component evaluation diff --git a/src/trspecfit/eval_jax.py b/src/trspecfit/eval_jax.py new file mode 100644 index 0000000..286ef06 --- /dev/null +++ b/src/trspecfit/eval_jax.py @@ -0,0 +1,707 @@ +"""Experimental JAX evaluator backend for 2D scheduled plans. + +Covers the full lowered 2D surface (docs/design/jax-planning.md, +Phases B + C): static component ops, dynamics groups, arithmetic +expressions, subcycle dynamics, profile-varying parameters, +kernel-matrix convolution, and Voigt (via a Weideman rational +approximation of ``wofz``). ``can_lower_jax_2d`` gates entry; graphs +it rejects run on the compiled NumPy evaluator. + +The backend compiles one jitted function per plan via trace-time +unrolling: ``make_evaluator_2d_jax(plan)`` walks the schedule arrays +with host-side Python control flow while tracing, so the compiled +XLA program contains no dispatch, and only ``theta`` is traced. + +Value checks that the NumPy path performs on parameter-dependent +quantities (LinBack ordering, convolution kernel positivity) cannot +run on traced values and are omitted here; keep fit bounds ordered +and kernel widths positive. + +Importing this module enables JAX 64-bit mode globally +(``jax_enable_x64``); parity with the float64 NumPy evaluator requires +it. + +JAX is an optional dependency: ``pip install "trspecfit[jax]"``. +""" + +from collections.abc import Callable +from typing import TYPE_CHECKING + +import numpy as np + +from trspecfit.graph_ir import ( + ConvKernelKind, + DynFuncKind, + ExprNodeKind, + OpKind, + ParamSourceKind, + ProfileFuncKind, + ScheduledPlan2D, +) + +if TYPE_CHECKING: + import jax + import jax.numpy as jnp + from jax.scipy.special import erf as _jax_erf + from jax.scipy.special import erfc as _jax_erfc + + _HAVE_JAX = True +else: + try: + import jax + import jax.numpy as jnp + from jax.scipy.special import erf as _jax_erf + from jax.scipy.special import erfc as _jax_erfc + except ImportError: # pragma: no cover - exercised only without [jax] + jax = None + jnp = None + _jax_erf = None + _jax_erfc = None + _HAVE_JAX = False + else: + # Parity contract with the float64 NumPy evaluator. + jax.config.update("jax_enable_x64", True) + _HAVE_JAX = True + + +# +def _require_jax() -> None: + """Raise a helpful ImportError when JAX is not installed.""" + + if not _HAVE_JAX: + raise ImportError( + "The JAX evaluator backend requires jax; " + 'install it with: pip install "trspecfit[jax]"' + ) + + +# --------------------------------------------------------------------------- +# Component-op kernels (jnp mirrors of functions/energy.py bodies) +# --------------------------------------------------------------------------- +# Same broadcasting contract as the NumPy evaluator: *x* is +# ``(1, n_energy)``, params are ``(n_time, 1)`` columns, *spectrum* is +# ``(n_time, n_energy)``. + + +# +def _Offset(x, y0, spectrum=None): + return jnp.ones_like(x) * y0 + + +# +def _Shirley(x, pShirley, spectrum): + flipped = jnp.flip(spectrum, axis=-1) + return pShirley * jnp.flip(jnp.cumsum(flipped, axis=-1), axis=-1) + + +# +def _LinBack(x, m, b, xStart, xStop, spectrum=None): + # No xStart < xStop validation here: parameters are traced values, + # so the ordering cannot be checked at trace time. The NumPy + # reference path validates; keep fit bounds ordered. + y = m * (x - xStart) + b + y_stop = m * (xStop - xStart) + b + return jnp.where(x < xStart, b, jnp.where(x > xStop, y_stop, y)) + + +# +def _Gauss(x, A, x0, SD): + return A * jnp.exp(-1 / 2 * ((x - x0) / SD) ** 2) + + +# +def _GaussAsym(x, A, x0, SD, ratio): + return jnp.where(x < x0, _Gauss(x, A, x0, SD), _Gauss(x, A, x0, SD * ratio)) + + +# +def _Lorentz(x, A, x0, W): + return A / (1 + ((x - x0) / W * 2) ** 2) + + +# +def _GLS(x, A, x0, F, m): + u2 = ((x - x0) / F) ** 2 + return A * ((1 - m) * jnp.exp(-u2 * 4 * jnp.log(2)) + m / (1 + 4 * u2)) + + +# +def _GLP(x, A, x0, F, m): + u2 = ((x - x0) / F) ** 2 + return A * jnp.exp(-u2 * 4 * jnp.log(2) * (1 - m)) / (1 + 4 * m * u2) + + +# +def _DS(x, A, x0, F, alpha): + dx = x - x0 + return ( + A + * jnp.cos(jnp.pi * alpha / 2 + (1 - alpha) * jnp.arctan(dx / F)) + / (F**2 + dx**2) ** ((1 - alpha) / 2) + ) + + +# +def _weideman_coeffs(n_terms: int) -> tuple[float, np.ndarray]: + """Precompute Weideman (1994) rational-approximation coefficients. + + Host-side NumPy, evaluated once at import. ``_wofz`` uses the + result to approximate the Faddeeva function ``w(z)`` for + ``Im(z) > 0`` (always true for Voigt: z carries ``+i W/2``). + """ + + m = 2 * n_terms + k = np.arange(-m + 1, m) + L = np.sqrt(n_terms / np.sqrt(2.0)) + t = L * np.tan(k * np.pi / (2 * m)) + f = np.concatenate(([0.0], np.exp(-(t**2)) * (L**2 + t**2))) + a = np.real(np.fft.fft(np.fft.fftshift(f))) / (2 * m) + return float(L), a[1 : n_terms + 1][::-1] + + +# 64 terms: max relative error vs scipy wofz < 1e-13 over the physical +# Voigt domain (checked in tests/test_evaluate_jax.py). +_WEIDEMAN_L, _WEIDEMAN_A = _weideman_coeffs(64) + + +# +def _wofz(z): + """Faddeeva function ``w(z)`` for ``Im(z) > 0`` (Weideman 1994).""" + + iz = 1j * z + Z = (_WEIDEMAN_L + iz) / (_WEIDEMAN_L - iz) + p = jnp.zeros_like(Z) + for coeff in _WEIDEMAN_A: # Horner, unrolled at trace time + p = p * Z + coeff + return 2 * p / (_WEIDEMAN_L - iz) ** 2 + (1 / jnp.sqrt(jnp.pi)) / (_WEIDEMAN_L - iz) + + +# +def _Voigt(x, A, x0, SD, W): + scale = SD * jnp.sqrt(2.0) + voigt = jnp.real(_wofz(((x - x0) + 1j * (W / 2)) / scale)) + peak_voigt = jnp.real(_wofz(1j * (W / 2) / scale)) + return A * voigt / peak_voigt + + +JAX_OP_DISPATCH: dict[int, Callable] = { + int(OpKind.GAUSS): _Gauss, + int(OpKind.GAUSS_ASYM): _GaussAsym, + int(OpKind.LORENTZ): _Lorentz, + int(OpKind.VOIGT): _Voigt, + int(OpKind.GLS): _GLS, + int(OpKind.GLP): _GLP, + int(OpKind.DS): _DS, + int(OpKind.OFFSET): _Offset, + int(OpKind.LINBACK): _LinBack, + int(OpKind.SHIRLEY): _Shirley, +} + + +# --------------------------------------------------------------------------- +# Dynamics kernels (jnp mirrors of functions/time.py bodies) +# --------------------------------------------------------------------------- + + +# +def _stepFun(t, A, t0): + return jnp.where(t < t0, 0.0, A) + + +# +def _linFun(t, m, t0): + return jnp.where(t < t0, 0.0, m * (t - t0)) + + +# +def _expFun(t, A, tau, t0): + return jnp.where(t < t0, 0.0, A * jnp.exp(-1 / tau * (t - t0))) + + +# +def _sinFun(t, A, f, phi, t0): + return jnp.where(t < t0, 0.0, A * jnp.sin(2 * jnp.pi * f * (t - t0) + phi)) + + +# +def _sinDivX(t, A, f, t0): + return jnp.where(t < t0, 0.0, A * jnp.sinc(2 * f * (t - t0))) + + +# +def _erfFun(t, A, SD, t0): + return A / 2 * (1 + _jax_erf((t - t0) / (SD * jnp.sqrt(2.0)))) + + +# +def _sqrtFun(t, A, t0): + # Double-where instead of sqrt(clip(u, 0)): with a varying t0 the + # clip form differentiates as sqrt'(0) * clip'(u<0) = inf * 0 = NaN + # on every pre-onset sample, poisoning the whole Jacobian. The + # safe inner argument keeps both branch tangents finite; the outer + # where selects 0 pre-onset (values identical to the NumPy path). + u = t - t0 + u_safe = jnp.where(u > 0, u, 1.0) + return jnp.where(u > 0, A * jnp.sqrt(u_safe), 0.0) + + +JAX_DYNAMICS_DISPATCH: dict[int, Callable] = { + int(DynFuncKind.EXPFUN): _expFun, + int(DynFuncKind.SINFUN): _sinFun, + int(DynFuncKind.LINFUN): _linFun, + int(DynFuncKind.SINDIVX): _sinDivX, + int(DynFuncKind.ERFFUN): _erfFun, + int(DynFuncKind.SQRTFUN): _sqrtFun, + int(DynFuncKind.STEPFUN): _stepFun, +} + + +# --------------------------------------------------------------------------- +# Profile kernels (jnp mirrors of functions/profile.py bodies) +# --------------------------------------------------------------------------- + + +# +def _pExpDecay(x, A, tau): + return A * jnp.exp(-x / tau) + + +# +def _pLinear(x, m, b): + return m * x + b + + +# +def _pGauss(x, A, x0, SD): + return A * jnp.exp(-0.5 * ((x - x0) / SD) ** 2) + + +JAX_PROFILE_DISPATCH: dict[int, Callable] = { + int(ProfileFuncKind.PEXPDECAY): _pExpDecay, + int(ProfileFuncKind.PLINEAR): _pLinear, + int(ProfileFuncKind.PGAUSS): _pGauss, +} + + +# --------------------------------------------------------------------------- +# Convolution kernels and edge-mass companions (functions/time.py mirrors) +# --------------------------------------------------------------------------- +# The NumPy companions validate kernel parameters (strictly positive, +# finite) on every call; traced values cannot be validated, so the JAX +# mirrors omit that backstop. + + +# +def _gaussCONV(x, SD): + return jnp.exp(-1 / 2 * (x / SD) ** 2) + + +# +def _expSymCONV(x, tau): + return jnp.exp(-1 / tau * jnp.abs(x)) + + +# +def _expDecayCONV(x, tau): + return jnp.where(x < 0, 0.0, _expSymCONV(x, tau)) + + +# +def _expRiseCONV(x, tau): + return jnp.where(x > 0, 0.0, _expSymCONV(x, tau)) + + +# +def _boxCONV(x, width): + return jnp.where(jnp.abs(x) <= width / 2, 1.0, 0.0) + + +# +def _gaussCONV_edge_mass(dt_left, dt_right, SD): + scale = SD * jnp.sqrt(jnp.pi / 2) + M_L = scale * _jax_erfc(dt_left / (jnp.sqrt(2.0) * SD)) + M_R = scale * _jax_erfc(-dt_right / (jnp.sqrt(2.0) * SD)) + return M_L, M_R + + +# +def _expSymCONV_edge_mass(dt_left, dt_right, tau): + return tau * jnp.exp(-dt_left / tau), tau * jnp.exp(dt_right / tau) + + +# +def _expDecayCONV_edge_mass(dt_left, dt_right, tau): + return tau * jnp.exp(-dt_left / tau), jnp.zeros_like(dt_right) + + +# +def _expRiseCONV_edge_mass(dt_left, dt_right, tau): + return jnp.zeros_like(dt_left), tau * jnp.exp(dt_right / tau) + + +# +def _boxCONV_edge_mass(dt_left, dt_right, width): + M_L = jnp.clip(width / 2 - dt_left, 0.0, width) + M_R = jnp.clip(dt_right + width / 2, 0.0, width) + return M_L, M_R + + +JAX_CONV_KERNEL_DISPATCH: dict[int, Callable] = { + int(ConvKernelKind.GAUSSCONV): _gaussCONV, + int(ConvKernelKind.EXPSYMCONV): _expSymCONV, + int(ConvKernelKind.EXPDECAYCONV): _expDecayCONV, + int(ConvKernelKind.EXPRISECONV): _expRiseCONV, + int(ConvKernelKind.BOXCONV): _boxCONV, +} + +JAX_CONV_EDGE_MASS_DISPATCH: dict[int, Callable] = { + int(ConvKernelKind.GAUSSCONV): _gaussCONV_edge_mass, + int(ConvKernelKind.EXPSYMCONV): _expSymCONV_edge_mass, + int(ConvKernelKind.EXPDECAYCONV): _expDecayCONV_edge_mass, + int(ConvKernelKind.EXPRISECONV): _expRiseCONV_edge_mass, + int(ConvKernelKind.BOXCONV): _boxCONV_edge_mass, +} + + +# --------------------------------------------------------------------------- +# Expression evaluation (trace-time unrolled RPN) +# --------------------------------------------------------------------------- + + +# +def _eval_expr_rows(instructions: np.ndarray, rows: list, shape: tuple): + """Evaluate one packed RPN program over per-parameter trace rows. + + Mirrors ``eval_expr_program``; PARAM_REF pushes the *shape*-shaped + row (a traced value), constants stay Python floats until an + operator combines them. ``rows`` entries are ``(n_time,)`` for + plan expressions and ``(n_time, n_aux)`` for per-sample profile + expressions. + """ + + stack: list = [] + n_instr = len(instructions) // 2 + + for i in range(n_instr): + kind = int(instructions[2 * i]) + operand = instructions[2 * i + 1] + + if kind == ExprNodeKind.CONST: + stack.append(float(np.int64(operand).view(np.float64))) + elif kind == ExprNodeKind.PARAM_REF: + stack.append(rows[int(operand)]) + elif kind == ExprNodeKind.ADD: + b, a = stack.pop(), stack.pop() + stack.append(a + b) + elif kind == ExprNodeKind.SUB: + b, a = stack.pop(), stack.pop() + stack.append(a - b) + elif kind == ExprNodeKind.MUL: + b, a = stack.pop(), stack.pop() + stack.append(a * b) + elif kind == ExprNodeKind.DIV: + b, a = stack.pop(), stack.pop() + stack.append(a / b) + elif kind == ExprNodeKind.NEG: + stack.append(-stack.pop()) + elif kind == ExprNodeKind.POW: + b, a = stack.pop(), stack.pop() + stack.append(a**b) + + assert len(stack) == 1 + result = stack[0] + if isinstance(result, float): # constant-only program + return jnp.full(shape, result, dtype=jnp.float64) + return result + + +# --------------------------------------------------------------------------- +# Evaluator factory +# --------------------------------------------------------------------------- + + +# +def _check_plan_supported(plan: ScheduledPlan2D) -> None: + """Reject plan features the JAX backend has no kernel for. + + Callers should gate at the graph level with ``can_lower_jax_2d``; + this is the defensive plan-level check for plans built directly. + """ + + unsupported: list[str] = [] + for op_idx in range(plan.n_ops): + if int(plan.op_kinds[op_idx]) not in JAX_OP_DISPATCH: + unsupported.append(f"op kind {OpKind(int(plan.op_kinds[op_idx])).name}") + for step in range(plan.n_conv_steps): + if int(plan.conv_func_ids[step]) not in JAX_CONV_KERNEL_DISPATCH: + unsupported.append( + f"conv kernel {ConvKernelKind(int(plan.conv_func_ids[step])).name}" + ) + if unsupported: + raise ValueError( + "Plan is not supported by the JAX evaluator: " + + ", ".join(sorted(set(unsupported))) + + ". Gate with can_lower_jax_2d(graph) and use the NumPy " + "evaluator for unsupported models." + ) + + +# +def _build_evaluate_2d(plan: ScheduledPlan2D) -> Callable: + """Build the pure traced function ``theta -> (n_time, n_energy)``. + + Shared by ``make_evaluator_2d_jax`` (jit) and + ``make_jacobian_2d_jax`` (jit of jacfwd). + """ + + n_opt = len(plan.opt_indices) + n_time = plan.n_time + + # + def _evaluate(theta): + # Per-parameter trace rows; list mutation is trace-time only. + rows = [jnp.asarray(plan.param_traces_init[i]) for i in range(plan.n_params)] + for k in range(n_opt): + rows[int(plan.opt_indices[k])] = jnp.broadcast_to(theta[k], (n_time,)) + + # Interleaved dynamics/expression/convolution resolution in + # schedule order. + for step in range(len(plan.resolution_kinds)): + kind = int(plan.resolution_kinds[step]) + idx = int(plan.resolution_indices[step]) + if kind == 0: # dynamics group + target = int(plan.dyn_group_target_row[idx]) + acc = rows[int(plan.dyn_group_base_row[idx])] + s_start = int(plan.dyn_group_indptr[idx]) + s_end = int(plan.dyn_group_indptr[idx + 1]) + for s in range(s_start, s_end): + func = JAX_DYNAMICS_DISPATCH[int(plan.dyn_sub_func_id[s])] + n_par = int(plan.dyn_sub_n_params[s]) + # t=0 read is exact: substep param rows are + # time-constant by construction (see eval_2d). + dyn_params = [ + rows[int(row)][0] for row in plan.dyn_sub_param_rows[s, :n_par] + ] + acc = acc + func( + jnp.asarray(plan.dyn_sub_time_axes[s]), *dyn_params + ) * jnp.asarray(plan.dyn_sub_masks[s]) + rows[target] = acc + elif kind == 1: # expression + target = int(plan.expr_target_rows[idx]) + start = int(plan.expr_indptr[idx]) + end = int(plan.expr_indptr[idx + 1]) + rows[target] = _eval_expr_rows( + plan.expr_instructions[start:end], rows, (n_time,) + ) + else: # kind == 2: resolved-trace convolution + operator = plan.conv_operator + assert operator is not None # type guard: set when steps exist + target = int(plan.conv_target_rows[idx]) + func_id = int(plan.conv_func_ids[idx]) + p_start = int(plan.conv_param_indptr[idx]) + p_end = int(plan.conv_param_indptr[idx + 1]) + # t=0 read is exact; same time-constant invariant as above + kernel_params = [ + rows[int(plan.conv_param_rows[j])][0] for j in range(p_start, p_end) + ] + mass_left, mass_right = JAX_CONV_EDGE_MASS_DISPATCH[func_id]( + jnp.asarray(operator.dt_left), + jnp.asarray(operator.dt_right), + *kernel_params, + ) + kernel_values = JAX_CONV_KERNEL_DISPATCH[func_id]( + jnp.asarray(operator.dt_unique), *kernel_params + ) + # conv_matrix_apply without its host-side value checks + interior = kernel_values[operator.gather_idx] * jnp.asarray( + operator.quad_weights + ) + row_sums = interior.sum(axis=1) + mass_left + mass_right + y = rows[target] + y_conv = interior @ y + mass_left * y[0] + mass_right * y[-1] + rows[target] = y_conv / row_sums + + # Profile sample groups -> (n_time, n_aux) per group, then + # per-sample profile expressions over broadcast virtual rows. + n_aux = plan.n_aux + aux_2d = jnp.asarray(plan.aux_axis)[jnp.newaxis, :] + profile_samples: list = [] + for g in range(plan.n_profile_samples): + base_row = int(plan.profile_sample_base_rows[g]) + value = jnp.broadcast_to(rows[base_row][:, jnp.newaxis], (n_time, n_aux)) + c_start = int(plan.profile_sample_component_indptr[g]) + c_end = int(plan.profile_sample_component_indptr[g + 1]) + for c in range(c_start, c_end): + func = JAX_PROFILE_DISPATCH[int(plan.profile_component_func_ids[c])] + p_start = int(plan.profile_component_param_indptr[c]) + p_end = int(plan.profile_component_param_indptr[c + 1]) + params = [ + rows[int(row)][:, jnp.newaxis] + for row in plan.profile_component_param_rows[p_start:p_end] + ] + value = value + func(aux_2d, *params) + profile_samples.append(value) + + profile_exprs: list = [] + if plan.n_profile_exprs > 0: + virtual_rows = [ + jnp.broadcast_to(row[:, jnp.newaxis], (n_time, n_aux)) for row in rows + ] + profile_samples + for e in range(plan.n_profile_exprs): + start = int(plan.profile_expr_indptr[e]) + end = int(plan.profile_expr_indptr[e + 1]) + profile_exprs.append( + _eval_expr_rows( + plan.profile_expr_instructions[start:end], + virtual_rows, + (n_time, n_aux), + ) + ) + + # Component evaluation, unrolled in schedule order. + energy = jnp.asarray(plan.energy)[jnp.newaxis, :] + result = jnp.asarray(plan.cached_result) + peak_sum = jnp.asarray(plan.cached_peak_sum) + for op_idx in range(plan.n_ops): + if plan.op_is_constant[op_idx]: + continue + func = JAX_OP_DISPATCH[int(plan.op_kinds[op_idx])] + start = int(plan.op_param_indptr[op_idx]) + end = int(plan.op_param_indptr[op_idx + 1]) + if plan.op_is_profiled[op_idx]: + # Vectorized over aux (axis 1), then averaged. The + # NumPy path loops per aux point to avoid materialized + # (n_time, n_aux, n_energy) temporaries; under XLA the + # fused form wins on simplicity and lets the compiler + # decide. + sources = [] + for sk, si in zip( + plan.op_param_source_kinds[start:end], + plan.op_param_indices[start:end], + strict=True, + ): + if int(sk) == int(ParamSourceKind.SCALAR): + source = jnp.broadcast_to( + rows[int(si)][:, jnp.newaxis], (n_time, n_aux) + ) + elif int(sk) == int(ParamSourceKind.PROFILE_SAMPLE): + source = profile_samples[int(si)] + else: + source = profile_exprs[int(si)] + sources.append(source[:, :, jnp.newaxis]) # (n_time, n_aux, 1) + if plan.op_needs_spectrum[op_idx]: + component = func( + energy[jnp.newaxis, :, :], + *sources, + peak_sum[:, jnp.newaxis, :], + ).mean(axis=1) + else: + component = func(energy[jnp.newaxis, :, :], *sources).mean(axis=1) + else: + params = [ + rows[int(row)][:, jnp.newaxis] + for row in plan.op_param_indices[start:end] + ] + if plan.op_needs_spectrum[op_idx]: + component = func(energy, *params, peak_sum) + else: + component = func(energy, *params) + result = result + component + if plan.op_is_pre_spectrum[op_idx]: + peak_sum = peak_sum + component + + return result + + return _evaluate + + +# +def _wrap_jitted(jitted: Callable, n_opt: int) -> Callable[[np.ndarray], np.ndarray]: + """Host-side wrapper: theta validation outside the jitted region.""" + + # + def evaluate(theta: np.ndarray) -> np.ndarray: + theta_arr = np.asarray(theta, dtype=np.float64) + if theta_arr.shape != (n_opt,): + raise ValueError( + f"theta shape {theta_arr.shape} does not match " + f"plan.opt_indices length {n_opt}" + ) + return np.asarray(jitted(theta_arr)) + + return evaluate + + +# +def make_evaluator_2d_jax( + plan: ScheduledPlan2D, +) -> Callable[[np.ndarray], np.ndarray]: + """Compile a jitted JAX evaluator for a 2D scheduled plan. + + Parameters + ---------- + plan : ScheduledPlan2D + Compiled 2D execution schedule (from ``schedule_2d``). The + source graph must pass ``can_lower_jax_2d``. + + Returns + ------- + Callable[[np.ndarray], np.ndarray] + ``evaluate(theta) -> (n_time, n_energy)`` with the same theta + contract as ``evaluate_2d(plan, theta)``. The first call + triggers XLA compilation; subsequent calls reuse it. + + Raises + ------ + ImportError + If jax is not installed. + ValueError + If the plan uses features outside the JAX slice. + """ + + _require_jax() + _check_plan_supported(plan) + return _wrap_jitted(jax.jit(_build_evaluate_2d(plan)), len(plan.opt_indices)) + + +# +def make_jacobian_2d_jax( + plan: ScheduledPlan2D, +) -> Callable[[np.ndarray], np.ndarray]: + """Compile a jitted analytic Jacobian of the 2D model output. + + Forward-mode (``jax.jacfwd``): the optimizer vector is short and the + output grid is large, so one JVP pass per theta entry is the right + direction. + + Derivative caveats (all measure-zero or by-construction flat): + ``boxCONV``'s derivative w.r.t. its width is 0 almost everywhere + (hard edges), and step-like ``where`` branches (stepFun onset, + GaussAsym at x0) contribute subgradient-style zeros at the exact + switching point. + + Parameters + ---------- + plan : ScheduledPlan2D + Compiled 2D execution schedule (from ``schedule_2d``). The + source graph must pass ``can_lower_jax_2d``. + + Returns + ------- + Callable[[np.ndarray], np.ndarray] + ``jacobian(theta) -> (n_time, n_energy, n_opt)`` — derivative + of the model spectrum w.r.t. each optimizer parameter, columns + in ``plan.opt_param_names`` order. + + Raises + ------ + ImportError + If jax is not installed. + ValueError + If the plan uses features outside the JAX slice. + """ + + _require_jax() + _check_plan_supported(plan) + jitted = jax.jit(jax.jacfwd(_build_evaluate_2d(plan))) + return _wrap_jitted(jitted, len(plan.opt_indices)) diff --git a/src/trspecfit/fitlib.py b/src/trspecfit/fitlib.py index b219763..b79cf13 100644 --- a/src/trspecfit/fitlib.py +++ b/src/trspecfit/fitlib.py @@ -23,7 +23,7 @@ import math import pathlib import time -from collections.abc import Sequence +from collections.abc import Callable, Sequence from typing import Any, cast import corner @@ -287,6 +287,70 @@ def residual_fun( raise ValueError(f"Unknown res_type '{res_type}'") +# +def jacobian_fun( + 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 Jacobian of :func:`residual_fun` for lmfit's ``Dfun``. + + The signature mirrors ``residual_fun`` because lmfit calls the + Jacobian with the same ``fcn_args``. Requires the + ``fit_model_jax`` dispatch convention: + ``args = (evaluator, jacobian, theta_indices, model, dim)`` with + *jacobian* from ``eval_jax.make_jacobian_2d_jax``. + + Returns + ------- + ndarray + ``d(residual)/d(varying params)``, shape + ``(n_residuals, n_varys)``, columns in lmfit varying-parameter + order (``col_deriv=0``). Residual is ``data - fit``, so this + is the negated model Jacobian over the fit window. + """ + + if e_lim is None: + e_lim = [] + if t_lim is None: + t_lim = [] + if args is None or not callable(args[1]): + raise ValueError( + "jacobian_fun requires the fit_model_jax dispatch args " + "(evaluator, jacobian, theta_indices, model, dim)." + ) + jacobian = args[1] + theta_indices: np.ndarray = args[2] + model = args[3] + + par_values = np.asarray( + ulmfit.par_extract(par, return_type="list"), dtype=np.float64 + ) + # (n_time, n_energy, n_opt) + jac = np.asarray(jacobian(par_values[theta_indices]), dtype=np.float64) + + window = _fit_window_slices(2, e_lim, t_lim) + n_opt = jac.shape[-1] + d_res = -jac[window].reshape(-1, n_opt) + + # Column order: plan opt order -> lmfit varying-parameter order. + opt_names = [model.parameter_names[int(i)] for i in theta_indices] + var_names = [name for name in par if par[name].vary] + if sorted(opt_names) != sorted(var_names): + raise RuntimeError( + "JAX Jacobian column mismatch: plan optimizer parameters " + f"{opt_names} do not match lmfit varying parameters {var_names}." + ) + columns = [opt_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 @@ -451,6 +515,7 @@ def fit_wrapper( mc_settings: ulmfit.MC | None = None, fit_alg_1: str = "Nelder", fit_alg_2: str = "leastsq", + jac_fun: Callable[..., np.ndarray] | None = None, show_output: int = 0, save_output: int = 0, save_path: PathLike = "", @@ -523,6 +588,11 @@ def fit_wrapper( fit_alg_2 : str, default='leastsq' Second optimization method (stages=2 only). Typically 'leastsq' for accurate local optimization and error bars. + jac_fun : callable, optional + Analytic Jacobian with the same signature as ``residual_fun`` + (e.g. :func:`jacobian_fun` on the JAX backend). Passed to + lmfit as ``Dfun`` for stages whose method is ``'leastsq'``; + ignored for gradient-free methods. show_output : {0, 1}, default=0 Output mode: @@ -678,13 +748,20 @@ def fit_wrapper( # construct lmfit minimizer mini = lmfit.Minimizer(residual_fun, par_ini, fcn_args=(*const, "lmfit", args)) + + # analytic Jacobian: only lmfit's leastsq accepts a Dfun + def _method_kws(method: str) -> dict[str, Any]: + if jac_fun is not None and method == "leastsq": + return {"Dfun": jac_fun, "col_deriv": 0} + return {} + # perform fit(s) if show_output >= 1: t_ini = time.time() print(f"\nTime initialize: {t_ini - t_0} s") # if stages == 1: # one fit only - par_fin = mini.minimize(method=fit_alg_1) + par_fin = mini.minimize(method=fit_alg_1, **_method_kws(fit_alg_1)) par_fin_params = _result_params(par_fin) if show_output >= 1: print(f"\nResults fit (method={fit_alg_1}): ") @@ -693,7 +770,7 @@ def fit_wrapper( print(f"Time fit: {t_fit - t_ini} s") # if stages == 2: # find global minimum + local optimization - par_fin_gm = mini.minimize(method=fit_alg_1) + par_fin_gm = mini.minimize(method=fit_alg_1, **_method_kws(fit_alg_1)) par_fin_gm_params = _result_params(par_fin_gm) if show_output >= 1: print(f"\nResults global minumum fit (method={fit_alg_1}): ") @@ -701,7 +778,9 @@ def fit_wrapper( t_fit0 = time.time() print(f"Time fit (global minimum): {t_fit0 - t_ini} s") # - par_fin = mini.minimize(method=fit_alg_2, params=par_fin_gm_params) + par_fin = mini.minimize( + method=fit_alg_2, params=par_fin_gm_params, **_method_kws(fit_alg_2) + ) par_fin_params = _result_params(par_fin) if show_output >= 1: print(f"\nResults local optimization fit (method={fit_alg_2}): ") diff --git a/src/trspecfit/graph_ir.py b/src/trspecfit/graph_ir.py index 2540e18..303ef09 100644 --- a/src/trspecfit/graph_ir.py +++ b/src/trspecfit/graph_ir.py @@ -441,29 +441,13 @@ def to_dot(self, *, collapse_profiles: bool = True) -> str: return "\n".join(lines) -# -# -@dataclass(frozen=True) -class ExprProgram: - """Compiled expression: flat int array encoding an RPN program. - - Encoding: pairs of ``(node_kind, operand)``. - - - ``CONST``: operand is float bits (``np.float64.view(np.int64)``) - - ``PARAM_REF``: operand is row index into trace matrix - - Operators: operand is 0 (unused) - """ - - instructions: np.ndarray # (2 * n_instructions,) int64 - - # # @dataclass(frozen=True) class ScheduledPlan2D: """Compiled 2D execution schedule. - No Python objects in the hot path (except ``expr_programs``). + No Python objects in the hot path. """ energy: np.ndarray # (n_energy,) @@ -497,9 +481,16 @@ class ScheduledPlan2D: dyn_sub_masks: np.ndarray # (n_substeps, n_time) float64 # --- Expression evaluation --- + # Packed RPN programs, CSR-style: program ``i`` occupies + # ``expr_instructions[expr_indptr[i]:expr_indptr[i + 1]]``. + # Instructions are ``(node_kind, operand)`` int64 pairs: + # CONST -> operand is float bits (np.float64.view(np.int64)) + # PARAM_REF -> operand is row index into the trace matrix + # operators -> operand is 0 (unused) n_expressions: int expr_target_rows: np.ndarray # (n_expressions,) int - expr_programs: list[ExprProgram] + expr_instructions: np.ndarray # (total_expr_words,) int64 + expr_indptr: np.ndarray # (n_expressions + 1,) int # --- Interleaved parameter resolution schedule --- # Dynamics groups, expressions, and convolution steps may depend on @@ -508,7 +499,7 @@ class ScheduledPlan2D: # resolution_kinds / resolution_indices arrays encode the correct # topological execution order: # kind=0 -> dynamics group step, index into dyn_group_* arrays - # kind=1 -> expression step, index into expr_* arrays / expr_programs + # kind=1 -> expression step, index into expr_* arrays # kind=2 -> convolution step, index into conv_* arrays resolution_kinds: np.ndarray # (n_dyn_groups + n_expressions + n_conv_steps,) int8 resolution_indices: np.ndarray # (n_dyn_groups + n_expressions + n_conv_steps,) int @@ -539,8 +530,12 @@ class ScheduledPlan2D: profile_component_func_ids: np.ndarray # (n_profile_components,) int profile_component_param_indptr: np.ndarray # (n_profile_components + 1,) int profile_component_param_rows: np.ndarray # (total_profile_component_params,) int + # Packed per-sample profile expressions (same encoding as + # expr_instructions; PARAM_REF operands >= n_params index profile + # sample groups). n_profile_exprs: int - profile_expr_programs: list[ExprProgram] + profile_expr_instructions: np.ndarray # (total_profile_expr_words,) int64 + profile_expr_indptr: np.ndarray # (n_profile_exprs + 1,) int # --- Scheduled component ops --- n_ops: int @@ -577,9 +572,12 @@ class ScheduledPlan1D: opt_param_names: list[str] # (n_opt,) canonical optimizer param names # --- Expression evaluation (topological order, no dynamics) --- + # Packed RPN programs, CSR-style (see ScheduledPlan2D for the + # encoding); PARAM_REF operands index the scalar parameter vector. n_expressions: int expr_target_indices: np.ndarray # (n_expressions,) int - expr_programs: list[ExprProgram] + expr_instructions: np.ndarray # (total_expr_words,) int64 + expr_indptr: np.ndarray # (n_expressions + 1,) int # --- Profile-varying parameter groups (fixed aux_axis shape) --- n_aux: int @@ -591,7 +589,8 @@ class ScheduledPlan1D: profile_component_param_indptr: np.ndarray # (n_profile_components + 1,) int profile_component_param_indices: np.ndarray # (total_profile_component_params,) int n_profile_exprs: int - profile_expr_programs: list[ExprProgram] + profile_expr_instructions: np.ndarray # (total_profile_expr_words,) int64 + profile_expr_indptr: np.ndarray # (n_profile_exprs + 1,) int # --- Scheduled component ops --- n_ops: int @@ -1672,6 +1671,53 @@ def can_lower_2d(graph: GraphIR) -> bool: return True +# --------------------------------------------------------------------------- +# can_lower_jax_2d +# --------------------------------------------------------------------------- + +# JAX backend (docs/design/jax-planning.md, Phases B + C): covers the +# full lowered 2D surface, including profiles, convolution, subcycle +# dynamics, and Voigt (Weideman wofz approximation in eval_jax). The +# sets exist so future NumPy-side widening does not silently imply JAX +# support — carve out here first, widen JAX after kernels land. +_LOWERABLE_JAX_2D_FUNCTIONS: frozenset[str] = _LOWERABLE_2D_FUNCTIONS + +_NON_LOWERABLE_JAX_2D_NODE_KINDS: frozenset[NodeKind] = _NON_LOWERABLE_2D_NODE_KINDS + + +# +def can_lower_jax_2d(graph: GraphIR) -> bool: + """Check whether the experimental 2D JAX backend can compile this graph. + + At most as wide as :func:`can_lower_2d` by construction, so a graph + that fails this gate but passes ``can_lower_2d`` falls back to the + compiled NumPy evaluator (never straight to the interpreter). + + Parameters + ---------- + graph : GraphIR + The model graph to check. + + Returns + ------- + bool + True if the JAX evaluator supports this graph. + """ + + if not can_lower_2d(graph): + return False + + for node in graph.nodes: + if node.kind in _NON_LOWERABLE_JAX_2D_NODE_KINDS: + return False + + if node.kind in (NodeKind.COMPONENT_EVAL, NodeKind.SPECTRUM_FED_OP): + if node.function_name not in _LOWERABLE_JAX_2D_FUNCTIONS: + return False + + return True + + # Node kinds that are never valid in 1D energy models. 1D models have # no time axis, so DYNAMICS_TRACE / PARAM_PLUS_TRACE should not appear. # Start from the 2D blocklist so future unsupported node kinds propagate @@ -1778,7 +1824,7 @@ class SymbolicRPN: This is the *frontend* output of the expression compiler. ``schedule_2d`` binds names to trace-matrix row indices and - produces the final ``ExprProgram``. + produces the final packed instruction array. Each instruction is a ``(ExprNodeKind, operand)`` pair: @@ -1888,8 +1934,8 @@ def _walk(node: ast.AST) -> None: def _bind_expr_to_rows( symbolic: SymbolicRPN, name_to_row: dict[str, int], -) -> ExprProgram: - """Convert a symbolic RPN program to a row-bound ExprProgram. +) -> np.ndarray: + """Convert a symbolic RPN program to a row-bound instruction array. Parameters ---------- @@ -1900,8 +1946,10 @@ def _bind_expr_to_rows( Returns ------- - ExprProgram - Row-bound RPN program ready for the evaluator. + np.ndarray + ``(2 * n_instructions,)`` int64 array of ``(node_kind, operand)`` + pairs ready for the evaluator (see ``ScheduledPlan2D`` for the + operand encoding). """ flat: list[int] = [] @@ -1915,7 +1963,31 @@ def _bind_expr_to_rows( flat.append(name_to_row[operand]) else: flat.append(0) - return ExprProgram(instructions=np.array(flat, dtype=np.int64)) + return np.array(flat, dtype=np.int64) + + +# +def _pack_expr_programs( + programs: list[np.ndarray], +) -> tuple[np.ndarray, np.ndarray]: + """Concatenate row-bound RPN programs into CSR-style packed arrays. + + Returns + ------- + tuple[np.ndarray, np.ndarray] + ``(instructions, indptr)``: flat int64 instruction words plus + ``(n_programs + 1,)`` offsets; program ``i`` occupies + ``instructions[indptr[i]:indptr[i + 1]]``. + """ + + indptr = np.zeros(len(programs) + 1, dtype=np.intp) + for i, program in enumerate(programs): + indptr[i + 1] = indptr[i] + len(program) + if programs: + instructions = np.concatenate(programs) + else: + instructions = np.zeros(0, dtype=np.int64) + return instructions, indptr # --------------------------------------------------------------------------- @@ -2327,7 +2399,7 @@ def schedule_2d(graph: GraphIR) -> ScheduledPlan2D: expr_ref_maps[expr_node.id] = ref_map # Compile and bind expressions - expr_programs: list[ExprProgram] = [] + expr_programs: list[np.ndarray] = [] expr_target_rows_list: list[int] = [] for expr_node in expr_nodes_topo: assert expr_node.expr_string is not None @@ -2347,6 +2419,7 @@ def schedule_2d(graph: GraphIR) -> ScheduledPlan2D: ) expr_target_rows = np.array(expr_target_rows_list, dtype=np.intp) + expr_instructions, expr_indptr = _pack_expr_programs(expr_programs) # Build resolution schedule: map DYNAMICS_TRACE node ids to their # group index, and emit each group exactly once (on the *last* trace @@ -2480,7 +2553,8 @@ def schedule_2d(graph: GraphIR) -> ScheduledPlan2D: profile_component_param_indptr = profiles.component_param_indptr profile_component_param_rows = profiles.component_param_indices n_profile_exprs = profiles.n_exprs - profile_expr_programs_2d = profiles.expr_programs + profile_expr_instructions_2d = profiles.expr_instructions + profile_expr_indptr_2d = profiles.expr_indptr # ------------------------------------------------------------------ # # 5. Schedule component ops # @@ -2534,7 +2608,8 @@ def schedule_2d(graph: GraphIR) -> ScheduledPlan2D: dyn_sub_time_axes, dyn_sub_masks, expr_target_rows, - expr_programs, + expr_instructions, + expr_indptr, conv_target_rows, conv_func_ids, conv_param_indptr, @@ -2564,7 +2639,8 @@ def schedule_2d(graph: GraphIR) -> ScheduledPlan2D: param_traces_init, profile_sample_values_init, n_params, - profile_expr_programs_2d, + profile_expr_instructions_2d, + profile_expr_indptr_2d, ) energy = graph.energy[np.newaxis, :] @@ -2628,7 +2704,8 @@ def schedule_2d(graph: GraphIR) -> ScheduledPlan2D: dyn_sub_masks=dyn_sub_masks, n_expressions=n_expressions, expr_target_rows=expr_target_rows, - expr_programs=expr_programs, + expr_instructions=expr_instructions, + expr_indptr=expr_indptr, resolution_kinds=resolution_kinds, resolution_indices=resolution_indices, n_aux=n_aux, @@ -2640,7 +2717,8 @@ def schedule_2d(graph: GraphIR) -> ScheduledPlan2D: profile_component_param_indptr=profile_component_param_indptr, profile_component_param_rows=profile_component_param_rows, n_profile_exprs=n_profile_exprs, - profile_expr_programs=profile_expr_programs_2d, + profile_expr_instructions=profile_expr_instructions_2d, + profile_expr_indptr=profile_expr_indptr_2d, n_ops=n_ops, op_schedule=op_schedule, op_kinds=op_kinds, @@ -2668,13 +2746,14 @@ def schedule_2d(graph: GraphIR) -> ScheduledPlan2D: # -def _eval_expr_scalar(program: ExprProgram, values: np.ndarray) -> float: - """Evaluate an RPN ExprProgram against a scalar parameter vector. +def _eval_expr_scalar(instructions: np.ndarray, values: np.ndarray) -> float: + """Evaluate a compiled RPN program against a scalar parameter vector. Parameters ---------- - program - Compiled RPN instruction array. + instructions + Compiled RPN instruction array (one program's slice of the + packed ``expr_instructions``). values ``(n_params,)`` scalar parameter vector. @@ -2685,7 +2764,7 @@ def _eval_expr_scalar(program: ExprProgram, values: np.ndarray) -> float: """ stack: list[float] = [] - instr = program.instructions + instr = instructions n_instr = len(instr) // 2 for i in range(n_instr): @@ -2797,7 +2876,8 @@ class _CompiledProfileGroups(NamedTuple): sample_is_constant: np.ndarray sample_group_idx: dict[str, int] n_exprs: int - expr_programs: list[ExprProgram] + expr_instructions: np.ndarray + expr_indptr: np.ndarray expr_is_constant: np.ndarray expr_group_idx: dict[str, int] @@ -2950,7 +3030,7 @@ def _compile_profile_groups( group_name = _profile_group_name(node.name, "profile_expr") profile_expr_groups.setdefault(group_name, []).append(node) - expr_programs: list[ExprProgram] = [] + expr_programs: list[np.ndarray] = [] expr_is_constant_list: list[bool] = [] expr_group_idx: dict[str, int] = {} for group_name, p_expr_nodes in profile_expr_groups.items(): @@ -3015,6 +3095,8 @@ def _compile_profile_groups( expr_is_constant_list.append(is_constant) expr_group_idx[group_name] = len(expr_programs) - 1 + expr_instructions, expr_indptr = _pack_expr_programs(expr_programs) + return _CompiledProfileGroups( aux_axis=plan_aux_axis, n_aux=n_aux, @@ -3027,7 +3109,8 @@ def _compile_profile_groups( sample_is_constant=sample_is_constant, sample_group_idx=sample_group_idx, n_exprs=len(expr_programs), - expr_programs=expr_programs, + expr_instructions=expr_instructions, + expr_indptr=expr_indptr, expr_is_constant=np.array(expr_is_constant_list, dtype=np.bool_), expr_group_idx=expr_group_idx, ) @@ -3270,12 +3353,12 @@ def _schedule_component_ops( # -def _eval_expr_vector(program: ExprProgram, traces: np.ndarray) -> np.ndarray: - """Evaluate an RPN ExprProgram against an aux-resolved trace matrix.""" +def _eval_expr_vector(instructions: np.ndarray, traces: np.ndarray) -> np.ndarray: + """Evaluate a compiled RPN program against an aux-resolved trace matrix.""" from trspecfit.eval_2d import eval_expr_program - return eval_expr_program(program, traces) + return eval_expr_program(instructions, traces) # @@ -3322,11 +3405,12 @@ def _evaluate_profile_expr_values( scalar_values: np.ndarray, profile_sample_values: np.ndarray, n_params: int, - profile_expr_programs: list[ExprProgram], + profile_expr_instructions: np.ndarray, + profile_expr_indptr: np.ndarray, ) -> np.ndarray: """Evaluate lowered per-sample profile expressions over the aux axis.""" - n_exprs = len(profile_expr_programs) + n_exprs = len(profile_expr_indptr) - 1 if n_exprs == 0: n_aux = profile_sample_values.shape[1] if profile_sample_values.size else 0 return np.zeros((0, n_aux), dtype=np.float64) @@ -3340,8 +3424,12 @@ def _evaluate_profile_expr_values( traces[n_params:, :] = profile_sample_values expr_values = np.empty((n_exprs, n_aux), dtype=np.float64) - for expr_idx, program in enumerate(profile_expr_programs): - expr_values[expr_idx, :] = _eval_expr_vector(program, traces) + for expr_idx in range(n_exprs): + start = int(profile_expr_indptr[expr_idx]) + end = int(profile_expr_indptr[expr_idx + 1]) + expr_values[expr_idx, :] = _eval_expr_vector( + profile_expr_instructions[start:end], traces + ) return expr_values @@ -3486,7 +3574,7 @@ def schedule_1d(graph: GraphIR) -> ScheduledPlan1D: if id_to_node[nid].kind == NodeKind.EXPRESSION and not _is_profile_expr_node(id_to_node[nid]) ] - expr_programs: list[ExprProgram] = [] + expr_programs: list[np.ndarray] = [] expr_target_indices_list: list[int] = [] for expr_node in expr_nodes_topo: assert expr_node.expr_string is not None @@ -3513,6 +3601,7 @@ def schedule_1d(graph: GraphIR) -> ScheduledPlan1D: n_expressions = len(expr_programs) expr_target_indices = np.array(expr_target_indices_list, dtype=np.intp) + expr_instructions, expr_indptr = _pack_expr_programs(expr_programs) # ------------------------------------------------------------------ # # 4. Compile profile groups (samples + expressions) # @@ -3536,7 +3625,8 @@ def schedule_1d(graph: GraphIR) -> ScheduledPlan1D: profile_component_param_indptr = profiles.component_param_indptr profile_component_param_indices = profiles.component_param_indices n_profile_exprs = profiles.n_exprs - profile_expr_programs = profiles.expr_programs + profile_expr_instructions = profiles.expr_instructions + profile_expr_indptr = profiles.expr_indptr # ------------------------------------------------------------------ # # 5. Schedule component ops # @@ -3574,7 +3664,8 @@ def schedule_1d(graph: GraphIR) -> ScheduledPlan1D: for i in range(n_expressions): target_idx = int(expr_target_indices[i]) param_values_init[target_idx] = _eval_expr_scalar( - expr_programs[i], param_values_init + expr_instructions[expr_indptr[i] : expr_indptr[i + 1]], + param_values_init, ) profile_sample_values_init = _evaluate_profile_sample_values( @@ -3590,7 +3681,8 @@ def schedule_1d(graph: GraphIR) -> ScheduledPlan1D: param_values_init, profile_sample_values_init, n_params, - profile_expr_programs, + profile_expr_instructions, + profile_expr_indptr, ) # ------------------------------------------------------------------ # @@ -3634,7 +3726,8 @@ def schedule_1d(graph: GraphIR) -> ScheduledPlan1D: opt_param_names=opt_param_names, n_expressions=n_expressions, expr_target_indices=expr_target_indices, - expr_programs=expr_programs, + expr_instructions=expr_instructions, + expr_indptr=expr_indptr, n_aux=n_aux, aux_axis=plan_aux_axis, n_profile_samples=n_profile_samples, @@ -3644,7 +3737,8 @@ def schedule_1d(graph: GraphIR) -> ScheduledPlan1D: profile_component_param_indptr=profile_component_param_indptr, profile_component_param_indices=profile_component_param_indices, n_profile_exprs=n_profile_exprs, - profile_expr_programs=profile_expr_programs, + profile_expr_instructions=profile_expr_instructions, + profile_expr_indptr=profile_expr_indptr, n_ops=n_ops, op_kinds=op_kinds, op_param_indptr=op_param_indptr, diff --git a/src/trspecfit/spectra.py b/src/trspecfit/spectra.py index 781b038..6294c3f 100644 --- a/src/trspecfit/spectra.py +++ b/src/trspecfit/spectra.py @@ -201,6 +201,50 @@ def fit_model_gir( return fit_model_mcp(x, par, plot_sum, *args) +# +def fit_model_jax( + x: Sequence[float] | np.ndarray, + par: Sequence[float] | np.ndarray, + plot_sum: bool, + *args: Any, +) -> np.ndarray | list[np.ndarray]: + """Generate spectrum using a compiled JAX evaluator when available. + + When the first element of *args* is a callable it is the jitted + evaluator from ``eval_jax.make_evaluator_2d_jax``; otherwise the + call is forwarded to :func:`fit_model_gir` (NumPy compiled plan or + interpreter fallback). + + Parameters + ---------- + x : array-like + Independent variable axis (energy or time). + par : array-like + Full parameter vector (all params, fixed + varying). + plot_sum : bool + Component return mode (2D JAX path always returns the sum). + *args + ``(evaluator, jacobian, theta_indices, model, dim)`` for the + JAX path — *jacobian* is carried for ``fitlib.jacobian_fun`` + (lmfit ``Dfun``), not used here. Otherwise the + :func:`fit_model_gir` conventions apply. + + Notes + ----- + The evaluator/jacobian entries are per-plan closures and do not + pickle; MCMC via ``lmfit.emcee`` with ``workers > 1`` is not + supported on this path (single-worker MCMC works). + """ + + if callable(args[0]): + evaluator = args[0] + theta_indices: np.ndarray = args[2] + par_arr = np.asarray(par, dtype=np.float64) + return np.asarray(evaluator(par_arr[theta_indices])) + + return fit_model_gir(x, par, plot_sum, *args) + + # def fit_model_compare( x: Sequence[float] | np.ndarray, diff --git a/src/trspecfit/trspecfit.py b/src/trspecfit/trspecfit.py index 14b47e0..71ec7e3 100644 --- a/src/trspecfit/trspecfit.py +++ b/src/trspecfit/trspecfit.py @@ -176,7 +176,11 @@ class Project: save plots and data spec_fun_str : str Name of fitting function in ``trspecfit.spectra`` (e.g. - ``'fit_model_gir'``, ``'fit_model_mcp'``, ``'fit_model_compare'``) + ``'fit_model_gir'``, ``'fit_model_mcp'``, ``'fit_model_compare'``, + ``'fit_model_jax'``). ``'fit_model_jax'`` runs 2D fits on the + experimental JAX backend with an analytic Jacobian for leastsq + stages (requires the ``[jax]`` extra); models the JAX gate + rejects fall back to the compiled NumPy path. Notes ----- @@ -3738,9 +3742,12 @@ def _build_1d_dispatch_args( runtime still evaluates against the original model object. """ - if fit_fun_str not in ("fit_model_gir", "fit_model_compare"): + if fit_fun_str not in ("fit_model_gir", "fit_model_compare", "fit_model_jax"): return (model, 1) + # fit_model_jax lowers 1D to the compiled NumPy plan (the JAX + # backend is 2D-only); fit_model_jax delegates plan args to + # fit_model_gir. from trspecfit.graph_ir import build_graph, can_lower_1d, schedule_1d model_1d = copy.copy(model) @@ -3996,8 +4003,13 @@ def fit_2d(self, model_name: str, stages: int = 1, **fit_wrapper_kwargs) -> None # --- dispatch: GIR fast path vs interpreter --- _fun_str = self.p.spec_fun_str - if _fun_str in ("fit_model_gir", "fit_model_compare"): - from trspecfit.graph_ir import build_graph, can_lower_2d, schedule_2d + if _fun_str in ("fit_model_gir", "fit_model_compare", "fit_model_jax"): + from trspecfit.graph_ir import ( + build_graph, + can_lower_2d, + can_lower_jax_2d, + schedule_2d, + ) _graph = build_graph(self.model_2d) if can_lower_2d(_graph): @@ -4015,6 +4027,23 @@ def fit_2d(self, model_name: str, stages: int = 1, **fit_wrapper_kwargs) -> None self.model_2d, 2, ) + if _fun_str == "fit_model_jax" and can_lower_jax_2d(_graph): + from trspecfit.eval_jax import ( + make_evaluator_2d_jax, + make_jacobian_2d_jax, + ) + + _args = ( + make_evaluator_2d_jax(_plan), + make_jacobian_2d_jax(_plan), + _theta_indices, + self.model_2d, + 2, + ) + # Analytic Jacobian for leastsq stages (lmfit Dfun) + fit_wrapper_kwargs.setdefault("jac_fun", fitlib.jacobian_fun) + # JAX-non-lowerable stays on the compiled NumPy plan: + # fit_model_jax delegates plan args to fit_model_gir. else: # Not lowerable — fit_model_gir delegates to fit_model_mcp _args = (self.model_2d, 2) diff --git a/tests/models/file_time.yaml b/tests/models/file_time.yaml index a23c3a7..c2bdf31 100644 --- a/tests/models/file_time.yaml +++ b/tests/models/file_time.yaml @@ -81,6 +81,15 @@ MonoSqrt: A: [0.2, True, 0, 5] t0: [0, False, 0, 1] +# sqrtFun with a *varying* onset: regression fixture for the JAX +# Jacobian (the sqrt(clip) form produced NaN derivatives on every +# pre-onset sample when t0 varies). t0 sits off the shared test grid +# so finite differences stay well-defined near the onset. +MonoSqrtVaryT0: + sqrtFun: + A: [0.2, True, 0, 5] + t0: [0.5, True, -2, 5] + # SD ~ 0.8*dt of the shared test time axis: wide enough that the width is # identifiable from the data, narrow enough not to bury the expFun dynamics MonoExpPosIRF: diff --git a/tests/test_evaluate_2d.py b/tests/test_evaluate_2d.py index 745c431..84b6366 100644 --- a/tests/test_evaluate_2d.py +++ b/tests/test_evaluate_2d.py @@ -84,6 +84,41 @@ def _extract_theta(plan, model): ) +try: + import jax as _jax # noqa: F401 + + _HAVE_JAX = True +except ImportError: # pragma: no cover - CI min-versions job has no jax + _HAVE_JAX = False + +# plan id -> (plan, jitted evaluator); the plan reference keeps the id +# from being recycled after garbage collection +_jax_evaluator_cache: dict = {} + + +# +def _assert_jax_parity(plan, theta, reference): + """Assert the JAX evaluator matches the NumPy GIR result. + + Piggybacks on every evaluator-vs-interpreter comparison so the + whole 2D matrix (including regression fixtures) covers the JAX + backend. No-op without jax installed. If the JAX gate ever + narrows below can_lower_2d again, unsupported plans raise + ValueError here — split the affected test rather than skipping + silently. + """ + + if not _HAVE_JAX: + return + key = id(plan) + if key not in _jax_evaluator_cache: + from trspecfit.eval_jax import make_evaluator_2d_jax + + _jax_evaluator_cache[key] = (plan, make_evaluator_2d_jax(plan)) + _, evaluate_jax = _jax_evaluator_cache[key] + np.testing.assert_allclose(evaluate_jax(theta), reference, rtol=1e-12, atol=1e-12) + + # def _compare_evaluator_vs_interpreter(model, plan, *, rtol=1e-10): """Run evaluate_2d and interpreter, assert they match.""" @@ -93,6 +128,7 @@ def _compare_evaluator_vs_interpreter(model, plan, *, rtol=1e-10): model.create_value_2d() slow = model.value_2d np.testing.assert_allclose(fast, slow, rtol=rtol, atol=1e-10) + _assert_jax_parity(plan, theta, fast) return theta @@ -111,6 +147,7 @@ def _perturb_theta(plan, model, theta, indices, deltas): model.create_value_2d() slow = model.value_2d np.testing.assert_allclose(fast, slow, rtol=1e-10, atol=1e-10) + _assert_jax_parity(plan, theta_new, fast) # --------------------------------------------------------------------------- diff --git a/tests/test_evaluate_jax.py b/tests/test_evaluate_jax.py new file mode 100644 index 0000000..fb2d4f4 --- /dev/null +++ b/tests/test_evaluate_jax.py @@ -0,0 +1,372 @@ +"""JAX evaluator backend: gate, Jacobian, wofz accuracy, and e2e fit. + +Broad evaluator parity is NOT here: every evaluator-vs-interpreter +comparison in ``test_evaluate_2d.py`` also asserts JAX parity (via +``_assert_jax_parity``), so the full 2D matrix — including regression +fixtures — covers the JAX backend without duplication. This module +keeps only JAX-specific coverage: the capability gate, the analytic +Jacobian, the Weideman wofz approximation, the chained-convolution +fixture (absent from the 2D matrix), error contracts, and an +end-to-end ``File.fit_2d`` run on ``spec_fun_str="fit_model_jax"``. + +Parity target is ``evaluate_2d`` (the compiled NumPy evaluator), not the +interpreter — the JAX backend's contract is bit-level agreement with the +reference backend it may eventually replace (float64 mode, tight rtol). + +Skips entirely when jax is not installed (optional ``[jax]`` extra). +""" + +import numpy as np +import pytest + +pytest.importorskip("jax") + +from test_evaluate_2d import ( # noqa: E402 + _extract_theta, + _make_2d_model, + _make_2d_profile_model, + _make_energy_model, +) + +from trspecfit.eval_2d import evaluate_2d # noqa: E402 +from trspecfit.eval_jax import ( # noqa: E402 + make_evaluator_2d_jax, + make_jacobian_2d_jax, +) +from trspecfit.graph_ir import ( # noqa: E402 + build_graph, + can_lower_2d, + can_lower_jax_2d, + schedule_2d, +) + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +# +def _make_plan(model_info, dynamics_params, **kwargs): + """Build model + graph + plan; graph must pass the JAX gate.""" + + _file, model = _make_2d_model(model_info, dynamics_params, **kwargs) + graph = build_graph(model) + assert can_lower_jax_2d(graph) + plan = schedule_2d(graph) + return plan, model + + +# +def _assert_parity(plan, model, *, rtol=1e-12, atol=1e-12): + """Compare the jitted JAX evaluator against evaluate_2d. + + Checks the initial theta and a perturbed theta through the same + compiled evaluator (exercises jit reuse, not just first-trace + correctness). + """ + + evaluate_jax = make_evaluator_2d_jax(plan) + theta = _extract_theta(plan, model) + + ref = evaluate_2d(plan, theta) + got = evaluate_jax(theta) + np.testing.assert_allclose(got, ref, rtol=rtol, atol=atol) + + theta_new = theta * 1.05 + 0.01 + ref_new = evaluate_2d(plan, theta_new) + got_new = evaluate_jax(theta_new) + np.testing.assert_allclose(got_new, ref_new, rtol=rtol, atol=atol) + + +# --------------------------------------------------------------------------- +# Gate behavior +# --------------------------------------------------------------------------- + + +# +# +class TestJaxGate: + """can_lower_jax_2d covers the lowered 2D surface, rejects 1D/MCP-only.""" + + _COVERED = [ + (["cached_shirley_peak"], [], {}), + (["glp_only"], [("GLP_01_A", ["MonoExpPos"])], {}), + (["voigt_only"], [], {}), + (["glp_only"], [("GLP_01_A", ["MonoExpPosIRF"])], {}), + ( + ["glp_only"], + [("GLP_01_A", ["ModelNone", "MonoExpNeg"])], + {"frequency": 10}, + ), + ] + + # + @pytest.mark.parametrize( + "model_info,dynamics,kwargs", + _COVERED, + ids=["static", "dynamics", "voigt", "convolution", "subcycle"], + ) + def test_lowered_surface_passes(self, model_info, dynamics, kwargs): + _file, model = _make_2d_model(model_info, dynamics, **kwargs) + graph = build_graph(model) + assert can_lower_2d(graph) + assert can_lower_jax_2d(graph) + + # + def test_profile_model_passes(self): + _file, model = _make_2d_profile_model( + ["single_gauss"], + [("Gauss_01_x0", ["MonoExpPos"])], + [("Gauss_01_A", ["profile_pExpDecay"])], + ) + graph = build_graph(model) + assert can_lower_2d(graph) + assert can_lower_jax_2d(graph) + + # + def test_1d_graph_rejected(self): + _file, model = _make_energy_model(["glp_only"]) + assert not can_lower_jax_2d(build_graph(model)) + + +# --------------------------------------------------------------------------- +# Profile-model plan helper (used by the Jacobian tests) +# --------------------------------------------------------------------------- + + +# +def _make_profile_plan(model_info, dynamics_params, profiles, **kwargs): + _file, model = _make_2d_profile_model( + model_info, dynamics_params, profiles, **kwargs + ) + graph = build_graph(model) + assert can_lower_jax_2d(graph) + plan = schedule_2d(graph) + return plan, model + + +# --------------------------------------------------------------------------- +# Parity: chained convolution (fixture absent from the 2D matrix) +# --------------------------------------------------------------------------- + + +# +# +class TestJaxParityChainedConvolution: + """MonoExpPosDoubleIRF is not exercised by test_evaluate_2d.""" + + # + def test_chained_convolution(self): + """Two CONVOLUTION nodes on one trace, applied in order.""" + + plan, model = _make_plan(["glp_only"], [("GLP_01_A", ["MonoExpPosDoubleIRF"])]) + _assert_parity(plan, model) + + +# --------------------------------------------------------------------------- +# wofz approximation accuracy +# --------------------------------------------------------------------------- + + +# +# +class TestWofzAccuracy: + """Weideman wofz matches scipy over the physical Voigt domain.""" + + # + def test_against_scipy(self): + from scipy.special import wofz as scipy_wofz + + from trspecfit.eval_jax import _wofz + + # z = (dx + i W/2) / (SD sqrt(2)): wide dx range, W/SD from + # Lorentzian-dominated to Gaussian-dominated + dx = np.linspace(-200.0, 200.0, 2001) + for im in [1e-3, 1e-1, 1.0, 10.0, 100.0]: + z = dx + 1j * im + got = np.asarray(_wofz(z)) + ref = scipy_wofz(z) + np.testing.assert_allclose(got.real, ref.real, rtol=1e-12, atol=1e-15) + + +# --------------------------------------------------------------------------- +# Analytic Jacobian vs central finite differences of evaluate_2d +# --------------------------------------------------------------------------- + + +# +def _assert_jacobian_matches_fd(plan, model): + """Compare make_jacobian_2d_jax against central differences. + + Finite differences run through ``evaluate_2d`` (the NumPy + reference), so this cross-validates the Jacobian against the other + backend, not against the JAX evaluator differentiating itself. + Central differences carry O(h^2) truncation error scaling with the + derivative magnitude and O(eps_machine/h) cancellation error + scaling with the *model* magnitude, so the bound uses both scales + (elementwise rtol on near-zero entries is meaningless). + """ + + jacobian = make_jacobian_2d_jax(plan) + theta = _extract_theta(plan, model) + jac = jacobian(theta) + assert jac.shape == (len(plan.time), len(plan.energy), len(theta)) + + f_scale = np.max(np.abs(evaluate_2d(plan, theta))) + for i in range(len(theta)): + h = 1e-6 * max(1.0, abs(theta[i])) + theta_plus = theta.copy() + theta_plus[i] += h + theta_minus = theta.copy() + theta_minus[i] -= h + fd = (evaluate_2d(plan, theta_plus) - evaluate_2d(plan, 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"Jacobian column {i} ({plan.opt_param_names[i]}): " + f"max |analytic - fd| = {max_err:.3e} exceeds {tol:.3e}" + ) + + +# +# +class TestJaxJacobian: + """Analytic Jacobian across the lowered feature surface.""" + + # + def test_dynamics_and_expression(self): + plan, model = _make_plan(["glp_expression"], [("GLP_01_A", ["MonoExpPos"])]) + _assert_jacobian_matches_fd(plan, model) + + # + def test_convolution(self): + plan, model = _make_plan(["glp_only"], [("GLP_01_A", ["MonoExpPosIRF"])]) + _assert_jacobian_matches_fd(plan, model) + + # + def test_subcycles(self): + plan, model = _make_plan( + ["glp_only"], + [("GLP_01_A", ["ModelNone", "MonoExpNeg"])], + frequency=10, + ) + _assert_jacobian_matches_fd(plan, model) + + # + def test_profiled_parameter(self): + plan, model = _make_profile_plan( + ["single_gauss"], + [("Gauss_01_x0", ["MonoExpPos"])], + [("Gauss_01_A", ["profile_pExpDecay"])], + ) + _assert_jacobian_matches_fd(plan, model) + + # + def test_voigt(self): + """Differentiates through the complex Weideman wofz.""" + + plan, model = _make_plan(["voigt_only"], []) + _assert_jacobian_matches_fd(plan, model) + + # + def test_sqrt_with_varying_onset(self): + """Regression: sqrt(clip) gave NaN derivatives pre-onset when + t0 varies (inf * 0 through the clip chain rule), poisoning the + whole Jacobian and failing leastsq with lmfit's non-finite + error.""" + + plan, model = _make_plan(["glp_only"], [("GLP_01_A", ["MonoSqrtVaryT0"])]) + _assert_parity(plan, model) + jacobian = make_jacobian_2d_jax(plan) + jac = jacobian(_extract_theta(plan, model)) + assert np.all(np.isfinite(jac)) + _assert_jacobian_matches_fd(plan, model) + + +# --------------------------------------------------------------------------- +# Plan-level rejection and theta contract +# --------------------------------------------------------------------------- + + +# +# +class TestJaxEvaluatorErrors: + """make_evaluator_2d_jax rejects bad theta shapes.""" + + # + def test_theta_length_mismatch_raises(self): + plan, model = _make_plan(["glp_only"], []) + evaluate_jax = make_evaluator_2d_jax(plan) + theta = _extract_theta(plan, model) + with pytest.raises(ValueError, match="theta shape"): + evaluate_jax(np.append(theta, 1.0)) + + +# --------------------------------------------------------------------------- +# End-to-end: File.fit_2d on the JAX backend with analytic Jacobian +# --------------------------------------------------------------------------- + + +# +# +class TestJaxFit2D: + """spec_fun_str='fit_model_jax' fits through the public API.""" + + # + @pytest.mark.slow + def test_fit_recovers_truth_with_analytic_jacobian(self): + """Two-stage fit (Nelder + leastsq/Dfun) recovers clean-data truth.""" + + from _utils import extract_truth_pars, make_project, simulate_clean + + from trspecfit import File + + project = make_project(name="jax_e2e", spec_fun_str="fit_model_jax") + + energy = np.linspace(83, 87, 30) + time = np.linspace(-2, 10, 24) + truth_file = File(parent_project=project, name="truth") + truth_file.energy = energy + truth_file.time = time + truth_file.dim = 2 + truth_file.load_model( + model_yaml="models/file_energy.yaml", model_info="single_glp" + ) + truth_file.add_time_dependence( + target_model="single_glp", + target_parameter="GLP_01_A", + dynamics_yaml="models/file_time.yaml", + dynamics_model=["MonoExpPos"], + ) + truth_pars = extract_truth_pars(truth_file.model_active) + clean = simulate_clean(truth_file.model_active) + + fit_file = File( + parent_project=project, + name="fit", + data=clean, + energy=energy.copy(), + time=time.copy(), + ) + fit_file.load_model( + model_yaml="models/file_energy.yaml", model_info="single_glp" + ) + fit_file.define_baseline( + time_start=0, time_stop=3, time_type="ind", show_plot=False + ) + fit_file.fit_baseline(model_name="single_glp", stages=2, try_ci=0) + fit_file.add_time_dependence( + target_model="single_glp", + target_parameter="GLP_01_A", + dynamics_yaml="models/file_time.yaml", + dynamics_model=["MonoExpPos"], + ) + fit_file.fit_2d(model_name="single_glp", stages=2, try_ci=0) + + assert fit_file.model_2d is not None # type guard + result_params = fit_file.model_2d.result[1].params + for name, true_val in truth_pars.items(): + fit_val = result_params[name].value + assert np.isclose(true_val, fit_val, rtol=1e-8, atol=1e-10), ( + f"{name}: true={true_val:.6f}, fit={fit_val:.6f}" + ) diff --git a/tests/test_graph_ir.py b/tests/test_graph_ir.py index 66508cd..125d4a5 100644 --- a/tests/test_graph_ir.py +++ b/tests/test_graph_ir.py @@ -2123,7 +2123,7 @@ def test_expressions_compiled(self): plan, _graph, _model = self._make_plan() assert plan.n_expressions == 4 - assert len(plan.expr_programs) == 4 + assert len(plan.expr_indptr) == 5 # def test_expression_target_rows_valid(self): @@ -2140,8 +2140,8 @@ def test_expression_programs_nonempty(self): plan, _graph, _model = self._make_plan() - for prog in plan.expr_programs: - assert len(prog.instructions) > 0 + for i in range(plan.n_expressions): + assert plan.expr_indptr[i + 1] > plan.expr_indptr[i] # def test_expression_a_reads_resolved(self): @@ -2180,7 +2180,9 @@ def test_expression_a_reads_resolved(self): assert expr_target_idx is not None # type guard # Check that the program contains a PARAM_REF to the resolved row - prog = plan.expr_programs[expr_target_idx] + prog = plan.expr_instructions[ + plan.expr_indptr[expr_target_idx] : plan.expr_indptr[expr_target_idx + 1] + ] resolved_row = _find_row_for_name(graph, plan, resolved_node.name) _assert_program_references_row(prog, resolved_row) @@ -2517,17 +2519,16 @@ def _find_row_for_name(graph, plan, param_name): # -def _assert_program_references_row(program, expected_row): - """Assert that an ExprProgram contains a PARAM_REF to the given row.""" +def _assert_program_references_row(instructions, expected_row): + """Assert that an RPN program contains a PARAM_REF to the given row.""" - instr = program.instructions - n_instr = len(instr) // 2 + n_instr = len(instructions) // 2 for i in range(n_instr): - kind = ExprNodeKind(instr[2 * i]) - operand = instr[2 * i + 1] + kind = ExprNodeKind(instructions[2 * i]) + operand = instructions[2 * i + 1] if kind == ExprNodeKind.PARAM_REF and operand == expected_row: return - raise AssertionError(f"ExprProgram does not reference row {expected_row}") + raise AssertionError(f"RPN program does not reference row {expected_row}") # ===================================================================== #