diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 5c7966cc..48e8fbcd 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -149,7 +149,7 @@ jobs: assert float((tensor * tensor).sum().item()) == 1240.0 print("CuPy and PyTorch CUDA kernels and DLPack passed") PY - - name: Run xFit GPU parity tests + - name: Run xFit and xRay GPU parity tests run: >- uv run --no-sync pytest -q -rs --junitxml="${RUNNER_TEMP}/xfit-gpu.xml" @@ -157,6 +157,7 @@ jobs: tests/xfit/test_models.py::test_cupy_finite_difference_theta_fix_matches_numpy tests/xfit/test_models.py::test_fit_dipoles_device_matches_portable_result_without_bulk_d2h tests/xfit/test_solver.py::test_cupy_tall_jacobian_diagnostics_match_numpy + tests/xray/test_mode_refinement_batched.py::test_batched_refinement_on_cupy_matches_numpy - name: Require all selected tests to execute run: | uv run --no-sync python - <<'PY' @@ -166,11 +167,11 @@ jobs: report = ET.parse(Path(os.environ["RUNNER_TEMP"]) / "xfit-gpu.xml") cases = report.findall(".//testcase") - assert len(cases) == 6, len(cases) + assert len(cases) == 9, len(cases) assert not report.findall(".//skipped"), "GPU tests must not skip" assert not report.findall(".//failure") assert not report.findall(".//error") - print("All 6 selected tests passed without skips") + print("All 9 selected tests passed without skips") PY - name: Run native xFit GPU tests run: >- diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 1e597343..2d338a88 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -132,8 +132,9 @@ checks in separate fresh containers. Image pulls have bounded retries; installation and test failures fail the job without retrying. The GPU job uses one L40G with Python 3.12 and the locked CUDA 13 dependencies. -It checks CuPy/PyTorch execution, runs six xFit GPU parity cases, and runs -all five synthetic quickstarts with `--require-gpu`. Missing CUDA support or +It checks CuPy/PyTorch execution, runs six xFit GPU parity cases and three +xRay refinement cases with different time units, and runs all five synthetic +quickstarts with `--require-gpu`. Missing CUDA support or skipped parity cases fail the job. JUnit results and quickstart summaries are uploaded with the tested commit SHA. GPU jobs run one at a time; a new revision cancels the previous workflow for the same branch. diff --git a/docs/xray/LINEAR-PREDICTION-REFINEMENT.md b/docs/xray/LINEAR-PREDICTION-REFINEMENT.md new file mode 100644 index 00000000..094d373e --- /dev/null +++ b/docs/xray/LINEAR-PREDICTION-REFINEMENT.md @@ -0,0 +1,123 @@ +# Nonlinear refinement of linear-prediction modes + +These experimental, opt-in tools refine linear-prediction modes by nonlinear +least squares. `cuphoton.xray.mode_refinement` fits one trace with SciPy; +`cuphoton.xray.mode_refinement_batched` fits a batch on NumPy or CuPy. Both use +an analytic Jacobian for the model + +``` +trace(t) = constant + sum_k amplitude_k exp(-decay_k t) cos(angular_frequency_k t + phase_k) +``` + +The fits are unconstrained: decay can become negative, and frequencies have +no bounds. They are separate from the existing production +`cuphoton.xray.iterative_fit` solver, which uses positive decay and bounded +frequencies. These tools have no detector or distributed pipeline integration. + +## Per-trace refinement + +`linear_prediction_refined(time, trace, n_components, n_modes)` runs +`linear_prediction_numpy`, keeps the `n_modes` strongest oscillating modes as +the starting point, and refines them and the constant with +`scipy.optimize.least_squares` (Levenberg-Marquardt, `xtol` and `ftol` +1e-12). When the linear-prediction step returned fewer than `n_modes` +oscillating modes, the missing ones are seeded one at a time from the +strongest peak of the residual spectrum (amplitude from the peak height, zero +decay and phase) and refined before the next seed. `refine_modes` is the +same solver from caller-supplied starting modes. + +`RefinedModes` carries the refined parameters ordered by amplitude, local +one-sigma uncertainty estimates, the reconstruction, the residual rms, the +solver status, and how many modes came from linear prediction +(`initial_mode_count`) versus the residual spectrum (`seeded_mode_count`). +The uncertainty calculation uses an SVD of the column-normalized Jacobian +and the estimated residual variance. A rank-deficient Jacobian produces +`NaN` uncertainties. Finite estimates describe the local fit under the model +and independent, equal-variance noise assumptions; they do not establish +identifiability away from that solution or account for model mismatch. + +A large residual can result from model mismatch, a poor local solution, or +failure to converge. Check the solver status and the reconstruction before +interpreting fitted parameters or uncertainties. + +`linear-prediction-validate --refine` runs the validation sweep a second time +with this estimator (`estimator: cpu-refined` in the summary). Compare mode +loss, bias, scatter, and residuals across noise levels and distortions. +The summary records solver convergence separately from mode recovery: +matching the true frequencies does not establish that optimization converged. +The Python entry point is +`synthetic_validation.validate_linear_prediction(refine=True)`, which returns +a report with the same text and JSON data as the command. + +Results depend on the fixture, initialization, and noise realization; +refinement does not guarantee recovery of every mode or attainment of the +Cramer-Rao bound. The comparison supplies the fixture's known mode count to +refinement; it does not test automatic model-order selection or establish +recovery on detector data with an unknown number of modes. + +```bash +uv run cuphoton xray lpv --refine --snr-db 40,20,10,5 --trials 200 --output-dir /tmp/lpv-refined +uv run cuphoton xray lpv --refine --distortion chirp:0.05 --snr-db 30 +uv run --extra viz cuphoton xray lpmr --output /tmp/review.html +``` + +`linear-prediction-modes-review` (`lpmr`) writes an offline Bokeh page with +the noisy trace, noise-free input, linear-prediction and refined +reconstructions, and individual modes. The page shows refinement convergence; +with a distortion, the individual true modes are labeled as reference modes. +Use `--snr-db`, `--distortion`, `--seed`, and `--samples` to configure the +comparison, or `--json` to report its artifact path and convergence count. +The reusable Python API is `mode_refinement_viz.write_modes_review(...)`. +The existing example script forwards its options to this command. + +## Batched refinement + +`refine_modes_batched(xp, time, traces, theta0, n_modes)` takes traces shaped +`(B, N)` and starting parameters `(B, 4K+1)` in the same layout as the +per-trace solver (amplitude, decay, angular frequency, phase per mode, then +the constant). It normalizes the Jacobian columns and solves the damped +normal equations as a batch of small `(4K+1)` systems with `xp.linalg.solve`. +A step that lowers the cost is accepted and its damping divided by 3; +otherwise the damping is multiplied by 5 and that trace keeps its parameters. + +Convergence requires the largest absolute component of the +column-normalized Jacobian's residual gradient, divided by +`max(residual norm, 1)`, to be at most `tol` (default 1e-9). This criterion +is independent of the choice of time units. A tiny damped step alone does +not establish convergence. Iteration stops when every trace has converged +or after `max_iter` (default 60). Check each trace's convergence flag; +stationarity does not guarantee a good fit or agreement with another solver. + +At tight tolerances, rounding can prevent further cost reduction before the +gradient criterion is met. Such traces remain marked unconverged even when +their parameters agree closely with another fit. Choose a tolerance suited +to the required accuracy and inspect residuals as well as convergence flags. + +## Benchmark + +`linear-prediction-refine-benchmark` (`lprb`) builds `--traces` noisy copies +of the fixture (two modes, 96 samples, noise sigma 0.02), perturbs the true +parameters per trace to make identical starting points, and times three +paths: the SciPy solver in a Python loop over traces, the batched solver on +NumPy, and the batched solver on CuPy. Each path is timed `--repeat` times +(default 3) and the best is reported. For the batched paths the clock covers +the solver only: the arrays are on the device before the clock starts and +the device is synchronised before and after, so host-to-device transfer and +the copy of the result back are not included. One untimed CuPy call on up to +16 traces runs first to warm the device path. The command also +reports the largest parameter difference between each batched result and +the SciPy result, the number of converged traces, and the largest residual +RMS for each path. Diagnostics belong to that path's fastest repetition. +Parameter agreement alone does not establish convergence. + +```bash +uv run cuphoton xray lprb --traces 256 --repeat 3 --json +uv run cuphoton xray lprb --traces 2048 --repeat 3 --json +uv run cuphoton xray lprb --traces 8192 --repeat 3 --json +``` + +Use `--no-gpu` for the CPU comparison alone. These timings cover refinement +from supplied starts and exclude linear-prediction initialization, input I/O, +and detector processing. They cannot establish an end-to-end detector +speedup. Performance and parameter agreement need measurement on the target +hardware and representative traces. diff --git a/docs/xray/LINEAR-PREDICTION-VALIDATION.md b/docs/xray/LINEAR-PREDICTION-VALIDATION.md new file mode 100644 index 00000000..a3a8e369 --- /dev/null +++ b/docs/xray/LINEAR-PREDICTION-VALIDATION.md @@ -0,0 +1,162 @@ +# Linear prediction validation on synthetic damped modes + +`linear-prediction-validate` (`lpv`) checks the linear-prediction fit against +a known answer. It builds traces from a small set of damped sinusoids, adds +white Gaussian noise at chosen signal-to-noise levels, runs the fit many +times, and compares the scatter and bias of every recovered parameter with +the Cramer-Rao lower bound for the same model. It also counts how often a +true mode is missing from the fit. No external data is needed. + +```bash +uv run cuphoton xray lpv --snr-db 40,30,20,10 --trials 200 --output-dir /tmp/lpv +uv run cuphoton xray lpv --distortion chirp:0.05 --snr-db 30 --output-dir /tmp/lpv-chirp +uv run cuphoton xray lpv --json +``` + +Keep `--output-dir` outside the checkout. It receives `summary.json` and, when +the `viz` extra is installed, `validation.html`. The HTML includes its Bokeh +resources and can be viewed offline. + +## Model and fixture + +The trace is + +``` +trace(t) = constant + sum_k amplitude_k exp(-decay_k t) cos(angular_frequency_k t + phase_k) +``` + +sampled at `samples` points from `t = 0` to `duration`. Time is in the input +time unit; angular frequency is in radians per time unit and decay in inverse +time units. The default fixture has two modes, `(1.25, 0.09, 2.4, 0.30)` and +`(0.45, 0.03, 0.9, -0.75)` as `(amplitude, decay, angular_frequency, phase)`, +constant `0.15`, 96 samples over 9.5 time units. With zero noise it is the +same trace as `synthetic_trace`, and the fit recovers the modes to better than +`1e-6`. `cuphoton.xray.synthetic_validation.synthetic_modes_trace` builds the +fixture; `DampedMode` describes a mode. + +Noise is white Gaussian, independent per sample. A level's `snr_db` is +`20 log10(rms of the mean-removed clean trace / sigma)`. One `numpy` generator +seeded by `--seed` draws every trial in order, so a run is reproducible from +the summary's `config`. +SNR levels must be finite and produce finite, positive noise scales; +`--components` and `--trials` must be positive integers. + +## Bounds + +`cramer_rao_bounds` uses the analytic Jacobian of the model at the true +parameters (all parameters free, including the constant). A column-scaled +singular-value decomposition computes the bounds associated with the Fisher +information `J^T J / sigma^2` without forming or inverting that matrix. +Numerically rank-deficient models, including modes with identical frequency +and decay but different phases, raise `ValueError`: the individual parameters +cannot be identified. These are standard-deviation bounds; the summary +also carries the squared value as `crlb_variance`. For the single undamped +sinusoid test fixture, the angular-frequency bound matches the square root +of the closed-form variance `24 sigma^2 / (A^2 dt^2 N (N^2 - 1))` +(Kay, Estimation Theory, 1993) to within 3 percent. The calculation also frees +the decay and constant. The legacy `step` argument remains accepted for +compatibility; analytic derivatives do not use a finite-difference step. + +## Mode matching and statistics + +A true mode is recovered in a trial when a fitted mode has an angular +frequency within `match_tolerance` (default 0.3 rad per time unit) of it; +the assignment first maximizes the number of recovered modes, then minimizes +their total frequency error, using total decay error to break ties. +Each fitted mode is used at most once. The +per-mode `loss_rate` is the fraction of trials in which the mode was not +recovered. A trial is successful when the estimator returned and every true +mode was recovered; a trial in which the estimator raised counts in +`estimator_errors` and as a loss of every mode. The text report includes +the error count so estimator failures can be distinguished from mode loss. + +For each mode and parameter (`amplitude`, `decay`, `angular_frequency`, +`phase`) the summary gives `bias`, `std` (ddof 1) and `rmse` computed over +the recovered trials only, the bound (`crlb_std`, `crlb_variance`) and +`std_over_crlb_std`. Truth and fitted modes use nonnegative amplitudes and +frequencies. A negative frequency is reflected along with its phase; a +negative amplitude is folded into the phase. Both phases and phase errors +are wrapped to `[-pi, pi)`; the summary records truth in this convention. +Undefined statistics are written as JSON `null`: standard deviation requires +at least two recovered trials, and bias and RMSE require at least one. + +Read the decay statistics of a lightly damped mode together with its +`loss_rate`. The root filter keeps decaying roots only, so when noise pushes +the fitted decay of such a mode through zero the mode is dropped rather than +reported with a negative decay; the surviving decay estimates are truncated +at zero, and their scatter and bias are conditional on survival. The summary +labels these statistics `statistics_over: recovered_trials`. Conditional +statistics and finite-sample scatter can fall below the unconditional bound. + +Every level also reports `residual_rms_over_sigma_median`, the median over +trials of the rms residual of the reconstruction divided by sigma. It is +near 1 for an accurate reconstruction at the specified noise level. A large +value can indicate model mismatch, missed modes, or estimator error; it does +not by itself distinguish these causes. + +## Model mismatch + +`--distortion kind:amount` replaces the clean trace with one outside the +model class: + +| kind | trace | +| --- | --- | +| `chirp` | every frequency rises by the fraction `amount` across the record | +| `gaussian_envelope` | the first mode decays as a Gaussian with 1/e time `amount` | +| `baseline_drift` | a linear ramp of `amount` across the record | +| `clip` | the trace is clipped above `amount` | +| `glitch` | one sample at mid-record is offset by `amount` | + +Amounts must be finite. The Gaussian 1/e time must be positive and is measured +from the first sample. + +The truth used for bias and bounds stays the undistorted mode set, so the +statistics measure the estimator's response to the mismatch. A chirp or a +clipped trace can return modes with a confident, biased frequency and no +loss; the residual ratio supplies an additional reconstruction diagnostic. + +## Summary schema + +`summary.json` has `schema_version` 1 and five top-level keys. + +- `command`: `linear-prediction-validate`. +- `config`: the model string and units, the true modes and constant, the + sampling (`samples`, `duration`, `dt`, `start`), the noise definition and + levels, `seed`, `trials_per_level`, `n_components` (model order), the + mode-matching criterion and tolerance, how the statistics are computed, and + the estimators and distortions covered. +- `runtime`: `cuphoton.core.runtime.runtime_metadata` for the CPU backend + (package, Python, platform, NumPy versions, backend, device, dtype) plus + `scipy_version`, `source_revision`, and `source_dirty`. A revision is + recorded only when this module is tracked by its Git checkout; + `source_dirty` reports tracked changes anywhere in that checkout. + Both are `null` for an installed or untracked module, or when Git metadata + is unavailable. Ignored and untracked run artifacts do not mark the source + dirty. +- `results`: one entry per sweep condition (estimator by distortion by + signal-to-noise level): `snr_db`, `signal_rms`, `noise_sigma`, `trials` with + `attempted`, `successful`, `failed` and `estimator_errors`, + `any_mode_lost_rate`, `residual_rms_over_sigma_median`, and `modes`, each + with `recovered_trials`, `loss_rate`, `statistics_over` and the four + parameter blocks described above. +- `artifacts`: filenames relative to the summary, `summary` and `figure` + (`null` when Bokeh is not installed). + +The figure shows, per estimator and per parameter, the sample standard +deviation against the bound versus signal-to-noise, and the loss rates. +The signal RMS in `config.noise` describes the first sweep; each result +records its own signal RMS because distortions can change it. + +## Python entry points + +`cuphoton.xray.synthetic_validation` exposes `synthetic_modes_trace`, +`distort_trace`, `cramer_rao_bounds`, `match_modes`, `validation_sweep`, +`build_summary`, `write_validation_figure` and `write_validation_run`. +`validation_sweep` accepts another estimator callable with the signature +`estimator(time, trace, n_components)` returning an object with +`angular_frequency`, `decay`, `amplitude`, `phase` and `reconstruction` +(NumPy or CuPy arrays), so the same sweep can compare estimators. +`build_summary` combines sweeps with shared truth, sampling, seed, matching, +SNR levels, trial counts and model order. It rejects incompatible settings +instead of describing later sweeps with the first sweep's configuration. +Write separate summaries for changes to those settings. diff --git a/docs/xray/README.md b/docs/xray/README.md index 62b4b167..815dd483 100644 --- a/docs/xray/README.md +++ b/docs/xray/README.md @@ -34,6 +34,7 @@ Choose a new output directory for each quickstart run. | Input inspection | `data-probe`, `roi-candidates` | | Trace work | `extract-trace`, `trace-smoke` | | Linear prediction | `linear-prediction-*`, `prediction-roots-benchmark`, `model-order-sweep` | +| Synthetic validation | `linear-prediction-validate` (`lpv`), `linear-prediction-refine-benchmark` (`lprb`), `linear-prediction-modes-review` (`lpmr`) | | Detector products | `detector-mask`, `detector-artifacts`, `detector-artifact-normalize`, `detector-artifact-compare` | | Distributed detector work | `detector-artifact-distributed`, `detector-artifact-merge` | | Review | `report`, `validation-viz`, `workflow-viz`, `phonon-viz` | @@ -295,3 +296,5 @@ See [Data and artifact contracts](../data-artifacts.md#xray-hdf5-and-trace-produ - [GPU-first behavior](GPU-FIRST.md) - [Distributed detector artifacts](DISTRIBUTED-DETECTOR-ARTIFACTS.md) - [Validation visualization](VALIDATION-VIZ.md) +- [Linear prediction validation](LINEAR-PREDICTION-VALIDATION.md) +- [Experimental nonlinear refinement of linear-prediction modes](LINEAR-PREDICTION-REFINEMENT.md) diff --git a/examples/xray_lp_modes_review.py b/examples/xray_lp_modes_review.py new file mode 100644 index 00000000..8598e2ce --- /dev/null +++ b/examples/xray_lp_modes_review.py @@ -0,0 +1,28 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Write a synthetic-mode review through the shared XRay command. + + uv run python examples/xray_lp_modes_review.py --output review.html + uv run python examples/xray_lp_modes_review.py --distortion chirp:0.05 + +Python callers can use +``cuphoton.xray.mode_refinement_viz.write_modes_review``. +Requires the ``viz`` extra. +""" + +from __future__ import annotations + +import sys + +from cuphoton.core.cli import run_component + + +def main(argv=None) -> int: + args = sys.argv[1:] if argv is None else argv + return run_component("xray", ["linear-prediction-modes-review", *args]) + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/src/cuphoton/xray/commands.py b/src/cuphoton/xray/commands.py index d1c6861e..4499d815 100644 --- a/src/cuphoton/xray/commands.py +++ b/src/cuphoton/xray/commands.py @@ -21,6 +21,7 @@ NonNegativeIntegerInvariant, PairInvariant, PathValueInvariant, + PositiveIntegerInvariant, SequenceInvariant, SetInvariant, StringInvariant, @@ -596,6 +597,196 @@ class JsonArg(BoolInvariant): _required = False +class LinearPredictionValidateCommand(_XRayCommand): + _description_ = ( + "Validate linear prediction on synthetic damped modes against " + "the Cramer-Rao bound." + ) + _shortname_ = "lpv" + _handler_name_ = "_linear_prediction_validate" + + samples = 96 + + class SamplesArg(IntegerInvariant): + _arg = "--samples" + _help = "Number of synthetic trace samples." + _required = False + _default = 96 + + components = 6 + + class ComponentsArg(PositiveIntegerInvariant): + _arg = "--components" + _help = "Number of SVD components to fit." + _required = False + _default = 6 + + trials = 200 + + class TrialsArg(PositiveIntegerInvariant): + _arg = "--trials" + _help = "Monte Carlo trials per signal-to-noise level." + _required = False + _default = 200 + + snr_db = "40,30,20,10" + + class SnrDbArg(StringInvariant): + _arg = "--snr-db" + _help = "Comma-separated signal-to-noise levels in dB." + _required = False + _default = "40,30,20,10" + + seed = 20260914 + + class SeedArg(IntegerInvariant): + _arg = "--seed" + _help = "Random seed for the noise draws." + _required = False + _default = 20260914 + + distortion = None + + class DistortionArg(StringInvariant): + _arg = "--distortion" + _help = ( + "Optional model mismatch as kind:amount, one of chirp, " + "gaussian_envelope, baseline_drift, clip, glitch " + "(for example chirp:0.05)." + ) + _required = False + _default = None + + refine = False + + class RefineArg(BoolInvariant): + _arg = "--refine" + _help = ( + "Also run the sweep with nonlinear least-squares refinement " + "of the linear-prediction modes." + ) + _required = False + + output_dir = None + + class OutputDirArg(PathValueInvariant): + _arg = "--output-dir" + _help = ( + "Directory for summary.json and the validation figure " + "(created if needed; keep it outside the checkout)." + ) + _required = False + + json = False + + class JsonArg(BoolInvariant): + _arg = "--json" + _help = "Print the summary as JSON instead of the text report." + _required = False + + +class LinearPredictionModesReviewCommand(_XRayCommand): + _description_ = ( + "Write a visual comparison of synthetic modes and refinements." + ) + _shortname_ = "lpmr" + _handler_name_ = "_linear_prediction_modes_review" + + output_path = "xray_lp_modes_review.html" + + class OutputPathArg(PathValueInvariant): + _arg = "--output" + _help = "Output HTML path (requires the viz extra)." + _required = False + _default = "xray_lp_modes_review.html" + + snr_db = "40,20,10,5" + + class SnrDbArg(StringInvariant): + _arg = "--snr-db" + _help = "Comma-separated signal-to-noise levels in dB." + _required = False + _default = "40,20,10,5" + + distortion = None + + class DistortionArg(StringInvariant): + _arg = "--distortion" + _help = "Model mismatch as kind:amount, for example chirp:0.05." + _required = False + _default = None + + seed = 11 + + class SeedArg(IntegerInvariant): + _arg = "--seed" + _help = "Random seed for reproducible noise." + _required = False + _default = 11 + + samples = 96 + + class SamplesArg(PositiveIntegerInvariant): + _arg = "--samples" + _help = "Number of synthetic trace samples." + _required = False + _default = 96 + + json = False + + class JsonArg(BoolInvariant): + _arg = "--json" + _help = "Emit machine-readable artifact information." + _required = False + + +class LinearPredictionRefineBenchmarkCommand(_XRayCommand): + _description_ = ( + "Benchmark per-trace SciPy refinement against the batched " + "Levenberg-Marquardt solver on NumPy and CuPy." + ) + _shortname_ = "lprb" + _handler_name_ = "_linear_prediction_refine_benchmark" + + samples = 96 + + class SamplesArg(PositiveIntegerInvariant): + _arg = "--samples" + _help = "Number of synthetic trace samples." + _required = False + _default = 96 + + traces = 256 + + class TracesArg(PositiveIntegerInvariant): + _arg = "--traces" + _help = "Number of traces in the batch." + _required = False + _default = 256 + + repeat = 3 + + class RepeatArg(PositiveIntegerInvariant): + _arg = "--repeat" + _help = "Timed repetitions per path; the best time is reported." + _required = False + _default = 3 + + no_gpu = False + + class NoGpuArg(BoolInvariant): + _arg = "--no-gpu" + _help = "Skip the CuPy comparison." + _required = False + + json = False + + class JsonArg(BoolInvariant): + _arg = "--json" + _help = "Emit machine-readable JSON." + _required = False + + class LinearPredictionBenchmarkCommand(_XRayCommand): _description_ = "Benchmark serial and batched linear-prediction P1." _shortname_ = "lpb" @@ -3686,10 +3877,9 @@ def _data_probe(args): result = probe_hdf5_pair(h5dir=args.h5dir, fon=args.fon, foff=args.foff) payload = result.to_dict() if args.json: - print(json.dumps(payload, indent=2, sort_keys=True)) + args._out(json.dumps(payload, indent=2, sort_keys=True)) else: - _print_file_probe("on", result.on) - _print_file_probe("off", result.off) + args._out(result) return 0 @@ -4241,25 +4431,6 @@ def _detector_mask(args): return 0 -def _print_file_probe(label, probe): - print(f"{label}_path={probe.path}") - print(f"{label}_size_bytes={probe.size_bytes}") - print(f"{label}_schema={probe.schema}") - print(f"{label}_ipm_pairs={','.join(probe.ipm_pairs) or '-'}") - print(f"{label}_keys={','.join(probe.keys)}") - for dataset in probe.datasets: - shape = "x".join(str(dim) for dim in dataset.shape) - chunks = ( - "-" - if dataset.chunks is None - else "x".join(str(dim) for dim in dataset.chunks) - ) - print( - f"{label}_dataset={dataset.name} " - f"shape={shape} dtype={dataset.dtype} chunks={chunks}" - ) - - def _linear_prediction_smoke(args): from .linear_prediction import ( LinearPredictionComparison, @@ -4342,6 +4513,88 @@ def _linear_prediction_smoke(args): return 0 +def _parse_validation_distortion(value): + distortion = None + if value: + try: + kind, amount = str(value).split(":", maxsplit=1) + if not kind.strip(): + raise ValueError("missing distortion kind") + distortion = (kind.strip(), float(amount)) + except ValueError as exc: + raise ValueError( + "--distortion must be kind:amount with a numeric amount " + "(for example chirp:0.05)" + ) from exc + return distortion + + +def _linear_prediction_validate(args): + from .synthetic_validation import validate_linear_prediction + + snr = tuple(float(x) for x in str(args.snr_db).split(",") if x.strip()) + distortion = _parse_validation_distortion(args.distortion) + report = validate_linear_prediction( + samples=args.samples, + snr_db=snr, + trials=args.trials, + n_components=args.components, + seed=args.seed, + distortion=distortion, + output_dir=args.output_dir, + refine=args.refine, + ) + args._out( + json.dumps( + report.to_dict(), indent=2, sort_keys=True, allow_nan=False + ) + if args.json + else report + ) + return 0 + + +def _linear_prediction_modes_review(args): + from .mode_refinement_viz import write_modes_review + + report = write_modes_review( + args.output_path, + snr_db=tuple( + float(x) for x in str(args.snr_db).split(",") if x.strip() + ), + distortion=_parse_validation_distortion(args.distortion), + seed=args.seed, + samples=args.samples, + ) + args._out( + json.dumps( + report.to_dict(), indent=2, sort_keys=True, allow_nan=False + ) + if args.json + else report + ) + return 0 + + +def _linear_prediction_refine_benchmark(args): + from .mode_refinement_batched import benchmark_refinement + + result = benchmark_refinement( + samples=args.samples, + traces=args.traces, + run_gpu=not args.no_gpu, + repeat=args.repeat, + ) + args._out( + json.dumps( + result.to_dict(), indent=2, sort_keys=True, allow_nan=False + ) + if args.json + else result + ) + return 0 + + def _linear_prediction_benchmark(args): from .linear_prediction import benchmark_linear_prediction_p1_batch @@ -5325,188 +5578,54 @@ def _prediction_roots_benchmark(args): def _model_order_sweep(args): - from .linear_prediction import model_order_sweep + from .linear_prediction import model_order_sweep, model_order_sweep_batch components = _component_sweep_values(args) trace_paths = _trace_npz_batch_paths_from_args(args, required=False) + source = None if trace_paths: _reject_trace_npz_hdf5_options(args) - if len(trace_paths) == 1: - time, trace, source = _load_trace_npz(trace_paths[0]) - payload = _single_model_order_sweep_payload( - model_order_sweep( - time, - trace, - components, - roots_backend=args.roots_backend, - relative_tolerance=args.relative_tolerance, - ), - source=source, - ) - else: - time, trace_rows, sources = _load_trace_npz_batch(trace_paths) - payload = _model_order_sweep_batch_payload( - time, - trace_rows, - sources, - components=components, - roots_backend=args.roots_backend, - relative_tolerance=args.relative_tolerance, - ) - else: - time, trace, source = _trace_for_workbench(args) - payload = _single_model_order_sweep_payload( - model_order_sweep( - time, - trace, - components, - roots_backend=args.roots_backend, - relative_tolerance=args.relative_tolerance, - ), - source=source, + if len(trace_paths) > 1: + time, trace_rows, sources = _load_trace_npz_batch(trace_paths) + result = model_order_sweep_batch( + time, + trace_rows, + sources, + components=components, + roots_backend=args.roots_backend, + relative_tolerance=args.relative_tolerance, ) - - if args.json: - print(json.dumps(payload, indent=2, sort_keys=True)) - elif payload["source"]["kind"] == "trace-npz-batch": - _print_model_order_sweep_batch(payload) else: - _print_model_order_sweep_payload(payload) - return 0 - - -def _single_model_order_sweep_payload(result, *, source): - payload = dataclass_asdict(result) - payload["source"] = source - return payload - - -def _model_order_sweep_batch_payload( - time, - trace_rows, - sources, - *, - components, - roots_backend: str, - relative_tolerance: float, -): - from .linear_prediction import model_order_sweep - - trace_payloads = [] - best_components = [] - best_selected_model_orders = [] - best_rms_residuals = [] - best_reconstruction_errors = [] - for trace_index, (trace, source) in enumerate( - zip(trace_rows, sources, strict=True) - ): + time, trace, source = ( + _load_trace_npz(trace_paths[0]) + if trace_paths + else _trace_for_workbench(args) + ) result = model_order_sweep( time, trace, components, - roots_backend=roots_backend, - relative_tolerance=relative_tolerance, - ) - trace_payload = _single_model_order_sweep_payload( - result, - source=source, - ) - trace_payload["trace_index"] = trace_index - trace_payloads.append(trace_payload) - best_components.append(int(result.best_components)) - best_selected_model_orders.append( - int(result.best_selected_model_order) - ) - best_rms_residuals.append(float(result.best_rms_residual)) - best_reconstruction_errors.append( - float(result.best_reconstruction_rms_error) + roots_backend=args.roots_backend, + relative_tolerance=args.relative_tolerance, ) - return { - "source": { - "kind": "trace-npz-batch", - "trace_count": len(sources), - "traces": sources, - }, - "trace_count": len(sources), - "samples": int(time.shape[0]), - "roots_backend": roots_backend, - "relative_tolerance": float(relative_tolerance), - "component_counts": tuple(int(item) for item in components), - "best_components_by_trace": tuple(best_components), - "best_components_unique": tuple(sorted(set(best_components))), - "best_selected_model_orders_by_trace": tuple( - best_selected_model_orders - ), - "best_selected_model_orders_unique": tuple( - sorted(set(best_selected_model_orders)) - ), - "best_rms_residual_min": min(best_rms_residuals), - "best_rms_residual_max": max(best_rms_residuals), - "best_reconstruction_rms_error_min": min(best_reconstruction_errors), - "best_reconstruction_rms_error_max": max(best_reconstruction_errors), - "traces": tuple(trace_payloads), - } - - -def _print_model_order_sweep_payload(payload): - print(f"source={payload['source']['kind']}") - print(f"samples={payload['samples']}") - print(f"roots_backend={payload['roots_backend']}") - print(f"best_components={payload['best_components']}") - print(f"best_selected_model_order={payload['best_selected_model_order']}") - print(f"best_rms_residual={payload['best_rms_residual']:.6g}") - print( - "best_reconstruction_rms_error=" - f"{payload['best_reconstruction_rms_error']:.6g}" - ) - for entry in payload["entries"]: - print( - f"components={entry['components']} " - f"selected_model_order={entry['selected_model_order']} " - f"rms_residual={entry['rms_residual']:.6g} " - "reconstruction_rms_error=" - f"{entry['reconstruction_rms_error']:.6g} " - f"chi2={entry['chi2']:.6g} " - f"selected_roots={entry['selected_root_count']} " - f"decaying_roots={entry['decaying_root_count']}" - ) - - -def _print_model_order_sweep_batch(payload): - print(f"source={payload['source']['kind']}") - print(f"trace_count={payload['trace_count']}") - print(f"samples={payload['samples']}") - print(f"roots_backend={payload['roots_backend']}") - component_counts = ",".join( - str(item) for item in payload["component_counts"] - ) - print(f"component_counts={component_counts}") - print( - "best_components_unique=" - + ",".join(str(item) for item in payload["best_components_unique"]) - ) - print( - "best_selected_model_orders_unique=" - + ",".join( - str(item) for item in payload["best_selected_model_orders_unique"] - ) - ) - for trace_payload in payload["traces"]: - print(f"trace_index={trace_payload['trace_index']}") - print(f" best_components={trace_payload['best_components']}") - print( - " best_selected_model_order=" - f"{trace_payload['best_selected_model_order']}" - ) - print( - " best_reconstruction_rms_error=" - f"{trace_payload['best_reconstruction_rms_error']:.6g}" - ) + if args.json: + payload = result.to_dict() + if source is not None: + payload["source"] = source + args._out(json.dumps(payload, indent=2, sort_keys=True)) + else: + if source is not None: + args._out(f"source={source['kind']}") + args._out(result) + return 0 def _subspace_benchmark(args): - from .subspace import compare_subspace_methods + from .subspace import ( + compare_subspace_methods, + compare_subspace_methods_batch, + ) trace_paths = _trace_npz_batch_paths_from_args(args, required=False) svd_backends = _subspace_svd_backends(args) @@ -5515,217 +5634,54 @@ def _subspace_benchmark(args): if args.method is None else tuple(args.method) ) + source = None if trace_paths: _reject_trace_npz_hdf5_options(args) - if len(trace_paths) == 1: - time, trace, source = _load_trace_npz(trace_paths[0]) - payload = _single_subspace_benchmark_payload( - compare_subspace_methods( - time, - trace, - model_order=args.model_order, - components=args.components, - methods=methods, - svd_backends=svd_backends, - pencil_rows=args.pencil_rows, - randomized_oversamples=args.randomized_oversamples, - power_iterations=args.power_iterations, - random_seed=args.random_seed, - ), - source=source, - svd_backends=svd_backends, - ) - else: - time, trace_rows, sources = _load_trace_npz_batch(trace_paths) - payload = _subspace_benchmark_batch_payload( - time, - trace_rows, - sources, - model_order=args.model_order, - components=args.components, - methods=methods, - svd_backends=svd_backends, - pencil_rows=args.pencil_rows, - randomized_oversamples=args.randomized_oversamples, - power_iterations=args.power_iterations, - random_seed=args.random_seed, - ) - else: - time, trace, source = _trace_for_workbench(args) - payload = _single_subspace_benchmark_payload( - compare_subspace_methods( - time, - trace, - model_order=args.model_order, - components=args.components, - methods=methods, - svd_backends=svd_backends, - pencil_rows=args.pencil_rows, - randomized_oversamples=args.randomized_oversamples, - power_iterations=args.power_iterations, - random_seed=args.random_seed, - ), - source=source, + if len(trace_paths) > 1: + time, trace_rows, sources = _load_trace_npz_batch(trace_paths) + result = compare_subspace_methods_batch( + time, + trace_rows, + sources, + model_order=args.model_order, + components=args.components, + methods=methods, svd_backends=svd_backends, + pencil_rows=args.pencil_rows, + randomized_oversamples=args.randomized_oversamples, + power_iterations=args.power_iterations, + random_seed=args.random_seed, ) - - if args.json: - print(json.dumps(payload, indent=2, sort_keys=True)) - elif payload["source"]["kind"] == "trace-npz-batch": - _print_subspace_benchmark_batch(payload) else: - print(f"source={payload['source']['kind']}") - print(f"samples={payload['samples']}") - print(f"model_order={payload['model_order']}") - print(f"components={payload['components']}") - print(f"baseline_chi2={payload['baseline_chi2']:.6g}") - print(f"baseline_rms_residual={payload['baseline_rms_residual']:.6g}") - for method in payload["methods"]: - print(f"method={method['method']}") - print(f" svd_backend={method['svd_backend']}") - print(f" svd_rank={method['svd_rank']}") - print(f" rms_residual={method['rms_residual']:.6g}") - print( - " max_abs_reconstruction_diff=" - f"{method['max_abs_reconstruction_diff']:.6g}" - ) - print(f" elapsed_s={method['elapsed_s']:.6g}") - return 0 - - -def _single_subspace_benchmark_payload(result, *, source, svd_backends): - payload = dataclass_asdict(result) - payload["source"] = source - payload["svd_backends"] = svd_backends - return payload - - -def _subspace_benchmark_batch_payload( - time, - trace_rows, - sources, - *, - model_order: int, - components: int, - methods, - svd_backends, - pencil_rows: int | None, - randomized_oversamples: int, - power_iterations: int, - random_seed: int, -): - from .subspace import compare_subspace_methods - - method_summary = { - (method, svd_backend): { - "method": method, - "svd_backend": svd_backend, - "trace_count": 0, - "rms_residual_min": float("inf"), - "rms_residual_max": 0.0, - "max_abs_reconstruction_diff_max": 0.0, - "elapsed_s_total": 0.0, - } - for method in methods - for svd_backend in svd_backends - } - trace_payloads = [] - baseline_chi2_values = [] - baseline_rms_values = [] - - for trace_index, (trace, source) in enumerate( - zip(trace_rows, sources, strict=True) - ): + time, trace, source = ( + _load_trace_npz(trace_paths[0]) + if trace_paths + else _trace_for_workbench(args) + ) result = compare_subspace_methods( time, trace, - model_order=model_order, - components=components, + model_order=args.model_order, + components=args.components, methods=methods, svd_backends=svd_backends, - pencil_rows=pencil_rows, - randomized_oversamples=randomized_oversamples, - power_iterations=power_iterations, - random_seed=random_seed, - ) - trace_payload = _single_subspace_benchmark_payload( - result, - source=source, - svd_backends=svd_backends, + pencil_rows=args.pencil_rows, + randomized_oversamples=args.randomized_oversamples, + power_iterations=args.power_iterations, + random_seed=args.random_seed, ) - trace_payload["trace_index"] = trace_index - trace_payloads.append(trace_payload) - baseline_chi2_values.append(float(result.baseline_chi2)) - baseline_rms_values.append(float(result.baseline_rms_residual)) - - for method_result in result.methods: - summary = method_summary[ - (method_result.method, method_result.svd_backend) - ] - summary["trace_count"] += 1 - summary["rms_residual_min"] = min( - float(summary["rms_residual_min"]), - float(method_result.rms_residual), - ) - summary["rms_residual_max"] = max( - float(summary["rms_residual_max"]), - float(method_result.rms_residual), - ) - summary["max_abs_reconstruction_diff_max"] = max( - float(summary["max_abs_reconstruction_diff_max"]), - float(method_result.max_abs_reconstruction_diff), - ) - summary["elapsed_s_total"] += float(method_result.elapsed_s) - - summaries = tuple( - method_summary[(method, svd_backend)] - for method in methods - for svd_backend in svd_backends - ) - return { - "source": { - "kind": "trace-npz-batch", - "trace_count": len(sources), - "traces": sources, - }, - "trace_count": len(sources), - "samples": int(time.shape[0]), - "model_order": int(model_order), - "components": int(components), - "svd_backends": svd_backends, - "baseline_chi2_min": min(baseline_chi2_values), - "baseline_chi2_max": max(baseline_chi2_values), - "baseline_rms_residual_min": min(baseline_rms_values), - "baseline_rms_residual_max": max(baseline_rms_values), - "method_summary": summaries, - "traces": tuple(trace_payloads), - } - -def _print_subspace_benchmark_batch(payload): - print(f"source={payload['source']['kind']}") - print(f"trace_count={payload['trace_count']}") - print(f"samples={payload['samples']}") - print(f"model_order={payload['model_order']}") - print(f"components={payload['components']}") - print( - "baseline_rms_residual_range=" - f"{payload['baseline_rms_residual_min']:.6g}:" - f"{payload['baseline_rms_residual_max']:.6g}" - ) - for summary in payload["method_summary"]: - print(f"method={summary['method']}") - print(f" svd_backend={summary['svd_backend']}") - print(f" trace_count={summary['trace_count']}") - print( - " rms_residual_range=" - f"{summary['rms_residual_min']:.6g}:" - f"{summary['rms_residual_max']:.6g}" - ) - print( - " max_abs_reconstruction_diff_max=" - f"{summary['max_abs_reconstruction_diff_max']:.6g}" - ) + if args.json: + payload = result.to_dict() + if source is not None: + payload["source"] = source + payload["svd_backends"] = svd_backends + args._out(json.dumps(payload, indent=2, sort_keys=True)) + else: + if source is not None: + args._out(f"source={source['kind']}") + args._out(result) + return 0 def _subspace_acceptance(args): diff --git a/src/cuphoton/xray/hdf5.py b/src/cuphoton/xray/hdf5.py index bbae346f..3e0df47e 100644 --- a/src/cuphoton/xray/hdf5.py +++ b/src/cuphoton/xray/hdf5.py @@ -39,6 +39,15 @@ class DatasetProbe: dtype: str chunks: tuple[int, ...] | None + def __str__(self) -> str: + shape = "x".join(str(dim) for dim in self.shape) + chunks = ( + "-" + if self.chunks is None + else "x".join(str(dim) for dim in self.chunks) + ) + return f"{self.name} shape={shape} dtype={self.dtype} chunks={chunks}" + @dataclass(frozen=True) class FileProbe: @@ -50,12 +59,38 @@ class FileProbe: ipm_pairs: tuple[str, ...] load_plan: Hdf5LoadPlan | None + def __str__(self) -> str: + return self.format_text() + + def format_text(self, *, label: str | None = None) -> str: + """Summarize the file, optionally prefixing fields with a label.""" + prefix = "" if label is None else f"{label}_" + lines = [ + f"{prefix}path={self.path}", + f"{prefix}size_bytes={self.size_bytes}", + f"{prefix}schema={self.schema}", + f"{prefix}ipm_pairs={','.join(self.ipm_pairs) or '-'}", + f"{prefix}keys={','.join(self.keys)}", + ] + lines.extend( + f"{prefix}dataset={dataset}" for dataset in self.datasets + ) + return "\n".join(lines) + @dataclass(frozen=True) class Hdf5PairProbe: on: FileProbe off: FileProbe + def __str__(self) -> str: + return "\n".join( + ( + self.on.format_text(label="on"), + self.off.format_text(label="off"), + ) + ) + def to_dict(self) -> dict[str, Any]: return { "on": _file_probe_to_dict(self.on), diff --git a/src/cuphoton/xray/linear_prediction.py b/src/cuphoton/xray/linear_prediction.py index 8d522a3b..d93bd85e 100644 --- a/src/cuphoton/xray/linear_prediction.py +++ b/src/cuphoton/xray/linear_prediction.py @@ -4,7 +4,8 @@ from __future__ import annotations -from dataclasses import dataclass +from collections.abc import Sequence +from dataclasses import asdict, dataclass from math import pi from time import perf_counter from typing import Any @@ -431,6 +432,17 @@ class ModelOrderSweepEntry: singular_value_ratio: float | None elapsed_s: float + def __str__(self) -> str: + return ( + f"components={self.components} " + f"selected_model_order={self.selected_model_order} " + f"rms_residual={self.rms_residual:.6g} " + f"reconstruction_rms_error={self.reconstruction_rms_error:.6g} " + f"chi2={self.chi2:.6g} " + f"selected_roots={self.selected_root_count} " + f"decaying_roots={self.decaying_root_count}" + ) + @dataclass(frozen=True) class ModelOrderSweep: @@ -443,6 +455,115 @@ class ModelOrderSweep: best_reconstruction_rms_error: float entries: tuple[ModelOrderSweepEntry, ...] + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + def __str__(self) -> str: + return "\n".join( + ( + f"samples={self.samples}", + f"roots_backend={self.roots_backend}", + f"best_components={self.best_components}", + f"best_selected_model_order={self.best_selected_model_order}", + f"best_rms_residual={self.best_rms_residual:.6g}", + "best_reconstruction_rms_error=" + f"{self.best_reconstruction_rms_error:.6g}", + *(str(entry) for entry in self.entries), + ) + ) + + +@dataclass(frozen=True) +class ModelOrderSweepBatch: + """Per-trace model-order sweeps and their input provenance.""" + + samples: int + roots_backend: str + relative_tolerance: float + component_counts: tuple[int, ...] + sources: tuple[dict[str, Any], ...] + traces: tuple[ModelOrderSweep, ...] + + def to_dict(self) -> dict[str, Any]: + best_components = tuple( + trace.best_components for trace in self.traces + ) + best_orders = tuple( + trace.best_selected_model_order for trace in self.traces + ) + best_residuals = tuple( + trace.best_rms_residual for trace in self.traces + ) + best_errors = tuple( + trace.best_reconstruction_rms_error for trace in self.traces + ) + trace_payloads = [] + for index, (trace, source) in enumerate( + zip(self.traces, self.sources, strict=True) + ): + payload = trace.to_dict() + payload["source"] = source + payload["trace_index"] = index + trace_payloads.append(payload) + return { + "source": { + "kind": "trace-npz-batch", + "trace_count": len(self.traces), + "traces": self.sources, + }, + "trace_count": len(self.traces), + "samples": self.samples, + "roots_backend": self.roots_backend, + "relative_tolerance": self.relative_tolerance, + "component_counts": self.component_counts, + "best_components_by_trace": best_components, + "best_components_unique": tuple(sorted(set(best_components))), + "best_selected_model_orders_by_trace": best_orders, + "best_selected_model_orders_unique": tuple( + sorted(set(best_orders)) + ), + "best_rms_residual_min": min(best_residuals), + "best_rms_residual_max": max(best_residuals), + "best_reconstruction_rms_error_min": min(best_errors), + "best_reconstruction_rms_error_max": max(best_errors), + "traces": tuple(trace_payloads), + } + + def __str__(self) -> str: + lines = [ + "source=trace-npz-batch", + f"trace_count={len(self.traces)}", + f"samples={self.samples}", + f"roots_backend={self.roots_backend}", + "component_counts=" + ",".join(map(str, self.component_counts)), + "best_components_unique=" + + ",".join( + str(item) + for item in sorted( + {trace.best_components for trace in self.traces} + ) + ), + "best_selected_model_orders_unique=" + + ",".join( + str(item) + for item in sorted( + {trace.best_selected_model_order for trace in self.traces} + ) + ), + ] + for index, trace in enumerate(self.traces): + lines.extend( + ( + f"trace_index={index}", + f" best_components={trace.best_components}", + " best_selected_model_order=" + f"{trace.best_selected_model_order}", + " best_reconstruction_rms_error=" + f"{trace.best_reconstruction_rms_error:.6g}", + ) + ) + return "\n".join(lines) + @dataclass(frozen=True) class _PredictionRootsResult: @@ -591,6 +712,7 @@ def amplitudes_and_phases_cupy(a0: Any, a1: Any): "LinearPredictionVariableArtifactsBenchmark", "LinearPredictionRuntimeBridgeBenchmark", "ModelOrderSweep", + "ModelOrderSweepBatch", "ModelOrderSweepEntry", "PredictionRootsBackendBenchmark", "PredictionRootsBenchmark", @@ -619,6 +741,7 @@ def amplitudes_and_phases_cupy(a0: Any, a1: Any): "linear_prediction_cupy", "linear_prediction_numpy", "model_order_sweep", + "model_order_sweep_batch", "synthetic_prediction_coefficients", "synthetic_trace", "synthetic_trace_batch", @@ -1268,6 +1391,39 @@ def model_order_sweep( ) +def model_order_sweep_batch( + time, + trace_rows, + sources: Sequence[dict[str, Any]], + *, + components: tuple[int, ...], + roots_backend: str = "eigvals", + relative_tolerance: float = 0.01, +) -> ModelOrderSweepBatch: + """Sweep component counts for rows sharing one time axis.""" + + results = tuple( + model_order_sweep( + time, + trace, + components, + roots_backend=roots_backend, + relative_tolerance=relative_tolerance, + ) + for trace, _source in zip(trace_rows, sources, strict=True) + ) + if not results: + raise ValueError("trace_rows must contain at least one trace") + return ModelOrderSweepBatch( + samples=results[0].samples, + roots_backend=roots_backend, + relative_tolerance=float(relative_tolerance), + component_counts=tuple(int(item) for item in components), + sources=tuple(sources), + traces=results, + ) + + def benchmark_linear_prediction_p1_batch( *, samples: int = 96, diff --git a/src/cuphoton/xray/mode_refinement.py b/src/cuphoton/xray/mode_refinement.py new file mode 100644 index 00000000..337ce930 --- /dev/null +++ b/src/cuphoton/xray/mode_refinement.py @@ -0,0 +1,279 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Nonlinear least-squares refinement of damped modes, initialised by +linear prediction. + +Linear prediction gives frequencies and decays without an initial guess, +but it is not a maximum-likelihood estimator and its root filter can drop a +lightly damped mode under noise. Refining the same damped-mode model by +nonlinear least squares from the linear-prediction result can improve +parameter estimates, and a mode the filter dropped can be seeded from the +peak of the residual spectrum before refinement. The model is + + trace(t) = c + sum_k A_k exp(-d_k t) cos(w_k t + p_k) + +with an analytic Jacobian; the solver is ``scipy.optimize.least_squares``. +NumPy only. +""" + +from __future__ import annotations + +from dataclasses import dataclass, replace +from numbers import Integral + +import numpy as np +from scipy.optimize import least_squares + +from .linear_prediction import linear_prediction_numpy + + +@dataclass(frozen=True) +class RefinedModes: + """Refined damped modes and the fit diagnostics. + + Arrays are ordered by descending amplitude. ``sigma_*`` are one-sigma + local uncertainties from the Jacobian and residual variance. They are + NaN when the scaled Jacobian is rank deficient or no residual degrees + of freedom remain. ``initial_mode_count`` is how many oscillating modes + the linear-prediction step supplied and ``seeded_mode_count`` how many had + to be seeded from the residual spectrum. + """ + + time: np.ndarray + amplitude: np.ndarray + decay: np.ndarray + angular_frequency: np.ndarray + phase: np.ndarray + constant: float + sigma_amplitude: np.ndarray + sigma_decay: np.ndarray + sigma_angular_frequency: np.ndarray + sigma_phase: np.ndarray + reconstruction: np.ndarray + residual_rms: float + converged: bool + iterations: int + initial_mode_count: int + seeded_mode_count: int + + +def _fit_inputs(time, trace) -> tuple[np.ndarray, np.ndarray]: + time, trace = np.asarray(time), np.asarray(trace) + if np.iscomplexobj(time) or np.iscomplexobj(trace): + raise ValueError("time and trace must be real") + time = np.asarray(time, dtype=np.float64) + trace = np.asarray(trace, dtype=np.float64) + if ( + time.ndim != 1 + or trace.ndim != 1 + or time.shape != trace.shape + or time.size < 2 + or not np.all(np.isfinite(time)) + or not np.all(np.isfinite(trace)) + ): + raise ValueError( + "time and trace must be matching finite one-dimensional arrays " + "with at least two samples" + ) + return time, trace + + +def _wrap_phase(phase): + return np.where( + (phase >= -np.pi) & (phase < np.pi), + phase, + (phase + np.pi) % (2.0 * np.pi) - np.pi, + ) + + +def _model_and_jacobian(theta: np.ndarray, time: np.ndarray, n: int): + model = np.full(time.shape, float(theta[4 * n])) + jac = np.empty((time.size, 4 * n + 1)) + jac[:, 4 * n] = 1.0 + for k in range(n): + a, d, w, p = theta[4 * k : 4 * k + 4] + env = np.exp(-d * time) + arg = w * time + _wrap_phase(p) + cos = np.cos(arg) + sin = np.sin(arg) + model += a * env * cos + jac[:, 4 * k] = env * cos + jac[:, 4 * k + 1] = -a * time * env * cos + jac[:, 4 * k + 2] = -a * time * env * sin + jac[:, 4 * k + 3] = -a * env * sin + return model, jac + + +def seed_from_residual( + time: np.ndarray, residual: np.ndarray, pad: int = 4096 +) -> tuple[float, float, float, float]: + """Seed one damped mode from the strongest peak of the residual + spectrum: amplitude from the peak height, zero decay, zero phase.""" + time, residual = _fit_inputs(time, residual) + dt = float(time[1] - time[0]) + r = residual - residual.mean() + pad = max(pad, time.size) + spec = np.abs(np.fft.rfft(r, pad)) + freq = 2.0 * np.pi * np.fft.rfftfreq(pad, d=dt) + k = int(np.argmax(spec[1:])) + 1 + return float(2.0 * spec[k] / time.size), 0.0, float(freq[k]), 0.0 + + +def refine_modes( + time: np.ndarray, + trace: np.ndarray, + initial_modes: list[tuple[float, float, float, float]], + constant: float, + *, + max_nfev: int = 2000, + initial_mode_count: int | None = None, + seeded_mode_count: int = 0, +) -> RefinedModes: + """Refine ``initial_modes`` (amplitude, decay, angular frequency, phase) + and the constant by nonlinear least squares. Initial and trial phases + are reduced to ``[-pi, pi)`` before evaluating the model.""" + time, trace = _fit_inputs(time, trace) + n = len(initial_modes) + if n == 0: + raise ValueError("at least one initial mode is required") + theta0 = np.array( + [v for m in initial_modes for v in m] + [constant], dtype=np.float64 + ) + theta0[3 : 4 * n : 4] = _wrap_phase(theta0[3 : 4 * n : 4]) + + def residual(theta): + return _model_and_jacobian(theta, time, n)[0] - trace + + def jacobian(theta): + return _model_and_jacobian(theta, time, n)[1] + + res = least_squares( + residual, + theta0, + jac=jacobian, + method="lm", + xtol=1e-12, + ftol=1e-12, + max_nfev=max_nfev, + ) + theta = res.x + model, jac = _model_and_jacobian(theta, time, n) + resid = trace - model + dof = time.size - theta.size + sigma = np.full(theta.size, np.nan) + try: + # Normalize columns before testing rank so parameter units do not + # determine identifiability. SVD avoids squaring the condition number. + scale = np.linalg.norm(jac, axis=0) + scale = np.where(scale > 0, scale, 1.0) + _, singular, vh = np.linalg.svd(jac / scale, full_matrices=False) + cutoff = np.finfo(float).eps * max(jac.shape) * singular[0] + if dof > 0 and np.all(singular > cutoff): + var = float(resid @ resid) / dof + sigma = np.sqrt(np.sum((vh / singular[:, None]) ** 2, axis=0)) + sigma *= np.sqrt(var) / scale + except np.linalg.LinAlgError: + pass + amp = theta[0 : 4 * n : 4].copy() + dec = theta[1 : 4 * n : 4].copy() + freq = theta[2 : 4 * n : 4].copy() + ph = _wrap_phase(theta[3 : 4 * n : 4]) + # a negative amplitude or frequency is the same mode with a shifted phase + neg = amp < 0 + amp[neg] = -amp[neg] + ph[neg] = ph[neg] + np.pi + negf = freq < 0 + freq[negf] = -freq[negf] + ph[negf] = -ph[negf] + ph = _wrap_phase(ph) + order = np.argsort(-amp) + sig = sigma.reshape(-1) + return RefinedModes( + time=time, + amplitude=amp[order], + decay=dec[order], + angular_frequency=freq[order], + phase=ph[order], + constant=float(theta[4 * n]), + sigma_amplitude=sig[0 : 4 * n : 4][order], + sigma_decay=sig[1 : 4 * n : 4][order], + sigma_angular_frequency=sig[2 : 4 * n : 4][order], + sigma_phase=sig[3 : 4 * n : 4][order], + reconstruction=model, + residual_rms=float(np.sqrt(np.mean(resid**2))), + converged=bool(res.success), + iterations=int(res.nfev), + initial_mode_count=( + n if initial_mode_count is None else initial_mode_count + ), + seeded_mode_count=seeded_mode_count, + ) + + +def linear_prediction_refined( + time, + trace, + n_components: int, + n_modes: int, + *, + roots_backend: str = "eigvals", +) -> RefinedModes: + """Linear prediction, then nonlinear refinement of the ``n_modes`` + strongest oscillating modes; modes the linear-prediction step did not + return are seeded one at a time from the residual spectrum. + + Amplitudes and phases use the first input sample as the time origin, + matching linear prediction; the returned time retains the input axis. + """ + if ( + isinstance(n_modes, bool) + or not isinstance(n_modes, Integral) + or n_modes < 1 + ): + raise ValueError("n_modes must be a positive integer") + time, trace = _fit_inputs(time, trace) + lp = linear_prediction_numpy( + time, trace, n_components, roots_backend=roots_backend + ) + fit_time = time - time[0] + w = np.asarray(lp.angular_frequency, dtype=float) + d = np.asarray(lp.decay, dtype=float) + a = np.asarray(lp.amplitude, dtype=float) + p = np.asarray(lp.phase, dtype=float) + modes = [ + (float(a[i]), float(d[i]), float(w[i]), float(p[i])) + for i in range(w.size) + if w[i] > 0 + ] + modes.sort(key=lambda m: -m[0]) + modes = modes[:n_modes] + found = len(modes) + constant = float(np.mean(trace)) + residual = trace - np.asarray(lp.reconstruction, dtype=float) + seeded = 0 + while len(modes) < n_modes: + modes.append(seed_from_residual(fit_time, residual)) + seeded += 1 + partial = refine_modes(fit_time, trace, modes, constant) + residual = trace - partial.reconstruction + modes = [ + ( + float(partial.amplitude[i]), + float(partial.decay[i]), + float(partial.angular_frequency[i]), + float(partial.phase[i]), + ) + for i in range(len(modes)) + ] + constant = partial.constant + result = refine_modes( + fit_time, + trace, + modes, + constant, + initial_mode_count=found, + seeded_mode_count=seeded, + ) + return replace(result, time=time) diff --git a/src/cuphoton/xray/mode_refinement_batched.py b/src/cuphoton/xray/mode_refinement_batched.py new file mode 100644 index 00000000..d764b113 --- /dev/null +++ b/src/cuphoton/xray/mode_refinement_batched.py @@ -0,0 +1,417 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Batched Levenberg-Marquardt refinement of damped modes on NumPy or CuPy. + +Many detector-row traces share one time axis and one model shape (K damped +modes plus a constant), so the refinement of :mod:`mode_refinement` can be +run for a whole batch at once: analytic Jacobians for all traces by +broadcasting, normal equations solved as a batch of small (4K+1) systems, +and per-trace damping with accept/reject. The array module is a parameter, +so the same code runs on NumPy and CuPy; the parameters are the same +(amplitude, decay, angular frequency, phase per mode, then the constant). +Both solvers find local optima; convergence alone does not establish that +they found the same optimum. +""" + +from __future__ import annotations + +from dataclasses import asdict, dataclass +from numbers import Integral, Real +from time import perf_counter +from typing import Any + +import numpy as np + + +def _sync(xp) -> None: + dev = getattr(getattr(xp, "cuda", None), "Device", None) + if dev is not None: + dev().synchronize() + + +def _wrap_phase(xp, phase): + return xp.where( + (phase >= -xp.pi) & (phase < xp.pi), + phase, + (phase + xp.pi) % (2.0 * xp.pi) - xp.pi, + ) + + +def model_and_jacobian_batched(xp, theta, time, n_modes: int): + """theta (B, 4K+1), time (N,) -> model (B, N), jacobian (B, N, 4K+1).""" + b = theta.shape[0] + n = time.shape[0] + p = 4 * n_modes + 1 + t = time[None, :] + model = xp.broadcast_to(theta[:, 4 * n_modes][:, None], (b, n)).copy() + jac = xp.empty((b, n, p), dtype=theta.dtype) + jac[:, :, 4 * n_modes] = 1.0 + for k in range(n_modes): + a = theta[:, 4 * k][:, None] + d = theta[:, 4 * k + 1][:, None] + w = theta[:, 4 * k + 2][:, None] + ph = _wrap_phase(xp, theta[:, 4 * k + 3][:, None]) + env = xp.exp(-d * t) + arg = w * t + ph + cos = xp.cos(arg) + sin = xp.sin(arg) + model += a * env * cos + jac[:, :, 4 * k] = env * cos + jac[:, :, 4 * k + 1] = -a * t * env * cos + jac[:, :, 4 * k + 2] = -a * t * env * sin + jac[:, :, 4 * k + 3] = -a * env * sin + return model, jac + + +@dataclass(frozen=True) +class BatchedRefinement: + """Refined parameters for a batch: ``theta`` (B, 4K+1) in the same + layout as the input, per-trace residual rms, convergence flags, the + number of iterations used and the wall time including any device + synchronisation.""" + + theta: Any + residual_rms: Any + converged: Any + iterations: int + elapsed_s: float + + +def refine_modes_batched( + xp, + time, + traces, + theta0, + n_modes: int, + *, + max_iter: int = 60, + tol: float = 1e-9, + lambda0: float = 1e-3, +) -> BatchedRefinement: + """Levenberg-Marquardt over a batch of traces (B, N) from starts + ``theta0`` (B, 4K+1). Damping is per trace: a step that lowers the + cost is accepted and its damping divided by 3, otherwise the damping + is multiplied by 5 and the trace keeps its parameters. Iteration + stops when every trace has converged or ``max_iter`` is reached. + Convergence requires a small gradient in Jacobian-scaled coordinates; + a step made tiny by damping does not imply stationarity. Scaling the + Jacobian columns makes the solve and stopping criterion independent of + parameter units. Positive damping also handles zero Jacobian columns + from an initial zero-amplitude mode. + """ + for name, count in (("n_modes", n_modes), ("max_iter", max_iter)): + if ( + isinstance(count, bool) + or not isinstance(count, Integral) + or count < 1 + ): + raise ValueError(f"{name} must be a positive integer") + for name, value in (("tol", tol), ("lambda0", lambda0)): + if ( + isinstance(value, bool) + or not isinstance(value, Real) + or not np.isfinite(value) + or value <= 0 + ): + raise ValueError(f"{name} must be finite and positive") + for name, values in ( + ("time", time), + ("traces", traces), + ("theta0", theta0), + ): + if xp.iscomplexobj(values): + raise ValueError(f"{name} must contain only real values") + t = xp.asarray(time, dtype=xp.float64) + y = xp.asarray(traces, dtype=xp.float64) + theta = xp.asarray(theta0, dtype=xp.float64).copy() + if y.ndim != 2 or theta.ndim != 2 or theta.shape[0] != y.shape[0]: + raise ValueError("traces must be (B, N) and theta0 (B, 4K+1)") + if theta.shape[1] != 4 * n_modes + 1: + raise ValueError("theta0 must have 4 * n_modes + 1 columns") + if t.ndim != 1 or t.shape[0] != y.shape[1]: + raise ValueError("time must be (N,) matching traces (B, N)") + if y.shape[0] == 0 or y.shape[1] == 0: + raise ValueError("traces must contain at least one trace and sample") + for name, values in (("time", t), ("traces", y), ("theta0", theta)): + if not bool(xp.all(xp.isfinite(values))): + raise ValueError(f"{name} must contain only finite values") + theta[:, 3 : 4 * n_modes : 4] = _wrap_phase( + xp, theta[:, 3 : 4 * n_modes : 4] + ) + _sync(xp) + start = perf_counter() + b = y.shape[0] + lam = xp.full(b, lambda0, dtype=xp.float64) + model, jac = model_and_jacobian_batched(xp, theta, t, n_modes) + resid = model - y + cost = xp.sum(resid * resid, axis=1) + converged = xp.zeros(b, dtype=bool) + eye = xp.eye(theta.shape[1], dtype=xp.float64)[None, :, :] + iterations = 0 + for _ in range(max_iter): + iterations += 1 + scale = xp.sqrt(xp.sum(jac * jac, axis=1)) + scale = xp.where(scale > 0, scale, 1.0) + scaled_jac = jac / scale[:, None, :] + jtj = xp.einsum("bnp,bnq->bpq", scaled_jac, scaled_jac) + jtr = xp.einsum("bnp,bn->bp", scaled_jac, resid) + stationary = xp.max(xp.abs(jtr), axis=1) <= ( + tol * xp.maximum(xp.sqrt(cost), 1.0) + ) + converged = converged | stationary + if bool(xp.all(converged)): + break + step = ( + -xp.linalg.solve(jtj + lam[:, None, None] * eye, jtr[:, :, None])[ + :, :, 0 + ] + / scale + ) + trial = theta + step + model_t, jac_t = model_and_jacobian_batched(xp, trial, t, n_modes) + resid_t = model_t - y + cost_t = xp.sum(resid_t * resid_t, axis=1) + accept = (cost_t < cost) & ~converged + theta = xp.where(accept[:, None], trial, theta) + resid = xp.where(accept[:, None], resid_t, resid) + jac = xp.where(accept[:, None, None], jac_t, jac) + cost = xp.where(accept, cost_t, cost) + lam = xp.where(accept, lam / 3.0, lam * 5.0) + _sync(xp) + rms = xp.sqrt(cost / y.shape[1]) + return BatchedRefinement( + theta=theta, + residual_rms=rms, + converged=converged, + iterations=iterations, + elapsed_s=perf_counter() - start, + ) + + +@dataclass(frozen=True) +class RefinementBenchmark: + """Timings and fit quality from each path's fastest repetition. + + Convergence counts and worst residual RMS accompany parameter agreement; + a fast or matching result does not establish that every fit converged. + """ + + traces: int + samples: int + n_modes: int + repeat: int + cpu_serial_scipy_s: float + numpy_batched_s: float + cupy_batched_s: float | None + max_abs_theta_diff_numpy: float + max_abs_theta_diff_cupy: float | None + gpu_error: str | None + cpu_serial_converged: int + numpy_batched_converged: int + cupy_batched_converged: int | None + cpu_serial_max_residual_rms: float + numpy_batched_max_residual_rms: float + cupy_batched_max_residual_rms: float | None + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + def __str__(self) -> str: + lines = [ + f"traces={self.traces} samples={self.samples} " + f"repeat={self.repeat}", + f"cpu_serial_scipy_s={self.cpu_serial_scipy_s:.6g}", + f"numpy_batched_s={self.numpy_batched_s:.6g}", + f"max_abs_theta_diff_numpy={self.max_abs_theta_diff_numpy:.3g}", + f"cpu_serial_converged={self.cpu_serial_converged}/{self.traces}", + f"numpy_batched_converged={self.numpy_batched_converged}/{self.traces}", + f"cpu_serial_max_residual_rms={self.cpu_serial_max_residual_rms:.6g}", + f"numpy_batched_max_residual_rms={self.numpy_batched_max_residual_rms:.6g}", + ] + if self.gpu_error: + lines.extend( + ("gpu_status=unavailable", f"gpu_error={self.gpu_error}") + ) + elif self.cupy_batched_s is None: + lines.append("gpu_status=skipped") + else: + lines.extend( + ( + f"cupy_batched_s={self.cupy_batched_s:.6g}", + f"max_abs_theta_diff_cupy={self.max_abs_theta_diff_cupy:.3g}", + f"cupy_batched_converged={self.cupy_batched_converged}/{self.traces}", + f"cupy_batched_max_residual_rms={self.cupy_batched_max_residual_rms:.6g}", + ) + ) + return "\n".join(lines) + + +def benchmark_refinement( + *, + samples: int = 96, + traces: int = 256, + n_modes: int = 2, + noise_sigma: float = 0.02, + seed: int = 0, + run_gpu: bool = True, + repeat: int = 3, +) -> RefinementBenchmark: + """Time the per-trace SciPy refinement against the batched solver on + NumPy and on CuPy, from identical perturbed starts, and report the + largest parameter difference between the batched and SciPy results. + + Every path is timed ``repeat`` times and the best time is reported. + The batched timings cover the solver only: the arrays are on the + device before the clock starts and the device is synchronised before + and after. The SciPy timing covers the Python loop over traces.""" + for name, value in ( + ("samples", samples), + ("traces", traces), + ("n_modes", n_modes), + ("repeat", repeat), + ): + if ( + isinstance(value, bool) + or not isinstance(value, Integral) + or value < 1 + ): + raise ValueError(f"{name} must be a positive integer") + if not np.isfinite(noise_sigma) or noise_sigma < 0: + raise ValueError("noise_sigma must be finite and nonnegative") + from .mode_refinement import refine_modes + from .synthetic_validation import ( + DEFAULT_MODES, + modes_to_theta, + synthetic_modes_trace, + ) + + if n_modes > len(DEFAULT_MODES): + raise ValueError(f"n_modes must be at most {len(DEFAULT_MODES)}") + modes = DEFAULT_MODES[:n_modes] + fx = synthetic_modes_trace(samples, modes=modes) + rng = np.random.default_rng(seed) + y = fx.clean[None, :] + rng.normal(0.0, noise_sigma, (traces, samples)) + truth = modes_to_theta(modes, fx.constant) + theta0 = truth[None, :] * (1.0 + 0.05 * rng.standard_normal((traces, 1))) + theta0 = theta0 + 0.02 * rng.standard_normal(theta0.shape) + + # warm the device path once so compile and allocation costs are not + # charged to the timed run + if run_gpu: + try: + import cupy as cp + + refine_modes_batched(cp, fx.time, y[:16], theta0[:16], n_modes) + except Exception: # pragma: no cover - reported below + pass + + def serial_scipy(): + out = np.empty_like(theta0) + converged = 0 + max_residual_rms = 0.0 + for i in range(traces): + init = [ + tuple(theta0[i, 4 * k : 4 * k + 4]) for k in range(n_modes) + ] + r = refine_modes(fx.time, y[i], init, float(theta0[i, -1])) + converged += int(r.converged) + max_residual_rms = max(max_residual_rms, r.residual_rms) + # refine_modes normalises signs and orders by amplitude; undo + # the ordering by matching frequencies to the start + for k in range(n_modes): + j = int( + np.argmin( + np.abs(r.angular_frequency - theta0[i, 4 * k + 2]) + ) + ) + out[i, 4 * k : 4 * k + 4] = ( + r.amplitude[j], + r.decay[j], + r.angular_frequency[j], + r.phase[j], + ) + out[i, -1] = r.constant + return out, converged, max_residual_rms + + cpu_serial_s = float("inf") + for _ in range(repeat): + start = perf_counter() + candidate = serial_scipy() + elapsed = perf_counter() - start + if elapsed < cpu_serial_s: + cpu_serial_s = elapsed + serial, serial_converged, serial_max_residual_rms = candidate + + res_np = min( + ( + refine_modes_batched(np, fx.time, y, theta0, n_modes) + for _ in range(repeat) + ), + key=lambda r: r.elapsed_s, + ) + diff_np = float( + np.max(np.abs(_canonical(res_np.theta, n_modes) - serial)) + ) + + cupy_s = None + diff_cp = None + gpu_error = None + cp_converged = None + cp_max_residual_rms = None + if run_gpu: + try: + import cupy as cp + + res_cp = min( + ( + refine_modes_batched(cp, fx.time, y, theta0, n_modes) + for _ in range(repeat) + ), + key=lambda r: r.elapsed_s, + ) + theta_cp = cp.asnumpy(res_cp.theta) + cupy_s = res_cp.elapsed_s + diff_cp = float( + np.max(np.abs(_canonical(theta_cp, n_modes) - serial)) + ) + cp_converged = int(cp.count_nonzero(res_cp.converged)) + cp_max_residual_rms = float(cp.max(res_cp.residual_rms)) + except Exception as exc: # pragma: no cover - depends on hardware + gpu_error = f"{type(exc).__name__}: {exc}" + return RefinementBenchmark( + traces=traces, + samples=samples, + n_modes=n_modes, + repeat=repeat, + cpu_serial_scipy_s=cpu_serial_s, + numpy_batched_s=res_np.elapsed_s, + cupy_batched_s=cupy_s, + max_abs_theta_diff_numpy=diff_np, + max_abs_theta_diff_cupy=diff_cp, + gpu_error=gpu_error, + cpu_serial_converged=serial_converged, + numpy_batched_converged=int(np.count_nonzero(res_np.converged)), + cupy_batched_converged=cp_converged, + cpu_serial_max_residual_rms=serial_max_residual_rms, + numpy_batched_max_residual_rms=float(np.max(res_np.residual_rms)), + cupy_batched_max_residual_rms=cp_max_residual_rms, + ) + + +def _canonical(theta: np.ndarray, n_modes: int) -> np.ndarray: + """Same sign and phase conventions as :func:`refine_modes`.""" + out = np.array(theta, dtype=np.float64, copy=True) + for k in range(n_modes): + a = out[:, 4 * k] + w = out[:, 4 * k + 2] + ph = out[:, 4 * k + 3] + neg = a < 0 + a[neg] = -a[neg] + ph[neg] += np.pi + negf = w < 0 + w[negf] = -w[negf] + ph[negf] = -ph[negf] + out[:, 4 * k + 3] = (ph + np.pi) % (2.0 * np.pi) - np.pi + return out diff --git a/src/cuphoton/xray/mode_refinement_viz.py b/src/cuphoton/xray/mode_refinement_viz.py new file mode 100644 index 00000000..a18101fe --- /dev/null +++ b/src/cuphoton/xray/mode_refinement_viz.py @@ -0,0 +1,193 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Visual comparison of linear prediction and experimental mode refinement.""" + +from __future__ import annotations + +from collections.abc import Sequence +from dataclasses import dataclass +from pathlib import Path + +import numpy as np + +from .linear_prediction import linear_prediction_numpy +from .mode_refinement import linear_prediction_refined +from .synthetic_validation import distort_trace, synthetic_modes_trace + + +@dataclass(frozen=True) +class ModesReviewReport: + """Location and fit status of a standalone synthetic-mode review.""" + + html_path: Path + trace_count: int + converged_count: int + + def __str__(self) -> str: + return ( + f"modes_review={self.html_path}\n" + f"traces={self.trace_count}\n" + f"refinement_converged={self.converged_count}/{self.trace_count}" + ) + + def to_dict(self) -> dict[str, object]: + return { + "html_path": str(self.html_path), + "trace_count": self.trace_count, + "converged_count": self.converged_count, + } + + +def _mode_curve(amplitude, decay, frequency, phase, time): + return ( + amplitude * np.exp(-decay * time) * np.cos(frequency * time + phase) + ) + + +def write_modes_review( + output: Path | str, + *, + snr_db: Sequence[float] = (40.0, 20.0, 10.0, 5.0), + distortion: tuple[str, float] | None = None, + seed: int = 11, + samples: int = 96, +) -> ModesReviewReport: + """Write an offline HTML comparison at each requested noise level. + + SNR uses the standard deviation of the noise-free signal. Distortions + change that signal; the mode panel still shows the undistorted reference + modes. Refinement convergence is an optimizer status, not scientific + acceptance. Requires the ``viz`` extra. + """ + snr_levels = tuple(float(level) for level in snr_db) + if not snr_levels: + raise ValueError("snr_db must contain at least one level") + if not np.all(np.isfinite(snr_levels)): + raise ValueError("snr_db levels must be finite") + fixture = synthetic_modes_trace(samples) + time = fixture.time + clean = fixture.clean + if distortion is not None: + clean = distort_trace( + time, fixture.modes, fixture.constant, *distortion + ) + with np.errstate(over="ignore", under="ignore", divide="ignore"): + signal_rms = np.std(clean - clean.mean()) + sigmas = signal_rms / np.power(10.0, np.asarray(snr_levels) / 20) + if not np.all(np.isfinite(sigmas)) or np.any(sigmas <= 0): + raise ValueError("snr_db must produce finite, positive noise levels") + + try: + from bokeh.embed import file_html + from bokeh.layouts import gridplot + from bokeh.plotting import figure + from bokeh.resources import INLINE + except ImportError as exc: # pragma: no cover - depends on the viz extra + raise RuntimeError( + "Bokeh is required for modes review; install 'cuphoton[viz]' " + "or run 'uv sync --extra viz' for development" + ) from exc + + fine = np.linspace(time[0], time[-1], 2000) + rng = np.random.default_rng(seed) + rows = [] + converged_count = 0 + for snr, sigma in zip(snr_levels, sigmas, strict=True): + trace = clean + rng.normal(0.0, sigma, size=time.shape) + lp = linear_prediction_numpy(time, trace, 6) + refined = linear_prediction_refined( + time, trace, 6, len(fixture.modes) + ) + converged_count += int(refined.converged) + lp_ratio = np.sqrt(np.mean((trace - lp.reconstruction) ** 2)) / sigma + status = "converged" if refined.converged else "not converged" + left = figure( + width=760, + height=280, + title=( + f"{snr:g} dB, sigma {sigma:.4g}: residual/noise " + f"lp {lp_ratio:.1f}, refined " + f"{refined.residual_rms / sigma:.1f}; {status}" + ), + x_axis_label="time", + y_axis_label="trace", + ) + left.scatter(time, trace, size=4, color="black", legend_label="trace") + # Plot the generated samples: a glitch is defined on that sample grid. + left.line(time, clean, color="gray", legend_label="noise-free signal") + left.line( + time, + np.asarray(lp.reconstruction), + color="red", + legend_label="linear prediction", + ) + left.line( + time, + refined.reconstruction, + color="blue", + line_dash="dashed", + legend_label="refined", + ) + left.legend.location = "top_right" + left.legend.label_text_font_size = "8pt" + right = figure( + width=520, + height=280, + title=( + "modes: reference (solid), LP (dashed), refined (dotted)" + + ( + f"; signal distortion {distortion[0]}" + if distortion + else "" + ) + ), + x_axis_label="time", + ) + colors = ("#1f77b4", "#2ca02c", "#ff7f0e", "#9467bd") + for index, mode in enumerate(fixture.modes): + right.line( + fine, + _mode_curve( + mode.amplitude, + mode.decay, + mode.angular_frequency, + mode.phase, + fine, + ), + color=colors[index % len(colors)], + alpha=0.5, + legend_label=f"reference w={mode.angular_frequency:g}", + ) + for result, label, color, dash in ( + (lp, "lp", "red", "dashed"), + (refined, "refined", "blue", "dotted"), + ): + for index, frequency in enumerate(result.angular_frequency): + if frequency == 0: + continue + decay = float(result.decay[index]) + right.line( + fine, + _mode_curve( + float(result.amplitude[index]), + decay, + float(frequency), + float(result.phase[index]), + fine, + ), + color=color, + line_dash=dash, + legend_label=f"{label} w={frequency:.3f} d={decay:+.3f}", + ) + right.legend.label_text_font_size = "7pt" + rows.append([left, right]) + title = "XRay linear prediction and mode refinement" + if distortion is not None: + title += f" ({distortion[0]}:{distortion[1]:g})" + output = Path(output).expanduser() + html = file_html(gridplot(rows), INLINE, title) + output.parent.mkdir(parents=True, exist_ok=True) + output.write_text(html, encoding="utf-8") + return ModesReviewReport(output, len(snr_levels), converged_count) diff --git a/src/cuphoton/xray/subspace.py b/src/cuphoton/xray/subspace.py index 3088aa1d..98afb635 100644 --- a/src/cuphoton/xray/subspace.py +++ b/src/cuphoton/xray/subspace.py @@ -4,7 +4,7 @@ from __future__ import annotations -from dataclasses import dataclass +from dataclasses import asdict, dataclass from time import perf_counter import numpy as np @@ -27,6 +27,22 @@ class SubspaceMethodBenchmark: roots_real: tuple[float, ...] roots_imag: tuple[float, ...] + def __str__(self) -> str: + return "\n".join( + ( + f"method={self.method}", + f" svd_backend={self.svd_backend}", + f" svd_rank={self.svd_rank}", + f" rms_residual={self.rms_residual:.6g}", + " max_abs_reconstruction_diff=" + f"{self.max_abs_reconstruction_diff:.6g}", + f" elapsed_s={self.elapsed_s:.6g}", + ) + ) + + def to_dict(self) -> dict[str, object]: + return asdict(self) + @dataclass(frozen=True) class SubspaceBenchmark: @@ -37,6 +53,159 @@ class SubspaceBenchmark: baseline_rms_residual: float methods: tuple[SubspaceMethodBenchmark, ...] + def __str__(self) -> str: + return "\n".join( + ( + f"samples={self.samples}", + f"model_order={self.model_order}", + f"components={self.components}", + f"baseline_chi2={self.baseline_chi2:.6g}", + f"baseline_rms_residual={self.baseline_rms_residual:.6g}", + *(str(method) for method in self.methods), + ) + ) + + def to_dict(self) -> dict[str, object]: + return asdict(self) + + +@dataclass(frozen=True) +class _SubspaceMethodSummary: + method: str + svd_backend: str + trace_count: int + rms_residual_min: float + rms_residual_max: float + max_abs_reconstruction_diff_max: float + elapsed_s_total: float + + def __str__(self) -> str: + return "\n".join( + ( + f"method={self.method}", + f" svd_backend={self.svd_backend}", + f" trace_count={self.trace_count}", + " rms_residual_range=" + f"{self.rms_residual_min:.6g}:{self.rms_residual_max:.6g}", + " max_abs_reconstruction_diff_max=" + f"{self.max_abs_reconstruction_diff_max:.6g}", + ) + ) + + +@dataclass(frozen=True) +class SubspaceBenchmarkBatch: + samples: int + model_order: int + components: int + methods: tuple[str, ...] + svd_backends: tuple[str, ...] + sources: tuple[dict[str, object], ...] + traces: tuple[SubspaceBenchmark, ...] + + @property + def method_summary(self) -> tuple[_SubspaceMethodSummary, ...]: + summaries = [] + for method in self.methods: + for svd_backend in self.svd_backends: + results = tuple( + result + for trace in self.traces + for result in trace.methods + if (result.method, result.svd_backend) + == (method, svd_backend) + ) + summaries.append( + _SubspaceMethodSummary( + method=method, + svd_backend=svd_backend, + trace_count=len(results), + rms_residual_min=min( + [ + float("inf"), + *(result.rms_residual for result in results), + ] + ), + rms_residual_max=max( + [ + 0.0, + *(result.rms_residual for result in results), + ] + ), + max_abs_reconstruction_diff_max=max( + [ + 0.0, + *( + result.max_abs_reconstruction_diff + for result in results + ), + ] + ), + elapsed_s_total=sum( + (result.elapsed_s for result in results), 0.0 + ), + ) + ) + return tuple(summaries) + + def __str__(self) -> str: + baseline_rms = tuple( + trace.baseline_rms_residual for trace in self.traces + ) + return "\n".join( + ( + "source=trace-npz-batch", + f"trace_count={len(self.traces)}", + f"samples={self.samples}", + f"model_order={self.model_order}", + f"components={self.components}", + "baseline_rms_residual_range=" + f"{min(baseline_rms):.6g}:{max(baseline_rms):.6g}", + *(str(summary) for summary in self.method_summary), + ) + ) + + def to_dict(self) -> dict[str, object]: + traces = tuple( + { + **trace.to_dict(), + "source": source, + "svd_backends": self.svd_backends, + "trace_index": index, + } + for index, (trace, source) in enumerate( + zip(self.traces, self.sources, strict=True) + ) + ) + return { + "source": { + "kind": "trace-npz-batch", + "trace_count": len(self.sources), + "traces": self.sources, + }, + "trace_count": len(self.sources), + "samples": self.samples, + "model_order": self.model_order, + "components": self.components, + "svd_backends": self.svd_backends, + "baseline_chi2_min": min( + trace.baseline_chi2 for trace in self.traces + ), + "baseline_chi2_max": max( + trace.baseline_chi2 for trace in self.traces + ), + "baseline_rms_residual_min": min( + trace.baseline_rms_residual for trace in self.traces + ), + "baseline_rms_residual_max": max( + trace.baseline_rms_residual for trace in self.traces + ), + "method_summary": tuple( + asdict(summary) for summary in self.method_summary + ), + "traces": traces, + } + @dataclass(frozen=True) class _SubspaceSvd: @@ -137,6 +306,53 @@ def compare_subspace_methods( ) +def compare_subspace_methods_batch( + time, + trace_rows, + sources, + *, + model_order: int, + components: int, + methods: tuple[str, ...] = ("matrix-pencil", "esprit"), + svd_backends: tuple[str, ...] = ("full",), + pencil_rows: int | None = None, + randomized_oversamples: int = 8, + power_iterations: int = 1, + random_seed: int = 0, +) -> SubspaceBenchmarkBatch: + """Compare trace rows while retaining their individual source metadata.""" + + results = [] + trace_sources = [] + for trace, source in zip(trace_rows, sources, strict=True): + results.append( + compare_subspace_methods( + time, + trace, + model_order=model_order, + components=components, + methods=methods, + svd_backends=svd_backends, + pencil_rows=pencil_rows, + randomized_oversamples=randomized_oversamples, + power_iterations=power_iterations, + random_seed=random_seed, + ) + ) + trace_sources.append(source) + if not results: + raise ValueError("at least one trace is required") + return SubspaceBenchmarkBatch( + samples=int(np.asarray(time).shape[0]), + model_order=int(model_order), + components=int(components), + methods=tuple(methods), + svd_backends=tuple(svd_backends), + sources=tuple(trace_sources), + traces=tuple(results), + ) + + def synthetic_subspace_benchmark( *, samples: int = 96, @@ -458,8 +674,10 @@ def _sort_roots(roots: np.ndarray): __all__ = [ "SubspaceBenchmark", + "SubspaceBenchmarkBatch", "SubspaceMethodBenchmark", "compare_subspace_methods", + "compare_subspace_methods_batch", "esprit_roots", "matrix_pencil_roots", "synthetic_subspace_benchmark", diff --git a/src/cuphoton/xray/synthetic_validation.py b/src/cuphoton/xray/synthetic_validation.py new file mode 100644 index 00000000..7bafd241 --- /dev/null +++ b/src/cuphoton/xray/synthetic_validation.py @@ -0,0 +1,1152 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Synthetic damped-mode fixtures and Cramer-Rao validation for the +linear-prediction fit. + +The fixture generates traces with known modes so the workflow can be +validated without external datasets. The validation sweep compares the +estimator's scatter and bias with the Cramer-Rao lower bound computed from +the Fisher information of the same model, together with the rate at +which true modes are lost from the fit. Everything here runs on NumPy; the +estimator under test is :func:`linear_prediction_numpy` unless another +estimator callable is supplied. +""" + +from __future__ import annotations + +import json +import subprocess +from collections.abc import Callable, Sequence +from dataclasses import asdict, dataclass, field, replace +from pathlib import Path +from typing import Any + +import numpy as np + +from cuphoton.core.runtime import runtime_metadata + +from .linear_prediction import linear_prediction_numpy + +SUMMARY_SCHEMA_VERSION = 1 +PARAMETERS = ("amplitude", "decay", "angular_frequency", "phase") + + +@dataclass(frozen=True) +class DampedMode: + """One real damped sinusoid. + + ``amplitude * exp(-decay * t) * cos(angular_frequency * t + phase)`` + """ + + amplitude: float + decay: float + angular_frequency: float + phase: float = 0.0 + + +@dataclass(frozen=True) +class SyntheticModesTrace: + time: np.ndarray + trace: np.ndarray + clean: np.ndarray + modes: tuple[DampedMode, ...] + constant: float + noise_sigma: float + seed: int | None + + +DEFAULT_MODES = ( + DampedMode(1.25, 0.09, 2.4, 0.30), + DampedMode(0.45, 0.03, 0.9, -0.75), +) + + +def modes_model( + theta: np.ndarray, time: np.ndarray, n_modes: int +) -> np.ndarray: + """Evaluate the multi-mode model from a flat parameter vector. + + Four parameters per mode, then the constant. + """ + out = np.full(time.shape, float(theta[4 * n_modes])) + for k in range(n_modes): + a, d, w, p = theta[4 * k : 4 * k + 4] + out = out + a * np.exp(-d * time) * np.cos(w * time + p) + return out + + +def modes_to_theta( + modes: Sequence[DampedMode], constant: float +) -> np.ndarray: + return np.array( + [ + v + for m in modes + for v in (m.amplitude, m.decay, m.angular_frequency, m.phase) + ] + + [constant], + dtype=np.float64, + ) + + +def _modes_jacobian( + theta: np.ndarray, time: np.ndarray, n_modes: int +) -> np.ndarray: + """Analytic model derivatives in the order used by ``modes_to_theta``.""" + jac = np.ones((time.size, theta.size)) + for k in range(n_modes): + amplitude, decay, frequency, phase = theta[4 * k : 4 * k + 4] + envelope = np.exp(-decay * time) + angle = frequency * time + phase + cosine = envelope * np.cos(angle) + sine = envelope * np.sin(angle) + jac[:, 4 * k] = cosine + jac[:, 4 * k + 1] = -amplitude * time * cosine + jac[:, 4 * k + 2] = -amplitude * time * sine + jac[:, 4 * k + 3] = -amplitude * sine + return jac + + +def synthetic_modes_trace( + samples: int = 96, + *, + modes: Sequence[DampedMode] = DEFAULT_MODES, + constant: float = 0.15, + duration: float = 9.5, + noise_sigma: float = 0.0, + seed: int | None = None, +) -> SyntheticModesTrace: + """Build a trace with known modes plus optional white Gaussian noise. + + With the defaults and ``noise_sigma=0`` this reproduces + :func:`synthetic_trace` exactly. + """ + if samples < 16: + raise ValueError("samples must be at least 16") + if not np.isfinite(duration) or duration <= 0: + raise ValueError("duration must be finite and positive") + if not np.isfinite(noise_sigma) or noise_sigma < 0: + raise ValueError("noise_sigma must be finite and nonnegative") + time = np.linspace(0.0, duration, samples, dtype=np.float64) + clean = modes_model(modes_to_theta(modes, constant), time, len(modes)) + trace = clean + if noise_sigma > 0: + rng = np.random.default_rng(seed) + trace = clean + rng.normal(0.0, noise_sigma, size=time.shape) + return SyntheticModesTrace( + time, trace, clean, tuple(modes), constant, float(noise_sigma), seed + ) + + +DISTORTIONS = ( + "chirp", + "gaussian_envelope", + "baseline_drift", + "clip", + "glitch", +) + + +def distort_trace( + time: np.ndarray, + modes: Sequence[DampedMode], + constant: float, + kind: str, + amount: float, +) -> np.ndarray: + """Return a clean trace that departs from the damped-mode model. + + ``chirp``: every frequency rises by the fraction ``amount`` across the + record. ``gaussian_envelope``: the first mode decays as a Gaussian with + 1/e time ``amount`` from the first sample instead of an exponential. + ``baseline_drift``: a + linear ramp of ``amount`` across the record. ``clip``: the trace is + clipped above ``amount``. ``glitch``: one sample at mid-record is + offset by ``amount``. These are outside the model class on purpose: + the sweep reports how the fit responds and the residual-to-noise ratio + that can help diagnose them. + """ + if kind not in DISTORTIONS: + raise ValueError(f"unknown distortion {kind!r}") + if not np.isfinite(amount): + raise ValueError("distortion amount must be finite") + if kind == "gaussian_envelope" and amount <= 0: + raise ValueError("gaussian_envelope amount must be positive") + span = float(time[-1] - time[0]) + out = np.full(time.shape, float(constant)) + for k, m in enumerate(modes): + phase = m.angular_frequency * time + m.phase + if kind == "chirp": + phase = phase + ( + 0.5 + * m.angular_frequency + * amount + * (time - time[0]) ** 2 + / span + ) + if kind == "gaussian_envelope" and k == 0: + with np.errstate(over="ignore", under="ignore"): + env = np.exp(-(((time - time[0]) / amount) ** 2)) + else: + env = np.exp(-m.decay * time) + out = out + m.amplitude * env * np.cos(phase) + if kind == "baseline_drift": + out = out + amount * (time - time[0]) / span + elif kind == "clip": + out = np.minimum(out, amount) + elif kind == "glitch": + out = out.copy() + out[time.size // 2] += amount + return out + + +def cramer_rao_bounds( + modes: Sequence[DampedMode], + constant: float, + time: np.ndarray, + noise_sigma: float, + step: float = 1e-6, +) -> dict[str, np.ndarray]: + """Cramer-Rao lower bound (standard deviation) per parameter for white + Gaussian noise of ``noise_sigma``. + + The Fisher information is J^T J / sigma^2 with J the analytic Jacobian + of :func:`modes_model`. A column-scaled SVD gives the bounds without + forming or inverting the normal matrix. Models with numerically + unidentifiable parameters raise ``ValueError``. Returned per mode as + arrays ``amplitude``, ``decay``, ``angular_frequency``, ``phase`` plus + the scalar ``constant``. ``step`` remains accepted for compatibility + but is unused by the analytic calculation. + """ + if not np.isfinite(noise_sigma) or noise_sigma <= 0: + raise ValueError("noise_sigma must be finite and positive") + with np.errstate(over="ignore", under="ignore"): + noise_variance = np.square(np.float64(noise_sigma)) + if not np.isfinite(noise_variance) or noise_variance <= 0: + raise ValueError("noise_sigma must have a finite, positive variance") + if not np.isfinite(step) or step <= 0: + raise ValueError("step must be finite and positive") + theta = modes_to_theta(modes, constant) + time = np.asarray(time, dtype=np.float64) + if time.ndim != 1 or not np.all(np.isfinite(time)): + raise ValueError("time must be a finite one-dimensional array") + if not np.all(np.isfinite(theta)): + raise ValueError("modes and constant must be finite") + n = len(modes) + try: + with np.errstate(over="raise", divide="raise", invalid="raise"): + jac = _modes_jacobian(theta, time, n) + scales = np.linalg.norm(jac, axis=0) + if time.size < theta.size or np.any(scales == 0): + raise ValueError( + "Cramer-Rao bounds require identifiable modes" + ) + _, singular, vh = np.linalg.svd(jac / scales, full_matrices=False) + cutoff = np.finfo(float).eps * max(jac.shape) * singular[0] + if singular[-1] <= cutoff: + raise ValueError( + "Cramer-Rao bounds require identifiable modes" + ) + sd = noise_sigma * np.linalg.norm(vh / singular[:, None], axis=0) + sd /= scales + except (FloatingPointError, np.linalg.LinAlgError) as exc: + raise ValueError( + "noise_sigma or modes produce singular or nonfinite " + "Cramer-Rao bounds" + ) from exc + with np.errstate(over="ignore", under="ignore"): + variance = np.square(sd) + if not np.all(np.isfinite(variance)) or np.any(variance <= 0): + raise ValueError( + "noise_sigma or modes produce nonfinite or zero Cramer-Rao bounds" + ) + return { + "amplitude": sd[0 : 4 * n : 4], + "decay": sd[1 : 4 * n : 4], + "angular_frequency": sd[2 : 4 * n : 4], + "phase": sd[3 : 4 * n : 4], + "constant": sd[4 * n], + } + + +def _to_host(values) -> np.ndarray: + """Return a float NumPy array; CuPy arrays are copied to the host.""" + get = getattr(values, "get", None) + if get is not None and not isinstance(values, np.ndarray): + values = get() + return np.asarray(values, dtype=float) + + +def _integer_errors(errors: np.ndarray, n_modes: int) -> list[int]: + """Scale binary floats exactly, ranking nonfinite errors last.""" + ratios = [ + float(error).as_integer_ratio() if np.isfinite(error) else None + for error in errors + ] + denominator = max((r[1] for r in ratios if r is not None), default=1) + scaled = [ + r[0] * (denominator // r[1]) if r is not None else None + for r in ratios + ] + unavailable = ( + n_modes * max((x for x in scaled if x is not None), default=0) + 1 + ) + return [unavailable if x is None else x for x in scaled] + + +def _minimum_mode_assignment(cost: list[list[int]]) -> list[int]: + """Rectangular Hungarian assignment without rounding the integer costs.""" + n_rows, n_columns = len(cost), len(cost[0]) + row_price = [0] * (n_rows + 1) + column_price = [0] * (n_columns + 1) + owner = [0] * (n_columns + 1) + previous = [0] * (n_columns + 1) + for row in range(1, n_rows + 1): + owner[0] = row + column = 0 + slack: list[int | None] = [None] * (n_columns + 1) + visited = [False] * (n_columns + 1) + while True: + visited[column] = True + current_row = owner[column] + step = None + next_column = 0 + for candidate in range(1, n_columns + 1): + if visited[candidate]: + continue + reduced = ( + cost[current_row - 1][candidate - 1] + - row_price[current_row] + - column_price[candidate] + ) + best = slack[candidate] + if best is None or reduced < best: + best = slack[candidate] = reduced + previous[candidate] = column + if step is None or best < step: + step, next_column = best, candidate + assert step is not None + for candidate, best in enumerate(slack): + if visited[candidate]: + row_price[owner[candidate]] += step + column_price[candidate] -= step + elif best is not None: + slack[candidate] = best - step + column = next_column + if owner[column] == 0: + break + while column: + predecessor = previous[column] + owner[column] = owner[predecessor] + column = predecessor + assignment = [0] * n_rows + for column, row in enumerate(owner[1:]): + if row: + assignment[row - 1] = column + return assignment + + +def match_modes( + modes: Sequence[DampedMode], + angular_frequency: np.ndarray, + decay: np.ndarray, + *, + tolerance: float, +) -> list[int | None]: + """Match within the frequency ``tolerance`` (rad per time unit). + + Maximize recovered modes, then minimize total frequency error, then + total decay error. Return fitted indices or ``None`` for missing modes. + Nonfinite decay errors rank last without preventing frequency recovery. + """ + if not np.isfinite(tolerance) or tolerance < 0: + raise ValueError("tolerance must be finite and nonnegative") + out: list[int | None] = [None] * len(modes) + if not modes or angular_frequency.size == 0: + return out + err = np.abs( + np.abs([m.angular_frequency for m in modes])[:, None] + - np.abs(angular_frequency)[None, :] + ) + valid = err <= tolerance + if not np.any(valid): + return out + with np.errstate(over="ignore", invalid="ignore"): + decay_err = np.abs( + np.array([m.decay for m in modes])[:, None] - decay[None, :] + ) + frequency_cost = _integer_errors(err[valid], len(modes)) + decay_cost = _integer_errors(decay_err[valid], len(modes)) + # Integer bounds make each priority dominate every possible total of + # the lower priorities. No epsilon lets decay outweigh frequency. + frequency_scale = len(modes) * max(decay_cost) + 1 + real_cost = [ + frequency * frequency_scale + damping + for frequency, damping in zip(frequency_cost, decay_cost, strict=True) + ] + missing_cost = len(modes) * max(real_cost) + 1 + invalid_cost = (len(modes) + 1) * missing_cost + cost = [ + [invalid_cost] * angular_frequency.size + [missing_cost] * len(modes) + for _ in modes + ] + for (row, column), value in zip( + np.argwhere(valid), real_cost, strict=True + ): + cost[row][column] = value + for row, column in enumerate(_minimum_mode_assignment(cost)): + if column < angular_frequency.size: + out[row] = column + return out + + +@dataclass(frozen=True) +class ParameterStats: + """Statistics of one estimated parameter over the trials in which the + mode was recovered, against the unconditional Cramer-Rao bound. + + ``crlb_std`` is the standard-deviation bound and ``crlb_variance`` its + square; ``std_over_crlb_std`` compares the sample standard deviation + with the former. + """ + + truth: float + bias: float + std: float + rmse: float + crlb_std: float + crlb_variance: float + std_over_crlb_std: float + + +@dataclass(frozen=True) +class ModeValidation: + """One true mode: recovery count, loss rate and per-parameter + statistics computed over the recovered trials only.""" + + angular_frequency: ParameterStats + decay: ParameterStats + amplitude: ParameterStats + phase: ParameterStats + recovered_trials: int + loss_rate: float + + def __str__(self) -> str: + w = self.angular_frequency + d = self.decay + return ( + f"mode w={w.truth:g} decay={d.truth:g}: " + f"freq std/crlb={w.std:.3g}/{w.crlb_std:.3g} " + f"({w.std_over_crlb_std:.2f}x) bias={w.bias:+.3g}; " + f"decay std/crlb={d.std:.3g}/{d.crlb_std:.3g} " + f"({d.std_over_crlb_std:.2f}x) bias={d.bias:+.3g}; " + f"recovered={self.recovered_trials} " + f"loss_rate={self.loss_rate:.3f}" + ) + + @property + def frequency_ratio(self) -> float: + return self.angular_frequency.std_over_crlb_std + + @property + def decay_ratio(self) -> float: + return self.decay.std_over_crlb_std + + +@dataclass(frozen=True) +class ValidationLevel: + """One sweep condition. A trial is successful when the estimator + returned and every true mode was matched; ``estimator_errors`` counts + trials in which the estimator raised, which are also failed trials. + Solver convergence is recorded separately from frequency-based recovery; + estimators without convergence diagnostics leave both counts at zero. + """ + + snr_db: float + noise_sigma: float + trials_attempted: int + trials_successful: int + trials_failed: int + estimator_errors: int + n_components: int + modes: tuple[ModeValidation, ...] + any_mode_lost_rate: float + residual_ratio: float + solver_converged: int = 0 + solver_unconverged: int = 0 + + def __str__(self) -> str: + lines = [ + f"snr_db={self.snr_db:g} sigma={self.noise_sigma:.4g} " + f"trials={self.trials_successful}/{self.trials_attempted} " + f"estimator_errors={self.estimator_errors} " + f"any_mode_lost_rate={self.any_mode_lost_rate:.3f} " + f"residual_ratio={self.residual_ratio:.2f}" + ] + if self.solver_converged + self.solver_unconverged: + lines.append( + f"solver_converged={self.solver_converged} " + f"solver_unconverged={self.solver_unconverged}" + ) + lines.extend(f" {mode}" for mode in self.modes) + return "\n".join(lines) + + +@dataclass(frozen=True) +class ValidationSweep: + samples: int + duration: float + signal_rms: float + levels: tuple[ValidationLevel, ...] + match_tolerance: float + backend: str + modes: tuple[DampedMode, ...] = DEFAULT_MODES + constant: float = 0.15 + seed: int = 0 + distortion: tuple[str, float] | None = None + extra: dict = field(default_factory=dict) + + def __str__(self) -> str: + lines = [f"estimator={self.backend}"] + if self.distortion is not None: + lines.append( + f"distortion={self.distortion[0]}:{self.distortion[1]:g}" + ) + lines.extend(str(level) for level in self.levels) + return "\n".join(lines) + + def to_dict(self) -> dict: + return asdict(self) + + +@dataclass(frozen=True) +class ValidationReport: + """Validation sweeps with their serializable summary and artifacts.""" + + sweeps: tuple[ValidationSweep, ...] + summary: dict[str, Any] + + def __str__(self) -> str: + first = self.sweeps[0] + lines = [ + f"samples={first.samples}", + f"signal_rms={first.signal_rms:.6g}", + ] + lines.extend(str(sweep) for sweep in self.sweeps) + lines.extend( + f"{key}={name if name else 'not written'}" + for key, name in self.summary["artifacts"].items() + ) + return "\n".join(lines) + + def to_dict(self) -> dict[str, Any]: + return self.summary + + +def _wrap_phase(values: np.ndarray) -> np.ndarray: + return (values + np.pi) % (2.0 * np.pi) - np.pi + + +def _canonical( + amplitude: np.ndarray, angular_frequency: np.ndarray, phase: np.ndarray +): + """Normalize amplitude/frequency signs and wrap equivalent phases.""" + amplitude = np.array(amplitude, dtype=float, copy=True) + angular_frequency = np.array(angular_frequency, dtype=float, copy=True) + phase = np.array(phase, dtype=float, copy=True) + negative_frequency = angular_frequency < 0 + angular_frequency[negative_frequency] *= -1 + phase[negative_frequency] *= -1 + neg = amplitude < 0 + amplitude[neg] = -amplitude[neg] + phase[neg] = phase[neg] + np.pi + return amplitude, angular_frequency, _wrap_phase(phase) + + +def _stats( + samples: np.ndarray, truth: float, crlb_std: float, *, wrap: bool +) -> ParameterStats: + err = samples - truth + if wrap: + err = _wrap_phase(err) + n = err.size + bias = float(err.mean()) if n else float("nan") + std = float(err.std(ddof=1)) if n > 1 else float("nan") + rmse = float(np.sqrt(np.mean(err**2))) if n else float("nan") + return ParameterStats( + truth=float(truth), + bias=bias, + std=std, + rmse=rmse, + crlb_std=float(crlb_std), + crlb_variance=float(crlb_std) ** 2, + std_over_crlb_std=std / float(crlb_std), + ) + + +def validation_sweep( + *, + samples: int = 96, + modes: Sequence[DampedMode] = DEFAULT_MODES, + constant: float = 0.15, + duration: float = 9.5, + snr_db: Sequence[float] = (40.0, 30.0, 20.0, 10.0), + trials: int = 200, + n_components: int = 6, + match_tolerance: float = 0.3, + seed: int = 20260914, + estimator: Callable | None = None, + backend: str = "cpu", + distortion: tuple[str, float] | None = None, +) -> ValidationSweep: + """Monte Carlo sweep of the linear-prediction estimator against the + Cramer-Rao bound. + + ``snr_db`` is defined against the rms of the mean-removed clean trace. + For each level the per-mode scatter, bias and rmse of every parameter + are compared with the bound, and ``loss_rate`` is the fraction of + trials in which the true mode had no fitted mode within + ``match_tolerance`` (radians per time unit) of its angular frequency; + each fitted mode can match one true mode. + + The statistics are computed over the recovered trials only and the + bound is unconditional. The root filter keeps decaying roots, so the + surviving decay estimates are truncated at zero; for a lightly damped + mode the reported decay statistics are therefore conditional on + survival and can move non-monotonically with noise. Read them together + with ``loss_rate``. + + ``estimator`` may return CuPy arrays; they are moved to the host. A + trial in which the estimator raises is counted in ``estimator_errors`` + and as a loss of every mode. A returned ``converged`` flag is counted + separately: frequency recovery and conditional parameter statistics do + not imply solver convergence. + + Truth and fitted modes use nonnegative amplitudes and frequencies, + with equivalent phases in ``[-pi, pi)``. The returned truth uses the + same convention as the parameter statistics. + + ``distortion`` = (kind, amount) replaces the clean trace with one from + :func:`distort_trace`, outside the model class; ``residual_ratio`` is + then the median rms residual of the reconstruction divided by the + noise sigma. Large values can reflect model mismatch, missed modes or + estimator error; this ratio alone does not distinguish those causes. + """ + if trials < 1: + raise ValueError("trials must be at least 1") + if n_components <= 0: + raise ValueError("n_components must be positive") + if not np.isfinite(match_tolerance) or match_tolerance < 0: + raise ValueError("match_tolerance must be finite and nonnegative") + if len(snr_db) == 0: + raise ValueError("snr_db must contain at least one level") + if not np.all(np.isfinite(snr_db)): + raise ValueError("snr_db levels must be finite") + amplitudes, frequencies, phases = _canonical( + np.array([m.amplitude for m in modes]), + np.array([m.angular_frequency for m in modes]), + np.array([m.phase for m in modes]), + ) + modes = tuple( + replace( + m, amplitude=float(a), angular_frequency=float(w), phase=float(p) + ) + for m, a, w, p in zip( + modes, amplitudes, frequencies, phases, strict=True + ) + ) + est = estimator or linear_prediction_numpy + base = synthetic_modes_trace( + samples, modes=modes, constant=constant, duration=duration + ) + clean = base.clean + if distortion is not None: + clean = distort_trace( + base.time, modes, constant, distortion[0], distortion[1] + ) + signal_rms = float(np.std(clean - clean.mean())) + with np.errstate(over="ignore", under="ignore", divide="ignore"): + noise_sigmas = signal_rms / np.power(10.0, np.asarray(snr_db) / 20) + # Validate every level before trials, including underflow/overflow in + # the SNR conversion, so a later invalid level cannot yield partial work. + bounds_per_level = [ + cramer_rao_bounds(modes, constant, base.time, float(sigma)) + for sigma in noise_sigmas + ] + rng = np.random.default_rng(seed) + levels = [] + for snr, sigma, bounds in zip( + snr_db, noise_sigmas, bounds_per_level, strict=True + ): + found: list[dict[str, list[float]]] = [ + {name: [] for name in PARAMETERS} for _ in modes + ] + lost = np.zeros(len(modes), dtype=int) + any_lost = 0 + errors = 0 + solver_converged = 0 + solver_unconverged = 0 + ratios = [] + for _ in range(trials): + y = clean + rng.normal(0.0, sigma, size=base.time.shape) + try: + r = est(base.time, y, n_components) + except Exception: + errors += 1 + any_lost += 1 + lost += 1 + continue + status = getattr(r, "converged", None) + if status is not None: + solver_converged += int(bool(status)) + solver_unconverged += int(not bool(status)) + d = _to_host(r.decay) + a, w, p = _canonical( + _to_host(r.amplitude), + _to_host(r.angular_frequency), + _to_host(r.phase), + ) + rec = _to_host(r.reconstruction) + ratios.append(float(np.sqrt(np.mean((y - rec) ** 2)) / sigma)) + hit = match_modes(modes, w, d, tolerance=match_tolerance) + if any(h is None for h in hit): + any_lost += 1 + for k, h in enumerate(hit): + if h is None: + lost[k] += 1 + continue + found[k]["amplitude"].append(a[h]) + found[k]["decay"].append(d[h]) + found[k]["angular_frequency"].append(w[h]) + found[k]["phase"].append(p[h]) + per_mode = [] + for k, m in enumerate(modes): + stats = { + name: _stats( + np.asarray(found[k][name]), + getattr(m, name), + bounds[name][k], + wrap=name == "phase", + ) + for name in PARAMETERS + } + per_mode.append( + ModeValidation( + angular_frequency=stats["angular_frequency"], + decay=stats["decay"], + amplitude=stats["amplitude"], + phase=stats["phase"], + recovered_trials=int(trials - lost[k]), + loss_rate=float(lost[k] / trials), + ) + ) + levels.append( + ValidationLevel( + snr_db=float(snr), + noise_sigma=float(sigma), + trials_attempted=trials, + trials_successful=trials - any_lost, + trials_failed=any_lost, + estimator_errors=errors, + solver_converged=solver_converged, + solver_unconverged=solver_unconverged, + n_components=n_components, + modes=tuple(per_mode), + any_mode_lost_rate=float(any_lost / trials), + residual_ratio=float(np.median(ratios)) + if ratios + else float("nan"), + ) + ) + return ValidationSweep( + samples=samples, + duration=duration, + signal_rms=signal_rms, + levels=tuple(levels), + match_tolerance=match_tolerance, + backend=backend, + modes=tuple(modes), + constant=constant, + seed=seed, + distortion=distortion, + ) + + +def validate_linear_prediction( + *, + samples: int = 96, + snr_db: Sequence[float] = (40.0, 30.0, 20.0, 10.0), + trials: int = 200, + n_components: int = 6, + seed: int = 20260914, + distortion: tuple[str, float] | None = None, + output_dir: Path | str | None = None, + refine: bool = False, +) -> ValidationReport: + """Validate linear prediction and optionally write its run artifacts. + + The returned report provides text and strict-JSON-compatible summaries + of the same sweep. Set ``output_dir`` to write ``summary.json`` and, + when Bokeh is available, a validation figure. Set ``refine`` to compare + the experimental nonlinear estimator using the same noise realizations. + """ + sweeps: tuple[ValidationSweep, ...] = ( + validation_sweep( + samples=samples, + snr_db=snr_db, + trials=trials, + n_components=n_components, + seed=seed, + distortion=distortion, + ), + ) + if refine: + from .mode_refinement import linear_prediction_refined + + n_modes = len(sweeps[0].modes) + sweeps += ( + validation_sweep( + samples=samples, + snr_db=snr_db, + trials=trials, + n_components=n_components, + seed=seed, + distortion=distortion, + estimator=lambda time, trace, components: ( + linear_prediction_refined( + time, trace, components, n_modes + ) + ), + backend="cpu-refined", + ), + ) + summary = ( + build_summary(sweeps) + if output_dir is None + else write_validation_run(output_dir, sweeps) + ) + return ValidationReport(sweeps=sweeps, summary=summary) + + +def _source_provenance() -> tuple[str | None, bool | None]: + """Revision and tracked-file dirtiness for this module's checkout.""" + try: + source = Path(__file__).resolve() + options: dict[str, Any] = { + "cwd": source.parent, + "capture_output": True, + "text": True, + "timeout": 5, + "check": True, + } + # A wheel installed in another project's .venv can have an enclosing + # Git checkout without belonging to that checkout's source tree. + subprocess.run( + ["git", "ls-files", "--error-unmatch", "--", source.name], + **options, + ) + revision = subprocess.run( + ["git", "rev-parse", "HEAD"], **options + ).stdout.strip() + status = subprocess.run( + ["git", "status", "--porcelain", "--untracked-files=no"], + **options, + ) + except (OSError, subprocess.SubprocessError): + return None, None + return revision, bool(status.stdout.strip()) + + +def source_revision() -> str | None: + """Git revision when this module belongs to a tracked source checkout.""" + return _source_provenance()[0] + + +def build_summary( + sweeps: Sequence[ValidationSweep], + *, + command: str = "linear-prediction-validate", + artifacts: dict[str, str | None] | None = None, +) -> dict[str, Any]: + """Assemble the ``schema_version`` 1 summary of one or more sweeps. + + ``config`` records the truth, the sampling, the noise definition, the + seed, the trial count, the model order and the mode-matching rule; + ``runtime`` the package, Python, NumPy and SciPy versions, the source + revision and the backend; ``results`` one entry per sweep condition + (estimator by distortion by signal-to-noise level); ``artifacts`` the + relative filenames written next to the summary. + """ + if not sweeps: + raise ValueError("at least one sweep is required") + first = sweeps[0] + shared_fields = ( + "samples", + "duration", + "modes", + "constant", + "seed", + "match_tolerance", + ) + first_levels = [ + (level.snr_db, level.trials_attempted, level.n_components) + for level in first.levels + ] + for sweep in sweeps[1:]: + if ( + any( + getattr(sweep, name) != getattr(first, name) + for name in shared_fields + ) + or [ + (level.snr_db, level.trials_attempted, level.n_components) + for level in sweep.levels + ] + != first_levels + ): + raise ValueError( + "summary sweeps must share truth, sampling, seed, " + "matching and trial settings" + ) + dt = first.duration / (first.samples - 1) + config = { + "model": ( + "trace(t) = constant + sum_k amplitude_k exp(-decay_k t) " + "cos(angular_frequency_k t + phase_k)" + ), + "units": { + "time": "input time unit", + "amplitude": "trace units", + "decay": "1 / time unit", + "angular_frequency": "rad / time unit", + "phase": "rad", + }, + "truth": { + "modes": [asdict(m) for m in first.modes], + "constant": first.constant, + }, + "sampling": { + "samples": first.samples, + "duration": first.duration, + "dt": dt, + "start": 0.0, + }, + "noise": { + "kind": "white Gaussian, independent per sample", + "snr_db_definition": ( + "20 log10(rms of the mean-removed clean trace / sigma)" + ), + "signal_rms": first.signal_rms, + "snr_db": [lvl.snr_db for lvl in first.levels], + }, + "seed": first.seed, + "trials_per_level": first.levels[0].trials_attempted, + "n_components": first.levels[0].n_components, + "mode_matching": { + "criterion": ( + "a true mode is recovered when a fitted mode lies within " + "the tolerance of its angular frequency; fitted modes are " + "assigned to maximize matches, then minimize total frequency " + "error, then total decay error, and each is used at most once" + ), + "tolerance_rad_per_time": first.match_tolerance, + }, + "statistics": ( + "bias, std (ddof 1) and rmse are computed over the recovered " + "trials of each mode; crlb_std and crlb_variance are the " + "unconditional Cramer-Rao bounds from the model's Fisher " + "information using an analytic Jacobian at the true " + "parameters; the phase " + "error is wrapped to [-pi, pi)" + ), + "estimators": sorted({s.backend for s in sweeps}), + "distortions": sorted( + { + f"{s.distortion[0]}:{s.distortion[1]:g}" + if s.distortion + else "none" + for s in sweeps + } + ), + } + runtime = runtime_metadata(backend="cpu", dtype="float64") + try: + from scipy import __version__ as scipy_version + except ImportError: # pragma: no cover - scipy is a dependency + scipy_version = None + runtime["scipy_version"] = scipy_version + runtime["source_revision"], runtime["source_dirty"] = _source_provenance() + results = [] + for sweep in sweeps: + for lvl in sweep.levels: + results.append( + { + "estimator": sweep.backend, + "distortion": ( + { + "kind": sweep.distortion[0], + "amount": sweep.distortion[1], + } + if sweep.distortion + else None + ), + "snr_db": lvl.snr_db, + "signal_rms": sweep.signal_rms, + "noise_sigma": lvl.noise_sigma, + "trials": { + "attempted": lvl.trials_attempted, + "successful": lvl.trials_successful, + "failed": lvl.trials_failed, + "estimator_errors": lvl.estimator_errors, + }, + **( + { + "solver_convergence": { + "converged": lvl.solver_converged, + "unconverged": lvl.solver_unconverged, + } + } + if lvl.solver_converged + lvl.solver_unconverged + else {} + ), + "any_mode_lost_rate": lvl.any_mode_lost_rate, + "residual_rms_over_sigma_median": ( + lvl.residual_ratio + if np.isfinite(lvl.residual_ratio) + else None + ), + "modes": [ + { + "recovered_trials": m.recovered_trials, + "loss_rate": m.loss_rate, + "statistics_over": "recovered_trials", + **{ + name: { + key: value if np.isfinite(value) else None + for key, value in asdict( + getattr(m, name) + ).items() + } + for name in PARAMETERS + }, + } + for m in lvl.modes + ], + } + ) + return { + "schema_version": SUMMARY_SCHEMA_VERSION, + "command": command, + "config": config, + "runtime": runtime, + "results": results, + "artifacts": dict(artifacts or {}), + } + + +def write_validation_figure( + sweeps: Sequence[ValidationSweep], path: Path, *, title: str +) -> bool: + """Standalone Bokeh HTML: per estimator and parameter the sample + standard deviation against the bound versus signal-to-noise, and the + loss rate. Returns False when the ``viz`` extra is not installed.""" + try: + from bokeh.embed import file_html + from bokeh.layouts import gridplot + from bokeh.plotting import figure + from bokeh.resources import INLINE + except ImportError: + return False + palette = ["#1f77b4", "#d62728", "#2ca02c", "#9467bd", "#ff7f0e"] + rows = [] + for sweep in sweeps: + snr = [lvl.snr_db for lvl in sweep.levels] + label = sweep.backend + ( + f" {sweep.distortion[0]}:{sweep.distortion[1]:g}" + if sweep.distortion + else "" + ) + row = [] + for name in ("angular_frequency", "decay"): + fig = figure( + width=360, + height=260, + y_axis_type="log", + title=f"{label}: {name} std (points) vs bound (dashed)", + x_axis_label="snr (dB)", + y_axis_label="std", + ) + for k, mode in enumerate(sweep.modes): + std = [ + getattr(lvl.modes[k], name).std for lvl in sweep.levels + ] + crlb = [ + getattr(lvl.modes[k], name).crlb_std + for lvl in sweep.levels + ] + color = palette[k % len(palette)] + fig.line(snr, crlb, color=color, line_dash="dashed") + fig.scatter( + snr, + std, + color=color, + size=7, + legend_label=f"mode w={mode.angular_frequency:g}", + ) + fig.legend.label_text_font_size = "8pt" + row.append(fig) + fig = figure( + width=360, + height=260, + title=f"{label}: loss rate", + x_axis_label="snr (dB)", + y_axis_label="fraction of trials", + ) + for k, mode in enumerate(sweep.modes): + fig.line( + snr, + [lvl.modes[k].loss_rate for lvl in sweep.levels], + color=palette[k % len(palette)], + legend_label=f"mode w={mode.angular_frequency:g}", + ) + fig.line( + snr, + [lvl.any_mode_lost_rate for lvl in sweep.levels], + color="black", + line_dash="dotted", + legend_label="any mode", + ) + fig.legend.label_text_font_size = "8pt" + row.append(fig) + rows.append(row) + path.write_text( + file_html(gridplot(rows), INLINE, title), encoding="utf-8" + ) + return True + + +def write_validation_run( + output_dir: Path | str, + sweeps: Sequence[ValidationSweep], + *, + command: str = "linear-prediction-validate", + figure_name: str = "validation.html", +) -> dict[str, Any]: + """Write ``summary.json`` and, when Bokeh is available, the figure + into ``output_dir`` (created if needed). Returns the summary.""" + out = Path(output_dir) + out.mkdir(parents=True, exist_ok=True) + have_figure = write_validation_figure( + sweeps, out / figure_name, title="xray linear prediction validation" + ) + artifacts = { + "summary": "summary.json", + "figure": figure_name if have_figure else None, + } + summary = build_summary(sweeps, command=command, artifacts=artifacts) + with open(out / "summary.json", "w", encoding="utf-8", newline="\n") as f: + json.dump(summary, f, indent=2, sort_keys=True, allow_nan=False) + f.write("\n") + return summary diff --git a/tests/core/test_cli_contract.py b/tests/core/test_cli_contract.py index ad4fdec1..5a9fdc95 100644 --- a/tests/core/test_cli_contract.py +++ b/tests/core/test_cli_contract.py @@ -75,13 +75,13 @@ def test_public_command_surface_counts_are_exact() -> None: ("xpois", 7, 7, 162, 1), ("xscan", 44, 44, 229, 1), ("xrep", 6, 6, 123, 1), - ("xray", 33, 31, 392, 1), + ("xray", 36, 34, 412, 1), ] assert len(per_group) == 6 - assert sum(item[1] for item in per_group) == 94 - assert sum(item[1] + item[4] for item in per_group) == 99 - assert sum(item[2] for item in per_group) == 91 - assert sum(item[3] for item in per_group) == 961 + assert sum(item[1] for item in per_group) == 97 + assert sum(item[1] + item[4] for item in per_group) == 102 + assert sum(item[2] for item in per_group) == 94 + assert sum(item[3] for item in per_group) == 981 def test_public_registry_order_and_component_derivations() -> None: diff --git a/tests/xray/test_cli.py b/tests/xray/test_cli.py index 10b78bdf..338a9e9a 100644 --- a/tests/xray/test_cli.py +++ b/tests/xray/test_cli.py @@ -4,6 +4,7 @@ from __future__ import annotations +import io import json import numpy as np @@ -37,6 +38,59 @@ def test_cli_help(capsys): assert "GPU-accelerated X-ray detector analysis tools" in captured.out +@pytest.mark.parametrize("as_json", [False, True]) +@pytest.mark.parametrize( + ("command_name", "marker"), + [ + ("DataProbeCommand", "on"), + ("LinearPredictionValidateCommand", "samples"), + ("LinearPredictionRefineBenchmarkCommand", "traces"), + ("LinearPredictionModesReviewCommand", "html"), + ("ModelOrderSweepCommand", "trace_count"), + ("SubspaceBenchmarkCommand", "trace_count"), + ], +) +def test_report_commands_use_injected_output( + command_name, marker, as_json, tmp_path, capsys +): + from cuphoton.xray import commands + + options = {"json": as_json} + if command_name == "DataProbeCommand": + h5py = pytest.importorskip("h5py") + with h5py.File(tmp_path / "cube.h5", "w") as handle: + handle.create_dataset("imgs", data=np.zeros((2, 3, 4))) + options.update(h5dir=tmp_path, fon="cube.h5", foff="cube.h5") + elif command_name == "LinearPredictionValidateCommand": + options.update(samples=48, snr_db="30", trials=2) + elif command_name == "LinearPredictionRefineBenchmarkCommand": + options.update(traces=2, repeat=1, no_gpu=True) + elif command_name == "LinearPredictionModesReviewCommand": + pytest.importorskip("bokeh") + options.update(output_path=tmp_path / "review.html", snr_db="30") + else: + time, rows = synthetic_trace_batch(samples=48, traces=2) + paths = [tmp_path / f"trace-{index}.npz" for index in range(2)] + for path, trace in zip(paths, rows, strict=True): + np.savez(path, time=time, trace=trace) + options["trace_npz"] = paths + if command_name == "ModelOrderSweepCommand": + options["component_values"] = [2, 4] + + output, error = io.StringIO(), io.StringIO() + command = getattr(commands, command_name)(output=output, error=error) + command.start(options) + + captured = capsys.readouterr() + assert captured.out == captured.err == error.getvalue() == "" + text = output.getvalue() + assert text.endswith("\n") + if as_json: + assert isinstance(json.loads(text), dict) + else: + assert marker in text + + def test_app_spec_uses_canonical_namespace(): spec = get_component("xray") @@ -64,6 +118,9 @@ def test_core_command_registry_loads_xray_commands(): "extract-trace": "et", "gpu-policy": "gp", "linear-prediction-benchmark": "lpb", + "linear-prediction-validate": "lpv", + "linear-prediction-refine-benchmark": "lprb", + "linear-prediction-modes-review": "lpmr", "linear-prediction-fixed-stages-benchmark": "lpfsb", "linear-prediction-p2-benchmark": "lppb", "linear-prediction-profile-summary": "lpps", @@ -205,6 +262,122 @@ def test_detector_artifacts_cli_rejects_unknown_diagnostic_level( assert "invalid choice: 'verbose'" in captured.err +def test_linear_prediction_validate_text_json_and_output_dir( + tmp_path, capsys +): + # the optional --distortion string must be allowed to be absent + assert main(["lpv", "--trials", "3", "--snr-db", "30"]) == 0 + assert "loss_rate=" in capsys.readouterr().out + assert ( + main( + [ + "lpv", + "--trials", + "3", + "--snr-db", + "30", + "--distortion", + "chirp:0.05", + "--json", + ] + ) + == 0 + ) + payload = json.loads(capsys.readouterr().out) + assert payload["schema_version"] == 1 + assert payload["results"][0]["distortion"] == { + "kind": "chirp", + "amount": 0.05, + } + assert payload["artifacts"] == {} + out = tmp_path / "lpv" + assert ( + main( + [ + "lpv", + "--trials", + "3", + "--snr-db", + "30", + "--output-dir", + str(out), + ] + ) + == 0 + ) + text = capsys.readouterr().out + assert "summary=summary.json" in text + written = json.loads((out / "summary.json").read_text()) + assert written["results"][0]["trials"]["attempted"] == 3 + assert written["artifacts"]["summary"] == "summary.json" + + +@pytest.mark.parametrize("option", ["--components", "--trials"]) +@pytest.mark.parametrize("value", ["0", "-1"]) +def test_linear_prediction_validate_requires_positive_counts( + capsys, option, value +): + assert main(["lpv", option, value]) == 2 + captured = capsys.readouterr() + assert f"argument {option}: must be at least 1" in captured.err + assert not captured.out + + +@pytest.mark.parametrize("json_args", [[], ["--json"]]) +@pytest.mark.parametrize( + ("args", "message"), + [ + (["--snr-db", "nan"], "snr_db"), + (["--snr-db", "inf"], "snr_db"), + (["--distortion", "gaussian_envelope:0"], "gaussian_envelope"), + (["--distortion", "chirp"], "--distortion must be kind:amount"), + (["--distortion", "chirp:bad"], "--distortion must be kind:amount"), + ], +) +def test_linear_prediction_validate_rejects_invalid_reports( + capsys, args, message, json_args +): + _assert_cli_error(capsys, ["lpv", *args, *json_args], message) + + +def test_linear_prediction_validate_text_distinguishes_estimator_errors( + capsys, monkeypatch +): + from cuphoton.xray import synthetic_validation + + def broken(t, y, k): + raise RuntimeError("estimator failed") + + monkeypatch.setattr( + synthetic_validation, "linear_prediction_numpy", broken + ) + assert main(["lpv", "--trials", "2", "--snr-db", "30"]) == 0 + output = capsys.readouterr().out + assert "trials=0/2 estimator_errors=2" in output + assert "any_mode_lost_rate=1.000" in output + + +def test_linear_prediction_validate_refine_adds_a_second_estimator(capsys): + assert ( + main(["lpv", "--trials", "3", "--snr-db", "30", "--refine", "--json"]) + == 0 + ) + payload = json.loads(capsys.readouterr().out) + assert payload["config"]["estimators"] == ["cpu", "cpu-refined"] + assert [r["estimator"] for r in payload["results"]] == [ + "cpu", + "cpu-refined", + ] + + +def test_linear_prediction_refine_benchmark_cli_cpu_only(capsys): + assert main(["lprb", "--traces", "4", "--no-gpu", "--json"]) == 0 + payload = json.loads(capsys.readouterr().out) + assert payload["traces"] == 4 + assert payload["cupy_batched_s"] is None + assert payload["max_abs_theta_diff_numpy"] < 1e-6 + + def test_gpu_policy(capsys): assert main(["gpu-policy"]) == 0 captured = capsys.readouterr() @@ -2219,3 +2392,28 @@ def test_subspace_acceptance_rejects_with_text_status(tmp_path, capsys): assert "status=rejected" in output assert "trace_count=1" in output assert "accepted=False" in output + + +def test_modes_review_cli_writes_an_offline_artifact(tmp_path, capsys): + pytest.importorskip("bokeh") + path = tmp_path / "review.html" + assert ( + main(["lpmr", "--output", str(path), "--snr-db", "30", "--json"]) == 0 + ) + report = json.loads(capsys.readouterr().out) + assert report["html_path"] == str(path) + assert report["trace_count"] == 1 + assert "cdn.bokeh.org" not in path.read_text() + + +@pytest.mark.parametrize( + "args", [["--snr-db", "nan"], ["--distortion", "chirp"]] +) +def test_modes_review_cli_rejects_invalid_inputs(tmp_path, capsys, args): + path = tmp_path / "review.html" + _assert_cli_error( + capsys, + ["lpmr", "--output", str(path), *args], + "snr_db" if args[0] == "--snr-db" else "--distortion", + ) + assert not path.exists() diff --git a/tests/xray/test_linear_prediction.py b/tests/xray/test_linear_prediction.py index 0a9d9072..d78ab3c2 100644 --- a/tests/xray/test_linear_prediction.py +++ b/tests/xray/test_linear_prediction.py @@ -35,6 +35,7 @@ linear_prediction_p2_artifacts_numpy, linear_prediction_variable_artifacts_cupy_batched, model_order_sweep, + model_order_sweep_batch, synthetic_prediction_coefficients, synthetic_trace, synthetic_trace_batch, @@ -1910,6 +1911,55 @@ def test_model_order_sweep_chooses_smallest_qualified_order(): assert result.best_selected_model_order == 2 +def test_model_order_sweep_batch_retains_results_and_provenance(): + time, rows = synthetic_trace_batch(samples=48, traces=2) + sources = [ + {"kind": "trace-npz", "path": "first.npz"}, + {"kind": "trace-npz", "path": "second.npz"}, + ] + batch = model_order_sweep_batch( + time.tolist(), + rows, + sources, + components=(6, 2, 4), + relative_tolerance=0.05, + ) + expected = [ + model_order_sweep(time, row, (6, 2, 4), relative_tolerance=0.05) + for row in rows + ] + + payload = batch.to_dict() + assert batch.sources == tuple(sources) + assert payload["source"]["traces"] == tuple(sources) + assert payload["trace_count"] == payload["source"]["trace_count"] == 2 + assert payload["component_counts"] == (6, 2, 4) + assert payload["best_components_by_trace"] == tuple( + result.best_components for result in expected + ) + assert payload["best_selected_model_orders_unique"] == tuple( + sorted({result.best_selected_model_order for result in expected}) + ) + residuals = [result.best_rms_residual for result in expected] + assert payload["best_rms_residual_min"] == pytest.approx(min(residuals)) + assert payload["best_rms_residual_max"] == pytest.approx(max(residuals)) + for index, trace in enumerate(payload["traces"]): + assert trace["trace_index"] == index + assert trace["source"] == sources[index] + assert trace["best_components"] == expected[index].best_components + assert trace["best_rms_residual"] == pytest.approx(residuals[index]) + assert trace["entries"] == batch.traces[index].to_dict()["entries"] + + +@pytest.mark.parametrize("source_count", [1, 3]) +def test_model_order_sweep_batch_rejects_unmatched_provenance(source_count): + time, rows = synthetic_trace_batch(samples=32, traces=2) + sources = [{"kind": "trace-npz"}] * source_count + + with pytest.raises(ValueError, match="zip"): + model_order_sweep_batch(time, rows, sources, components=(2,)) + + def test_compare_cpu_gpu_can_skip_gpu(): comparison = compare_cpu_gpu( samples=64, diff --git a/tests/xray/test_mode_refinement.py b/tests/xray/test_mode_refinement.py new file mode 100644 index 00000000..49810d2f --- /dev/null +++ b/tests/xray/test_mode_refinement.py @@ -0,0 +1,229 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +import numpy as np +import pytest + +from cuphoton.xray.mode_refinement import ( + _model_and_jacobian, + linear_prediction_refined, + refine_modes, + seed_from_residual, +) +from cuphoton.xray.synthetic_validation import ( + cramer_rao_bounds, + synthetic_modes_trace, + validation_sweep, +) + + +def test_analytic_jacobian_matches_finite_differences(): + fx = synthetic_modes_trace(96) + theta = np.array([1.25, 0.09, 2.4, 0.3, 0.45, 0.03, 0.9, -0.75, 0.15]) + _, jac = _model_and_jacobian(theta, fx.time, 2) + h = 1e-6 + for k in range(theta.size): + e = np.zeros_like(theta) + e[k] = h + num = ( + _model_and_jacobian(theta + e, fx.time, 2)[0] + - _model_and_jacobian(theta - e, fx.time, 2)[0] + ) / (2 * h) + np.testing.assert_allclose(jac[:, k], num, atol=1e-7) + + +def test_refinement_recovers_exact_modes_from_a_perturbed_start(): + fx = synthetic_modes_trace(96) + start = [(1.1, 0.05, 2.5, 0.0), (0.5, 0.0, 0.8, 0.0)] + r = refine_modes(fx.time, fx.trace, start, 0.0) + assert r.converged + np.testing.assert_allclose(r.angular_frequency, [2.4, 0.9], atol=1e-8) + np.testing.assert_allclose(r.decay, [0.09, 0.03], atol=1e-8) + np.testing.assert_allclose(r.amplitude, [1.25, 0.45], atol=1e-8) + assert abs(r.constant - 0.15) < 1e-8 + assert r.residual_rms < 1e-10 + + +def test_lp_refinement_preserves_the_first_sample_time_origin(): + fx = synthetic_modes_trace(96) + base = linear_prediction_refined(fx.time, fx.trace, 6, 2) + shifted_time = fx.time + 1000.0 + shifted = linear_prediction_refined(shifted_time, fx.trace, 6, 2) + assert shifted.converged + assert shifted.residual_rms < 1e-10 + np.testing.assert_array_equal(shifted.time, shifted_time) + for field in ( + "amplitude", + "decay", + "angular_frequency", + "phase", + "reconstruction", + ): + np.testing.assert_allclose( + getattr(shifted, field), getattr(base, field), atol=1e-10 + ) + + +def test_seed_from_residual_finds_the_missing_mode(): + fx = synthetic_modes_trace(256, duration=25.5) + weak = 0.45 * np.exp(-0.03 * fx.time) * np.cos(0.9 * fx.time - 0.75) + amp, dec, freq, _ = seed_from_residual(fx.time, weak) + assert abs(freq - 0.9) < 0.05 + assert dec == 0.0 and amp > 0 + + +def test_seed_from_residual_uses_samples_beyond_the_default_fft_length(): + time = np.arange(8192) * 0.1 + frequency = 2 * np.pi * 64 / (4096 * 0.1) + residual = np.cos(frequency * time) + amplitude, _, actual_frequency, _ = seed_from_residual(time, residual) + assert actual_frequency == pytest.approx(frequency) + assert amplitude == pytest.approx(1.0) + + residual[:4096] = 0.0 + amplitude, _, actual_frequency, _ = seed_from_residual(time, residual) + assert actual_frequency == pytest.approx(frequency) + assert amplitude == pytest.approx(0.5) + + +@pytest.mark.parametrize("phase", [1e10, 1e16]) +def test_large_initial_phase_is_reduced_before_evaluation(phase): + time = np.linspace(0, 10, 50) + wrapped = (phase + np.pi) % (2 * np.pi) - np.pi + trace = np.exp(-0.2 * time) * np.cos(2 * time + wrapped) + 0.1 + result = refine_modes(time, trace, [(1.0, 0.2, 2.0, phase)], 0.1) + assert result.converged + assert result.residual_rms < 1e-12 + reconstructed = result.constant + np.sum( + result.amplitude[:, None] + * np.exp(-result.decay[:, None] * time) + * np.cos( + result.angular_frequency[:, None] * time + result.phase[:, None] + ), + axis=0, + ) + np.testing.assert_allclose( + reconstructed, result.reconstruction, atol=1e-14 + ) + np.testing.assert_allclose(reconstructed, trace, atol=1e-12) + + +@pytest.mark.parametrize("invalid", [0.15, [0.15], np.ones((96, 1))]) +def test_refinement_rejects_trace_broadcasting(invalid): + time = np.linspace(0, 9.5, 96) + with pytest.raises(ValueError, match="matching finite one-dimensional"): + refine_modes(time, invalid, [(1.0, 0.1, 2.0, 0.0)], 0.0) + + +@pytest.mark.parametrize("field", ["time", "trace"]) +def test_refinement_rejects_complex_inputs_before_discarding_values(field): + fixture = synthetic_modes_trace(96) + inputs = {"time": fixture.time, "trace": fixture.trace} + inputs[field] = inputs[field].astype(complex) + 1j + with pytest.raises(ValueError, match="must be real"): + refine_modes( + **inputs, initial_modes=[(1.0, 0.1, 2.0, 0.0)], constant=0 + ) + with pytest.raises(ValueError, match="must be real"): + linear_prediction_refined(**inputs, n_components=6, n_modes=2) + + +def test_refined_estimator_keeps_the_weak_mode_under_noise(): + fx = synthetic_modes_trace(96) + sigma = float(np.std(fx.clean - fx.clean.mean())) / 100 + rng = np.random.default_rng(0) + lost = 0 + for _ in range(60): + y = fx.clean + rng.normal(0.0, sigma, size=fx.time.shape) + r = linear_prediction_refined(fx.time, y, 6, 2) + if np.min(np.abs(r.angular_frequency - 0.9)) > 0.05: + lost += 1 + assert lost == 0 + + +def test_refined_sweep_sits_near_the_cramer_rao_bound(): + sweep = validation_sweep( + snr_db=(30.0,), + trials=60, + seed=2, + estimator=lambda t, y, k: linear_prediction_refined(t, y, k, 2), + backend="cpu-refined", + ) + level = sweep.levels[0] + for mode in level.modes: + assert mode.loss_rate == 0.0 + assert 0.7 < mode.frequency_ratio < 1.4 + assert 0.7 < mode.decay_ratio < 1.4 + + +def test_refinement_rejects_empty_initial_modes(): + fx = synthetic_modes_trace(96) + with pytest.raises(ValueError): + refine_modes(fx.time, fx.trace, [], 0.0) + + +@pytest.mark.parametrize("n_modes", [0, -1, 1.5, True]) +def test_lp_refinement_rejects_invalid_mode_counts(n_modes): + fixture = synthetic_modes_trace(96) + with pytest.raises( + ValueError, match="n_modes must be a positive integer" + ): + linear_prediction_refined(fixture.time, fixture.trace, 6, n_modes) + + +def test_bounds_are_the_reference_for_the_refined_uncertainties(): + # the per-fit one-sigma from the Jacobian should agree with the bound + # to within the scatter of a single residual-variance estimate + fx = synthetic_modes_trace(96) + sigma = float(np.std(fx.clean - fx.clean.mean())) / 1000 + b = cramer_rao_bounds(fx.modes, fx.constant, fx.time, sigma) + y = fx.clean + np.random.default_rng(5).normal(0.0, sigma, fx.time.shape) + r = linear_prediction_refined(fx.time, y, 6, 2) + ratio = r.sigma_angular_frequency / b["angular_frequency"] + assert np.all(ratio > 0.6) and np.all(ratio < 1.6) + + +def test_rank_deficient_fit_reports_unavailable_uncertainty(): + fx = synthetic_modes_trace(96) + start = [(1.0, 0.09, 2.4, 0.3), (0.3, 0.09, 2.41, 0.3)] + theta = np.array([v for mode in start for v in mode] + [0.15]) + clean, _ = _model_and_jacobian(theta, fx.time, 2) + y = clean + np.random.default_rng(6).normal(0.0, 0.02, fx.time.size) + result = refine_modes(fx.time, y, start, 0.15) + assert result.converged + assert result.residual_rms < 0.025 + for field in ( + "sigma_amplitude", + "sigma_decay", + "sigma_angular_frequency", + "sigma_phase", + ): + assert np.all(np.isnan(getattr(result, field))) + + +@pytest.mark.parametrize("time_scale", [1e-9, 1e9]) +def test_identifiable_uncertainty_is_independent_of_time_units(time_scale): + fx = synthetic_modes_trace(96, noise_sigma=0.002, seed=5) + start = [ + (m.amplitude, m.decay, m.angular_frequency, m.phase) for m in fx.modes + ] + base = refine_modes(fx.time, fx.trace, start, fx.constant) + scaled_start = [ + (a, d / time_scale, w / time_scale, p) for a, d, w, p in start + ] + scaled = refine_modes( + fx.time * time_scale, fx.trace, scaled_start, fx.constant + ) + for field in ( + "sigma_amplitude", + "sigma_decay", + "sigma_angular_frequency", + "sigma_phase", + ): + actual = getattr(scaled, field) + if field in ("sigma_decay", "sigma_angular_frequency"): + actual = actual * time_scale + assert np.all(np.isfinite(actual)) + assert np.all(actual > 0) + np.testing.assert_allclose(actual, getattr(base, field), rtol=1e-4) diff --git a/tests/xray/test_mode_refinement_batched.py b/tests/xray/test_mode_refinement_batched.py new file mode 100644 index 00000000..649425d1 --- /dev/null +++ b/tests/xray/test_mode_refinement_batched.py @@ -0,0 +1,260 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +import json + +import numpy as np +import pytest + +from cuphoton.xray.mode_refinement import _model_and_jacobian, refine_modes +from cuphoton.xray.mode_refinement_batched import ( + _canonical, + benchmark_refinement, + model_and_jacobian_batched, + refine_modes_batched, +) +from cuphoton.xray.synthetic_validation import ( + modes_to_theta, + synthetic_modes_trace, +) + + +def test_batched_model_and_jacobian_match_the_serial_ones(): + fx = synthetic_modes_trace(96) + theta = modes_to_theta(fx.modes, fx.constant) + batch = np.stack([theta, theta * 1.01, theta * 0.98]) + model, jac = model_and_jacobian_batched(np, batch, fx.time, 2) + for i in range(3): + m, j = _model_and_jacobian(batch[i], fx.time, 2) + np.testing.assert_allclose(model[i], m, rtol=0, atol=1e-12) + np.testing.assert_allclose(jac[i], j, rtol=0, atol=1e-12) + + +def test_batched_large_phase_matches_wrapped_scalar_model(): + time = np.linspace(0, 10, 50) + start = np.array([[1.0, 0.2, 2.0, 1e16, 0.1]]) + wrapped = start.copy() + wrapped[0, 3] = (start[0, 3] + np.pi) % (2 * np.pi) - np.pi + trace, expected_jac = _model_and_jacobian(wrapped[0], time, 1) + model, jac = model_and_jacobian_batched(np, start, time, 1) + np.testing.assert_allclose(model[0], trace, atol=1e-14) + np.testing.assert_allclose(jac[0], expected_jac, atol=1e-14) + result = refine_modes_batched(np, time, trace[None, :], start, 1) + assert result.converged.tolist() == [True] + assert result.residual_rms[0] < 1e-14 + np.testing.assert_allclose(result.theta, wrapped, atol=1e-14) + + +def test_batched_refinement_matches_scipy_per_trace(): + fx = synthetic_modes_trace(96) + rng = np.random.default_rng(1) + traces = 12 + y = fx.clean[None, :] + rng.normal(0.0, 0.02, (traces, fx.time.size)) + truth = modes_to_theta(fx.modes, fx.constant) + theta0 = truth[None, :] * (1.0 + 0.05 * rng.standard_normal((traces, 1))) + res = refine_modes_batched(np, fx.time, y, theta0, 2) + assert bool(np.all(res.converged)) + got = _canonical(res.theta, 2) + for i in range(traces): + init = [tuple(theta0[i, 4 * k : 4 * k + 4]) for k in range(2)] + r = refine_modes(fx.time, y[i], init, float(theta0[i, -1])) + j = int(np.argmin(np.abs(r.angular_frequency - 2.4))) + assert abs(got[i, 2] - r.angular_frequency[j]) < 1e-7 + assert abs(got[i, 1] - r.decay[j]) < 1e-7 + assert abs(got[i, -1] - r.constant) < 1e-7 + + +def test_batched_refinement_rejects_bad_shapes(): + fx = synthetic_modes_trace(96) + with pytest.raises(ValueError): + refine_modes_batched(np, fx.time, fx.trace, np.zeros((1, 9)), 2) + with pytest.raises(ValueError): + refine_modes_batched( + np, fx.time, fx.trace[None, :], np.zeros((1, 7)), 2 + ) + + +@pytest.mark.parametrize("time_size,trace_size", [(1, 96), (96, 1), (0, 0)]) +def test_batched_refinement_rejects_broadcast_or_empty_samples( + time_size, trace_size +): + with pytest.raises(ValueError, match="time|sample"): + refine_modes_batched( + np, + np.arange(time_size), + np.zeros((1, trace_size)), + np.zeros((1, 9)), + 2, + ) + + +@pytest.mark.parametrize("input_name", ["time", "traces", "theta0"]) +@pytest.mark.parametrize("invalid_value", [np.nan, 1j]) +def test_batched_refinement_rejects_invalid_input(input_name, invalid_value): + fx = synthetic_modes_trace(96) + inputs = { + "time": fx.time.copy(), + "traces": fx.trace[None, :].copy(), + "theta0": modes_to_theta(fx.modes, fx.constant)[None, :], + } + inputs[input_name] = inputs[input_name].astype(type(invalid_value)) + inputs[input_name].flat[0] = invalid_value + with pytest.raises(ValueError, match=input_name): + refine_modes_batched(np, **inputs, n_modes=2) + + +@pytest.mark.parametrize( + "options", [{"tol": np.inf}, {"tol": 0}, {"max_iter": -1}, {"n_modes": 0}] +) +def test_batched_refinement_rejects_invalid_stopping_options(options): + fx = synthetic_modes_trace(96) + arguments = {"n_modes": 2, **options} + with pytest.raises(ValueError, match=next(iter(options))): + refine_modes_batched( + np, + fx.time, + fx.trace[None, :], + modes_to_theta(fx.modes, fx.constant)[None, :], + **arguments, + ) + + +def test_batched_refinement_accepts_an_exact_start_without_improvement(): + fx = synthetic_modes_trace(96) + truth = modes_to_theta(fx.modes, fx.constant) + result = refine_modes_batched( + np, fx.time, fx.trace[None, :], truth[None, :], 2 + ) + assert result.converged.tolist() == [True] + assert result.iterations == 1 + assert result.residual_rms[0] < 1e-14 + + +def test_zero_amplitude_start_does_not_abort_other_rows(): + fx = synthetic_modes_trace(96) + truth = modes_to_theta(fx.modes, fx.constant) + starts = np.stack([truth, truth]) + starts[1, 4] = 0.0 + traces = np.stack([fx.trace, fx.trace]) + result = refine_modes_batched(np, fx.time, traces, starts, 2) + assert result.converged.tolist() == [True, True] + np.testing.assert_allclose( + result.theta, np.stack([truth, truth]), atol=1e-7 + ) + assert np.all(result.residual_rms < 1e-8) + + +@pytest.mark.parametrize("lambda0", [1e10, 1e30]) +def test_large_damping_does_not_make_a_nonstationary_start_converged(lambda0): + fx = synthetic_modes_trace(96) + start = modes_to_theta(fx.modes, fx.constant)[None, :] + start[0, -1] += 1.0 + result = refine_modes_batched( + np, fx.time, fx.trace[None, :], start, 2, lambda0=lambda0, max_iter=1 + ) + assert result.converged.tolist() == [False] + + +@pytest.mark.parametrize("time_scale", [1e-9, 1.0, 1e9]) +@pytest.mark.parametrize("backend", ["numpy", "cupy"]) +def test_batched_convergence_is_independent_of_time_units( + time_scale, backend +): + xp = np + if backend == "cupy": + xp = pytest.importorskip("cupy") + try: + xp.cuda.runtime.getDeviceCount() + except Exception: # pragma: no cover - no device + pytest.skip("no CUDA device") + fx = synthetic_modes_trace(96) + truth = modes_to_theta(fx.modes, fx.constant) + start = truth.copy() + start[[1, 2, 5, 6]] /= time_scale + start[[2, 6]] *= 1.05 + result = refine_modes_batched( + xp, fx.time * time_scale, fx.trace[None, :], start[None, :], 2 + ) + assert result.converged.tolist() == [True] + assert result.residual_rms[0] < 1e-8 + recovered = result.theta[0].copy() + if backend == "cupy": + recovered = xp.asnumpy(recovered) + recovered[[1, 2, 5, 6]] *= time_scale + np.testing.assert_allclose(recovered, truth, atol=1e-8) + + +@pytest.mark.parametrize("lambda0", [0.0, -1.0, np.nan, np.inf]) +def test_batched_refinement_requires_positive_finite_damping(lambda0): + fx = synthetic_modes_trace(96) + truth = modes_to_theta(fx.modes, fx.constant)[None, :] + with pytest.raises(ValueError, match="lambda0"): + refine_modes_batched( + np, fx.time, fx.trace[None, :], truth, 2, lambda0=lambda0 + ) + + +def test_benchmark_runs_on_cpu_and_reports_agreement(): + r = benchmark_refinement(traces=8, run_gpu=False) + assert r.traces == 8 + assert r.max_abs_theta_diff_numpy < 1e-6 + assert r.cupy_batched_s is None and r.gpu_error is None + + +def test_benchmark_reports_unconverged_fits_and_residuals(monkeypatch): + from cuphoton.xray import mode_refinement_batched + + def one_iteration(*args, **kwargs): + return refine_modes_batched(*args, **kwargs, max_iter=1) + + monkeypatch.setattr( + mode_refinement_batched, "refine_modes_batched", one_iteration + ) + result = benchmark_refinement(traces=2, repeat=1, run_gpu=False) + payload = result.to_dict() + assert payload["cpu_serial_converged"] == 2 + assert payload["numpy_batched_converged"] == 0 + assert payload["cupy_batched_converged"] is None + assert ( + payload["numpy_batched_max_residual_rms"] + > payload["cpu_serial_max_residual_rms"] + ) + json.dumps(payload, allow_nan=False) + assert "numpy_batched_converged=0/2" in str(result) + assert "gpu_status=skipped" in str(result) + + +@pytest.mark.parametrize("options", [{"traces": 0}, {"n_modes": 3}]) +def test_benchmark_rejects_empty_or_unavailable_fixtures(options): + with pytest.raises(ValueError, match=next(iter(options))): + benchmark_refinement(**options, run_gpu=False) + + +@pytest.mark.parametrize("time_scale", [1e-9, 1.0, 1e9]) +def test_batched_refinement_on_cupy_matches_numpy(time_scale): + cp = pytest.importorskip("cupy") + try: + cp.cuda.runtime.getDeviceCount() + except Exception: # pragma: no cover - no device + pytest.skip("no CUDA device") + fx = synthetic_modes_trace(96) + rng = np.random.default_rng(2) + y = fx.clean[None, :] + rng.normal(0.0, 0.02, (32, fx.time.size)) + truth = modes_to_theta(fx.modes, fx.constant) + theta0 = truth[None, :] * (1.0 + 0.05 * rng.standard_normal((32, 1))) + theta0[:, [1, 2, 5, 6]] /= time_scale + # Resolve noisy optima to a practical gradient tolerance; at 1e-9, + # rounding of cost differences can prevent further accepted steps. + a = refine_modes_batched(np, fx.time * time_scale, y, theta0, 2, tol=1e-8) + b = refine_modes_batched(cp, fx.time * time_scale, y, theta0, 2, tol=1e-8) + got = cp.asnumpy(b.theta) + expected = a.theta.copy() + got[:, [1, 2, 5, 6]] *= time_scale + expected[:, [1, 2, 5, 6]] *= time_scale + assert np.all(a.converged) + assert np.all(cp.asnumpy(b.converged)) + np.testing.assert_allclose(got, expected, rtol=0, atol=1e-7) + np.testing.assert_allclose( + cp.asnumpy(b.residual_rms), a.residual_rms, rtol=1e-10 + ) diff --git a/tests/xray/test_mode_refinement_viz.py b/tests/xray/test_mode_refinement_viz.py new file mode 100644 index 00000000..38b3710e --- /dev/null +++ b/tests/xray/test_mode_refinement_viz.py @@ -0,0 +1,105 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +import re +from dataclasses import replace + +import numpy as np +import pytest + +from cuphoton.xray import mode_refinement_viz as viz +from cuphoton.xray.synthetic_validation import ( + distort_trace, + synthetic_modes_trace, +) + + +def test_modes_review_writes_offline_html_and_reports_fit_status(tmp_path): + pytest.importorskip("bokeh") + output = tmp_path / "review" / "modes.html" + + report = viz.write_modes_review(output, snr_db=(40.0, 20.0), samples=48) + + assert report.html_path == output + assert report.trace_count == 2 + assert report.converged_count == 2 + assert report.to_dict() == { + "html_path": str(output), + "trace_count": 2, + "converged_count": 2, + } + assert "refinement_converged=2/2" in str(report) + html = output.read_text(encoding="utf-8") + assert "Bokeh" in html + assert "noise-free signal" in html + assert "residual/noise" in html + assert re.search(r"]*\bsrc\s*=", html) is None + + +@pytest.mark.parametrize("distortion", [("chirp", 0.3), ("glitch", 1.0)]) +def test_modes_review_plots_distorted_truth_and_failed_convergence( + tmp_path, monkeypatch, distortion +): + pytest.importorskip("bokeh") + from bokeh.models import Plot + + captured = {} + + def capture_html(layout, resources, title): + captured["layout"] = layout + return "" + + original_refine = viz.linear_prediction_refined + + def unconverged(*args, **kwargs): + return replace(original_refine(*args, **kwargs), converged=False) + + monkeypatch.setattr("bokeh.embed.file_html", capture_html) + monkeypatch.setattr(viz, "linear_prediction_refined", unconverged) + report = viz.write_modes_review( + tmp_path / "modes.html", + snr_db=(20.0,), + distortion=distortion, + samples=48, + ) + + plots = list(captured["layout"].select({"type": Plot})) + trace_plot = next(p for p in plots if "residual/noise" in p.title.text) + assert "not converged" in trace_plot.title.text + assert report.converged_count == 0 + truth = next( + item.renderers[0] + for legend in trace_plot.legend + for item in legend.items + if item.label.value == "noise-free signal" + ) + fixture = synthetic_modes_trace(48) + expected = distort_trace( + fixture.time, fixture.modes, fixture.constant, *distortion + ) + np.testing.assert_array_equal(truth.data_source.data["x"], fixture.time) + np.testing.assert_array_equal(truth.data_source.data["y"], expected) + assert not np.array_equal(expected, fixture.clean) + assert any("signal distortion" in p.title.text for p in plots) + + +@pytest.mark.parametrize( + ("kwargs", "message"), + [ + ({"snr_db": ()}, "at least one"), + ({"snr_db": (np.nan,)}, "must be finite"), + ({"snr_db": (10000.0,)}, "positive noise"), + ({"snr_db": (-10000.0,)}, "positive noise"), + ({"samples": 8}, "at least 16"), + ], +) +def test_modes_review_rejects_invalid_fixture_before_writing( + tmp_path, kwargs, message +): + output = tmp_path / "modes.html" + + with pytest.raises(ValueError, match=message): + viz.write_modes_review(output, **kwargs) + + assert not output.exists() diff --git a/tests/xray/test_subspace.py b/tests/xray/test_subspace.py index 4445a0a3..56ce18ee 100644 --- a/tests/xray/test_subspace.py +++ b/tests/xray/test_subspace.py @@ -12,6 +12,7 @@ from cuphoton.xray.linear_prediction import synthetic_trace from cuphoton.xray.subspace import ( compare_subspace_methods, + compare_subspace_methods_batch, esprit_roots, matrix_pencil_roots, ) @@ -120,6 +121,104 @@ def test_compare_subspace_methods_can_use_partial_svd(): assert len(method.roots_real) == 5 +def test_subspace_batch_preserves_trace_results_and_source_order(): + time, trace = synthetic_trace(samples=64) + trace_rows = np.stack((trace, trace + 0.1 * np.cos(3.0 * time))) + sources = ( + {"kind": "trace-npz", "path": "trace-y2.npz"}, + { + "kind": "trace-npz", + "path": "trace-y10.npz", + "summary": {"detector_row": 10}, + "time_matches_first": True, + }, + ) + options = { + "model_order": 4, + "components": 6, + "methods": ("esprit", "matrix-pencil"), + "svd_backends": ("randomized", "full"), + "pencil_rows": 24, + "randomized_oversamples": 3, + "power_iterations": 2, + "random_seed": 9, + } + batch = compare_subspace_methods_batch( + time, trace_rows, sources, **options + ) + payload = batch.to_dict() + + assert payload["source"]["traces"] == sources + assert payload["trace_count"] == 2 + for index, (trace_result, trace_payload) in enumerate( + zip(batch.traces, payload["traces"], strict=True) + ): + expected = compare_subspace_methods( + time, trace_rows[index], **options + ) + assert trace_payload["trace_index"] == index + assert trace_payload["source"] == sources[index] + assert trace_payload["svd_backends"] == options["svd_backends"] + assert trace_result.baseline_chi2 == expected.baseline_chi2 + for result, expected_method in zip( + trace_result.methods, expected.methods, strict=True + ): + assert result.pencil_rows == 24 + assert result.rms_residual == pytest.approx( + expected_method.rms_residual + ) + np.testing.assert_allclose( + result.roots_real, expected_method.roots_real + ) + np.testing.assert_allclose( + result.roots_imag, expected_method.roots_imag + ) + + assert [ + (summary.method, summary.svd_backend) + for summary in batch.method_summary + ] == [ + ("esprit", "randomized"), + ("esprit", "full"), + ("matrix-pencil", "randomized"), + ("matrix-pencil", "full"), + ] + for summary in batch.method_summary: + methods = [ + method + for result in batch.traces + for method in result.methods + if (method.method, method.svd_backend) + == (summary.method, summary.svd_backend) + ] + assert summary.trace_count == 2 + assert summary.rms_residual_min == min( + method.rms_residual for method in methods + ) + assert summary.rms_residual_max == max( + method.rms_residual for method in methods + ) + assert summary.max_abs_reconstruction_diff_max == max( + method.max_abs_reconstruction_diff for method in methods + ) + assert summary.elapsed_s_total == pytest.approx( + sum(method.elapsed_s for method in methods) + ) + + assert payload["baseline_chi2_min"] == min( + result.baseline_chi2 for result in batch.traces + ) + assert payload["baseline_chi2_max"] == max( + result.baseline_chi2 for result in batch.traces + ) + assert payload["baseline_rms_residual_min"] == min( + result.baseline_rms_residual for result in batch.traces + ) + assert payload["baseline_rms_residual_max"] == max( + result.baseline_rms_residual for result in batch.traces + ) + + def test_matrix_pencil_rejects_rankless_trace(): with pytest.raises(ValueError, match="rank"): matrix_pencil_roots(np.zeros(32), model_order=5) diff --git a/tests/xray/test_synthetic_validation.py b/tests/xray/test_synthetic_validation.py new file mode 100644 index 00000000..c71a8f51 --- /dev/null +++ b/tests/xray/test_synthetic_validation.py @@ -0,0 +1,674 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +import json +from dataclasses import replace +from types import SimpleNamespace + +import numpy as np +import pytest + +from cuphoton.xray.linear_prediction import ( + linear_prediction_numpy, + synthetic_trace, +) +from cuphoton.xray.synthetic_validation import ( + PARAMETERS, + DampedMode, + _modes_jacobian, + build_summary, + cramer_rao_bounds, + distort_trace, + match_modes, + modes_model, + modes_to_theta, + synthetic_modes_trace, + validate_linear_prediction, + validation_sweep, + write_validation_run, +) + + +def test_fixture_reproduces_synthetic_trace_at_zero_noise(): + time, trace = synthetic_trace(96) + fx = synthetic_modes_trace(96) + np.testing.assert_allclose(fx.time, time) + np.testing.assert_allclose(fx.trace, trace) + assert fx.noise_sigma == 0.0 + + +def test_fixture_noise_is_reproducible_and_has_the_requested_scale(): + a = synthetic_modes_trace(256, noise_sigma=0.05, seed=7) + b = synthetic_modes_trace(256, noise_sigma=0.05, seed=7) + np.testing.assert_array_equal(a.trace, b.trace) + resid = a.trace - a.clean + assert abs(resid.std() - 0.05) < 0.01 + + +def test_estimator_recovers_known_modes_exactly_at_zero_noise(): + fx = synthetic_modes_trace(96) + r = linear_prediction_numpy(fx.time, fx.trace, 6) + hit = match_modes( + fx.modes, + np.asarray(r.angular_frequency), + np.asarray(r.decay), + tolerance=0.05, + ) + assert all(h is not None for h in hit) + for m, h in zip(fx.modes, hit, strict=True): + assert abs(float(r.angular_frequency[h]) - m.angular_frequency) < 1e-6 + assert abs(float(r.decay[h]) - m.decay) < 1e-6 + + +@pytest.mark.parametrize("sigma", [0.01, 1e-160]) +def test_cramer_rao_bound_scales_with_noise_and_shrinks_with_trace_length( + sigma, +): + fx = synthetic_modes_trace(96) + b1 = cramer_rao_bounds(fx.modes, fx.constant, fx.time, sigma) + b2 = cramer_rao_bounds(fx.modes, fx.constant, fx.time, 2 * sigma) + np.testing.assert_allclose( + b2["angular_frequency"], + 2 * b1["angular_frequency"], + rtol=1e-6, + atol=0, + ) + longer = synthetic_modes_trace(192, duration=19.0) + b3 = cramer_rao_bounds(longer.modes, longer.constant, longer.time, sigma) + assert np.all(b3["angular_frequency"] < b1["angular_frequency"]) + + +def test_single_mode_bound_matches_the_undamped_closed_form(): + # Real sinusoid A cos(w t + phi) in white Gaussian noise, unknown A, w and + # phi, samples at t = n dt for n = 0..N-1 (Kay, Fundamentals of + # Statistical Signal Processing: Estimation Theory, 1993, sinusoidal + # parameter estimation): var(w) >= 24 sigma^2 / (A^2 dt^2 N (N^2 - 1)). + # The complex-exponential form has 12 in place of 24. The numerical + # bound also frees the decay and the constant, so it may sit slightly + # above the closed form, never below it. + mode = (DampedMode(1.0, 0.0, 2.0, 0.3),) + fx = synthetic_modes_trace(256, modes=mode, constant=0.0, duration=25.5) + dt = fx.time[1] - fx.time[0] + n = fx.time.size + sigma = 0.01 + closed = np.sqrt(24 * sigma**2 / (1.0**2 * dt**2 * n * (n**2 - 1))) + numerical = cramer_rao_bounds(mode, 0.0, fx.time, sigma)[ + "angular_frequency" + ][0] + assert numerical >= closed * 0.999 + assert numerical < closed * 1.03 + + +def test_integer_mode_parameters_have_the_same_bounds_as_floats(): + time = np.linspace(0, 25.5, 256) + actual = cramer_rao_bounds((DampedMode(1, 0, 2, 0),), 0, time, 0.01) + expected = cramer_rao_bounds( + (DampedMode(1.0, 0.0, 2.0, 0.0),), 0.0, time, 0.01 + ) + for name in (*PARAMETERS, "constant"): + np.testing.assert_allclose(actual[name], expected[name]) + assert np.all(np.isfinite(actual[name])) + + +def test_analytic_jacobian_and_bounds_match_finite_differences(): + fixture = synthetic_modes_trace() + theta = modes_to_theta(fixture.modes, fixture.constant) + n_modes = len(fixture.modes) + step = 1e-6 + numerical = np.column_stack( + [ + ( + modes_model(theta + delta, fixture.time, n_modes) + - modes_model(theta - delta, fixture.time, n_modes) + ) + / (2 * step) + for delta in step * np.eye(theta.size) + ] + ) + np.testing.assert_allclose( + _modes_jacobian(theta, fixture.time, n_modes), + numerical, + rtol=1e-7, + atol=3e-10, + ) + expected = 0.01 * np.sqrt(np.diag(np.linalg.inv(numerical.T @ numerical))) + actual = cramer_rao_bounds( + fixture.modes, fixture.constant, fixture.time, 0.01 + ) + for index, name in enumerate(PARAMETERS): + np.testing.assert_allclose( + actual[name], expected[index:-1:4], rtol=1e-8 + ) + np.testing.assert_allclose(actual["constant"], expected[-1], rtol=1e-8) + + +@pytest.mark.parametrize("time_scale", [1e-9, 1.0, 1e9]) +def test_bounds_resolve_nearby_modes_in_different_time_units(time_scale): + modes = (DampedMode(1, 0.01, 2, 0.3), DampedMode(0.5, 0.01, 2.02, 1.5)) + time = np.linspace(0, 100, 512) + expected = cramer_rao_bounds(modes, 0.15, time, 0.01) + scaled = tuple( + replace( + m, + decay=m.decay / time_scale, + angular_frequency=m.angular_frequency / time_scale, + ) + for m in modes + ) + actual = cramer_rao_bounds(scaled, 0.15, time * time_scale, 0.01) + for name in (*PARAMETERS, "constant"): + factor = time_scale if name in ("decay", "angular_frequency") else 1 + np.testing.assert_allclose( + actual[name] * factor, expected[name], rtol=1e-8 + ) + assert np.all(np.isfinite(actual[name])) + assert np.all(actual[name] > 0) + + +@pytest.mark.parametrize( + "modes", + [ + (DampedMode(0.0, 0.1, 2.0, 0.3),), + (DampedMode(1.0, 0.1, 2.0, 0.3),) * 2, + (DampedMode(1.0, 0.1, 2.0, 0.3), DampedMode(0.5, 0.1, 2.0, 1.5)), + ], + ids=[ + "zero-amplitude", + "duplicate-modes", + "coincident-phase-distinct-modes", + ], +) +def test_bounds_reject_unidentifiable_modes(modes): + with pytest.raises(ValueError, match="Cramer-Rao bounds"): + cramer_rao_bounds(modes, 0.0, np.linspace(0.0, 9.5, 96), 0.01) + + +def test_bounds_reject_more_parameters_than_samples(): + fixture = synthetic_modes_trace() + with pytest.raises(ValueError, match="Cramer-Rao bounds"): + cramer_rao_bounds( + fixture.modes, fixture.constant, fixture.time[:8], 0.01 + ) + + +@pytest.mark.parametrize("start", [0.0, 3.0]) +def test_chirp_integrates_the_requested_frequency_ramp(start): + time = np.linspace(start, start + 10.0, 256) + # Integrating omega(t) = 2 + .04 * (t - start) gives this phase. + expected = np.cos(2 * time + 0.02 * (time - start) ** 2 + 0.3) + actual = distort_trace(time, (DampedMode(1, 0, 2, 0.3),), 0, "chirp", 0.2) + np.testing.assert_allclose(actual, expected) + + +@pytest.mark.parametrize("start", [0.0, 3.0]) +def test_gaussian_envelope_starts_at_the_first_sample(start): + time = np.linspace(start, start + 10.0, 256) + expected = np.exp(-(((time - start) / 5.0) ** 2)) * np.cos(2 * time + 0.3) + actual = distort_trace( + time, (DampedMode(1, 0.1, 2, 0.3),), 0, "gaussian_envelope", 5.0 + ) + np.testing.assert_allclose(actual, expected) + + +@pytest.mark.parametrize("frequency_sign", [1, -1]) +def test_matching_recovers_close_modes_independent_of_truth_order( + frequency_sign, +): + modes = (DampedMode(1, 0, frequency_sign * 1.0), DampedMode(1, 0, 1.15)) + fitted = np.array([1.05, frequency_sign * 0.85]) + decay = np.zeros(2) + assert match_modes(modes, fitted, decay, tolerance=0.2) == [1, 0] + assert match_modes(modes[::-1], fitted, decay, tolerance=0.2) == [0, 1] + assert match_modes(modes, fitted[:1], decay[:1], tolerance=0.2) == [ + 0, + None, + ] + assert match_modes( + modes, np.array([np.nan]), decay[:1], tolerance=0.2 + ) == [None, None] + + +@pytest.mark.parametrize( + "truth_frequency,truth_decay,fitted_frequency,fitted_decay,tolerance,expected", + [ + # Different frequencies can also have equal total frequency error. + ([1, 2], [0.1, 0.8], [3, 4], [0.8, 0.1], 3, [1, 0]), + # Include fitted modes left unused by the primary assignment. + ([2], [0.1], [1, 3], [0.8, 0.1], 1, [1]), + # A partial recovery goes to the true mode with the matching decay. + ([2, 2], [0.1, 0.8], [2], [0.8], 0, [None, 0]), + # Better decay must not sacrifice an otherwise recoverable mode. + ([1, 1.15], [0, 100], [1.05, 0.85], [0, 100], 0.2, [1, 0]), + # Nor may it outweigh even one representable step in frequency error. + ([1], [0], [np.nextafter(1.0, 2.0), 1], [0, 1e100], 1, [1]), + # Rounded floating-point totals hide this genuine frequency advantage. + ( + [1, 2, 10], + [0, 1, 2], + [1.5, np.nextafter(1.5, 2.0), 1000], + [1, 0, 2], + 1000, + [0, 1, 2], + ), + ], +) +def test_matching_uses_decay_only_after_recovery_and_frequency( + truth_frequency, + truth_decay, + fitted_frequency, + fitted_decay, + tolerance, + expected, +): + modes = tuple( + DampedMode(1, d, w) + for w, d in zip(truth_frequency, truth_decay, strict=True) + ) + for reverse in (False, True): + order = slice(None, None, -1 if reverse else 1) + actual = match_modes( + modes, + np.array(fitted_frequency)[order], + np.array(fitted_decay)[order], + tolerance=tolerance, + ) + assert actual == [ + len(fitted_frequency) - 1 - h if reverse and h is not None else h + for h in expected + ] + + +@pytest.mark.parametrize("unavailable_decay", [np.nan, np.inf, -np.inf]) +def test_matching_nonfinite_decay_does_not_prevent_frequency_recovery( + unavailable_decay, +): + modes = (DampedMode(1, 0.1, 2),) + assert match_modes( + modes, np.array([2]), np.array([unavailable_decay]), tolerance=0 + ) == [0] + assert match_modes( + modes, + np.array([2, 2]), + np.array([unavailable_decay, 0.1]), + tolerance=0, + ) == [1] + + +@pytest.mark.parametrize("reverse", [False, True]) +def test_sweep_matches_exact_modes_with_shared_frequency_by_decay(reverse): + modes = (DampedMode(1, 0.05, 2, 0.3), DampedMode(0.4, 0.5, 2, 1)) + fitted_modes = modes[::-1] if reverse else modes + fixture = synthetic_modes_trace(modes=modes) + + def exact_estimator(t, y, k): + return SimpleNamespace( + **{ + name: np.array([getattr(m, name) for m in fitted_modes]) + for name in PARAMETERS + }, + reconstruction=fixture.clean, + ) + + level = validation_sweep( + modes=modes, + snr_db=(80,), + trials=4, + seed=314, + estimator=exact_estimator, + ).levels[0] + assert level.trials_successful == 4 + for mode in level.modes: + assert mode.recovered_trials == 4 + assert mode.loss_rate == 0 + for name in PARAMETERS: + stats = getattr(mode, name) + assert stats.bias == pytest.approx(0, abs=1e-14) + assert stats.rmse == pytest.approx(0, abs=1e-14) + + +def test_validation_sweep_reports_bound_ratios_and_loss_rates(): + sweep = validation_sweep(snr_db=(40.0, 20.0), trials=40, seed=1) + assert len(sweep.levels) == 2 + strong = sweep.levels[0].modes[0] + w = strong.angular_frequency + assert w.crlb_std > 0 and w.std > 0 + assert w.crlb_variance == pytest.approx(w.crlb_std**2) + # A coarse plausibility check: finite-sample, recovered-trial scatter + # from a biased estimator need not exceed the unconditional CRLB. + assert 0.2 < w.std_over_crlb_std < 10.0 + assert abs(w.bias) < 0.01 + assert w.rmse >= abs(w.bias) + for name in PARAMETERS: + assert getattr(strong, name).truth == getattr(sweep.modes[0], name) + assert abs(strong.phase.bias) < 0.05 + for level in sweep.levels: + assert level.trials_attempted == 40 + assert level.trials_successful + level.trials_failed == 40 + assert level.estimator_errors == 0 + for m in level.modes: + assert 0.0 <= m.loss_rate <= 1.0 + assert m.recovered_trials == round(40 * (1 - m.loss_rate)) + d = sweep.to_dict() + assert d["levels"][0]["snr_db"] == 40.0 + + +def test_estimator_errors_count_as_failed_trials(): + def broken(t, y, k): + raise RuntimeError("no fit") + + sweep = validation_sweep(snr_db=(30.0,), trials=3, estimator=broken) + level = sweep.levels[0] + assert level.estimator_errors == 3 + assert level.trials_failed == 3 and level.trials_successful == 0 + assert all(m.loss_rate == 1.0 for m in level.modes) + assert np.isnan(level.residual_ratio) + + +@pytest.mark.parametrize("amplitude", [-1.0, 1.0]) +@pytest.mark.parametrize("phase", [0.3, 0.3 + 4 * np.pi]) +@pytest.mark.parametrize("frequency_sign", [1, -1]) +def test_sweep_uses_the_same_amplitude_phase_convention_for_truth_and_fit( + amplitude, phase, frequency_sign +): + mode = DampedMode( + amplitude, 0.1, frequency_sign * 2.0, frequency_sign * phase + ) + fixture = synthetic_modes_trace(modes=(mode,)) + + def exact_estimator(t, y, k): + return SimpleNamespace( + amplitude=np.array([mode.amplitude]), + decay=np.array([mode.decay]), + angular_frequency=np.array([-mode.angular_frequency]), + phase=np.array([-mode.phase]), + reconstruction=fixture.clean, + ) + + sweep = validation_sweep( + modes=(mode,), + snr_db=(30.0,), + trials=2, + n_components=1, + estimator=exact_estimator, + ) + result = sweep.levels[0].modes[0] + assert result.recovered_trials == 2 + for name in PARAMETERS: + stats = getattr(result, name) + assert stats.bias == pytest.approx(0.0, abs=1e-14) + assert stats.rmse == pytest.approx(0.0, abs=1e-14) + assert stats.truth == getattr(sweep.modes[0], name) + + canonical = synthetic_modes_trace(modes=sweep.modes) + np.testing.assert_allclose(canonical.clean, fixture.clean, atol=1e-14) + summary = build_summary([sweep]) + truth = summary["config"]["truth"]["modes"][0] + assert truth["amplitude"] == 1.0 + assert truth["angular_frequency"] == 2.0 + expected_phase = 0.3 - np.pi if amplitude < 0 else 0.3 + assert truth["phase"] == pytest.approx(expected_phase) + for name in PARAMETERS: + assert summary["results"][0]["modes"][0][name]["truth"] == truth[name] + + +def test_equivalent_frequency_signs_have_the_same_mode_recovery(): + positive = DampedMode(1, 0.1, 2, 0.3) + negative = replace(positive, angular_frequency=-2, phase=-0.3) + sweeps = [ + validation_sweep( + modes=(mode,), snr_db=(80,), trials=2, n_components=4 + ) + for mode in (positive, negative) + ] + assert sweeps[0].levels == sweeps[1].levels + assert sweeps[0].modes == sweeps[1].modes + assert sweeps[0].levels[0].modes[0].recovered_trials == 2 + + +def test_summary_schema_and_run_artifacts(tmp_path): + sweep = validation_sweep(snr_db=(30.0,), trials=6, seed=3) + summary = build_summary([sweep]) + assert summary["schema_version"] == 1 + assert summary["command"] == "linear-prediction-validate" + config = summary["config"] + assert config["truth"]["modes"][0]["angular_frequency"] == 2.4 + assert config["sampling"]["samples"] == 96 + assert config["seed"] == 3 and config["trials_per_level"] == 6 + assert config["n_components"] == 6 + assert config["mode_matching"]["tolerance_rad_per_time"] == 0.3 + assert summary["runtime"]["backend"] == "cpu" + assert summary["runtime"]["numpy_version"] + result = summary["results"][0] + assert result["trials"]["attempted"] == 6 + mode = result["modes"][0] + assert mode["statistics_over"] == "recovered_trials" + for name in PARAMETERS: + assert set(mode[name]) == { + "truth", + "bias", + "std", + "rmse", + "crlb_std", + "crlb_variance", + "std_over_crlb_std", + } + out = tmp_path / "run" + written = write_validation_run(out, [sweep]) + assert (out / "summary.json").exists() + assert written["artifacts"]["summary"] == "summary.json" + figure = written["artifacts"]["figure"] + assert figure is None or (out / figure).exists() + with pytest.raises(ValueError): + build_summary([]) + + +@pytest.mark.parametrize("write_artifacts", [False, True]) +def test_validation_report_summarizes_the_run(tmp_path, write_artifacts): + out = tmp_path / "run" if write_artifacts else None + report = validate_linear_prediction( + samples=64, + snr_db=(30.0,), + trials=1, + n_components=4, + seed=7, + distortion=("chirp", 0.05), + output_dir=out, + ) + summary = report.to_dict() + assert json.loads(json.dumps(summary, allow_nan=False)) == summary + config = summary["config"] + assert config["sampling"]["samples"] == 64 + assert config["seed"] == 7 and config["trials_per_level"] == 1 + assert config["n_components"] == 4 + assert config["distortions"] == ["chirp:0.05"] + assert report.sweeps[0].levels[0].snr_db == 30.0 + text = str(report) + assert text.startswith("samples=64\nsignal_rms=") + assert "estimator=cpu\ndistortion=chirp:0.05\nsnr_db=30 " in text + if out is not None: + assert json.loads((out / "summary.json").read_text()) == summary + assert "summary=summary.json" in text + for name in summary["artifacts"].values(): + assert name is None or (out / name).is_file() + else: + assert summary["artifacts"] == {} + assert not (tmp_path / "run").exists() + + +@pytest.mark.parametrize("trials", [1, 2]) +def test_small_sweeps_write_strict_json_and_two_trials_have_std( + tmp_path, trials +): + sweep = validation_sweep(snr_db=(30.0,), trials=trials, seed=3) + summary = write_validation_run(tmp_path, [sweep]) + encoded = json.dumps(summary, allow_nan=False) + assert json.loads(encoded) == json.loads( + (tmp_path / "summary.json").read_text() + ) + strong = summary["results"][0]["modes"][0] + assert strong["recovered_trials"] == trials + assert (strong["angular_frequency"]["std"] is None) == (trials == 1) + + +def test_failed_sweep_statistics_are_null_in_summary(): + def broken(t, y, k): + raise RuntimeError("no fit") + + sweep = validation_sweep(snr_db=(30.0,), trials=2, estimator=broken) + summary = build_summary([sweep]) + json.dumps(summary, allow_nan=False) + result = summary["results"][0] + assert result["residual_rms_over_sigma_median"] is None + assert result["modes"][0]["angular_frequency"]["bias"] is None + + +@pytest.mark.parametrize( + "change", [{"samples": 192}, {"seed": 77}, {"match_tolerance": 0.1}] +) +def test_summary_rejects_incompatible_sweep_provenance(change): + sweep = validation_sweep(snr_db=(30.0,), trials=2) + with pytest.raises(ValueError, match="summary sweeps must share"): + build_summary([sweep, replace(sweep, **change)]) + + +def test_summary_preserves_distortion_signal_scale(): + sweep = validation_sweep(snr_db=(30.0,), trials=2) + distorted = replace( + sweep, backend="other", distortion=("clip", 0.5), signal_rms=0.1 + ) + summary = build_summary([sweep, distorted]) + assert summary["results"][0]["signal_rms"] == sweep.signal_rms + assert summary["results"][1]["signal_rms"] == 0.1 + + +def test_invalid_inputs_are_rejected(): + fx = synthetic_modes_trace(96) + with pytest.raises(ValueError): + cramer_rao_bounds(fx.modes, fx.constant, fx.time, 0.0) + with pytest.raises(ValueError): + validation_sweep(trials=0, snr_db=(40.0,)) + with pytest.raises(ValueError): + validation_sweep(trials=5, snr_db=()) + with pytest.raises(ValueError): + synthetic_modes_trace(8) + + +@pytest.mark.parametrize("sigma", [-0.1, np.nan, np.inf, -np.inf]) +def test_fixture_rejects_invalid_noise_sigma(sigma): + with pytest.raises(ValueError, match="noise_sigma"): + synthetic_modes_trace(noise_sigma=sigma) + + +@pytest.mark.parametrize("sigma", [np.nan, np.inf, -0.1, 1e-200, 1e200]) +def test_bounds_reject_invalid_noise_scales_without_runtime_warnings(sigma): + fx = synthetic_modes_trace() + with ( + np.errstate(all="raise"), + pytest.raises(ValueError, match="noise_sigma"), + ): + cramer_rao_bounds(fx.modes, fx.constant, fx.time, sigma) + + +@pytest.mark.parametrize( + "kwargs, message", + [ + ({"snr_db": (30.0, np.nan)}, "snr_db"), + ({"snr_db": (30.0, np.inf)}, "snr_db"), + ({"snr_db": (30.0, -np.inf)}, "snr_db"), + ({"snr_db": (30.0, 1e308)}, "noise_sigma"), + ({"snr_db": (30.0, -1e308)}, "noise_sigma"), + ({"snr_db": (30.0, 4000.0)}, "noise_sigma"), + ({"snr_db": (30.0, -4000.0)}, "noise_sigma"), + ({"n_components": 0}, "n_components"), + ({"distortion": ("gaussian_envelope", 0.0)}, "must be positive"), + ({"distortion": ("gaussian_envelope", -1.0)}, "must be positive"), + ({"distortion": ("glitch", np.inf)}, "must be finite"), + ({"match_tolerance": np.nan}, "match_tolerance"), + ({"duration": 0.0}, "duration"), + ], +) +def test_invalid_sweep_input_is_rejected_before_any_estimator_calls( + kwargs, message +): + calls = [] + + def estimator(*args): + calls.append(args) + raise RuntimeError("invalid input must not reach the estimator") + + with np.errstate(all="raise"), pytest.raises(ValueError, match=message): + validation_sweep(trials=1, estimator=estimator, **kwargs) + assert not calls + + +def test_undistorted_residual_ratio_is_near_one(): + sweep = validation_sweep(snr_db=(30.0,), trials=30, seed=9) + assert 0.7 < sweep.levels[0].residual_ratio < 1.6 + + +def test_distortions_raise_the_residual_ratio(): + for kind, amount in [ + ("chirp", 0.2), + ("gaussian_envelope", 6.0), + ("baseline_drift", 0.3), + ("clip", 1.0), + ]: + sweep = validation_sweep( + snr_db=(30.0,), trials=20, seed=9, distortion=(kind, amount) + ) + assert sweep.levels[0].residual_ratio > 4.0, kind + assert sweep.distortion == (kind, amount) + + +def test_distort_trace_kinds_and_rejection(): + fx = synthetic_modes_trace(96) + for kind, amount in [ + ("chirp", 0.1), + ("gaussian_envelope", 5.0), + ("baseline_drift", 0.2), + ("clip", 1.0), + ("glitch", 0.5), + ]: + y = distort_trace(fx.time, fx.modes, fx.constant, kind, amount) + assert y.shape == fx.time.shape + assert not np.allclose(y, fx.clean), kind + with pytest.raises(ValueError): + distort_trace(fx.time, fx.modes, fx.constant, "square", 1.0) + + +def test_validation_reports_solver_failure_separately_from_mode_recovery(): + def unconverged(time, trace, components): + fixture = synthetic_modes_trace(time.size) + return SimpleNamespace( + amplitude=np.array([m.amplitude for m in fixture.modes]), + decay=np.array([m.decay for m in fixture.modes]), + angular_frequency=np.array( + [m.angular_frequency for m in fixture.modes] + ), + phase=np.array([m.phase for m in fixture.modes]), + reconstruction=fixture.clean, + converged=False, + ) + + sweep = validation_sweep(snr_db=(30,), trials=2, estimator=unconverged) + level = sweep.levels[0] + assert level.trials_successful == 2 + assert level.any_mode_lost_rate == 0 + result = build_summary([sweep])["results"][0] + assert result["solver_convergence"] == {"converged": 0, "unconverged": 2} + assert "solver_unconverged=2" in str(sweep) + + +def test_refined_validation_is_available_without_the_cli(tmp_path): + report = validate_linear_prediction( + snr_db=(30,), trials=2, refine=True, output_dir=tmp_path + ) + assert [s.backend for s in report.sweeps] == ["cpu", "cpu-refined"] + result = report.to_dict()["results"][1] + status = result["solver_convergence"] + assert status["converged"] + status["unconverged"] == 2 + assert ( + json.loads((tmp_path / "summary.json").read_text()) + == report.to_dict() + ) diff --git a/tests/xray/test_synthetic_validation_provenance.py b/tests/xray/test_synthetic_validation_provenance.py new file mode 100644 index 00000000..958416c6 --- /dev/null +++ b/tests/xray/test_synthetic_validation_provenance.py @@ -0,0 +1,115 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +import subprocess + +import pytest + +from cuphoton.xray import synthetic_validation + + +@pytest.fixture +def source_checkout(tmp_path, monkeypatch): + def git(*args): + return subprocess.run( + [ + "git", + "-c", + "user.name=cuPhoton tests", + "-c", + "user.email=test@nvidia.com", + "-c", + "commit.gpgsign=false", + "-c", + "core.hooksPath=/dev/null", + *args, + ], + cwd=tmp_path, + capture_output=True, + text=True, + check=True, + timeout=5, + ).stdout.strip() + + git("init", "-q") + source = tmp_path / "src/cuphoton/xray/synthetic_validation.py" + source.parent.mkdir(parents=True) + source.write_text("# source checkout\n") + (tmp_path / ".gitignore").write_text(".venv/\n") + git("add", "src", ".gitignore") + git("commit", "-qm", "Initial test source") + monkeypatch.setattr(synthetic_validation, "__file__", str(source)) + return source, git + + +def test_clean_tracked_source_reports_its_revision(source_checkout): + source, git = source_checkout + revision = git("rev-parse", "HEAD") + assert synthetic_validation._source_provenance() == (revision, False) + assert synthetic_validation.source_revision() == revision + # Generated run files do not change the tracked source provenance. + (source.parent / "summary.json").write_text("{}\n") + assert synthetic_validation._source_provenance() == (revision, False) + + +@pytest.mark.parametrize("staged", [False, True]) +def test_changed_tracked_source_is_dirty(source_checkout, staged): + source, git = source_checkout + source.write_text("# changed source\n") + if staged: + git("add", "src") + assert synthetic_validation._source_provenance() == ( + git("rev-parse", "HEAD"), + True, + ) + + +def test_dirty_tracks_changes_outside_the_module_directory(source_checkout): + source, git = source_checkout + checkout = source.parents[3] + (checkout / ".gitignore").write_text(".venv/\n*.json\n") + assert synthetic_validation._source_provenance() == ( + git("rev-parse", "HEAD"), + True, + ) + + +@pytest.mark.parametrize("directory", [".venv", "untracked"]) +def test_untracked_install_does_not_report_enclosing_checkout( + source_checkout, monkeypatch, directory +): + source, _ = source_checkout + installed = ( + source.parents[3] + / directory + / "lib/python3.12/site-packages/cuphoton/xray/synthetic_validation.py" + ) + installed.parent.mkdir(parents=True) + installed.write_text("# installed package\n") + monkeypatch.setattr(synthetic_validation, "__file__", str(installed)) + assert synthetic_validation._source_provenance() == (None, None) + assert synthetic_validation.source_revision() is None + + +def test_non_git_install_has_unknown_provenance(tmp_path, monkeypatch): + source = tmp_path / "synthetic_validation.py" + source.write_text("# installed package\n") + monkeypatch.setattr(synthetic_validation, "__file__", str(source)) + assert synthetic_validation._source_provenance() == (None, None) + + +def test_git_unavailable_has_unknown_provenance(source_checkout, monkeypatch): + monkeypatch.setenv("PATH", "") + assert synthetic_validation._source_provenance() == (None, None) + + +@pytest.mark.parametrize("dirty", [False, True]) +def test_summary_includes_tracked_source_provenance(source_checkout, dirty): + source, git = source_checkout + if dirty: + source.write_text("# changed source\n") + sweep = synthetic_validation.validation_sweep(snr_db=(30,), trials=2) + runtime = synthetic_validation.build_summary([sweep])["runtime"] + assert runtime["source_revision"] == git("rev-parse", "HEAD") + assert runtime["source_dirty"] is dirty diff --git a/tests/xray/test_synthetic_validation_viz.py b/tests/xray/test_synthetic_validation_viz.py new file mode 100644 index 00000000..fe7d0ca1 --- /dev/null +++ b/tests/xray/test_synthetic_validation_viz.py @@ -0,0 +1,79 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +import sys +from html.parser import HTMLParser + +import pytest + +from cuphoton.xray.synthetic_validation import ( + validation_sweep, + write_validation_figure, +) + + +class _ResourceParser(HTMLParser): + def __init__(self): + super().__init__() + self.resource_urls = [] + + def handle_starttag(self, tag, attrs): + resource_attr = {"script": "src", "link": "href"}.get(tag) + if resource_attr is not None: + url = dict(attrs).get(resource_attr) + if url is not None: + self.resource_urls.append(url) + + +@pytest.fixture +def sweep(): + return validation_sweep(snr_db=(30.0,), trials=2) + + +def test_validation_figure_embeds_resources_for_offline_use(tmp_path, sweep): + pytest.importorskip("bokeh") + path = tmp_path / "validation.html" + + assert write_validation_figure([sweep], path, title="Validation figure") + + html = path.read_text(encoding="utf-8") + resources = _ResourceParser() + resources.feed(html) + assert resources.resource_urls == [] + assert "/* BEGIN bokeh.min.js */" in html + assert "Validation figure" in html + assert "angular_frequency std (points) vs bound (dashed)" in html + assert "decay std (points) vs bound (dashed)" in html + assert "loss rate" in html + + +def test_validation_figure_preserves_bokeh_output_state( + tmp_path, sweep, monkeypatch +): + pytest.importorskip("bokeh") + from bokeh.io import state + + current = state.State() + monkeypatch.setattr(state, "_STATE", current) + current.output_file(tmp_path / "existing.html", title="Existing output") + existing_file = current.file + existing_document = current.document + + assert write_validation_figure( + [sweep], tmp_path / "validation.html", title="Validation figure" + ) + + assert current.file is existing_file + assert current.document is existing_document + assert current.document.roots == [] + + +def test_validation_figure_returns_false_without_bokeh(tmp_path, monkeypatch): + monkeypatch.setitem(sys.modules, "bokeh", None) + monkeypatch.setitem(sys.modules, "bokeh.embed", None) + monkeypatch.setitem(sys.modules, "bokeh.layouts", None) + path = tmp_path / "validation.html" + + assert not write_validation_figure([], path, title="Validation figure") + assert not path.exists()