From a1df0cc5e12476a29ba6afd769a34312c7974f6a Mon Sep 17 00:00:00 2001 From: Johannes Mahl Date: Thu, 9 Jul 2026 13:54:46 -0700 Subject: [PATCH 01/36] fix conv kernel support truncation by rebuilding from current parameter values MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The kernel time axis was built once at model construction from the initial width (±4·SD_init for gaussCONV) and never rebuilt, silently truncating the kernel and biasing the fitted width once it grew past its init. Both eval paths now rebuild the support per evaluation via a shared conv_kernel_support() helper (symmetric, odd-length, guarded): the mcp path in Component.value, the GIR path in the eval_2d conv step and the schedule_2d precompute. The frozen conv_support_* plan arrays and kernel_time node snapshots are removed. Adds a deterministic regression test (fit init 16x below truth, verified to fail on the old code), a grown-width GIR/mcp parity test, and support-builder unit tests. Drops the "init conv widths generously" workaround guidance from docs and example 04. Benchmarked: no performance change on examples 01/04. --- TODO.md | 1 - docs/ai/benchmark.md | 2 +- docs/ai/check-example.md | 4 -- docs/design/lowered_evaluator.md | 1 - .../04_parameter_profiles/example.ipynb | 1 - .../04_parameter_profiles/models_time.yaml | 5 +- pyproject.toml | 2 +- src/trspecfit/eval_2d.py | 44 ++++++++---- src/trspecfit/graph_ir.py | 60 ++++------------ src/trspecfit/mcp.py | 25 +++++-- src/trspecfit/utils/arrays.py | 35 +++++++++ tests/models/file_time.yaml | 10 +++ tests/test_arrays.py | 71 ++++++++++++++++++- tests/test_evaluate_2d.py | 41 ++++++----- tests/test_gir_integration.py | 55 ++++++++++++++ tests/test_graph_ir.py | 45 +----------- 16 files changed, 260 insertions(+), 142 deletions(-) diff --git a/TODO.md b/TODO.md index e6c3c22..2612170 100644 --- a/TODO.md +++ b/TODO.md @@ -3,7 +3,6 @@ ## Fitting - [ ] **Enable 1D time-trace fitting (post-SbS kinetics)**: wire standalone `TIME_1D` dynamics models (already evaluable on the mcp path per `docs/design/supported_models.md`, but connected to no fit method) to a fit entry point, so parameter-vs-time traces extracted from SbS results can be fit to `functions/time.py` dynamics (`expFun` sums, IRF convolution) inside the package instead of in external scripts. Frame it as the diagnostic/initialization rung before `fit_2d`, not a statistical equivalent (per-slice correlations and unpropagated SbS uncertainties make two-step inferior). Design direction: promote SbS results into a time-axis `File` — the promoted object inherits limits/CI/MCMC/save/export machinery, and trace fits land in a `SavedFitSlot` like every other fit type (no parallel fitting pipeline). Requires `File` to support time as the primary axis, which today it assumes is energy. Weighting traces by per-slice stderr ties into the `sigma_type` expansion item below. -- [ ] **Conv kernel support is sized from the initial parameter value**: `Component.create_t_kernel` builds the kernel time axis once at model construction (`t_range = par_init * kernel_width`, e.g. ±4·SD for `gaussCONV`) and never rebuilds it. When the fitted width grows past its init, the kernel is silently truncated, biasing the recovered width (observed while building `04_parameter_profiles`, 2026-06-11: truth SD=10, init 5 → fitted SD≈10.8 at any SNR; 03's SD 0.148 vs truth 0.15 with init 0.1 is likely the same effect, mild). The GIR path snapshots the same static axis (`graph_ir.py`, `kernel_time`). Fix candidates: rebuild the kernel axis when the kernel parameter value changes; size the support from the parameter's max bound; or at minimum warn when `fitted_value * kernel_width` exceeds the kernel range. Workaround used in the examples: initialize conv widths generously above the expected value (commented in the YAMLs). When fixing, add a deterministic regression test that inits the conv width well below truth (the failure mode above); broad robustness-vs-start checking deliberately lives in the benchmark skill (`--par-variability`, added 2026-07-06), not the test suite, to avoid weak-or-flaky convergence asserts. ## Noise and simulation diff --git a/docs/ai/benchmark.md b/docs/ai/benchmark.md index 1d2163c..f8d17be 100644 --- a/docs/ai/benchmark.md +++ b/docs/ai/benchmark.md @@ -90,7 +90,7 @@ The report separates the two failure signals: reached the same optimum. Nonzero spread here indicates a flat objective direction (the parameter is not identifiable from the data) or init-dependent machinery: state derived once from initial parameter values - and never rebuilt during the fit, such as a convolution-kernel support. + and never rebuilt during the fit. Parameters above 1% relative spread are flagged `start-sensitive`. This is a diagnostic, not a pass/fail test — it deliberately lives here rather diff --git a/docs/ai/check-example.md b/docs/ai/check-example.md index 1c0b1e7..df937a9 100644 --- a/docs/ai/check-example.md +++ b/docs/ai/check-example.md @@ -211,10 +211,6 @@ Hard-won gotchas worth re-checking when a criterion looks borderline: - **`%%capture` path quoting (IPython 9.x):** a bare path in `%cd -q ../dir` tokenizes as a malformed number and crashes. Quote it. -- **Convolution kernel sizing:** `create_t_kernel` sizes the kernel from the - *initial* parameter value and never rebuilds it — a too-small init silently - truncates the kernel and biases the fitted width. Initialize conv widths - generously above the expected value. - **Subcycle-boundary time samples:** generate synthetic multi-cycle data on the *reloaded* CSV axes, not the in-memory `np.arange` axes — boundary-exact samples flip subcycle assignment under `%.6e` rounding and bias the fit. diff --git a/docs/design/lowered_evaluator.md b/docs/design/lowered_evaluator.md index db49784..d028a7f 100644 --- a/docs/design/lowered_evaluator.md +++ b/docs/design/lowered_evaluator.md @@ -199,7 +199,6 @@ class GraphNode: # SUBCYCLE_MASK: {"time_n_sub": array} # SUBCYCLE_REMAP: {"time_norm": array} # PROFILE_SAMPLE: {"aux_axis": array} - # CONVOLUTION: {"kernel_time": array} # Empty dict for nodes that need no array payload. # This is graph-level data, not backend-specific -- # the scheduler copies what it needs into the plan. diff --git a/examples/fitting_workflows/04_parameter_profiles/example.ipynb b/examples/fitting_workflows/04_parameter_profiles/example.ipynb index 718dc64..364fda5 100644 --- a/examples/fitting_workflows/04_parameter_profiles/example.ipynb +++ b/examples/fitting_workflows/04_parameter_profiles/example.ipynb @@ -443,7 +443,6 @@ "\n", "**Time-dependent profile parameters**\n", "- Attach the profile first, then call `add_time_dependence` on the profile parameter (`{target_par}_{component}_{par}` name); the model is promoted to 2D automatically.\n", - "- Initialize convolution widths generously — the kernel support is sized from the initial value (see the comment in `models_time.yaml`).\n", "\n", "**Available profile functions** ([`src/trspecfit/functions/profile.py`](../../../src/trspecfit/functions/profile.py))\n", "- `pExpDecay(x, A, tau)` — IMFP weighting, Beer-Lambert absorption\n", diff --git a/examples/fitting_workflows/04_parameter_profiles/models_time.yaml b/examples/fitting_workflows/04_parameter_profiles/models_time.yaml index 70c9b11..425514f 100644 --- a/examples/fitting_workflows/04_parameter_profiles/models_time.yaml +++ b/examples/fitting_workflows/04_parameter_profiles/models_time.yaml @@ -18,10 +18,7 @@ # blurs the step with the Gaussian instrument response (IRF). BandBendingRecovery: gaussCONV: - # IRF width (ps). Init generously: the convolution kernel support is - # sized from the initial value (+-4*SD), so a too-small init truncates - # the kernel when the fitted SD grows beyond it. - SD: [20, True, 1, 50] + SD: [20, True, 1, 50] # IRF width (ps) expFun: A: [0.3, True, 0, 1] # collapse amplitude (eV/nm) tau: [50, True, 5, 500] # recovery time constant (ps) diff --git a/pyproject.toml b/pyproject.toml index b06690d..f91d28f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "trspecfit" -version = "0.10.0" +version = "0.10.1" authors = [ {name = "Johannes Mahl", email = "infinitymonkeyatwork@gmail.com"}, ] diff --git a/src/trspecfit/eval_2d.py b/src/trspecfit/eval_2d.py index 88310bf..c137fe6 100644 --- a/src/trspecfit/eval_2d.py +++ b/src/trspecfit/eval_2d.py @@ -22,7 +22,7 @@ ParamSourceKind, ScheduledPlan2D, ) -from trspecfit.utils.arrays import my_conv +from trspecfit.utils.arrays import conv_kernel_support, my_conv # --------------------------------------------------------------------------- # Dynamics dispatch table @@ -38,17 +38,30 @@ DynFuncKind.STEPFUN: (fcts_time.stepFun, 2), } -# Convolution kernel dispatch: kernel function evaluated on the frozen -# kernel-time support with per-theta kernel parameters. Mirrors MCP's -# Model.combine(...) path. +# Convolution kernel dispatch: kernel function plus its *_kernel_width +# helper. The kernel support is rebuilt from the current kernel +# parameters at every evaluation so it tracks the fitted width. +# Mirrors MCP's Component.value(...) conv path. CONV_KERNEL_DISPATCH: dict[int, tuple] = { - ConvKernelKind.GAUSSCONV: (fcts_time.gaussCONV, 1), - ConvKernelKind.LORENTZCONV: (fcts_time.lorentzCONV, 1), - ConvKernelKind.VOIGTCONV: (fcts_time.voigtCONV, 2), - ConvKernelKind.EXPSYMCONV: (fcts_time.expSymCONV, 1), - ConvKernelKind.EXPDECAYCONV: (fcts_time.expDecayCONV, 1), - ConvKernelKind.EXPRISECONV: (fcts_time.expRiseCONV, 1), - ConvKernelKind.BOXCONV: (fcts_time.boxCONV, 1), + ConvKernelKind.GAUSSCONV: (fcts_time.gaussCONV, fcts_time.gaussCONV_kernel_width), + ConvKernelKind.LORENTZCONV: ( + fcts_time.lorentzCONV, + fcts_time.lorentzCONV_kernel_width, + ), + ConvKernelKind.VOIGTCONV: (fcts_time.voigtCONV, fcts_time.voigtCONV_kernel_width), + ConvKernelKind.EXPSYMCONV: ( + fcts_time.expSymCONV, + fcts_time.expSymCONV_kernel_width, + ), + ConvKernelKind.EXPDECAYCONV: ( + fcts_time.expDecayCONV, + fcts_time.expDecayCONV_kernel_width, + ), + ConvKernelKind.EXPRISECONV: ( + fcts_time.expRiseCONV, + fcts_time.expRiseCONV_kernel_width, + ), + ConvKernelKind.BOXCONV: (fcts_time.boxCONV, fcts_time.boxCONV_kernel_width), } @@ -337,16 +350,17 @@ def evaluate_2d(plan: ScheduledPlan2D, theta: np.ndarray) -> np.ndarray: else: # kind == 2: resolved-trace convolution target = int(plan.conv_target_rows[idx]) func_id = int(plan.conv_func_ids[idx]) - kernel_func, _k_par = CONV_KERNEL_DISPATCH[func_id] + kernel_func, width_func = CONV_KERNEL_DISPATCH[func_id] p_start = int(plan.conv_param_indptr[idx]) p_end = int(plan.conv_param_indptr[idx + 1]) kernel_params = [ float(traces[int(plan.conv_param_rows[j]), 0]) for j in range(p_start, p_end) ] - s_start = int(plan.conv_support_indptr[idx]) - s_end = int(plan.conv_support_indptr[idx + 1]) - support = plan.conv_support_values[s_start:s_end] + support = conv_kernel_support( + kernel_params[0] * width_func(*kernel_params), + float(plan.time[1] - plan.time[0]), + ) kernel = kernel_func(support, *kernel_params) traces[target, :] = my_conv(plan.time, traces[target, :], kernel) diff --git a/src/trspecfit/graph_ir.py b/src/trspecfit/graph_ir.py index 64f5866..7869906 100644 --- a/src/trspecfit/graph_ir.py +++ b/src/trspecfit/graph_ir.py @@ -518,21 +518,15 @@ class ScheduledPlan2D: # --- Resolved-trace convolution program --- # Each conv step rewrites a trace row in-place after its PARAM_PLUS_TRACE # is fully resolved. Chained CONVOLUTION nodes emit multiple steps - # targeting the same row, executed in order. Kernel values are - # recomputed per theta from conv_param_rows; support values are frozen - # at plan-build time from ``node.arrays["kernel_time"]``. Only - # ``package == "time"`` kernels are lowered. + # targeting the same row, executed in order. Kernel values and the + # kernel support are recomputed per theta from conv_param_rows, so the + # support tracks the fitted width. Only ``package == "time"`` kernels + # are lowered. n_conv_steps: int conv_target_rows: np.ndarray # (n_conv_steps,) int -- trace row rewritten conv_func_ids: np.ndarray # (n_conv_steps,) int -- kernel function registry id conv_param_indptr: np.ndarray # (n_conv_steps + 1,) int -- CSR row pointers conv_param_rows: np.ndarray # (total_conv_params,) int -- kernel param trace rows - conv_support_indptr: ( - np.ndarray - ) # (n_conv_steps + 1,) int -- CSR into support values - conv_support_values: ( - np.ndarray - ) # (total_support,) float -- kernel time axis samples # --- Profile-varying parameter groups (fixed aux_axis shape) --- n_aux: int @@ -841,9 +835,6 @@ def build_graph(model: Model) -> GraphIR: function_name=comp.fct_str, package=pkg_name, ) - # Store kernel-related arrays if available - if comp.time is not None: - b.nodes[nid].arrays["kernel_time"] = comp.time elif is_shirley: nid = b.add_node( NodeKind.SPECTRUM_FED_OP, @@ -1054,8 +1045,6 @@ def _emit_dynamics_subgraph( function_name=dyn_comp.fct_str, package="time", ) - if dyn_comp.time is not None: - b.nodes[nid].arrays["kernel_time"] = dyn_comp.time # Wire dynamics params -> conv node for pos, dnid in enumerate(dyn_param_nids[i]): b.add_edge(dnid, nid, EdgeKind.PARAM_INPUT, position=pos) @@ -1523,8 +1512,7 @@ def _is_lowerable_convolution_2d(node: GraphNode, graph: GraphIR) -> bool: Lowering contract -- all of: 1. Time-domain kernel (``package == "time"``) with a registered - kernel function and a ``kernel_time`` array populated at - graph-build time. + kernel function. 2. Resolved-trace shape: the node has exactly one ``TRACE_INPUT`` ancestor, and walking that chain terminates at a ``PARAM_PLUS_TRACE``. @@ -1541,7 +1529,6 @@ def _is_lowerable_convolution_2d(node: GraphNode, graph: GraphIR) -> bool: node.kind != NodeKind.CONVOLUTION or node.package != "time" or node.function_name not in _FUNCTION_NAME_TO_CONV_KERNEL - or "kernel_time" not in node.arrays ): return False @@ -2391,8 +2378,6 @@ def schedule_2d(graph: GraphIR) -> ScheduledPlan2D: conv_func_ids_list: list[int] = [] conv_param_indptr_list: list[int] = [0] conv_param_rows_list: list[int] = [] - conv_support_indptr_list: list[int] = [0] - conv_support_values_list: list[float] = [] for conv_node in conv_nodes_topo: assert conv_node.function_name is not None @@ -2419,21 +2404,10 @@ def schedule_2d(graph: GraphIR) -> ScheduledPlan2D: conv_param_rows_list.append(name_to_row[src_node.name]) conv_param_indptr_list.append(len(conv_param_rows_list)) - support = conv_node.arrays.get("kernel_time") - if support is None: - raise ValueError( - f"CONVOLUTION node {conv_node.name!r} is missing" - " kernel_time array (required for lowered convolution)" - ) - conv_support_values_list.extend(float(v) for v in np.asarray(support)) - conv_support_indptr_list.append(len(conv_support_values_list)) - conv_target_rows = np.array(conv_target_rows_list, dtype=np.intp) conv_func_ids = np.array(conv_func_ids_list, dtype=np.intp) conv_param_indptr = np.array(conv_param_indptr_list, dtype=np.intp) conv_param_rows = np.array(conv_param_rows_list, dtype=np.intp) - conv_support_indptr = np.array(conv_support_indptr_list, dtype=np.intp) - conv_support_values = np.array(conv_support_values_list, dtype=np.float64) resolution_kinds_list: list[int] = [] resolution_indices_list: list[int] = [] @@ -2899,17 +2873,8 @@ def schedule_2d(graph: GraphIR) -> ScheduledPlan2D: int(DynFuncKind.STEPFUN): fcts_time.stepFun, } - _CONV_KERNEL_DISPATCH: dict[int, Callable[..., Any]] = { - int(ConvKernelKind.GAUSSCONV): fcts_time.gaussCONV, - int(ConvKernelKind.LORENTZCONV): fcts_time.lorentzCONV, - int(ConvKernelKind.VOIGTCONV): fcts_time.voigtCONV, - int(ConvKernelKind.EXPSYMCONV): fcts_time.expSymCONV, - int(ConvKernelKind.EXPDECAYCONV): fcts_time.expDecayCONV, - int(ConvKernelKind.EXPRISECONV): fcts_time.expRiseCONV, - int(ConvKernelKind.BOXCONV): fcts_time.boxCONV, - } - from trspecfit.eval_2d import eval_expr_program - from trspecfit.utils.arrays import my_conv + from trspecfit.eval_2d import CONV_KERNEL_DISPATCH, eval_expr_program + from trspecfit.utils.arrays import conv_kernel_support, my_conv # Dynamics groups, expressions, and resolved-trace convolutions are # interleaved in topological order so that downstream consumers see @@ -2939,16 +2904,17 @@ def schedule_2d(graph: GraphIR) -> ScheduledPlan2D: ) else: # kind == 2: resolved-trace convolution target_row = int(conv_target_rows[idx]) - kernel_func = _CONV_KERNEL_DISPATCH[int(conv_func_ids[idx])] + kernel_func, width_func = CONV_KERNEL_DISPATCH[int(conv_func_ids[idx])] p_start = int(conv_param_indptr[idx]) p_end = int(conv_param_indptr[idx + 1]) kernel_params = [ float(param_traces_init[int(conv_param_rows[j]), 0]) for j in range(p_start, p_end) ] - s_start = int(conv_support_indptr[idx]) - s_end = int(conv_support_indptr[idx + 1]) - support = conv_support_values[s_start:s_end] + support = conv_kernel_support( + kernel_params[0] * width_func(*kernel_params), + float(graph.time[1] - graph.time[0]), + ) kernel = kernel_func(support, *kernel_params) param_traces_init[target_row, :] = my_conv( graph.time, param_traces_init[target_row, :], kernel @@ -3070,8 +3036,6 @@ def schedule_2d(graph: GraphIR) -> ScheduledPlan2D: conv_func_ids=conv_func_ids, conv_param_indptr=conv_param_indptr, conv_param_rows=conv_param_rows, - conv_support_indptr=conv_support_indptr, - conv_support_values=conv_support_values, ) diff --git a/src/trspecfit/mcp.py b/src/trspecfit/mcp.py index 810a806..95716ac 100644 --- a/src/trspecfit/mcp.py +++ b/src/trspecfit/mcp.py @@ -1617,7 +1617,7 @@ def describe(self, detail: int = 1) -> None: print() # - def create_t_kernel(self) -> np.ndarray: + def create_t_kernel(self, *, par_values: list[Any] | None = None) -> np.ndarray: """ Create time axis for convolution kernel. @@ -1625,6 +1625,13 @@ def create_t_kernel(self) -> np.ndarray: time axis to properly handle edge effects. This method creates an appropriately sized kernel axis based on the kernel width. + Parameters + ---------- + par_values : list, optional + Current kernel parameter values. Defaults to the initial + values from par_dict (model construction). Component.value + passes the live values so the support tracks the fitted width. + Returns ------- ndarray @@ -1632,15 +1639,18 @@ def create_t_kernel(self) -> np.ndarray: """ # get kernel parameters i.e. component parameters - par_k = cast("list[Any]", ulmfit.par_extract(self.par_dict, return_type="list")) + if par_values is None: + par_values = cast( + "list[Any]", ulmfit.par_extract(self.par_dict, return_type="list") + ) # define kernel time axis. Kernel-width helpers may inspect the # full parameter list for multi-parameter kernels such as Voigt. - kernel_width = getattr(fcts_time, self.fct_str + "_kernel_width")(*par_k) - t_range = par_k[0] * kernel_width + kernel_width = getattr(fcts_time, self.fct_str + "_kernel_width")(*par_values) + t_range = par_values[0] * kernel_width if self.time is None or len(self.time) < 2: raise ValueError(f"time axis of component {self.fct_str} not defined") t_step = self.time[1] - self.time[0] - return np.arange(-t_range, t_range + t_step, t_step) + return uarr.conv_kernel_support(t_range, t_step) # def value(self, t_ind: int = 0, **kwargs) -> np.ndarray: @@ -1711,6 +1721,11 @@ def value(self, t_ind: int = 0, **kwargs) -> np.ndarray: raise ValueError( f"Time axis not defined for component '{self.comp_name}'" ) + # conv kernels: rebuild the support from the current parameter + # values so the axis tracks the fitted width (a frozen support + # would truncate the kernel once the width grows past its init) + if self.comp_type == "conv": + self.time = self.create_t_kernel(par_values=pars) if self.subcycle == 0: # single cycle return np.asarray(self.fct(self.time, *pars, **kwargs)) # multi-cycle diff --git a/src/trspecfit/utils/arrays.py b/src/trspecfit/utils/arrays.py index dcdcecf..5e12faf 100644 --- a/src/trspecfit/utils/arrays.py +++ b/src/trspecfit/utils/arrays.py @@ -308,6 +308,41 @@ def pad_x_y( return x_pad, y_pad +# +def conv_kernel_support(t_range: float, t_step: float) -> NDArray[np.float64]: + """ + Build a symmetric, odd-length time axis for a convolution kernel. + + Rebuilt from the current kernel parameter values at every model + evaluation so the support tracks the fitted width — a frozen support + silently truncates the kernel when the width grows past its initial + value. + + Parameters + ---------- + t_range : float + Half-width of the kernel support, typically + ``kernel_par * kernel_width(...)`` from the ``*_kernel_width`` + helpers in trspecfit.functions.time + t_step : float + Time axis step size (matches the data time axis) + + Returns + ------- + ndarray + Kernel time axis ``[-n*t_step, ..., 0, ..., n*t_step]`` with + ``n = max(1, ceil(|t_range| / t_step))``, guaranteed symmetric + and odd-length so ``mode='same'`` convolution stays centered + """ + + if not np.isfinite(t_range): + raise ValueError(f"Kernel support range is not finite: {t_range}") + if not np.isfinite(t_step) or t_step <= 0: + raise ValueError(f"Kernel time step must be positive: {t_step}") + n = max(1, int(np.ceil(abs(t_range) / t_step))) + return np.arange(-n, n + 1, dtype=np.float64) * t_step + + # def my_conv( x: ArrayLike, diff --git a/tests/models/file_time.yaml b/tests/models/file_time.yaml index 034defe..adb8f10 100644 --- a/tests/models/file_time.yaml +++ b/tests/models/file_time.yaml @@ -89,6 +89,16 @@ MonoExpPosIRF: tau: [2.5, True, 1, 10] t0: [0, False, 0, 1] +# truth-side twin of MonoExpPosIRF for the kernel-support regression test: +# SD initializes at the truth value, 16x above MonoExpPosIRF's init +MonoExpPosIRFWide: + gaussCONV: + SD: [0.8, True, 0, 1] + expFun: + A: [1, True, 0, 5] + tau: [2.5, True, 1, 10] + t0: [0, False, 0, 1] + MonoExpPosLorentzIRF: lorentzCONV: W: [0.5, True, 0.01, 5] diff --git a/tests/test_arrays.py b/tests/test_arrays.py index 920c29d..0388cf8 100644 --- a/tests/test_arrays.py +++ b/tests/test_arrays.py @@ -1,8 +1,75 @@ -"""Tests for trspecfit.utils.arrays — running_mean.""" +"""Tests for trspecfit.utils.arrays — running_mean, conv_kernel_support.""" import numpy as np +import pytest -from trspecfit.utils.arrays import running_mean +from trspecfit.utils.arrays import conv_kernel_support, running_mean + + +# +# +class TestConvKernelSupport: + """Tests for the convolution kernel support builder.""" + + # + def test_symmetric_and_odd_length(self): + """Support is symmetric around 0 with an odd number of samples.""" + + axis = conv_kernel_support(3.7, 0.5) + assert axis.size % 2 == 1 + np.testing.assert_allclose(axis, -axis[::-1]) + assert axis[axis.size // 2] == 0.0 + + # + def test_covers_t_range(self): + """Support extends to at least ±t_range.""" + + axis = conv_kernel_support(3.7, 0.5) + assert axis.max() >= 3.7 + assert axis.min() <= -3.7 + + # + def test_exact_multiple(self): + """t_range on the grid gives endpoints exactly at ±t_range.""" + + axis = conv_kernel_support(4.0, 0.5) + np.testing.assert_allclose(axis[0], -4.0) + np.testing.assert_allclose(axis[-1], 4.0) + assert axis.size == 17 + + # + def test_minimum_support(self): + """t_range below one step still yields a 3-sample support.""" + + axis = conv_kernel_support(0.01, 0.5) + np.testing.assert_allclose(axis, [-0.5, 0.0, 0.5]) + + # + def test_grows_with_t_range(self): + """Doubling t_range widens the support accordingly.""" + + narrow = conv_kernel_support(2.0, 0.5) + wide = conv_kernel_support(4.0, 0.5) + assert wide.size > narrow.size + assert wide.max() >= 2 * narrow.max() - 0.5 + + # + def test_nonfinite_range_raises(self): + """Non-finite t_range raises ValueError.""" + + with pytest.raises(ValueError, match="not finite"): + conv_kernel_support(np.nan, 0.5) + with pytest.raises(ValueError, match="not finite"): + conv_kernel_support(np.inf, 0.5) + + # + def test_bad_step_raises(self): + """Non-positive or non-finite t_step raises ValueError.""" + + with pytest.raises(ValueError, match="positive"): + conv_kernel_support(1.0, 0.0) + with pytest.raises(ValueError, match="positive"): + conv_kernel_support(1.0, -0.5) # diff --git a/tests/test_evaluate_2d.py b/tests/test_evaluate_2d.py index 6aefea4..2df6360 100644 --- a/tests/test_evaluate_2d.py +++ b/tests/test_evaluate_2d.py @@ -714,8 +714,8 @@ def test_profile_with_time_dep_profile_params(self): class TestDynamicsConvolution: """Lowered IRF dynamics: CONVOLUTION is compiled into a kind=2 step. - Covers plan encoding (conv program layout, frozen support, resolution - ordering) and numerical parity against ``Model.create_value_2d()``. + Covers plan encoding (conv program layout, resolution ordering) and + numerical parity against ``Model.create_value_2d()``. The IRF path rewrites a resolved trace row in place via ``my_conv`` -- the same code MCP calls -- so tolerance matches the other OpKind parity tests. @@ -763,10 +763,6 @@ def test_conv_program_populated(self): n_kernel_params = int(plan.conv_param_indptr[1]) assert n_kernel_params >= 1 # at least SD assert plan.conv_param_rows.shape == (n_kernel_params,) - assert plan.conv_support_indptr.shape == (2,) - assert plan.conv_support_indptr[0] == 0 - n_support = int(plan.conv_support_indptr[1]) - assert plan.conv_support_values.shape == (n_support,) # def test_conv_target_row_valid(self): @@ -785,18 +781,30 @@ def test_conv_param_rows_valid(self): assert 0 <= int(row) < plan.n_params # - def test_conv_support_frozen_from_kernel_time(self): - """Plan's conv support matches the CONVOLUTION node's kernel_time array.""" + def test_conv_support_tracks_grown_kernel_width(self): + """Parity holds when the kernel width grows well past its init. - plan, graph, _model = self._make_irf_plan() - from trspecfit.graph_ir import NodeKind + Regression guard for the frozen-support bug: the kernel support + used to be sized once from the initial SD, silently truncating + the kernel (and breaking parity) once the fitted SD grew past + it. Both paths now rebuild the support from the current value. + """ - conv_nodes = [n for n in graph.nodes if n.kind == NodeKind.CONVOLUTION] - assert len(conv_nodes) == 1 - expected = conv_nodes[0].arrays["kernel_time"] - start = int(plan.conv_support_indptr[0]) - end = int(plan.conv_support_indptr[1]) - np.testing.assert_array_equal(plan.conv_support_values[start:end], expected) + plan, _graph, model = self._make_irf_plan() + theta = _compare_evaluator_vs_interpreter(model, plan) + sd_idx = plan.opt_param_names.index("GLP_01_A_gaussCONV_SD") + # SD init is 5e-2 on a dt=2.2 axis; +0.6 grows the support from + # a sub-sample kernel to a multi-sample one + _perturb_theta(plan, model, theta, [sd_idx], [0.6]) + + # interpreter side: the kernel axis was rebuilt to span the + # grown width (±4*SD for gaussCONV) + A_par = next(p for p in model.components[0].pars if p.name == "GLP_01_A") + assert A_par.t_model is not None # type guard + conv_comp = next(c for c in A_par.t_model.components if c.comp_type == "conv") + assert conv_comp.time is not None # type guard + SD_grown = theta[sd_idx] + 0.6 + assert conv_comp.time.max() >= 4 * SD_grown # def test_conv_step_runs_after_dynamics(self): @@ -837,7 +845,6 @@ def test_non_irf_has_empty_conv_program(self): assert plan.conv_target_rows.shape == (0,) assert plan.conv_func_ids.shape == (0,) assert plan.conv_param_rows.shape == (0,) - assert plan.conv_support_values.shape == (0,) assert 2 not in plan.resolution_kinds.tolist() # Per-kernel parity: every lowerable kernel must match MCP at diff --git a/tests/test_gir_integration.py b/tests/test_gir_integration.py index 2fd5a3a..687997f 100644 --- a/tests/test_gir_integration.py +++ b/tests/test_gir_integration.py @@ -695,6 +695,61 @@ def test_gir_fit_writes_back_to_model(self): f"{name}: true={true_val:.6f}, fit={fit_val:.6f}" ) + # + @pytest.mark.slow + def test_kernel_width_recovered_when_init_below_truth(self): + """Conv kernel width is recovered when initialized well below truth. + + Regression for the frozen-kernel-support bug: the support was + sized once from the initial SD, so once the fitted SD grew past + its init the kernel was silently truncated and the fit converged + to a biased width at any SNR. The support now tracks the current + parameter value, so a fit started far below truth must recover it. + """ + + project = make_project(name="gir_kernel_regrow") + SD_name = "GLP_01_A_gaussCONV_SD" + SD_truth = 0.8 + + energy = np.linspace(83, 87, 30) + time = np.linspace(-2, 10, 61) # dt = 0.2 + # truth model: MonoExpPosIRFWide inits SD at the truth value, so + # the simulated data carries a correctly sized kernel regardless + # of when the support is built + 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=_FILE_ENERGY_YAML, model_info="single_glp") + truth_file.add_time_dependence( + target_model="single_glp", + target_parameter="GLP_01_A", + dynamics_yaml=_TIME_YAML, + dynamics_model=["MonoExpPosIRFWide"], + ) + truth_model = truth_file.model_active + assert truth_model is not None # type guard + assert truth_model.lmfit_pars[SD_name].value == SD_truth + clean = simulate_clean(truth_model) + + # fit model: same shape, but SD inits at 5e-2, 16x below truth -- + # a frozen ±4*init support could never represent the true kernel + fit_file = _make_fit_file(project, clean, energy, time) + 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=_TIME_YAML, + dynamics_model=["MonoExpPosIRF"], + ) + fit_file.fit_2d(model_name="single_glp", stages=2, try_ci=0) + + assert fit_file.model_2d is not None # type guard + SD_fit = fit_file.model_2d.result[1].params[SD_name].value + assert np.isclose(SD_fit, SD_truth, rtol=2e-2), ( + f"kernel width not recovered: truth={SD_truth}, fit={SD_fit:.4f}" + ) + # @pytest.mark.slow def test_compare_mode_through_fit_2d(self): diff --git a/tests/test_graph_ir.py b/tests/test_graph_ir.py index 060fdf8..4538999 100644 --- a/tests/test_graph_ir.py +++ b/tests/test_graph_ir.py @@ -1421,32 +1421,15 @@ def test_expfun_still_emitted_as_dynamics_trace(self): def test_irf_dynamics_lowerable_in_2d(self): """Time-domain IRF dynamics are lowerable on the 2D backend. - The CONVOLUTION node carries ``package == "time"``, a registered - kernel function, and a populated ``kernel_time`` array, so - ``can_lower_2d`` accepts the graph. + The CONVOLUTION node carries ``package == "time"`` and a + registered kernel function, so ``can_lower_2d`` accepts the + graph. """ _file, model = _make_irf_dynamics_model() graph = build_graph(model) assert can_lower_2d(graph) - # - def test_irf_without_kernel_time_not_lowerable(self): - """CONVOLUTION missing ``kernel_time`` falls back to MCP. - - Safety guard for the lowering contract: if the graph builder - omits the frozen support metadata (e.g. a future code path that - doesn't populate it), the 2D backend must reject the graph - rather than silently producing wrong numerics. - """ - - _file, model = _make_irf_dynamics_model() - graph = build_graph(model) - conv_nodes = _nodes_by_kind(graph, NodeKind.CONVOLUTION) - assert len(conv_nodes) == 1 - del conv_nodes[0].arrays["kernel_time"] - assert not can_lower_2d(graph) - # def test_conv_without_ppt_chain_not_lowerable(self): """Structural gate: conv must wrap a PARAM_PLUS_TRACE to be lowered. @@ -1532,28 +1515,6 @@ def test_convolution_is_final_resolved(self): source = graph.nodes[A_edge.source] assert source.kind == NodeKind.CONVOLUTION - # - def test_kernel_time_populated_on_dynamics_conv(self): - """Dynamics CONVOLUTION nodes carry the kernel time axis in arrays. - - Lowering contract: lowered time-domain convolution requires - ``node.arrays["kernel_time"]`` to be present at graph-build time so - the scheduler can freeze kernel support without re-deriving it from - MCP helpers. Mirrors the top-level handling for conv components. - """ - - _file, model = _make_irf_dynamics_model() - graph = build_graph(model) - - conv_nodes = _nodes_by_kind(graph, NodeKind.CONVOLUTION) - assert len(conv_nodes) == 1 - conv = conv_nodes[0] - assert conv.package == "time" - assert "kernel_time" in conv.arrays - kernel_time = conv.arrays["kernel_time"] - assert kernel_time.ndim == 1 - assert kernel_time.size > 0 - # # From 3da2be7e2474f40d95e47326ed8d5d93d58c19b6 Mon Sep 17 00:00:00 2001 From: Johannes Mahl Date: Thu, 9 Jul 2026 15:31:33 -0700 Subject: [PATCH 02/36] make conv IRF test fixture identifiable on its time grid MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The MonoExpPosIRF gaussCONV SD (5e-2) was ~10x below the shared test time step, so the kernel was numerically a delta and chi² was bit-flat in SD: the F4 roundtrip fit landed wherever optimizer round-off left it, which differs across scipy/lmfit versions and machines (failed on min-versions CI only, deterministically). The old frozen-support code had masked this with a spurious gradient from its asymmetric kernel. Raise MonoExpPosIRF SD to 0.4 (~0.8*dt): identifiable (exact recovery from 0.5x/2x starts, r(SD,A)=0.47) without burying the expFun dynamics (peak keeps 77%). Keep the sub-sample variant as MonoExpPosIRFNarrow for kernel-support tests that need a far-below-truth fit init. --- tests/models/file_time.yaml | 14 +++++++++----- tests/test_evaluate_2d.py | 10 +++++----- tests/test_gir_integration.py | 14 +++++++------- tests/test_model_parser.py | 4 ++-- 4 files changed, 23 insertions(+), 19 deletions(-) diff --git a/tests/models/file_time.yaml b/tests/models/file_time.yaml index adb8f10..dd9fac7 100644 --- a/tests/models/file_time.yaml +++ b/tests/models/file_time.yaml @@ -81,19 +81,23 @@ MonoSqrt: A: [0.2, True, 0, 5] t0: [0, False, 0, 1] +# 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: gaussCONV: - SD: [5.0E-2, True, 0, 1] + SD: [0.4, True, 0, 1] expFun: A: [1, True, 0, 5] tau: [2.5, True, 1, 10] t0: [0, False, 0, 1] -# truth-side twin of MonoExpPosIRF for the kernel-support regression test: -# SD initializes at the truth value, 16x above MonoExpPosIRF's init -MonoExpPosIRFWide: +# sub-sample IRF: SD sits far below the shared test time steps, so the +# kernel is numerically a delta and SD is not identifiable from the data. +# Use as the far-below-truth fit init in kernel-support tests, never as a +# recovery target. +MonoExpPosIRFNarrow: gaussCONV: - SD: [0.8, True, 0, 1] + SD: [5.0E-2, True, 0, 1] expFun: A: [1, True, 0, 5] tau: [2.5, True, 1, 10] diff --git a/tests/test_evaluate_2d.py b/tests/test_evaluate_2d.py index 2df6360..18e350b 100644 --- a/tests/test_evaluate_2d.py +++ b/tests/test_evaluate_2d.py @@ -722,10 +722,10 @@ class TestDynamicsConvolution: """ # - def _make_irf_plan(self): + def _make_irf_plan(self, dyn_model="MonoExpPosIRF"): file, model = _make_2d_model( ["glp_only"], - [("GLP_01_A", ["MonoExpPosIRF"])], + [("GLP_01_A", [dyn_model])], ) graph = build_graph(model) assert can_lower_2d(graph) @@ -790,11 +790,11 @@ def test_conv_support_tracks_grown_kernel_width(self): it. Both paths now rebuild the support from the current value. """ - plan, _graph, model = self._make_irf_plan() + plan, _graph, model = self._make_irf_plan(dyn_model="MonoExpPosIRFNarrow") theta = _compare_evaluator_vs_interpreter(model, plan) sd_idx = plan.opt_param_names.index("GLP_01_A_gaussCONV_SD") - # SD init is 5e-2 on a dt=2.2 axis; +0.6 grows the support from - # a sub-sample kernel to a multi-sample one + # Narrow SD init is 5e-2 on a dt=2.2 axis; +0.6 grows the support + # from a sub-sample kernel to a multi-sample one (within bounds) _perturb_theta(plan, model, theta, [sd_idx], [0.6]) # interpreter side: the kernel axis was rebuilt to span the diff --git a/tests/test_gir_integration.py b/tests/test_gir_integration.py index 687997f..421990c 100644 --- a/tests/test_gir_integration.py +++ b/tests/test_gir_integration.py @@ -709,13 +709,13 @@ def test_kernel_width_recovered_when_init_below_truth(self): project = make_project(name="gir_kernel_regrow") SD_name = "GLP_01_A_gaussCONV_SD" - SD_truth = 0.8 + SD_truth = 0.4 energy = np.linspace(83, 87, 30) time = np.linspace(-2, 10, 61) # dt = 0.2 - # truth model: MonoExpPosIRFWide inits SD at the truth value, so - # the simulated data carries a correctly sized kernel regardless - # of when the support is built + # truth model: MonoExpPosIRF inits SD at the truth value, so the + # simulated data carries a correctly sized kernel regardless of + # when the support is built truth_file = File(parent_project=project, name="truth") truth_file.energy = energy truth_file.time = time @@ -725,14 +725,14 @@ def test_kernel_width_recovered_when_init_below_truth(self): target_model="single_glp", target_parameter="GLP_01_A", dynamics_yaml=_TIME_YAML, - dynamics_model=["MonoExpPosIRFWide"], + dynamics_model=["MonoExpPosIRF"], ) truth_model = truth_file.model_active assert truth_model is not None # type guard assert truth_model.lmfit_pars[SD_name].value == SD_truth clean = simulate_clean(truth_model) - # fit model: same shape, but SD inits at 5e-2, 16x below truth -- + # fit model: same shape, but SD inits at 5e-2, 8x below truth -- # a frozen ±4*init support could never represent the true kernel fit_file = _make_fit_file(project, clean, energy, time) fit_file.fit_baseline(model_name="single_glp", stages=2, try_ci=0) @@ -740,7 +740,7 @@ def test_kernel_width_recovered_when_init_below_truth(self): target_model="single_glp", target_parameter="GLP_01_A", dynamics_yaml=_TIME_YAML, - dynamics_model=["MonoExpPosIRF"], + dynamics_model=["MonoExpPosIRFNarrow"], ) fit_file.fit_2d(model_name="single_glp", stages=2, try_ci=0) diff --git a/tests/test_model_parser.py b/tests/test_model_parser.py index b9c5475..4668480 100644 --- a/tests/test_model_parser.py +++ b/tests/test_model_parser.py @@ -231,7 +231,7 @@ def test_IRF_model(self): # check the components assert model.components[0].fct_str == "gaussCONV" assert model.components[0].comp_name == "gaussCONV" - assert model.components[0].par_dict["SD"] == [5.0e-2, True, 0, 1] + assert model.components[0].par_dict["SD"] == [0.4, True, 0, 1] assert model.components[1].fct_str == "expFun" assert model.components[1].comp_name == "expFun_01" @@ -342,7 +342,7 @@ def test_simple_2D_model(self): td_par_model = model.components[2].pars[1].t_model assert td_par_model is not None # type guard assert td_par_model.components[0].comp_name == "gaussCONV" - assert td_par_model.components[0].par_dict["SD"] == [5.0e-2, True, 0, 1] + assert td_par_model.components[0].par_dict["SD"] == [0.4, True, 0, 1] assert td_par_model.components[1].fct_str == "expFun" assert td_par_model.components[1].comp_name == "expFun_01" assert td_par_model.components[1].par_dict["A"] == [1, True, 0, 5] From aeb1c623d0dd78d461164c7ec598377b2f54683e Mon Sep 17 00:00:00 2001 From: Johannes Mahl Date: Thu, 9 Jul 2026 19:10:44 -0700 Subject: [PATCH 03/36] fix plot_comparison on a fresh Simulator The SNR plot title called get_snr() before the data_clean auto-simulate guard, so plot_comparison without a prior simulate_1d/2d raised ValueError instead of simulating first. Hoist the auto-simulate above the title construction (code-review-2026-07, FAIL 1/4). --- src/trspecfit/simulator.py | 12 ++++++++---- tests/test_plotting.py | 24 ++++++++++++++++++++++++ 2 files changed, 32 insertions(+), 4 deletions(-) diff --git a/src/trspecfit/simulator.py b/src/trspecfit/simulator.py index d542b79..5c76cfb 100644 --- a/src/trspecfit/simulator.py +++ b/src/trspecfit/simulator.py @@ -1300,6 +1300,14 @@ def plot_comparison( get_snr : SNR calculation shown in title """ + # auto-simulate before building the title: get_snr raises on a + # fresh Simulator with no simulated data + if self.data_clean is None: + if dim == 1: + self.simulate_1d(t_ind) + elif dim == 2: + self.simulate_2d() + detection_str = f" [{self.detection}]" plt_title = ( f"Simulated Data (SNR: {self.get_snr(scale=snr_scale):.1f}" @@ -1307,8 +1315,6 @@ def plot_comparison( ) if dim == 1: - if self.data_clean is None: - self.simulate_1d(t_ind) if self.data_clean is None or self.data_noisy is None or self.noise is None: raise RuntimeError("Simulation data not available for plotting") @@ -1333,8 +1339,6 @@ def plot_comparison( ) elif dim == 2: - if self.data_clean is None: - self.simulate_2d() if self.data_clean is None or self.data_noisy is None or self.noise is None: raise RuntimeError("Simulation data not available for plotting") diff --git a/tests/test_plotting.py b/tests/test_plotting.py index 0eb1a10..769a4ba 100644 --- a/tests/test_plotting.py +++ b/tests/test_plotting.py @@ -794,6 +794,30 @@ def test_simulator_2d_respects_axis_labels(self): assert any(ax.get_xlabel() == "Binding Energy (eV)" for ax in main_axes) plt.close("all") + # + def test_simulator_plot_comparison_fresh_auto_simulates(self): + """plot_comparison on a fresh Simulator auto-simulates, not raises. + + Regression: the SNR plot title was built before the auto-simulate + guard, so calling plot_comparison without a prior simulate_1d/2d + raised ValueError from get_snr. + """ + + from trspecfit import Simulator + + file = self._make_file_with_model() + model = file.model_active + assert model is not None # type guard + + sim_1d = Simulator(model, noise_level=0.05) + sim_1d.plot_comparison(dim=1, save_img=-2) + assert sim_1d.data_clean is not None + + sim_2d = Simulator(model, noise_level=0.05) + sim_2d.plot_comparison(dim=2, save_img=-2) + assert sim_2d.data_clean is not None + plt.close("all") + # # From a9324faa92ecce702b5625eee3387733112615ee Mon Sep 17 00:00:00 2001 From: Johannes Mahl Date: Thu, 9 Jul 2026 19:10:58 -0700 Subject: [PATCH 04/36] fix sign_change infinite loop on all-zero input The ignore_zeros propagation loop rolls the previous non-zero sign through zeros, but all-zero input has no sign to propagate, so 'while sz.any()' never terminated. Skip propagation when asign has no non-zero entries; the result is correctly all-zeros (code-review-2026-07, FAIL 2/4). --- src/trspecfit/utils/arrays.py | 4 ++- tests/test_arrays.py | 46 +++++++++++++++++++++++++++++++++-- 2 files changed, 47 insertions(+), 3 deletions(-) diff --git a/src/trspecfit/utils/arrays.py b/src/trspecfit/utils/arrays.py index 5e12faf..0038467 100644 --- a/src/trspecfit/utils/arrays.py +++ b/src/trspecfit/utils/arrays.py @@ -236,7 +236,9 @@ def sign_change(array: ArrayLike, *, ignore_zeros: bool = True) -> NDArray[np.in asign = np.sign(array) - if ignore_zeros: + # all-zero input has no sign to propagate (and the loop below would + # never terminate); the roll-difference then correctly yields no changes + if ignore_zeros and asign.any(): sz = asign == 0 while sz.any(): asign[sz] = np.roll(asign, 1)[sz] diff --git a/tests/test_arrays.py b/tests/test_arrays.py index 0388cf8..271e0d1 100644 --- a/tests/test_arrays.py +++ b/tests/test_arrays.py @@ -1,9 +1,51 @@ -"""Tests for trspecfit.utils.arrays — running_mean, conv_kernel_support.""" +"""Tests for trspecfit.utils.arrays — running_mean, conv_kernel_support, +sign_change.""" import numpy as np import pytest -from trspecfit.utils.arrays import conv_kernel_support, running_mean +from trspecfit.utils.arrays import conv_kernel_support, running_mean, sign_change + + +# +# +class TestSignChange: + """Tests for sign_change zero-crossing detection.""" + + # + def test_all_zero_input_returns_no_changes(self): + """All-zero input terminates and reports no sign changes. + + Regression: the zero-propagation loop never terminated on + all-zero input (np.roll can't introduce a non-zero sign). + """ + + result = sign_change(np.zeros(5), ignore_zeros=True) + np.testing.assert_array_equal(result, np.zeros(5, dtype=int)) + + # + def test_zero_crossing_with_ignore_zeros(self): + """Zeros between opposite signs count as one crossing.""" + + np.testing.assert_array_equal( + sign_change([1, 0, -1], ignore_zeros=True), [0, 0, 1] + ) + + # + def test_zero_crossing_without_ignore_zeros(self): + """ignore_zeros=False treats zero as its own sign.""" + + np.testing.assert_array_equal( + sign_change([1, 0, -1], ignore_zeros=False), [0, 1, 1] + ) + + # + def test_no_crossing(self): + """Same-sign input reports no changes.""" + + np.testing.assert_array_equal( + sign_change([1, 2, 0, 3], ignore_zeros=True), [0, 0, 0, 0] + ) # From 83064a90d2786d8921335f2f74eeed32cfac175e Mon Sep 17 00:00:00 2001 From: Johannes Mahl Date: Thu, 9 Jul 2026 19:12:16 -0700 Subject: [PATCH 05/36] raise a clear error for my_conv on a degenerate x axis my_conv computed x_arr[1] - x_arr[0] with no length guard, so a single-element x raised a bare IndexError from the hot path. Guard with a ValueError naming the precondition; benchmark unchanged (code-review-2026-07, FAIL 3/4). --- src/trspecfit/utils/arrays.py | 5 +++++ tests/test_arrays.py | 42 +++++++++++++++++++++++++++++++++-- 2 files changed, 45 insertions(+), 2 deletions(-) diff --git a/src/trspecfit/utils/arrays.py b/src/trspecfit/utils/arrays.py index 0038467..a18d562 100644 --- a/src/trspecfit/utils/arrays.py +++ b/src/trspecfit/utils/arrays.py @@ -379,6 +379,11 @@ def my_conv( x_arr = np.asarray(x, dtype=float) y_arr = np.asarray(y, dtype=float) kernel_arr = np.asarray(kernel, dtype=float) + if x_arr.size < 2: + raise ValueError( + "my_conv requires at least 2 x samples to determine the " + f"step size, got {x_arr.size}" + ) pad_size = int(kernel_arr.size / 2) # Add padding to minimize edge artifacts diff --git a/tests/test_arrays.py b/tests/test_arrays.py index 271e0d1..10345ae 100644 --- a/tests/test_arrays.py +++ b/tests/test_arrays.py @@ -1,10 +1,48 @@ """Tests for trspecfit.utils.arrays — running_mean, conv_kernel_support, -sign_change.""" +sign_change, my_conv.""" import numpy as np import pytest -from trspecfit.utils.arrays import conv_kernel_support, running_mean, sign_change +from trspecfit.utils.arrays import ( + conv_kernel_support, + my_conv, + running_mean, + sign_change, +) + + +# +# +class TestMyConv: + """Tests for my_conv padded convolution.""" + + # + def test_delta_kernel_is_identity(self): + """A single-sample kernel returns the input signal unchanged.""" + + x = np.linspace(0, 10, 21) + y = np.sin(x) + np.testing.assert_allclose(my_conv(x, y, np.array([1.0])), y) + + # + def test_normalization(self): + """Kernel normalization preserves the level of a constant signal.""" + + x = np.linspace(0, 10, 21) + y = np.full(21, 3.0) + kernel = np.array([1.0, 2.0, 1.0]) + np.testing.assert_allclose(my_conv(x, y, kernel), y) + + # + def test_single_element_x_raises(self): + """Degenerate x axis raises a clear error instead of IndexError. + + Regression: x_arr[1] - x_arr[0] had no length guard. + """ + + with pytest.raises(ValueError, match="at least 2 x samples"): + my_conv(np.array([0.0]), np.array([1.0]), np.array([1.0])) # From 7ced405ff67e00791334badb0a4ebd7edc4ef948 Mon Sep 17 00:00:00 2001 From: Johannes Mahl Date: Thu, 9 Jul 2026 19:23:49 -0700 Subject: [PATCH 06/36] guard plot_1d y_norm against constant traces y_norm=1 normalized each trace by its own range, which is zero for a constant trace, silently plotting all-NaN data. Map zero-range traces to baseline 0 (a constant trace has no amplitude to normalize) (code-review-2026-07, FAIL 4/4). --- src/trspecfit/utils/plot.py | 9 +++++++-- tests/test_plotting.py | 20 ++++++++++++++++++++ 2 files changed, 27 insertions(+), 2 deletions(-) diff --git a/src/trspecfit/utils/plot.py b/src/trspecfit/utils/plot.py index 7acb552..0c9926d 100644 --- a/src/trspecfit/utils/plot.py +++ b/src/trspecfit/utils/plot.py @@ -698,9 +698,14 @@ def plot_1d( raise ValueError("x axis could not be determined") y_data = data_series[i] - # Normalize if requested + # Normalize if requested; a constant trace has zero amplitude and + # maps to baseline 0 (dividing by its zero range would yield NaNs) if y_norm == 1: - y_plot = (y_data - np.min(y_data)) / (np.max(y_data - np.min(y_data))) + y_range = np.max(y_data) - np.min(y_data) + if y_range == 0: + y_plot = np.zeros_like(y_data, dtype=float) + else: + y_plot = (y_data - np.min(y_data)) / y_range y_plot = y_plot + i * waterfall else: y_plot = y_scale_arr[i] * y_data + i * waterfall diff --git a/tests/test_plotting.py b/tests/test_plotting.py index 769a4ba..2491fa5 100644 --- a/tests/test_plotting.py +++ b/tests/test_plotting.py @@ -135,6 +135,26 @@ def test_basic_plot(self): assert ax.get_xlabel() == "x axis" # default label plt.close("all") + # + def test_y_norm_constant_trace_stays_finite(self): + """y_norm=1 maps a constant trace to baseline 0 instead of NaN. + + Regression: normalization divided by the trace's own range, + which is zero for a constant trace (silent inf/NaN plot). + """ + + x = np.linspace(0, 10, 50) + y_list = [np.full(50, 5.0), np.sin(x)] + config = PlotConfig() + + plot_1d(y_list, x=x, config=config, y_norm=1, save_img=0) + ax = plt.gca() + y_const, y_sine = (line.get_ydata() for line in ax.get_lines()[:2]) + assert np.all(np.isfinite(y_const)) + np.testing.assert_array_equal(y_const, np.zeros(50)) + assert np.all(np.isfinite(y_sine)) + plt.close("all") + # def test_plot_with_custom_config(self): """Test 1D plot with custom config""" From 894772afd3955c10a829c7c5d8da715313ace231 Mon Sep 17 00:00:00 2001 From: Johannes Mahl Date: Thu, 9 Jul 2026 19:26:35 -0700 Subject: [PATCH 07/36] update claude commit settings --- CLAUDE.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CLAUDE.md b/CLAUDE.md index 2aa7aa6..e0bf3ac 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -1,6 +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. - **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 5ba8a1a67ed420147c49424573140867bf4ec543 Mon Sep 17 00:00:00 2001 From: Johannes Mahl Date: Thu, 9 Jul 2026 19:49:38 -0700 Subject: [PATCH 08/36] fix partial-range create_value_2d using slice-relative time indices Model.create_value_2d(t_ind=[start, stop]) passed the loop index to create_value_1d, which expects an absolute index into self.time. Any partial-range evaluation of a time-dependent model therefore computed the dynamics for t[0:stop-start] instead of t[start:stop]. Pass t_start + ti and add a regression test comparing a partial-range evaluation against the matching rows of the full evaluation. --- src/trspecfit/mcp.py | 4 +++- tests/test_mcp_eval.py | 26 ++++++++++++++++++++++++++ 2 files changed, 29 insertions(+), 1 deletion(-) diff --git a/src/trspecfit/mcp.py b/src/trspecfit/mcp.py index 95716ac..342ccf6 100644 --- a/src/trspecfit/mcp.py +++ b/src/trspecfit/mcp.py @@ -987,10 +987,12 @@ def create_value_2d(self, t_ind: list[int] | None = None) -> None: if self.time is None or self.energy is None: raise ValueError("Model time and energy axes required for 2D evaluation") + t_start = 0 if t_ind is None else t_ind[0] time_slice = self.time if t_ind is None else self.time[t_ind[0] : t_ind[1]] self.value_2d = np.empty((len(time_slice), len(self.energy))) for ti, _t in enumerate(time_slice): - val = self.create_value_1d(t_ind=ti, return_1d=1) + # create_value_1d expects an absolute index into self.time + val = self.create_value_1d(t_ind=t_start + ti, return_1d=1) if val is None: raise RuntimeError("create_value_1d returned None during 2D eval") self.value_2d[ti, :] = val diff --git a/tests/test_mcp_eval.py b/tests/test_mcp_eval.py index 15f5629..681226d 100644 --- a/tests/test_mcp_eval.py +++ b/tests/test_mcp_eval.py @@ -107,6 +107,32 @@ def test_eval_time_dependent_expression_value2d(self): v2_late = p_x0_2.value(t_ind=50) assert np.isclose(v2_late, v1_late + 3.6) + # + def test_eval_partial_time_range_uses_absolute_indices(self): + """Partial-range 2D eval matches the same rows of the full eval. + + Regression: create_value_2d(t_ind=[start, stop]) passed the + slice-relative loop index to create_value_1d, so time-dependent + parameters were evaluated at t[0:stop-start] instead of + t[start:stop]. + """ + + file, model = self._make_file_with_model(["energy_expression"]) + file.add_time_dependence( + target_model="energy_expression", + target_parameter="GLP_01_x0", + dynamics_yaml="models/file_time.yaml", + dynamics_model=["MonoExpPosIRF"], + ) + + model.create_value_2d() + value_2d_full = model.value_2d.copy() + + start, stop = 40, 60 + model.create_value_2d(t_ind=[start, stop]) + assert model.value_2d.shape == (stop - start, len(file.energy)) + np.testing.assert_allclose(model.value_2d, value_2d_full[start:stop]) + # def test_eval_expression_fan_out(self): """Fan-out: GLP_02 and GLP_03 both reference GLP_01 directly.""" From 2f50fe6a13490be09dd65870552000f109d1ef3e Mon Sep 17 00:00:00 2001 From: Johannes Mahl Date: Thu, 9 Jul 2026 20:04:21 -0700 Subject: [PATCH 09/36] raise instead of fabricating axes in describe/define_baseline/set_fit_limits These File methods assigned index axes to self.energy/self.time as a side effect when axes were missing. Since File(data=...) always creates index axes itself, data without axes means the object was corrupted by direct attribute assignment - raise a clear ValueError instead of silently persisting fabricated axes from an inspection or setup call. describe on a File with no data at all still warns and returns, since an empty File is a normal lifecycle state. --- src/trspecfit/trspecfit.py | 44 +++++++++++++++++----------------- tests/test_file.py | 49 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 71 insertions(+), 22 deletions(-) diff --git a/src/trspecfit/trspecfit.py b/src/trspecfit/trspecfit.py index f53c0ae..4250b72 100644 --- a/src/trspecfit/trspecfit.py +++ b/src/trspecfit/trspecfit.py @@ -1699,11 +1699,15 @@ def describe(self, *, waterfall: float | None = None) -> None: warnings.warn("No data loaded; nothing to describe.", stacklevel=2) return if self.energy is None: - self.energy = np.arange(self.data.shape[-1]) - warnings.warn("Energy axis missing; using index axis.", stacklevel=2) + raise ValueError( + "Energy axis missing; cannot describe data. " + "Pass energy= when constructing File." + ) if self.dim == 2 and self.time is None: - self.time = np.arange(self.data.shape[0]) - warnings.warn("Time axis missing; using index axis.", stacklevel=2) + raise ValueError( + "Time axis missing; cannot describe 2D data. " + "Pass time= when constructing File." + ) config = self.plot_config @@ -2310,14 +2314,15 @@ def define_baseline( if self.data is None: raise ValueError("No data loaded; cannot define baseline.") if self.time is None: - self.time = np.arange(self.data.shape[0]) - warnings.warn( - "Time axis missing; using index axis for baseline definition.", - stacklevel=2, + raise ValueError( + "Time axis missing; cannot define baseline. " + "Pass time= when constructing File." ) if self.energy is None: - self.energy = np.arange(self.data.shape[1]) - warnings.warn("Energy axis missing; using index axis.", stacklevel=2) + raise ValueError( + "Energy axis missing; cannot define baseline. " + "Pass energy= when constructing File." + ) self.base_t_ind = self._resolve_time_selection( time_start, time_stop, time_type=time_type ) @@ -2378,13 +2383,11 @@ def set_fit_limits( If True, plot data with fit limits indicated """ - if self.data is None and self.energy is None: - raise ValueError("No data/energy axis loaded; cannot set fit limits.") - if self.energy is None and self.data is not None: - self.energy = np.arange(self.data.shape[-1]) - warnings.warn("Energy axis missing; using index axis.", stacklevel=2) if self.energy is None: - raise ValueError("Energy axis unavailable; cannot set fit limits.") + raise ValueError( + "Energy axis missing; cannot set fit limits. " + "Pass energy= when constructing File." + ) energy = self.energy if energy_limits is None: energy_limits = [float(np.min(energy)), float(np.max(energy))] @@ -2410,12 +2413,9 @@ def set_fit_limits( time_limits = [float(np.min(self.time)), float(np.max(self.time))] if time_limits is not None: if self.time is None: - if self.data is None or self.dim != 2: - raise ValueError("Time axis missing; cannot apply time limits.") - self.time = np.arange(self.data.shape[0]) - warnings.warn( - "Time axis missing; using index axis for time limits.", - stacklevel=2, + raise ValueError( + "Time axis missing; cannot apply time limits. " + "Pass time= when constructing File." ) self.t_lim_abs = list(time_limits) self.t_lim = self._resolve_time_selection( diff --git a/tests/test_file.py b/tests/test_file.py index 27872a8..a9cf019 100644 --- a/tests/test_file.py +++ b/tests/test_file.py @@ -520,6 +520,55 @@ def test_define_baseline_invalid_time_type_raises(self): with pytest.raises(ValueError, match="Unknown time_type"): file.define_baseline(-10, 0, time_type="bogus", show_plot=False) + # + def _make_axisless_file_with_data(self, *, show_output: int = 0): + """Data present but axes missing — only reachable by bypassing __init__. + + File(data=...) always fabricates index axes, so this corrupted + state indicates direct attribute assignment. + """ + + project = make_project(show_output=show_output) + file = File(parent_project=project) + file.data = np.zeros((5, 7)) + file.dim = 2 + return file + + # + def test_describe_missing_axes_raises_without_mutating(self): + """describe raises on missing axes instead of fabricating them. + + Regression: describe assigned index axes to self.energy/self.time + as a side effect of an inspection call. + """ + + file = self._make_axisless_file_with_data(show_output=1) + with pytest.raises(ValueError, match="Energy axis missing"): + file.describe() + assert file.energy is None + file.energy = np.arange(7.0) + with pytest.raises(ValueError, match="Time axis missing"): + file.describe() + assert file.time is None + + # + def test_define_baseline_missing_time_axis_raises_without_mutating(self): + """define_baseline raises on a missing time axis instead of fabricating it.""" + + file = self._make_axisless_file_with_data() + with pytest.raises(ValueError, match="Time axis missing"): + file.define_baseline(0, 2, time_type="ind", show_plot=False) + assert file.time is None + + # + def test_set_fit_limits_missing_energy_axis_raises_without_mutating(self): + """set_fit_limits raises on a missing energy axis instead of fabricating it.""" + + file = self._make_axisless_file_with_data() + with pytest.raises(ValueError, match="Energy axis missing"): + file.set_fit_limits([1, 3], show_plot=False) + assert file.energy is None + # # From 158bf2cc870ad26f113d8f76f082e82c836ad305 Mon Sep 17 00:00:00 2001 From: Johannes Mahl Date: Thu, 9 Jul 2026 20:13:03 -0700 Subject: [PATCH 10/36] raise clear errors for convolution on a single-point time axis The kernel step size comes from time[1] - time[0], which raised a bare IndexError in the schedule_2d conv precompute (and would again in evaluate_2d) for a 1-point axis. Validate once at scheduling when conv steps are present. The mcp-layer guard in create_t_kernel already caught this at model construction but blamed an undefined time axis; give the too-short case its own message. Tests cover both layers. --- src/trspecfit/graph_ir.py | 5 +++++ src/trspecfit/mcp.py | 8 +++++++- tests/test_graph_ir.py | 16 ++++++++++++++++ tests/test_mcp_eval.py | 25 +++++++++++++++++++++++++ 4 files changed, 53 insertions(+), 1 deletion(-) diff --git a/src/trspecfit/graph_ir.py b/src/trspecfit/graph_ir.py index 7869906..d630ef6 100644 --- a/src/trspecfit/graph_ir.py +++ b/src/trspecfit/graph_ir.py @@ -2373,6 +2373,11 @@ def schedule_2d(graph: GraphIR) -> ScheduledPlan2D: ] _conv_id_to_idx: dict[int, int] = {n.id: i for i, n in enumerate(conv_nodes_topo)} n_conv_steps = len(conv_nodes_topo) + if n_conv_steps > 0 and n_time < 2: + raise ValueError( + "Convolution requires a time axis with at least 2 points " + f"to determine the kernel step size, got {n_time}." + ) conv_target_rows_list: list[int] = [] conv_func_ids_list: list[int] = [] diff --git a/src/trspecfit/mcp.py b/src/trspecfit/mcp.py index 342ccf6..3bfd07d 100644 --- a/src/trspecfit/mcp.py +++ b/src/trspecfit/mcp.py @@ -1649,8 +1649,14 @@ def create_t_kernel(self, *, par_values: list[Any] | None = None) -> np.ndarray: # full parameter list for multi-parameter kernels such as Voigt. kernel_width = getattr(fcts_time, self.fct_str + "_kernel_width")(*par_values) t_range = par_values[0] * kernel_width - if self.time is None or len(self.time) < 2: + if self.time is None: raise ValueError(f"time axis of component {self.fct_str} not defined") + if len(self.time) < 2: + raise ValueError( + f"Convolution component {self.fct_str} requires a time axis " + f"with at least 2 points to determine the kernel step size, " + f"got {len(self.time)}." + ) t_step = self.time[1] - self.time[0] return uarr.conv_kernel_support(t_range, t_step) diff --git a/tests/test_graph_ir.py b/tests/test_graph_ir.py index 4538999..6f6f1c9 100644 --- a/tests/test_graph_ir.py +++ b/tests/test_graph_ir.py @@ -1417,6 +1417,22 @@ def test_expfun_still_emitted_as_dynamics_trace(self): trace_fns = {n.function_name for n in trace_nodes} assert "expFun" in trace_fns + # + def test_single_point_time_axis_raises_at_schedule(self): + """schedule_2d raises a clear error for a conv on a 1-point time axis. + + Regression: the conv precompute read time[1] - time[0] and raised + a bare IndexError. Unreachable through the public API (model + construction guards the kernel axis), so exercised by mutating + graph.time directly as a GIR-layer invariant check. + """ + + _file, model = _make_irf_dynamics_model() + graph = build_graph(model) + graph.time = np.array([0.0]) + with pytest.raises(ValueError, match="at least 2 points"): + schedule_2d(graph) + # def test_irf_dynamics_lowerable_in_2d(self): """Time-domain IRF dynamics are lowerable on the 2D backend. diff --git a/tests/test_mcp_eval.py b/tests/test_mcp_eval.py index 681226d..a0215d2 100644 --- a/tests/test_mcp_eval.py +++ b/tests/test_mcp_eval.py @@ -133,6 +133,31 @@ def test_eval_partial_time_range_uses_absolute_indices(self): assert model.value_2d.shape == (stop - start, len(file.energy)) np.testing.assert_allclose(model.value_2d, value_2d_full[start:stop]) + # + def test_conv_on_single_point_time_axis_raises(self): + """IRF dynamics on a 1-point time axis raise a clear error. + + The convolution kernel step size comes from time[1] - time[0], + so the kernel axis cannot be built. The old message claimed the + time axis was "not defined" even when it existed with 1 point. + """ + + project = make_project() + file = File(parent_project=project) + file.energy = np.linspace(80, 90, 201) + file.time = np.array([0.0]) + file.load_model( + model_yaml="models/file_energy.yaml", + model_info=["energy_expression"], + ) + with pytest.raises(ValueError, match="at least 2 points"): + file.add_time_dependence( + target_model="energy_expression", + target_parameter="GLP_01_x0", + dynamics_yaml="models/file_time.yaml", + dynamics_model=["MonoExpPosIRF"], + ) + # def test_eval_expression_fan_out(self): """Fan-out: GLP_02 and GLP_03 both reference GLP_01 directly.""" From 10f065383135e249b876ef6823cf95813e1e72d7 Mon Sep 17 00:00:00 2001 From: Johannes Mahl Date: Thu, 9 Jul 2026 20:16:47 -0700 Subject: [PATCH 11/36] raise instead of returning -1.0 for t_vary Par without t_model Par.value printed a warning and returned -1.0 when t_vary was set but no t_model was attached, silently poisoning every spectrum evaluated from the corrupted parameter. Raise a RuntimeError naming the parameter instead, matching the adjacent inconsistent-state guards. --- src/trspecfit/mcp.py | 6 ++++-- tests/test_mcp_eval.py | 14 ++++++++++++++ 2 files changed, 18 insertions(+), 2 deletions(-) diff --git a/src/trspecfit/mcp.py b/src/trspecfit/mcp.py index 3bfd07d..4796cf1 100644 --- a/src/trspecfit/mcp.py +++ b/src/trspecfit/mcp.py @@ -2254,8 +2254,10 @@ def value( value = float(base[0] + self.t_model.value_1d[t_ind]) else: - value = -1.0 - print(f't_vary attribute of Par "{self.name}" is not valid') + raise RuntimeError( + f'Par "{self.name}" has t_vary set but no t_model attached; ' + "add time dependence via File.add_time_dependence." + ) return value diff --git a/tests/test_mcp_eval.py b/tests/test_mcp_eval.py index a0215d2..b27e4f6 100644 --- a/tests/test_mcp_eval.py +++ b/tests/test_mcp_eval.py @@ -158,6 +158,20 @@ def test_conv_on_single_point_time_axis_raises(self): dynamics_model=["MonoExpPosIRF"], ) + # + def test_t_vary_without_t_model_raises(self): + """Par.value raises on t_vary without a t_model. + + Regression: it printed a warning and returned -1.0, silently + poisoning any spectrum evaluated from the corrupted parameter. + """ + + _file, model = self._make_file_with_model(["energy_expression"]) + p_A = self._par(model, "GLP_01_A") + p_A.t_vary = True # corrupt state: no t_model attached + with pytest.raises(RuntimeError, match="no t_model"): + p_A.value(t_ind=0) + # def test_eval_expression_fan_out(self): """Fan-out: GLP_02 and GLP_03 both reference GLP_01 directly.""" From fa81a4d12301770a2de8ca702952b01f6ea6268d Mon Sep 17 00:00:00 2001 From: Johannes Mahl Date: Thu, 9 Jul 2026 20:51:13 -0700 Subject: [PATCH 12/36] validate fit-window data is finite at the fit_wrapper entry NaN/Inf in the data surfaced as lmfit's generic error blaming "input data or the output of your objective/model function", leaving the user to figure out which. Check the fit-window slice once in fit_wrapper - the choke point for all fit entry points - and raise a message with the non-finite count and a pointer to set_fit_limits(). Data outside the e_lim/t_lim window never reaches the residual and stays legal. Factor the residual_fun window slicing into a shared helper, and add test_fit_validation.py closing the check-17 gaps: NaN/Inf at the public fit level and single-element energy/time axes through the pipeline. --- src/trspecfit/fitlib.py | 59 ++++++----- tests/test_fit_validation.py | 196 +++++++++++++++++++++++++++++++++++ 2 files changed, 231 insertions(+), 24 deletions(-) create mode 100644 tests/test_fit_validation.py diff --git a/src/trspecfit/fitlib.py b/src/trspecfit/fitlib.py index 1cae54d..130e5df 100644 --- a/src/trspecfit/fitlib.py +++ b/src/trspecfit/fitlib.py @@ -149,6 +149,25 @@ def compute_fit_metrics( } +# +def _fit_window_slices( + ndim: int, e_lim: list[int] | None, t_lim: list[int] | None +) -> tuple[slice, ...]: + """Build array slices selecting the user-defined fit window. + + Empty or None limits select the full axis. 1D data is indexed as + [energy]; 2D data as [time, energy]. + """ + + e_slice = slice(e_lim[0], e_lim[1]) if e_lim else slice(None) + if ndim == 1: + return (e_slice,) + if ndim == 2: + t_slice = slice(t_lim[0], t_lim[1]) if t_lim else slice(None) + return (t_slice, e_slice) + raise ValueError("data must be 1D or 2D") + + # def residual_fun( par: Any, @@ -248,30 +267,8 @@ def residual_fun( data_arr = np.asarray(data) # select user-defined region to consider for residual computation - if len(data_arr.shape) == 1: # 1D data - if len(e_lim) != 0: - residual = data_arr[e_lim[0] : e_lim[1]] - fit_arr[e_lim[0] : e_lim[1]] - else: # use entire data and fit array to compute RSS - residual = data_arr - fit_arr - elif len(data_arr.shape) == 2: # 2D data - if (len(e_lim) != 0) and (len(t_lim) == 0): - residual = ( - data_arr[:, e_lim[0] : e_lim[1]] - fit_arr[:, e_lim[0] : e_lim[1]] - ) - elif (len(e_lim) == 0) and (len(t_lim) != 0): - residual = ( - data_arr[t_lim[0] : t_lim[1], :] - fit_arr[t_lim[0] : t_lim[1], :] - ) - elif (len(e_lim) != 0) and (len(t_lim) != 0): - residual = ( - data_arr[t_lim[0] : t_lim[1], e_lim[0] : e_lim[1]] - - fit_arr[t_lim[0] : t_lim[1], e_lim[0] : e_lim[1]] - ) - # or use entire data and fit array to compute RSS - else: - residual = data_arr - fit_arr - else: - raise ValueError("data must be 1D or 2D") + window = _fit_window_slices(data_arr.ndim, e_lim, t_lim) + residual = data_arr[window] - fit_arr[window] # type of residual to return if res_type == "RSS": @@ -654,6 +651,20 @@ def fit_wrapper( if stages not in (1, 2): raise ValueError(f"stages must be 1 or 2, got {stages}") + # Fail fast on NaN/Inf inside the fit window: lmfit raises a generic + # error that blames "input data or the objective/model function", + # leaving the user to figure out which. Non-finite data outside the + # e_lim/t_lim window never reaches the residual and stays legal. + data_arr = np.asarray(const[1], dtype=float) + window = data_arr[_fit_window_slices(data_arr.ndim, const[4], const[5])] + n_bad = int(np.size(window) - np.count_nonzero(np.isfinite(window))) + if n_bad > 0: + raise ValueError( + f"Data contains {n_bad} non-finite value(s) (NaN/Inf) inside " + "the fit window. Clean the data or exclude the affected " + "region with set_fit_limits() before fitting." + ) + # construct the lmfit parameters if necessary if isinstance(par, lmfit.parameter.Parameters): par_ini = copy.deepcopy(par) diff --git a/tests/test_fit_validation.py b/tests/test_fit_validation.py new file mode 100644 index 0000000..e8f7743 --- /dev/null +++ b/tests/test_fit_validation.py @@ -0,0 +1,196 @@ +"""Fit-entry validation: non-finite data and degenerate axes. + +Covers the public File fit entry points with NaN/Inf-contaminated data +(inside vs outside the fit window) and single-element energy/time axes. +The non-finite check lives at the fit_wrapper choke point, so one +entry point per dimensionality is exercised rather than every method. +""" + +import matplotlib + +matplotlib.use("Agg") + +import numpy as np +import pytest +from _utils import make_project, simulate_clean + +from trspecfit import File + + +# +def _simulate_truth(project): + """Simulate clean 2D data (single GLP + MonoExpPos on its amplitude).""" + + energy = np.linspace(83, 87, 30) + time = np.linspace(-2, 10, 24) + + file = File(parent_project=project, name="truth") + file.energy = energy + file.time = time + file.dim = 2 + file.load_model( + model_yaml="models/file_energy.yaml", + model_info="single_glp", + ) + file.add_time_dependence( + target_model="single_glp", + target_parameter="GLP_01_A", + dynamics_yaml="models/file_time.yaml", + dynamics_model=["MonoExpPos"], + ) + assert file.model_active is not None # type guard + return simulate_clean(file.model_active), energy, time + + +# +def _make_fit_file(project, data, energy, time): + """Create a file loaded with data and the 1D single_glp model.""" + + file = File( + parent_project=project, + name="fit", + data=data, + energy=energy.copy(), + time=time.copy(), + ) + file.load_model( + model_yaml="models/file_energy.yaml", + model_info="single_glp", + ) + return file + + +# +# +class TestNonFiniteData: + """NaN/Inf in data fail fast with a clear error at the fit entry.""" + + # + def test_nan_in_fit_window_raises(self): + """fit_spectrum on data with a NaN in the fit window raises. + + Regression guard on the message: lmfit's own error blames "input + data or the output of your objective/model function", leaving the + user to figure out which one is broken. + """ + + project = make_project(name="nan_window") + data, energy, time = _simulate_truth(project) + data[5, 10] = np.nan + + fit_file = _make_fit_file(project, data, energy, time) + with pytest.raises(ValueError, match="non-finite"): + fit_file.fit_spectrum( + "single_glp", time_point=float(time[5]), show_plot=False, try_ci=0 + ) + + # + def test_inf_in_fit_window_raises(self): + """fit_spectrum treats Inf the same as NaN.""" + + project = make_project(name="inf_window") + data, energy, time = _simulate_truth(project) + data[5, 10] = np.inf + + fit_file = _make_fit_file(project, data, energy, time) + with pytest.raises(ValueError, match="non-finite"): + fit_file.fit_spectrum( + "single_glp", time_point=float(time[5]), show_plot=False, try_ci=0 + ) + + # + def test_nan_outside_fit_window_is_allowed(self): + """A NaN excluded by set_fit_limits never reaches the residual. + + Cutting a contaminated detector region out of the fit window is a + supported workflow; validation must only inspect the window. + """ + + project = make_project(name="nan_outside") + data, energy, time = _simulate_truth(project) + data[5, 0] = np.nan # energy[0] = 83.0, below the limit chosen next + + fit_file = _make_fit_file(project, data, energy, time) + fit_file.set_fit_limits([83.5, 87.0], show_plot=False) + fit_file.fit_spectrum( + "single_glp", time_point=float(time[5]), show_plot=False, try_ci=0 + ) + + result = fit_file.model_spec.result + assert result[1] != [] + assert result[1].success + assert np.isfinite(result[1].chisqr) + + # + def test_nan_in_fit_window_raises_2d(self): + """fit_2d validates the (t_lim, e_lim) window of the 2D data.""" + + project = make_project(name="nan_2d") + data, energy, time = _simulate_truth(project) + data[15, 10] = np.nan + + fit_file = _make_fit_file(project, data, energy, time) + fit_file.define_baseline(-2, -1, show_plot=False) + fit_file.fit_baseline(model_name="single_glp", stages=1, 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"], + ) + with pytest.raises(ValueError, match="non-finite"): + fit_file.fit_2d("single_glp", stages=1, try_ci=0) + + +# +# +class TestDegenerateAxes: + """Single-element energy/time axes through the public fit pipeline.""" + + # + def test_single_element_energy_axis_fits(self): + """fit_spectrum on a 1-point energy axis completes. + + Documents current behavior: the pipeline tolerates the degenerate + axis and Nelder-Mead returns a result. With more free parameters + than data points the fit is underdetermined (negative nfree), and + judging that is left to the user. + """ + + project = make_project(name="one_energy") + data, energy, time = _simulate_truth(project) + + fit_file = _make_fit_file(project, data[:, :1], energy[:1], time) + fit_file.fit_spectrum( + "single_glp", time_point=float(time[5]), show_plot=False, try_ci=0 + ) + + result = fit_file.model_spec.result + assert result[1] != [] + assert result[1].nfree < 0 # underdetermined, documented not endorsed + + # + def test_single_element_time_axis_fit_2d(self): + """fit_2d on a 1-point time axis completes for conv-free dynamics. + + Convolution models raise at model construction (kernel step size + needs 2 points); plain dynamics stay evaluable at a single time. + """ + + project = make_project(name="one_time") + data, energy, time = _simulate_truth(project) + + fit_file = _make_fit_file(project, data[:1, :], energy, time[:1]) + fit_file.define_baseline(0, 0, time_type="ind", show_plot=False) + fit_file.fit_baseline(model_name="single_glp", stages=1, 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("single_glp", stages=1, try_ci=0) + + result = fit_file.model_2d.result + assert result[1] != [] + assert result[1].success From f0e5e66e704bc0e91e13301f1cba4b0a7f8cf7bd Mon Sep 17 00:00:00 2001 From: Johannes Mahl Date: Thu, 9 Jul 2026 21:37:17 -0700 Subject: [PATCH 13/36] honor silent mode in config errors, MCMC, and fit/setup display paths One themed pass over the silent-mode violations (review checks 3, 12): - _load_config: any error other than a missing file was swallowed with a print gated on show_output, so a broken config silently fell back to defaults; raise ValueError instead. - fit_wrapper MCMC: the emcee banner and progress bar ran unconditionally, and with show_output=0, save_output=0 the walker and corner figures reached plt.show() and were left open; gate the prints on show_output and close unsaved figures in silent mode. - fit_2d / fit_slice_by_slice: time_display and display(params) ran whenever stages >= 1; gate on show_output like fit_baseline and fit_spectrum. - define_baseline / set_fit_limits: plotted on their show_plot=True default without consulting Project.show_output; now suppressed in silent mode. --- src/trspecfit/fitlib.py | 19 ++++--- src/trspecfit/trspecfit.py | 27 +++++----- tests/test_auto_export.py | 45 +++++++++++++++++ tests/test_file.py | 101 +++++++++++++++++++++++++++++++++++-- 4 files changed, 168 insertions(+), 24 deletions(-) diff --git a/src/trspecfit/fitlib.py b/src/trspecfit/fitlib.py index 130e5df..465898a 100644 --- a/src/trspecfit/fitlib.py +++ b/src/trspecfit/fitlib.py @@ -756,11 +756,11 @@ def fit_wrapper( min=np.log(mc_settings.sigma_min), max=np.log(mc_settings.sigma_max), ) - # always print progress bar - print( - "\nProgress of lmfit.emcee confidence interval determination\n" - "(based on Markov chain Monte Carlo parameter space sampling):" - ) + if show_output >= 1: + print( + "\nProgress of lmfit.emcee confidence interval determination\n" + "(based on Markov chain Monte Carlo parameter space sampling):" + ) # burn necessary if starting point not close to max(probability distribution) # i.e. not close to the optimized parameter set, so burn=0 is ok here! emcee_fin = mini.emcee( @@ -772,7 +772,7 @@ def fit_wrapper( ntemps=mc_settings.ntemps, workers=mc_settings.workers, is_weighted=mc_settings.is_weighted, - progress=True, + progress=show_output >= 1, ) emcee_fin_params = _result_params(emcee_fin) emcee_flatchain = cast( @@ -789,7 +789,12 @@ def fit_wrapper( t_emcee1 = time.time() print(f"Time lmfit.emcee: {t_emcee1 - t_emcee0} s") # acceptence fraction of all walkers (plot) - emcee_save = save_output if show_output >= 1 else -abs(save_output) + # display per show_output, save per save_output (_finalize_plot + # semantics: >= 0 shows, abs == 1 saves, so -2 means neither) + if show_output >= 1: + emcee_save = save_output + else: + emcee_save = -1 if abs(save_output) == 1 else -2 fig_emcee_walker, _ax = plt.subplots(1, 1, dpi=75) plt.plot(emcee_acceptance_fraction, "o") plt.xlabel("Walker number") diff --git a/src/trspecfit/trspecfit.py b/src/trspecfit/trspecfit.py index 4250b72..0ee970f 100644 --- a/src/trspecfit/trspecfit.py +++ b/src/trspecfit/trspecfit.py @@ -820,12 +820,11 @@ def _load_config(self, config_file: PathLike) -> None: self._config_file = config_path except FileNotFoundError: + # missing config is a designed fallback: all settings optional if self.show_output >= 1: print(f"Config file {config_path} not found, using defaults") except Exception as e: # noqa: BLE001 - if self.show_output >= 1: - print(f"Error loading config: {e}") - print("Using default settings") + raise ValueError(f"Failed to load config file {config_path}: {e}") from e # ------------------------------------------------------------------ # Project-level model loading and baseline fitting @@ -2306,7 +2305,8 @@ def define_baseline( - 'ind': Time array indices show_plot : bool, default=True - If True, plot the resulting baseline spectrum + If True, plot the resulting baseline spectrum. Suppressed + when ``Project.show_output < 1``. """ if self.dim == 1: @@ -2337,7 +2337,7 @@ def define_baseline( ) # plot - if show_plot: + if show_plot and self.p.show_output >= 1: if self.data_base is None: warnings.warn( "Baseline data is unavailable; skipping baseline plot.", @@ -2380,7 +2380,8 @@ def set_fit_limits( Time range for fitting ``[min, max]`` in absolute values. If None, no time limits are applied. show_plot : bool, default=True - If True, plot data with fit limits indicated + If True, plot data with fit limits indicated. Suppressed + when ``Project.show_output < 1``. """ if self.energy is None: @@ -2422,7 +2423,7 @@ def set_fit_limits( float(np.min(time_limits)), float(np.max(time_limits)) ) - if show_plot: # show data with limits + if show_plot and self.p.show_output >= 1: # show data with limits if self.dim == 1: if self.data is None: warnings.warn("Data missing; cannot plot fit limits.", stacklevel=2) @@ -3319,7 +3320,7 @@ def _slice_path(s_i: int) -> pathlib.Path: ) self.model_sbs.update_value(new_par_values=seed_template, par_select="all") self.model_sbs.args = _args_sbs - if stages >= 1: + if stages >= 1 and self.p.show_output >= 1: fitlib.time_display( t_start=t_sbs, print_str="Time elapsed for Slice-by-Slice fit: " ) @@ -4058,10 +4059,12 @@ def fit_2d(self, model_name: str, stages: int = 1, **fit_wrapper_kwargs) -> None self._save_2d_fit_legacy( save_path=path_2d_results, save_files=self.p.auto_export ) - fitlib.time_display( - t_start=t_2d, print_str="Time elapsed for 2D model fit: " - ) - display(self.model_2d.result[1].params) # display final pars below figure + if self.p.show_output >= 1: + fitlib.time_display( + t_start=t_2d, print_str="Time elapsed for 2D model fit: " + ) + # display final pars below figure + display(self.model_2d.result[1].params) # def _save_2d_fit_legacy( diff --git a/tests/test_auto_export.py b/tests/test_auto_export.py index bad6041..421cab0 100644 --- a/tests/test_auto_export.py +++ b/tests/test_auto_export.py @@ -12,10 +12,13 @@ matplotlib.use("Agg") +import matplotlib.pyplot as plt import numpy as np +import pytest from _utils import make_project, simulate_noisy from trspecfit import File, fitlib +from trspecfit.utils.lmfit import MC # @@ -207,6 +210,48 @@ def test_baseline_plots_when_verbose_even_without_export( # disk write was suppressed. assert mock.call_count == 1 + # + def test_fit_2d_silent_mode_prints_nothing(self, tmp_path, capsys): + """fit_2d honors show_output=0: no timing line, no params display. + + Regression: time_display and display(params) ran whenever + stages >= 1, regardless of show_output. + """ + + project, file = _baseline_setup(tmp_path, auto_export=False) + file.fit_baseline(model_name="single_glp", stages=1, try_ci=0) + file.add_time_dependence( + target_model="single_glp", + target_parameter="GLP_01_A", + dynamics_yaml="models/file_time.yaml", + dynamics_model=["MonoExpPos"], + ) + capsys.readouterr() # drop setup output + file.fit_2d("single_glp", stages=1, try_ci=0) + assert capsys.readouterr().out == "" + + # + @pytest.mark.slow + def test_mcmc_silent_mode_no_output_no_figures(self, tmp_path, capsys): + """MCMC honors silent mode: no progress banner, no shown figures. + + Regression: the emcee progress banner and progress=True ran + unconditionally, and with show_output=0, save_output=0 the walker + and corner figures reached _finalize_plot(0), i.e. plt.show(), + and were left open. + """ + + project, file = _baseline_setup(tmp_path, auto_export=False) + mc = MC(use_mc=1, steps=20, nwalkers=32, burn=5, thin=1) + n_figs = len(plt.get_fignums()) + capsys.readouterr() # drop setup output + file.fit_baseline(model_name="single_glp", stages=1, try_ci=0, mc_settings=mc) + + # emcee prints its own short-chain autocorrelation notice for the + # deliberately tiny chain; only our banner is under test here + assert "Progress of lmfit.emcee" not in capsys.readouterr().out + assert len(plt.get_fignums()) == n_figs + # def test_sbs_skips_per_slice_plot_when_no_export(self, tmp_path, monkeypatch): mock = MagicMock() diff --git a/tests/test_file.py b/tests/test_file.py index a9cf019..25c6e22 100644 --- a/tests/test_file.py +++ b/tests/test_file.py @@ -359,10 +359,10 @@ class TestFitLimitsAndBaseline: """Test fit limits and baseline.""" # - def _make_file_with_data(self): + def _make_file_with_data(self, *, show_output: int = 0): """Create file with axes and 2D data.""" - project = make_project() + project = make_project(show_output=show_output) file = File(parent_project=project) file.energy = np.linspace(80, 90, 201) file.time = np.linspace(-10, 100, 111) @@ -569,6 +569,58 @@ def test_set_fit_limits_missing_energy_axis_raises_without_mutating(self): file.set_fit_limits([1, 3], show_plot=False) assert file.energy is None + # + def test_define_baseline_plot_suppressed_when_silent(self, monkeypatch): + """show_plot=True defers to Project.show_output (silent = no plot). + + Regression: define_baseline plotted on its show_plot=True default + even with show_output=0. + """ + + import trspecfit.utils.plot as uplt + + mock = unittest.mock.MagicMock() + monkeypatch.setattr(uplt, "plot_1d", mock) + file = self._make_file_with_data() # make_project() is silent + file.define_baseline(-10, 0, show_plot=True) + assert mock.call_count == 0 + + # + def test_define_baseline_plots_when_verbose(self, monkeypatch): + """define_baseline still plots with show_plot=True and show_output=1.""" + + import trspecfit.utils.plot as uplt + + mock = unittest.mock.MagicMock() + monkeypatch.setattr(uplt, "plot_1d", mock) + file = self._make_file_with_data(show_output=1) + file.define_baseline(-10, 0, show_plot=True) + assert mock.call_count == 1 + + # + def test_set_fit_limits_plot_suppressed_when_silent(self, monkeypatch): + """show_plot=True defers to Project.show_output (silent = no plot).""" + + import trspecfit.utils.plot as uplt + + mock = unittest.mock.MagicMock() + monkeypatch.setattr(uplt, "plot_2d", mock) + file = self._make_file_with_data() # make_project() is silent + file.set_fit_limits([82, 88], show_plot=True) + assert mock.call_count == 0 + + # + def test_set_fit_limits_plots_when_verbose(self, monkeypatch): + """set_fit_limits still plots with show_plot=True and show_output=1.""" + + import trspecfit.utils.plot as uplt + + mock = unittest.mock.MagicMock() + monkeypatch.setattr(uplt, "plot_2d", mock) + file = self._make_file_with_data(show_output=1) + file.set_fit_limits([82, 88], show_plot=True) + assert mock.call_count == 1 + # # @@ -1577,9 +1629,48 @@ def test_project_yaml_default_nan_roundtrips(self, tmp_path): project = Project(path=tmp_path, config_file="project.yaml") assert np.isnan(project.sigma_data) assert project.noise_type == NOISE_TYPE_UNKNOWN - # Proves _load_config didn't bail early — show_output overrides - # the default of 1. - assert project.show_output == 0 + + +# +# +class TestProjectConfigLoading: + """Project config-file loading error behavior.""" + + # + def test_malformed_config_raises(self, tmp_path): + """A config file that fails to parse raises instead of using defaults. + + Regression: any error after FileNotFoundError was swallowed with a + print gated on show_output, so with show_output=0 a broken config + silently fell back to defaults. + """ + + from trspecfit import Project + + config = tmp_path / "project.yaml" + config.write_text("show_output: [unclosed\n") + with pytest.raises(ValueError, match="Failed to load config"): + Project(path=tmp_path, config_file="project.yaml") + + # + def test_missing_config_falls_back_to_defaults(self, tmp_path): + """A missing config file falls back to defaults (designed behavior).""" + + from trspecfit import Project + + project = Project(path=tmp_path, config_file="does_not_exist.yaml") + assert project.show_output == 1 # default intact + + # + def test_valid_config_applies(self, tmp_path): + """A valid config file overrides the defaults.""" + + from trspecfit import Project + + config = tmp_path / "project.yaml" + config.write_text("show_output: 0\n") + project = Project(path=tmp_path, config_file="project.yaml") + assert project.show_output == 0 # override applied, not default 1 if __name__ == "__main__": From b139601e93bc8ea874da52e5e9ce42a331a428a1 Mon Sep 17 00:00:00 2001 From: Johannes Mahl Date: Fri, 10 Jul 2026 07:46:04 -0700 Subject: [PATCH 14/36] speed up 2D evaluator and my_conv hot paths - eval_expr_program: push scalar constants and trace-row views instead of allocating per-instruction arrays; broadcast constant-only results - profile sample/expr evaluation: write into preallocated buffers instead of broadcast_to().copy() and np.repeat temporaries - profiled ops: vectorize over the aux axis instead of a per-aux Python loop (~4.5x faster on a profiled 2D model) - my_conv: pad y directly (the padded x grid was built and discarded) and normalize the kernel instead of the padded signal --- src/trspecfit/eval_2d.py | 94 ++++++++++++++++++----------------- src/trspecfit/graph_ir.py | 1 - src/trspecfit/utils/arrays.py | 16 +++--- tests/test_graph_ir.py | 3 ++ 4 files changed, 59 insertions(+), 55 deletions(-) diff --git a/src/trspecfit/eval_2d.py b/src/trspecfit/eval_2d.py index c137fe6..b314cb9 100644 --- a/src/trspecfit/eval_2d.py +++ b/src/trspecfit/eval_2d.py @@ -78,8 +78,10 @@ def eval_expr_program( """Evaluate an RPN ExprProgram against the trace matrix. Works for both plan initialization and hot-path evaluation. - Each PARAM_REF reads a full ``(n_time,)`` row from *traces*; - constants are broadcast to ``(n_time,)`` via ``np.full``. + Each PARAM_REF pushes a *view* of its ``(n_time,)`` trace row and + constants stay scalar; every operator allocates a fresh array, so + the views are never written to. Callers must not mutate the result + in place (it may alias a *traces* row). Parameters ---------- @@ -95,20 +97,19 @@ def eval_expr_program( """ n_time = traces.shape[1] - stack: list[np.ndarray] = [] + stack: list[np.ndarray | np.float64] = [] instr = program.instructions n_instr = len(instr) // 2 for i in range(n_instr): - kind = ExprNodeKind(instr[2 * i]) + kind = int(instr[2 * i]) operand = instr[2 * i + 1] if kind == ExprNodeKind.CONST: - val = np.int64(operand).view(np.float64) - stack.append(np.full(n_time, val, dtype=np.float64)) + stack.append(np.int64(operand).view(np.float64)) elif kind == ExprNodeKind.PARAM_REF: - stack.append(traces[int(operand), :].copy()) + stack.append(traces[int(operand), :]) elif kind == ExprNodeKind.ADD: b, a = stack.pop(), stack.pop() @@ -134,7 +135,10 @@ def eval_expr_program( stack.append(a**b) assert len(stack) == 1 - return stack[0] + result = stack[0] + if not isinstance(result, np.ndarray): # constant-only program + return np.full(n_time, float(result), dtype=np.float64) + return result # --------------------------------------------------------------------------- @@ -170,10 +174,9 @@ def _evaluate_profile_sample_values_2d( for group_idx in range(n_groups): base_row = int(profile_sample_base_rows[group_idx]) - # base trace -> (n_time, 1) -> broadcast to (n_time, n_aux) - values = np.broadcast_to( - traces[base_row, :][:, np.newaxis], (n_time, n_aux) - ).copy() + # base trace -> (n_time, 1) broadcast into the output row + values = sample_values[group_idx] + values[:] = traces[base_row, :][:, np.newaxis] comp_start = int(profile_sample_component_indptr[group_idx]) comp_end = int(profile_sample_component_indptr[group_idx + 1]) @@ -187,8 +190,6 @@ def _evaluate_profile_sample_values_2d( ] values += np.asarray(func(aux_2d, *params), dtype=np.float64) - sample_values[group_idx] = values - return sample_values @@ -219,7 +220,9 @@ def _evaluate_profile_expr_values_2d( # Virtual trace: regular params repeated across aux, profile samples # flattened from (n_groups, n_time, n_aux) -> (n_groups, n_time*n_aux). virtual = np.empty((n_params + n_groups, n_cols), dtype=np.float64) - virtual[:n_params, :] = np.repeat(traces, n_aux, axis=1) + # broadcast-write params across aux in place (no np.repeat temporary) + virtual_params = virtual[:n_params, :].reshape(n_params, n_time, n_aux) + virtual_params[:] = traces[:, :, np.newaxis] if n_groups > 0: virtual[n_params:, :] = profile_sample_values.reshape(n_groups, n_cols) @@ -243,41 +246,41 @@ def _evaluate_profiled_op_2d( peak_sum: np.ndarray, *, needs_spectrum: bool, - n_aux: int, ) -> np.ndarray: - """Evaluate one profiled 2D op: loop over aux points, average.""" + """Evaluate one profiled 2D op vectorized over aux, then average. - func, _needs = OP_DISPATCH[kind] - n_time = traces.shape[1] - n_energy = energy.shape[-1] - accumulated = np.zeros((n_time, n_energy), dtype=np.float64) - - for aux_i in range(n_aux): - params: list[np.ndarray] = [] - for source_kind, source_idx in zip( - param_source_kinds, - param_indices, - strict=True, - ): - sk = int(source_kind) - si = int(source_idx) - if sk == int(ParamSourceKind.SCALAR): - param = traces[si, :][:, np.newaxis] # (n_time, 1) - elif sk == int(ParamSourceKind.PROFILE_SAMPLE): - param = profile_sample_values[si, :, aux_i][ - :, np.newaxis - ] # (n_time, 1) - else: - param = profile_expr_values[si, :, aux_i][:, np.newaxis] # (n_time, 1) - params.append(param) + Broadcasts to ``(n_time, n_aux, n_energy)`` in a single call: energy + is ``(1, 1, n_energy)``, scalar params ``(n_time, 1, 1)``, and + profiled params ``(n_time, n_aux, 1)``. All energy functions are + pure ufunc arithmetic or reduce along ``axis=-1`` (Shirley), so the + extra leading axis broadcasts through unchanged. + """ - if needs_spectrum: - accumulated += func(energy, *params, peak_sum) + func, _needs = OP_DISPATCH[kind] + params: list[np.ndarray] = [] + for source_kind, source_idx in zip( + param_source_kinds, + param_indices, + strict=True, + ): + sk = int(source_kind) + si = int(source_idx) + if sk == int(ParamSourceKind.SCALAR): + param = traces[si, :][:, np.newaxis, np.newaxis] # (n_time, 1, 1) + elif sk == int(ParamSourceKind.PROFILE_SAMPLE): + param = profile_sample_values[si][:, :, np.newaxis] # (n_time, n_aux, 1) else: - accumulated += func(energy, *params) + param = profile_expr_values[si][:, :, np.newaxis] # (n_time, n_aux, 1) + params.append(param) + + energy_3d = energy[np.newaxis, :, :] if energy.ndim == 2 else energy + if needs_spectrum: + stacked = func(energy_3d, *params, peak_sum[:, np.newaxis, :]) + else: + stacked = func(energy_3d, *params) - accumulated /= n_aux - return accumulated + averaged: np.ndarray = np.asarray(stacked, dtype=np.float64).mean(axis=1) + return averaged # --------------------------------------------------------------------------- @@ -406,7 +409,6 @@ def evaluate_2d(plan: ScheduledPlan2D, theta: np.ndarray) -> np.ndarray: profile_expr_values, peak_sum, needs_spectrum=needs_spectrum, - n_aux=plan.n_aux, ) else: param_rows = plan.op_param_indices[start:end] diff --git a/src/trspecfit/graph_ir.py b/src/trspecfit/graph_ir.py index d630ef6..099b34c 100644 --- a/src/trspecfit/graph_ir.py +++ b/src/trspecfit/graph_ir.py @@ -2972,7 +2972,6 @@ def schedule_2d(graph: GraphIR) -> ScheduledPlan2D: profile_expr_values_init, cached_peak_sum, needs_spectrum=bool(op_needs_spectrum[op_idx]), - n_aux=n_aux, ) else: param_rows = op_param_indices[start:end] diff --git a/src/trspecfit/utils/arrays.py b/src/trspecfit/utils/arrays.py index a18d562..ca68c3d 100644 --- a/src/trspecfit/utils/arrays.py +++ b/src/trspecfit/utils/arrays.py @@ -386,19 +386,19 @@ def my_conv( ) pad_size = int(kernel_arr.size / 2) - # Add padding to minimize edge artifacts - _x_pad, y_pad = pad_x_y(x_arr, y_arr, float(x_arr[1] - x_arr[0]), pad_size) + # Add padding to minimize edge artifacts (y only; the x grid is not + # needed for the convolution itself) + y_pad = np.pad(y_arr, pad_size, mode="edge") + + # Normalize the kernel (cheaper than dividing the padded signal) + kernel_norm = kernel_arr / np.sum(kernel_arr) # Compute convolution with normalized kernel if method == "scipy": - y_conv_pad = np.asarray( - convolve(y_pad, kernel_arr, mode="same") / np.sum(kernel_arr), - dtype=float, - ) + y_conv_pad = np.asarray(convolve(y_pad, kernel_norm, mode="same"), dtype=float) elif method == "numpy": y_conv_pad = np.asarray( - np.convolve(y_pad, kernel_arr, mode="same") / np.sum(kernel_arr), - dtype=float, + np.convolve(y_pad, kernel_norm, mode="same"), dtype=float ) else: raise ValueError(f"Unknown method '{method}'") diff --git a/tests/test_graph_ir.py b/tests/test_graph_ir.py index 6f6f1c9..163c886 100644 --- a/tests/test_graph_ir.py +++ b/tests/test_graph_ir.py @@ -2434,6 +2434,9 @@ def test_matches_asteval(self, expr_string, variables): program = _bind_expr_to_rows(symbolic, name_to_row) result = eval_expr_program(program, traces) + # Always a full (n_time,) row, even for constant-only programs + # (callers assign or reshape the result, so a scalar would break) + assert np.shape(result) == (n_time,) # All time steps should match the scalar asteval result assert np.allclose(result, expected, rtol=1e-12), ( f"Expression {expr_string!r}: RPN gave {result[0]}, asteval gave {expected}" From 7030f3cd3777d9e98a5a6b6b0b1982155a362455 Mon Sep 17 00:00:00 2001 From: Johannes Mahl Date: Fri, 10 Jul 2026 09:51:20 -0700 Subject: [PATCH 15/36] fix benchmark harness to exercise par profiles in example 04 - parse notebook calls parenthesis-matched instead of per cell: the multi-call cell in example 04 replayed add_time_dependence with the wrong target_parameter (GLP_01_A instead of GLP_01_x0_pLinear_01_m) - replay add_par_profile calls (before dynamics) and load data/aux_axis.csv; the profile example previously compiled with plan.n_aux == 0, never reaching the profiled-op path - filter replayed calls to the benchmarked "2D" model and list attached profiles in the preamble --- .claude/skills/benchmark/benchmark_gir.py | 111 +++++++++++++++++----- 1 file changed, 85 insertions(+), 26 deletions(-) diff --git a/.claude/skills/benchmark/benchmark_gir.py b/.claude/skills/benchmark/benchmark_gir.py index 9239341..3ec2749 100644 --- a/.claude/skills/benchmark/benchmark_gir.py +++ b/.claude/skills/benchmark/benchmark_gir.py @@ -74,39 +74,79 @@ def _extract_float(src, key): # -def _parse_dynamics_from_notebook(notebook_path): - """Extract add_time_dependence kwargs from example.ipynb. +def _iter_call_args(src, func_name): + """Yield the argument text of each ``func_name(...)`` call in *src*. + + Matches parentheses so multiple calls per cell stay separate (a + cell-level regex would mix kwargs of neighboring calls). + """ + + start = 0 + while True: + idx = src.find(func_name + "(", start) + if idx == -1: + return + depth = 0 + for j in range(idx + len(func_name), len(src)): + if src[j] == "(": + depth += 1 + elif src[j] == ")": + depth -= 1 + if depth == 0: + yield src[idx + len(func_name) + 1 : j] + start = j + 1 + break + else: + return + + +# +def _parse_calls_from_notebook(notebook_path): + """Extract add_par_profile / add_time_dependence kwargs from example.ipynb. Returns ------- - list of dict - Each dict has keys: target_model, target_parameter, + profile_calls : list of dict + add_par_profile kwargs: target_model, target_parameter, + profile_yaml, profile_model. + dynamics_calls : list of dict + add_time_dependence kwargs: target_model, target_parameter, dynamics_yaml, dynamics_model, and optionally frequency. """ with notebook_path.open() as f: nb = json.load(f) - calls = [] + profile_calls = [] + dynamics_calls = [] for cell in nb["cells"]: if cell["cell_type"] != "code": continue src = "".join(cell["source"]) - if "add_time_dependence" not in src: - continue - call = { - "target_model": _extract_str(src, "target_model"), - "target_parameter": _extract_str(src, "target_parameter"), - "dynamics_yaml": _extract_str(src, "dynamics_yaml"), - "dynamics_model": _extract_str_or_list(src, "dynamics_model"), - } - freq = _extract_float(src, "frequency") - if freq is not None: - call["frequency"] = freq - calls.append(call) + for args in _iter_call_args(src, "add_par_profile"): + profile_calls.append( + { + "target_model": _extract_str(args, "target_model"), + "target_parameter": _extract_str(args, "target_parameter"), + "profile_yaml": _extract_str(args, "profile_yaml"), + "profile_model": _extract_str_or_list(args, "profile_model"), + } + ) + + for args in _iter_call_args(src, "add_time_dependence"): + call = { + "target_model": _extract_str(args, "target_model"), + "target_parameter": _extract_str(args, "target_parameter"), + "dynamics_yaml": _extract_str(args, "dynamics_yaml"), + "dynamics_model": _extract_str_or_list(args, "dynamics_model"), + } + freq = _extract_float(args, "frequency") + if freq is not None: + call["frequency"] = freq + dynamics_calls.append(call) - return calls + return profile_calls, dynamics_calls # ------------------------------------------------------------------ @@ -140,7 +180,10 @@ def load_example(example_num, *, add_dynamics=True): ------- file : File dynamics_calls : list of dict - Parsed add_time_dependence kwargs from the notebook. + Parsed add_time_dependence kwargs from the notebook ("2D" model + only). + profile_calls : list of dict + Parsed add_par_profile kwargs already attached to the model. """ folder = _find_example_folder(example_num) @@ -152,6 +195,8 @@ def load_example(example_num, *, add_dynamics=True): energy = np.loadtxt(data_dir / "energy.csv") time_ax = np.loadtxt(data_dir / "time.csv") data = np.loadtxt(data_dir / "data.csv", delimiter=",") + aux_path = data_dir / "aux_axis.csv" + aux_axis = np.loadtxt(aux_path) if aux_path.exists() else None file = File( parent_project=project, @@ -159,16 +204,26 @@ def load_example(example_num, *, add_dynamics=True): data=data, energy=energy, time=time_ax, + aux_axis=aux_axis, ) file.load_model(model_yaml="models_energy.yaml", model_info="2D") - dynamics_calls = _parse_dynamics_from_notebook(folder / "example.ipynb") + profile_calls, dynamics_calls = _parse_calls_from_notebook(folder / "example.ipynb") + # The notebooks also attach to their baseline/SbS models; only calls + # targeting the benchmarked "2D" model apply here. + profile_calls = [c for c in profile_calls if c["target_model"] == "2D"] + dynamics_calls = [c for c in dynamics_calls if c["target_model"] == "2D"] + + # Profiles are model structure, not time dependence: attach always, + # and before dynamics (dynamics may target a profile parameter). + for call in profile_calls: + file.add_par_profile(**call) if add_dynamics: for call in dynamics_calls: file.add_time_dependence(**call) - return file, dynamics_calls + return file, dynamics_calls, profile_calls # @@ -333,7 +388,7 @@ def _snapshot(label): print(f" {label:32s}{call_count[0]:6d} (+{call_count[0] - prev})") try: - file, dynamics_calls = load_example(example_num, add_dynamics=False) + file, dynamics_calls, _ = load_example(example_num, add_dynamics=False) file.define_baseline( time_start=0, time_stop=10, time_type="ind", show_plot=False ) @@ -397,7 +452,7 @@ def timed_sched(*args, **kwargs): graph_ir.schedule_2d = timed_sched try: - file, dynamics_calls = load_example(example_num, add_dynamics=False) + file, dynamics_calls, _ = load_example(example_num, add_dynamics=False) file.define_baseline( time_start=0, time_stop=10, time_type="ind", show_plot=False ) @@ -454,7 +509,7 @@ def capture_par_variability(example_num, *, n_starts=4): redchis = [] for run in range(n_starts + 1): - file, dynamics_calls = load_example(example_num, add_dynamics=False) + file, dynamics_calls, _ = load_example(example_num, add_dynamics=False) file.define_baseline( time_start=0, time_stop=10, time_type="ind", show_plot=False ) @@ -563,7 +618,7 @@ def bench_fit(example_num, dynamics_calls, *, n_reps=3): ("fit_model_gir", gir_times), ("fit_model_mcp", mcp_times), ]: - file, _ = load_example(example_num, add_dynamics=False) + file, _, _ = load_example(example_num, add_dynamics=False) file.define_baseline( time_start=0, time_stop=10, time_type="ind", show_plot=False ) @@ -682,11 +737,15 @@ def bench_fit(example_num, dynamics_calls, *, n_reps=3): folder = _find_example_folder(args.example) print(f"Example: {folder.name}") - file, dynamics_calls = load_example(args.example) + file, dynamics_calls, profile_attaches = load_example(args.example) model = file.model_active assert model is not None graph = build_graph(model) print(f" lowerable: {can_lower_2d(graph)}") + for call in profile_attaches: + target_parameter = call["target_parameter"] + profile_model = call["profile_model"] + print(f" profile: {target_parameter} <- {profile_model}") for call in dynamics_calls: target_parameter = call["target_parameter"] dynamics_model = call["dynamics_model"] From 7b43441fdc1d35e34855d45ac8ecd96a2ab76373 Mon Sep 17 00:00:00 2001 From: Johannes Mahl Date: Fri, 10 Jul 2026 09:53:39 -0700 Subject: [PATCH 16/36] return profiled-op evaluation to a per-aux loop with hoisted sources The aux vectorization only wins when profiled params enter the energy function linearly (amplitude-only profiles keep the transcendentals at (n_time, 1, n_energy)); with a profiled position it materializes full (n_time, n_aux, n_energy) temporaries and measures ~60% slower on the real example 04 workload (39 -> 64 ms/call), exposed by the fixed benchmark harness. Param sources are still resolved once outside the loop. --- src/trspecfit/eval_2d.py | 51 ++++++++++++++++++++++++--------------- src/trspecfit/graph_ir.py | 1 + 2 files changed, 32 insertions(+), 20 deletions(-) diff --git a/src/trspecfit/eval_2d.py b/src/trspecfit/eval_2d.py index b314cb9..1dcf428 100644 --- a/src/trspecfit/eval_2d.py +++ b/src/trspecfit/eval_2d.py @@ -246,18 +246,26 @@ def _evaluate_profiled_op_2d( peak_sum: np.ndarray, *, needs_spectrum: bool, + n_aux: int, ) -> np.ndarray: - """Evaluate one profiled 2D op vectorized over aux, then average. - - Broadcasts to ``(n_time, n_aux, n_energy)`` in a single call: energy - is ``(1, 1, n_energy)``, scalar params ``(n_time, 1, 1)``, and - profiled params ``(n_time, n_aux, 1)``. All energy functions are - pure ufunc arithmetic or reduce along ``axis=-1`` (Shirley), so the - extra leading axis broadcasts through unchanged. + """Evaluate one profiled 2D op: loop over aux points, average. + + Param sources are resolved to ``(n_time, n_aux)`` views once, + outside the loop. The per-aux loop is deliberate: vectorizing over + aux in a single call only wins when profiled params enter the + function linearly (amplitude-only profiles, where broadcasting + keeps the transcendental part at ``(n_time, 1, n_energy)``); with a + profiled position or width the energy function materializes full + ``(n_time, n_aux, n_energy)`` temporaries and measures ~60% slower + (example 04: profiled x0, n_aux=50, 175x280 grid). """ func, _needs = OP_DISPATCH[kind] - params: list[np.ndarray] = [] + n_time = traces.shape[1] + n_energy = energy.shape[-1] + + # Resolve each param source once (scalars as no-copy broadcast views) + sources: list[np.ndarray] = [] for source_kind, source_idx in zip( param_source_kinds, param_indices, @@ -266,21 +274,23 @@ def _evaluate_profiled_op_2d( sk = int(source_kind) si = int(source_idx) if sk == int(ParamSourceKind.SCALAR): - param = traces[si, :][:, np.newaxis, np.newaxis] # (n_time, 1, 1) + source = np.broadcast_to(traces[si, :][:, np.newaxis], (n_time, n_aux)) elif sk == int(ParamSourceKind.PROFILE_SAMPLE): - param = profile_sample_values[si][:, :, np.newaxis] # (n_time, n_aux, 1) + source = profile_sample_values[si] # (n_time, n_aux) else: - param = profile_expr_values[si][:, :, np.newaxis] # (n_time, n_aux, 1) - params.append(param) - - energy_3d = energy[np.newaxis, :, :] if energy.ndim == 2 else energy - if needs_spectrum: - stacked = func(energy_3d, *params, peak_sum[:, np.newaxis, :]) - else: - stacked = func(energy_3d, *params) + source = profile_expr_values[si] # (n_time, n_aux) + sources.append(source) + + accumulated = np.zeros((n_time, n_energy), dtype=np.float64) + for aux_i in range(n_aux): + params = [s[:, aux_i, np.newaxis] for s in sources] + if needs_spectrum: + accumulated += func(energy, *params, peak_sum) + else: + accumulated += func(energy, *params) - averaged: np.ndarray = np.asarray(stacked, dtype=np.float64).mean(axis=1) - return averaged + accumulated /= n_aux + return accumulated # --------------------------------------------------------------------------- @@ -409,6 +419,7 @@ def evaluate_2d(plan: ScheduledPlan2D, theta: np.ndarray) -> np.ndarray: profile_expr_values, peak_sum, needs_spectrum=needs_spectrum, + n_aux=plan.n_aux, ) else: param_rows = plan.op_param_indices[start:end] diff --git a/src/trspecfit/graph_ir.py b/src/trspecfit/graph_ir.py index 099b34c..d630ef6 100644 --- a/src/trspecfit/graph_ir.py +++ b/src/trspecfit/graph_ir.py @@ -2972,6 +2972,7 @@ def schedule_2d(graph: GraphIR) -> ScheduledPlan2D: profile_expr_values_init, cached_peak_sum, needs_spectrum=bool(op_needs_spectrum[op_idx]), + n_aux=n_aux, ) else: param_rows = op_param_indices[start:end] From 992bd29664a399c9d7ca565124bf537f41a843c0 Mon Sep 17 00:00:00 2001 From: Johannes Mahl Date: Fri, 10 Jul 2026 10:18:55 -0700 Subject: [PATCH 17/36] close GIR/MCP parity coverage gaps (review check 18) - compare-mode coverage for all lowerable energy shapes and all 7 IRF kernels; residual parity for the six non-exp dynamics functions - new fixtures: profile_pGauss, chained-conv MonoExpPosDoubleIRF, pinned_gauss_offset + profile_pExpDecayFixed (constant profiled op) - pipeline parity for multi-substep single-cycle dynamics (BiExpSharedT0) and chained convolution (pins n_conv_steps == 2) - TIME_1D standalone dynamics: fit_model_gir falls back to MCP and matches at the residual level (can_lower_1d rejects the domain) - constant profiled op folds into cached_result at plan build (compile-time branch previously reached by no fixture) --- tests/models/file_energy.yaml | 10 ++ tests/models/file_profile.yaml | 14 ++ tests/models/file_time.yaml | 12 ++ tests/test_evaluate_2d.py | 23 +++ tests/test_gir_integration.py | 254 ++++++++++++++++++++++++++++++--- 5 files changed, 295 insertions(+), 18 deletions(-) diff --git a/tests/models/file_energy.yaml b/tests/models/file_energy.yaml index 2253c8e..a7d484e 100644 --- a/tests/models/file_energy.yaml +++ b/tests/models/file_energy.yaml @@ -118,6 +118,16 @@ single_glp: F: [1.0, True, 0.75, 2.5] m: [0.3, True, 0, 1] +# all-pinned Gauss + free Offset — with a fixed profile on Gauss_01_A the +# profiled op is constant and folds into the cached result at plan build +pinned_gauss_offset: + Offset: + y0: [1.0, True, 0, 5] + Gauss: + A: [10, False] + x0: [85.0, False] + SD: [1.2, False] + # single GLP with the shape factor F pinned — exercises fixed-parameter # handling in the MCMC quantile table (fixed params get no posterior row) glp_one_fixed: diff --git a/tests/models/file_profile.yaml b/tests/models/file_profile.yaml index 07f1941..7ac7003 100644 --- a/tests/models/file_profile.yaml +++ b/tests/models/file_profile.yaml @@ -12,6 +12,20 @@ profile_pLinear: m: [-0.5, True, -2, 2] b: [0.0, False, -1.0, 1.0] +# pGauss profile for an amplitude parameter (centered on the test aux axis) +profile_pGauss: + pGauss: + A: [3.0, True, 0.5, 10] + x0: [2.0, True, 0, 4] + SD: [1.5, True, 0.5, 5] + +# fully-fixed profile: attached to a pinned parameter it makes the profiled +# op constant, exercising the compile-time constant-op path in schedule_2d +profile_pExpDecayFixed: + pExpDecay: + A: [200, False] + tau: [2.0, False] + # round-trip test profiles (moderate values for reliable convergence) # linear shift of x0 over depth: x0(z) = 0.5*z, max shift 2.0 eV diff --git a/tests/models/file_time.yaml b/tests/models/file_time.yaml index dd9fac7..fb04b02 100644 --- a/tests/models/file_time.yaml +++ b/tests/models/file_time.yaml @@ -152,6 +152,18 @@ MonoExpPosBoxIRF: tau: [2.5, True, 1, 10] t0: [0, False, 0, 1] +# chained convolution: both kernels rewrite the same resolved trace in +# sequence (two conv steps targeting one row in the lowered plan) +MonoExpPosDoubleIRF: + gaussCONV: + SD: [0.4, True, 0, 1] + boxCONV: + width: [2.0, True, 0.1, 20] + expFun: + A: [1, True, 0, 5] + tau: [2.5, True, 1, 10] + t0: [0, False, 0, 1] + # Empty subcycle placeholder (applies to all times in multi-cycle) ModelNone: none: {} diff --git a/tests/test_evaluate_2d.py b/tests/test_evaluate_2d.py index 18e350b..343f8d3 100644 --- a/tests/test_evaluate_2d.py +++ b/tests/test_evaluate_2d.py @@ -654,6 +654,29 @@ def test_profiled_shirley(self): plan = schedule_2d(graph) _compare_evaluator_vs_interpreter(model, plan) + # + def test_constant_profiled_op_folds_into_cache(self): + """A fully-fixed profiled component compiles into the cached result. + + Exercises the compile-time constant-op branch in schedule_2d: the + profiled op is evaluated once at plan build and folded into + ``cached_result`` instead of re-evaluating per theta. + """ + + _file, model = _make_2d_profile_model( + ["pinned_gauss_offset"], + [("Offset_y0", ["MonoExpPos"])], + [("Gauss_01_A", ["profile_pExpDecayFixed"])], + ) + graph = build_graph(model) + assert can_lower_2d(graph) + plan = schedule_2d(graph) + + # the pinned profiled Gauss op is constant; the free Offset is not + assert plan.op_is_constant.any() + assert not plan.op_is_constant.all() + _compare_evaluator_vs_interpreter(model, plan) + # def test_profile_with_time_dep_profile_params(self): """Profile function params themselves have dynamics (the hard case). diff --git a/tests/test_gir_integration.py b/tests/test_gir_integration.py index 421990c..b10f5f7 100644 --- a/tests/test_gir_integration.py +++ b/tests/test_gir_integration.py @@ -28,6 +28,17 @@ _TIME_YAML = "models/file_time.yaml" _PROFILE_YAML = "models/file_profile.yaml" +# One dynamics model per lowerable convolution kernel. +_IRF_KERNEL_MODELS = [ + "MonoExpPosIRF", + "MonoExpPosLorentzIRF", + "MonoExpPosVoigtIRF", + "MonoExpPosExpSymIRF", + "MonoExpPosExpDecayIRF", + "MonoExpPosExpRiseIRF", + "MonoExpPosBoxIRF", +] + # --------------------------------------------------------------------------- # Helpers @@ -458,18 +469,7 @@ def test_residual_with_slicing(self): np.testing.assert_allclose(res_gir, res_mcp, rtol=1e-10, atol=1e-10) # - @pytest.mark.parametrize( - "dyn_model", - [ - "MonoExpPosIRF", - "MonoExpPosLorentzIRF", - "MonoExpPosVoigtIRF", - "MonoExpPosExpSymIRF", - "MonoExpPosExpDecayIRF", - "MonoExpPosExpRiseIRF", - "MonoExpPosBoxIRF", - ], - ) + @pytest.mark.parametrize("dyn_model", _IRF_KERNEL_MODELS) def test_residual_same_gir_vs_mcp_irf(self, dyn_model): """residual_fun parity for IRF / CONVOLUTION dynamics across all lowerable kernel functions. @@ -515,14 +515,15 @@ def test_residual_same_gir_vs_mcp_irf(self, dyn_model): np.testing.assert_allclose(res_gir, res_mcp, rtol=1e-10, atol=1e-10) # - def test_compare_mode_irf(self): + @pytest.mark.parametrize("dyn_model", _IRF_KERNEL_MODELS) + def test_compare_mode_irf(self, dyn_model): """fit_model_compare exercises both paths for an IRF 2D fit.""" project = _make_project(spec_fun_str="fit_model_compare") file, model = _make_2d_model( project, ["glp_only"], - [("GLP_01_A", ["MonoExpPosIRF"])], + [("GLP_01_A", [dyn_model])], ) graph = build_graph(model) @@ -607,6 +608,178 @@ def test_subcycle_compare_mode(self): ) assert result.shape == (len(file.time), len(file.energy)) + # + @pytest.mark.parametrize( + "model_info,target_par", + [ + (["gauss_asym_only"], "GaussAsym_01_A"), + (["lorentz_only"], "Lorentz_01_A"), + (["voigt_only"], "Voigt_01_A"), + (["gls_only"], "GLS_01_A"), + (["ds_only"], "DS_01_A"), + (["linback_peak"], "GLP_01_A"), + (["shirley_peak"], "GLP_01_A"), + ], + ) + def test_compare_mode_energy_shapes(self, model_info, target_par): + """Every lowerable energy shape matches GIR vs MCP through the + compare pipeline (with dynamics attached), not just at the + evaluator level. + """ + + project = _make_project(spec_fun_str="fit_model_compare") + file, model = _make_2d_model( + project, + model_info, + [(target_par, ["MonoExpPos"])], + ) + + graph = build_graph(model) + assert can_lower_2d(graph) + plan = schedule_2d(graph) + name_to_idx = {n: i for i, n in enumerate(model.parameter_names)} + theta_indices = np.array( + [name_to_idx[n] for n in plan.opt_param_names], dtype=np.intp + ) + par = _extract_par_list(model) + + result = spectra.fit_model_compare( + file.energy, par, True, plan, theta_indices, model, 2 + ) + assert result.shape == (len(file.time), len(file.energy)) + + # + @pytest.mark.parametrize( + "dyn_model", + ["MonoStep", "MonoSin", "MonoLin", "MonoSinDivX", "MonoErf", "MonoSqrt"], + ) + def test_residual_same_gir_vs_mcp_dynamics(self, dyn_model): + """residual_fun parity across all non-exp dynamics functions + (only expFun had parity coverage; the others had pure-math tests). + """ + + project = _make_project() + file, model = _make_2d_model( + project, + ["glp_only"], + [("GLP_01_A", [dyn_model])], + ) + + model.create_value_2d() + assert model.value_2d is not None + data = model.value_2d + 0.01 + + graph = build_graph(model) + assert can_lower_2d(graph) + plan = schedule_2d(graph) + name_to_idx = {n: i for i, n in enumerate(model.parameter_names)} + theta_indices = np.array( + [name_to_idx[n] for n in plan.opt_param_names], dtype=np.intp + ) + + par = model.lmfit_pars + res_gir = fitlib.residual_fun( + par=par, + x=file.energy, + data=data, + fit_fun_str="fit_model_gir", + args=(plan, theta_indices, model, 2), + ) + res_mcp = fitlib.residual_fun( + par=par, + x=file.energy, + data=data, + fit_fun_str="fit_model_mcp", + args=(model, 2), + ) + np.testing.assert_allclose(res_gir, res_mcp, rtol=1e-10, atol=1e-10) + + # + def test_residual_same_gir_vs_mcp_multi_substep_single_cycle(self): + """Multi-substep dynamics without subcycles (frequency omitted) + match through the residual pipeline. + """ + + project = _make_project() + file, model = _make_2d_model( + project, + ["glp_only"], + [("GLP_01_A", ["BiExpSharedT0"])], + ) + + model.create_value_2d() + assert model.value_2d is not None + data = model.value_2d + 0.01 + + graph = build_graph(model) + assert can_lower_2d(graph) + plan = schedule_2d(graph) + assert plan.n_dyn_groups == 1 # two substeps, one group, no subcycle + name_to_idx = {n: i for i, n in enumerate(model.parameter_names)} + theta_indices = np.array( + [name_to_idx[n] for n in plan.opt_param_names], dtype=np.intp + ) + + par = model.lmfit_pars + res_gir = fitlib.residual_fun( + par=par, + x=file.energy, + data=data, + fit_fun_str="fit_model_gir", + args=(plan, theta_indices, model, 2), + ) + res_mcp = fitlib.residual_fun( + par=par, + x=file.energy, + data=data, + fit_fun_str="fit_model_mcp", + args=(model, 2), + ) + np.testing.assert_allclose(res_gir, res_mcp, rtol=1e-10, atol=1e-10) + + # + def test_residual_same_gir_vs_mcp_chained_conv(self): + """Chained CONVOLUTION nodes (two kernels on one trace) match + between GIR and MCP through the residual pipeline. + """ + + project = _make_project() + file, model = _make_2d_model( + project, + ["glp_only"], + [("GLP_01_A", ["MonoExpPosDoubleIRF"])], + ) + + model.create_value_2d() + assert model.value_2d is not None + data = model.value_2d + 0.01 + + graph = build_graph(model) + assert can_lower_2d(graph) + plan = schedule_2d(graph) + assert plan.n_conv_steps == 2 # actually chained, not merged + name_to_idx = {n: i for i, n in enumerate(model.parameter_names)} + theta_indices = np.array( + [name_to_idx[n] for n in plan.opt_param_names], dtype=np.intp + ) + + par = model.lmfit_pars + res_gir = fitlib.residual_fun( + par=par, + x=file.energy, + data=data, + fit_fun_str="fit_model_gir", + args=(plan, theta_indices, model, 2), + ) + res_mcp = fitlib.residual_fun( + par=par, + x=file.energy, + data=data, + fit_fun_str="fit_model_mcp", + args=(model, 2), + ) + np.testing.assert_allclose(res_gir, res_mcp, rtol=1e-10, atol=1e-10) + # --------------------------------------------------------------------------- # End-to-end through File.fit_2d @@ -937,14 +1110,15 @@ def test_compare_mode_1d(self): assert result.shape == (len(file.energy),) # - def test_compare_mode_profile_1d(self): + @pytest.mark.parametrize("profile_model", ["profile_pExpDecay", "profile_pGauss"]) + def test_compare_mode_profile_1d(self, profile_model): """Profile-aware 1D models run through compare mode without mismatch.""" project = _make_project(spec_fun_str="fit_model_compare") file, model = _make_1d_profile_model( project, ["single_gauss"], - [("Gauss_01_A", ["profile_pExpDecay"])], + [("Gauss_01_A", [profile_model])], ) graph = build_graph(model) @@ -1004,7 +1178,8 @@ def test_residual_same_gir_vs_mcp_profile_2d(self): np.testing.assert_allclose(res_gir, res_mcp, rtol=1e-10, atol=1e-10) # - def test_compare_mode_profile_2d(self): + @pytest.mark.parametrize("profile_model", ["profile_pExpDecay", "profile_pGauss"]) + def test_compare_mode_profile_2d(self, profile_model): """2D profiled models run through compare mode without mismatch.""" project = _make_project(spec_fun_str="fit_model_compare") @@ -1012,7 +1187,7 @@ def test_compare_mode_profile_2d(self): project, ["single_gauss"], [("Gauss_01_x0", ["MonoExpPos"])], - [("Gauss_01_A", ["profile_pExpDecay"])], + [("Gauss_01_A", [profile_model])], ) graph = build_graph(model) @@ -1029,6 +1204,49 @@ def test_compare_mode_profile_2d(self): ) assert result.shape == (len(file.time), len(file.energy)) + # + def test_time_1d_dynamics_model_mcp_fallback(self): + """Standalone TIME_1D dynamics models evaluate via the MCP fallback. + + ``can_lower_1d`` rejects the TIME_1D domain (outside lowered-backend + scope, see docs/design/supported_models.md), so ``fit_model_gir`` + called with interpreter args must fall through to MCP and produce + the same residual. + """ + + project = _make_project() + file = File(parent_project=project) + file.time = np.linspace(-10, 100, 111) + dyn = file.load_model( + model_yaml=_TIME_YAML, + model_info=["MonoExpPos"], + par_name="parTEST", + model_type="dynamics", + ) + + graph = build_graph(dyn) + assert not can_lower_1d(graph) + + dyn.create_value_1d() + assert dyn.value_1d is not None + data = np.asarray(dyn.value_1d) + 0.01 + + res_mcp = fitlib.residual_fun( + par=dyn.lmfit_pars, + x=file.time, + data=data, + fit_fun_str="fit_model_mcp", + args=(dyn, 1), + ) + res_gir = fitlib.residual_fun( + par=dyn.lmfit_pars, + x=file.time, + data=data, + fit_fun_str="fit_model_gir", + args=(dyn, 1), + ) + np.testing.assert_allclose(res_gir, res_mcp, rtol=1e-10, atol=1e-10) + # --------------------------------------------------------------------------- # End-to-end through File.fit_baseline / File.fit_spectrum (1D) From 2cef348ae73c22af477eded4d6667d15c4597712 Mon Sep 17 00:00:00 2001 From: Johannes Mahl Date: Fri, 10 Jul 2026 10:44:56 -0700 Subject: [PATCH 18/36] add Returns sections to profile function docstrings --- src/trspecfit/functions/profile.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/src/trspecfit/functions/profile.py b/src/trspecfit/functions/profile.py index 4f6a52a..7910ae3 100644 --- a/src/trspecfit/functions/profile.py +++ b/src/trspecfit/functions/profile.py @@ -57,6 +57,11 @@ def pExpDecay(x: np.ndarray, A: float, tau: float) -> np.ndarray: tau : float Decay constant (same units as x). + Returns + ------- + ndarray + Profile values ``A * exp(-x / tau)``, same shape as *x*. + Notes ----- Primary use cases: IMFP-weighted depth profiles in XPS, @@ -80,6 +85,11 @@ def pLinear(x: np.ndarray, m: float, b: float) -> np.ndarray: b : float Intercept. + Returns + ------- + ndarray + Profile values ``m * x + b``, same shape as *x*. + Notes ----- Primary use cases: band bending over depth, center offset over position. @@ -104,6 +114,11 @@ def pGauss(x: np.ndarray, A: float, x0: float, SD: float) -> np.ndarray: SD : float Standard deviation (width). + Returns + ------- + ndarray + Gaussian profile values (peak *A* at *x0*), same shape as *x*. + Notes ----- Primary use cases: fluence averaging, inhomogeneous broadening. From 6a0388d937c09da72c07826d1086b900406e9c6b Mon Sep 17 00:00:00 2001 From: Johannes Mahl Date: Fri, 10 Jul 2026 10:45:14 -0700 Subject: [PATCH 19/36] skip MCMC walker/corner figure construction when neither shown nor saved With show_output=0 and save_output=0 the figures were still built and closed on every MCMC run; the silent-mode test now also asserts corner.corner is never called (fails on the old code). --- src/trspecfit/fitlib.py | 47 +++++++++++++++++++++------------------ tests/test_auto_export.py | 13 +++++++---- 2 files changed, 34 insertions(+), 26 deletions(-) diff --git a/src/trspecfit/fitlib.py b/src/trspecfit/fitlib.py index 465898a..f8e4cb8 100644 --- a/src/trspecfit/fitlib.py +++ b/src/trspecfit/fitlib.py @@ -788,33 +788,36 @@ def fit_wrapper( lmfit.report_fit(emcee_fin_params) t_emcee1 = time.time() print(f"Time lmfit.emcee: {t_emcee1 - t_emcee0} s") - # acceptence fraction of all walkers (plot) # display per show_output, save per save_output (_finalize_plot - # semantics: >= 0 shows, abs == 1 saves, so -2 means neither) + # semantics: >= 0 shows, abs == 1 saves, so -2 means neither); + # skip figure construction entirely when neither shows nor saves if show_output >= 1: emcee_save = save_output else: emcee_save = -1 if abs(save_output) == 1 else -2 - fig_emcee_walker, _ax = plt.subplots(1, 1, dpi=75) - plt.plot(emcee_acceptance_fraction, "o") - plt.xlabel("Walker number") - plt.ylabel("Acceptance fraction") - uplt._finalize_plot( - emcee_save, f"{save_path}_emcee_walker_acceptance_ratio.png" - ) - # draw all combinations of the typically ellipsoidal chi plot - # [ plot] - emcee_truths = [ - emcee_fin_params.valuesdict().get(par_name) for par_name in emcee_var_names - ] - fig_emcee_corner = plt.figure(figsize=(10, 10)) - corner.corner( - emcee_flatchain, - labels=emcee_var_names, - truths=emcee_truths, - fig=fig_emcee_corner, - ) - uplt._finalize_plot(emcee_save, f"{save_path}_emcee_corner_plot.png") + if emcee_save != -2: + # acceptance fraction of all walkers (plot) + fig_emcee_walker, _ax = plt.subplots(1, 1, dpi=75) + plt.plot(emcee_acceptance_fraction, "o") + plt.xlabel("Walker number") + plt.ylabel("Acceptance fraction") + uplt._finalize_plot( + emcee_save, f"{save_path}_emcee_walker_acceptance_ratio.png" + ) + # draw all combinations of the typically ellipsoidal chi plot + # [ plot] + emcee_truths = [ + emcee_fin_params.valuesdict().get(par_name) + for par_name in emcee_var_names + ] + fig_emcee_corner = plt.figure(figsize=(10, 10)) + corner.corner( + emcee_flatchain, + labels=emcee_var_names, + truths=emcee_truths, + fig=fig_emcee_corner, + ) + uplt._finalize_plot(emcee_save, f"{save_path}_emcee_corner_plot.png") # get percentage borders to categorize emcee.flatchain data sigma_borders = sigma_start_stop_percent(ci_sigmas) # one row per sampled parameter (varying model params + the __lnsigma diff --git a/tests/test_auto_export.py b/tests/test_auto_export.py index 421cab0..cbedcfe 100644 --- a/tests/test_auto_export.py +++ b/tests/test_auto_export.py @@ -232,15 +232,19 @@ def test_fit_2d_silent_mode_prints_nothing(self, tmp_path, capsys): # @pytest.mark.slow - def test_mcmc_silent_mode_no_output_no_figures(self, tmp_path, capsys): - """MCMC honors silent mode: no progress banner, no shown figures. + def test_mcmc_silent_mode_no_output_no_figures(self, tmp_path, capsys, monkeypatch): + """MCMC honors silent mode: no progress banner, no figures built. Regression: the emcee progress banner and progress=True ran unconditionally, and with show_output=0, save_output=0 the walker - and corner figures reached _finalize_plot(0), i.e. plt.show(), - and were left open. + and corner figures were built and reached _finalize_plot(0), + i.e. plt.show(), and were left open. With neither display nor + save requested the figures must not be constructed at all. """ + mock_corner = MagicMock() + monkeypatch.setattr(fitlib, "corner", mock_corner) + project, file = _baseline_setup(tmp_path, auto_export=False) mc = MC(use_mc=1, steps=20, nwalkers=32, burn=5, thin=1) n_figs = len(plt.get_fignums()) @@ -251,6 +255,7 @@ def test_mcmc_silent_mode_no_output_no_figures(self, tmp_path, capsys): # deliberately tiny chain; only our banner is under test here assert "Progress of lmfit.emcee" not in capsys.readouterr().out assert len(plt.get_fignums()) == n_figs + assert mock_corner.corner.call_count == 0 # def test_sbs_skips_per_slice_plot_when_no_export(self, tmp_path, monkeypatch): From 64e61c8d4ff095c7ae10eeddd746a0608450b8d1 Mon Sep 17 00:00:00 2001 From: Johannes Mahl Date: Fri, 10 Jul 2026 11:03:27 -0700 Subject: [PATCH 20/36] dedupe scalar RPN evaluator: eval_1d reuses graph_ir._eval_expr_scalar --- src/trspecfit/eval_1d.py | 71 ++-------------------------------------- 1 file changed, 2 insertions(+), 69 deletions(-) diff --git a/src/trspecfit/eval_1d.py b/src/trspecfit/eval_1d.py index 372fab4..e1fd39e 100644 --- a/src/trspecfit/eval_1d.py +++ b/src/trspecfit/eval_1d.py @@ -10,80 +10,13 @@ import numpy as np from trspecfit.graph_ir import ( - ExprNodeKind, - ExprProgram, ScheduledPlan1D, + _eval_expr_scalar, _evaluate_profile_expr_values, _evaluate_profile_sample_values, _evaluate_scheduled_op_1d, ) -# --------------------------------------------------------------------------- -# Scalar RPN expression evaluator -# --------------------------------------------------------------------------- - - -# -def eval_expr_program_1d( - program: ExprProgram, - values: np.ndarray, -) -> float: - """Evaluate an RPN ExprProgram against a scalar parameter vector. - - Parameters - ---------- - program - Compiled RPN instruction array. - values - ``(n_params,)`` scalar parameter vector. - - Returns - ------- - float - Scalar result. - """ - - stack: list[float] = [] - instr = program.instructions - n_instr = len(instr) // 2 - - for i in range(n_instr): - kind = ExprNodeKind(instr[2 * i]) - operand = instr[2 * i + 1] - - if kind == ExprNodeKind.CONST: - stack.append(float(np.int64(operand).view(np.float64))) - - elif kind == ExprNodeKind.PARAM_REF: - stack.append(float(values[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 - return stack[0] - - # --------------------------------------------------------------------------- # Core 1D evaluator # --------------------------------------------------------------------------- @@ -127,7 +60,7 @@ 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_program_1d(plan.expr_programs[i], values) + values[target] = _eval_expr_scalar(plan.expr_programs[i], values) profile_sample_values = _evaluate_profile_sample_values( plan.aux_axis, From 9b797f47ab58a44df6b1637db159fb49e075dd26 Mon Sep 17 00:00:00 2001 From: Johannes Mahl Date: Fri, 10 Jul 2026 11:07:28 -0700 Subject: [PATCH 21/36] extract shared trace-resolution loop into eval_2d.resolve_param_traces --- src/trspecfit/eval_2d.py | 141 ++++++++++++++++++++++++++------------ src/trspecfit/graph_ir.py | 83 +++++++--------------- 2 files changed, 122 insertions(+), 102 deletions(-) diff --git a/src/trspecfit/eval_2d.py b/src/trspecfit/eval_2d.py index 1dcf428..56a26d3 100644 --- a/src/trspecfit/eval_2d.py +++ b/src/trspecfit/eval_2d.py @@ -141,6 +141,83 @@ def eval_expr_program( return result +# --------------------------------------------------------------------------- +# Shared trace resolution +# --------------------------------------------------------------------------- + + +# +def resolve_param_traces( + traces: np.ndarray, + time: np.ndarray, + resolution_kinds: np.ndarray, + resolution_indices: np.ndarray, + dyn_group_target_row: np.ndarray, + dyn_group_base_row: np.ndarray, + dyn_group_indptr: np.ndarray, + dyn_sub_func_id: np.ndarray, + dyn_sub_n_params: np.ndarray, + dyn_sub_param_rows: np.ndarray, + dyn_sub_time_axes: np.ndarray, + dyn_sub_masks: np.ndarray, + expr_target_rows: np.ndarray, + expr_programs: list[ExprProgram], + conv_target_rows: np.ndarray, + conv_func_ids: np.ndarray, + conv_param_indptr: np.ndarray, + conv_param_rows: np.ndarray, +) -> None: + """Resolve dynamics, expressions, and convolutions into *traces* in place. + + Dynamics groups, expressions, and resolved-trace convolutions are + interleaved in topological order so that downstream consumers see the + fully resolved trace (base + dynamics + expressions + IRF). A dynamics + group evaluates all substeps (e.g. two expFun in a bi-exponential) and + sums them: target = base + sum(traces). Expression-valued dynamics + params are resolved before the group that consumes them. + + Shared between the hot path (``evaluate_2d``) and compile-time trace + initialization in ``schedule_2d``. + """ + + for step in range(len(resolution_kinds)): + kind = int(resolution_kinds[step]) + idx = int(resolution_indices[step]) + if kind == 0: # dynamics group + target = int(dyn_group_target_row[idx]) + base = int(dyn_group_base_row[idx]) + traces[target, :] = traces[base, :] + s_start = int(dyn_group_indptr[idx]) + s_end = int(dyn_group_indptr[idx + 1]) + for s in range(s_start, s_end): + func_id = int(dyn_sub_func_id[s]) + func, _n_par = DYNAMICS_DISPATCH[func_id] + n_par = int(dyn_sub_n_params[s]) + param_rows = dyn_sub_param_rows[s, :n_par] + dyn_params = [float(traces[int(row), 0]) for row in param_rows] + traces[target, :] += ( + func(dyn_sub_time_axes[s], *dyn_params) * dyn_sub_masks[s] + ) + elif kind == 1: # expression + target = int(expr_target_rows[idx]) + traces[target, :] = eval_expr_program(expr_programs[idx], traces) + else: # kind == 2: resolved-trace convolution + target = int(conv_target_rows[idx]) + func_id = int(conv_func_ids[idx]) + kernel_func, width_func = CONV_KERNEL_DISPATCH[func_id] + p_start = int(conv_param_indptr[idx]) + p_end = int(conv_param_indptr[idx + 1]) + kernel_params = [ + float(traces[int(conv_param_rows[j]), 0]) for j in range(p_start, p_end) + ] + support = conv_kernel_support( + kernel_params[0] * width_func(*kernel_params), + float(time[1] - time[0]), + ) + kernel = kernel_func(support, *kernel_params) + traces[target, :] = my_conv(time, traces[target, :], kernel) + + # --------------------------------------------------------------------------- # Profile evaluation helpers (2D) # --------------------------------------------------------------------------- @@ -334,48 +411,28 @@ def evaluate_2d(plan: ScheduledPlan2D, theta: np.ndarray) -> np.ndarray: # 1b. Broadcast optimizer params traces[plan.opt_indices, :] = theta[:, np.newaxis] - # 1c+d. Resolve dynamics groups and expressions in interleaved topo - # order. A dynamics group evaluates all substeps (e.g. two expFun - # in a bi-exponential) and sums them: target = base + sum(traces). - # Expression-valued dynamics params are resolved before the group - # that consumes them. - 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]) - base = int(plan.dyn_group_base_row[idx]) - traces[target, :] = traces[base, :] - 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_id = int(plan.dyn_sub_func_id[s]) - func, _n_par = DYNAMICS_DISPATCH[func_id] - n_par = int(plan.dyn_sub_n_params[s]) - param_rows = plan.dyn_sub_param_rows[s, :n_par] - dyn_params = [float(traces[int(row), 0]) for row in param_rows] - traces[target, :] += ( - func(plan.dyn_sub_time_axes[s], *dyn_params) * plan.dyn_sub_masks[s] - ) - elif kind == 1: # expression - target = int(plan.expr_target_rows[idx]) - traces[target, :] = eval_expr_program(plan.expr_programs[idx], traces) - else: # kind == 2: resolved-trace convolution - target = int(plan.conv_target_rows[idx]) - func_id = int(plan.conv_func_ids[idx]) - kernel_func, width_func = CONV_KERNEL_DISPATCH[func_id] - p_start = int(plan.conv_param_indptr[idx]) - p_end = int(plan.conv_param_indptr[idx + 1]) - kernel_params = [ - float(traces[int(plan.conv_param_rows[j]), 0]) - for j in range(p_start, p_end) - ] - support = conv_kernel_support( - kernel_params[0] * width_func(*kernel_params), - float(plan.time[1] - plan.time[0]), - ) - kernel = kernel_func(support, *kernel_params) - traces[target, :] = my_conv(plan.time, traces[target, :], kernel) + # 1c+d. Resolve dynamics groups, expressions, and trace convolutions + # in interleaved topological order. + resolve_param_traces( + traces, + plan.time, + plan.resolution_kinds, + plan.resolution_indices, + plan.dyn_group_target_row, + plan.dyn_group_base_row, + plan.dyn_group_indptr, + plan.dyn_sub_func_id, + plan.dyn_sub_n_params, + plan.dyn_sub_param_rows, + plan.dyn_sub_time_axes, + plan.dyn_sub_masks, + plan.expr_target_rows, + plan.expr_programs, + plan.conv_target_rows, + plan.conv_func_ids, + plan.conv_param_indptr, + plan.conv_param_rows, + ) # 1e. Profile evaluation (after parameter resolution). profile_sample_values = _evaluate_profile_sample_values_2d( diff --git a/src/trspecfit/graph_ir.py b/src/trspecfit/graph_ir.py index d630ef6..e96f4ca 100644 --- a/src/trspecfit/graph_ir.py +++ b/src/trspecfit/graph_ir.py @@ -2863,67 +2863,30 @@ def schedule_2d(graph: GraphIR) -> ScheduledPlan2D: row = name_to_row[node.name] param_traces_init[row, :] = node.value if node.value is not None else 0.0 - # PARAM_PLUS_TRACE: base + dynamics trace at initial values. - # We need to evaluate dynamics functions at initial parameter values - # to populate these rows. - from trspecfit.functions import time as fcts_time - - _DYN_DISPATCH: dict[int, Callable[..., Any]] = { - int(DynFuncKind.EXPFUN): fcts_time.expFun, - int(DynFuncKind.SINFUN): fcts_time.sinFun, - int(DynFuncKind.LINFUN): fcts_time.linFun, - int(DynFuncKind.SINDIVX): fcts_time.sinDivX, - int(DynFuncKind.ERFFUN): fcts_time.erfFun, - int(DynFuncKind.SQRTFUN): fcts_time.sqrtFun, - int(DynFuncKind.STEPFUN): fcts_time.stepFun, - } + # PARAM_PLUS_TRACE: base + dynamics trace at initial values, resolved + # by the shared hot-path routine. + from trspecfit.eval_2d import resolve_param_traces - from trspecfit.eval_2d import CONV_KERNEL_DISPATCH, eval_expr_program - from trspecfit.utils.arrays import conv_kernel_support, my_conv - - # Dynamics groups, expressions, and resolved-trace convolutions are - # interleaved in topological order so that downstream consumers see - # the fully resolved trace (base + dynamics + expressions + IRF). - for step in range(len(resolution_kinds)): - kind = int(resolution_kinds[step]) - idx = int(resolution_indices[step]) - if kind == 0: # dynamics group - target = int(dyn_group_target_row[idx]) - base = int(dyn_group_base_row[idx]) - param_traces_init[target, :] = param_traces_init[base, :] - for s in range(int(dyn_group_indptr[idx]), int(dyn_group_indptr[idx + 1])): - n_dp = int(dyn_sub_n_params[s]) - p_vals = [ - float(param_traces_init[dyn_sub_param_rows[s, j], 0]) - for j in range(n_dp) - ] - func = _DYN_DISPATCH[int(dyn_sub_func_id[s])] - param_traces_init[target, :] += ( - func(dyn_sub_time_axes[s], *p_vals) * dyn_sub_masks[s] - ) - elif kind == 1: # expression - target_row = int(expr_target_rows[idx]) - program = expr_programs[idx] - param_traces_init[target_row, :] = eval_expr_program( - program, param_traces_init - ) - else: # kind == 2: resolved-trace convolution - target_row = int(conv_target_rows[idx]) - kernel_func, width_func = CONV_KERNEL_DISPATCH[int(conv_func_ids[idx])] - p_start = int(conv_param_indptr[idx]) - p_end = int(conv_param_indptr[idx + 1]) - kernel_params = [ - float(param_traces_init[int(conv_param_rows[j]), 0]) - for j in range(p_start, p_end) - ] - support = conv_kernel_support( - kernel_params[0] * width_func(*kernel_params), - float(graph.time[1] - graph.time[0]), - ) - kernel = kernel_func(support, *kernel_params) - param_traces_init[target_row, :] = my_conv( - graph.time, param_traces_init[target_row, :], kernel - ) + resolve_param_traces( + param_traces_init, + graph.time, + resolution_kinds, + resolution_indices, + dyn_group_target_row, + dyn_group_base_row, + dyn_group_indptr, + dyn_sub_func_id, + dyn_sub_n_params, + dyn_sub_param_rows, + dyn_sub_time_axes, + dyn_sub_masks, + expr_target_rows, + expr_programs, + conv_target_rows, + conv_func_ids, + conv_param_indptr, + conv_param_rows, + ) # ------------------------------------------------------------------ # # 6b. Precompute constant component contributions # From 1c049f2302915ebcb044c17d0e299dd84cfe725c Mon Sep 17 00:00:00 2001 From: Johannes Mahl Date: Fri, 10 Jul 2026 12:29:21 -0700 Subject: [PATCH 22/36] plan kernel-matrix convolution for non-uniform time axes --- TODO.md | 1 + docs/design/jax-planning.md | 15 +- docs/design/kernel-matrix-convolution.md | 168 +++++++++++++++++++++++ 3 files changed, 181 insertions(+), 3 deletions(-) create mode 100644 docs/design/kernel-matrix-convolution.md diff --git a/TODO.md b/TODO.md index 2612170..9e8120e 100644 --- a/TODO.md +++ b/TODO.md @@ -3,6 +3,7 @@ ## Fitting - [ ] **Enable 1D time-trace fitting (post-SbS kinetics)**: wire standalone `TIME_1D` dynamics models (already evaluable on the mcp path per `docs/design/supported_models.md`, but connected to no fit method) to a fit entry point, so parameter-vs-time traces extracted from SbS results can be fit to `functions/time.py` dynamics (`expFun` sums, IRF convolution) inside the package instead of in external scripts. Frame it as the diagnostic/initialization rung before `fit_2d`, not a statistical equivalent (per-slice correlations and unpropagated SbS uncertainties make two-step inferior). Design direction: promote SbS results into a time-axis `File` — the promoted object inherits limits/CI/MCMC/save/export machinery, and trace fits land in a `SavedFitSlot` like every other fit type (no parallel fitting pipeline). Requires `File` to support time as the primary axis, which today it assumes is energy. Weighting traces by per-slice stderr ties into the `sigma_type` expansion item below. +- [ ] **Kernel-matrix convolution (non-uniform time axes)**: replace the 1D-kernel convolution (`my_conv` + per-theta `conv_kernel_support`) with a quadrature-weighted kernel-matrix operator on both the mcp and GIR paths (one branch — parity tests couple them). Fixes silently wrong IRF convolution on non-uniform time axes (measured 2026-07-10 on example 21's 0.5→2.0 step axis: deviations up to ~7% of trace max vs the exact continuous convolution, worst just past the step change; the shipped example round-trips only because its data was generated through the same operator) and removes the theta-dependent kernel-shape jit blocker for the JAX track. One shared helper serves standalone 1D time traces and dynamics traces inside 2D models (single mcp conv site in `Model._combine_component`, plus two GIR sites), so it also enables the 1D time-trace fitting item above on measured delay axes. Requires regenerating example 21's data and one golden-value update. Plan: [docs/design/kernel-matrix-convolution.md](docs/design/kernel-matrix-convolution.md). Sequencing: after `fix-conv-kernels` merges, before the JAX backend item below. ## Noise and simulation diff --git a/docs/design/jax-planning.md b/docs/design/jax-planning.md index 4b4df38..53993ab 100644 --- a/docs/design/jax-planning.md +++ b/docs/design/jax-planning.md @@ -73,6 +73,13 @@ useful cleanup before or during a JAX port: `evaluate(plan, theta)` is already the right API. New work should protect that contract rather than pushing JAX concerns back into model objects or fit-time parsing. +- **Land kernel-matrix convolution first.** The lowered convolution path + builds a theta-dependent 1D kernel whose support length tracks the fitted + width — silently wrong on non-uniform time axes and a jit blocker + (theta-dependent array shapes). Replacing it with the quadrature-weighted + kernel-matrix operator + ([kernel-matrix-convolution.md](kernel-matrix-convolution.md)) on the + NumPy paths removes that blocker before the port begins. Of these, expression flattening is the one prep item most worth doing even if the eventual backend choice changes again. @@ -94,8 +101,9 @@ The main technical work is in the evaluator itself: the compiled path. - **SciPy-dependent kernels need JAX-compatible replacements.** The current code still depends on SciPy for Voigt/Faddeeva (`wofz`) and for convolution - utilities used by the lowered convolution path. A JAX backend needs - compatible implementations for those pieces. + utilities used by the lowered convolution path. The kernel-matrix change + ([kernel-matrix-convolution.md](kernel-matrix-convolution.md)) retires the + convolution dependency, leaving Voigt/`wofz` as the remaining gap. ### 3. Jacobian and optimizer work @@ -174,7 +182,8 @@ Add the remaining lowered features incrementally: - profile-varying parameters, - subcycle-aware dynamics, -- resolved-trace convolution, +- resolved-trace convolution (kernel-matrix form; see + [kernel-matrix-convolution.md](kernel-matrix-convolution.md)), - Voigt / special-function support. Each widening step should ship with direct parity tests against the existing diff --git a/docs/design/kernel-matrix-convolution.md b/docs/design/kernel-matrix-convolution.md new file mode 100644 index 0000000..a287324 --- /dev/null +++ b/docs/design/kernel-matrix-convolution.md @@ -0,0 +1,168 @@ +--- +orphan: true +--- + +# Planning Note: Kernel-Matrix Convolution + +## Summary + +Replace the 1D-kernel-array convolution +([`my_conv`](../../src/trspecfit/utils/arrays.py) plus the per-theta +`conv_kernel_support` rebuild) with a quadrature-weighted kernel-matrix +operator, on both the mcp and GIR paths, in one branch. Two motivations: + +1. **Correctness.** Sample-index convolution is silently wrong on + non-uniform time axes. Measured on example 21's axis + (steps 0.5 → 2.0) with its truth parameters (2026-07-10): up to ~7% of + the trace maximum versus the exact continuous convolution, worst just + past the step change. +2. **Architecture.** It removes theta-dependent kernel array shapes — the + main jit blocker for the JAX track + ([jax-planning.md](jax-planning.md)) — and retires the SciPy + convolution dependency in the lowered path. + +The mcp and GIR changes cannot be split across branches: parity tests +assert the two paths agree to 1e-10, so the convolution semantics must +change everywhere at once. + + +## The defect today + +- Kernels are sampled at `t_step = time[1] - time[0]` + ([`Component.create_t_kernel`](../../src/trspecfit/mcp.py), + `resolve_param_traces` in + [`eval_2d.py`](../../src/trspecfit/eval_2d.py), and the schedule-time + trace init in [`graph_ir.py`](../../src/trspecfit/graph_ir.py)). +- `my_conv` convolves by sample index; the axis itself never enters the + computation. On a non-uniform axis the effective IRF width scales with + the local step (4x wider in example 21's coarse region), and points + near the step change see a blend of both regimes. +- Nothing guards against this: a non-uniform time axis flows into the + convolution path silently. +- The shipped example is self-consistent — its data was generated through + the same operator, so parameter recovery round-trips. Real measured + data on a fine-around-t0 + coarse-tail axis would produce biased fits, + concentrated in the IRF width and everything near the step change. + + +## The operator + +For a monotonic time axis `t` (length `n_t`) and kernel function `g` +with fitted parameters `theta_k`: + +- `dt[i, j] = t_i - t_j` — theta-independent, built once. +- `w_j` = trapezoid quadrature weights (`np.gradient(t)`) — + theta-independent, built once. +- Per evaluation: `K[i, j] = g(dt[i, j]; theta_k) * w_j`, rows + normalized to sum to 1; then `y_conv = K @ y`. + +Properties: + +- Exact on any monotonic axis; on a uniform axis the weights cancel in + the row normalization and the result matches the current path up to + the (small) truncation of the current finite support. +- Static shapes, no support truncation at all — this supersedes the + dynamic-support machinery from the `fix-conv-kernels` branch rather + than layering on top of it. +- Differentiable and jit-friendly: per-theta work is an elementwise + kernel evaluation plus a matmul. +- O(n_t^2) per convolution node per evaluation. At the current + n_t ~ hundreds this is sub-millisecond; see the benchmark gate below. + + +## Generality: one operator for 1D and 2D + +Call-site inventory — every convolution in the package acts on a 1D +trace sampled on the time axis: + +- **mcp (single site).** `Model._combine_component` in + [`mcp.py`](../../src/trspecfit/mcp.py) handles `comp_type == "conv"` + for standalone `TIME_1D` models and for dynamics traces inside 2D + models alike — parameter time dependence evaluates via + `par.t_model.create_value_1d()`, i.e. the 2D case reuses the 1D + time-trace evaluator. Replacing this one site covers both. +- **GIR (two sites).** The `kind == 2` branch of + `resolve_param_traces` in + [`eval_2d.py`](../../src/trspecfit/eval_2d.py), and the schedule-time + trace initialization in + [`graph_ir.py`](../../src/trspecfit/graph_ir.py). The plan gains the + precomputed `dt` matrix and weights; the per-theta support recompute + disappears. +- **Future.** The 1D time-trace fitting item in `TODO.md` inherits + correct convolution on measured (typically non-uniform) delay axes + automatically — the matrix operator is an enabler for that feature, + not just compatible with it. +- Energy-domain convolution remains undefined (explicitly rejected in + `_combine_component`); out of scope here, though the same helper + would apply if it is ever defined. + + +## Design decisions + +1. **Edge policy.** `my_conv` pads the signal with edge values + (`mode="edge"`), convolves, and crops. Matrix equivalents: + (a) row normalization only — kernel mass falling outside the window + is renormalized over interior samples; or (b) edge-mass + accumulation — kernel mass beyond each end of the axis is added to + the first/last column, reproducing edge-padding semantics. + Recommendation: (b). It matches current behavior on uniform axes + (signal assumed constant beyond the window — the right assumption + for saturating or decaying traces) and minimizes golden-value churn. + Decide before implementation and test both edge rows explicitly. +2. **Kernel bodies evaluate on `dt` directly.** Registry kernel + functions (`gaussCONV`, …) are elementwise in their first argument, + so they apply to the `dt` matrix unchanged; no new numeric bodies. + The `*_kernel_width` helpers and `conv_kernel_support` existed only + to size the 1D support and are removed together with the + dynamic-support recompute and its tests. +3. **Normalization.** Per-row, generalizing the current + kernel-divided-by-sum normalization to non-uniform sampling. +4. **Validation.** Keep a monotonic-axis check at the authoring layer; + non-uniform spacing stops being a (silent) error condition and + becomes supported. + + +## Implementation order and blast radius + +1. `utils/arrays.py`: matrix-convolution helper (theta-independent + builder + per-theta apply), unit-tested against an analytic + Gaussian-times-exponential reference, uniform-axis agreement with + the current path within truncation tolerance, and a dense + continuous-convolution reference on a non-uniform axis. +2. mcp path: `_combine_component` and the `Component` kernel handling + (`create_t_kernel` path). +3. GIR: plan fields (`dt`, weights; kernel function ids stay), the + `kind == 2` branch of `resolve_param_traces`, and the schedule-time + init convolution. +4. Parity: mcp vs GIR tests keep asserting 1e-10 — both paths must + consume the same helper. +5. Regenerate example 21's data (`generate_data.ipynb`; the old + operator's artifact is baked into the CSVs) and re-run affected + examples per `docs/ai/check-example.md`. Update roundtrip/golden + values once. +6. Benchmark before/after per `docs/ai/benchmark.md`. The O(n_t^2) + cost is expected to be noise at current sizes; if long time axes + ever make it matter, the escape hatch is a banded matrix with a + static support cap derived from parameter bounds (shapes stay + static). + + +## Relationship to the JAX track + +Land this before Phase B/C of [jax-planning.md](jax-planning.md). It +removes two blockers listed there: the SciPy convolution utilities in +the lowered path, and the theta-dependent kernel shapes introduced by +the dynamic-support fix. After this change, porting convolution to JAX +is an elementwise kernel evaluation plus a matmul. + + +## Success criteria + +- Non-uniform-axis convolution matches a dense continuous-convolution + reference to floating-point tolerance. +- Uniform-axis results match the previous path within the documented + truncation tolerance; examples and golden values updated exactly once. +- mcp/GIR parity retained at 1e-10. +- No per-theta array shapes remain anywhere in the convolution path. +- Benchmarks show no regression beyond noise, or document the banded + fallback and its trigger. From 55cd1e85229dd2885580d53e5d2905bb7c155693 Mon Sep 17 00:00:00 2001 From: Johannes Mahl Date: Fri, 10 Jul 2026 12:38:24 -0700 Subject: [PATCH 23/36] dedupe schedule_2d/schedule_1d profile compilation and op scheduling --- src/trspecfit/graph_ir.py | 1717 ++++++++++++++++--------------------- 1 file changed, 747 insertions(+), 970 deletions(-) diff --git a/src/trspecfit/graph_ir.py b/src/trspecfit/graph_ir.py index e96f4ca..0db9126 100644 --- a/src/trspecfit/graph_ir.py +++ b/src/trspecfit/graph_ir.py @@ -21,7 +21,7 @@ from collections.abc import Callable from dataclasses import dataclass, field from enum import IntEnum -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, NamedTuple import numpy as np @@ -2435,7 +2435,7 @@ def schedule_2d(graph: GraphIR) -> ScheduledPlan2D: resolution_indices = np.array(resolution_indices_list, dtype=np.intp) # ------------------------------------------------------------------ # - # 4b. Compile PROFILE_SAMPLE groups # + # 4b. Compile profile groups (samples + expressions) # # ------------------------------------------------------------------ # # Build edge indexes for profile compilation (mirrors schedule_1d). param_edges_by_target: dict[int, list[GraphEdge]] = {} @@ -2452,406 +2452,51 @@ def schedule_2d(graph: GraphIR) -> ScheduledPlan2D: elif edge.kind == EdgeKind.SPECTRUM_INPUT: spectrum_input_targets.add(edge.target) - profile_sample_groups: dict[str, list[GraphNode]] = {} - for nid in topo_order: - node = id_to_node[nid] - if node.kind != NodeKind.PROFILE_SAMPLE: - continue - group_name = _profile_group_name(node.name, "profile_sample") - profile_sample_groups.setdefault(group_name, []).append(node) - - plan_aux_axis = np.zeros(0, dtype=np.float64) - n_aux = 0 - profile_sample_base_rows_list: list[int] = [] - profile_sample_component_indptr_list: list[int] = [0] - profile_component_func_ids_list: list[int] = [] - profile_component_param_indptr_list: list[int] = [0] - profile_component_param_rows_list: list[int] = [] - profile_sample_is_constant_list: list[bool] = [] - profile_sample_group_idx: dict[str, int] = {} - - for group_name, sample_nodes in profile_sample_groups.items(): - sample_nodes_sorted = sorted( - sample_nodes, - key=lambda node: _profile_group_index(node.name, "profile_sample"), - ) - aux_indices = [ - _profile_group_index(node.name, "profile_sample") - for node in sample_nodes_sorted - ] - if aux_indices != list(range(len(sample_nodes_sorted))): - raise ValueError( - f"PROFILE_SAMPLE nodes for {group_name!r} do not cover " - "a contiguous aux-axis range" - ) - - aux_axis = sample_nodes_sorted[0].arrays.get("aux_axis") - if aux_axis is None: - raise ValueError(f"PROFILE_SAMPLE {group_name!r} is missing aux_axis") - aux_axis = np.asarray(aux_axis, dtype=np.float64) - if n_aux == 0: - n_aux = len(aux_axis) - plan_aux_axis = aux_axis.copy() - elif len(aux_axis) != n_aux or not np.array_equal(aux_axis, plan_aux_axis): - raise ValueError("All lowered profile groups must share one fixed aux_axis") - if len(sample_nodes_sorted) != n_aux: - raise ValueError( - f"PROFILE_SAMPLE group {group_name!r} has " - f"{len(sample_nodes_sorted)} samples but aux_axis length {n_aux}" - ) - - rep_node = sample_nodes_sorted[0] - rep_param_edges = sorted( - param_edges_by_target.get(rep_node.id, []), - key=lambda edge: edge.position or 0, - ) - if not rep_param_edges: - raise ValueError(f"PROFILE_SAMPLE {group_name!r} has no PARAM_INPUT edges") - - base_node = id_to_node[rep_param_edges[0].source] - base_row = name_to_row[base_node.name] - is_constant = bool(row_is_constant[base_row]) - profile_sample_base_rows_list.append(base_row) - - component_func_by_name: dict[str, int] = {} - component_param_rows_by_name: dict[str, list[int]] = {} - component_order: list[str] = [] - for edge in rep_param_edges[1:]: - src_node = id_to_node[edge.source] - src_row = name_to_row[src_node.name] - # PARAM_PLUS_TRACE nodes have a "_resolved" suffix; CONVOLUTION - # nodes wrap a PPT for profile-time-dynamics and carry their own - # "__dynamics" name. In both cases, walk back to - # the underlying PPT so profile component parsing sees the - # original profile parameter name. - parse_name = src_node.name - if src_node.kind == NodeKind.PARAM_PLUS_TRACE: - parse_name = parse_name.removesuffix("_resolved") - elif src_node.kind == NodeKind.CONVOLUTION: - ppt_node = _walk_convolution_to_param_plus_trace( - src_node, graph.edges, id_to_node - ) - parse_name = ppt_node.name.removesuffix("_resolved") - comp_name, func_name = _parse_profile_component_param_name( - group_name, - parse_name, - ) - prof_func_kind = _FUNCTION_NAME_TO_PROFILE_FUNC.get(func_name) - if prof_func_kind is None: - raise ValueError(f"Unknown profile function: {func_name!r}") - - if comp_name not in component_func_by_name: - component_order.append(comp_name) - component_func_by_name[comp_name] = int(prof_func_kind) - component_param_rows_by_name[comp_name] = [] - - component_param_rows_by_name[comp_name].append(src_row) - is_constant = is_constant and bool(row_is_constant[src_row]) - - for comp_name in component_order: - profile_component_func_ids_list.append(component_func_by_name[comp_name]) - profile_component_param_rows_list.extend( - component_param_rows_by_name[comp_name] - ) - profile_component_param_indptr_list.append( - len(profile_component_param_rows_list) - ) - profile_sample_component_indptr_list.append( - len(profile_component_func_ids_list) - ) - - group_idx = len(profile_sample_base_rows_list) - 1 - profile_sample_group_idx[group_name] = group_idx - profile_sample_is_constant_list.append(is_constant) - - n_profile_samples = len(profile_sample_base_rows_list) - profile_sample_base_rows = np.array(profile_sample_base_rows_list, dtype=np.intp) - profile_sample_component_indptr = np.array( - profile_sample_component_indptr_list, dtype=np.intp - ) - profile_component_func_ids = np.array( - profile_component_func_ids_list, dtype=np.intp - ) - profile_component_param_indptr = np.array( - profile_component_param_indptr_list, dtype=np.intp - ) - profile_component_param_rows = np.array( - profile_component_param_rows_list, dtype=np.intp - ) - profile_sample_is_constant = np.array( - profile_sample_is_constant_list, dtype=np.bool_ - ) - - # ------------------------------------------------------------------ # - # 4c. Compile per-sample profile expressions # - # ------------------------------------------------------------------ # - profile_expr_groups: dict[str, list[GraphNode]] = {} - for nid in topo_order: - node = id_to_node[nid] - if _is_profile_expr_node(node): - group_name = _profile_group_name(node.name, "profile_expr") - profile_expr_groups.setdefault(group_name, []).append(node) - - profile_expr_programs_2d: list[ExprProgram] = [] - profile_expr_is_constant_list: list[bool] = [] - profile_expr_group_idx: dict[str, int] = {} - for group_name, p_expr_nodes in profile_expr_groups.items(): - p_expr_nodes_sorted = sorted( - p_expr_nodes, - key=lambda node: _profile_group_index(node.name, "profile_expr"), - ) - aux_indices = [ - _profile_group_index(node.name, "profile_expr") - for node in p_expr_nodes_sorted - ] - if aux_indices != list(range(len(p_expr_nodes_sorted))): - raise ValueError( - f"Profile expression nodes for {group_name!r} do not cover " - "a contiguous aux-axis range" - ) - if len(p_expr_nodes_sorted) != n_aux: - raise ValueError( - f"Profile expression group {group_name!r} has " - f"{len(p_expr_nodes_sorted)} samples but aux_axis length {n_aux}" - ) - - rep_node = p_expr_nodes_sorted[0] - if rep_node.expr_string is None: - raise ValueError( - f"Profile expression {group_name!r} is missing expr_string" - ) - expr_refs = set(_extract_expression_references(rep_node.expr_string)) - - prof_ref_map: dict[str, int] = {} - for edge in expr_ref_edges_by_target.get(rep_node.id, []): - src_node = id_to_node[edge.source] - if src_node.kind == NodeKind.PROFILE_SAMPLE: - sample_name = _profile_group_name(src_node.name, "profile_sample") - src_idx = n_params + profile_sample_group_idx[sample_name] - match_name = sample_name - else: - src_idx = name_to_row[src_node.name] - match_name = src_node.name - # PARAM_PLUS_TRACE nodes carry a "_resolved" suffix; - # the expression string uses the bare param name. - if match_name not in expr_refs and src_node.name.endswith("_resolved"): - match_name = src_node.name.removesuffix("_resolved") - - if match_name in expr_refs: - prof_ref_map[match_name] = src_idx - - symbolic = compile_expr_symbolic(rep_node.expr_string) - prof_binding: dict[str, int] = dict(name_to_row) - prof_binding.update(prof_ref_map) - program = _bind_expr_to_rows(symbolic, prof_binding) - profile_expr_programs_2d.append(program) - - is_constant = True - for name in symbolic.referenced_names: - bound_idx = int(prof_binding[name]) - if bound_idx < n_params: - is_constant = is_constant and bool(row_is_constant[bound_idx]) - else: - is_constant = is_constant and bool( - profile_sample_is_constant[bound_idx - n_params] - ) - profile_expr_is_constant_list.append(is_constant) - profile_expr_group_idx[group_name] = len(profile_expr_programs_2d) - 1 - - n_profile_exprs = len(profile_expr_programs_2d) - profile_expr_is_constant = np.array(profile_expr_is_constant_list, dtype=np.bool_) - - # ------------------------------------------------------------------ # - # 5. Schedule component ops # - # ------------------------------------------------------------------ # - # Identify peak_sum contributors: nodes with ADDEND edges into the - # "peak_sum" SUM node (if it exists). - peak_sum_sources: set[int] = set() - peak_sum_nid = graph.node_by_name.get("peak_sum") - if peak_sum_nid is not None: - for e in addend_edges_by_target.get(peak_sum_nid, []): - peak_sum_sources.add(e.source) - - # Collect per-sample component inputs for PROFILE_AVERAGE nodes. - profile_avg_sample_inputs: dict[int, list[GraphNode]] = {} - sample_component_ids: set[int] = set() - for nid in topo_order: - node = id_to_node[nid] - if node.kind != NodeKind.PROFILE_AVERAGE: - continue - sample_nodes = [ - id_to_node[edge.source] - for edge in addend_edges_by_target.get(node.id, []) - if id_to_node[edge.source].kind - in (NodeKind.COMPONENT_EVAL, NodeKind.SPECTRUM_FED_OP) - ] - profile_avg_sample_inputs[node.id] = sample_nodes - sample_component_ids.update(sample.id for sample in sample_nodes) - - comp_op_kinds = {NodeKind.COMPONENT_EVAL, NodeKind.SPECTRUM_FED_OP} - comp_nodes_topo = [ - id_to_node[nid] - for nid in topo_order - if ( - (id_to_node[nid].kind in comp_op_kinds and nid not in sample_component_ids) - or id_to_node[nid].kind == NodeKind.PROFILE_AVERAGE - ) - ] - n_ops = len(comp_nodes_topo) - - op_schedule = np.arange(n_ops, dtype=np.intp) - op_kinds_list: list[int] = [] - op_param_indptr_list: list[int] = [0] - op_param_source_kinds_list: list[int] = [] - op_param_indices_list: list[int] = [] - op_needs_spectrum_list: list[bool] = [] - op_is_pre_spectrum_list: list[bool] = [] - op_is_profiled_list: list[bool] = [] - op_is_constant_list: list[bool] = [] - - for comp_node in comp_nodes_topo: - if comp_node.kind == NodeKind.PROFILE_AVERAGE: - # Profiled component: gather from per-sample COMPONENT_EVAL inputs. - sample_nodes = sorted( - profile_avg_sample_inputs.get(comp_node.id, []), - key=lambda node: _profile_component_sample_index(node.name), - ) - if not sample_nodes: - raise ValueError( - f"PROFILE_AVERAGE {comp_node.name!r} has no sample component inputs" - ) - if len(sample_nodes) != n_aux: - raise ValueError( - f"PROFILE_AVERAGE {comp_node.name!r} has " - f"{len(sample_nodes)} samples " - f"but aux_axis length {n_aux}" - ) - - rep_node = sample_nodes[0] - assert rep_node.function_name is not None - op = _FUNCTION_NAME_TO_OP.get(rep_node.function_name) - if op is None: - raise ValueError( - f"Unknown component function: {rep_node.function_name!r}" - ) - op_kinds_list.append(int(op)) - op_is_profiled_list.append(True) - - rep_param_edges = sorted( - param_edges_by_target.get(rep_node.id, []), - key=lambda edge: edge.position or 0, - ) - sample_param_edges = [ - sorted( - param_edges_by_target.get(sample_node.id, []), - key=lambda edge: edge.position or 0, - ) - for sample_node in sample_nodes - ] - is_constant = True - for pos, rep_edge in enumerate(rep_param_edges): - src_node = id_to_node[rep_edge.source] - if src_node.kind == NodeKind.PROFILE_SAMPLE: - group_name = _profile_group_name(src_node.name, "profile_sample") - source_kind = int(ParamSourceKind.PROFILE_SAMPLE) - source_idx = profile_sample_group_idx[group_name] - is_constant = is_constant and bool( - profile_sample_is_constant[source_idx] - ) - for aux_i, edges in enumerate(sample_param_edges): - sample_src = id_to_node[edges[pos].source] - if sample_src.kind != NodeKind.PROFILE_SAMPLE: - raise ValueError( - "Mixed parameter source kinds " - f"in profiled op {comp_node.name!r}" - ) - if ( - _profile_group_name(sample_src.name, "profile_sample") - != group_name - or _profile_group_index(sample_src.name, "profile_sample") - != aux_i - ): - raise ValueError( - "Inconsistent PROFILE_SAMPLE wiring " - f"in {comp_node.name!r}" - ) - elif _is_profile_expr_node(src_node): - group_name = _profile_group_name(src_node.name, "profile_expr") - source_kind = int(ParamSourceKind.PROFILE_EXPR) - source_idx = profile_expr_group_idx[group_name] - is_constant = is_constant and bool( - profile_expr_is_constant[source_idx] - ) - for aux_i, edges in enumerate(sample_param_edges): - sample_src = id_to_node[edges[pos].source] - if not _is_profile_expr_node(sample_src): - raise ValueError( - "Mixed expression source kinds " - f"in profiled op {comp_node.name!r}" - ) - if ( - _profile_group_name(sample_src.name, "profile_expr") - != group_name - or _profile_group_index(sample_src.name, "profile_expr") - != aux_i - ): - raise ValueError( - "Inconsistent profile-expression " - f"wiring in {comp_node.name!r}" - ) - else: - source_kind = int(ParamSourceKind.SCALAR) - source_idx = name_to_row[src_node.name] - is_constant = is_constant and bool(row_is_constant[source_idx]) - for edges in sample_param_edges[1:]: - if id_to_node[edges[pos].source].id != src_node.id: - raise ValueError( - "Scalar parameter source changed " - f"across samples in {comp_node.name!r}" - ) - - op_param_source_kinds_list.append(source_kind) - op_param_indices_list.append(source_idx) - - op_param_indptr_list.append(len(op_param_indices_list)) - else: - # Non-profiled component: standard param row wiring. - assert comp_node.function_name is not None - op = _FUNCTION_NAME_TO_OP.get(comp_node.function_name) - if op is None: - raise ValueError( - f"Unknown component function: {comp_node.function_name!r}" - ) - op_kinds_list.append(int(op)) - op_is_profiled_list.append(False) - - param_edges = sorted( - param_edges_by_target.get(comp_node.id, []), - key=lambda edge: edge.position or 0, - ) - is_constant = True - for pe in param_edges: - src_node = id_to_node[pe.source] - src_row = name_to_row[src_node.name] - op_param_source_kinds_list.append(int(ParamSourceKind.SCALAR)) - op_param_indices_list.append(src_row) - is_constant = is_constant and bool(row_is_constant[src_row]) - - op_param_indptr_list.append(len(op_param_indices_list)) - - has_spec_input = comp_node.id in spectrum_input_targets - op_needs_spectrum_list.append(has_spec_input) - op_is_pre_spectrum_list.append(comp_node.id in peak_sum_sources) - op_is_constant_list.append((not has_spec_input) and is_constant) + profiles = _compile_profile_groups( + graph, + topo_order, + id_to_node, + param_edges_by_target, + expr_ref_edges_by_target, + name_to_row, + row_is_constant, + n_params, + ) + plan_aux_axis = profiles.aux_axis + n_aux = profiles.n_aux + n_profile_samples = profiles.n_samples + profile_sample_base_rows = profiles.sample_base_indices + profile_sample_component_indptr = profiles.sample_component_indptr + profile_component_func_ids = profiles.component_func_ids + 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 - op_kinds = np.array(op_kinds_list, dtype=np.intp) - op_param_indptr = np.array(op_param_indptr_list, dtype=np.intp) - op_param_source_kinds = np.array(op_param_source_kinds_list, dtype=np.int8) - op_param_indices = np.array(op_param_indices_list, dtype=np.intp) - op_needs_spectrum = np.array(op_needs_spectrum_list, dtype=np.bool_) - op_is_pre_spectrum = np.array(op_is_pre_spectrum_list, dtype=np.bool_) - op_is_profiled = np.array(op_is_profiled_list, dtype=np.bool_) - op_is_constant = np.array(op_is_constant_list, dtype=np.bool_) + # ------------------------------------------------------------------ # + # 5. Schedule component ops # + # ------------------------------------------------------------------ # + ops = _schedule_component_ops( + graph, + topo_order, + id_to_node, + param_edges_by_target, + addend_edges_by_target, + spectrum_input_targets, + name_to_row, + row_is_constant, + profiles, + ) + n_ops = ops.n_ops + op_schedule = np.arange(n_ops, dtype=np.intp) + op_kinds = ops.op_kinds + op_param_indptr = ops.op_param_indptr + op_param_source_kinds = ops.op_param_source_kinds + op_param_indices = ops.op_param_indices + op_needs_spectrum = ops.op_needs_spectrum + op_is_pre_spectrum = ops.op_is_pre_spectrum + op_is_profiled = ops.op_is_profiled + op_is_constant = ops.op_is_constant # ------------------------------------------------------------------ # # 6. Initialize trace matrix # @@ -3016,114 +2661,602 @@ 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. - Parameters - ---------- - program - Compiled RPN instruction array. - values - ``(n_params,)`` scalar parameter vector. + Parameters + ---------- + program + Compiled RPN instruction array. + values + ``(n_params,)`` scalar parameter vector. + + Returns + ------- + float + Scalar result. + """ + + stack: list[float] = [] + instr = program.instructions + n_instr = len(instr) // 2 + + for i in range(n_instr): + kind = ExprNodeKind(instr[2 * i]) + operand = instr[2 * i + 1] + + if kind == ExprNodeKind.CONST: + stack.append(float(np.int64(operand).view(np.float64))) + elif kind == ExprNodeKind.PARAM_REF: + stack.append(float(values[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 + return stack[0] + + +# +def _profile_group_name(node_name: str, label: str) -> str: + """Return the shared profile group base name from a per-sample node name.""" + + match = re.fullmatch(rf"(.+)_{label}_(\d+)", node_name) + if match is None: + raise ValueError(f"Malformed profile node name: {node_name!r}") + return match.group(1) + + +# +def _profile_group_index(node_name: str, label: str) -> int: + """Return the aux-axis sample index encoded in a per-sample node name.""" + + match = re.fullmatch(rf"(.+)_{label}_(\d+)", node_name) + if match is None: + raise ValueError(f"Malformed profile node name: {node_name!r}") + return int(match.group(2)) + + +# +def _profile_component_sample_index(node_name: str) -> int: + """Return the aux-axis sample index from ``_sample_``.""" + + return _profile_group_index(node_name, "sample") + + +# +def _is_profile_expr_node(node: GraphNode) -> bool: + """Return True for per-sample profile-expression nodes.""" + + return ( + node.kind == NodeKind.EXPRESSION + and re.fullmatch(r"(.+)_profile_expr_(\d+)", node.name) is not None + ) + + +# +def _parse_profile_component_param_name( + target_param_name: str, + source_param_name: str, +) -> tuple[str, str]: + """Parse ``__`` into component + function name.""" + + prefix = f"{target_param_name}_" + if not source_param_name.startswith(prefix): + raise ValueError( + f"Profile parameter {source_param_name!r} does not match " + f"target parameter {target_param_name!r}" + ) + + remainder = source_param_name[len(prefix) :] + comp_name, sep, _par_name = remainder.rpartition("_") + if not sep or not comp_name: + raise ValueError(f"Malformed profile parameter name: {source_param_name!r}") + + func_name, sep2, comp_idx = comp_name.rpartition("_") + if not sep2 or not comp_idx.isdigit(): + raise ValueError(f"Malformed profile component name: {comp_name!r}") + + return comp_name, func_name + + +# +# +class _CompiledProfileGroups(NamedTuple): + """Packed results of PROFILE_SAMPLE / profile-expression compilation.""" + + aux_axis: np.ndarray + n_aux: int + n_samples: int + sample_base_indices: np.ndarray + sample_component_indptr: np.ndarray + component_func_ids: np.ndarray + component_param_indptr: np.ndarray + component_param_indices: np.ndarray + sample_is_constant: np.ndarray + sample_group_idx: dict[str, int] + n_exprs: int + expr_programs: list[ExprProgram] + expr_is_constant: np.ndarray + expr_group_idx: dict[str, int] + + +# +def _compile_profile_groups( + graph: GraphIR, + topo_order: list[int], + id_to_node: dict[int, GraphNode], + param_edges_by_target: dict[int, list[GraphEdge]], + expr_ref_edges_by_target: dict[int, list[GraphEdge]], + name_to_index: dict[str, int], + index_is_constant: np.ndarray, + n_params: int, +) -> _CompiledProfileGroups: + """Compile PROFILE_SAMPLE groups and per-sample profile expressions. + + Shared by ``schedule_2d`` (indices are trace-matrix rows; every node + has one) and ``schedule_1d`` (indices are scalar parameter slots). + 2D-only source kinds (PARAM_PLUS_TRACE, CONVOLUTION) are walked back + to the bare profile parameter name for component parsing; the + scalar-lowerable errors can only fire on the 1D path, where + trace-valued nodes have no index. + """ + + # --- PROFILE_SAMPLE groups --- + profile_sample_groups: dict[str, list[GraphNode]] = {} + for nid in topo_order: + node = id_to_node[nid] + if node.kind != NodeKind.PROFILE_SAMPLE: + continue + group_name = _profile_group_name(node.name, "profile_sample") + profile_sample_groups.setdefault(group_name, []).append(node) + + plan_aux_axis = np.zeros(0, dtype=np.float64) + n_aux = 0 + sample_base_indices_list: list[int] = [] + sample_component_indptr_list: list[int] = [0] + component_func_ids_list: list[int] = [] + component_param_indptr_list: list[int] = [0] + component_param_indices_list: list[int] = [] + sample_is_constant_list: list[bool] = [] + sample_group_idx: dict[str, int] = {} + + for group_name, sample_nodes in profile_sample_groups.items(): + sample_nodes_sorted = sorted( + sample_nodes, + key=lambda node: _profile_group_index(node.name, "profile_sample"), + ) + aux_indices = [ + _profile_group_index(node.name, "profile_sample") + for node in sample_nodes_sorted + ] + if aux_indices != list(range(len(sample_nodes_sorted))): + raise ValueError( + f"PROFILE_SAMPLE nodes for {group_name!r} do not cover " + "a contiguous aux-axis range" + ) + + aux_axis = sample_nodes_sorted[0].arrays.get("aux_axis") + if aux_axis is None: + raise ValueError(f"PROFILE_SAMPLE {group_name!r} is missing aux_axis") + aux_axis = np.asarray(aux_axis, dtype=np.float64) + if n_aux == 0: + n_aux = len(aux_axis) + plan_aux_axis = aux_axis.copy() + elif len(aux_axis) != n_aux or not np.array_equal(aux_axis, plan_aux_axis): + raise ValueError("All lowered profile groups must share one fixed aux_axis") + if len(sample_nodes_sorted) != n_aux: + raise ValueError( + f"PROFILE_SAMPLE group {group_name!r} has " + f"{len(sample_nodes_sorted)} samples but aux_axis length {n_aux}" + ) + + rep_node = sample_nodes_sorted[0] + rep_param_edges = sorted( + param_edges_by_target.get(rep_node.id, []), + key=lambda edge: edge.position or 0, + ) + if not rep_param_edges: + raise ValueError(f"PROFILE_SAMPLE {group_name!r} has no PARAM_INPUT edges") + + base_node = id_to_node[rep_param_edges[0].source] + if base_node.name not in name_to_index: + raise ValueError( + f"PROFILE_SAMPLE base source {base_node.name!r} is not scalar-lowerable" + ) + base_idx = name_to_index[base_node.name] + is_constant = bool(index_is_constant[base_idx]) + sample_base_indices_list.append(base_idx) + + component_func_by_name: dict[str, int] = {} + component_param_indices_by_name: dict[str, list[int]] = {} + component_order: list[str] = [] + for edge in rep_param_edges[1:]: + src_node = id_to_node[edge.source] + if src_node.name not in name_to_index: + raise ValueError( + f"Profile parameter source {src_node.name!r} " + "is not scalar-lowerable" + ) + # PARAM_PLUS_TRACE nodes have a "_resolved" suffix; CONVOLUTION + # nodes wrap a PPT for profile-time-dynamics and carry their own + # "__dynamics" name. In both cases, walk back to + # the underlying PPT so profile component parsing sees the + # original profile parameter name. + parse_name = src_node.name + if src_node.kind == NodeKind.PARAM_PLUS_TRACE: + parse_name = parse_name.removesuffix("_resolved") + elif src_node.kind == NodeKind.CONVOLUTION: + ppt_node = _walk_convolution_to_param_plus_trace( + src_node, graph.edges, id_to_node + ) + parse_name = ppt_node.name.removesuffix("_resolved") + comp_name, func_name = _parse_profile_component_param_name( + group_name, + parse_name, + ) + prof_func_kind = _FUNCTION_NAME_TO_PROFILE_FUNC.get(func_name) + if prof_func_kind is None: + raise ValueError(f"Unknown profile function: {func_name!r}") + + if comp_name not in component_func_by_name: + component_order.append(comp_name) + component_func_by_name[comp_name] = int(prof_func_kind) + component_param_indices_by_name[comp_name] = [] + + src_idx = name_to_index[src_node.name] + component_param_indices_by_name[comp_name].append(src_idx) + is_constant = is_constant and bool(index_is_constant[src_idx]) + + for comp_name in component_order: + component_func_ids_list.append(component_func_by_name[comp_name]) + component_param_indices_list.extend( + component_param_indices_by_name[comp_name] + ) + component_param_indptr_list.append(len(component_param_indices_list)) + sample_component_indptr_list.append(len(component_func_ids_list)) + + sample_group_idx[group_name] = len(sample_base_indices_list) - 1 + sample_is_constant_list.append(is_constant) + + sample_is_constant = np.array(sample_is_constant_list, dtype=np.bool_) + + # --- Per-sample profile expressions --- + profile_expr_groups: dict[str, list[GraphNode]] = {} + for nid in topo_order: + node = id_to_node[nid] + if _is_profile_expr_node(node): + group_name = _profile_group_name(node.name, "profile_expr") + profile_expr_groups.setdefault(group_name, []).append(node) + + expr_programs: list[ExprProgram] = [] + expr_is_constant_list: list[bool] = [] + expr_group_idx: dict[str, int] = {} + for group_name, p_expr_nodes in profile_expr_groups.items(): + p_expr_nodes_sorted = sorted( + p_expr_nodes, + key=lambda node: _profile_group_index(node.name, "profile_expr"), + ) + aux_indices = [ + _profile_group_index(node.name, "profile_expr") + for node in p_expr_nodes_sorted + ] + if aux_indices != list(range(len(p_expr_nodes_sorted))): + raise ValueError( + f"Profile expression nodes for {group_name!r} do not cover " + "a contiguous aux-axis range" + ) + if len(p_expr_nodes_sorted) != n_aux: + raise ValueError( + f"Profile expression group {group_name!r} has " + f"{len(p_expr_nodes_sorted)} samples but aux_axis length {n_aux}" + ) + + rep_node = p_expr_nodes_sorted[0] + if rep_node.expr_string is None: + raise ValueError( + f"Profile expression {group_name!r} is missing expr_string" + ) + expr_refs = set(_extract_expression_references(rep_node.expr_string)) - Returns - ------- - float - Scalar result. - """ + prof_ref_map: dict[str, int] = {} + for edge in expr_ref_edges_by_target.get(rep_node.id, []): + src_node = id_to_node[edge.source] + if src_node.kind == NodeKind.PROFILE_SAMPLE: + sample_name = _profile_group_name(src_node.name, "profile_sample") + src_idx = n_params + sample_group_idx[sample_name] + match_name = sample_name + else: + src_idx = name_to_index[src_node.name] + match_name = src_node.name + # PARAM_PLUS_TRACE nodes carry a "_resolved" suffix; + # the expression string uses the bare param name. + if match_name not in expr_refs and src_node.name.endswith("_resolved"): + match_name = src_node.name.removesuffix("_resolved") - stack: list[float] = [] - instr = program.instructions - n_instr = len(instr) // 2 + if match_name in expr_refs: + prof_ref_map[match_name] = src_idx - for i in range(n_instr): - kind = ExprNodeKind(instr[2 * i]) - operand = instr[2 * i + 1] + symbolic = compile_expr_symbolic(rep_node.expr_string) + binding: dict[str, int] = dict(name_to_index) + binding.update(prof_ref_map) + expr_programs.append(_bind_expr_to_rows(symbolic, binding)) - if kind == ExprNodeKind.CONST: - stack.append(float(np.int64(operand).view(np.float64))) - elif kind == ExprNodeKind.PARAM_REF: - stack.append(float(values[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) + is_constant = True + for name in symbolic.referenced_names: + bound_idx = int(binding[name]) + if bound_idx < n_params: + is_constant = is_constant and bool(index_is_constant[bound_idx]) + else: + is_constant = is_constant and bool( + sample_is_constant[bound_idx - n_params] + ) + expr_is_constant_list.append(is_constant) + expr_group_idx[group_name] = len(expr_programs) - 1 - assert len(stack) == 1 - return stack[0] + return _CompiledProfileGroups( + aux_axis=plan_aux_axis, + n_aux=n_aux, + n_samples=len(sample_base_indices_list), + sample_base_indices=np.array(sample_base_indices_list, dtype=np.intp), + sample_component_indptr=np.array(sample_component_indptr_list, dtype=np.intp), + component_func_ids=np.array(component_func_ids_list, dtype=np.intp), + component_param_indptr=np.array(component_param_indptr_list, dtype=np.intp), + component_param_indices=np.array(component_param_indices_list, dtype=np.intp), + sample_is_constant=sample_is_constant, + sample_group_idx=sample_group_idx, + n_exprs=len(expr_programs), + expr_programs=expr_programs, + expr_is_constant=np.array(expr_is_constant_list, dtype=np.bool_), + expr_group_idx=expr_group_idx, + ) # -def _profile_group_name(node_name: str, label: str) -> str: - """Return the shared profile group base name from a per-sample node name.""" +# +class _ScheduledOps(NamedTuple): + """Packed component-op schedule shared by both plan dataclasses.""" - match = re.fullmatch(rf"(.+)_{label}_(\d+)", node_name) - if match is None: - raise ValueError(f"Malformed profile node name: {node_name!r}") - return match.group(1) + n_ops: int + op_kinds: np.ndarray + op_param_indptr: np.ndarray + op_param_source_kinds: np.ndarray + op_param_indices: np.ndarray + op_needs_spectrum: np.ndarray + op_is_pre_spectrum: np.ndarray + op_is_profiled: np.ndarray + op_is_constant: np.ndarray # -def _profile_group_index(node_name: str, label: str) -> int: - """Return the aux-axis sample index encoded in a per-sample node name.""" - - match = re.fullmatch(rf"(.+)_{label}_(\d+)", node_name) - if match is None: - raise ValueError(f"Malformed profile node name: {node_name!r}") - return int(match.group(2)) +def _schedule_component_ops( + graph: GraphIR, + topo_order: list[int], + id_to_node: dict[int, GraphNode], + param_edges_by_target: dict[int, list[GraphEdge]], + addend_edges_by_target: dict[int, list[GraphEdge]], + spectrum_input_targets: set[int], + name_to_index: dict[str, int], + index_is_constant: np.ndarray, + profiles: _CompiledProfileGroups, +) -> _ScheduledOps: + """Schedule component ops (plain, spectrum-fed, and profiled). + + Shared by ``schedule_2d`` and ``schedule_1d``; like + ``_compile_profile_groups``, *name_to_index* holds trace-matrix rows + on the 2D path and scalar parameter slots on the 1D path, and the + "Non-scalar parameter source" errors can only fire on the 1D path. + """ + # Identify peak_sum contributors: nodes with ADDEND edges into the + # "peak_sum" SUM node (if it exists). + peak_sum_sources: set[int] = set() + peak_sum_nid = graph.node_by_name.get("peak_sum") + if peak_sum_nid is not None: + for edge in addend_edges_by_target.get(peak_sum_nid, []): + peak_sum_sources.add(edge.source) -# -def _profile_component_sample_index(node_name: str) -> int: - """Return the aux-axis sample index from ``_sample_``.""" + # Collect per-sample component inputs for PROFILE_AVERAGE nodes. + profile_avg_sample_inputs: dict[int, list[GraphNode]] = {} + sample_component_ids: set[int] = set() + for nid in topo_order: + node = id_to_node[nid] + if node.kind != NodeKind.PROFILE_AVERAGE: + continue + sample_nodes = [ + id_to_node[edge.source] + for edge in addend_edges_by_target.get(node.id, []) + if id_to_node[edge.source].kind + in (NodeKind.COMPONENT_EVAL, NodeKind.SPECTRUM_FED_OP) + ] + profile_avg_sample_inputs[node.id] = sample_nodes + sample_component_ids.update(sample.id for sample in sample_nodes) - return _profile_group_index(node_name, "sample") + comp_op_kinds = {NodeKind.COMPONENT_EVAL, NodeKind.SPECTRUM_FED_OP} + comp_nodes_topo = [ + id_to_node[nid] + for nid in topo_order + if ( + (id_to_node[nid].kind in comp_op_kinds and nid not in sample_component_ids) + or id_to_node[nid].kind == NodeKind.PROFILE_AVERAGE + ) + ] + n_aux = profiles.n_aux + op_kinds_list: list[int] = [] + op_param_indptr_list: list[int] = [0] + op_param_source_kinds_list: list[int] = [] + op_param_indices_list: list[int] = [] + op_needs_spectrum_list: list[bool] = [] + op_is_pre_spectrum_list: list[bool] = [] + op_is_profiled_list: list[bool] = [] + op_is_constant_list: list[bool] = [] -# -def _is_profile_expr_node(node: GraphNode) -> bool: - """Return True for per-sample profile-expression nodes.""" + for comp_node in comp_nodes_topo: + if comp_node.kind == NodeKind.PROFILE_AVERAGE: + # Profiled component: gather from per-sample COMPONENT_EVAL inputs. + sample_nodes = sorted( + profile_avg_sample_inputs.get(comp_node.id, []), + key=lambda node: _profile_component_sample_index(node.name), + ) + if not sample_nodes: + raise ValueError( + f"PROFILE_AVERAGE {comp_node.name!r} has no sample component inputs" + ) + if len(sample_nodes) != n_aux: + raise ValueError( + f"PROFILE_AVERAGE {comp_node.name!r} has " + f"{len(sample_nodes)} samples " + f"but aux_axis length {n_aux}" + ) - return ( - node.kind == NodeKind.EXPRESSION - and re.fullmatch(r"(.+)_profile_expr_(\d+)", node.name) is not None - ) + rep_node = sample_nodes[0] + assert rep_node.function_name is not None + op = _FUNCTION_NAME_TO_OP.get(rep_node.function_name) + if op is None: + raise ValueError( + f"Unknown component function: {rep_node.function_name!r}" + ) + op_kinds_list.append(int(op)) + op_is_profiled_list.append(True) + rep_param_edges = sorted( + param_edges_by_target.get(rep_node.id, []), + key=lambda edge: edge.position or 0, + ) + sample_param_edges = [ + sorted( + param_edges_by_target.get(sample_node.id, []), + key=lambda edge: edge.position or 0, + ) + for sample_node in sample_nodes + ] + is_constant = True + for pos, rep_edge in enumerate(rep_param_edges): + src_node = id_to_node[rep_edge.source] + if src_node.kind == NodeKind.PROFILE_SAMPLE: + group_name = _profile_group_name(src_node.name, "profile_sample") + source_kind = int(ParamSourceKind.PROFILE_SAMPLE) + source_idx = profiles.sample_group_idx[group_name] + is_constant = is_constant and bool( + profiles.sample_is_constant[source_idx] + ) + for aux_i, edges in enumerate(sample_param_edges): + sample_src = id_to_node[edges[pos].source] + if sample_src.kind != NodeKind.PROFILE_SAMPLE: + raise ValueError( + "Mixed parameter source kinds " + f"in profiled op {comp_node.name!r}" + ) + if ( + _profile_group_name(sample_src.name, "profile_sample") + != group_name + or _profile_group_index(sample_src.name, "profile_sample") + != aux_i + ): + raise ValueError( + "Inconsistent PROFILE_SAMPLE wiring " + f"in {comp_node.name!r}" + ) + elif _is_profile_expr_node(src_node): + group_name = _profile_group_name(src_node.name, "profile_expr") + source_kind = int(ParamSourceKind.PROFILE_EXPR) + source_idx = profiles.expr_group_idx[group_name] + is_constant = is_constant and bool( + profiles.expr_is_constant[source_idx] + ) + for aux_i, edges in enumerate(sample_param_edges): + sample_src = id_to_node[edges[pos].source] + if not _is_profile_expr_node(sample_src): + raise ValueError( + "Mixed expression source kinds " + f"in profiled op {comp_node.name!r}" + ) + if ( + _profile_group_name(sample_src.name, "profile_expr") + != group_name + or _profile_group_index(sample_src.name, "profile_expr") + != aux_i + ): + raise ValueError( + "Inconsistent profile-expression " + f"wiring in {comp_node.name!r}" + ) + else: + if src_node.name not in name_to_index: + raise ValueError( + f"Non-scalar parameter source {src_node.name!r} in 1D op" + ) + source_kind = int(ParamSourceKind.SCALAR) + source_idx = name_to_index[src_node.name] + is_constant = is_constant and bool(index_is_constant[source_idx]) + for edges in sample_param_edges[1:]: + if id_to_node[edges[pos].source].id != src_node.id: + raise ValueError( + "Scalar parameter source changed " + f"across samples in {comp_node.name!r}" + ) -# -def _parse_profile_component_param_name( - target_param_name: str, - source_param_name: str, -) -> tuple[str, str]: - """Parse ``__`` into component + function name.""" + op_param_source_kinds_list.append(source_kind) + op_param_indices_list.append(source_idx) - prefix = f"{target_param_name}_" - if not source_param_name.startswith(prefix): - raise ValueError( - f"Profile parameter {source_param_name!r} does not match " - f"target parameter {target_param_name!r}" - ) + op_param_indptr_list.append(len(op_param_indices_list)) + else: + # Non-profiled component: standard param wiring. + assert comp_node.function_name is not None + op = _FUNCTION_NAME_TO_OP.get(comp_node.function_name) + if op is None: + raise ValueError( + f"Unknown component function: {comp_node.function_name!r}" + ) + op_kinds_list.append(int(op)) + op_is_profiled_list.append(False) - remainder = source_param_name[len(prefix) :] - comp_name, sep, _par_name = remainder.rpartition("_") - if not sep or not comp_name: - raise ValueError(f"Malformed profile parameter name: {source_param_name!r}") + param_edges = sorted( + param_edges_by_target.get(comp_node.id, []), + key=lambda edge: edge.position or 0, + ) + is_constant = True + for edge in param_edges: + src_node = id_to_node[edge.source] + if src_node.name not in name_to_index: + raise ValueError( + f"Non-scalar parameter source {src_node.name!r} in 1D op" + ) + src_idx = name_to_index[src_node.name] + op_param_source_kinds_list.append(int(ParamSourceKind.SCALAR)) + op_param_indices_list.append(src_idx) + is_constant = is_constant and bool(index_is_constant[src_idx]) - func_name, sep2, comp_idx = comp_name.rpartition("_") - if not sep2 or not comp_idx.isdigit(): - raise ValueError(f"Malformed profile component name: {comp_name!r}") + op_param_indptr_list.append(len(op_param_indices_list)) - return comp_name, func_name + has_spec_input = comp_node.id in spectrum_input_targets + op_needs_spectrum_list.append(has_spec_input) + op_is_pre_spectrum_list.append(comp_node.id in peak_sum_sources) + op_is_constant_list.append((not has_spec_input) and is_constant) + + return _ScheduledOps( + n_ops=len(comp_nodes_topo), + op_kinds=np.array(op_kinds_list, dtype=np.intp), + op_param_indptr=np.array(op_param_indptr_list, dtype=np.intp), + op_param_source_kinds=np.array(op_param_source_kinds_list, dtype=np.int8), + op_param_indices=np.array(op_param_indices_list, dtype=np.intp), + op_needs_spectrum=np.array(op_needs_spectrum_list, dtype=np.bool_), + op_is_pre_spectrum=np.array(op_is_pre_spectrum_list, dtype=np.bool_), + op_is_profiled=np.array(op_is_profiled_list, dtype=np.bool_), + op_is_constant=np.array(op_is_constant_list, dtype=np.bool_), + ) # @@ -3273,510 +3406,154 @@ def schedule_1d(graph: GraphIR) -> ScheduledPlan1D: If the graph cannot be lowered (domain, unsupported nodes, etc.). """ - if not can_lower_1d(graph): - raise ValueError("Graph cannot be lowered to 1D backend") - - assert graph.energy is not None - - # ------------------------------------------------------------------ # - # 1. Topological sort + helper lookups # - # ------------------------------------------------------------------ # - topo_order = _topological_sort(graph) - id_to_node: dict[int, GraphNode] = {n.id: n for n in graph.nodes} - param_edges_by_target: dict[int, list[GraphEdge]] = {} - expr_ref_edges_by_target: dict[int, list[GraphEdge]] = {} - addend_edges_by_target: dict[int, list[GraphEdge]] = {} - spectrum_input_targets: set[int] = set() - for edge in graph.edges: - if edge.kind == EdgeKind.PARAM_INPUT: - param_edges_by_target.setdefault(edge.target, []).append(edge) - elif edge.kind == EdgeKind.EXPR_REF: - expr_ref_edges_by_target.setdefault(edge.target, []).append(edge) - elif edge.kind == EdgeKind.ADDEND: - addend_edges_by_target.setdefault(edge.target, []).append(edge) - elif edge.kind == EdgeKind.SPECTRUM_INPUT: - spectrum_input_targets.add(edge.target) - - # ------------------------------------------------------------------ # - # 2. Assign scalar parameter indices # - # ------------------------------------------------------------------ # - _ROW_KINDS = frozenset( - { - NodeKind.STATIC_PARAM, - NodeKind.OPT_PARAM, - NodeKind.EXPRESSION, - } - ) - - opt_nodes: list[GraphNode] = [] - static_nodes: list[GraphNode] = [] - computed_nodes: list[GraphNode] = [] - - for nid in topo_order: - node = id_to_node[nid] - if node.kind not in _ROW_KINDS or _is_profile_expr_node(node): - continue - if node.kind == NodeKind.OPT_PARAM and node.vary: - opt_nodes.append(node) - elif node.kind in (NodeKind.STATIC_PARAM, NodeKind.OPT_PARAM): - static_nodes.append(node) - else: - computed_nodes.append(node) - - all_param_nodes = opt_nodes + static_nodes + computed_nodes - n_params = len(all_param_nodes) - name_to_idx = {node.name: idx for idx, node in enumerate(all_param_nodes)} - idx_is_constant = np.zeros(n_params, dtype=np.bool_) - for node in static_nodes: - idx_is_constant[name_to_idx[node.name]] = True - - n_opt = len(opt_nodes) - opt_indices = np.arange(n_opt, dtype=np.intp) - opt_param_names = [n.name for n in opt_nodes] - - # ------------------------------------------------------------------ # - # 3. Compile scalar expressions # - # ------------------------------------------------------------------ # - expr_nodes_topo = [ - id_to_node[nid] - for nid in topo_order - if id_to_node[nid].kind == NodeKind.EXPRESSION - and not _is_profile_expr_node(id_to_node[nid]) - ] - expr_programs: list[ExprProgram] = [] - expr_target_indices_list: list[int] = [] - for expr_node in expr_nodes_topo: - assert expr_node.expr_string is not None - symbolic = compile_expr_symbolic(expr_node.expr_string) - expr_refs = set(_extract_expression_references(expr_node.expr_string)) - - ref_map: dict[str, int] = {} - for edge in expr_ref_edges_by_target.get(expr_node.id, []): - src_node = id_to_node[edge.source] - src_idx = name_to_idx[src_node.name] - if src_node.name in expr_refs: - ref_map[src_node.name] = src_idx - - binding = dict(name_to_idx) - binding.update(ref_map) - program = _bind_expr_to_rows(symbolic, binding) - expr_programs.append(program) - - target_idx = name_to_idx[expr_node.name] - expr_target_indices_list.append(target_idx) - idx_is_constant[target_idx] = all( - idx_is_constant[int(binding[name])] for name in symbolic.referenced_names - ) - - n_expressions = len(expr_programs) - expr_target_indices = np.array(expr_target_indices_list, dtype=np.intp) - - # ------------------------------------------------------------------ # - # 4. Compile PROFILE_SAMPLE groups # - # ------------------------------------------------------------------ # - profile_sample_groups: dict[str, list[GraphNode]] = {} - for nid in topo_order: - node = id_to_node[nid] - if node.kind != NodeKind.PROFILE_SAMPLE: - continue - group_name = _profile_group_name(node.name, "profile_sample") - profile_sample_groups.setdefault(group_name, []).append(node) - - plan_aux_axis = np.zeros(0, dtype=np.float64) - n_aux = 0 - profile_sample_base_indices_list: list[int] = [] - profile_sample_component_indptr_list: list[int] = [0] - profile_component_func_ids_list: list[int] = [] - profile_component_param_indptr_list: list[int] = [0] - profile_component_param_indices_list: list[int] = [] - profile_sample_is_constant_list: list[bool] = [] - profile_sample_group_idx: dict[str, int] = {} - - for group_name, sample_nodes in profile_sample_groups.items(): - sample_nodes_sorted = sorted( - sample_nodes, - key=lambda node: _profile_group_index(node.name, "profile_sample"), - ) - aux_indices = [ - _profile_group_index(node.name, "profile_sample") - for node in sample_nodes_sorted - ] - if aux_indices != list(range(len(sample_nodes_sorted))): - raise ValueError( - f"PROFILE_SAMPLE nodes for {group_name!r} do not cover " - "a contiguous aux-axis range" - ) - - aux_axis = sample_nodes_sorted[0].arrays.get("aux_axis") - if aux_axis is None: - raise ValueError(f"PROFILE_SAMPLE {group_name!r} is missing aux_axis") - aux_axis = np.asarray(aux_axis, dtype=np.float64) - if n_aux == 0: - n_aux = len(aux_axis) - plan_aux_axis = aux_axis.copy() - elif len(aux_axis) != n_aux or not np.array_equal(aux_axis, plan_aux_axis): - raise ValueError("All lowered profile groups must share one fixed aux_axis") - if len(sample_nodes_sorted) != n_aux: - raise ValueError( - f"PROFILE_SAMPLE group {group_name!r} has {len(sample_nodes_sorted)} " - f"samples but aux_axis length {n_aux}" - ) - - rep_node = sample_nodes_sorted[0] - param_edges = sorted( - param_edges_by_target.get(rep_node.id, []), - key=lambda edge: edge.position or 0, - ) - if not param_edges: - raise ValueError(f"PROFILE_SAMPLE {group_name!r} has no PARAM_INPUT edges") - - base_node = id_to_node[param_edges[0].source] - if base_node.name not in name_to_idx: - raise ValueError( - f"PROFILE_SAMPLE base source {base_node.name!r} is not scalar-lowerable" - ) - base_idx = name_to_idx[base_node.name] - is_constant = bool(idx_is_constant[base_idx]) - profile_sample_base_indices_list.append(base_idx) - - component_func_by_name: dict[str, int] = {} - component_param_indices_by_name: dict[str, list[int]] = {} - component_order: list[str] = [] - for edge in param_edges[1:]: - src_node = id_to_node[edge.source] - if src_node.name not in name_to_idx: - raise ValueError( - f"Profile parameter source {src_node.name!r} " - "is not scalar-lowerable" - ) - comp_name, func_name = _parse_profile_component_param_name( - group_name, - src_node.name, - ) - prof_func_kind = _FUNCTION_NAME_TO_PROFILE_FUNC.get(func_name) - if prof_func_kind is None: - raise ValueError(f"Unknown profile function: {func_name!r}") - - if comp_name not in component_func_by_name: - component_order.append(comp_name) - component_func_by_name[comp_name] = int(prof_func_kind) - component_param_indices_by_name[comp_name] = [] - - src_idx = name_to_idx[src_node.name] - component_param_indices_by_name[comp_name].append(src_idx) - is_constant = is_constant and bool(idx_is_constant[src_idx]) - - for comp_name in component_order: - profile_component_func_ids_list.append(component_func_by_name[comp_name]) - profile_component_param_indices_list.extend( - component_param_indices_by_name[comp_name] - ) - profile_component_param_indptr_list.append( - len(profile_component_param_indices_list) - ) - profile_sample_component_indptr_list.append( - len(profile_component_func_ids_list) - ) - - group_idx = len(profile_sample_base_indices_list) - 1 - profile_sample_group_idx[group_name] = group_idx - profile_sample_is_constant_list.append(is_constant) - - n_profile_samples = len(profile_sample_base_indices_list) - profile_sample_base_indices = np.array( - profile_sample_base_indices_list, dtype=np.intp - ) - profile_sample_component_indptr = np.array( - profile_sample_component_indptr_list, dtype=np.intp - ) - profile_component_func_ids = np.array( - profile_component_func_ids_list, dtype=np.intp - ) - profile_component_param_indptr = np.array( - profile_component_param_indptr_list, dtype=np.intp - ) - profile_component_param_indices = np.array( - profile_component_param_indices_list, dtype=np.intp - ) - profile_sample_is_constant = np.array( - profile_sample_is_constant_list, - dtype=np.bool_, - ) - - # ------------------------------------------------------------------ # - # 5. Compile per-sample profile expressions # - # ------------------------------------------------------------------ # - profile_expr_groups: dict[str, list[GraphNode]] = {} - for nid in topo_order: - node = id_to_node[nid] - if _is_profile_expr_node(node): - group_name = _profile_group_name(node.name, "profile_expr") - profile_expr_groups.setdefault(group_name, []).append(node) - - profile_expr_programs: list[ExprProgram] = [] - profile_expr_is_constant_list: list[bool] = [] - profile_expr_group_idx: dict[str, int] = {} - for group_name, expr_nodes in profile_expr_groups.items(): - expr_nodes_sorted = sorted( - expr_nodes, - key=lambda node: _profile_group_index(node.name, "profile_expr"), - ) - aux_indices = [ - _profile_group_index(node.name, "profile_expr") - for node in expr_nodes_sorted - ] - if aux_indices != list(range(len(expr_nodes_sorted))): - raise ValueError( - f"Profile expression nodes for {group_name!r} do not cover " - "a contiguous aux-axis range" - ) - if len(expr_nodes_sorted) != n_aux: - raise ValueError( - f"Profile expression group {group_name!r} has {len(expr_nodes_sorted)} " - f"samples but aux_axis length {n_aux}" - ) - - rep_node = expr_nodes_sorted[0] - if rep_node.expr_string is None: - raise ValueError( - f"Profile expression {group_name!r} is missing expr_string" - ) - expr_refs = set(_extract_expression_references(rep_node.expr_string)) - - prof_ref_map: dict[str, int] = {} - for edge in expr_ref_edges_by_target.get(rep_node.id, []): - src_node = id_to_node[edge.source] - if src_node.kind == NodeKind.PROFILE_SAMPLE: - sample_name = _profile_group_name(src_node.name, "profile_sample") - src_idx = n_params + profile_sample_group_idx[sample_name] - match_name = sample_name - else: - src_idx = name_to_idx[src_node.name] - match_name = src_node.name - - if match_name in expr_refs: - prof_ref_map[match_name] = src_idx - - symbolic = compile_expr_symbolic(rep_node.expr_string) - binding = dict(name_to_idx) - binding.update(prof_ref_map) - program = _bind_expr_to_rows(symbolic, binding) - profile_expr_programs.append(program) + if not can_lower_1d(graph): + raise ValueError("Graph cannot be lowered to 1D backend") - is_constant = True - for name in symbolic.referenced_names: - bound_idx = int(binding[name]) - if bound_idx < n_params: - is_constant = is_constant and bool(idx_is_constant[bound_idx]) - else: - is_constant = is_constant and bool( - profile_sample_is_constant[bound_idx - n_params] - ) - profile_expr_is_constant_list.append(is_constant) - profile_expr_group_idx[group_name] = len(profile_expr_programs) - 1 + assert graph.energy is not None - n_profile_exprs = len(profile_expr_programs) - profile_expr_is_constant = np.array(profile_expr_is_constant_list, dtype=np.bool_) + # ------------------------------------------------------------------ # + # 1. Topological sort + helper lookups # + # ------------------------------------------------------------------ # + topo_order = _topological_sort(graph) + id_to_node: dict[int, GraphNode] = {n.id: n for n in graph.nodes} + param_edges_by_target: dict[int, list[GraphEdge]] = {} + expr_ref_edges_by_target: dict[int, list[GraphEdge]] = {} + addend_edges_by_target: dict[int, list[GraphEdge]] = {} + spectrum_input_targets: set[int] = set() + for edge in graph.edges: + if edge.kind == EdgeKind.PARAM_INPUT: + param_edges_by_target.setdefault(edge.target, []).append(edge) + elif edge.kind == EdgeKind.EXPR_REF: + expr_ref_edges_by_target.setdefault(edge.target, []).append(edge) + elif edge.kind == EdgeKind.ADDEND: + addend_edges_by_target.setdefault(edge.target, []).append(edge) + elif edge.kind == EdgeKind.SPECTRUM_INPUT: + spectrum_input_targets.add(edge.target) # ------------------------------------------------------------------ # - # 6. Schedule component ops # + # 2. Assign scalar parameter indices # # ------------------------------------------------------------------ # - peak_sum_sources: set[int] = set() - peak_sum_nid = graph.node_by_name.get("peak_sum") - if peak_sum_nid is not None: - for edge in addend_edges_by_target.get(peak_sum_nid, []): - peak_sum_sources.add(edge.source) + _ROW_KINDS = frozenset( + { + NodeKind.STATIC_PARAM, + NodeKind.OPT_PARAM, + NodeKind.EXPRESSION, + } + ) + + opt_nodes: list[GraphNode] = [] + static_nodes: list[GraphNode] = [] + computed_nodes: list[GraphNode] = [] - profile_avg_sample_inputs: dict[int, list[GraphNode]] = {} - sample_component_ids: set[int] = set() for nid in topo_order: node = id_to_node[nid] - if node.kind != NodeKind.PROFILE_AVERAGE: + if node.kind not in _ROW_KINDS or _is_profile_expr_node(node): continue - sample_nodes = [ - id_to_node[edge.source] - for edge in addend_edges_by_target.get(node.id, []) - if id_to_node[edge.source].kind - in (NodeKind.COMPONENT_EVAL, NodeKind.SPECTRUM_FED_OP) - ] - profile_avg_sample_inputs[node.id] = sample_nodes - sample_component_ids.update(sample.id for sample in sample_nodes) + if node.kind == NodeKind.OPT_PARAM and node.vary: + opt_nodes.append(node) + elif node.kind in (NodeKind.STATIC_PARAM, NodeKind.OPT_PARAM): + static_nodes.append(node) + else: + computed_nodes.append(node) - comp_nodes_topo = [ + all_param_nodes = opt_nodes + static_nodes + computed_nodes + n_params = len(all_param_nodes) + name_to_idx = {node.name: idx for idx, node in enumerate(all_param_nodes)} + idx_is_constant = np.zeros(n_params, dtype=np.bool_) + for node in static_nodes: + idx_is_constant[name_to_idx[node.name]] = True + + n_opt = len(opt_nodes) + opt_indices = np.arange(n_opt, dtype=np.intp) + opt_param_names = [n.name for n in opt_nodes] + + # ------------------------------------------------------------------ # + # 3. Compile scalar expressions # + # ------------------------------------------------------------------ # + expr_nodes_topo = [ id_to_node[nid] for nid in topo_order - if ( - ( - id_to_node[nid].kind - in (NodeKind.COMPONENT_EVAL, NodeKind.SPECTRUM_FED_OP) - and nid not in sample_component_ids - ) - or id_to_node[nid].kind == NodeKind.PROFILE_AVERAGE - ) + if id_to_node[nid].kind == NodeKind.EXPRESSION + and not _is_profile_expr_node(id_to_node[nid]) ] + expr_programs: list[ExprProgram] = [] + expr_target_indices_list: list[int] = [] + for expr_node in expr_nodes_topo: + assert expr_node.expr_string is not None + symbolic = compile_expr_symbolic(expr_node.expr_string) + expr_refs = set(_extract_expression_references(expr_node.expr_string)) - op_kinds_list: list[int] = [] - op_param_indptr_list: list[int] = [0] - op_param_source_kinds_list: list[int] = [] - op_param_indices_list: list[int] = [] - op_needs_spectrum_list: list[bool] = [] - op_is_pre_spectrum_list: list[bool] = [] - op_is_profiled_list: list[bool] = [] - op_is_constant_list: list[bool] = [] - - for comp_node in comp_nodes_topo: - if comp_node.kind == NodeKind.PROFILE_AVERAGE: - sample_nodes = sorted( - profile_avg_sample_inputs.get(comp_node.id, []), - key=lambda node: _profile_component_sample_index(node.name), - ) - if not sample_nodes: - raise ValueError( - f"PROFILE_AVERAGE {comp_node.name!r} has no sample component inputs" - ) - if len(sample_nodes) != n_aux: - raise ValueError( - f"PROFILE_AVERAGE {comp_node.name!r} has " - f"{len(sample_nodes)} samples " - f"but aux_axis length {n_aux}" - ) - - rep_node = sample_nodes[0] - assert rep_node.function_name is not None - op = _FUNCTION_NAME_TO_OP.get(rep_node.function_name) - if op is None: - raise ValueError( - f"Unknown component function: {rep_node.function_name!r}" - ) - op_kinds_list.append(int(op)) - op_is_profiled_list.append(True) - - rep_param_edges = sorted( - param_edges_by_target.get(rep_node.id, []), - key=lambda edge: edge.position or 0, - ) - sample_param_edges = [ - sorted( - param_edges_by_target.get(sample_node.id, []), - key=lambda edge: edge.position or 0, - ) - for sample_node in sample_nodes - ] - is_constant = True - for pos, rep_edge in enumerate(rep_param_edges): - src_node = id_to_node[rep_edge.source] - if src_node.kind == NodeKind.PROFILE_SAMPLE: - group_name = _profile_group_name(src_node.name, "profile_sample") - source_kind = int(ParamSourceKind.PROFILE_SAMPLE) - source_idx = profile_sample_group_idx[group_name] - is_constant = is_constant and bool( - profile_sample_is_constant[source_idx] - ) - for aux_i, edges in enumerate(sample_param_edges): - sample_src = id_to_node[edges[pos].source] - if sample_src.kind != NodeKind.PROFILE_SAMPLE: - raise ValueError( - "Mixed parameter source kinds " - f"in profiled op {comp_node.name!r}" - ) - if ( - _profile_group_name(sample_src.name, "profile_sample") - != group_name - or _profile_group_index(sample_src.name, "profile_sample") - != aux_i - ): - raise ValueError( - "Inconsistent PROFILE_SAMPLE wiring " - f"in {comp_node.name!r}" - ) - elif _is_profile_expr_node(src_node): - group_name = _profile_group_name(src_node.name, "profile_expr") - source_kind = int(ParamSourceKind.PROFILE_EXPR) - source_idx = profile_expr_group_idx[group_name] - is_constant = is_constant and bool( - profile_expr_is_constant[source_idx] - ) - for aux_i, edges in enumerate(sample_param_edges): - sample_src = id_to_node[edges[pos].source] - if not _is_profile_expr_node(sample_src): - raise ValueError( - "Mixed expression source kinds " - f"in profiled op {comp_node.name!r}" - ) - if ( - _profile_group_name(sample_src.name, "profile_expr") - != group_name - or _profile_group_index(sample_src.name, "profile_expr") - != aux_i - ): - raise ValueError( - "Inconsistent profile-expression " - f"wiring in {comp_node.name!r}" - ) - else: - if src_node.name not in name_to_idx: - raise ValueError( - f"Non-scalar parameter source {src_node.name!r} in 1D op" - ) - source_kind = int(ParamSourceKind.SCALAR) - source_idx = name_to_idx[src_node.name] - is_constant = is_constant and bool(idx_is_constant[source_idx]) - for edges in sample_param_edges[1:]: - if id_to_node[edges[pos].source].id != src_node.id: - raise ValueError( - "Scalar parameter source changed " - f"across samples in {comp_node.name!r}" - ) - - op_param_source_kinds_list.append(source_kind) - op_param_indices_list.append(source_idx) + ref_map: dict[str, int] = {} + for edge in expr_ref_edges_by_target.get(expr_node.id, []): + src_node = id_to_node[edge.source] + src_idx = name_to_idx[src_node.name] + if src_node.name in expr_refs: + ref_map[src_node.name] = src_idx - op_param_indptr_list.append(len(op_param_indices_list)) - else: - assert comp_node.function_name is not None - op = _FUNCTION_NAME_TO_OP.get(comp_node.function_name) - if op is None: - raise ValueError( - f"Unknown component function: {comp_node.function_name!r}" - ) - op_kinds_list.append(int(op)) - op_is_profiled_list.append(False) + binding = dict(name_to_idx) + binding.update(ref_map) + program = _bind_expr_to_rows(symbolic, binding) + expr_programs.append(program) - param_edges = sorted( - param_edges_by_target.get(comp_node.id, []), - key=lambda edge: edge.position or 0, - ) - is_constant = True - for edge in param_edges: - src_node = id_to_node[edge.source] - if src_node.name not in name_to_idx: - raise ValueError( - f"Non-scalar parameter source {src_node.name!r} in 1D op" - ) - src_idx = name_to_idx[src_node.name] - op_param_source_kinds_list.append(int(ParamSourceKind.SCALAR)) - op_param_indices_list.append(src_idx) - is_constant = is_constant and bool(idx_is_constant[src_idx]) + target_idx = name_to_idx[expr_node.name] + expr_target_indices_list.append(target_idx) + idx_is_constant[target_idx] = all( + idx_is_constant[int(binding[name])] for name in symbolic.referenced_names + ) - op_param_indptr_list.append(len(op_param_indices_list)) + n_expressions = len(expr_programs) + expr_target_indices = np.array(expr_target_indices_list, dtype=np.intp) - has_spec_input = comp_node.id in spectrum_input_targets - op_needs_spectrum_list.append(has_spec_input) - op_is_pre_spectrum_list.append(comp_node.id in peak_sum_sources) - op_is_constant_list.append((not has_spec_input) and is_constant) + # ------------------------------------------------------------------ # + # 4. Compile profile groups (samples + expressions) # + # ------------------------------------------------------------------ # + profiles = _compile_profile_groups( + graph, + topo_order, + id_to_node, + param_edges_by_target, + expr_ref_edges_by_target, + name_to_idx, + idx_is_constant, + n_params, + ) + plan_aux_axis = profiles.aux_axis + n_aux = profiles.n_aux + n_profile_samples = profiles.n_samples + profile_sample_base_indices = profiles.sample_base_indices + profile_sample_component_indptr = profiles.sample_component_indptr + profile_component_func_ids = profiles.component_func_ids + 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 - n_ops = len(comp_nodes_topo) - op_kinds = np.array(op_kinds_list, dtype=np.intp) - op_param_indptr = np.array(op_param_indptr_list, dtype=np.intp) - op_param_source_kinds = np.array(op_param_source_kinds_list, dtype=np.int8) - op_param_indices = np.array(op_param_indices_list, dtype=np.intp) - op_needs_spectrum = np.array(op_needs_spectrum_list, dtype=np.bool_) - op_is_pre_spectrum = np.array(op_is_pre_spectrum_list, dtype=np.bool_) - op_is_profiled = np.array(op_is_profiled_list, dtype=np.bool_) - op_is_constant = np.array(op_is_constant_list, dtype=np.bool_) + # ------------------------------------------------------------------ # + # 5. Schedule component ops # + # ------------------------------------------------------------------ # + ops = _schedule_component_ops( + graph, + topo_order, + id_to_node, + param_edges_by_target, + addend_edges_by_target, + spectrum_input_targets, + name_to_idx, + idx_is_constant, + profiles, + ) + n_ops = ops.n_ops + op_kinds = ops.op_kinds + op_param_indptr = ops.op_param_indptr + op_param_source_kinds = ops.op_param_source_kinds + op_param_indices = ops.op_param_indices + op_needs_spectrum = ops.op_needs_spectrum + op_is_pre_spectrum = ops.op_is_pre_spectrum + op_is_profiled = ops.op_is_profiled + op_is_constant = ops.op_is_constant # ------------------------------------------------------------------ # - # 7. Initialize scalar + profile values # + # 6. Initialize scalar + profile values # # ------------------------------------------------------------------ # param_values_init = np.zeros(n_params, dtype=np.float64) for node in opt_nodes + static_nodes: @@ -3807,7 +3584,7 @@ def schedule_1d(graph: GraphIR) -> ScheduledPlan1D: ) # ------------------------------------------------------------------ # - # 8. Precompute constant component contributions # + # 7. Precompute constant component contributions # # ------------------------------------------------------------------ # energy = graph.energy cached_result = np.zeros(len(energy), dtype=np.float64) @@ -3837,7 +3614,7 @@ def schedule_1d(graph: GraphIR) -> ScheduledPlan1D: cached_peak_sum += component # ------------------------------------------------------------------ # - # 9. Pack into ScheduledPlan1D # + # 8. Pack into ScheduledPlan1D # # ------------------------------------------------------------------ # return ScheduledPlan1D( energy=energy, From c50488c9440fd8cbafaec8034cad6d5b3f6b483c Mon Sep 17 00:00:00 2001 From: Johannes Mahl Date: Fri, 10 Jul 2026 12:47:25 -0700 Subject: [PATCH 24/36] dedupe simulator noise generation and HDF5 metadata serialization --- src/trspecfit/simulator.py | 174 ++++++++++++------------------------- 1 file changed, 54 insertions(+), 120 deletions(-) diff --git a/src/trspecfit/simulator.py b/src/trspecfit/simulator.py index 5c76cfb..f8c108f 100644 --- a/src/trspecfit/simulator.py +++ b/src/trspecfit/simulator.py @@ -409,13 +409,9 @@ def add_noise( if self.detection == "analog": # Use traditional noise addition - if dim == 1: - noise = self._generate_noise_analog_1d(clean_data) - elif dim == 2: - noise = self._generate_noise_analog_2d(clean_data) - else: + if dim not in (1, 2): raise ValueError(f"dim must be 1 or 2, got {dim}") - + noise = self._generate_noise_analog(clean_data) noisy_data = clean_data + noise elif self.detection == "photon_counting": @@ -766,9 +762,9 @@ def simulate_n( return clean_data, noisy_data_list, noise_list # - def _generate_noise_analog_1d(self, signal: np.ndarray) -> np.ndarray: + def _generate_noise_analog(self, signal: np.ndarray) -> np.ndarray: """ - Generate 1D noise array for analog detectors. + Generate noise array for analog detectors (1D or 2D signal). Parameters ---------- @@ -813,53 +809,6 @@ def _generate_noise_analog_1d(self, signal: np.ndarray) -> np.ndarray: "Use 'poisson', 'gaussian', or 'none'" ) - # - def _generate_noise_analog_2d(self, signal: np.ndarray) -> np.ndarray: - """ - Generate 2D noise array for analog detectors. - - Parameters - ---------- - signal : ndarray - Clean 2D signal array. - - Returns - ------- - ndarray - 2D noise array with same shape as signal. - """ - - if self.noise_type == "none": - return np.zeros_like(signal) - - if self.noise_type == "gaussian": - # Gaussian noise with amplitude proportional to noise_level - noise_amp = self.noise_level * np.max(np.abs(signal)) - return cast("np.ndarray", self.rng.normal(0, noise_amp, signal.shape)) - - if self.noise_type == "poisson": - # Poisson noise: scale signal to photon counts, add noise, scale back - signal_positive = np.abs(signal) - - # Scale signal to photon counts - scale_factor = 1.0 / (self.noise_level + 1e-10) - signal_scaled = signal_positive * scale_factor - - # Generate Poisson noise - noisy_scaled = self.rng.poisson(signal_scaled) - - # Scale back and compute noise component - signal_noisy = noisy_scaled / scale_factor - noise = signal_noisy - signal_positive - - # Restore original sign - return cast("np.ndarray", noise * np.sign(signal)) - - raise ValueError( - f"Unknown noise type: {self.noise_type}. " - "Use 'poisson', 'gaussian', or 'none'" - ) - # def _sample_photons_1d(self, signal: np.ndarray) -> np.ndarray: """ @@ -1606,6 +1555,50 @@ def save_data( if show_output >= 1: print(f"Data saved to: {filepath}") + # + def _write_axes_hdf5(self, f: h5py.File) -> None: + """Write energy and time axes at the file root (empty time if 1D).""" + + f.create_dataset("energy", data=self.model.energy) + if self.model.time is not None and len(self.model.time) > 0: + f.create_dataset("time", data=self.model.time) + else: + f.create_dataset("time", data=np.array([])) + + # + def _write_detection_metadata(self, meta: h5py.Group) -> None: + """Write detection/noise settings and seed as metadata attributes.""" + + meta.attrs["detection"] = self.detection + if self.detection == "analog": + meta.attrs["noise_level"] = self.noise_level + meta.attrs["noise_type"] = self.noise_type + elif self.detection == "photon_counting": + meta.attrs["counts_per_delay"] = self.counts_per_delay + if self.count_rate is not None: + meta.attrs["count_rate"] = self.count_rate + if self.integration_time is not None: + meta.attrs["integration_time"] = self.integration_time + + if self.seed is not None: + meta.attrs["seed"] = self.seed + + # + def _model_parameters_json(self) -> str: + """Serialize the model's lmfit parameters (full spec) to JSON.""" + + params_dict = {} + for par_name in self.model.lmfit_pars: + par = self.model.lmfit_pars[par_name] + params_dict[par_name] = { + "value": float(par.value), + "vary": bool(par.vary), + "min": float(par.min) if par.min is not None else None, + "max": float(par.max) if par.max is not None else None, + "expr": par.expr or None, + } + return json.dumps(params_dict, indent=2) + # def _save_hdf5(self, filepath: str, n_data: list[np.ndarray] | None = None) -> None: """ @@ -1623,14 +1616,7 @@ def _save_hdf5(self, filepath: str, n_data: list[np.ndarray] | None = None) -> N with h5py.File(filepath, "w") as f: # Save axes at root level - f.create_dataset("energy", data=self.model.energy) - - # Handle time axis - check if it exists and has valid data - if self.model.time is not None and len(self.model.time) > 0: - f.create_dataset("time", data=self.model.time) - else: - # For 1D simulations, save empty array - f.create_dataset("time", data=np.array([])) + self._write_axes_hdf5(f) # Save clean data at root level if self.data_clean is not None: @@ -1651,21 +1637,7 @@ def _save_hdf5(self, filepath: str, n_data: list[np.ndarray] | None = None) -> N # Save metadata group at root level meta = f.create_group("metadata") - meta.attrs["detection"] = self.detection - - # Save detection-specific parameters - if self.detection == "analog": - meta.attrs["noise_level"] = self.noise_level - meta.attrs["noise_type"] = self.noise_type - elif self.detection == "photon_counting": - meta.attrs["counts_per_delay"] = self.counts_per_delay - if self.count_rate is not None: - meta.attrs["count_rate"] = self.count_rate - if self.integration_time is not None: - meta.attrs["integration_time"] = self.integration_time - - if self.seed is not None: - meta.attrs["seed"] = self.seed + self._write_detection_metadata(meta) # Determine dimensionality from clean data if self.data_clean is not None: @@ -1681,19 +1653,7 @@ def _save_hdf5(self, filepath: str, n_data: list[np.ndarray] | None = None) -> N meta.attrs["n_datasets"] = 1 # Save model parameters as JSON in metadata - params_dict = {} - for par_name in self.model.lmfit_pars: - par = self.model.lmfit_pars[par_name] - params_dict[par_name] = { - "value": float(par.value), - "vary": bool(par.vary), - "min": float(par.min) if par.min is not None else None, - "max": float(par.max) if par.max is not None else None, - "expr": par.expr or None, - } - - # Save as JSON string in metadata - meta.attrs["model_parameters"] = json.dumps(params_dict, indent=2) + meta.attrs["model_parameters"] = self._model_parameters_json() meta.attrs["model_name"] = self.model.name # @@ -1894,11 +1854,7 @@ def _initialize_sweep_hdf5( with h5py.File(filepath, "w") as f: # Save axes (same for all configs) - f.create_dataset("energy", data=self.model.energy) - if self.model.time is not None and len(self.model.time) > 0: - f.create_dataset("time", data=self.model.time) - else: - f.create_dataset("time", data=np.array([])) + self._write_axes_hdf5(f) # Create groups for organization f.create_group("parameter_configs") @@ -1911,19 +1867,7 @@ def _initialize_sweep_hdf5( meta.attrs["total_datasets"] = n_configs * n_realizations # Simulator settings - meta.attrs["detection"] = self.detection - if self.detection == "analog": - meta.attrs["noise_level"] = self.noise_level - meta.attrs["noise_type"] = self.noise_type - elif self.detection == "photon_counting": - meta.attrs["counts_per_delay"] = self.counts_per_delay - if self.count_rate is not None: - meta.attrs["count_rate"] = self.count_rate - if self.integration_time is not None: - meta.attrs["integration_time"] = self.integration_time - - if self.seed is not None: - meta.attrs["seed"] = self.seed + self._write_detection_metadata(meta) # Parameter sweep settings meta.attrs["sweep_strategy"] = parameter_sweep.strategy @@ -1943,17 +1887,7 @@ def _initialize_sweep_hdf5( meta.attrs["parameter_space"] = json.dumps(param_space, indent=2) # Save full model definition once (static across all configs) - model_params = {} - for par_name in self.model.lmfit_pars: - par = self.model.lmfit_pars[par_name] - model_params[par_name] = { - "value": float(par.value), - "vary": bool(par.vary), - "min": float(par.min) if par.min is not None else None, - "max": float(par.max) if par.max is not None else None, - "expr": par.expr or None, - } - meta.attrs["model_parameters"] = json.dumps(model_params, indent=2) + meta.attrs["model_parameters"] = self._model_parameters_json() meta.attrs["model_name"] = self.model.name # Dimension matches the actual data written, not model capability From 8d5c6c42dd741e59244a798a702d9f84cf34072d Mon Sep 17 00:00:00 2001 From: Johannes Mahl Date: Fri, 10 Jul 2026 12:57:46 -0700 Subject: [PATCH 25/36] migrate fit_io reads to shared hdf5 helpers; scope contract to reads --- docs/design/repo_architecture.md | 10 ++-- src/trspecfit/utils/fit_io.py | 90 ++++++++++++++------------------ 2 files changed, 46 insertions(+), 54 deletions(-) diff --git a/docs/design/repo_architecture.md b/docs/design/repo_architecture.md index b7e34d1..d325f43 100644 --- a/docs/design/repo_architecture.md +++ b/docs/design/repo_architecture.md @@ -242,8 +242,12 @@ own `scipy.signal.convolve` wrapper. ### `utils/hdf5.py` Typed HDF5 helpers. `require_group`, `require_dataset`, `json_loads_attr`. -All HDF5 I/O in the repo should go through these rather than raw -`h5py` calls — they normalize attribute types across numpy/bytes/str. +All HDF5 *reads* should go through these rather than raw subscripting — +`require_*` narrow `Group | Dataset` lookups with a clear error naming the +archive path, and `json_loads_attr` normalizes JSON attributes across +numpy/bytes/str. Write-side calls (`h5py.File(...)`, `create_group`, +`create_dataset`, attribute assignment) have no wrapper and use `h5py` +directly. ### `utils/fit_io.py` @@ -327,5 +331,5 @@ evaluator. New features are generally prototyped on that slow path first. - **New fit-result post-processing (CI, MCMC, in-fit plots)** → `fitlib.py`. - **New fit-archive field, exporter format, or comparison metric** → `utils/fit_io.py` (data model + writer/reader + CSV exporter) and `fit_results.py` (query / `compare_models`). Slot extraction stays in `utils/fit_io.py`; the four `_append__slot` call sites in `trspecfit.py` should not be replicated elsewhere. - **New simulator feature / sampling strategy** → `simulator.py` / `utils/sweep.py`. -- **New HDF5 I/O** → go through `utils/hdf5.py` helpers. +- **New HDF5 reads** → go through `utils/hdf5.py` helpers (writes use `h5py` directly). - **Performance optimization of an existing feature** → lower into `graph_ir` / `eval_*`. Do **not** optimize mcp. diff --git a/src/trspecfit/utils/fit_io.py b/src/trspecfit/utils/fit_io.py index e802059..fe823d7 100644 --- a/src/trspecfit/utils/fit_io.py +++ b/src/trspecfit/utils/fit_io.py @@ -41,6 +41,7 @@ plt_fit_res_2d, plt_fit_res_pars, ) +from trspecfit.utils.hdf5 import require_dataset, require_group FitType = Literal["baseline", "spectrum", "sbs", "2d"] SCHEMA_VERSION = "2" @@ -141,32 +142,6 @@ def _compute_sigma_eff( return float(sigma_data) -# -def _as_group(obj: Any) -> h5py.Group: - """ - Narrow an h5py lookup result (``Group | Dataset | Datatype | Link``) to - ``h5py.Group``, raising if it isn't one. Used at archive-traversal sites - to give pyright a stable type without sprinkling ``cast`` everywhere. - """ - - if not isinstance(obj, h5py.Group): - raise TypeError( - f"expected h5py.Group at archive path, got {type(obj).__name__}" - ) - return obj - - -# -def _as_dataset(obj: Any) -> h5py.Dataset: - """``Dataset`` counterpart to :func:`_as_group` for read-side lookups.""" - - if not isinstance(obj, h5py.Dataset): - raise TypeError( - f"expected h5py.Dataset at archive path, got {type(obj).__name__}" - ) - return obj - - # @dataclass(frozen=True) class SavedFitSlot: @@ -901,10 +876,10 @@ def _find_file_by_fingerprint( files_obj = archive.get("files") if files_obj is None: return None - files_group = _as_group(files_obj) + files_group = require_group(files_obj, "files") for key in sorted(files_group.keys()): - fg = _as_group(files_group[key]) - meta = _as_group(fg["metadata"]) + fg = require_group(files_group[key], f"files/{key}") + meta = require_group(fg["metadata"], f"files/{key}/metadata") if str(meta.attrs.get("data_sha256", "")) != fingerprint["data_sha256"]: continue if str(meta.attrs.get("energy_sha256", "")) != fingerprint["energy_sha256"]: @@ -940,13 +915,13 @@ def _find_slot_by_archive_key( slots_obj = file_group.get("slots") if slots_obj is None: return None - slots_group = _as_group(slots_obj) + slots_group = require_group(slots_obj, "slots") for key in sorted(slots_group.keys()): - slot_group = _as_group(slots_group[key]) + slot_group = require_group(slots_group[key], f"slots/{key}") meta_obj = slot_group.get("metadata") if meta_obj is None: continue - meta = _as_group(meta_obj) + meta = require_group(meta_obj, f"slots/{key}/metadata") if str(meta.attrs.get("archive_slot_key", "")) == archive_slot_key: return slot_group return None @@ -1190,7 +1165,7 @@ def _classify_archive_for_write( """ meta_obj = archive.get("metadata") - meta = _as_group(meta_obj) if meta_obj is not None else None + meta = require_group(meta_obj, "metadata") if meta_obj is not None else None if meta is not None and "schema_version" in meta.attrs: existing = str(meta.attrs["schema_version"]) if existing != project.schema_version: @@ -1530,9 +1505,13 @@ def _read_mcmc_group(group: h5py.Group) -> dict[str, Any]: if flatchain_obj is None: flatchain: pd.DataFrame | None = None else: - flatchain = _decode_dataframe(_as_dataset(flatchain_obj)) + flatchain = _decode_dataframe(require_dataset(flatchain_obj, "mcmc/flatchain")) ci_obj = group.get("ci") - ci = _decode_dataframe(_as_dataset(ci_obj)) if ci_obj is not None else None + ci = ( + _decode_dataframe(require_dataset(ci_obj, "mcmc/ci")) + if ci_obj is not None + else None + ) lnsigma_attr = group.attrs.get("lnsigma") lnsigma: float | None if lnsigma_attr is None: @@ -1552,33 +1531,41 @@ def _read_slot( ) -> SavedFitSlot: """Decode one slot group into a ``SavedFitSlot``.""" - meta = _as_group(slot_group["metadata"]) + meta = require_group(slot_group["metadata"], "metadata") a = meta.attrs fit_type = cast(FitType, _attr_str(a["fit_type"])) selection_json = _attr_str(a["selection_json"]) model_name = _attr_str(a["model_name"]) - params = _decode_dataframe(_as_dataset(slot_group["params"])) + params = _decode_dataframe(require_dataset(slot_group["params"], "params")) if fit_type != "sbs": # Restore the schema's "" ↔ None / NaN ↔ None mappings for long-form # params. sbs params is wide-form numeric and carries no None # semantics, so this only applies to baseline / spectrum / 2d. _restore_long_params_nones(params) - observed = np.asarray(_as_dataset(slot_group["observed"])[...]) - fit_arr = np.asarray(_as_dataset(slot_group["fit"])[...]) + observed = np.asarray(require_dataset(slot_group["observed"], "observed")[...]) + fit_arr = np.asarray(require_dataset(slot_group["fit"], "fit")[...]) metrics: dict[str, Any] if fit_type == "sbs": - metrics = _read_metrics_per_slice(_as_dataset(slot_group["metrics_per_slice"])) + metrics = _read_metrics_per_slice( + require_dataset(slot_group["metrics_per_slice"], "metrics_per_slice") + ) else: metrics = {k: float(np.asarray(a[k]).item()) for k in _METRICS_KEYS} conf_ci_obj = slot_group.get("conf_ci") conf_ci = ( - _decode_dataframe(_as_dataset(conf_ci_obj)) if conf_ci_obj is not None else None + _decode_dataframe(require_dataset(conf_ci_obj, "conf_ci")) + if conf_ci_obj is not None + else None ) mcmc_obj = slot_group.get("mcmc") - mcmc = _read_mcmc_group(_as_group(mcmc_obj)) if mcmc_obj is not None else None + mcmc = ( + _read_mcmc_group(require_group(mcmc_obj, "mcmc")) + if mcmc_obj is not None + else None + ) yaml_filename = _attr_str(a["yaml_filename"]) if "yaml_filename" in a else None selection = json.loads(selection_json) @@ -1621,7 +1608,7 @@ def _read_slot( def _read_file(file_group: h5py.Group) -> SavedFile: """Decode one file group into a ``SavedFile``.""" - meta = _as_group(file_group["metadata"]) + meta = require_group(file_group["metadata"], "metadata") a = meta.attrs name = _attr_str(a["name"]) original_path = _attr_str(a["original_path"]) @@ -1636,16 +1623,16 @@ def _read_file(file_group: h5py.Group) -> SavedFile: e_lim = [int(x) for x in np.asarray(a["e_lim"]).ravel()] if "e_lim" in a else None t_lim = [int(x) for x in np.asarray(a["t_lim"]).ravel()] if "t_lim" in a else None - data = np.asarray(_as_dataset(file_group["data"])[...]) - energy = np.asarray(_as_dataset(file_group["energy"])[...]) - time = np.asarray(_as_dataset(file_group["time"])[...]) + data = np.asarray(require_dataset(file_group["data"], "data")[...]) + energy = np.asarray(require_dataset(file_group["energy"], "energy")[...]) + time = np.asarray(require_dataset(file_group["time"], "time")[...]) slot_records: list[SavedFitSlot] = [] slots_obj = file_group.get("slots") if slots_obj is not None: - slots_group = _as_group(slots_obj) + slots_group = require_group(slots_obj, "slots") for key in sorted(slots_group.keys()): - sg = _as_group(slots_group[key]) + sg = require_group(slots_group[key], f"slots/{key}") slot_records.append( _read_slot(sg, file_fingerprint=fingerprint, file_name=name) ) @@ -1680,7 +1667,7 @@ def read_archive(filepath: PathLike | str) -> SavedProject: path = Path(filepath) with h5py.File(path, "r") as archive: - meta = _as_group(archive["metadata"]) + meta = require_group(archive["metadata"], "metadata") ma = meta.attrs schema_version = _attr_str(ma["schema_version"]) if schema_version != SCHEMA_VERSION: @@ -1691,9 +1678,10 @@ def read_archive(filepath: PathLike | str) -> SavedProject: files_obj = archive.get("files") files: list[SavedFile] = [] if files_obj is not None: - files_group = _as_group(files_obj) + files_group = require_group(files_obj, "files") for key in sorted(files_group.keys()): - files.append(_read_file(_as_group(files_group[key]))) + fg = require_group(files_group[key], f"files/{key}") + files.append(_read_file(fg)) return SavedProject( name=_attr_str(ma["project_name"]), From 37cdf3354c54875ab2c0fbc738d241295007285e Mon Sep 17 00:00:00 2001 From: Johannes Mahl Date: Fri, 10 Jul 2026 13:46:00 -0700 Subject: [PATCH 26/36] move refline styling and panel size into PlotConfig; delegate plot_comparison 2D to plot_2d_grid --- src/trspecfit/config/plot.py | 11 +++++++++ src/trspecfit/simulator.py | 43 ++++++--------------------------- src/trspecfit/utils/plot.py | 47 ++++++++++++++++++++++++------------ 3 files changed, 50 insertions(+), 51 deletions(-) diff --git a/src/trspecfit/config/plot.py b/src/trspecfit/config/plot.py index c4a5204..3d72a41 100644 --- a/src/trspecfit/config/plot.py +++ b/src/trspecfit/config/plot.py @@ -78,8 +78,14 @@ class PlotConfig: X-coordinates for vertical lines hlines : list[float] | None Y-coordinates for horizontal lines + refline_color : str + Color for vlines/hlines reference lines + refline_style : str + Line style for vlines/hlines reference lines ticksize : float | None Font size for tick labels + panel_size : tuple[float, float] + Per-panel (width, height) in inches for multi-panel grid plots Examples -------- @@ -156,8 +162,13 @@ class PlotConfig: waterfall: float = 0 vlines: list[float] | None = None hlines: list[float] | None = None + refline_color: str = "#808080" + refline_style: str = ":" ticksize: float | None = None + # Per-panel size (width, height in inches) for multi-panel grid plots + panel_size: tuple[float, float] = (4.0, 3.0) + # Normalization y_norm: int = 0 # 0: no normalization, 1: normalize to [0,1] y_scale: list[float] | None = None diff --git a/src/trspecfit/simulator.py b/src/trspecfit/simulator.py index f8c108f..e83c5f5 100644 --- a/src/trspecfit/simulator.py +++ b/src/trspecfit/simulator.py @@ -60,7 +60,6 @@ from typing import cast import h5py -import matplotlib.pyplot as plt import numpy as np from trspecfit.config.plot import PlotConfig @@ -1296,42 +1295,14 @@ def plot_comparison( if plot_kwargs: resolved_config = resolved_config.copy(**plot_kwargs) - # Create 3-panel plot - _fig, axes = plt.subplots(1, 3, figsize=(15, 4)) - - panels = [ - (self.data_clean, "Clean Model Data"), - (self.data_noisy, plt_title), - (self.noise, "Noise (Simulated - Clean)"), - ] - for ax, (data, title) in zip(axes, panels, strict=True): - im = ax.pcolormesh( - self.model.energy, - self.model.time, - data, - shading="nearest", - cmap=resolved_config.z_colormap, - ) - ax.set_title(title) - ax.set_xlabel(resolved_config.x_label) - ax.set_ylabel(resolved_config.y_label) - if resolved_config.ticksize is not None: - ax.tick_params(labelsize=resolved_config.ticksize) - uplt._apply_axis_settings( - ax, - x_type=resolved_config.x_type, - x_dir=resolved_config.x_dir, - y_type=resolved_config.y_type, - y_dir=resolved_config.y_dir, - x_lim=resolved_config.x_lim, - y_lim=resolved_config.y_lim, - ) - plt.colorbar(im, ax=ax) - - plt.tight_layout() - uplt._finalize_plot( + uplt.plot_2d_grid( + [self.data_clean, self.data_noisy, self.noise], + x=self.model.energy, + y=self.model.time, + titles=["Clean Model Data", plt_title, "Noise (Simulated - Clean)"], + config=resolved_config, + columns=3, save_img=save_img, - save_path="", dpi_save=resolved_config.dpi_save, ) diff --git a/src/trspecfit/utils/plot.py b/src/trspecfit/utils/plot.py index 0c9926d..0916f46 100644 --- a/src/trspecfit/utils/plot.py +++ b/src/trspecfit/utils/plot.py @@ -181,6 +181,7 @@ def plot_2d_grid( y: ArrayLike | None = None, titles: Sequence[str] | None = None, config: "PlotConfig | None" = None, + columns: int | None = None, vlines: Sequence[Sequence[float]] | None = None, hlines: Sequence[Sequence[float]] | None = None, save_img: int = 0, @@ -201,7 +202,10 @@ def plot_2d_grid( titles : list of str, optional Title for each panel. If None, panels are untitled. config : PlotConfig, optional - Plot configuration (colormap, axis directions, labels). + Plot configuration (colormap, axis directions, labels, panel size, + reference-line styling). + columns : int, optional + Number of grid columns. If None, auto-selected from the panel count. vlines : list of list of float, optional Per-panel vertical reference lines. Each element is a list of x-coordinates for that panel (or empty list for none). @@ -224,8 +228,10 @@ def plot_2d_grid( if n == 0: return + if columns is not None: + cols = columns # Auto-select columns: 2 for <=4, 3 for <=9, 4 for <=16, 5 above - if n <= 4: + elif n <= 4: cols = 2 elif n <= 9: cols = 3 @@ -240,10 +246,11 @@ def plot_2d_grid( x_arr = None if x is None else np.asarray(x) y_arr = None if y is None else np.asarray(y) + panel_w, panel_h = config.panel_size fig, axs = plt.subplots( rows, cols, - figsize=(4.0 * cols, 3.0 * rows), + figsize=(panel_w * cols, panel_h * rows), squeeze=False, dpi=config.dpi_plot or 100, ) @@ -271,7 +278,11 @@ def plot_2d_grid( config.x_dir, config.y_type, config.y_dir, + config.x_lim, + config.y_lim, ) + if config.ticksize is not None: + ax.tick_params(labelsize=config.ticksize) # Reference lines if vlines is not None and idx < len(vlines) and vlines[idx]: @@ -279,16 +290,16 @@ def plot_2d_grid( x=np.asarray(vlines[idx]), ymin=np.min(yp), ymax=np.max(yp), - color="#000000", - linestyle=":", + color=config.refline_color, + linestyle=config.refline_style, ) if hlines is not None and idx < len(hlines) and hlines[idx]: ax.hlines( y=np.asarray(hlines[idx]), xmin=np.min(xp), xmax=np.max(xp), - color="#000000", - linestyle=":", + color=config.refline_color, + linestyle=config.refline_style, ) # Hide unused subplots @@ -342,6 +353,7 @@ def plot_2d( - z_colorbar : 'ver' or 'hor' for colorbar orientation - data_slice : [[x_start, x_stop], [y_start, y_stop]] for slicing by index - vlines, hlines : List of coordinates for reference lines + - refline_color, refline_style : Reference-line color and line style - ticksize : Font size for tick labels - dpi_plot, dpi_save : Display and save resolution - save_img : 0 (display), 1 (save+display), -1 (save only) @@ -403,6 +415,8 @@ def plot_2d( z_type = kwargs.get("z_type", config.z_type) vlines = kwargs.get("vlines", config.vlines) hlines = kwargs.get("hlines", config.hlines) + refline_color = kwargs.get("refline_color", config.refline_color) + refline_style = kwargs.get("refline_style", config.refline_style) ticksize = kwargs.get("ticksize", config.ticksize) save_img = kwargs.get("save_img", 0) save_path = kwargs.get("save_path", "") @@ -513,8 +527,8 @@ def plot_2d( y=np.asarray(hlines), xmin=np.min(x_plt), xmax=np.max(x_plt), - color="#000000", - linestyle=":", + color=refline_color, + linestyle=refline_style, ) if vlines is not None: @@ -522,8 +536,8 @@ def plot_2d( x=np.asarray(vlines), ymin=np.min(y_plt), ymax=np.max(y_plt), - color="#000000", - linestyle=":", + color=refline_color, + linestyle=refline_style, ) # Save/show/close @@ -572,6 +586,7 @@ def plot_1d( - y_norm : 0 (raw data) or 1 (normalize each trace to [0, 1]) - y_scale : List of scaling factors for each trace - vlines, hlines : List of coordinates for reference lines + - refline_color, refline_style : Reference-line color and line style - ticksize : Font size for tick labels - dpi_plot, dpi_save : Display and save resolution - save_img : 0 (display), 1 (save+display), -1 (save only) @@ -648,6 +663,8 @@ def plot_1d( legend = kwargs.get("legend", config.legend) vlines = kwargs.get("vlines", config.vlines) hlines = kwargs.get("hlines", config.hlines) + refline_color = kwargs.get("refline_color", config.refline_color) + refline_style = kwargs.get("refline_style", config.refline_style) y_scale = kwargs.get("y_scale", config.y_scale) # Determine number of plots @@ -745,8 +762,8 @@ def plot_1d( y=np.asarray(hlines), xmin=x_minmax[0], xmax=x_minmax[1], - color="#808080", - linestyle=":", + color=refline_color, + linestyle=refline_style, ) if vlines is not None: @@ -765,8 +782,8 @@ def plot_1d( x=np.asarray(vlines), ymin=y_minmax[0], ymax=y_minmax[1], - color="#808080", - linestyle="--", + color=refline_color, + linestyle=refline_style, ) # Axis settings From 931a6424f29669b34237e2b13f8c3e722ab4cb02 Mon Sep 17 00:00:00 2001 From: Johannes Mahl Date: Fri, 10 Jul 2026 13:46:32 -0700 Subject: [PATCH 27/36] add z_colormap_res: diverging, zero-centered colormap for 2D residual maps --- src/trspecfit/config/plot.py | 4 ++++ src/trspecfit/fitlib.py | 19 ++++++++++++++----- src/trspecfit/trspecfit.py | 2 ++ 3 files changed, 20 insertions(+), 5 deletions(-) diff --git a/src/trspecfit/config/plot.py b/src/trspecfit/config/plot.py index 3d72a41..ed35eee 100644 --- a/src/trspecfit/config/plot.py +++ b/src/trspecfit/config/plot.py @@ -56,6 +56,8 @@ class PlotConfig: DPI for saving plots z_colormap : str Colormap name for 2D plots + z_colormap_res : str + Diverging colormap name for 2D residual maps (centered on 0) data_slice : list[list[int]] | None Data slicing indices for 2D plots: [[x_start, x_stop], [y_start, y_stop]] colors : list[str] | None @@ -143,6 +145,8 @@ class PlotConfig: # 2D plot settings z_colormap: str = "viridis" + # Residual maps are signed and centered on 0 -> diverging colormap + z_colormap_res: str = "RdBu_r" z_colorbar: str = "ver" # 'ver' or 'hor' z_type: str = "lin" # 'lin' or 'log' for color scale # 2D data handling diff --git a/src/trspecfit/fitlib.py b/src/trspecfit/fitlib.py index f8e4cb8..6ae2d3d 100644 --- a/src/trspecfit/fitlib.py +++ b/src/trspecfit/fitlib.py @@ -1394,8 +1394,11 @@ def plt_fit_res_2d( - z_lim_top : Color scale ``[min, max]`` for data and fit panels. Synchronized scale enables direct comparison - z_lim_res : Color scale ``[min, max]`` for residual panel. - Independent scale optimizes residual visibility - - z_colormap : Colormap name (default 'viridis') + If None, symmetric around 0 so the diverging colormap's + midpoint marks zero residual + - z_colormap : Colormap name for data/fit panels (default 'viridis') + - z_colormap_res : Diverging colormap name for the residual panel + (default 'RdBu_r') - x_dir, y_dir : 'def' or 'rev' for axis direction - x_type, y_type : 'lin' or 'log' for axis scale - save_img : 0 (display), 1 (save+display), -1 (save only) @@ -1409,6 +1412,7 @@ def plt_fit_res_2d( x_label = kwargs.get("x_label", config.x_label) y_label = kwargs.get("y_label", config.y_label) z_colormap = kwargs.get("z_colormap", config.z_colormap) + z_colormap_res = kwargs.get("z_colormap_res", config.z_colormap_res) x_dir = kwargs.get("x_dir", config.x_dir) x_type = kwargs.get("x_type", config.x_type) y_dir = kwargs.get("y_dir", config.y_dir) @@ -1457,8 +1461,13 @@ def plt_fit_res_2d( else: range_dat_fit = z_lim_top - # Residual has independent scale - range_res = [np.min(res_cut), np.max(res_cut)] if z_lim_res is None else z_lim_res + # Residual has independent scale, symmetric around 0 by default so the + # diverging colormap's midpoint marks zero residual + if z_lim_res is None: + res_amp = np.max(np.abs(res_cut)) + range_res = [-res_amp, res_amp] + else: + range_res = z_lim_res # Create figure layout fig, axs = plt.subplot_mosaic( @@ -1508,7 +1517,7 @@ def plt_fit_res_2d( x_arr, y_arr, res, - cmap=z_colormap, + cmap=z_colormap_res, vmin=range_res[0], vmax=range_res[1], shading="nearest", diff --git a/src/trspecfit/trspecfit.py b/src/trspecfit/trspecfit.py index 0ee970f..9ea3ea5 100644 --- a/src/trspecfit/trspecfit.py +++ b/src/trspecfit/trspecfit.py @@ -238,6 +238,7 @@ def _set_defaults(self) -> None: self.y_dir = "def" self.y_type = "lin" self.z_colormap = "viridis" + self.z_colormap_res = "RdBu_r" self.z_colorbar = "ver" self.z_type = "lin" self.dpi_plt = 100 @@ -755,6 +756,7 @@ def describe(self, detail: int = 0) -> None: print(f" y_dir: {self.y_dir}") print(f" y_type: {self.y_type}") print(f" z_colormap: {self.z_colormap}") + print(f" z_colormap_res: {self.z_colormap_res}") print(f" z_colorbar: {self.z_colorbar}") print(f" z_type: {self.z_type}") print(f" dpi_plt: {self.dpi_plt}") From b0c91e0db04bd14f9a441d2be19090e300ce2150 Mon Sep 17 00:00:00 2001 From: Johannes Mahl Date: Fri, 10 Jul 2026 13:52:57 -0700 Subject: [PATCH 28/36] make every PlotConfig field settable via project.yaml; add coverage guard --- pyproject.toml | 2 +- src/trspecfit/config/plot.py | 5 +++-- src/trspecfit/trspecfit.py | 3 +++ tests/test_plotting.py | 29 +++++++++++++++++++++++++++++ 4 files changed, 36 insertions(+), 3 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index f91d28f..596a68c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "trspecfit" -version = "0.10.1" +version = "0.10.2" authors = [ {name = "Johannes Mahl", email = "infinitymonkeyatwork@gmail.com"}, ] diff --git a/src/trspecfit/config/plot.py b/src/trspecfit/config/plot.py index ed35eee..dd0750b 100644 --- a/src/trspecfit/config/plot.py +++ b/src/trspecfit/config/plot.py @@ -200,7 +200,8 @@ def from_project(cls, project, **overrides) -> "PlotConfig": "y_label": "t_label", "dpi_plot": "dpi_plt", } - limit_fields = {"x_lim", "y_lim", "z_lim"} + # Tuple-typed fields arrive as lists when set via project.yaml + tuple_fields = {"x_lim", "y_lim", "z_lim", "panel_size"} config_dict = {} for field in fields(cls): @@ -209,7 +210,7 @@ def from_project(cls, project, **overrides) -> "PlotConfig": continue value = cp.deepcopy(getattr(project, source_attr)) - if field.name in limit_fields and value is not None: + if field.name in tuple_fields and value is not None: value = tuple(value) config_dict[field.name] = value diff --git a/src/trspecfit/trspecfit.py b/src/trspecfit/trspecfit.py index 9ea3ea5..cce6d24 100644 --- a/src/trspecfit/trspecfit.py +++ b/src/trspecfit/trspecfit.py @@ -259,7 +259,10 @@ def _set_defaults(self) -> None: self.waterfall = 0 self.vlines = None self.hlines = None + self.refline_color = "#808080" + self.refline_style = ":" self.ticksize = None + self.panel_size = (4.0, 3.0) self.y_norm = 0 self.y_scale = None # File I/O settings diff --git a/tests/test_plotting.py b/tests/test_plotting.py index 2491fa5..936b9ca 100644 --- a/tests/test_plotting.py +++ b/tests/test_plotting.py @@ -53,6 +53,27 @@ def test_default_creation(self): assert config.dpi_plot == 100 assert config.z_colormap == "viridis" + # + def test_every_field_settable_via_project(self): + """Every PlotConfig field must have a Project counterpart. + + Project.yaml keys only apply to existing Project attributes, and + PlotConfig.from_project only copies fields the Project has — a + PlotConfig field without a Project default is silently unsettable + from project.yaml. + """ + + from dataclasses import fields + + project = make_project(name="cfg-coverage") + aliases = {"x_label": "e_label", "y_label": "t_label", "dpi_plot": "dpi_plt"} + missing = [ + f.name + for f in fields(PlotConfig) + if not hasattr(project, aliases.get(f.name, f.name)) + ] + assert missing == [], f"PlotConfig fields without Project defaults: {missing}" + # def test_custom_creation(self): """Test creating config with custom values""" @@ -559,6 +580,10 @@ class TestPlotConfigFromYAML: "z_label: 'Counts'\n" "x_dir: 'rev'\n" "z_colormap: 'RdBu'\n" + "z_colormap_res: 'coolwarm'\n" + "refline_color: '#ff0000'\n" + "refline_style: '--'\n" + "panel_size: [5.0, 3.5]\n" "x_lim: [63.0, -2.6]\n" "y_lim: [-0.5, 5.0]\n" ) @@ -569,6 +594,10 @@ class TestPlotConfigFromYAML: "z_label": "Counts", "x_dir": "rev", "z_colormap": "RdBu", + "z_colormap_res": "coolwarm", + "refline_color": "#ff0000", + "refline_style": "--", + "panel_size": (5.0, 3.5), "x_lim": (63.0, -2.6), "y_lim": (-0.5, 5.0), } From 9014377e69f216356cadabcd1f60ef0343b8daf4 Mon Sep 17 00:00:00 2001 From: Johannes Mahl Date: Fri, 10 Jul 2026 14:09:18 -0700 Subject: [PATCH 29/36] open sweep HDF5 file once instead of per config --- src/trspecfit/simulator.py | 197 +++++++++++++++++++------------------ 1 file changed, 102 insertions(+), 95 deletions(-) diff --git a/src/trspecfit/simulator.py b/src/trspecfit/simulator.py index e83c5f5..5124908 100644 --- a/src/trspecfit/simulator.py +++ b/src/trspecfit/simulator.py @@ -1727,43 +1727,49 @@ def simulate_parameter_sweep( f" Output file: {filepath}\n" ) - # Initialize HDF5 file with structure - self._initialize_sweep_hdf5( - filepath, parameter_sweep, n_realizations, n_configs, dim - ) - - # Process each configuration - for config_idx, param_config in enumerate(parameter_sweep): - if show_progress: - # Format parameters nicely - param_str = ", ".join(f"{k}={v:.3g}" for k, v in param_config.items()) - print( - f"Processing config {config_idx + 1}/{n_configs}: {{{param_str}}}" - ) - - # Update model parameters - param_names = list(param_config.keys()) - param_values = list(param_config.values()) - self.model.update_value(param_values, par_select=param_names) - - # Generate noisy realizations for this config - clean, noisy_list, _noise_list = self.simulate_n( - n=n_realizations, - dim=dim, - show_progress=False, # Don't clutter output + # Keep the file open for the whole sweep (per-config open/close is + # costly for large sweeps); each config is flushed after appending + with h5py.File(filepath, "w") as f: + # Initialize HDF5 file with structure + self._initialize_sweep_hdf5( + f, parameter_sweep, n_realizations, n_configs, dim ) - # Append to HDF5 immediately (memory-efficient) - self._append_config_to_hdf5( - filepath, config_idx, param_config, clean, noisy_list - ) + # Process each configuration + for config_idx, param_config in enumerate(parameter_sweep): + if show_progress: + # Format parameters nicely + param_str = ", ".join( + f"{k}={v:.3g}" for k, v in param_config.items() + ) + print( + f"Processing config {config_idx + 1}/{n_configs}:" + f" {{{param_str}}}" + ) + + # Update model parameters + param_names = list(param_config.keys()) + param_values = list(param_config.values()) + self.model.update_value(param_values, par_select=param_names) + + # Generate noisy realizations for this config + clean, noisy_list, _noise_list = self.simulate_n( + n=n_realizations, + dim=dim, + show_progress=False, # Don't clutter output + ) - if show_progress: - print( - f" ✓ Saved config {config_idx + 1}" - f" with {n_realizations} realizations" + # Append to HDF5 immediately (memory-efficient) + self._append_config_to_hdf5( + f, config_idx, param_config, clean, noisy_list ) + if show_progress: + print( + f" ✓ Saved config {config_idx + 1}" + f" with {n_realizations} realizations" + ) + if show_progress: print( f"\n{'=' * 60}\n" @@ -1777,7 +1783,7 @@ def simulate_parameter_sweep( # def _initialize_sweep_hdf5( self, - filepath: str, + f: h5py.File, parameter_sweep: ParameterSweep, n_realizations: int, n_configs: int, @@ -1811,8 +1817,8 @@ def _initialize_sweep_hdf5( Parameters ---------- - filepath : str - Path to HDF5 file to create + f : h5py.File + Freshly created HDF5 file, open for writing parameter_sweep : ParameterSweep Parameter sweep object (for metadata) n_realizations : int @@ -1823,51 +1829,50 @@ def _initialize_sweep_hdf5( Dimensionality of simulated data """ - with h5py.File(filepath, "w") as f: - # Save axes (same for all configs) - self._write_axes_hdf5(f) + # Save axes (same for all configs) + self._write_axes_hdf5(f) - # Create groups for organization - f.create_group("parameter_configs") - f.create_group("simulated_data") + # Create groups for organization + f.create_group("parameter_configs") + f.create_group("simulated_data") - # Save sweep metadata - meta = f.create_group("metadata") - meta.attrs["n_configs"] = n_configs - meta.attrs["n_realizations_per_config"] = n_realizations - meta.attrs["total_datasets"] = n_configs * n_realizations + # Save sweep metadata + meta = f.create_group("metadata") + meta.attrs["n_configs"] = n_configs + meta.attrs["n_realizations_per_config"] = n_realizations + meta.attrs["total_datasets"] = n_configs * n_realizations - # Simulator settings - self._write_detection_metadata(meta) + # Simulator settings + self._write_detection_metadata(meta) - # Parameter sweep settings - meta.attrs["sweep_strategy"] = parameter_sweep.strategy - meta.attrs["sweep_seed"] = ( - "None" if parameter_sweep.seed is None else parameter_sweep.seed - ) + # Parameter sweep settings + meta.attrs["sweep_strategy"] = parameter_sweep.strategy + meta.attrs["sweep_seed"] = ( + "None" if parameter_sweep.seed is None else parameter_sweep.seed + ) - # Save parameter space definition as JSON - param_space = {} - for par_name, spec in parameter_sweep.parameter_specs.items(): - # Convert numpy arrays to lists for JSON serialization - spec_copy = spec.copy() - if "values" in spec_copy: - spec_copy["values"] = spec_copy["values"].tolist() - param_space[par_name] = spec_copy + # Save parameter space definition as JSON + param_space = {} + for par_name, spec in parameter_sweep.parameter_specs.items(): + # Convert numpy arrays to lists for JSON serialization + spec_copy = spec.copy() + if "values" in spec_copy: + spec_copy["values"] = spec_copy["values"].tolist() + param_space[par_name] = spec_copy - meta.attrs["parameter_space"] = json.dumps(param_space, indent=2) + meta.attrs["parameter_space"] = json.dumps(param_space, indent=2) - # Save full model definition once (static across all configs) - meta.attrs["model_parameters"] = self._model_parameters_json() - meta.attrs["model_name"] = self.model.name + # Save full model definition once (static across all configs) + meta.attrs["model_parameters"] = self._model_parameters_json() + meta.attrs["model_name"] = self.model.name - # Dimension matches the actual data written, not model capability - meta.attrs["dimension"] = dim + # Dimension matches the actual data written, not model capability + meta.attrs["dimension"] = dim # def _append_config_to_hdf5( self, - filepath: str, + f: h5py.File, config_idx: int, param_config: dict[str, float], clean: np.ndarray, @@ -1878,8 +1883,8 @@ def _append_config_to_hdf5( Parameters ---------- - filepath : str - Path to HDF5 file + f : h5py.File + Sweep HDF5 file, open for writing config_idx : int Configuration index (for naming) param_config : dict @@ -1890,28 +1895,30 @@ def _append_config_to_hdf5( List of noisy realizations """ - with h5py.File(filepath, "a") as f: - # Create group for this configuration - config_name = f"config_{config_idx:06d}" - configs_group = require_group(f["parameter_configs"], "parameter_configs") - config_group = configs_group.create_group(config_name) - - # Save swept parameters as attributes - for par_name, value in param_config.items(): - config_group.attrs[par_name] = float(value) - - # Save all parameter values for this config (full model state) - param_values = { - par_name: float(self.model.lmfit_pars[par_name].value) - for par_name in self.model.lmfit_pars - } - config_group.attrs["all_parameter_values"] = json.dumps(param_values) - - # Save clean data for this configuration - config_group.create_dataset("clean", data=clean) - - # Save noisy realizations - simulated_group = require_group(f["simulated_data"], "simulated_data") - data_group = simulated_group.create_group(config_name) - for real_idx, noisy_data in enumerate(noisy_list): - data_group.create_dataset(f"{real_idx:06d}", data=noisy_data) + # Create group for this configuration + config_name = f"config_{config_idx:06d}" + configs_group = require_group(f["parameter_configs"], "parameter_configs") + config_group = configs_group.create_group(config_name) + + # Save swept parameters as attributes + for par_name, value in param_config.items(): + config_group.attrs[par_name] = float(value) + + # Save all parameter values for this config (full model state) + param_values = { + par_name: float(self.model.lmfit_pars[par_name].value) + for par_name in self.model.lmfit_pars + } + config_group.attrs["all_parameter_values"] = json.dumps(param_values) + + # Save clean data for this configuration + config_group.create_dataset("clean", data=clean) + + # Save noisy realizations + simulated_group = require_group(f["simulated_data"], "simulated_data") + data_group = simulated_group.create_group(config_name) + for real_idx, noisy_data in enumerate(noisy_list): + data_group.create_dataset(f"{real_idx:06d}", data=noisy_data) + + # Completed configs stay on disk if the sweep is interrupted + f.flush() From 3e39f4056e460fcb073ce3d708653c55e84148f6 Mon Sep 17 00:00:00 2001 From: Johannes Mahl Date: Fri, 10 Jul 2026 14:15:17 -0700 Subject: [PATCH 30/36] nudge toward parameter bounds in LinBack ordering error --- src/trspecfit/functions/energy.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/trspecfit/functions/energy.py b/src/trspecfit/functions/energy.py index e8646c1..2e03b9f 100644 --- a/src/trspecfit/functions/energy.py +++ b/src/trspecfit/functions/energy.py @@ -148,7 +148,8 @@ def LinBack( bad = np.argmax((xs >= xe).ravel()) raise ValueError( f"LinBack requires xStart < xStop, got xStart={xs.ravel()[bad]}, " - f"xStop={xe.ravel()[bad]}" + f"xStop={xe.ravel()[bad]}. If xStart/xStop are fit parameters, " + f"set min/max bounds that keep them ordered." ) y = m * (x - xStart) + b y_stop = m * (xStop - xStart) + b From d081d790537399710475f0131db8c780122e8e47 Mon Sep 17 00:00:00 2001 From: Johannes Mahl Date: Fri, 10 Jul 2026 14:38:11 -0700 Subject: [PATCH 31/36] label type-guard asserts in tests; document figure-inspection plot exception --- CLAUDE.md | 2 +- tests/test_evaluate_1d.py | 6 ++-- tests/test_evaluate_2d.py | 8 ++--- tests/test_file.py | 2 +- tests/test_gir_integration.py | 61 +++++++++++++++++++---------------- tests/test_graph_ir.py | 56 ++++++++++++++++---------------- tests/test_mcp_library.py | 2 +- 7 files changed, 71 insertions(+), 66 deletions(-) diff --git a/CLAUDE.md b/CLAUDE.md index e0bf3ac..63c0aa0 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -42,7 +42,7 @@ - **Pattern:** Use plain pytest. Avoid `unittest.TestCase` and fixtures; prefer explicit helper builders named by intent. - **API Usage:** Use the public API (`Project`, `File.load_model`, etc.) in tests to avoid masking bugs by skipping validation or setup. Use internals only for pure-math unit tests or explicit invariant checks. - **Execution:** Run `pytest -q`. Keep YAML test assets in `tests/models/`. -- **Plots:** Always suppress plot display in tests: pass `show_plot=False` where available, or `save_img=-2`. +- **Plots:** Always suppress plot display in tests: pass `show_plot=False` where available, or `save_img=-2`. Exception: figure-inspection tests that assert on the live axes cannot pass `save_img=-2` (it closes the figure); they rely on the module-level Agg backend and must call `plt.close("all")` after the assertions. - **Type Guards:** When `assert x is not None` narrows an `X | None` type, add a `# type guard` comment. - **Variable Naming:** For variables derived from registry parameters or components, keep original casing (e.g., `SD = 2.0`, `c_Shirley = Component("Shirley")`). Name derived variables as `{par}_{qualifier}` (e.g., `A_early`, `mean_A`). diff --git a/tests/test_evaluate_1d.py b/tests/test_evaluate_1d.py index 49b1073..6621ac8 100644 --- a/tests/test_evaluate_1d.py +++ b/tests/test_evaluate_1d.py @@ -32,7 +32,7 @@ def _make_energy_model(model_info): file = File(parent_project=project, energy=np.linspace(80, 90, 101)) file.load_model(model_yaml=_ENERGY_YAML, model_info=model_info) model = file.model_active - assert model is not None + assert model is not None # type guard return file, model @@ -65,7 +65,7 @@ def _make_profile_energy_model( profile_model=profile_model, ) model = file.model_active - assert model is not None + assert model is not None # type guard return file, model @@ -378,7 +378,7 @@ def test_2d_model_is_not_lowerable_1d(self): dynamics_model=["MonoExpPos"], ) model = file.model_active - assert model is not None + assert model is not None # type guard graph = build_graph(model) assert not can_lower_1d(graph) diff --git a/tests/test_evaluate_2d.py b/tests/test_evaluate_2d.py index 343f8d3..7e2985e 100644 --- a/tests/test_evaluate_2d.py +++ b/tests/test_evaluate_2d.py @@ -38,7 +38,7 @@ def _make_energy_model(model_info): file.energy = np.linspace(80, 90, 101) file.load_model(model_yaml=_ENERGY_YAML, model_info=model_info) model = file.model_active - assert model is not None + assert model is not None # type guard return file, model @@ -58,7 +58,7 @@ def _make_2d_model(model_info, dynamics_params, *, frequency=None, time=None): file.time = np.linspace(-10, 100, 51) if time is None else time file.load_model(model_yaml=_ENERGY_YAML, model_info=model_info) model = file.model_active - assert model is not None + assert model is not None # type guard for target_par, dyn_model in dynamics_params: kwargs = { @@ -518,7 +518,7 @@ def _make_2d_profile_model( file = File(parent_project=project, energy=energy, time=time, aux_axis=aux_axis) file.load_model(model_yaml=model_yaml, model_info=model_info) model = file.model_active - assert model is not None + assert model is not None # type guard for target_par, dyn_model in dynamics_params: file.add_time_dependence( @@ -1051,7 +1051,7 @@ def test_mixed_irf_and_subcycle_parity(self): file.time = np.linspace(-10, 100, 51) file.load_model(model_yaml=_ENERGY_YAML, model_info=["offset_only"]) model = file.model_active - assert model is not None + assert model is not None # type guard file.add_time_dependence( target_model="offset_only", diff --git a/tests/test_file.py b/tests/test_file.py index 25c6e22..21e5d3f 100644 --- a/tests/test_file.py +++ b/tests/test_file.py @@ -407,7 +407,7 @@ def test_set_fit_limits_none_uses_full_range(self): file = self._make_file_with_data() file.set_fit_limits(None, show_plot=False) - assert file.e_lim_abs is not None + assert file.e_lim_abs is not None # type guard assert np.isclose(file.e_lim_abs[0], 80.0) assert np.isclose(file.e_lim_abs[1], 90.0) diff --git a/tests/test_gir_integration.py b/tests/test_gir_integration.py index b10f5f7..002947e 100644 --- a/tests/test_gir_integration.py +++ b/tests/test_gir_integration.py @@ -65,7 +65,7 @@ def _make_2d_model(project, model_info, dynamics_params, *, frequency=None): file.time = np.linspace(-10, 100, 51) file.load_model(model_yaml=_ENERGY_YAML, model_info=model_info) model = file.model_active - assert model is not None + assert model is not None # type guard for target_par, dyn_model in dynamics_params: kwargs = { @@ -99,7 +99,7 @@ def _make_1d_profile_model(project, model_info, profiles): profile_model=profile_model, ) model = file.model_active - assert model is not None + assert model is not None # type guard return file, model @@ -132,7 +132,7 @@ def _make_2d_profile_model(project, model_info, dynamics_params, profiles): ) model = file.model_active - assert model is not None + assert model is not None # type guard return file, model @@ -228,7 +228,7 @@ def test_1d_fit_uses_gir_when_lowerable(self): file = File(parent_project=project, energy=np.linspace(80, 90, 101)) file.load_model(model_yaml=_ENERGY_YAML, model_info=["glp_only"]) model = file.model_active - assert model is not None + assert model is not None # type guard graph = build_graph(model) assert can_lower_1d(graph) @@ -309,7 +309,7 @@ def test_1d_plot_sum_false_falls_back(self): file = File(parent_project=project, energy=np.linspace(80, 90, 101)) file.load_model(model_yaml=_ENERGY_YAML, model_info=["offset_only"]) model = file.model_active - assert model is not None + assert model is not None # type guard graph = build_graph(model) plan = schedule_1d(graph) @@ -333,7 +333,7 @@ def test_1d_mcp_fallback_when_non_lowerable(self): file = File(parent_project=project, energy=np.linspace(80, 90, 101)) file.load_model(model_yaml=_ENERGY_YAML, model_info=["glp_only"]) model = file.model_active - assert model is not None + assert model is not None # type guard par = _extract_par_list(model) # Pass (model, dim=1) — no plan → delegates to MCP @@ -360,7 +360,7 @@ def test_residual_same_gir_vs_mcp(self): # Generate synthetic data from the model model.create_value_2d() - assert model.value_2d is not None + assert model.value_2d is not None # type guard data = model.value_2d + 0.01 # small offset so residual is non-zero # Compile GIR path @@ -432,7 +432,7 @@ def test_residual_with_slicing(self): ) model.create_value_2d() - assert model.value_2d is not None + assert model.value_2d is not None # type guard data = model.value_2d + 0.01 graph = build_graph(model) @@ -483,7 +483,7 @@ def test_residual_same_gir_vs_mcp_irf(self, dyn_model): ) model.create_value_2d() - assert model.value_2d is not None + assert model.value_2d is not None # type guard data = model.value_2d + 0.01 graph = build_graph(model) @@ -554,7 +554,7 @@ def test_subcycle_residual_gir_vs_mcp(self): ) model.create_value_2d() - assert model.value_2d is not None + assert model.value_2d is not None # type guard data = model.value_2d + 0.01 graph = build_graph(model) @@ -666,7 +666,7 @@ def test_residual_same_gir_vs_mcp_dynamics(self, dyn_model): ) model.create_value_2d() - assert model.value_2d is not None + assert model.value_2d is not None # type guard data = model.value_2d + 0.01 graph = build_graph(model) @@ -708,7 +708,7 @@ def test_residual_same_gir_vs_mcp_multi_substep_single_cycle(self): ) model.create_value_2d() - assert model.value_2d is not None + assert model.value_2d is not None # type guard data = model.value_2d + 0.01 graph = build_graph(model) @@ -751,7 +751,7 @@ def test_residual_same_gir_vs_mcp_chained_conv(self): ) model.create_value_2d() - assert model.value_2d is not None + assert model.value_2d is not None # type guard data = model.value_2d + 0.01 graph = build_graph(model) @@ -852,7 +852,7 @@ def test_gir_fit_writes_back_to_model(self): fit_file.fit_2d(model_name="single_glp", stages=2, try_ci=0) # Verify writeback: model_2d.lmfit_pars should match result params - assert fit_file.model_2d is not None + assert fit_file.model_2d is not None # type guard result_params = fit_file.model_2d.result[1].params for name in fit_file.model_2d.parameter_names: model_val = fit_file.model_2d.lmfit_pars[name].value @@ -963,11 +963,11 @@ def test_residual_same_gir_vs_mcp_1d(self): file = File(parent_project=project, energy=np.linspace(80, 90, 101)) file.load_model(model_yaml=_ENERGY_YAML, model_info=["glp_only"]) model = file.model_active - assert model is not None + assert model is not None # type guard # Generate synthetic data model.create_value_1d() - assert model.value_1d is not None + assert model.value_1d is not None # type guard data = model.value_1d + 0.01 # Compile GIR 1D path @@ -1009,7 +1009,7 @@ def test_residual_with_e_lim_1d(self): file = File(parent_project=project, energy=np.linspace(80, 90, 101)) file.load_model(model_yaml=_ENERGY_YAML, model_info=["offset_only"]) model = file.model_active - assert model is not None + assert model is not None # type guard model.create_value_1d() data = model.value_1d + 0.01 @@ -1056,7 +1056,7 @@ def test_residual_same_gir_vs_mcp_profile_1d(self): ) model.create_value_1d() - assert model.value_1d is not None + assert model.value_1d is not None # type guard data = model.value_1d + 0.01 graph = build_graph(model) @@ -1093,7 +1093,7 @@ def test_compare_mode_1d(self): file = File(parent_project=project, energy=np.linspace(80, 90, 101)) file.load_model(model_yaml=_ENERGY_YAML, model_info=["glp_expression"]) model = file.model_active - assert model is not None + assert model is not None # type guard graph = build_graph(model) assert can_lower_1d(graph) @@ -1148,7 +1148,7 @@ def test_residual_same_gir_vs_mcp_profile_2d(self): ) model.create_value_2d() - assert model.value_2d is not None + assert model.value_2d is not None # type guard data = model.value_2d + 0.01 graph = build_graph(model) @@ -1228,7 +1228,7 @@ def test_time_1d_dynamics_model_mcp_fallback(self): assert not can_lower_1d(graph) dyn.create_value_1d() - assert dyn.value_1d is not None + assert dyn.value_1d is not None # type guard data = np.asarray(dyn.value_1d) + 0.01 res_mcp = fitlib.residual_fun( @@ -1288,7 +1288,12 @@ class TestFileFitBaseline: # def test_1d_dispatch_args_lower_on_2d_file(self): - """1D workflow args compile on a 2D File with a plain energy model.""" + """1D workflow args compile on a 2D File with a plain energy model. + + Invariant check on the private dispatch-args contract that the + fit_baseline/fit_spectrum/fit_slice_by_slice call sites rely on; + end-to-end public-path coverage is test_gir_baseline_writes_back. + """ project = _make_project() energy = np.linspace(83, 87, 50) @@ -1324,7 +1329,7 @@ def test_gir_baseline_writes_back(self): # Tile 1D truth spectrum into 2D data (constant across time) truth.model_active.create_value_1d() - assert truth.model_active.value_1d is not None + assert truth.model_active.value_1d is not None # type guard spectrum_1d = truth.model_active.value_1d.copy() data_2d = np.tile(spectrum_1d, (len(truth.time), 1)) @@ -1332,7 +1337,7 @@ def test_gir_baseline_writes_back(self): fit_file.fit_baseline(model_name="single_glp", stages=2, try_ci=0) # Verify writeback - assert fit_file.model_base is not None + assert fit_file.model_base is not None # type guard result_params = fit_file.model_base.result[1].params for name in fit_file.model_base.parameter_names: model_val = fit_file.model_base.lmfit_pars[name].value @@ -1439,8 +1444,8 @@ def test_compare_mode_through_fit_slice_by_slice(self, n_workers): model_name="single_glp", stages=1, try_ci=0, n_workers=n_workers ) - assert fit_file.model_sbs is not None - assert fit_file.model_sbs.args is not None + assert fit_file.model_sbs is not None # type guard + assert fit_file.model_sbs.args is not None # type guard assert len(fit_file.model_sbs.args) == 4 assert isinstance(fit_file.model_sbs.args[0], ScheduledPlan1D) assert len(fit_file.results_sbs) == len(truth.time) @@ -1527,7 +1532,7 @@ def test_fit_slice_by_slice_restores_seed_template( fit_file = _make_1d_fit_file(project, data_2d, truth.energy, truth.time) fit_file.fit_baseline(model_name="single_glp", stages=1, try_ci=0) - assert fit_file.model_base is not None + assert fit_file.model_base is not None # type guard expected = ulmfit.par_extract(fit_file.model_base.result[1], return_type="list") seed_values = None @@ -1550,7 +1555,7 @@ def test_fit_slice_by_slice_restores_seed_template( seed_adapt=seed_adapt, ) - assert fit_file.model_sbs is not None + assert fit_file.model_sbs is not None # type guard for name, expected_value in zip( fit_file.model_sbs.parameter_names, expected, strict=True ): diff --git a/tests/test_graph_ir.py b/tests/test_graph_ir.py index 163c886..66508cd 100644 --- a/tests/test_graph_ir.py +++ b/tests/test_graph_ir.py @@ -33,7 +33,7 @@ def _make_energy_model(model_info): model_info=model_info, ) model = file.model_active - assert model is not None + assert model is not None # type guard return file, model @@ -58,7 +58,7 @@ def _make_2d_model(model_info, dynamics_params): model_info=model_info, ) model = file.model_active - assert model is not None + assert model is not None # type guard for target_par, _dyn_yaml, dyn_model in dynamics_params: file.add_time_dependence( @@ -159,20 +159,20 @@ def test_simple_energy_param_nodes(self): # Offset_y0: vary=True -> OPT_PARAM offset_y0 = _node_by_name(graph, "Offset_y0") - assert offset_y0 is not None + assert offset_y0 is not None # type guard assert offset_y0.kind == NodeKind.OPT_PARAM assert offset_y0.value == 2.0 assert offset_y0.vary is True # Shirley_pShirley: vary=False -> STATIC_PARAM shirley_p = _node_by_name(graph, "Shirley_pShirley") - assert shirley_p is not None + assert shirley_p is not None # type guard assert shirley_p.kind == NodeKind.STATIC_PARAM assert np.isclose(shirley_p.value, 4e-4) # GLP_01_A: vary=True with bounds glp_A = _node_by_name(graph, "GLP_01_A") - assert glp_A is not None + assert glp_A is not None # type guard assert glp_A.kind == NodeKind.OPT_PARAM assert glp_A.value == 20.0 assert glp_A.bounds == (5.0, 25.0) @@ -199,7 +199,7 @@ def test_offset_is_component_eval(self): graph = build_graph(model) offset_node = _node_by_name(graph, "Offset") - assert offset_node is not None + assert offset_node is not None # type guard assert offset_node.kind == NodeKind.COMPONENT_EVAL # No SPECTRUM_INPUT edge spec_edges = _edges_to(graph, offset_node.id, EdgeKind.SPECTRUM_INPUT) @@ -213,7 +213,7 @@ def test_shirley_has_spectrum_input(self): graph = build_graph(model) shirley_node = _node_by_name(graph, "Shirley") - assert shirley_node is not None + assert shirley_node is not None # type guard spec_edges = _edges_to(graph, shirley_node.id, EdgeKind.SPECTRUM_INPUT) assert len(spec_edges) == 1 # Source should be peak_sum SUM node @@ -335,7 +335,7 @@ def test_expression_nodes_created(self): assert len(expr_nodes) == 4 A_expr = _node_by_name(graph, "GLP_02_A") - assert A_expr is not None + assert A_expr is not None # type guard assert A_expr.kind == NodeKind.EXPRESSION assert A_expr.expr_string == "3/4*GLP_01_A" @@ -395,7 +395,7 @@ def test_forward_reference_expression(self): # GLP_01_A references GLP_02_A (forward ref) glp01_A = _node_by_name(graph, "GLP_01_A") - assert glp01_A is not None + assert glp01_A is not None # type guard assert glp01_A.kind == NodeKind.EXPRESSION @@ -481,7 +481,7 @@ def test_dynamics_edges(self): # PARAM_PLUS_TRACE has BASE_INPUT + TRACE_INPUT resolved = _node_by_name(graph, "GLP_01_A_resolved") - assert resolved is not None + assert resolved is not None # type guard base_edges = _edges_to(graph, resolved.id, EdgeKind.BASE_INPUT) trace_edges = _edges_to(graph, resolved.id, EdgeKind.TRACE_INPUT) assert len(base_edges) == 1 @@ -620,11 +620,11 @@ def test_energy_components_have_energy_package(self): graph = build_graph(model) glp01 = _node_by_name(graph, "GLP_01") - assert glp01 is not None + assert glp01 is not None # type guard assert glp01.package == "energy" offset = _node_by_name(graph, "Offset") - assert offset is not None + assert offset is not None # type guard assert offset.package == "energy" # @@ -661,7 +661,7 @@ def _make_profile_model(energy_model_info, target_par, profile_model_info): profile_model=profile_model_info, ) model = file.model_active - assert model is not None + assert model is not None # type guard return file, model @@ -712,7 +712,7 @@ def test_profile_average_has_addend_edges(self): graph = build_graph(model) comp_avg = _node_by_name(graph, "Gauss_01_profile_avg") - assert comp_avg is not None + assert comp_avg is not None # type guard addend_edges = _edges_to(graph, comp_avg.id, EdgeKind.ADDEND) assert len(addend_edges) == 5 # one per aux_axis point for edge in addend_edges: @@ -736,7 +736,7 @@ def test_profile_average_replaces_component_in_combination(self): # PROFILE_AVERAGE should exist and receive ADDEND from sample evals avg_node = _node_by_name(graph, "Gauss_01_profile_avg") - assert avg_node is not None + assert avg_node is not None # type guard assert avg_node.kind == NodeKind.PROFILE_AVERAGE # Each sample COMPONENT_EVAL feeds into PROFILE_AVERAGE @@ -790,7 +790,7 @@ def test_sample_component_eval_wiring(self): # Check the first sample's component eval sample_eval = _node_by_name(graph, "Gauss_01_sample_0") - assert sample_eval is not None + assert sample_eval is not None # type guard assert sample_eval.function_name == "Gauss" param_edges = _edges_to(graph, sample_eval.id, EdgeKind.PARAM_INPUT) @@ -909,7 +909,7 @@ def _make_subcycle_model(): frequency=10, ) model = file.model_active - assert model is not None + assert model is not None # type guard return file, model @@ -1087,12 +1087,12 @@ def test_dynamics_non_expression_params_unchanged(self): # expFun_02_t0 is [0, False, 0, 1] -> STATIC_PARAM t0 = _node_by_name(graph, "GLP_01_A_expFun_02_t0") - assert t0 is not None + assert t0 is not None # type guard assert t0.kind == NodeKind.STATIC_PARAM # expFun_01_A is [-1, True, -5, 0] -> OPT_PARAM A1 = _node_by_name(graph, "GLP_01_A_expFun_01_A") - assert A1 is not None + assert A1 is not None # type guard assert A1.kind == NodeKind.OPT_PARAM @@ -1195,7 +1195,7 @@ def test_expression_evaluated_per_sample(self): # GLP_02 should have per-sample COMPONENT_EVAL + EXPRESSION nodes glp02_avg = _node_by_name(graph, "GLP_02_profile_avg") - assert glp02_avg is not None + assert glp02_avg is not None # type guard assert glp02_avg.kind == NodeKind.PROFILE_AVERAGE # 5 sample component evals feed into GLP_02's profile average @@ -1205,7 +1205,7 @@ def test_expression_evaluated_per_sample(self): # Each per-sample EXPRESSION references a PROFILE_SAMPLE (not avg) for aux_i in range(5): expr_node = _node_by_name(graph, f"GLP_02_A_profile_expr_{aux_i}") - assert expr_node is not None + assert expr_node is not None # type guard assert expr_node.kind == NodeKind.EXPRESSION ref_edges = _edges_to(graph, expr_node.id, EdgeKind.EXPR_REF) @@ -1268,7 +1268,7 @@ def _make_time_dep_profile_model(dynamics_model=None): dynamics_model=dynamics_model or ["MonoExpPos"], ) model = file.model_active - assert model is not None + assert model is not None # type guard return file, model @@ -1380,7 +1380,7 @@ def _make_irf_dynamics_model(): dynamics_model=["MonoExpPosIRF"], ) model = file.model_active - assert model is not None + assert model is not None # type guard return file, model @@ -1709,8 +1709,8 @@ def test_visualize_passes_collapse(self): dot_collapsed = model.visualize(rendering="string", collapse_profiles=True) dot_full = model.visualize(rendering="string", collapse_profiles=False) - assert dot_collapsed is not None - assert dot_full is not None + assert dot_collapsed is not None # type guard + assert dot_full is not None # type guard assert "\u00d75" in dot_collapsed assert "profile_sample_1" in dot_full @@ -2177,7 +2177,7 @@ def test_expression_a_reads_resolved(self): if plan.expr_target_rows[i] == _find_row_for_name(graph, plan, "GLP_02_A"): expr_target_idx = i break - assert expr_target_idx is not None + 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] @@ -2343,8 +2343,8 @@ def test_chain_topological_order(self): if plan.expr_target_rows[i] == row_03_A: idx_03 = i - assert idx_02 is not None - assert idx_03 is not None + assert idx_02 is not None # type guard + assert idx_03 is not None # type guard assert idx_02 < idx_03 # diff --git a/tests/test_mcp_library.py b/tests/test_mcp_library.py index b652557..63fe05b 100644 --- a/tests/test_mcp_library.py +++ b/tests/test_mcp_library.py @@ -344,7 +344,7 @@ def test_voigt_kernel_axis_uses_both_width_parameters(self): t_mod.add_components([c_irf]) # Support should span the larger of 12*SD and 10*W. - assert c_irf.time is not None + assert c_irf.time is not None # type guard assert c_irf.time[0] == pytest.approx(-40.0) assert c_irf.time[-1] == pytest.approx(40.0) From 3fcc818c229a1fc613d51df8f08b91e360700d3e Mon Sep 17 00:00:00 2001 From: Johannes Mahl Date: Fri, 10 Jul 2026 14:56:24 -0700 Subject: [PATCH 32/36] reject cross-model dynamics expressions with a clear error; document t=0 invariant --- src/trspecfit/eval_2d.py | 5 +++++ src/trspecfit/mcp.py | 11 ++++++++++- tests/models/file_time.yaml | 10 +++++++++- tests/test_model_parser.py | 22 ++++++++++++++++++++++ 4 files changed, 46 insertions(+), 2 deletions(-) diff --git a/src/trspecfit/eval_2d.py b/src/trspecfit/eval_2d.py index 56a26d3..40481ba 100644 --- a/src/trspecfit/eval_2d.py +++ b/src/trspecfit/eval_2d.py @@ -194,6 +194,10 @@ def resolve_param_traces( func, _n_par = DYNAMICS_DISPATCH[func_id] n_par = int(dyn_sub_n_params[s]) param_rows = dyn_sub_param_rows[s, :n_par] + # Reading t=0 is exact: substep and kernel param rows are + # time-constant by construction — dynamics-model expressions + # cannot reference cross-model (potentially time-varying) + # parameters; add_dynamics rejects them. dyn_params = [float(traces[int(row), 0]) for row in param_rows] traces[target, :] += ( func(dyn_sub_time_axes[s], *dyn_params) * dyn_sub_masks[s] @@ -207,6 +211,7 @@ def resolve_param_traces( kernel_func, width_func = CONV_KERNEL_DISPATCH[func_id] p_start = int(conv_param_indptr[idx]) p_end = int(conv_param_indptr[idx + 1]) + # t=0 read is exact; same time-constant invariant as above kernel_params = [ float(traces[int(conv_param_rows[j]), 0]) for j in range(p_start, p_end) ] diff --git a/src/trspecfit/mcp.py b/src/trspecfit/mcp.py index 4796cf1..e7c6513 100644 --- a/src/trspecfit/mcp.py +++ b/src/trspecfit/mcp.py @@ -640,7 +640,16 @@ def add_dynamics(self, dynamics_model: "Dynamics", frequency: float = -1) -> Non ) # add Dynamics model and update corresponding parameter - target_par.update(dynamics_model) + try: + target_par.update(dynamics_model) + except NameError as e: + # lmfit evaluates dynamics expressions in the dynamics model's + # own parameter namespace; unknown names surface as NameError + raise ValueError( + f'Dynamics model for "{dynamics_model.name}" references an ' + f"unknown parameter ({e}). Expressions in a dynamics model " + "can only reference parameters of that same dynamics model." + ) from e # update model lmfit_par_list, parameter_names and components self.update() diff --git a/tests/models/file_time.yaml b/tests/models/file_time.yaml index fb04b02..5f2fb95 100644 --- a/tests/models/file_time.yaml +++ b/tests/models/file_time.yaml @@ -196,4 +196,12 @@ conv_last: tau: [2.5, True, 1, 10] t0: [0, False, 0, 1] gaussCONV: - SD: [5.0E-2, True, 0, 1] \ No newline at end of file + SD: [5.0E-2, True, 0, 1] +# expression referencing a parameter outside the dynamics model +# (must be rejected at add_time_dependence; energy-model parameters +# are not in a dynamics model's namespace) +CrossModelExpr: + expFun: + A: [1, True, 0, 5] + tau: ["GLP_01_x0 / 20"] + t0: [0, False, 0, 1] diff --git a/tests/test_model_parser.py b/tests/test_model_parser.py index 4668480..e5e9cab 100644 --- a/tests/test_model_parser.py +++ b/tests/test_model_parser.py @@ -397,6 +397,28 @@ def test_time_dependence_on_profiled_parameter_raises(self): dynamics_model=["MonoExpPos"], ) + # + def test_dynamics_expression_cross_model_reference_raises(self): + """Dynamics expressions cannot reference energy-model parameters. + + The lowered evaluator reads dynamics params as t=0 scalars + (eval_2d.resolve_param_traces); that is only exact because rows + feeding a dynamics substep are time-constant by construction. This + test pins the authoring-time rejection that guarantees it — if + cross-model references are ever allowed, the t=0 reads must be + revisited. + """ + + file = self._make_file_with_energy_model(model_energy=["simple_energy"]) + + with pytest.raises(ValueError, match="references an unknown parameter"): + file.add_time_dependence( + target_model="simple_energy", + target_parameter="GLP_01_A", + dynamics_yaml="models/file_time.yaml", + dynamics_model=["CrossModelExpr"], + ) + # # From ad68c38ed49138edc92acb6c623eac73bd20e42f Mon Sep 17 00:00:00 2001 From: Johannes Mahl Date: Fri, 10 Jul 2026 18:15:50 -0700 Subject: [PATCH 33/36] validate sweep specs, conv kernels, noise types, and DataFrame columns at the boundary --- TODO.md | 3 +- src/trspecfit/fitlib.py | 14 +++++++ src/trspecfit/simulator.py | 10 ++++- src/trspecfit/trspecfit.py | 4 +- src/trspecfit/utils/arrays.py | 9 ++++- src/trspecfit/utils/sweep.py | 19 +++++++++ tests/test_arrays.py | 18 +++++++++ tests/test_fitlib.py | 74 +++++++++++++++++++++++++++++++++++ tests/test_parameter_sweep.py | 63 +++++++++++++++++++++++++++++ 9 files changed, 208 insertions(+), 6 deletions(-) create mode 100644 tests/test_fitlib.py diff --git a/TODO.md b/TODO.md index 9e8120e..192e5c0 100644 --- a/TODO.md +++ b/TODO.md @@ -8,7 +8,7 @@ ## Noise and simulation -- [ ] **Simulator noise-language cleanup**: align simulator docs/metadata with the fit-results noise schema. Keep simulator `noise_type` meaning "noise distribution / random generator" (`gaussian`, `poisson`, `none`), not sigma shape. Fix the stale `Simulator.set_noise_type()` docstring that mentions `uniform`; clarify `detection` vs. `noise_type` vs. `noise_level`; and, for analog Gaussian simulations, consider saving the derived `sigma_data = noise_level * max(abs(clean_data))` alongside existing metadata. For parameter sweeps, store derived `sigma_data` per configuration when it depends on each clean dataset. +- [ ] **Simulator noise-language cleanup**: align simulator docs/metadata with the fit-results noise schema. Keep simulator `noise_type` meaning "noise distribution / random generator" (`gaussian`, `poisson`, `none`), not sigma shape (the stale `set_noise_type` docstring mentioning `uniform` was fixed and the setter validated, 2026-07-10). Clarify `detection` vs. `noise_type` vs. `noise_level`; and, for analog Gaussian simulations, consider saving the derived `sigma_data = noise_level * max(abs(clean_data))` alongside existing metadata. For parameter sweeps, store derived `sigma_data` per configuration when it depends on each clean dataset. - [ ] **Warn on subcycle-boundary time samples**: in `Dynamics.normalize_time`, emit a warning when a time sample lands within epsilon of a subcycle boundary (`|t * f_eff - round(t * f_eff)| < eps`). Subcycle masks switch discretely there, so the sample's assignment — and hence the model prediction for that whole row — is sensitive to floating-point representation of the time axis. This silently biased a multi-cycle fit when synthetic data was generated on `np.arange` axes but fit against their `%.6e`-rounded CSV reload (see `03_multi_cycle_dynamics/data/generate_data.ipynb`, 2026-06-11). A warning turns that silent bias into a one-line diagnosis; it also flags the physically ambiguous case of measuring exactly at the switching instant. - [ ] **Future `sigma_type` expansion in FitResults**: the constant, user-supplied sigma schema has landed (`SIGMA_TYPE_CONSTANT` in [fit_io.py](src/trspecfit/utils/fit_io.py) ~L54; `validate_noise_metadata` hard-locks `sigma_type` to `"constant"`), so this is now unblocked: extend uncertainty handling beyond scalar `sigma_data`. Keep `noise_type` for the statistical assumption/distribution and use `sigma_type` for sigma shape: initially `constant`, later `per_spectrum` and `per_point`. Add HDF5 storage, validation, baseline/SBS/2D alignment, `compare_models()` behavior, and tests for vector/matrix sigma. Defer automatic Poisson-derived sigma until residual-space variance propagation is explicit. @@ -21,6 +21,7 @@ - **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. - **Relocate the live accessors to `FitResults`**: 2026-06 added `File.get_correlations`, `File.get_conf_intervals`, `File.get_mcmc` (and the private `File._result_model` resolver) reading `model.result[...]` directly. These conceptually belong on `FitResults` (like `compare_models`, which already lives there with `File.compare_models` as thin sugar). The existing `File.get_fit_results` is in the same boat. Decide whether all of these should move into `FitResults` (with thin `File.*` sugar that delegates), and whether they read live `model.result` or persisted slots — then move them and update callers (notebook 12 reads them). - The raw list-index access (`result[1..4]`) and the deeper unified-results-object question are deferred to this item. +- [ ] **Decide how to guard/warn against in-place mutation of user-facing arrays**: internal machinery assumes `File.data`/`energy`/`time` and fit outputs are stable once set — e.g. `SavedFitSlot` stores `params`/`observed`/`fit`/`selection` by reference (`frozen=True` blocks reassignment, not in-place mutation; 2026-07 code review, check 1), and file fingerprints / `observed_sha256` are computed once at slot construction. A user mutating `file.data` in place instead of re-instantiating would desynchronize slots, fit limits, and cached evaluations in ways no single defensive copy fixes — so slot-level copies were considered and declined (2026-07-10) as papering over one symptom. Decide on a systemic stance instead: read-only views (`setflags(write=False)`) on public arrays, copy-on-set in setters, a documented ownership contract, and/or re-hash validation at save time. - [ ] **Disentangle plotting from saving/conversion in the fit pipeline**: figure rendering is currently entangled with data conversion and file IO. `fitlib.results_to_df` (results → DataFrame) and `File._save_2d_fit_legacy` / `_save_sbs_fit_legacy` (CSV writers) also render figures, and `fit_slice_by_slice` / `fit_2d` reach plotting only by calling that save-legacy path (`fit_baseline` is already disentangled — it calls `fitlib.plt_fit_res_1d` directly, gated by `_save_img_flag`; use it as the template). Blocks the legacy-shim removal item in the v1.0.0 checklist, since the `_save_*_legacy` impls are the live SbS/2D plotting path. Split into (a) compute/convert, (b) explicit save/export, (c) an explicit plotting API — cf. the `_save_img_flag` helper and `FitResults.plot_residuals`. Until then, the fit methods gate display on `show_output` and saving on `auto_export` via a `save_files` flag through the legacy methods. Scope: beyond the examples-upgrade branch. - [ ] **Decouple directory creation from path computation (`create_model_path`)**: `Project.create_model_path` ([trspecfit.py](src/trspecfit/trspecfit.py) ~L2167) `mkdir`s the `{path_results}/{file}/{fit_type}/{model}/` tree as a side effect of *computing* the path. Every fit method calls it to build the `save_path` handed to `fit_wrapper` (`fit_baseline` ~L2572, `fit_2d` ~L3983, spectrum ~L2804, sbs ~L3120), and the project plot-grid display builds image paths through it too (`fit_baselines` ~L1070, `fit_2d` ~L1432) — so empty `*_fits/` dir trees appear even when `auto_export: False` and nothing is written (observed in 21, 2026-06-27; every fitting example already does this). `fit_wrapper` only writes when `save_output=1 if auto_export else 0`. Fix: make `create_model_path` return the path *without* `mkdir`, and `mkdir`-on-write at the actual write sites (`fit_wrapper`'s save branch + the explicit `save_*`/`export_*` functions); the plot-grid readers then just build a path and `.exists()`-check. Do **not** gate on `auto_export` directly — explicit `save_fit()`/`export_fit()` legitimately need the dir and aren't `auto_export`-driven. Cosmetic only (the dirs are gitignored, no functional bug); blast radius spans baseline/spectrum/sbs/2d + savers, needs no-regression tests. Related to the plotting/saving disentanglement item above; scope: beyond the examples-upgrade branch. diff --git a/src/trspecfit/fitlib.py b/src/trspecfit/fitlib.py index 6ae2d3d..5f8db8f 100644 --- a/src/trspecfit/fitlib.py +++ b/src/trspecfit/fitlib.py @@ -1037,6 +1037,7 @@ def results_to_fit_2d( results: list[Any] | pd.DataFrame, const: tuple[Any, ...], args: tuple[Any, ...], + parameter_names: list[str] | None = None, num_fmt: str = "%.6e", delim: str = ",", save_2d: int = 0, @@ -1058,6 +1059,14 @@ def results_to_fit_2d( Each element: ``[par_ini, par_fin, conf_ci, emcee_fin, emcee_ci]`` - pd.DataFrame: From results_to_df() with parameters as columns + parameter_names : list of str, optional + For DataFrame results: select and order these columns as the + parameter vector before evaluation. Pass when the DataFrame may + carry extra non-parameter columns (e.g. the metrics columns in + ``results_to_df`` output); extra columns are otherwise passed to + the fit function as parameters. If None, all columns are used in + DataFrame order. Ignored for list results. + const : tuple Constants for residual_fun: (x, data, function_str, unpack, e_lim, t_lim) @@ -1095,6 +1104,11 @@ def results_to_fit_2d( e_lim_const, t_lim_const, ) = const + + # Select/order parameter columns; raises KeyError on missing names + if isinstance(results, pd.DataFrame) and parameter_names is not None: + results = results.loc[:, parameter_names] + lst = [] # intialize for i in range(len(results)): # list of lmfit_wrapper fit results diff --git a/src/trspecfit/simulator.py b/src/trspecfit/simulator.py index 5124908..fcefbc3 100644 --- a/src/trspecfit/simulator.py +++ b/src/trspecfit/simulator.py @@ -927,7 +927,7 @@ def set_noise_type(self, noise_type: str) -> None: Parameters ---------- noise_type : str - Noise distribution: ``'gaussian'`` or ``'uniform'``. + Noise distribution: ``'poisson'``, ``'gaussian'``, or ``'none'``. """ if self.detection != "analog": @@ -935,7 +935,13 @@ def set_noise_type(self, noise_type: str) -> None: "noise_type only applies to analog detection", stacklevel=2, ) - self.noise_type = noise_type.lower() + noise_type = noise_type.lower() + if noise_type not in ("poisson", "gaussian", "none"): + raise ValueError( + f"Unknown noise type: {noise_type}. " + "Use 'poisson', 'gaussian', or 'none'" + ) + self.noise_type = noise_type # def set_counts_per_delay(self, counts_per_delay: int) -> None: diff --git a/src/trspecfit/trspecfit.py b/src/trspecfit/trspecfit.py index cce6d24..bacc58b 100644 --- a/src/trspecfit/trspecfit.py +++ b/src/trspecfit/trspecfit.py @@ -3392,9 +3392,9 @@ def _save_sbs_fit_legacy( ) # get slice-by-slice fit spectra as a 2D map (write CSV only when saving) - df_sbs_pars = df_sbs.loc[:, self.model_sbs.parameter_names] fit_2d_sbs = fitlib.results_to_fit_2d( - results=df_sbs_pars, + results=df_sbs, + parameter_names=self.model_sbs.parameter_names, const=self.model_sbs.const, args=self.model_sbs.args, num_fmt=self.p.num_fmt, diff --git a/src/trspecfit/utils/arrays.py b/src/trspecfit/utils/arrays.py index ca68c3d..b255746 100644 --- a/src/trspecfit/utils/arrays.py +++ b/src/trspecfit/utils/arrays.py @@ -391,7 +391,14 @@ def my_conv( y_pad = np.pad(y_arr, pad_size, mode="edge") # Normalize the kernel (cheaper than dividing the padded signal) - kernel_norm = kernel_arr / np.sum(kernel_arr) + kernel_sum = np.sum(kernel_arr) + if kernel_sum == 0 or not np.isfinite(kernel_sum): + raise ValueError( + f"my_conv kernel sums to {kernel_sum}; cannot normalize. " + "The kernel likely collapsed below the axis step " + "(width parameter too small) or contains non-finite values." + ) + kernel_norm = kernel_arr / kernel_sum # Compute convolution with normalized kernel if method == "scipy": diff --git a/src/trspecfit/utils/sweep.py b/src/trspecfit/utils/sweep.py index faeef43..2f465f8 100644 --- a/src/trspecfit/utils/sweep.py +++ b/src/trspecfit/utils/sweep.py @@ -225,6 +225,16 @@ def add_lognormal( "n_samples": n_samples, } + # + def _require_parameters(self) -> None: + """Raise if no parameters have been added to the sweep.""" + + if not self.parameter_specs: + raise ValueError( + "No parameters added to this sweep. Use add_range / " + "add_uniform / add_normal / add_lognormal first." + ) + # def _determine_strategy(self) -> str: """ @@ -341,6 +351,12 @@ def _generate_random(self) -> Generator[dict[str, float], None, None]: config[par_name] = self.rng.normal(spec["mean"], spec["std"]) elif spec["type"] == "lognormal": config[par_name] = self.rng.lognormal(spec["mean"], spec["std"]) + else: + raise ValueError( + f"Unknown distribution type '{spec['type']}' for " + f"parameter '{par_name}'. Must be 'range', " + "'uniform', 'normal', or 'lognormal'." + ) yield config # @@ -354,6 +370,8 @@ def __iter__(self) -> Generator[dict[str, float], None, None]: Parameter configuration {par_name: value, ...} """ + self._require_parameters() + # Reset seed at start of iteration for reproducibility if self.seed is not None: self.rng = np.random.default_rng(self.seed) @@ -384,6 +402,7 @@ def get_n_configs(self) -> int: 12 """ + self._require_parameters() strategy = self._determine_strategy() if strategy == "grid": diff --git a/tests/test_arrays.py b/tests/test_arrays.py index 10345ae..d10edf1 100644 --- a/tests/test_arrays.py +++ b/tests/test_arrays.py @@ -44,6 +44,24 @@ def test_single_element_x_raises(self): with pytest.raises(ValueError, match="at least 2 x samples"): my_conv(np.array([0.0]), np.array([1.0]), np.array([1.0])) + # + def test_zero_sum_kernel_raises(self): + """A kernel that sums to zero raises instead of yielding NaN/inf.""" + + x = np.linspace(0, 10, 21) + y = np.sin(x) + with pytest.raises(ValueError, match="cannot normalize"): + my_conv(x, y, np.zeros(5)) + + # + def test_nonfinite_kernel_raises(self): + """A kernel with non-finite entries raises instead of propagating.""" + + x = np.linspace(0, 10, 21) + y = np.sin(x) + with pytest.raises(ValueError, match="cannot normalize"): + my_conv(x, y, np.array([1.0, np.nan, 1.0])) + # # diff --git a/tests/test_fitlib.py b/tests/test_fitlib.py new file mode 100644 index 0000000..3e93a3b --- /dev/null +++ b/tests/test_fitlib.py @@ -0,0 +1,74 @@ +"""Unit tests for fitlib bridge functions.""" + +import numpy as np +import pandas as pd +import pytest +from _utils import make_project + +from trspecfit import File, fitlib + + +# +def _make_1d_model_file(): + """File with a loaded 1D energy model (public API).""" + + project = make_project(name="fitlib") + file = File(parent_project=project, energy=np.linspace(80, 90, 101)) + file.load_model(model_yaml="models/file_energy.yaml", model_info="single_glp") + assert file.model_active is not None # type guard + return file + + +# +# +class TestResultsToFit2D: + """DataFrame-path column handling in results_to_fit_2d.""" + + # + def _const_args(self, file): + """const/args as the SbS fit path builds them (per-slice, MCP).""" + + model = file.model_active + assert model is not None # type guard + data = np.zeros_like(np.asarray(file.energy)) + const = (file.energy, data, "fit_model_mcp", 0, [], []) + args = (model, 1) + return model, const, args + + # + def test_parameter_names_selects_and_orders_columns(self): + """Extra columns and scrambled order are handled via parameter_names.""" + + file = _make_1d_model_file() + model, const, args = self._const_args(file) + + values = [model.lmfit_pars[n].value for n in model.parameter_names] + df_pars = pd.DataFrame([values, values], columns=model.parameter_names) + df_extra = df_pars.copy() + df_extra["chi2"] = [0.1, 0.2] # non-parameter column + # scramble column order on top of the extra column + df_extra = df_extra[["chi2", *reversed(model.parameter_names)]] + + fit_ref = fitlib.results_to_fit_2d(df_pars, const, args) + fit_sel = fitlib.results_to_fit_2d( + df_extra, const, args, parameter_names=model.parameter_names + ) + + np.testing.assert_allclose(fit_sel, fit_ref) + assert fit_ref.shape == (2, len(file.energy)) + + # + def test_missing_parameter_column_raises(self): + """Requesting a parameter column absent from the DataFrame fails.""" + + file = _make_1d_model_file() + model, const, args = self._const_args(file) + + values = [model.lmfit_pars[n].value for n in model.parameter_names] + df_pars = pd.DataFrame([values], columns=model.parameter_names) + df_missing = df_pars.drop(columns=model.parameter_names[:1]) + + with pytest.raises(KeyError): + fitlib.results_to_fit_2d( + df_missing, const, args, parameter_names=model.parameter_names + ) diff --git a/tests/test_parameter_sweep.py b/tests/test_parameter_sweep.py index 79d20ea..ef94885 100644 --- a/tests/test_parameter_sweep.py +++ b/tests/test_parameter_sweep.py @@ -221,6 +221,69 @@ def test_iteration_multiple_times(self): assert configs1 == configs2 assert len(configs1) == 5 + # + def test_empty_sweep_raises(self): + """A sweep with no parameters must fail loudly, not yield {}.""" + + sweep = ParameterSweep(strategy="grid") + + with pytest.raises(ValueError, match="No parameters added"): + sweep.get_n_configs() + with pytest.raises(ValueError, match="No parameters added"): + list(sweep) + + # + def test_unknown_spec_type_raises(self): + """Random generation must reject unrecognized distribution types.""" + + sweep = ParameterSweep(strategy="random") + sweep.add_uniform("param_A", 0, 10, n_samples=3) + # specs are only created by add_*; inject directly to pin the guard + sweep.parameter_specs["param_B"] = {"type": "cauchy", "n_samples": 3} + + with pytest.raises(ValueError, match="Unknown distribution type"): + list(sweep) + + +# +# +class TestSimulatorNoiseType: + """Test Simulator.set_noise_type validation.""" + + # + def _make_simulator(self): + """Analog Simulator on a minimal 1D energy model.""" + + project = make_project(name="test") + file = File( + parent_project=project, + energy=np.arange(0, 20, 0.5), + time=np.arange(-10, 100, 5), + ) + file.load_model( + model_yaml="models/file_energy.yaml", model_info="simple_energy" + ) + assert file.model_active is not None # type guard + return Simulator( + model=file.model_active, + detection="analog", + noise_level=0.05, + noise_type="gaussian", + seed=42, + ) + + # + def test_set_noise_type_normalizes_case(self): + sim = self._make_simulator() + sim.set_noise_type("Poisson") + assert sim.noise_type == "poisson" + + # + def test_set_noise_type_unknown_raises(self): + sim = self._make_simulator() + with pytest.raises(ValueError, match="Unknown noise type"): + sim.set_noise_type("uniform") + # # From d7f0e27887a55498ab90a04de5b27cfb79c039d7 Mon Sep 17 00:00:00 2001 From: Johannes Mahl Date: Fri, 10 Jul 2026 18:42:34 -0700 Subject: [PATCH 34/36] validate noise_type in Simulator constructor via shared helper --- src/trspecfit/simulator.py | 24 ++++++++++++++++-------- tests/test_parameter_sweep.py | 14 ++++++++++++++ 2 files changed, 30 insertions(+), 8 deletions(-) diff --git a/src/trspecfit/simulator.py b/src/trspecfit/simulator.py index fcefbc3..a1393ec 100644 --- a/src/trspecfit/simulator.py +++ b/src/trspecfit/simulator.py @@ -232,6 +232,7 @@ def __init__( ------ ValueError If detection type is invalid + If noise_type is invalid If counts_per_delay ≤ 0 (after estimation) Examples @@ -260,7 +261,7 @@ def __init__( # Analog detector parameters self.noise_level = noise_level - self.noise_type = noise_type.lower() + self.noise_type = self._validated_noise_type(noise_type) # Photon counting parameters self.counts_per_delay: int | None = counts_per_delay @@ -285,6 +286,19 @@ def __init__( self.data_noisy: np.ndarray | None = None # With noise self.noise: np.ndarray | None = None # Just the noise component + # + @staticmethod + def _validated_noise_type(noise_type: str) -> str: + """Normalize and validate a noise_type string.""" + + noise_type = noise_type.lower() + if noise_type not in ("poisson", "gaussian", "none"): + raise ValueError( + f"Unknown noise type: {noise_type}. " + "Use 'poisson', 'gaussian', or 'none'" + ) + return noise_type + # def __repr__(self) -> str: return f"Simulator(model='{self.model.name}', detection='{self.detection}')" @@ -935,13 +949,7 @@ def set_noise_type(self, noise_type: str) -> None: "noise_type only applies to analog detection", stacklevel=2, ) - noise_type = noise_type.lower() - if noise_type not in ("poisson", "gaussian", "none"): - raise ValueError( - f"Unknown noise type: {noise_type}. " - "Use 'poisson', 'gaussian', or 'none'" - ) - self.noise_type = noise_type + self.noise_type = self._validated_noise_type(noise_type) # def set_counts_per_delay(self, counts_per_delay: int) -> None: diff --git a/tests/test_parameter_sweep.py b/tests/test_parameter_sweep.py index ef94885..d8e7b5b 100644 --- a/tests/test_parameter_sweep.py +++ b/tests/test_parameter_sweep.py @@ -284,6 +284,20 @@ def test_set_noise_type_unknown_raises(self): with pytest.raises(ValueError, match="Unknown noise type"): sim.set_noise_type("uniform") + # + def test_constructor_unknown_noise_type_raises(self): + """The constructor shares the setter's validation.""" + + sim = self._make_simulator() + with pytest.raises(ValueError, match="Unknown noise type"): + Simulator( + model=sim.model, + detection="analog", + noise_level=0.05, + noise_type="uniform", + seed=42, + ) + # # From 3dee6beb48b0bf7036461e466f6b8976d521e03e Mon Sep 17 00:00:00 2001 From: Johannes Mahl Date: Fri, 10 Jul 2026 20:50:54 -0700 Subject: [PATCH 35/36] archive the July 2026 code review; document full-scope review protocol --- docs/ai/code-review.md | 32 ++ docs/design/archive/code-review-2026-07.md | 454 +++++++++++++++++++++ 2 files changed, 486 insertions(+) create mode 100644 docs/design/archive/code-review-2026-07.md diff --git a/docs/ai/code-review.md b/docs/ai/code-review.md index 995153d..3c6b939 100644 --- a/docs/ai/code-review.md +++ b/docs/ai/code-review.md @@ -31,6 +31,38 @@ For each item report one of: Work through the checklist in order. Use parallel agent/tool calls where items are independent. +## Full-scope protocol + +For `full` runs (all of `src/`, ~24k lines — too much for one honest pass), +split the work as follows: + +- **Repo-wide grep checks** — items 3 (broad exceptions), 5 (dead code), + 6 (typing), 9 (numpy anti-patterns), 10 (float `==`), 11 (ignored + warnings), 15 (global state), and 16 (security) run **once globally** + in a single grep session, with findings binned by file. They are cheap; + chunking them would duplicate work. +- **Read-the-code checks** — items 1, 2, 4, 7, 8, 12, 13, 14, 18, 19, + and 20 run **per chunk** via parallel subagents. Chunks follow the + two-layer architecture (see `docs/design/repo_architecture.md`): + - A. Authoring layer: `trspecfit.py`, `mcp.py`, `utils/parsing.py` + - B. Compiled hot path: `graph_ir.py`, `eval_1d.py`, `eval_2d.py`, + `functions/` (checks 18–20 concentrate here) + - C. Fitting & bridge: `fitlib.py`, `spectra.py`, `utils/lmfit.py`, + `utils/sbs.py` + - D. Persistence & results: `utils/fit_io.py`, `fit_results.py`, + `utils/hdf5.py` + - E. Simulation, plotting, config: `simulator.py`, `utils/sweep.py`, + `utils/plot.py`, `utils/arrays.py`, `config/` + - F. Tests light pass: check 17 plus CLAUDE.md test-pattern rules + (plain pytest, public API usage, `show_plot=False`), and + verification that parity coverage claimed in chunk B exists in + `tests/test_gir_integration.py`. Runs after B. +- Each subagent returns only concise findings: file:line, severity + (PASS/INFO/WARN/FAIL), one-line description. +- Findings are collected into a report file (e.g. + `docs/design/code-review-.md`) for triage — never fixed inline + during the review. + ## 1. Bugs and correctness Read the code in scope looking for: diff --git a/docs/design/archive/code-review-2026-07.md b/docs/design/archive/code-review-2026-07.md new file mode 100644 index 0000000..70276db --- /dev/null +++ b/docs/design/archive/code-review-2026-07.md @@ -0,0 +1,454 @@ +--- +orphan: true +--- + +# Full-Repo Code Review — July 2026 + +> Archived on 2026-07-10 after every finding was resolved: fixed on the +> `fix-conv-kernels` branch, declined with rationale recorded inline, or +> moved to a `TODO.md` item. The full-scope review protocol developed for +> this run now lives in `docs/ai/code-review.md`. + +Reviewed at commit `a1df0cc` using the checklist in `docs/ai/code-review.md` +(full scope). Method: repo-wide grep checks run once globally; read-the-code +checks run per architecture chunk (A authoring, B compiled hot path, +C fitting/bridge, D persistence, E simulation/plotting/config, F tests light +pass) by parallel subagents. Findings verified by spot-checks where marked. + +Scope: deep review of `src/`, light pass on `tests/` (edge-case coverage, +patterns, parity verification). Notebooks, YAML assets, docs excluded. + +**Triage key** (fill in during triage): `[F]` fix now, `[L]` later, `[W]` won't fix. + +## Summary table + +| # | Check | Status | Issues | +|---|-------|--------|--------| +| 1 | Bugs & correctness | **FAIL** | 4 FAIL, 12 WARN, 6 INFO | +| 2 | Performance | WARN | 12 WARN, 4 INFO | +| 3 | Broad exceptions | WARN | 1 WARN, 1 INFO | +| 4 | God classes / long methods | INFO | size mostly justified | +| 5 | Dead code | PASS | 1 INFO | +| 6 | Typing / modern Python syntax | PASS | no `Optional[`/`Union[` | +| 7 | Docstring coverage | WARN | 1 WARN, 2 INFO | +| 8 | Abstractions & duplication | WARN | 9 WARN, 5 INFO | +| 9 | Numpy anti-patterns | INFO | 1 INFO | +| 10 | Fragile array comparisons | PASS | only zero-guards found | +| 11 | Ignored warnings | PASS | none in src; 1 scoped pytest mark | +| 12 | Plotting mixed with logic | WARN | 8 WARN, 1 INFO | +| 13 | Separation of concerns | WARN | 8 WARN, 4 INFO | +| 14 | Missing `__repr__` | INFO | 8 dataclasses/classes | +| 15 | Global mutable state | INFO | 1 documented worker pattern | +| 16 | Security | PASS | safe YAML loaders, no eval/pickle/shell | +| 17 | Edge-case tests | INFO | 3 concrete gaps | +| 18 | GIR / MCP parity | WARN | core parity strong; 6 verified gaps | +| 19 | Two-layer design compliance | WARN | 4 WARN, 3 INFO | +| 20 | `can_lower_*` hygiene | INFO | 2 implicit handlings, otherwise complete | + +## FAIL findings (check 1 unless noted) + +All four spot-verified against source. + +- [x] `src/trspecfit/simulator.py:1304` — `plot_comparison` builds the plot + title via `self.get_snr(...)` *before* the `data_clean is None` + auto-simulate guard at :1310, so calling it on a fresh Simulator raises + instead of simulating first. Fix: move title construction below the + auto-simulate block. +- [x] `src/trspecfit/utils/arrays.py:241` — `sign_change(ignore_zeros=True)` + infinite-loops on all-zero input: `asign` stays all-zero so the + `while sz.any()` loop never terminates. Fix: bail out (return zeros) when + `asign.any()` is False. +- [x] `src/trspecfit/utils/arrays.py:383` — `my_conv` computes + `x_arr[1] - x_arr[0]` with no `len(x) >= 2` guard; single-element input + raises `IndexError`. Called from the hot-path 2D evaluator. Fix: guard or + document the precondition at the call boundary. +- [x] `src/trspecfit/utils/plot.py:703` — `plot_1d(y_norm=1)` divides by + `np.max(y_data - np.min(y_data))`, which is 0 for a constant trace + (produces inf/NaN plot silently). Fix: guard the zero-range case. + +## WARN findings + +### Check 1 — Bugs and correctness + +- [x] `src/trspecfit/mcp.py:992` — `create_value_2d(t_ind=[start,stop])` + passes loop index `ti` (0..n-1) to `create_value_1d` instead of + `start+ti`; partial-range evaluation is wrong. Latent: never called with + `t_ind` in-repo. +- [x] `src/trspecfit/mcp.py:2248` — `Par.value` returns `-1.0` and prints + when `t_vary` is set but `t_model` is None, instead of raising. +- [x] `src/trspecfit/trspecfit.py:1702`, `:2313`, `:2383` — `File.describe`, + `define_baseline`, and `set_fit_limits` mutate `self.energy`/`self.time` + in-place when axes are missing (side effects in inspection/setup paths). +- [x] `src/trspecfit/eval_2d.py:362` and `src/trspecfit/graph_ir.py:2916` — + `time[1] - time[0]` assumes `n_time >= 2`; a single-point time axis raises + `IndexError`. Cross-reference: no test covers this (check 17). +- [x] `src/trspecfit/eval_2d.py:343` — dynamics substeps read only + `traces[row, 0]`; a time-varying expression row would silently use its + t=0 value. + (Investigated 2026-07-10: latent — rows feeding dynamics substeps and + conv kernels are time-constant by construction. Dynamics-model + expressions are evaluated in the dynamics model's own lmfit namespace, + so cross-model references (the only route to a time-varying row) fail + at `add_time_dependence` before any graph build or fit; probed all + three routes (direct t_vary ref, static ref, via top-level expression + param). Hardened: `add_dynamics` now re-raises the deep lmfit + `NameError` as a clear `ValueError`; the invariant is documented at + both t=0 read sites; and + `test_dynamics_expression_cross_model_reference_raises` pins the + rejection so the reads get revisited if cross-model refs are ever + allowed. Follow-up considered and declined 2026-07-10: the rejection is + not atomic (target par keeps t_vary/t_model if the caller catches the + ValueError), but the leaked state re-raises on every subsequent + evaluation (verified: create_value_1d/2d, par.value) — loud, not + silent. The authoring API is non-atomic throughout (add_profile, + load_model), so a one-off rollback would imply a transactional + guarantee the layer doesn't have; recovery is reloading the model.) +- [x] `src/trspecfit/functions/energy.py:144` — `LinBack` raises `ValueError` + (xStart >= xStop) inside the hot-path numeric body; validation belongs in + the authoring layer (also flagged under check 19). + (Resolved under check 19, 2026-07-10: guard kept by design — see the + check-19 note.) +- [x] `src/trspecfit/fitlib.py:1097` — `results_to_fit_2d` DataFrame path + passes `iloc[i].values` (all columns); extra non-parameter columns break + the fit unless the caller pre-filters (the public API does). + (2026-07-10: added optional `parameter_names` kwarg that selects and + orders DataFrame columns (KeyError on missing); the SbS caller now + passes it instead of pre-filtering. Tested incl. scrambled column + order.) +- [x] `src/trspecfit/utils/fit_io.py:796` and `:171` — `SavedFitSlot` stores + `params`/`conf_ci`/`selection`/`metrics` by reference; `frozen=True` + blocks reassignment but not in-place DataFrame/dict mutation, weakening + the snapshot contract. + (Declined 2026-07-10: a user mutating `file.data`/results in place breaks + far more than slots (fingerprints, fit limits, cached evaluations), so a + slot-level defensive copy papers over one symptom. Deferred to a TODO.md + item on a systemic stance — read-only arrays, copy-on-set, ownership + contract, or save-time re-hash.) +- [x] `src/trspecfit/utils/sweep.py:334` — `_generate_random` silently skips + unknown spec types, yielding incomplete configs. + (2026-07-10: `else: raise ValueError` naming parameter and type, + mirroring the grid path's `_sample_distribution` which already raised.) +- [x] `src/trspecfit/utils/sweep.py:400` and `:309` — `get_n_configs()` on + empty `parameter_specs` raises `ValueError` from `max()`; empty sweep + yields one `{}` config. + (2026-07-10: `_require_parameters()` guard in both `get_n_configs` and + `__iter__` with a clear "No parameters added" error.) +- [x] `src/trspecfit/utils/arrays.py:388` — `my_conv` normalizes by + `np.sum(kernel_arr)` with no zero-sum guard. + (2026-07-10: raises `ValueError` on zero or non-finite kernel sum, + naming the likely cause (kernel width collapsed below the axis step); + one scalar comparison on an already-computed sum, no hot-path cost.) +- [x] `src/trspecfit/simulator.py:982` — `set_noise_type` docstring lists + `'uniform'` but only poisson/gaussian/none are handled. + (2026-07-10: docstring fixed to poisson/gaussian/none; setter AND + constructor now validate via a shared `_validated_noise_type` helper, so + a typo fails at construction/set time instead of at the first + `add_noise`. The broader noise-language cleanup stays a TODO.md item.) + +### Check 2 — Performance (hot path and bulk loops) + +- [x] `src/trspecfit/eval_2d.py:111`, `:108` — RPN evaluator copies the full + trace row on every `PARAM_REF` and allocates `np.full(n_time, ...)` per + constant instruction, per residual evaluation. + (Fixed in `b139601`: parameter rows are views, constants remain scalar, + and only constant-only programs allocate the required output row.) +- [x] `src/trspecfit/eval_2d.py:174-176`, `:222` — per-eval + `broadcast_to(...).copy()` per profile-sample group; `np.repeat(traces, + n_aux, axis=1)` builds a large matrix for profile expressions each eval. + (Fixed in `b139601`: both broadcasts now write directly into their + preallocated destination buffers.) +- [x] `src/trspecfit/eval_2d.py:255-279` — profiled ops loop `n_aux` times + calling energy functions instead of a vectorized aux broadcast. + (Evaluated both ways 2026-07-10 after fixing the benchmark harness to + actually attach example 04's profiles: aux vectorization is ~4.5x faster + only when profiled params enter the function linearly (amplitude-only), + and ~60% *slower* with a profiled position (example 04, n_aux=50) because + the full `(n_time, n_aux, n_energy)` temporaries materialize inside the + energy functions. Kept the loop; param-source resolution hoisted out of + it. See `_evaluate_profiled_op_2d` docstring.) +- [x] `src/trspecfit/fitlib.py:233`, `:238` — `residual_fun` resolves + `fit_fun_str` via `getattr(spectra, ...)` and calls `par_extract` on every + residual evaluation; both hoistable to setup. + (Evaluated 2026-07-10, not changed: measured 0.07 + 2.5 us per call vs + 1-8.5 ms per model eval, i.e. <0.3%; hoisting would change the `const` + contract across all five fit entry points for no observable gain.) +- [x] `src/trspecfit/spectra.py:292` — `fit_project_mcp` rebuilds the + `par_lookup` dict and runs full MCP `create_value_2d` per file on every + project-level residual (also check 19). + (Declined 2026-07-10: MCP stays the readable reference layer; no + micro-optimization there. The real fix is a lowered GIR project + evaluator — deliberately a later TODO.) +- [x] `src/trspecfit/fitlib.py:782` — MCMC walker/corner figures always built + when `use_emcee==1`, even with `show_output=0, save_output=0`. + (Figure construction now skipped entirely when neither shown nor saved.) +- [x] `src/trspecfit/utils/arrays.py:383`, `:388` — `my_conv` copies + x/y/kernel and recomputes `np.sum(kernel)` on every hot-path call; + padded x array is built then discarded. + (`b139601` removes the discarded padded-x construction and normalizes the + short kernel instead of dividing the padded signal. The live kernel changes + each evaluation, so its sum is still computed once per call; evaluator + float64 inputs pass through `np.asarray` without copies. `pad_x_y` is now + unused repo-wide — decision 2026-07-10: keeping it.) +- [x] `src/trspecfit/simulator.py:1984`, `:748`, `:798` — per-config HDF5 + open/close in sweep append; `simulate_n` appends per iteration instead of + pre-allocating; Poisson scale factor recomputed per `add_noise` call. + (2026-07-10: fixed the HDF5 part — `simulate_parameter_sweep` opens the + file once and passes the handle to `_initialize_sweep_hdf5` / + `_append_config_to_hdf5`, with a per-config `flush()` preserving + interrupted-sweep durability. The `simulate_n` pre-allocation and Poisson + scale-factor items were declined: they only speed up the analog path + and/or couple `simulate_n` to photon-sampling internals.) + +### Check 3 — Broad exceptions + +- [x] `src/trspecfit/trspecfit.py:825` — config-loading `except Exception` + only prints when `show_output >= 1`; with `show_output=0` errors are + swallowed silently. (`mcp.py:1545`, `:2368` re-raise as `ValueError` — + fine. `trspecfit.py:3290` `except BaseException` cancels futures and + re-raises — fine. `utils/lmfit.py:125` warns — acceptable.) + (Fixed in `f0e5e66`: missing files retain the designed fallback; all other + load/validation failures are re-raised as `ValueError`.) + +### Check 7 — Docstring coverage + +- [x] `src/trspecfit/functions/profile.py:47`, `:70`, `:92` — `pExpDecay`, + `pLinear`, `pGauss` lack `Returns` sections (user-facing `functions/`). + +### Check 8 — Abstractions and duplication + +- [x] `src/trspecfit/eval_1d.py:27` vs `graph_ir.py:3048` — scalar RPN + evaluator implemented twice (`eval_expr_program_1d` vs + `_eval_expr_scalar`). + (Deduped 2026-07-10: `eval_1d` now imports `_eval_expr_scalar` from + `graph_ir`, matching the other runtime helpers it already imports; + `eval_expr_program_1d` deleted.) +- [x] `src/trspecfit/graph_ir.py:2433-2657` vs `:3407-3608`, and + `:2707-2835` vs `:3659-3793` — PROFILE_SAMPLE/EXPR compilation and + PROFILE_AVERAGE op-wiring largely copy-pasted between `schedule_2d` and + `schedule_1d`. + (2026-07-10: sample/expr compilation extracted into + `_compile_profile_groups` and op scheduling into + `_schedule_component_ops` (NamedTuple returns); each shared body is the + superset — 2D-only PPT/CONVOLUTION walk-backs never fire on 1D graphs, + 1D-only "not scalar-lowerable"/"Non-scalar parameter source" errors + never fire on 2D, where every node has a row. Net −223 lines.) +- [x] `src/trspecfit/graph_ir.py:2866` vs `eval_2d.py:31` — `_DYN_DISPATCH` + mirrors `DYNAMICS_DISPATCH` with separate function references. + (Resolved 2026-07-10 together with the item below — the mirror only + existed to serve the duplicated loop.) +- [x] `src/trspecfit/graph_ir.py:2882-2921` vs `eval_2d.py:329-365` — + resolution loop duplicated at compile-time init vs hot-path eval. + (Extracted 2026-07-10 into `eval_2d.resolve_param_traces`; `evaluate_2d` + and `schedule_2d` init both call it. Array-argument signature keeps it + usable at compile time before the plan exists.) +- [x] `src/trspecfit/simulator.py:769`, `:1650`, `:1680` — near-identical + 1D/2D analog-noise generators; HDF5 metadata serialization duplicated + across `_save_hdf5` / `_initialize_sweep_hdf5`; params-to-JSON loop + repeated in three HDF5 writers. + (2026-07-10: `_generate_noise_analog_1d/2d` were byte-identical (all + numpy ops shape-agnostic) — merged into `_generate_noise_analog`; axes + write, detection/seed attrs, and full params-to-JSON extracted into + `_write_axes_hdf5` / `_write_detection_metadata` / + `_model_parameters_json`. The value-only params loop in + `_append_config_to_hdf5` is a different output and stays. The + `_sample_photons_1d/2d` pair genuinely differs — total vs per-row + normalization — and stays split.) + +### Check 12 — Plotting mixed with logic + +- [x] `src/trspecfit/fitlib.py:749`, `:764`, `:781` — MCMC progress print + and `progress=True` run unconditionally; with `show_output=0, + save_output=0` MCMC plots still reach `plt.show()` via + `_finalize_plot(0, ...)`. + (Fixed in `f0e5e66`; the separate cost of constructing unsaved MCMC figures + remains open under check 2.) +- [x] `src/trspecfit/trspecfit.py:4061` — `File.fit_2d` calls `time_display` + and `display(params)` whenever `stages>=1`, ignoring `show_output` + (fixed in `f0e5e66` together with the same pattern in + `fit_slice_by_slice`). +- [x] `src/trspecfit/trspecfit.py:2335`, `:2425` — `define_baseline` / + `set_fit_limits` plot on `show_plot=True` default without consulting + `Project.show_output`. (Fixed in `f0e5e66`.) +- [x] `src/trspecfit/simulator.py:1304`, `:1310` — `plot_comparison` mixes + SNR computation and auto-simulate side effects into the plot path + (see also the FAIL above). + (2026-07-10: the auto-simulate is deliberate (documented, regression + test) and the SNR title is the feature. The actionable part was the + hand-rolled 3-panel loop duplicating `plot_2d_grid` — now delegated + (`columns=3`; grid gained `columns`, x/y-lim, and ticksize support so + nothing was lost). Panel size now follows the grid convention.) + +### Check 13 — Separation of concerns + +- [x] `src/trspecfit/utils/fit_io.py:35`, `:1139` — raw `h5py` throughout + instead of the `utils/hdf5.py` helpers the architecture doc mandates + (`require_group`/`require_dataset`/`json_loads_attr`). Either migrate or + amend the contract in `repo_architecture.md`. + (Both, 2026-07-10: fit_io's local `_as_group`/`_as_dataset` duplicated + the shared helpers — deleted; all 25 read sites now use + `require_group`/`require_dataset` with real archive paths in errors. + Contract amended: helpers cover *reads* (type narrowing + JSON attrs); + writes have no wrapper and use `h5py` directly.) +- [x] `src/trspecfit/simulator.py:1620` and `utils/sweep.py:482` — same raw + `h5py.File` usage in `_save_hdf5`/`_initialize_sweep_hdf5` and + `SweepDataset` reads. + (Stale finding, 2026-07-10: `SweepDataset` reads already go through + `require_group`/`require_dataset`/`json_loads_attr`, and the simulator + sites are write-only — covered by the amended contract above.) +- [x] `src/trspecfit/fitlib.py:446` — `fit_wrapper` combines optimization, + CI, MCMC, plotting, CSV/TXT I/O, and notebook display in one ~207-line + function. + (Declined 2026-07-10: decomposition not wanted; it stays one function + unless a concrete need arises.) +- [x] `src/trspecfit/simulator.py:1347`, `utils/plot.py:246`, `:516` — + hardcoded `figsize`s and reference-line colors outside `PlotConfig` + (PlotConfig is documented as the single source of truth for styling). + (2026-07-10: added `refline_color`/`refline_style` (unified default + grey `#808080` dotted — 2D reflines were black, 1D vlines dashed) and + `panel_size` (default 4×3 per panel) to PlotConfig; `plot_1d`/`plot_2d` + accept them as per-call kwargs, `plot_2d_grid` reads them from config. + The simulator figsize disappeared via the plot_comparison delegation + below.) + +### Check 18 — GIR/MCP parity (verified against tests by chunk F) + +Core parity is strong: 2D baseline/IRF/subcycle/profile residual+compare, +1D compare/residual, and end-to-end `fit_model_compare` through `fit_2d`, +`fit_baseline`, `fit_spectrum`, `fit_slice_by_slice` (serial + 2 workers) +all exist in `tests/test_gir_integration.py`; roundtrip matrix backend `C` +adds broad coverage. Verified gaps at the `fit_model_compare` level: + +- [x] Energy shapes `GaussAsym`, `Lorentz`, `Voigt`, `GLS`, `DS`, `LinBack`: + evaluator-level parity only (`tests/test_evaluate_2d.py:146-171`), no + pipeline/compare coverage. (`Gauss` and `Shirley` partially covered via + roundtrip families and profile tests.) + (`test_compare_mode_energy_shapes`, parametrized over all 7 incl. Shirley.) +- [x] Dynamics `sinFun`, `linFun`, `sinDivX`, `erfFun`, `sqrtFun`, + `stepFun`: pure-math tests only (`tests/test_functions_time.py`); only + `expFun` has parity coverage. + (`test_residual_same_gir_vs_mcp_dynamics`, parametrized over all 6.) +- [x] Profile `pGauss`: pure-math only; no profile YAML fixture. + (`profile_pGauss` fixture; compare-mode profile 1D/2D tests parametrized.) +- [x] Multi-substep dynamics without subcycles (`frequency=-1`/omitted): + no parity test (all multi-substep tests use `frequency=10`). + (`test_residual_same_gir_vs_mcp_multi_substep_single_cycle` via + `BiExpSharedT0`.) +- [x] Chained CONVOLUTION nodes: structural gate test only + (`tests/test_graph_ir.py:1434`); no compare-backend coverage. + (`MonoExpPosDoubleIRF` fixture; parity test pins `n_conv_steps == 2`.) +- [x] TIME_1D-only domain models: graph tests only + (`tests/test_graph_ir.py:1101`); MCP fallback untested through the fit + pipeline. (`test_time_1d_dynamics_model_mcp_fallback` at the + residual_fun/fit_model_gir-fallback level; standalone trace *fitting* is + itself still a deferred TODO, so no deeper pipeline entry exists to test.) +- [x] `test_compare_mode_irf` covers only `MonoExpPosIRF`; the residual + variant is parametrized over all 7 kernels but compare mode is not. + (Both now share the `_IRF_KERNEL_MODELS` list.) +- Extra (2026-07-10): `test_constant_profiled_op_folds_into_cache` — a + fully-fixed profiled op (compile-time constant-op branch in `schedule_2d`) + had no fixture; only the type checkers caught a stale call there during + the check-2 batch. + +### Check 19 — Two-layer design compliance + +- [x] `src/trspecfit/spectra.py:244` — `fit_project_mcp` is fully + interpreter/MCP with per-call dict distribution; belongs on a + setup/compile path or a lowered multi-file evaluator. + (Declined as an MCP change 2026-07-10 — see the check-2 note; the + lowered multi-file evaluator is the future fix.) +- [x] `src/trspecfit/fitlib.py:233` — string-based `getattr` dispatch per + residual call on the bridge (same fix as the check 2 item). + (Closed 2026-07-10 by the check-2 measurement above: <0.3% of a residual + call; hoisting would ripple through the `const` contract of all five fit + entry points for no observable gain.) +- [x] `src/trspecfit/functions/energy.py:144-152` — `LinBack` validation and + error formatting in a numeric body called from the residual loop. + (Closed 2026-07-10, guard kept by design: xStart/xStop can be fit + parameters or expressions, so the optimizer can violate the ordering + mid-fit and setup-time validation cannot replace the runtime check. The + happy path costs one `np.any` comparison; formatting only runs on the + failure branch. No clamping — the error now nudges users to set min/max + bounds on xStart/xStop instead.) +- Hot evaluators themselves are clean: no `isinstance`/model-structure + branching in `eval_1d.py`/`eval_2d.py` inner loops (PASS). + +### Chunk F — test-pattern compliance (CLAUDE.md) + +- [x] `tests/test_gir_integration.py:1087` — uses internal + `File._build_1d_dispatch_args`; could mask validation bugs on the public + dispatch path. + (Closed 2026-07-10 as within the CLAUDE.md invariant-check exception: + the test pins the private dispatch-args contract used by three + trspecfit.py call sites, while `test_gir_baseline_writes_back` covers + the same machinery through the public path. Docstring now says so.) +- [x] `tests/test_plotting.py:671` (also `:713`, `:726`, `:740`, `:753`) — + `component.plot()` without `save_img=-2`/`show_plot=False`; mitigated by + Agg backend + `plt.close("all")` but not per convention. + (Closed 2026-07-10 as by-design: these tests assert on the live axes, so + `save_img=-2` would close the figure before the assertions; Agg + + `plt.close("all")` is the suppression. CLAUDE.md's Plots rule now carves + out figure-inspection tests. The `# type guard` INFO item was swept + suite-wide in the same pass.) + +## INFO findings (no action required; note during triage) + +- Check 1: `spectra.py:132` GIR fast path leaves `model.lmfit_pars` stale + mid-fit; `fitlib.py:840` docstring promises `emcee_fin=[]` but returns + `None`; `fit_io.py:807` `observed`/`fit` rely on callers copying; + `fit_io.py:1601` reader trusts on-disk sha256; `fit_io.py:345`/`:336` + stale docstrings (schema version "1" vs "2"; nonexistent builder name); + `mcp.py:1728` `Component.value` reassigns `self.time` for conv kernels. +- Check 4: `File` has 36 public methods; `fit_slice_by_slice` 234 lines; + `fit_wrapper` ~207 lines; `schedule_2d` ~1020 / `schedule_1d` ~616 lines + (justified compiler monoliths); `fit_io.py` ~2200-line module. All judged + justified by the respective reviewers. +- Check 5: `trspecfit.py:4132` single commented-out `dpi_plot` line ("NOT + AVAILABLE YET"). +- Check 7: 12 public items in `trspecfit.py` and 2 in `mcp.py` lack + Parameters/Returns sections (list in chunk A transcript); worst offenders + `load_fits:559`, `save_fit:2907`, `export_fit:2935`, `compare_models:4340`. +- Check 8: `_append_baseline_slot`/`_append_spectrum_slot` share ~30-line + extraction block (`trspecfit.py:3419`); GIR dispatch duplicated between + 1D helper and `File.fit_2d:3993`; `par_create`/`par_construct` share + branching (`utils/lmfit.py:113`); `_as_group`/`_as_dataset` duplicate + `hdf5.py` helpers (`fit_io.py:145`); `_attr_str`/`_to_str_value` + near-identical (`fit_io.py:1434`). +- Check 9: `utils/plot.py:670` `np.arange(0, n, 1)` → `np.arange(n)`. +- Check 13: `graph_ir.py:3171-3285` runtime helpers called from `eval_1d` + (layer blur); lazy `eval_2d` import inside `schedule_2d` (`:2876`); + CSV export invokes `fitlib` plotting (`fit_io.py:1953`); `fitlib.py:34` + imports `IPython.display`. +- Check 14: missing `__repr__` on `GraphIR`, `ScheduledPlan1D/2D`, + `GraphNode`, `ExprProgram`, `SavedFitSlot`, `SavedFile`, `SavedProject`, + `PlotConfig`, `par_dummy`. +- Check 15: `utils/sbs.py:155` worker-process globals — documented, + standard `ProcessPoolExecutor` initializer pattern; acceptable. +- Check 17 (test gaps): no single-point time axis test anywhere (directly + covers the `time[1]-time[0]` WARN above) — covered since 2026-07-09 + (`test_graph_ir.py`, `test_mcp_eval.py`); no single-element energy/time + arrays through the fit pipeline; no NaN/Inf-in-data tests at the public + `File`/`Project` fit level — both covered since 2026-07-09 + (`test_fit_validation.py`; non-finite fit-window data now raises a + clear error from `fit_wrapper`). +- Check 20: `NodeKind.SUM` has no explicit gate (always structural); + `DomainKind.TIME_1D` rejected by domain check rather than node-kind set. + Both implicit but correct; all other kinds handled explicitly. +- Chunk F: several `assert model is not None` without `# type guard` + comment (`test_gir_integration.py:57` etc.); many `test_plotting.py` + calls use `save_img=0` instead of `-2`. + +## Suggested triage order + +1. The four FAILs (small, local fixes; `sign_change` and `my_conv` are + pure-math and easy to test). +2. Check 1 WARNs that affect correctness of public behavior: + `mcp.py:992` t_ind indexing, axis-mutating describe/setup methods, + single-point time axis (fix together with the check 17 test gap). +3. Silent-mode violations (checks 3 and 12): config-load swallow, + MCMC prints/plots, `fit_2d` display calls — one themed pass. +4. Hot-path performance batch (check 2, eval_2d + residual_fun + my_conv). +5. Parity-coverage batch (check 18): extend the roundtrip/parametrized + fixtures; cheap wins are parametrizing `test_compare_mode_irf` and adding + a pGauss profile fixture. +6. Structural/duplication items (checks 8, 13) as deliberate refactors, + one per PR. From a0a6db61c8afea05a7746822e6382e1a2932163a Mon Sep 17 00:00:00 2001 From: Johannes Mahl Date: Fri, 10 Jul 2026 21:01:29 -0700 Subject: [PATCH 36/36] add 0.10.2 changelog entry --- CHANGELOG.md | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 47792b7..ce136bd 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,29 @@ and this project adheres to [Semantic Versioning](https://semver.org/). This file is maintained using the shared changelog workflow in [`docs/ai/changelog.md`](docs/ai/changelog.md). +## [0.10.2] - 2026-07-10 + +### Added + +- **Diverging colormap for residual maps**: new `z_colormap_res` plot setting (default `'RdBu_r'`). 2D data+fit+residual panels now render the residual with its own zero-centered diverging colormap and symmetric color limits, instead of reusing the data colormap. +- **New plot styling settings**: `refline_color` / `refline_style` for `vlines`/`hlines` reference lines (unified default: grey dotted — previously inconsistent between 1D and 2D plots) and `panel_size` for multi-panel grid plots; `plot_2d_grid` gained a `columns` argument. +- **Every `PlotConfig` field is now settable via `project.yaml`**, with tuple-valued fields (`x_lim`, `panel_size`, ...) accepting YAML lists; a coverage test guards future fields. +- `results_to_fit_2d` accepts `parameter_names` to select and order DataFrame columns, so extra non-parameter columns (e.g. from `results_to_df` output) cannot silently corrupt the reconstructed 2D fit. + +### Changed + +- **Silent mode is honored throughout**: with `show_output=0`, MCMC no longer prints its banner/progress bar or leaves walker/corner figures open, and `fit_2d`/`fit_slice_by_slice` no longer display timing and parameter tables. A broken `project.yaml` now raises `ValueError` instead of silently falling back to defaults. +- **Clear errors instead of silent misbehavior**: fits reject NaN/Inf inside the fit window with a message naming the count and pointing to `set_fit_limits()`; `describe`/`define_baseline`/`set_fit_limits` raise when axes are missing instead of fabricating index axes; convolution on a single-point time axis, zero-sum convolution kernels, unknown `Simulator` noise types (constructor and setter), unknown sweep distribution types, empty parameter sweeps, and dynamics-model expressions referencing parameters outside their own dynamics model all raise immediately with actionable messages; a `t_vary` parameter without a dynamics model raises instead of returning `-1.0`; `LinBack` ordering errors suggest setting bounds on `xStart`/`xStop`. +- **Performance**: the compiled 2D evaluator avoids per-instruction array allocation and copies (preallocated profile buffers, hoisted profiled-op parameter sources), `my_conv` no longer builds a discarded padded axis, MCMC walker/corner figures are only constructed when shown or saved, and parameter sweeps write to one open HDF5 file instead of reopening it per configuration. + +### Fixed + +- **Convolution kernel support no longer truncates during fits**: the IRF kernel time axis was built once from the initial width parameter and never rebuilt, silently truncating the kernel and biasing the fitted width once it grew past its initial value. Both evaluation paths now rebuild the kernel support from current parameter values on every evaluation. +- **Partial-range 2D evaluation used wrong time indices**: `Model.create_value_2d(t_ind=[start, stop])` computed dynamics for `t[0:stop-start]` instead of `t[start:stop]`. +- `Simulator.plot_comparison` on a fresh simulator now auto-simulates with current settings instead of failing. +- `sign_change` no longer hangs on all-zero input. +- `plot_1d` with `y_norm=1` no longer produces all-NaN plots for constant traces. + ## [0.10.0] - 2026-07-08 ### Added