From 536a41b03ad2a4e5f5d6108e8b67ab82e0e023ae Mon Sep 17 00:00:00 2001 From: Johannes Mahl Date: Sat, 11 Jul 2026 17:42:56 -0700 Subject: [PATCH 01/12] add version-increment check to commit-message guidance --- CLAUDE.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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. From 3d2c2f5215aa2becee372bbdfd4dedb58a7451ee Mon Sep 17 00:00:00 2001 From: Johannes Mahl Date: Sat, 11 Jul 2026 17:56:05 -0700 Subject: [PATCH 02/12] flatten expression programs into packed CSR instruction arrays --- PLAN.md | 76 +++++++++++++-- TODO.md | 2 +- docs/design/lowered_evaluator.md | 32 ++++--- pyproject.toml | 2 +- src/trspecfit/eval_1d.py | 8 +- src/trspecfit/eval_2d.py | 39 +++++--- src/trspecfit/graph_ir.py | 157 ++++++++++++++++++++----------- tests/test_graph_ir.py | 23 ++--- 8 files changed, 232 insertions(+), 107 deletions(-) diff --git a/PLAN.md b/PLAN.md index 0d24ddf..031165f 100644 --- a/PLAN.md +++ b/PLAN.md @@ -1,11 +1,71 @@ -# Active Plan +# Active Plan: JAX backend track -No active multi-step feature in progress. +Roadmap and rationale: [docs/design/jax-planning.md](docs/design/jax-planning.md) +(scope, sequencing, open technical constraints). This file tracks execution. -- Backlog lives in [`TODO.md`](TODO.md). -- Shipped-feature design notes live in [`docs/design/`](docs/design/) (and - archived deep-dives in [`docs/design/archive/`](docs/design/archive/)). +Working branch: `jax-backend`. -Populate this file when starting the next multi-step feature; clear it on -completion per `CLAUDE.md` (archive the design into `docs/design/` if the -decisions are durable, otherwise let the changelog stand as the record). +## Guiding decisions + +- The NumPy GIR evaluator stays the reference backend throughout; JAX must + earn default status via parity + fit-level benchmarks. +- **Backend capability gating**: mirror 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`). `Voigt` + stays out of the JAX set until a JAX-native `wofz` (Humlíček-style) lands. +- **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. +- "JAX evaluator + analytic Jacobian (keep lmfit)" and "replace lmfit" are + separate milestones; the latter is optional and evidence-driven (Phase E). + +## Phase A: backend-agnostic prep (no JAX dependency; can land alone) + +- [x] Flatten `ScheduledPlan2D.expr_programs` into packed arrays + (`expr_instructions` + `expr_indptr`, CSR-style); `ExprProgram` + dataclass removed, plans are now fully array-native. +- [x] Same for `profile_expr_programs` (both 1D and 2D plans). +- [x] Keep the evaluator API unchanged: `evaluate(plan, theta) -> ndarray`; + GIR/evaluator parity tests pass (272) and full suite passes (937). +- [ ] Small helper refactors that ease the port without behavior changes + (only if needed once Phase B starts). + +## Phase B: first JAX evaluator slice (file-level 2D fits) + +- [ ] Decide packaging: optional `[jax]` extra; import guarded. +- [ ] JAX capability gate (`can_lower_jax_2d` + function/node-kind sets). +- [ ] Functional-style JAX 2D evaluator covering: static component ops, + dynamics groups, arithmetic expressions. + Excluded for now: profiles, convolution, subcycles, Voigt. +- [ ] Dispatch strategy: trace-time unrolling over scheduled ops (or + explicit JAX control flow) replacing `OP_DISPATCH` etc. +- [ ] Host-side shape checks / exceptions stay outside the jitted region. +- [ ] Parity tests vs. the NumPy GIR evaluator (not only MCP). + +## Phase C: widen JAX coverage to the existing GIR surface + +- [ ] Profile-varying parameters. +- [ ] Subcycle-aware dynamics. +- [ ] Resolved-trace convolution (kernel-matrix form; shapes are already + theta-independent). +- [ ] Voigt via a JAX-implementable `wofz` approximation, with parity + tests against the SciPy `wofz` path. +- [ ] Each widening step ships with direct NumPy-GIR parity tests. + +## Phase D: analytic Jacobian (keep lmfit) + +- [ ] Theta-level residual function. +- [ ] JAX-derived Jacobian wired into `lmfit.Minimizer` (`Dfun` plumbing + in `fitlib.py`). +- [ ] Fit-level benchmarks (not just per-call evaluator speed). + +## Phase E (optional, decide after D): JAX-native optimizer + +Only if measurements show lmfit itself is the bottleneck. Not scoped here. + +## Success criteria (from the design note) + +- JAX matches NumPy GIR to practical float tolerance on the supported subset. +- Coverage matches the NumPy GIR surface or documents temporary exclusions. +- Fit-level benchmarks show a meaningful win (JAX alone or JAX + Jacobian). +- NumPy GIR path remains available as the conservative reference. diff --git a/TODO.md b/TODO.md index b025118..7fc14f8 100644 --- a/TODO.md +++ b/TODO.md @@ -13,7 +13,7 @@ ## 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. +- [ ] **[ACTIVE] 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. - [ ] **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/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/pyproject.toml b/pyproject.toml index 2466263..613eab6 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.11.1" authors = [ {name = "Johannes Mahl", email = "infinitymonkeyatwork@gmail.com"}, ] 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/graph_ir.py b/src/trspecfit/graph_ir.py index 2540e18..c63a4c4 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 @@ -1778,7 +1777,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 +1887,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 +1899,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 +1916,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 +2352,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 +2372,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 +2506,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 +2561,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 +2592,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 +2657,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 +2670,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 +2699,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 +2717,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 +2829,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 +2983,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 +3048,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 +3062,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 +3306,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 +3358,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 +3377,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 +3527,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 +3554,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 +3578,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 +3617,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 +3634,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 +3679,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 +3690,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/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}") # ===================================================================== # From 30b152d1896cd54d2fe321b709c2dd0e404950a7 Mon Sep 17 00:00:00 2001 From: Johannes Mahl Date: Sat, 11 Jul 2026 20:21:55 -0700 Subject: [PATCH 03/12] add experimental JAX evaluator backend for the 2D lowered path --- PLAN.md | 30 ++- pyproject.toml | 9 + src/trspecfit/eval_jax.py | 383 +++++++++++++++++++++++++++++++++++++ src/trspecfit/graph_ir.py | 49 +++++ tests/test_evaluate_jax.py | 287 +++++++++++++++++++++++++++ 5 files changed, 749 insertions(+), 9 deletions(-) create mode 100644 src/trspecfit/eval_jax.py create mode 100644 tests/test_evaluate_jax.py diff --git a/PLAN.md b/PLAN.md index 031165f..5e2b4d5 100644 --- a/PLAN.md +++ b/PLAN.md @@ -32,15 +32,27 @@ Working branch: `jax-backend`. ## Phase B: first JAX evaluator slice (file-level 2D fits) -- [ ] Decide packaging: optional `[jax]` extra; import guarded. -- [ ] JAX capability gate (`can_lower_jax_2d` + function/node-kind sets). -- [ ] Functional-style JAX 2D evaluator covering: static component ops, - dynamics groups, arithmetic expressions. - Excluded for now: profiles, convolution, subcycles, Voigt. -- [ ] Dispatch strategy: trace-time unrolling over scheduled ops (or - explicit JAX control flow) replacing `OP_DISPATCH` etc. -- [ ] Host-side shape checks / exceptions stay outside the jitted region. -- [ ] Parity tests vs. the NumPy GIR evaluator (not only MCP). +- [x] 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` (float64 parity). +- [x] JAX capability gate: `can_lower_jax_2d` composes with + `can_lower_2d`, so JAX-rejected graphs land on the compiled NumPy + path by construction. Voigt / profiles / convolution / subcycles + excluded. +- [x] Functional JAX 2D evaluator (`make_evaluator_2d_jax`): static + component ops, dynamics groups, arithmetic expressions; jnp kernel + mirrors of the functions/ bodies (LinBack drops its host-side + ordering ValueError — untraceable; fit bounds must keep order). +- [x] Dispatch strategy: trace-time unrolling — one jitted XLA program + per plan, no dispatch inside the compiled path. +- [x] Host-side checks outside jit (theta shape; plan-level feature + check in `_check_plan_supported`). +- [x] Parity tests vs. the NumPy GIR evaluator (tests/test_evaluate_jax.py, + 33 tests, rtol/atol 1e-12): all 9 static ops, all 7 dynamics kinds, + multi-substep + expression-valued dynamics, gate/rejection cases. +- 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: widen JAX coverage to the existing GIR surface diff --git a/pyproject.toml b/pyproject.toml index 613eab6..bd18cf6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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_jax.py b/src/trspecfit/eval_jax.py new file mode 100644 index 0000000..2a4d161 --- /dev/null +++ b/src/trspecfit/eval_jax.py @@ -0,0 +1,383 @@ +"""Experimental JAX evaluator backend for 2D scheduled plans. + +Phase B slice of the JAX track (docs/design/jax-planning.md): static +component ops, dynamics groups, and arithmetic expressions. Profiles, +convolution, subcycle dynamics, and Voigt are gated out by +``can_lower_jax_2d``; such graphs run on the compiled NumPy evaluator +instead. + +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. + +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 ( + DynFuncKind, + ExprNodeKind, + OpKind, + ScheduledPlan2D, +) + +if TYPE_CHECKING: + import jax + import jax.numpy as jnp + from jax.scipy.special import erf as _jax_erf + + _HAVE_JAX = True +else: + try: + import jax + import jax.numpy as jnp + from jax.scipy.special import erf as _jax_erf + except ImportError: # pragma: no cover - exercised only without [jax] + jax = None + jnp = None + _jax_erf = 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) + ) + + +JAX_OP_DISPATCH: dict[int, Callable] = { + int(OpKind.GAUSS): _Gauss, + int(OpKind.GAUSS_ASYM): _GaussAsym, + int(OpKind.LORENTZ): _Lorentz, + 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): + return A * jnp.sqrt(jnp.clip(t - t0, 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, +} + + +# --------------------------------------------------------------------------- +# Expression evaluation (trace-time unrolled RPN) +# --------------------------------------------------------------------------- + + +# +def _eval_expr_rows(instructions: np.ndarray, rows: list, n_time: int): + """Evaluate one packed RPN program over per-parameter trace rows. + + Mirrors ``eval_expr_program``; PARAM_REF pushes the ``(n_time,)`` + row (a traced value), constants stay Python floats until an + operator combines them. + """ + + 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((n_time,), result, dtype=jnp.float64) + return result + + +# --------------------------------------------------------------------------- +# Evaluator factory +# --------------------------------------------------------------------------- + + +# +def _check_plan_supported(plan: ScheduledPlan2D) -> None: + """Reject plan features outside the Phase B JAX slice. + + 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] = [] + if plan.n_conv_steps > 0: + unsupported.append("convolution steps") + if plan.n_profile_samples > 0 or plan.n_profile_exprs > 0: + unsupported.append("profile-varying parameters") + 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}") + 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 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) + + 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 resolution in schedule order. + # Convolution steps (kind 2) are excluded by _check_plan_supported. + 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 + else: # 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 + ) + + # 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]) + 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 + + jitted = jax.jit(_evaluate) + + # + def evaluate(theta: np.ndarray) -> np.ndarray: + # Host-side checks stay outside the jitted region. + 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 diff --git a/src/trspecfit/graph_ir.py b/src/trspecfit/graph_ir.py index c63a4c4..c3b9ae1 100644 --- a/src/trspecfit/graph_ir.py +++ b/src/trspecfit/graph_ir.py @@ -1671,6 +1671,55 @@ def can_lower_2d(graph: GraphIR) -> bool: return True +# --------------------------------------------------------------------------- +# can_lower_jax_2d +# --------------------------------------------------------------------------- + +# JAX backend, Phase B slice (docs/design/jax-planning.md): static +# component ops, dynamics groups, and arithmetic expressions. Profiles, +# convolution, and subcycle dynamics are excluded until the JAX +# evaluator widens (Phase C); Voigt is additionally excluded because +# jax.scipy.special has no wofz. +_LOWERABLE_JAX_2D_FUNCTIONS: frozenset[str] = _LOWERABLE_2D_FUNCTIONS - frozenset( + {"Voigt"} +) + +_NON_LOWERABLE_JAX_2D_NODE_KINDS: frozenset[NodeKind] = _NON_LOWERABLE_NODE_KINDS_BASE + + +# +def can_lower_jax_2d(graph: GraphIR) -> bool: + """Check whether the experimental 2D JAX backend can compile this graph. + + Strictly narrower than :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 diff --git a/tests/test_evaluate_jax.py b/tests/test_evaluate_jax.py new file mode 100644 index 0000000..061f074 --- /dev/null +++ b/tests/test_evaluate_jax.py @@ -0,0 +1,287 @@ +"""JAX evaluator backend: gate behavior and parity vs the NumPy GIR path. + +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 make_evaluator_2d_jax # noqa: E402 +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 accepts the Phase B slice, rejects the rest.""" + + # + def test_static_model_passes(self): + _file, model = _make_2d_model(["cached_shirley_peak"], []) + assert can_lower_jax_2d(build_graph(model)) + + # + def test_dynamics_model_passes(self): + _file, model = _make_2d_model(["glp_only"], [("GLP_01_A", ["MonoExpPos"])]) + assert can_lower_jax_2d(build_graph(model)) + + # + def test_voigt_rejected_but_numpy_lowerable(self): + """Voigt fails the JAX gate yet stays on the compiled NumPy path.""" + + _file, model = _make_2d_model(["voigt_only"], []) + graph = build_graph(model) + assert can_lower_2d(graph) + assert not can_lower_jax_2d(graph) + + # + def test_convolution_rejected_but_numpy_lowerable(self): + _file, model = _make_2d_model(["glp_only"], [("GLP_01_A", ["MonoExpPosIRF"])]) + graph = build_graph(model) + assert can_lower_2d(graph) + assert not can_lower_jax_2d(graph) + + # + def test_subcycle_rejected_but_numpy_lowerable(self): + _file, model = _make_2d_model( + ["glp_only"], + [("GLP_01_A", ["ModelNone", "MonoExpNeg"])], + frequency=10, + ) + graph = build_graph(model) + assert can_lower_2d(graph) + assert not can_lower_jax_2d(graph) + + # + def test_profile_rejected_but_numpy_lowerable(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 not 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)) + + +# --------------------------------------------------------------------------- +# Parity: static component ops +# --------------------------------------------------------------------------- + + +# +# +class TestJaxParityStatic: + """Every supported OpKind, all-static parameters.""" + + _STATIC_MODELS = [ + ["gauss_only"], + ["gauss_asym_only"], + ["lorentz_only"], + ["gls_only"], + ["glp_only"], + ["ds_only"], + ["offset_only"], + ["linback_peak"], + ["shirley_peak"], + ] + + # + @pytest.mark.parametrize( + "model_info", _STATIC_MODELS, ids=[m[0] for m in _STATIC_MODELS] + ) + def test_static_parity(self, model_info): + plan, model = _make_plan(model_info, []) + _assert_parity(plan, model) + + # + def test_constant_component_caching(self): + """Constant ops are served from the cached arrays, not re-evaluated.""" + + plan, model = _make_plan(["cached_shirley_peak"], []) + _assert_parity(plan, model) + + +# --------------------------------------------------------------------------- +# Parity: dynamics groups +# --------------------------------------------------------------------------- + + +# +# +class TestJaxParityDynamics: + """Every DynFuncKind through a dynamics group on GLP_01_A.""" + + _DYN_MODELS = [ + "MonoExpPos", + "MonoExpNeg", + "MonoStep", + "MonoSin", + "MonoLin", + "MonoSinDivX", + "MonoErf", + "MonoSqrt", + ] + + # + @pytest.mark.parametrize("dyn_model", _DYN_MODELS) + def test_single_dynamics_parity(self, dyn_model): + plan, model = _make_plan(["glp_only"], [("GLP_01_A", [dyn_model])]) + _assert_parity(plan, model) + + # + def test_multi_substep_group(self): + """Bi-exponential: two substeps summed into one dynamics group. + + The second expFun's t0 is an expression (``expFun_01_t0``), so + this also covers expression-valued dynamics parameters. + """ + + plan, model = _make_plan(["glp_only"], [("GLP_01_A", ["BiExpSharedT0"])]) + _assert_parity(plan, model) + + # + def test_two_dynamic_parameters(self): + """Independent dynamics groups on two parameters.""" + + plan, model = _make_plan( + ["glp_only"], + [("GLP_01_A", ["MonoExpPos"]), ("GLP_01_x0", ["MonoSin"])], + ) + _assert_parity(plan, model) + + +# --------------------------------------------------------------------------- +# Parity: arithmetic expressions +# --------------------------------------------------------------------------- + + +# +# +class TestJaxParityExpressions: + """Compiled RPN expressions, static and interleaved with dynamics.""" + + # + def test_static_expression(self): + plan, model = _make_plan(["glp_expression"], []) + _assert_parity(plan, model) + + # + def test_expression_reads_resolved_trace(self): + """Expression referencing a time-dependent (resolved) parameter.""" + + plan, model = _make_plan( + ["glp_expression"], + [("GLP_01_A", ["MonoExpPos"])], + ) + _assert_parity(plan, model) + + +# --------------------------------------------------------------------------- +# Plan-level rejection and theta contract +# --------------------------------------------------------------------------- + + +# +# +class TestJaxEvaluatorErrors: + """make_evaluator_2d_jax raises on unsupported plans and bad theta.""" + + # + def test_voigt_plan_raises(self): + _file, model = _make_2d_model(["voigt_only"], []) + plan = schedule_2d(build_graph(model)) + with pytest.raises(ValueError, match="VOIGT"): + make_evaluator_2d_jax(plan) + + # + def test_convolution_plan_raises(self): + _file, model = _make_2d_model(["glp_only"], [("GLP_01_A", ["MonoExpPosIRF"])]) + plan = schedule_2d(build_graph(model)) + with pytest.raises(ValueError, match="convolution"): + make_evaluator_2d_jax(plan) + + # + def test_profile_plan_raises(self): + _file, model = _make_2d_profile_model( + ["single_gauss"], + [("Gauss_01_x0", ["MonoExpPos"])], + [("Gauss_01_A", ["profile_pExpDecay"])], + ) + plan = schedule_2d(build_graph(model)) + with pytest.raises(ValueError, match="profile"): + make_evaluator_2d_jax(plan) + + # + 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)) From 1c7809a4d71404b5d7685032f6614db4fc9bcf71 Mon Sep 17 00:00:00 2001 From: Johannes Mahl Date: Sat, 11 Jul 2026 20:51:49 -0700 Subject: [PATCH 04/12] widen the JAX evaluator to the full lowered 2D surface --- PLAN.md | 32 +++- src/trspecfit/eval_jax.py | 310 ++++++++++++++++++++++++++++++++++--- src/trspecfit/graph_ir.py | 22 ++- tests/test_evaluate_jax.py | 230 ++++++++++++++++++++------- 4 files changed, 496 insertions(+), 98 deletions(-) diff --git a/PLAN.md b/PLAN.md index 5e2b4d5..76385c0 100644 --- a/PLAN.md +++ b/PLAN.md @@ -56,13 +56,31 @@ Working branch: `jax-backend`. ## Phase C: widen JAX coverage to the existing GIR surface -- [ ] Profile-varying parameters. -- [ ] Subcycle-aware dynamics. -- [ ] Resolved-trace convolution (kernel-matrix form; shapes are already - theta-independent). -- [ ] Voigt via a JAX-implementable `wofz` approximation, with parity - tests against the SciPy `wofz` path. -- [ ] Each widening step ships with direct NumPy-GIR parity tests. +- [x] 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; under + XLA the fused form is simpler and still 2.6x faster at n_aux=50). +- [x] 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. +- [x] 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. +- [x] 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. +- [x] Parity tests vs NumPy GIR for every widening step (44 tests + total): all conv kernels + chained double IRF, 2- and 3-subcycle + models with cross-subcycle expressions, profiled amplitude / + position / Shirley / mixed profile-dynamics expressions, Voigt. +- `can_lower_jax_2d` now 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 (keep lmfit) diff --git a/src/trspecfit/eval_jax.py b/src/trspecfit/eval_jax.py index 2a4d161..910afc3 100644 --- a/src/trspecfit/eval_jax.py +++ b/src/trspecfit/eval_jax.py @@ -1,16 +1,22 @@ """Experimental JAX evaluator backend for 2D scheduled plans. -Phase B slice of the JAX track (docs/design/jax-planning.md): static -component ops, dynamics groups, and arithmetic expressions. Profiles, -convolution, subcycle dynamics, and Voigt are gated out by -``can_lower_jax_2d``; such graphs run on the compiled NumPy evaluator -instead. +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. @@ -24,9 +30,12 @@ import numpy as np from trspecfit.graph_ir import ( + ConvKernelKind, DynFuncKind, ExprNodeKind, OpKind, + ParamSourceKind, + ProfileFuncKind, ScheduledPlan2D, ) @@ -34,6 +43,7 @@ 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: @@ -41,10 +51,12 @@ 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. @@ -129,10 +141,54 @@ def _DS(x, A, x0, F, alpha): ) +# +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, @@ -193,18 +249,127 @@ def _sqrtFun(t, A, t0): } +# --------------------------------------------------------------------------- +# 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, n_time: int): +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 ``(n_time,)`` + Mirrors ``eval_expr_program``; PARAM_REF pushes the *shape*-shaped row (a traced value), constants stay Python floats until an - operator combines them. + operator combines them. ``rows`` entries are ``(n_time,)`` for + plan expressions and ``(n_time, n_aux)`` for per-sample profile + expressions. """ stack: list = [] @@ -239,7 +404,7 @@ def _eval_expr_rows(instructions: np.ndarray, rows: list, n_time: int): assert len(stack) == 1 result = stack[0] if isinstance(result, float): # constant-only program - return jnp.full((n_time,), result, dtype=jnp.float64) + return jnp.full(shape, result, dtype=jnp.float64) return result @@ -250,20 +415,21 @@ def _eval_expr_rows(instructions: np.ndarray, rows: list, n_time: int): # def _check_plan_supported(plan: ScheduledPlan2D) -> None: - """Reject plan features outside the Phase B JAX slice. + """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] = [] - if plan.n_conv_steps > 0: - unsupported.append("convolution steps") - if plan.n_profile_samples > 0 or plan.n_profile_exprs > 0: - unsupported.append("profile-varying parameters") 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: " @@ -313,8 +479,8 @@ def _evaluate(theta): for k in range(n_opt): rows[int(plan.opt_indices[k])] = jnp.broadcast_to(theta[k], (n_time,)) - # Interleaved dynamics/expression resolution in schedule order. - # Convolution steps (kind 2) are excluded by _check_plan_supported. + # 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]) @@ -335,12 +501,76 @@ def _evaluate(theta): jnp.asarray(plan.dyn_sub_time_axes[s]), *dyn_params ) * jnp.asarray(plan.dyn_sub_masks[s]) rows[target] = acc - else: # kind == 1: expression + 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 + 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. @@ -353,14 +583,44 @@ def _evaluate(theta): 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]) - 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) + 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: - component = func(energy, *params) + 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 diff --git a/src/trspecfit/graph_ir.py b/src/trspecfit/graph_ir.py index c3b9ae1..303ef09 100644 --- a/src/trspecfit/graph_ir.py +++ b/src/trspecfit/graph_ir.py @@ -1675,25 +1675,23 @@ def can_lower_2d(graph: GraphIR) -> bool: # can_lower_jax_2d # --------------------------------------------------------------------------- -# JAX backend, Phase B slice (docs/design/jax-planning.md): static -# component ops, dynamics groups, and arithmetic expressions. Profiles, -# convolution, and subcycle dynamics are excluded until the JAX -# evaluator widens (Phase C); Voigt is additionally excluded because -# jax.scipy.special has no wofz. -_LOWERABLE_JAX_2D_FUNCTIONS: frozenset[str] = _LOWERABLE_2D_FUNCTIONS - frozenset( - {"Voigt"} -) +# 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_NODE_KINDS_BASE +_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. - Strictly narrower than :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). + 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 ---------- diff --git a/tests/test_evaluate_jax.py b/tests/test_evaluate_jax.py index 061f074..75bb87b 100644 --- a/tests/test_evaluate_jax.py +++ b/tests/test_evaluate_jax.py @@ -74,47 +74,34 @@ def _assert_parity(plan, model, *, rtol=1e-12, atol=1e-12): # # class TestJaxGate: - """can_lower_jax_2d accepts the Phase B slice, rejects the rest.""" - - # - def test_static_model_passes(self): - _file, model = _make_2d_model(["cached_shirley_peak"], []) - assert can_lower_jax_2d(build_graph(model)) - - # - def test_dynamics_model_passes(self): - _file, model = _make_2d_model(["glp_only"], [("GLP_01_A", ["MonoExpPos"])]) - assert can_lower_jax_2d(build_graph(model)) - - # - def test_voigt_rejected_but_numpy_lowerable(self): - """Voigt fails the JAX gate yet stays on the compiled NumPy path.""" - - _file, model = _make_2d_model(["voigt_only"], []) - graph = build_graph(model) - assert can_lower_2d(graph) - assert not can_lower_jax_2d(graph) - - # - def test_convolution_rejected_but_numpy_lowerable(self): - _file, model = _make_2d_model(["glp_only"], [("GLP_01_A", ["MonoExpPosIRF"])]) - graph = build_graph(model) - assert can_lower_2d(graph) - assert not can_lower_jax_2d(graph) - - # - def test_subcycle_rejected_but_numpy_lowerable(self): - _file, model = _make_2d_model( + """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, - ) + {"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 not can_lower_jax_2d(graph) + assert can_lower_jax_2d(graph) # - def test_profile_rejected_but_numpy_lowerable(self): + def test_profile_model_passes(self): _file, model = _make_2d_profile_model( ["single_gauss"], [("Gauss_01_x0", ["MonoExpPos"])], @@ -122,7 +109,7 @@ def test_profile_rejected_but_numpy_lowerable(self): ) graph = build_graph(model) assert can_lower_2d(graph) - assert not can_lower_jax_2d(graph) + assert can_lower_jax_2d(graph) # def test_1d_graph_rejected(self): @@ -144,6 +131,7 @@ class TestJaxParityStatic: ["gauss_only"], ["gauss_asym_only"], ["lorentz_only"], + ["voigt_only"], ["gls_only"], ["glp_only"], ["ds_only"], @@ -244,39 +232,173 @@ def test_expression_reads_resolved_trace(self): # --------------------------------------------------------------------------- -# Plan-level rejection and theta contract +# Parity: subcycle dynamics # --------------------------------------------------------------------------- # # -class TestJaxEvaluatorErrors: - """make_evaluator_2d_jax raises on unsupported plans and bad theta.""" +class TestJaxParitySubcycles: + """Subcycle time axes / masks (compiled schedule data) in parity.""" # - def test_voigt_plan_raises(self): - _file, model = _make_2d_model(["voigt_only"], []) - plan = schedule_2d(build_graph(model)) - with pytest.raises(ValueError, match="VOIGT"): - make_evaluator_2d_jax(plan) + def test_two_subcycles(self): + plan, model = _make_plan( + ["glp_only"], + [("GLP_01_A", ["ModelNone", "MonoExpNeg"])], + frequency=10, + ) + _assert_parity(plan, model) # - def test_convolution_plan_raises(self): - _file, model = _make_2d_model(["glp_only"], [("GLP_01_A", ["MonoExpPosIRF"])]) - plan = schedule_2d(build_graph(model)) - with pytest.raises(ValueError, match="convolution"): - make_evaluator_2d_jax(plan) + def test_three_subcycles_with_expressions(self): + """Cross-subcycle expression refs (MonoExpPosExpr <- MonoExpNeg).""" + + plan, model = _make_plan( + ["glp_only"], + [("GLP_01_A", ["ModelNone", "MonoExpNeg", "MonoExpPosExpr"])], + frequency=10, + ) + _assert_parity(plan, model) + + +# --------------------------------------------------------------------------- +# Parity: profile-varying parameters +# --------------------------------------------------------------------------- + + +# +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 + + +# +# +class TestJaxParityProfiles: + """Profiled amplitude (linear) and position (nonlinear), exprs, Shirley.""" # - def test_profile_plan_raises(self): - _file, model = _make_2d_profile_model( + def test_profiled_amplitude_with_dynamics(self): + plan, model = _make_profile_plan( ["single_gauss"], [("Gauss_01_x0", ["MonoExpPos"])], [("Gauss_01_A", ["profile_pExpDecay"])], ) - plan = schedule_2d(build_graph(model)) - with pytest.raises(ValueError, match="profile"): - make_evaluator_2d_jax(plan) + _assert_parity(plan, model) + + # + def test_two_profiles_same_component(self): + """Profiled x0 (nonlinear) + profiled A on one component.""" + + plan, model = _make_profile_plan( + ["single_gauss"], + [("Gauss_01_SD", ["MonoExpPos"])], + [ + ("Gauss_01_x0", ["roundtrip_pLinear_x0"]), + ("Gauss_01_A", ["roundtrip_pExpDecay_A"]), + ], + ) + _assert_parity(plan, model) + + # + def test_profile_expr_with_time_dep_ref(self): + """Per-sample expression reading a profiled and a resolved param.""" + + plan, model = _make_profile_plan( + ["two_glp_mixed_profile_dynamics"], + [("GLP_01_x0", ["MonoExpPos"])], + [("GLP_01_A", ["profile_pLinear"])], + ) + _assert_parity(plan, model) + + # + def test_profiled_shirley(self): + """Profiled spectrum-fed op (Shirley background).""" + + plan, model = _make_profile_plan( + ["shirley_peak"], + [("GLP_01_A", ["MonoExpPos"])], + [("Shirley_pShirley", ["profile_pLinear"])], + model_yaml="models/eval_2d_energy.yaml", + energy=np.linspace(80, 90, 101), + ) + model.lmfit_pars["Shirley_pShirley_pLinear_01_m"].value = 1.0e-5 + _assert_parity(plan, model) + + +# --------------------------------------------------------------------------- +# Parity: resolved-trace convolution (kernel-matrix) +# --------------------------------------------------------------------------- + + +# +# +class TestJaxParityConvolution: + """Every lowered conv kernel, plus a chained double IRF.""" + + _IRF_MODELS = [ + "MonoExpPosIRF", + "MonoExpPosExpSymIRF", + "MonoExpPosExpDecayIRF", + "MonoExpPosExpRiseIRF", + "MonoExpPosBoxIRF", + ] + + # + @pytest.mark.parametrize("dyn_model", _IRF_MODELS) + def test_kernel_parity(self, dyn_model): + plan, model = _make_plan(["glp_only"], [("GLP_01_A", [dyn_model])]) + _assert_parity(plan, model) + + # + 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) + + +# --------------------------------------------------------------------------- +# Plan-level rejection and theta contract +# --------------------------------------------------------------------------- + + +# +# +class TestJaxEvaluatorErrors: + """make_evaluator_2d_jax rejects bad theta shapes.""" # def test_theta_length_mismatch_raises(self): From d0d6349eac55f02bf761c7d2660e36de59e89ced Mon Sep 17 00:00:00 2001 From: Johannes Mahl Date: Sat, 11 Jul 2026 21:32:24 -0700 Subject: [PATCH 05/12] add analytic JAX Jacobian and fit_model_jax backend to the fit path --- PLAN.md | 33 +++++++- src/trspecfit/eval_jax.py | 113 ++++++++++++++++++++------- src/trspecfit/fitlib.py | 87 ++++++++++++++++++++- src/trspecfit/spectra.py | 44 +++++++++++ src/trspecfit/trspecfit.py | 37 ++++++++- tests/test_evaluate_jax.py | 155 ++++++++++++++++++++++++++++++++++++- 6 files changed, 428 insertions(+), 41 deletions(-) diff --git a/PLAN.md b/PLAN.md index 76385c0..fb35594 100644 --- a/PLAN.md +++ b/PLAN.md @@ -84,10 +84,35 @@ Working branch: `jax-backend`. ## Phase D: analytic Jacobian (keep lmfit) -- [ ] Theta-level residual function. -- [ ] JAX-derived Jacobian wired into `lmfit.Minimizer` (`Dfun` plumbing - in `fitlib.py`). -- [ ] Fit-level benchmarks (not just per-call evaluator speed). +- [x] `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 documented: + boxCONV width derivative is 0 a.e.; step-like `where` switches + give subgradient zeros at the exact switching point. +- [x] `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. +- [x] `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). +- [x] Jacobian validated against central finite differences of + `evaluate_2d` (the NumPy backend — cross-backend check) for + dynamics/expressions, convolution, subcycles, profiles, Voigt; + plus an internal-space check against lmfit's own FD at a real + fit's start point (exact match), and a slow e2e fit test. +- [x] 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. +- Note: from aggressively perturbed starts, plain 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. ## Phase E (optional, decide after D): JAX-native optimizer diff --git a/src/trspecfit/eval_jax.py b/src/trspecfit/eval_jax.py index 910afc3..e39b529 100644 --- a/src/trspecfit/eval_jax.py +++ b/src/trspecfit/eval_jax.py @@ -440,35 +440,13 @@ def _check_plan_supported(plan: ScheduledPlan2D) -> None: # -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. +def _build_evaluate_2d(plan: ScheduledPlan2D) -> Callable: + """Build the pure traced function ``theta -> (n_time, n_energy)``. - Raises - ------ - ImportError - If jax is not installed. - ValueError - If the plan uses features outside the JAX slice. + Shared by ``make_evaluator_2d_jax`` (jit) and + ``make_jacobian_2d_jax`` (jit of jacfwd). """ - _require_jax() - _check_plan_supported(plan) - n_opt = len(plan.opt_indices) n_time = plan.n_time @@ -627,11 +605,15 @@ def _evaluate(theta): return result - jitted = jax.jit(_evaluate) + 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: - # Host-side checks stay outside the jitted region. theta_arr = np.asarray(theta, dtype=np.float64) if theta_arr.shape != (n_opt,): raise ValueError( @@ -641,3 +623,78 @@ def evaluate(theta: np.ndarray) -> np.ndarray: 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/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/test_evaluate_jax.py b/tests/test_evaluate_jax.py index 75bb87b..dde70e3 100644 --- a/tests/test_evaluate_jax.py +++ b/tests/test_evaluate_jax.py @@ -20,7 +20,10 @@ ) from trspecfit.eval_2d import evaluate_2d # noqa: E402 -from trspecfit.eval_jax import make_evaluator_2d_jax # 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, @@ -390,6 +393,86 @@ def test_against_scipy(self): 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) + + # --------------------------------------------------------------------------- # Plan-level rejection and theta contract # --------------------------------------------------------------------------- @@ -407,3 +490,73 @@ def test_theta_length_mismatch_raises(self): 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}" + ) From 29fc5b75796c254c87def738bb200a5e9d722da6 Mon Sep 17 00:00:00 2001 From: Johannes Mahl Date: Sat, 11 Jul 2026 22:38:14 -0700 Subject: [PATCH 06/12] assert JAX parity across the 2D evaluator matrix; run JAX tests in CI --- .github/workflows/ci.yaml | 6 +- tests/test_evaluate_2d.py | 37 ++++++ tests/test_evaluate_jax.py | 232 +++---------------------------------- 3 files changed, 56 insertions(+), 219 deletions(-) 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/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 index dde70e3..f656c8c 100644 --- a/tests/test_evaluate_jax.py +++ b/tests/test_evaluate_jax.py @@ -1,4 +1,13 @@ -"""JAX evaluator backend: gate behavior and parity vs the NumPy GIR path. +"""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 @@ -121,152 +130,7 @@ def test_1d_graph_rejected(self): # --------------------------------------------------------------------------- -# Parity: static component ops -# --------------------------------------------------------------------------- - - -# -# -class TestJaxParityStatic: - """Every supported OpKind, all-static parameters.""" - - _STATIC_MODELS = [ - ["gauss_only"], - ["gauss_asym_only"], - ["lorentz_only"], - ["voigt_only"], - ["gls_only"], - ["glp_only"], - ["ds_only"], - ["offset_only"], - ["linback_peak"], - ["shirley_peak"], - ] - - # - @pytest.mark.parametrize( - "model_info", _STATIC_MODELS, ids=[m[0] for m in _STATIC_MODELS] - ) - def test_static_parity(self, model_info): - plan, model = _make_plan(model_info, []) - _assert_parity(plan, model) - - # - def test_constant_component_caching(self): - """Constant ops are served from the cached arrays, not re-evaluated.""" - - plan, model = _make_plan(["cached_shirley_peak"], []) - _assert_parity(plan, model) - - -# --------------------------------------------------------------------------- -# Parity: dynamics groups -# --------------------------------------------------------------------------- - - -# -# -class TestJaxParityDynamics: - """Every DynFuncKind through a dynamics group on GLP_01_A.""" - - _DYN_MODELS = [ - "MonoExpPos", - "MonoExpNeg", - "MonoStep", - "MonoSin", - "MonoLin", - "MonoSinDivX", - "MonoErf", - "MonoSqrt", - ] - - # - @pytest.mark.parametrize("dyn_model", _DYN_MODELS) - def test_single_dynamics_parity(self, dyn_model): - plan, model = _make_plan(["glp_only"], [("GLP_01_A", [dyn_model])]) - _assert_parity(plan, model) - - # - def test_multi_substep_group(self): - """Bi-exponential: two substeps summed into one dynamics group. - - The second expFun's t0 is an expression (``expFun_01_t0``), so - this also covers expression-valued dynamics parameters. - """ - - plan, model = _make_plan(["glp_only"], [("GLP_01_A", ["BiExpSharedT0"])]) - _assert_parity(plan, model) - - # - def test_two_dynamic_parameters(self): - """Independent dynamics groups on two parameters.""" - - plan, model = _make_plan( - ["glp_only"], - [("GLP_01_A", ["MonoExpPos"]), ("GLP_01_x0", ["MonoSin"])], - ) - _assert_parity(plan, model) - - -# --------------------------------------------------------------------------- -# Parity: arithmetic expressions -# --------------------------------------------------------------------------- - - -# -# -class TestJaxParityExpressions: - """Compiled RPN expressions, static and interleaved with dynamics.""" - - # - def test_static_expression(self): - plan, model = _make_plan(["glp_expression"], []) - _assert_parity(plan, model) - - # - def test_expression_reads_resolved_trace(self): - """Expression referencing a time-dependent (resolved) parameter.""" - - plan, model = _make_plan( - ["glp_expression"], - [("GLP_01_A", ["MonoExpPos"])], - ) - _assert_parity(plan, model) - - -# --------------------------------------------------------------------------- -# Parity: subcycle dynamics -# --------------------------------------------------------------------------- - - -# -# -class TestJaxParitySubcycles: - """Subcycle time axes / masks (compiled schedule data) in parity.""" - - # - def test_two_subcycles(self): - plan, model = _make_plan( - ["glp_only"], - [("GLP_01_A", ["ModelNone", "MonoExpNeg"])], - frequency=10, - ) - _assert_parity(plan, model) - - # - def test_three_subcycles_with_expressions(self): - """Cross-subcycle expression refs (MonoExpPosExpr <- MonoExpNeg).""" - - plan, model = _make_plan( - ["glp_only"], - [("GLP_01_A", ["ModelNone", "MonoExpNeg", "MonoExpPosExpr"])], - frequency=10, - ) - _assert_parity(plan, model) - - -# --------------------------------------------------------------------------- -# Parity: profile-varying parameters +# Profile-model plan helper (used by the Jacobian tests) # --------------------------------------------------------------------------- @@ -281,83 +145,15 @@ def _make_profile_plan(model_info, dynamics_params, profiles, **kwargs): return plan, model -# -# -class TestJaxParityProfiles: - """Profiled amplitude (linear) and position (nonlinear), exprs, Shirley.""" - - # - def test_profiled_amplitude_with_dynamics(self): - plan, model = _make_profile_plan( - ["single_gauss"], - [("Gauss_01_x0", ["MonoExpPos"])], - [("Gauss_01_A", ["profile_pExpDecay"])], - ) - _assert_parity(plan, model) - - # - def test_two_profiles_same_component(self): - """Profiled x0 (nonlinear) + profiled A on one component.""" - - plan, model = _make_profile_plan( - ["single_gauss"], - [("Gauss_01_SD", ["MonoExpPos"])], - [ - ("Gauss_01_x0", ["roundtrip_pLinear_x0"]), - ("Gauss_01_A", ["roundtrip_pExpDecay_A"]), - ], - ) - _assert_parity(plan, model) - - # - def test_profile_expr_with_time_dep_ref(self): - """Per-sample expression reading a profiled and a resolved param.""" - - plan, model = _make_profile_plan( - ["two_glp_mixed_profile_dynamics"], - [("GLP_01_x0", ["MonoExpPos"])], - [("GLP_01_A", ["profile_pLinear"])], - ) - _assert_parity(plan, model) - - # - def test_profiled_shirley(self): - """Profiled spectrum-fed op (Shirley background).""" - - plan, model = _make_profile_plan( - ["shirley_peak"], - [("GLP_01_A", ["MonoExpPos"])], - [("Shirley_pShirley", ["profile_pLinear"])], - model_yaml="models/eval_2d_energy.yaml", - energy=np.linspace(80, 90, 101), - ) - model.lmfit_pars["Shirley_pShirley_pLinear_01_m"].value = 1.0e-5 - _assert_parity(plan, model) - - # --------------------------------------------------------------------------- -# Parity: resolved-trace convolution (kernel-matrix) +# Parity: chained convolution (fixture absent from the 2D matrix) # --------------------------------------------------------------------------- # # -class TestJaxParityConvolution: - """Every lowered conv kernel, plus a chained double IRF.""" - - _IRF_MODELS = [ - "MonoExpPosIRF", - "MonoExpPosExpSymIRF", - "MonoExpPosExpDecayIRF", - "MonoExpPosExpRiseIRF", - "MonoExpPosBoxIRF", - ] - - # - @pytest.mark.parametrize("dyn_model", _IRF_MODELS) - def test_kernel_parity(self, dyn_model): - plan, model = _make_plan(["glp_only"], [("GLP_01_A", [dyn_model])]) - _assert_parity(plan, model) +class TestJaxParityChainedConvolution: + """MonoExpPosDoubleIRF is not exercised by test_evaluate_2d.""" # def test_chained_convolution(self): From 813491b89e5772edd673200e6895eca685c42f74 Mon Sep 17 00:00:00 2001 From: Johannes Mahl Date: Sat, 11 Jul 2026 22:38:35 -0700 Subject: [PATCH 07/12] bump to 0.12.0; document the JAX backend in module map and API docs --- PLAN.md | 13 ++++++++++++- docs/api/eval_jax.rst | 7 +++++++ docs/api/index.rst | 1 + docs/design/jax-planning.md | 6 ++++++ docs/design/repo_architecture.md | 12 ++++++++++++ pyproject.toml | 2 +- 6 files changed, 39 insertions(+), 2 deletions(-) create mode 100644 docs/api/eval_jax.rst diff --git a/PLAN.md b/PLAN.md index fb35594..a09b85d 100644 --- a/PLAN.md +++ b/PLAN.md @@ -78,6 +78,14 @@ Working branch: `jax-backend`. - `can_lower_jax_2d` now spans the full `can_lower_2d` surface (the separate sets remain so future NumPy widening doesn't silently imply JAX support). +- Test coverage (2026-07-11): 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. - 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. @@ -116,7 +124,10 @@ Working branch: `jax-backend`. ## Phase E (optional, decide after D): JAX-native optimizer -Only if measurements show lmfit itself is the bottleneck. Not scoped here. +Deferred (2026-07-11). The Phase D benchmarks show the evaluator and +Jacobian — not lmfit — dominate fit wall time at current model sizes, +so replacing the optimizer stack is not justified by evidence. Revisit +only if profiling a real workload shows lmfit overhead dominating. ## Success criteria (from the design note) 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/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/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/pyproject.toml b/pyproject.toml index bd18cf6..bdd490f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "trspecfit" -version = "0.11.1" +version = "0.12.0" authors = [ {name = "Johannes Mahl", email = "infinitymonkeyatwork@gmail.com"}, ] From 2b5300c1f3fa43f52d22dcffe672a5390c4032d2 Mon Sep 17 00:00:00 2001 From: Johannes Mahl Date: Sat, 11 Jul 2026 22:45:50 -0700 Subject: [PATCH 08/12] add the JAX backend to the benchmark harness --- .claude/skills/benchmark/SKILL.md | 2 +- .claude/skills/benchmark/benchmark_gir.py | 80 ++++++++++++++++++----- docs/ai/benchmark.md | 17 +++-- 3 files changed, 78 insertions(+), 21 deletions(-) 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/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 From 2e24c9a507e7edb312d901251d5ba9a3a70ee333 Mon Sep 17 00:00:00 2001 From: Johannes Mahl Date: Sat, 11 Jul 2026 23:24:44 -0700 Subject: [PATCH 09/12] add design notes: interactive UI and project-level shared fits --- TODO.md | 2 +- docs/design/project-level-fits.md | 88 +++++++++++++++++++++++++++++++ docs/design/ui.md | 78 +++++++++++++++++++++++++++ 3 files changed, 167 insertions(+), 1 deletion(-) create mode 100644 docs/design/project-level-fits.md create mode 100644 docs/design/ui.md diff --git a/TODO.md b/TODO.md index 7fc14f8..f8931a9 100644 --- a/TODO.md +++ b/TODO.md @@ -12,7 +12,7 @@ ## 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. +- [ ] **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). - [ ] **[ACTIVE] 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. - [ ] **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: 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/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. From 03851a781ad154f6a6659f6d3976057159f7e806 Mon Sep 17 00:00:00 2001 From: Johannes Mahl Date: Sat, 11 Jul 2026 23:27:23 -0700 Subject: [PATCH 10/12] archive the JAX-track execution record; clear PLAN and untag TODO --- PLAN.md | 142 ++--------------------------- TODO.md | 6 +- docs/design/archive/jax-backend.md | 129 ++++++++++++++++++++++++++ 3 files changed, 142 insertions(+), 135 deletions(-) create mode 100644 docs/design/archive/jax-backend.md diff --git a/PLAN.md b/PLAN.md index a09b85d..0d24ddf 100644 --- a/PLAN.md +++ b/PLAN.md @@ -1,137 +1,11 @@ -# Active Plan: JAX backend track +# Active Plan -Roadmap and rationale: [docs/design/jax-planning.md](docs/design/jax-planning.md) -(scope, sequencing, open technical constraints). This file tracks execution. +No active multi-step feature in progress. -Working branch: `jax-backend`. +- Backlog lives in [`TODO.md`](TODO.md). +- Shipped-feature design notes live in [`docs/design/`](docs/design/) (and + archived deep-dives in [`docs/design/archive/`](docs/design/archive/)). -## Guiding decisions - -- The NumPy GIR evaluator stays the reference backend throughout; JAX must - earn default status via parity + fit-level benchmarks. -- **Backend capability gating**: mirror 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`). `Voigt` - stays out of the JAX set until a JAX-native `wofz` (Humlíček-style) lands. -- **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. -- "JAX evaluator + analytic Jacobian (keep lmfit)" and "replace lmfit" are - separate milestones; the latter is optional and evidence-driven (Phase E). - -## Phase A: backend-agnostic prep (no JAX dependency; can land alone) - -- [x] Flatten `ScheduledPlan2D.expr_programs` into packed arrays - (`expr_instructions` + `expr_indptr`, CSR-style); `ExprProgram` - dataclass removed, plans are now fully array-native. -- [x] Same for `profile_expr_programs` (both 1D and 2D plans). -- [x] Keep the evaluator API unchanged: `evaluate(plan, theta) -> ndarray`; - GIR/evaluator parity tests pass (272) and full suite passes (937). -- [ ] Small helper refactors that ease the port without behavior changes - (only if needed once Phase B starts). - -## Phase B: first JAX evaluator slice (file-level 2D fits) - -- [x] 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` (float64 parity). -- [x] JAX capability gate: `can_lower_jax_2d` composes with - `can_lower_2d`, so JAX-rejected graphs land on the compiled NumPy - path by construction. Voigt / profiles / convolution / subcycles - excluded. -- [x] Functional JAX 2D evaluator (`make_evaluator_2d_jax`): static - component ops, dynamics groups, arithmetic expressions; jnp kernel - mirrors of the functions/ bodies (LinBack drops its host-side - ordering ValueError — untraceable; fit bounds must keep order). -- [x] Dispatch strategy: trace-time unrolling — one jitted XLA program - per plan, no dispatch inside the compiled path. -- [x] Host-side checks outside jit (theta shape; plan-level feature - check in `_check_plan_supported`). -- [x] Parity tests vs. the NumPy GIR evaluator (tests/test_evaluate_jax.py, - 33 tests, rtol/atol 1e-12): all 9 static ops, all 7 dynamics kinds, - multi-substep + expression-valued dynamics, gate/rejection cases. -- 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: widen JAX coverage to the existing GIR surface - -- [x] 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; under - XLA the fused form is simpler and still 2.6x faster at n_aux=50). -- [x] 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. -- [x] 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. -- [x] 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. -- [x] Parity tests vs NumPy GIR for every widening step (44 tests - total): all conv kernels + chained double IRF, 2- and 3-subcycle - models with cross-subcycle expressions, profiled amplitude / - position / Shirley / mixed profile-dynamics expressions, Voigt. -- `can_lower_jax_2d` now spans the full `can_lower_2d` surface (the - separate sets remain so future NumPy widening doesn't silently imply - JAX support). -- Test coverage (2026-07-11): 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. -- 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 (keep lmfit) - -- [x] `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 documented: - boxCONV width derivative is 0 a.e.; step-like `where` switches - give subgradient zeros at the exact switching point. -- [x] `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. -- [x] `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). -- [x] Jacobian validated against central finite differences of - `evaluate_2d` (the NumPy backend — cross-backend check) for - dynamics/expressions, convolution, subcycles, profiles, Voigt; - plus an internal-space check against lmfit's own FD at a real - fit's start point (exact match), and a slow e2e fit test. -- [x] 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. -- Note: from aggressively perturbed starts, plain 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. - -## Phase E (optional, decide after D): JAX-native optimizer - -Deferred (2026-07-11). The Phase D benchmarks show the evaluator and -Jacobian — not lmfit — dominate fit wall time at current model sizes, -so replacing the optimizer stack is not justified by evidence. Revisit -only if profiling a real workload shows lmfit overhead dominating. - -## Success criteria (from the design note) - -- JAX matches NumPy GIR to practical float tolerance on the supported subset. -- Coverage matches the NumPy GIR surface or documents temporary exclusions. -- Fit-level benchmarks show a meaningful win (JAX alone or JAX + Jacobian). -- NumPy GIR path remains available as the conservative reference. +Populate this file when starting the next multi-step feature; clear it on +completion per `CLAUDE.md` (archive the design into `docs/design/` if the +decisions are durable, otherwise let the changelog stand as the record). diff --git a/TODO.md b/TODO.md index f8931a9..ad2aabc 100644 --- a/TODO.md +++ b/TODO.md @@ -13,7 +13,11 @@ ## Performance & architecture - [ ] **Project-level fit backend**: `Project.fit_2d()` already supports `Project`/`File`/`Static` vary levels, but it currently evaluates through `fit_project_mcp()` and `Model.create_value_2d()` rather than any compiled path. Design direction settled (2026-07-11): lower per-file plans and evaluate them through a fused JAX residual with a joint analytic Jacobian — see [docs/design/project-level-fits.md](docs/design/project-level-fits.md) for the rationale, implementation sketch, and open questions (compile-time scaling, mixed lowerability, vmap batching for homogeneous series). -- [ ] **[ACTIVE] 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. +- [ ] **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/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)). From 1e196c9810255ff22da95073abb6f11ebe0b4cc7 Mon Sep 17 00:00:00 2001 From: Johannes Mahl Date: Sat, 11 Jul 2026 23:31:20 -0700 Subject: [PATCH 11/12] changelog: document 0.12.0 (JAX backend) and retitle 0.11.0 entry --- CHANGELOG.md | 17 ++++++++++++++++- 1 file changed, 16 insertions(+), 1 deletion(-) 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). From c20c40499eff44dddb2dd1422783b9bf9a912539 Mon Sep 17 00:00:00 2001 From: Johannes Mahl Date: Sat, 11 Jul 2026 23:44:36 -0700 Subject: [PATCH 12/12] fix NaN JAX Jacobian for sqrtFun with varying onset --- src/trspecfit/eval_jax.py | 9 ++++++++- tests/models/file_time.yaml | 9 +++++++++ tests/test_evaluate_jax.py | 14 ++++++++++++++ 3 files changed, 31 insertions(+), 1 deletion(-) diff --git a/src/trspecfit/eval_jax.py b/src/trspecfit/eval_jax.py index e39b529..286ef06 100644 --- a/src/trspecfit/eval_jax.py +++ b/src/trspecfit/eval_jax.py @@ -235,7 +235,14 @@ def _erfFun(t, A, SD, t0): # def _sqrtFun(t, A, t0): - return A * jnp.sqrt(jnp.clip(t - t0, 0)) + # 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] = { 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_jax.py b/tests/test_evaluate_jax.py index f656c8c..fb2d4f4 100644 --- a/tests/test_evaluate_jax.py +++ b/tests/test_evaluate_jax.py @@ -268,6 +268,20 @@ def test_voigt(self): 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