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

@@ -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.

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

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