diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 71fec582..21abb3cb 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 46e42a84..c0f38248 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -210,8 +210,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/README.md b/docs/xray/README.md index da61e8d1..a9413632 100644 --- a/docs/xray/README.md +++ b/docs/xray/README.md @@ -34,7 +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`) | +| 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` | @@ -330,3 +330,4 @@ See [Data and artifact contracts](../data-artifacts.md#xray-hdf5-and-trace-produ - [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 e73fe23b..4499d815 100644 --- a/src/cuphoton/xray/commands.py +++ b/src/cuphoton/xray/commands.py @@ -657,6 +657,16 @@ class DistortionArg(StringInvariant): _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): @@ -675,6 +685,108 @@ class JsonArg(BoolInvariant): _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" @@ -4401,14 +4513,11 @@ def _linear_prediction_smoke(args): return 0 -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()) +def _parse_validation_distortion(value): distortion = None - if args.distortion: + if value: try: - kind, amount = str(args.distortion).split(":", maxsplit=1) + kind, amount = str(value).split(":", maxsplit=1) if not kind.strip(): raise ValueError("missing distortion kind") distortion = (kind.strip(), float(amount)) @@ -4417,6 +4526,14 @@ def _linear_prediction_validate(args): "--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, @@ -4425,6 +4542,7 @@ def _linear_prediction_validate(args): seed=args.seed, distortion=distortion, output_dir=args.output_dir, + refine=args.refine, ) args._out( json.dumps( @@ -4436,6 +4554,47 @@ def _linear_prediction_validate(args): 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 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..15afce24 --- /dev/null +++ b/src/cuphoton/xray/mode_refinement_batched.py @@ -0,0 +1,432 @@ +# 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 _scaled_normal_equations(xp, jac, resid, cost, tol): + """Column scales, scaled normal matrix, scaled gradient and the + per-trace stationarity flag for the current iterate.""" + 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) + ) + return scale, jtj, jtr, stationary + + +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, jtj, jtr, stationary = _scaled_normal_equations( + xp, jac, resid, cost, tol + ) + 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) + else: + # The returned parameters include the final accepted step, so the + # flag must describe that iterate rather than the one before it. + converged = ( + converged | _scaled_normal_equations(xp, jac, resid, cost, tol)[3] + ) + _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/synthetic_validation.py b/src/cuphoton/xray/synthetic_validation.py index e278306c..b0e3a1c6 100644 --- a/src/cuphoton/xray/synthetic_validation.py +++ b/src/cuphoton/xray/synthetic_validation.py @@ -466,7 +466,10 @@ def decay_ratio(self) -> float: 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.""" + 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 @@ -478,6 +481,8 @@ class ValidationLevel: modes: tuple[ModeValidation, ...] any_mode_lost_rate: float residual_ratio: float + solver_converged: int = 0 + solver_unconverged: int = 0 def __str__(self) -> str: lines = [ @@ -487,6 +492,11 @@ def __str__(self) -> str: 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) @@ -617,7 +627,9 @@ def validation_sweep( ``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. + 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 @@ -681,6 +693,8 @@ def validation_sweep( 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) @@ -691,6 +705,10 @@ def validation_sweep( 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), @@ -739,6 +757,8 @@ def validation_sweep( 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), @@ -770,14 +790,16 @@ def validate_linear_prediction( 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. + when Bokeh is available, a validation figure. Set ``refine`` to compare + the experimental nonlinear estimator using the same noise realizations. """ - sweeps = ( + sweeps: tuple[ValidationSweep, ...] = ( validation_sweep( samples=samples, snr_db=snr_db, @@ -787,6 +809,26 @@ def validate_linear_prediction( 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 @@ -966,6 +1008,16 @@ def build_summary( "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 diff --git a/tests/core/test_cli_contract.py b/tests/core/test_cli_contract.py index 3ef79c25..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", 34, 32, 400, 1), + ("xray", 36, 34, 412, 1), ] assert len(per_group) == 6 - assert sum(item[1] for item in per_group) == 95 - assert sum(item[1] + item[4] for item in per_group) == 100 - assert sum(item[2] for item in per_group) == 92 - assert sum(item[3] for item in per_group) == 969 + 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 eed6adbe..338a9e9a 100644 --- a/tests/xray/test_cli.py +++ b/tests/xray/test_cli.py @@ -44,6 +44,8 @@ def test_cli_help(capsys): [ ("DataProbeCommand", "on"), ("LinearPredictionValidateCommand", "samples"), + ("LinearPredictionRefineBenchmarkCommand", "traces"), + ("LinearPredictionModesReviewCommand", "html"), ("ModelOrderSweepCommand", "trace_count"), ("SubspaceBenchmarkCommand", "trace_count"), ], @@ -61,6 +63,11 @@ def test_report_commands_use_injected_output( 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)] @@ -112,6 +119,8 @@ def test_core_command_registry_loads_xray_commands(): "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", @@ -348,6 +357,27 @@ def broken(t, y, k): 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() @@ -2362,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_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..9fee5d1a --- /dev/null +++ b/tests/xray/test_mode_refinement_batched.py @@ -0,0 +1,282 @@ +# 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] + + +def test_convergence_flag_includes_the_final_step_at_the_iteration_limit(): + fx = synthetic_modes_trace(96) + start = modes_to_theta(fx.modes, fx.constant)[None, :] + start[:, [2, 6]] *= 1.05 + at_limit = refine_modes_batched( + np, fx.time, fx.clean[None, :], start, 2, max_iter=5 + ) + past_limit = refine_modes_batched( + np, fx.time, fx.clean[None, :], start, 2, max_iter=6 + ) + assert at_limit.iterations == 5 + np.testing.assert_array_equal(at_limit.theta, past_limit.theta) + model, jac = model_and_jacobian_batched(np, at_limit.theta, fx.time, 2) + resid = model - fx.clean + scale = np.linalg.norm(jac, axis=1) + scaled = jac / np.where(scale > 0, scale, 1.0)[:, None, :] + gradient = np.max(np.abs(np.einsum("bnp,bn->bp", scaled, resid))) + assert gradient <= 1e-9 * max(np.linalg.norm(resid), 1.0) + assert at_limit.converged.tolist() == [True] + assert past_limit.converged.tolist() == [True] + + +@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_synthetic_validation.py b/tests/xray/test_synthetic_validation.py index 729bb18a..b109c5db 100644 --- a/tests/xray/test_synthetic_validation.py +++ b/tests/xray/test_synthetic_validation.py @@ -637,6 +637,43 @@ def test_distort_trace_kinds_and_rejection(): 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() + ) + + @pytest.mark.parametrize("time", [[], [[0.0, 1.0]], [0.0, float("nan")]]) def test_distort_trace_rejects_invalid_time(time): with pytest.raises(ValueError, match="time must be"):