Add experimental per-trace and batched mode refinement

.github/workflows/ci.yml

@@ -149,14 +149,15 @@ 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"
           tests/xfit/test_models.py::test_cupy_fit_matches_numpy_when_cuda_is_available
           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: >-

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.

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.

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.

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)

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())