From 495f6a97b364d59bf0a4e1f93f0fda996bf85454 Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Sun, 4 Oct 2026 14:17:58 +0200 Subject: [PATCH] Rotation scaling: penalised per-frame scale of the fulls; fulls-only partials get flux and partiality Three defects in the per-frame scaling of sparse (small-molecule, weak) rotation sweeps: 1. The fulls' per-frame scale pooled sparse frames by their RAW full count, but in a high-symmetry group with many systematic absences most usable counts stayed under the minimum, so most frames were never fitted and kept corr = 1 beside pinned, fitted frames (two gauges in one reference). The release step then read the gauge offset (a constant 134x on a cubic Ia-3d small-molecule sweep) as signal and gave each frame exp(kept_f * 4.9) - the e^-5 errors on a quarter of the frames, R1 0.62. A fixed box window also cannot follow a 100x absorption ramp over a few degrees, and frames under the credible floor, exempt from the window, ran away to 1e-7. Now each round fits every frame on its own fulls (no pooling, no minimum), and the scale is a penalised second-difference smoother of log G (Whittaker/Eilers), each frame at the information of its fit, lambda by cross-validation over blocks one rocking curve wide (interleaved single frames leak through shared rocking curves and chose to follow every frame). After convergence the existing ShrinkToRestrained hands back the per-frame deviation its neighbour shares. Pooling and the box window are gone from the fulls loop; the partials loop is unchanged. 2. With partial scaling off (< 50 rocking events per frame) the partials kept the integration-time corr: no incident-flux correction and not the partiality of the ingest-smoothed geometry, because only the partial scaling loop rewrote corr. They now get corr = prescaling_corr / partiality at G = 1 (host, and a device kernel). 3. The flux meter (per-frame mean background) jumped 30x between neighbouring frames of a sparse sweep - on a few reflections it measures which reflections the frame holds. It is read through the same smoother at the precision of each frame's mean. SHELXL R1(>4sig) against the published structures, rc174-cand -> this, on the in-house small-molecule sweeps: cubic Ia-3d 0.615 -> 0.119 (XDS 0.088); four organic sweeps (monoclinic / orthorhombic) 0.0515 -> 0.0459, 0.0451 -> 0.0396, 0.0747 -> 0.0709, 0.0452 -> 0.0413; ISa up to 11.6 -> 31. Raw flux instead of smoothed costs 0.002-0.003 R1 on the first two. A weak, decaying protein sweep on the fulls-only path: ISa 6.7 -> 27.1. CPU and GPU paths agree. Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01K5K8jvPPbmCrbqnWkddTuB --- docs/ACKNOWLEDGEMENT.md | 7 + docs/CPU_DATA_ANALYSIS.md | 1 + docs/CPU_DATA_ANALYSIS_INTEGRATION.md | 8 +- .../scale_merge/RotationScaleMerge.cpp | 353 +++++++++++++----- .../scale_merge/RotationScaleMerge.h | 57 +-- .../scale_merge/RotationScaleMergeGPU.cu | 121 +++--- .../scale_merge/RotationScaleMergeGPU.h | 22 +- tests/RotationScaleWalkTest.cpp | 49 ++- 8 files changed, 427 insertions(+), 191 deletions(-) diff --git a/docs/ACKNOWLEDGEMENT.md b/docs/ACKNOWLEDGEMENT.md index 7a5c5a5e7..f7f6ae336 100644 --- a/docs/ACKNOWLEDGEMENT.md +++ b/docs/ACKNOWLEDGEMENT.md @@ -234,6 +234,13 @@ K. Diederichs, "Linking crystallographic model and data quality" (2012), Science P. A. Karplus, "Better models by discarding data?" (2013), Acta Cryst. D69, 1215-1222 [doi:10.1107/S0907444913001121](https://doi.org/10.1107/S0907444913001121). +**The Whittaker smoother** — the per-frame scale of the rotation fulls, and the incident-flux meter, are +penalised least-squares curves (a second-difference penalty, the smoothness chosen by cross-validation) +in the form P. H. C. Eilers gave Whittaker's graduation: P. H. C. Eilers, "A perfect smoother" (2003), +Anal. Chem. 75, 3631-3636 [doi:10.1021/ac034173t](https://doi.org/10.1021/ac034173t); E. T. Whittaker, +"On a new method of graduation" (1923), Proc. Edinburgh Math. Soc. 41, 63-75 (DOI not verified, so +not given). + **Fisher's z-transformation** — the cross-validation of the scaling correction surfaces averages the change of the half-set CC1/2 over resolution shells on atanh(CC), so that the shells near CC = 1, where a multiplicative error shows, are not outweighed by the noise of the shells without signal. R. A. diff --git a/docs/CPU_DATA_ANALYSIS.md b/docs/CPU_DATA_ANALYSIS.md index 147c15f8d..e16c8cd42 100644 --- a/docs/CPU_DATA_ANALYSIS.md +++ b/docs/CPU_DATA_ANALYSIS.md @@ -52,6 +52,7 @@ The methods draw on, and in places reimplement, solutions from: - A. Hennequin, B. Couturier, V. V. Gligorov & L. Lacassagne, "SparseCCL: Connected Components Labeling and Analysis for sparse images", DASIP 2019, 65-70 (the connected-component labelling of §3.4, used via ACTS/traccc). - S. French & K. Wilson, "On the treatment of negative intensity observations", *Acta Cryst.* **A34** (1978), 517-525 (Bayesian amplitude estimation from intensities). - A. T. Brünger, "Free R value: a novel statistical quantity for assessing the accuracy of crystal structures", *Nature* **355** (1992), 472-475 (R-free cross-validation). +- P. H. C. Eilers, "A perfect smoother", *Anal. Chem.* **75** (2003), 3631-3636, after E. T. Whittaker, "On a new method of graduation", *Proc. Edinburgh Math. Soc.* **41** (1923), 63-75 (the penalised smoother of the fulls' per-frame scale and of the incident-flux meter, §10.6). - R. A. Fisher, "Frequency distribution of the values of the correlation coefficient in samples from an indefinitely large population", *Biometrika* **10** (1915), 507-521 (the z-transformation on which the correction surfaces' held-out half-set CC1/2 is compared). - M. Wojdyr, "GEMMI: A library for structural biology", *J. Open Source Softw.* **7** (2022), 4200 (model / structure-factor / map machinery used in §14). - J. P. Wright, "Experiences with GPU decompression for bitshuffle + LZ4 data", HDF5 User Group meeting (2021), and [github.com/jonwright/bslz4decoders](https://github.com/jonwright/bslz4decoders) (device-side decoding of bitshuffle+LZ4 images, §0). diff --git a/docs/CPU_DATA_ANALYSIS_INTEGRATION.md b/docs/CPU_DATA_ANALYSIS_INTEGRATION.md index 4a55123ab..1ac700beb 100644 --- a/docs/CPU_DATA_ANALYSIS_INTEGRATION.md +++ b/docs/CPU_DATA_ANALYSIS_INTEGRATION.md @@ -106,7 +106,9 @@ Only the ring moves. The signal disk $r_1$ stays circular, deliberately: it sets What a circular $r_1$ loses is flux, and that loss is **not** a function of resolution alone: measured per reflection, it carries a directional component worth several Ų with a definite principal axis, on top of the isotropic part. Nor is there anything in the merge to absorb it. There is **no per-shell scale**, and there cannot usefully be one: every scale in §10 is fitted against a reference built from a reflection's own symmetry equivalents, and equivalents share $s^2$ exactly, so any function of $s^2$ lies in the exact null space of the whole scaling model — a per-shell parameter would have zero residual to fit against. (XDS and DIALS have the same null space, for the same reason.) The isotropic part of the loss is instead degenerate with the overall Wilson $B$ and is silently reported as part of it, so **the reported `WILSON_B` / `_reflns.B_iso_Wilson_estimate` carries an $r_1$-dependent contribution**: measured across a constant-ring-area radius sweep it falls monotonically as the disk grows, by 0.5 Ų on sharp strong data and by up to ~10 Ų on weak wide-spot data. What this costs the *data* is much less than what it costs the flux, because most of the loss is matched by a proportional $\sigma$: it moves no CC$_{1/2}$ and no $R_\text{meas}$, and — to within a few hundredths of an ångström — no resolution cut. -**Measured spot footprint (automatic).** The radii above are chosen from spots at 5 Å, which at high X-ray energy sit close to the beam. Away from it a spot can grow several times wider — radially from the sensor's parallax and the obliquity of the incidence, tangentially from the crystal's azimuthal spread, which rotates the diffracted beam about the incident one and smears the spot along its ring. On small-molecule data at 20–25 keV the standard deviation grows from ~1 px near the beam to ~5 px at the detector edge: the $r_1 = 4$ disk holds a quarter of the flux there, the $6\ldots13$ px ring a third of it, and the profile widths learned inside $r_1$ (§9.3) saturate near $r_1^2/4$. So the pre-scan measures every spot it finds with a window that follows the spot — three of its own standard deviations, iterated and re-centred — separately along and across the radius, and tabulates the median widths $\sigma_ ho,\sigma_ au$ against the distance from the beam. Wherever $3\max(\sigma_ ho,\sigma_ au)>r_1$ the integrator then (i) starts the background ring at $3\sigma$ along each axis, (ii) sums the reflection over the $r_1$ disk **and** the $3\sigma$ footprint ellipse, so the summation — the profile fit's seed and its fallback — holds the spot rather than its core, and (iii) builds the per-reflection Gaussian at the measured widths on a grid grown to hold them. Where every spot fits the disk nothing is installed and the integration is unchanged bit for bit, which is the case for compact protein spots; like the measured radius, the footprint applies to the canonical pass and not to the geometry pre-pass, and a canonical pass whose wider rings the neighbours starve falls back to the settings without it. Judged by refining the published structures with SHELXL, it removes the intensity loss that grew with resolution on the small-molecule sets (rugnux/model intensity in the outermost shell 0.81–0.91 → 0.98–1.02). +**Measured spot footprint (automatic).** The radii above are chosen from spots at 5 Å, which at high X-ray energy sit close to the beam. Away from it a spot can grow several times wider — radially from the sensor's parallax and the obliquity of the incidence, tangentially from the crystal's azimuthal spread, which rotates the diffracted beam about the incident one and smears the spot along its ring. On small-molecule data at 20–25 keV the standard deviation grows from ~1 px near the beam to ~5 px at the detector edge: the $r_1 = 4$ disk holds a quarter of the flux there, the $6\ldots13$ px ring a third of it, and the profile widths learned inside $r_1$ (§9.3) saturate near $r_1^2/4$. So the pre-scan measures every spot it finds with a window that follows the spot — three of its own standard deviations, iterated and re-centred — separately along and across the radius, and tabulates the median widths $\sigma_ +ho,\sigma_ au$ against the distance from the beam. Wherever $3\max(\sigma_ +ho,\sigma_ au)>r_1$ the integrator then (i) starts the background ring at $3\sigma$ along each axis, (ii) sums the reflection over the $r_1$ disk **and** the $3\sigma$ footprint ellipse, so the summation — the profile fit's seed and its fallback — holds the spot rather than its core, and (iii) builds the per-reflection Gaussian at the measured widths on a grid grown to hold them. Where every spot fits the disk nothing is installed and the integration is unchanged bit for bit, which is the case for compact protein spots; like the measured radius, the footprint applies to the canonical pass and not to the geometry pre-pass, and a canonical pass whose wider rings the neighbours starve falls back to the settings without it. Judged by refining the published structures with SHELXL, it removes the intensity loss that grew with resolution on the small-molecule sets (rugnux/model intensity in the outermost shell 0.81–0.91 → 0.98–1.02). ### 9.2 Box summation (seed and fallback) @@ -381,6 +383,10 @@ The combine groups each reflection's partials into rocking events (contiguous ru The fulls are then re-scaled in the XDS sense — a per-image scale refit directly on the complete reflections under the unity partiality model — and merged (§10.4). Because every merged observation is now a counting-statistics-limited full rather than a partiality-divided slice, the error model reaches a far higher asymptotic $I/\sigma$. +How smooth that scale is over the rotation is left to the data rather than to a fixed window. Each round fits every frame on its own fulls against the current reference, giving a scale $G_f$ and its information $D_f=\sum w^2c^2$; the scale is then the curve $x=\log G$ minimising $\sum_f J_f\,(x_f-y_f)^2+\lambda\sum_f(\Delta^2 x)_f^2$, with $y_f$ the frame's own fit in log scale and $J_f$ its information carried there — a penalised (Whittaker–Eilers) smoother, solved as a five-band linear system. $\lambda$ is chosen by cross-validation: blocks of frames one rocking curve wide are left out in turn and predicted from the curve through the rest (neighbours closer than a rocking curve share their measurement, since a full sums those frames). A frame of hundreds of fulls is then followed frame by frame, a frame of two or three is carried by its neighbours, a stretch with none is bridged by a straight line, and a scale that falls a hundredfold over a few degrees — an absorbing crystal turning edge-on — is followed where a window would average across it. Once the curve settles, one free fit is shrunk toward it frame by frame by how much of each frame's deviation its neighbour shares (the lag-1 covariance), which hands back a real per-frame systematic and discards fit noise. + +When a sweep holds too few rocking events per frame to scale the partials, the partials stand at $G=1$ with every other term of their correction applied — the deterministic corrections, the incident flux and the partiality of the smoothed geometry — exactly as the partial scaling would leave them at that scale. The incident flux itself (the per-frame mean background) is read through the same penalised smoother, each frame at the precision of its mean: on a frame of a handful of reflections the mean measures which reflections the frame holds as much as the flux. + After scale-fulls, four **correction surfaces** are fitted on the combined fulls (rotation path, **on by default**; disable all with `--no-scaling-corrections`), each an alternating multiplicative refinement of the per-full scale against the merged reference: - **Decay.** Radiation damage weakens later frames more at higher resolution — a resolution×time (Debye–Waller) systematic the resolution-flat per-image scale cannot capture. A single global relative-$B$ rate is fitted, $\ln(I_\mathrm{ref}/I_\mathrm{obs}) = 2\,(\mathrm{d}B/\mathrm{d}n)\,(n-\bar n)\,s^2$ (frame $n$, $s^2 = 1/4d^2$), and folded into the scale. It engages only when the total relative-$B$ over the run exceeds a physical floor (2 Ų); below that the decay is negligible and "correcting" it only spreads symmetry equivalents (same $s^2$, different frames). An optional **per-batch relative-$B$** (`--relative-b[=deg]`, off unless requested; 10°-of-rotation batches by default) extends the single global rate to a smooth $B(n)$ curve — the same $s^2$-weighted decay fit solved independently over short frame batches, curvature-penalized so it cannot over-fit and cross-validated like the surfaces below — for crystals whose decay is non-linear in dose. Its cross-validation splits on **ASU-group parity**, not the frame parity the surfaces below use: a per-batch parameter owns whole frames and so cannot be scored on a held-out frame, whereas splitting the symmetry equivalents tests whether a batch's $B$ generalises to reflections it was not fitted on. diff --git a/image_analysis/scale_merge/RotationScaleMerge.cpp b/image_analysis/scale_merge/RotationScaleMerge.cpp index b1371ee4e..400420aa3 100644 --- a/image_analysis/scale_merge/RotationScaleMerge.cpp +++ b/image_analysis/scale_merge/RotationScaleMerge.cpp @@ -52,10 +52,12 @@ namespace { constexpr double SCALING_TOLERANCE = 1e-3; // The per-frame scale is smoothed over at least this many rocking curves (see Run). constexpr int SMOOTH_G_MIN_ROCKING_EVENTS = 6; - // ShrinkToRestrained reads two variances off the frames of the sweep; on fewer frames than this // neither of them means anything, and the restraint is kept whole. constexpr size_t MIN_FRAMES_FOR_SHRINKAGE = 20; + // Interleaved folds of the cross-validation that sets the smoothness of the fulls' scale. + constexpr int CV_FOLDS = 5; + // A fitted per-frame scale below this fraction of the run's median is not a measurement of anything: // it would mean the frame received 2% of its neighbours' dose while still producing indexable spots. @@ -205,7 +207,7 @@ namespace { // the loop shrank every scale by 10-30% per round for thirty rounds under the robust fit and // settled in fourteen under the plain one. With the weights the reference uses, the plain // alternating fit is exact coordinate descent on one objective and cannot climb. - double SolveScale(const std::vector &obs) { + double SolveScale(const std::vector &obs, double &info) { double num = 0.0, den = 0.0; for (const auto &o : obs) { const double w2 = o.weight * o.weight; @@ -213,6 +215,7 @@ namespace { den += w2 * o.coeff * o.coeff; } const double G = den > 0.0 ? num / den : NAN; + info = den; return std::isfinite(G) ? std::max(0.0, G) : 1.0; } @@ -280,28 +283,44 @@ RotationScaleMerge::RotationScaleMerge(const DiffractionExperiment &experiment, // here the fitted G takes back out, because smoothing leaves G its low frequencies. Only the fast part - // the part no smoothed per-frame scale can represent anyway - is taken on trust. // +// The mean is over the frame's reflections, whose backgrounds differ with where on the detector they +// fall, so on a frame of a handful of reflections it measures which reflections the frame holds as +// much as the flux: on a sparse small-molecule sweep it jumped thirtyfold between neighbouring frames. +// So the meter is read through the same smoothing as the fulls' scale (SmoothLogScale, lambda by +// cross-validation), each frame at the precision of its mean: a frame of thousands of reflections +// keeps its own reading, a frame of five is carried by its neighbours. +// // Neither half of this costs a pass over the observations of its own: the per-frame mean background is // accumulated by the ingest fill loop, which already reads every reflection, and the division into prescaling_corr // rides on the finiteness pass that follows it. -void RotationScaleMerge::MeasureIncidentFlux(const std::vector &mean_bkg) { +void RotationScaleMerge::MeasureIncidentFlux(const std::vector &mean_bkg, + const std::vector &mean_bkg_info) { frame_flux.assign(n_frames, 1.0); + // log of the mean at the information on it: n mean^2 / var for n background readings. A frame's + // reading is its own measurement, so the frames are left out one at a time (block 1). + std::vector y(n_frames, 0.0), J(n_frames, 0.0); + for (int f = 0; f < n_frames; ++f) + if (std::isfinite(mean_bkg[f]) && mean_bkg[f] > 0.0) { + y[f] = std::log(mean_bkg[f]); + J[f] = mean_bkg_info[f]; + } + const std::vector x = SmoothLogScale(y, J, ChooseLogScaleSmoothing(y, J, 1)); + // Collected in frame order, serially: it is one value per frame, and the median only depends on // the set, but keeping the order removes the question. std::vector finite; finite.reserve(n_frames); for (int f = 0; f < n_frames; ++f) - if (std::isfinite(mean_bkg[f])) finite.push_back(mean_bkg[f]); + if (std::isfinite(mean_bkg[f])) finite.push_back(x[f]); if (finite.empty()) return; // Gauge: the run median. Only the frame-to-frame RATIO matters - a constant factor is degenerate // with the overall scale. - const size_t mid = finite.size() / 2; - std::nth_element(finite.begin(), finite.begin() + mid, finite.end()); - const double median = finite[mid]; + const double median = median_of(finite); for (int f = 0; f < n_frames; ++f) - if (std::isfinite(mean_bkg[f])) frame_flux[f] = mean_bkg[f] / median; + if (std::isfinite(mean_bkg[f])) frame_flux[f] = std::exp(x[f] - median); } void RotationScaleMerge::Ingest() { @@ -340,7 +359,7 @@ void RotationScaleMerge::Ingest() { // The per-frame mean background - the incident-flux meter, see MeasureIncidentFlux - is // accumulated here rather than in a pass of its own. One frame is still summed in its own order, // which is what keeps the mean bit-exact. - std::vector mean_bkg(n_frames, NAN); + std::vector mean_bkg(n_frames, NAN), mean_bkg_info(n_frames, 0.0); // The h range comes out of the same sweep - the sort below buckets by h and needs to know how many // buckets that is - and so does the narrow key of every observation. The fat Obs is NOT built here: @@ -355,19 +374,28 @@ void RotationScaleMerge::Ingest() { const auto cell = partials_out[o].latt.GetUnitCell(); frame_cell_ok[o] = cell.is_close(*reference_cell, dist_tol, ang_tol) ? 1 : 0; } - double bkg_sum = 0.0; + double bkg_sum = 0.0, bkg_sum2 = 0.0; int bkg_n = 0; int lmin = INT_HI, lmax = INT_LO; for (const auto &r : partials_out[o].reflections) { - if (std::isfinite(r.bkg) && r.bkg > 0.0f) { bkg_sum += r.bkg; ++bkg_n; } + if (std::isfinite(r.bkg) && r.bkg > 0.0f) { + bkg_sum += r.bkg; + bkg_sum2 += static_cast(r.bkg) * r.bkg; + ++bkg_n; + } lmin = std::min(lmin, r.h); lmax = std::max(lmax, r.h); } - if (bkg_n > 0) mean_bkg[o] = bkg_sum / bkg_n; + if (bkg_n > 0) { + const double mean = bkg_sum / bkg_n; + const double var = bkg_sum2 / bkg_n - mean * mean; + mean_bkg[o] = mean; + if (bkg_n > 1 && var > 0.0) mean_bkg_info[o] = bkg_n * mean * mean / var; + } h_lo[o] = lmin; h_hi[o] = lmax; }); - MeasureIncidentFlux(mean_bkg); + MeasureIncidentFlux(mean_bkg, mean_bkg_info); int h_min = 0, h_max = 0; { @@ -1417,16 +1445,69 @@ void RotationScaleMerge::ReduceScalingGroupMeans(int n_groups, const std::vector }); } -int RotationScaleMerge::PoolHalfWidth(const std::vector &fcount, int f) { - const int n = static_cast(fcount.size()); - int64_t total = fcount[f]; - int h = 0; - while (total < static_cast(MIN_REFLECTIONS) && h < n) { - ++h; - if (f - h >= 0) total += fcount[f - h]; - if (f + h < n) total += fcount[f + h]; +std::vector RotationScaleMerge::SmoothLogScale(const std::vector &y, const std::vector &J, + double lambda) { + // Following Eilers (2003) Anal. Chem. 75, 3631-3636 (Whittaker's smoother). + // The normal equations (diag(J) + lambda D2'D2) x = J y, D2 the second-difference operator: a + // symmetric band of half-width two, solved by its Cholesky factor in one pass down and one up. + // A curve with no second difference (a straight line in log scale) costs nothing, so two frames + // with information are enough to make the system definite. + const int n = static_cast(y.size()); + if (std::count_if(J.begin(), J.end(), [](double j) { return j > 0.0; }) < 2) + return y; + std::vector a0(J), a1(n, 0.0), a2(n, 0.0); // diagonal, first and second off-diagonal + for (int i = 0; i + 2 < n; ++i) { + a0[i] += lambda; a0[i + 1] += 4.0 * lambda; a0[i + 2] += lambda; + a1[i] -= 2.0 * lambda; a1[i + 1] -= 2.0 * lambda; + a2[i] += lambda; } - return h; + std::vector l0(n), l1(n, 0.0), l2(n, 0.0), z(n), x(n); + for (int j = 0; j < n; ++j) { + const double p1 = j >= 1 ? l1[j - 1] : 0.0, p2 = j >= 2 ? l2[j - 2] : 0.0; + l0[j] = std::sqrt(a0[j] - p1 * p1 - p2 * p2); + if (j + 1 < n) l1[j] = (a1[j] - (j >= 1 ? l2[j - 1] * l1[j - 1] : 0.0)) / l0[j]; + if (j + 2 < n) l2[j] = a2[j] / l0[j]; + } + for (int j = 0; j < n; ++j) + z[j] = (J[j] * y[j] - (j >= 1 ? l1[j - 1] * z[j - 1] : 0.0) - (j >= 2 ? l2[j - 2] * z[j - 2] : 0.0)) / l0[j]; + for (int j = n - 1; j >= 0; --j) + x[j] = (z[j] - (j + 1 < n ? l1[j] * x[j + 1] : 0.0) - (j + 2 < n ? l2[j] * x[j + 2] : 0.0)) / l0[j]; + return x; +} + +double RotationScaleMerge::ChooseLogScaleSmoothing(const std::vector &y, const std::vector &J, + int block) { + // Blocks of `block` frames are left out in turn, every fifth block, and predicted from the curve + // through the others; the lambda whose curve predicts them best (each frame at its own + // information) is the one the data ask for. The grid runs in half decades from a curve that + // follows every frame (lambda = J/100: a frame's neighbours barely pull on it) to one that is + // straight over hundreds of frames. + std::vector positive; + for (double j : J) + if (j > 0.0) positive.push_back(j); + if (positive.size() < 2 * CV_FOLDS) + return 0.0; + auto fold_of = [&](int f) { return (f / block) % CV_FOLDS; }; + const double j_typ = median_of(positive); + const int n = static_cast(y.size()); + double best_lambda = 0.0, best_err = INFINITY; + for (int k = -4; k <= 20; ++k) { + const double lambda = j_typ * std::pow(10.0, 0.5 * k); + double err = 0.0; + for (int fold = 0; fold < CV_FOLDS; ++fold) { + std::vector J_train(J); + for (int f = 0; f < n; ++f) + if (fold_of(f) == fold) J_train[f] = 0.0; + const std::vector x = SmoothLogScale(y, J_train, lambda); + for (int f = 0; f < n; ++f) + if (fold_of(f) == fold) err += J[f] * (x[f] - y[f]) * (x[f] - y[f]); + } + if (err < best_err) { + best_err = err; + best_lambda = lambda; + } + } + return best_lambda; } template @@ -1435,16 +1516,11 @@ void RotationScaleMerge::FitPerFrameG(const std::vector &obs, const std::vect const std::vector &group_mean_in, bool unity, std::vector &g) { std::vector scaled(fstart.size(), 0); + std::vector info(fstart.size(), 0.0); const int n_fr = static_cast(fstart.size()); ParallelFor(n_fr, nthreads, [&](int f) { - // A full is a whole rocking event filed under one frame, so on a sparse sweep a frame holds a - // handful of them - too few to fit a scale on, and the frame was left unscaled, which on such a - // sweep was every frame. Those frames are fitted over the nearest frames on either side that - // together hold enough; a frame with enough of its own is fitted on its own, as before. - const int h = unity && pool_sparse_fulls ? PoolHalfWidth(fcount, f) : 0; std::vector so; - for (int j = std::max(0, f - h); j <= std::min(n_fr - 1, f + h); ++j) - for (int i = fstart[j]; i < fstart[j] + fcount[j]; ++i) { + for (int i = fstart[f]; i < fstart[f] + fcount[f]; ++i) { const auto &o = obs[i]; if (o.group < 0) continue; if (o.on_ice) continue; @@ -1471,27 +1547,17 @@ void RotationScaleMerge::FitPerFrameG(const std::vector &obs, const std::vect } so.push_back({coeff, static_cast(o.I), SafeInv(o.sigma, 1.0)}); } - if (so.size() < MIN_REFLECTIONS) return; // leave g[f]/corr untouched (as ScaleOnTheFly does) - g[f] = SolveScale(so); + // A partial frame needs MIN_REFLECTIONS of its own (corr is left untouched otherwise, as + // ScaleOnTheFly does). A full is a whole rocking event filed under one frame, so on a sparse + // sweep a frame holds a handful; its fit goes in at the information it has, and the smoothing + // in RunFullsScalingLoop decides how far its neighbours carry it. + if (so.empty() || (!unity && so.size() < MIN_REFLECTIONS)) return; + g[f] = SolveScale(so, info[f]); scaled[f] = 1; }); // Remember which frames were fitted this call (so the caller updates corr only there). frame_scaled_scratch = std::move(scaled); -} - -void RotationScaleMerge::UpdateCorr(std::vector &obs, const std::vector &g, - const std::vector &frame_scaled) const { - ParallelChunks(static_cast(obs.size()), nthreads, [&](int lo, int hi) { - for (int i = lo; i < hi; ++i) { - auto &o = obs[i]; - if (!frame_scaled[o.frame]) continue; - const double denom = static_cast(o.partiality) * g[o.frame]; // B_term = 1 (no B refine) - if (std::isfinite(o.prescaling_corr) && std::isfinite(denom) && denom > 0.0) - o.corr = static_cast(o.prescaling_corr / denom); - else - o.corr = NAN; - } - }); + frame_info_scratch = std::move(info); } void RotationScaleMerge::ShrinkToRestrained(const char *what, const std::vector &g_restrained, @@ -1573,7 +1639,6 @@ void RotationScaleMerge::ShrinkToRestrained(const char *what, const std::vector< RotationScaleMerge::ScalingLoopOutcome RotationScaleMerge::RunScalingLoop( const char *what, std::vector &g, const std::vector &frame_obs_count, int window, - bool release_restraint, const std::function &iterate, const std::function &, const std::vector &)> &rescale) { // Alternating least squares on a bilinear model: the reference intensities given the scales, the @@ -1597,14 +1662,9 @@ RotationScaleMerge::ScalingLoopOutcome RotationScaleMerge::RunScalingLoop( // mean, so no frame can lose its vote alone, and what a frame's own observations cannot determine // its neighbourhood's do. The post-loop smoothing is then this same operation and is not repeated. // - // The restraint is a prior, though, not a measurement, and on a merge where a frame carries - // hundreds of observations the frame's own data outvote it. So a loop that is asked to - // (release_restraint) ends with ONE free fit against the settled reference, shrunk frame by - // frame toward the restrained scale by how much of the difference is real - see - // ShrinkToRestrained. Only the loop over the fulls is asked to: a partial's scale is the - // direction the collapse runs in, and handing any of it back re-opens it (measured: ISa 23.8 -> - // 22.1 on a 360 deg, 14x merge), while a full is a whole rocking event and cannot enter it. And - // only off a loop that settled, since a loop still walking has measured nothing to hand back. + // The restraint is never released here: a partial's scale is the direction the collapse runs in, + // and handing any of it back re-opens it (measured: ISa 23.8 -> 22.1 on a 360 deg, 14x merge). The + // fulls, each a whole rocking event, have their own loop (RunFullsScalingLoop). // // The loop stops when the relative scales stop moving - the rms of |log(G_new/G_old)| over the // frames, each weighted by its observation count and nothing else (a weight that vanishes with G @@ -1684,27 +1744,12 @@ RotationScaleMerge::ScalingLoopOutcome RotationScaleMerge::RunScalingLoop( g_prev = g; } frame_scaled_scratch = scaled_any; - // Only off a settled loop: the deviation below is a measurement of the frames, and a loop that - // walked until it was stopped has not measured them. - if (release_restraint && window > 1 && out.converged) { - const std::vector g_restrained = g; - iterate(); // one free fit against the settled reference - for (int f = 0; f < n_frames; ++f) - scaled_any[f] |= frame_scaled_scratch[f]; - const std::vector g_free = g; - ShrinkToRestrained(what, g_restrained, scaled_any, frame_obs_count, g); - for (int f = 0; f < n_frames; ++f) - n_obs[f] = scaled_any[f] && std::isfinite(g[f]) && g[f] > 0.0 ? frame_obs_count[f] : 0; - const double g_ref = TypicalFrameScale(g, n_obs); - for (int f = 0; f < n_frames; ++f) { - if (!scaled_any[f] || !(std::isfinite(g[f]) && g[f] > 0.0)) { ratio[f] = 1.0; continue; } - ratio[f] = g_free[f]; - g[f] /= g_ref; - ratio[f] /= g[f]; - } - rescale(scaled_any, ratio); - frame_scaled_scratch = scaled_any; - } + out.best_step = best_step; + LogScalingOutcome(what, out); + return out; +} + +void RotationScaleMerge::LogScalingOutcome(const char *what, const ScalingLoopOutcome &out) const { if (out.converged) logger.Info("Per-frame scaling ({}): settled after {} iterations (rms |dlogG| {:.1e})", what, out.iterations, out.step); @@ -1712,12 +1757,121 @@ RotationScaleMerge::ScalingLoopOutcome RotationScaleMerge::RunScalingLoop( logger.Warning("Per-frame scaling ({}): stopped after {} iterations - the scales stopped settling " "(rms |dlogG| {:.1e}, not below {:.1e} for five rounds); more rounds would only walk " "them, so everything read off this merge is read off an unsettled state", what, - out.iterations, out.step, best_step); + out.iterations, out.step, out.best_step); else logger.Warning("Per-frame scaling ({}): NOT settled after the cap of {} iterations - the scales " "were still moving by rms |dlogG| {:.1e} (tolerance {:.0e}); everything read off " "this merge is read off an unsettled state", what, out.iterations, out.step, SCALING_TOLERANCE); +} + +RotationScaleMerge::ScalingLoopOutcome RotationScaleMerge::RunFullsScalingLoop( + std::vector &g, std::vector &g_fit, const std::vector &frame_obs_count, + const std::function &fit, + const std::function &, const std::vector &)> &rescale) { + // A full is a whole rocking event, so its frame's scale is free of the partiality model, and how + // smooth that scale is over the rotation is left to the data rather than fixed in advance. Each + // round fits every frame on its own fulls against the current reference, and the scale is then + // the curve in log G that stays closest to those fits, each at its information, while bending as + // little as possible (SmoothLogScale) - with how little chosen by how well the curve predicts + // frames it was not fitted to (ChooseLogScaleSmoothing). A frame of hundreds of fulls is then + // followed frame by frame, a frame of two or three is carried by its neighbours, a stretch with + // none is bridged, and a scale that falls by a hundred over a few degrees - an absorbing crystal + // turning edge-on - is followed where a fixed window would average across it. + // + // In log G about the scale s the fulls stand at, a frame's fit G says y = log s + log(G/s) at + // information s^2 sum w^2 c^2 (the fit's information on G, carried to log G at s). A fit at or + // below zero - a frame whose fulls hold nothing the reference expects - moves its frame one e-fold + // down per round and no further, so that a frame the crystal stopped diffracting on walks toward + // zero, where DropCollapsedFullScales finds it, without one round's noise throwing it there. + // + // A frame below the credible floor is left out of the step, as in RunScalingLoop: it walks + // toward zero for as long as the loop runs, which is the answer, not a failure to settle. + ScalingLoopOutcome out; + std::vector y(n_frames), J(n_frames), g_new(n_frames), ratio(n_frames); + std::vector n_obs(n_frames); + std::vector has_fulls(n_frames); + for (int f = 0; f < n_frames; ++f) { + n_obs[f] = frame_obs_count[f]; + has_fulls[f] = frame_obs_count[f] > 0 ? 1 : 0; + } + double best_step = INFINITY, lambda = 0.0; + int rounds_since_best = 0; + for (int it = 0; it < scaling_iter; ++it) { + fit(); + out.iterations = it + 1; + for (int f = 0; f < n_frames; ++f) { + const bool fitted = frame_scaled_scratch[f] && frame_info_scratch[f] > 0.0; + J[f] = fitted ? frame_info_scratch[f] * g[f] * g[f] : 0.0; + y[f] = std::log(g[f]); + if (fitted) + y[f] += g_fit[f] > 0.0 ? std::max(std::log(g_fit[f] / g[f]), -1.0) : -1.0; + } + lambda = ChooseLogScaleSmoothing(y, J, std::max(1, rocking_event_frames_at_start)); + const std::vector x = SmoothLogScale(y, J, lambda); + for (int f = 0; f < n_frames; ++f) + g_new[f] = std::exp(x[f]); + const double g_ref = TypicalFrameScale(g_new, n_obs); + std::vector scales; + for (int f = 0; f < n_frames; ++f) { + g_new[f] /= g_ref; + ratio[f] = g[f] / g_new[f]; + if (has_fulls[f]) scales.push_back(g_new[f]); + } + rescale(has_fulls, ratio); + const double g_floor = scales.empty() ? 0.0 : median_of(scales) * MIN_CREDIBLE_SCALE_RATIO; + double ss = 0.0, sw = 0.0; + for (int f = 0; f < n_frames; ++f) + if (has_fulls[f] && g_new[f] >= g_floor) { + const double d = std::log(g_new[f] / g[f]); + ss += frame_obs_count[f] * d * d; + sw += frame_obs_count[f]; + } + g = g_new; + out.step = sw > 0.0 ? std::sqrt(ss / sw) : 0.0; + logger.Debug("Per-frame scaling (fulls): iteration {}, rms |dlogG| {:.2e}, lambda {:.3g}, gauge {:.4f}", + it + 1, out.step, lambda, g_ref); + if (out.step < SCALING_TOLERANCE) { + out.converged = true; + break; + } + if (out.step < 0.9 * best_step) { + best_step = out.step; + rounds_since_best = 0; + } else if (++rounds_since_best >= 5) { + out.stalled = true; + break; + } + } + // The curve is the part of the scale the frames share over a rocking curve and more; what a + // frame's own fulls say beyond it is handed back by how much of it is real (ShrinkToRestrained) - + // only off a loop that settled, since a loop still walking has measured nothing to hand back. + if (out.converged) { + fit(); + std::vector fitted(n_frames); + for (int f = 0; f < n_frames; ++f) + fitted[f] = frame_scaled_scratch[f] && frame_info_scratch[f] > 0.0 ? 1 : 0; + g_new = g; + ShrinkToRestrained("fulls", g, fitted, frame_obs_count, g_fit); + for (int f = 0; f < n_frames; ++f) + if (fitted[f] && std::isfinite(g_fit[f]) && g_fit[f] > 0.0) g_new[f] = g_fit[f]; + const double g_ref = TypicalFrameScale(g_new, n_obs); + for (int f = 0; f < n_frames; ++f) { + g_new[f] /= g_ref; + ratio[f] = g[f] / g_new[f]; + } + rescale(has_fulls, ratio); + g = g_new; + } + frame_scaled_scratch = has_fulls; + out.best_step = best_step; + std::vector positive; + for (double j : J) + if (j > 0.0) positive.push_back(j); + if (!positive.empty() && lambda > 0.0) + logger.Info("Per-frame scaling (fulls): smoothness chosen by cross-validation over the frames - " + "a frame's scale is carried over about {:.0f} frames", std::pow(lambda / median_of(positive), 0.25)); + LogScalingOutcome("fulls", out); return out; } @@ -5660,7 +5814,6 @@ RotationScaleMerge::Result RotationScaleMerge::Run(bool for_search, bool full_st rocking_event_frames_at_start = RockingEventFrames(&rocking_events_at_start); const double events_per_frame = n_frames > 0 ? static_cast(rocking_events_at_start) / n_frames : 0.0; const bool scale_partials = events_per_frame >= MIN_EVENTS_PER_FRAME_SCALE_PARTIALS; - pool_sparse_fulls = !scale_partials; if (!scale_partials) logger.Info("Per-frame scale from the fulls alone: {:.1f} rocking events per frame, under the {:.0f} " "a partial's scale needs to stay apart from the partiality model", events_per_frame, @@ -5676,7 +5829,7 @@ RotationScaleMerge::Result RotationScaleMerge::Run(bool for_search, bool full_st std::vector g_dev(n_frames); std::vector scaled_dev(n_frames); partial_loop = RunScalingLoop("partials", g_partial, frame_count, smooth_window, - /*release_restraint=*/false, [&] { + [&] { gpu_->ScalePartials(1, min_partiality, d_min_limit.has_value()); gpu_->GetG(g_dev.data(), scaled_dev.data()); for (int f = 0; f < n_frames; ++f) @@ -5689,9 +5842,28 @@ RotationScaleMerge::Result RotationScaleMerge::Run(bool for_search, bool full_st } #endif if (!scale_partials) { + // corr as the partial scaling would leave it at G = 1: the deterministic corrections, the + // incident flux and the partiality of the smoothed geometry. std::fill(g_partial.begin(), g_partial.end(), 1.0); frame_scaled_scratch.assign(n_frames, 1); partial_loop.converged = true; + bool unit_on_gpu = false; +#ifdef JFJOCH_USE_CUDA + if (gpu_active_) { + gpu_->SetUnitPartialScale(); + unit_on_gpu = true; + } +#endif + if (!unit_on_gpu) + ParallelChunks(static_cast(partials.size()), ThreadsForWork(partials.size(), nthreads), + [&](int lo, int hi) { + for (int i = lo; i < hi; ++i) { + Obs &o = partials[i]; + const double denom = static_cast(o.partiality); + o.corr = std::isfinite(o.prescaling_corr) && std::isfinite(denom) && denom > 0.0 + ? static_cast(o.prescaling_corr / denom) : NAN; + } + }); } else if (!scaled_on_gpu) { const PartialLoopKey key{x.GetSpaceGroupOrP1().xhm(), merge_friedel, d_min_limit, d_max_limit, min_partiality, smooth_window, scaling_iter}; @@ -5740,7 +5912,7 @@ RotationScaleMerge::Result RotationScaleMerge::Run(bool for_search, bool full_st std::vector g_fitted; std::vector frame_fitted; partial_loop = RunScalingLoop("partials", g_partial, frame_count, smooth_window, - /*release_restraint=*/false, [&] { + [&] { ReduceScalingGroupMeans(n_groups, scaling_refs, scaling_corr, partial_mean); FitPerFrameG(scaling_obs, frame_start, frame_count, partial_mean, /*unity=*/false, g_partial); g_fitted = g_partial; @@ -5978,15 +6150,14 @@ RotationScaleMerge::Result RotationScaleMerge::Run(bool for_search, bool full_st gpu_->ResetFullsScale(); if (static_cast(g_full.size()) != n_frames) g_full.assign(n_frames, 1.0); - std::vector g_dev(n_frames); + std::vector g_fit(n_frames), info(n_frames); std::vector scaled_dev(n_frames); - fulls_loop = RunScalingLoop("fulls", g_full, f_count, smooth_window, - /*release_restraint=*/true, [&] { - gpu_->ScaleFulls(1, min_partiality, pool_sparse_fulls); - gpu_->GetG(g_dev.data(), scaled_dev.data()); - for (int f = 0; f < n_frames; ++f) - if (scaled_dev[f]) g_full[f] = g_dev[f]; + fulls_loop = RunFullsScalingLoop(g_full, g_fit, f_count, [&] { + gpu_->FitFullsScale(min_partiality); + gpu_->GetG(g_fit.data(), scaled_dev.data()); + gpu_->GetInfo(info.data()); frame_scaled_scratch = scaled_dev; + frame_info_scratch = info; }, [&](const std::vector &apply, const std::vector &ratio) { gpu_->SmoothFullsCorr(apply.data(), ratio.data()); }); @@ -6041,12 +6212,10 @@ RotationScaleMerge::Result RotationScaleMerge::Run(bool for_search, bool full_st // --- 4. Scale the fulls (XDS order, Unity model). --- if (scale_fulls && !scaled_fulls_on_gpu) { - std::vector full_mean; - fulls_loop = RunScalingLoop("fulls", g_full, fulls_frame_count, smooth_window, - /*release_restraint=*/true, [&] { + std::vector full_mean, g_fit(n_frames); + fulls_loop = RunFullsScalingLoop(g_full, g_fit, fulls_frame_count, [&] { ReduceGroupMeans(fulls, n_groups, full_mean); - FitPerFrameG(fulls, fulls_frame_start, fulls_frame_count, full_mean, /*unity=*/true, g_full); - UpdateCorr(fulls, g_full, frame_scaled_scratch); + FitPerFrameG(fulls, fulls_frame_start, fulls_frame_count, full_mean, /*unity=*/true, g_fit); }, [&](const std::vector &apply, const std::vector &ratio) { for (auto &o : fulls) if (apply[o.frame] && std::isfinite(o.corr)) diff --git a/image_analysis/scale_merge/RotationScaleMerge.h b/image_analysis/scale_merge/RotationScaleMerge.h index 842ad5b65..2d9d75a0e 100644 --- a/image_analysis/scale_merge/RotationScaleMerge.h +++ b/image_analysis/scale_merge/RotationScaleMerge.h @@ -43,10 +43,18 @@ // Stills use the per-image ScaleOnTheFly (fixed partiality) instead. class RotationScaleMerge { public: - // How many frames on either side a frame's fulls scale is fitted over (used without partial - // scaling): none when the frame holds MIN_REFLECTIONS fulls itself, else the fewest that bring the - // pooled count there. - static int PoolHalfWidth(const std::vector &fcount, int f); + // The log scale x minimising sum_f J_f (x_f - y_f)^2 + lambda * sum_f (x_f - 2 x_{f+1} + x_{f+2})^2: + // each frame's own estimate y_f at its information J_f, held to a curve that bends only as far as + // the data ask. A frame with J_f = 0 takes the curve through its neighbours. Needs at least two + // frames with J_f > 0; returns y unchanged otherwise. + static std::vector SmoothLogScale(const std::vector &y, const std::vector &J, + double lambda); + // The lambda of SmoothLogScale that best predicts frames left out of the fit, over a fixed grid + // relative to the typical J. The frames are left out in blocks of `block` (five interleaved folds of + // them): a full sums the frames of its rocking curve, so neighbours closer than that share their + // measurement and would predict each other whatever the scale does. + static double ChooseLogScaleSmoothing(const std::vector &y, const std::vector &J, + int block); struct Result { std::vector merged; MergeStatistics statistics; @@ -223,9 +231,6 @@ private: // this is how far they may go before giving up and saying so. int scaling_iter = 100; bool scale_fulls = true; - // Without partial scaling the fulls carry the whole per-frame scale, and a sparse frame is fitted - // over its neighbours (PoolHalfWidth). - bool pool_sparse_fulls = false; bool refine_decay_b = false; // per-time-block Debye-Waller decay correction (radiation damage) int absorption_iter = 0; // >0: fit a goniometer-frame absorption surface over this many iterations int modulation_iter = 0; // >0: fit a detector-plane modulation (flat-field) surface, this many iterations @@ -392,8 +397,10 @@ private: std::vector em_fit_pool; std::vector em_chi2; - // Set by FitPerFrameG: which frames were fitted this call (so corr/G is updated only there). + // Set by FitPerFrameG: which frames were fitted this call (so corr/G is updated only there), and + // the information each fit had on G, sum w^2 c^2 (0 where the frame was not fitted). std::vector frame_scaled_scratch; + std::vector frame_info_scratch; // Per-frame mosaicity smoothed in frame order (deterministic); used to recompute partiality and // written back for the per-image scaling table. Empty if there is no per-frame mosaicity. @@ -442,7 +449,7 @@ private: // the per-frame incident flux, which the finiteness pass then folds into prescaling_corr so that // corr = prescaling_corr / (partiality * G) divides it out and G fits only the residual. See the .cpp for why a // background is a usable flux meter and what it costs when it is not. - void MeasureIncidentFlux(const std::vector &mean_bkg); + void MeasureIncidentFlux(const std::vector &mean_bkg, const std::vector &mean_bkg_info); // Build the flat `partials` array (and the per-frame CSR, the finiteness mask and `perm`) from the // source reflections, skipping the observations whose resolution can never be in range: @@ -512,7 +519,15 @@ private: // over `window` frames and pins the gauge after every iteration; stops when the scales settle, stop // settling, or the cap is reached. On return frame_scaled_scratch flags every frame fitted in any // iteration. - struct ScalingLoopOutcome { int iterations = 0; bool converged = false; bool stalled = false; double step = 0.0; }; + struct ScalingLoopOutcome { int iterations = 0; bool converged = false; bool stalled = false; double step = 0.0; + double best_step = 0.0; }; + void LogScalingOutcome(const char *what, const ScalingLoopOutcome &out) const; + // Shrink a free per-frame scale back toward the restrained one, frame by frame, by how much of + // the difference between them is real (see ShrinkToRestrained in the .cpp). + void ShrinkToRestrained(const char *what, const std::vector &g_restrained, + const std::vector &fitted, + const std::vector &frame_obs_count, + std::vector &g) const; // The CPU partial scaling loop of an earlier Run, for a later Run that would run it on the same // inputs - which the P1 merges of a pass do. The loop restarts from corr_ingested, fits each frame // from its observations and the group means alone, and writes G and corr only on the frames it @@ -539,21 +554,21 @@ private: static constexpr size_t MAX_PARTIAL_LOOP_MEMOS = 3; ScalingLoopOutcome RunScalingLoop(const char *what, std::vector &g, const std::vector &frame_obs_count, int window, - bool release_restraint, const std::function &iterate, const std::function &, const std::vector &)> &rescale); - // Shrink a free per-frame scale back toward the restrained one, frame by frame, by how much of - // the difference between them is real (see ShrinkToRestrained in the .cpp). - void ShrinkToRestrained(const char *what, const std::vector &g_restrained, - const std::vector &fitted, - const std::vector &frame_obs_count, - std::vector &g) const; - - // corr = prescaling_corr / (partiality * G[frame]); leaves corr unchanged for frames that could not be fit. - void UpdateCorr(std::vector &obs, const std::vector &g, - const std::vector &frame_scaled) const; + // The per-frame scale of the fulls, run to convergence: `fit` fits every frame once against the + // current reference without touching corr (leaving the free G in `g_fit`, the fitted frames in + // frame_scaled_scratch and their information in frame_info_scratch); the scale is then the + // penalised smooth of those fits (SmoothLogScale, lambda by ChooseLogScaleSmoothing), gauge-pinned, + // and `rescale` multiplies the corr of each frame's fulls by g_old / g_new. `g` is the scale corr + // stands at (corr = 1/g) on entry and on return. + ScalingLoopOutcome RunFullsScalingLoop(std::vector &g, std::vector &g_fit, + const std::vector &frame_obs_count, + const std::function &fit, + const std::function &, + const std::vector &)> &rescale); // The windowed geometric mean of G over frames; the scaling loop applies the resulting ratio to // the observations' corr (on the GPU, to the resident corr in a kernel). diff --git a/image_analysis/scale_merge/RotationScaleMergeGPU.cu b/image_analysis/scale_merge/RotationScaleMergeGPU.cu index 0ecf08707..5bdfb7695 100644 --- a/image_analysis/scale_merge/RotationScaleMergeGPU.cu +++ b/image_analysis/scale_merge/RotationScaleMergeGPU.cu @@ -113,63 +113,51 @@ namespace { // One block per frame: the per-frame scale G as the weighted least-squares slope over the frame's // contiguous obs - the same plain fit as the host SolveScale (which says why it is not robust). - // Leaves g/scaled untouched for under-populated frames. `perm` (null for the partials, whose - // arrays are already frame-contiguous) maps a position in the frame's [lo,hi) range to the obs - // index, so the same kernel scales the fulls (emit-ordered) through a frame-grouping permutation - // without physically reordering the fulls arrays. + // Leaves g/scaled untouched for frames with fewer than min_obs usable observations. `perm` (null + // for the partials, whose arrays are already frame-contiguous) maps a position in the frame's + // [lo,hi) range to the obs index, so the same kernel scales the fulls (emit-ordered) through a + // frame-grouping permutation without physically reordering the fulls arrays. info[f] receives the + // fit's information on G, sum w^2 c^2 (0 for a frame left alone). __global__ void FitPerFrameGKernel(int n_frames, const int32_t *__restrict__ frame_start, const int32_t *__restrict__ frame_count, const float *__restrict__ I, const double *__restrict__ inv_sigma, const float *__restrict__ sco_coeff, const uint8_t *__restrict__ sco_ok, - const int32_t *__restrict__ perm, bool pool, - double *__restrict__ g, uint8_t *__restrict__ scaled) { + const int32_t *__restrict__ perm, long min_obs, + double *__restrict__ g, uint8_t *__restrict__ scaled, + double *__restrict__ info) { const int f = blockIdx.x; if (f >= n_frames) return; - // pool (the fulls): a frame holding fewer than MIN_REFLECTIONS is fitted over the fewest frames - // on either side that bring the count there - RotationScaleMerge::PoolHalfWidth, same rule. - __shared__ int s_h; - if (threadIdx.x == 0) { - long total = frame_count[f]; - int h = 0; - while (pool && total < long(MIN_REFLECTIONS) && h < n_frames) { - ++h; - if (f - h >= 0) total += frame_count[f - h]; - if (f + h < n_frames) total += frame_count[f + h]; - } - s_h = h; - } - __syncthreads(); - const int j_lo = max(0, f - s_h), j_hi = min(n_frames - 1, f + s_h); + const int lo = frame_start[f], hi = frame_start[f] + frame_count[f]; __shared__ double sh[BLK]; + if (threadIdx.x == 0) info[f] = 0.0; long cnt_local = 0; - for (int j = j_lo; j <= j_hi; ++j) - for (int i = frame_start[j] + threadIdx.x; i < frame_start[j] + frame_count[j]; i += blockDim.x) - if (sco_ok[perm ? perm[i] : i]) ++cnt_local; + for (int i = lo + threadIdx.x; i < hi; i += blockDim.x) + if (sco_ok[perm ? perm[i] : i]) ++cnt_local; const double cnt = BlockReduceSum(double(cnt_local), sh); __shared__ double s_cnt; if (threadIdx.x == 0) s_cnt = cnt; __syncthreads(); - if (s_cnt < MIN_REFLECTIONS) return; // leave g[f]/scaled[f] as-is + if (s_cnt < min_obs) return; // leave g[f]/scaled[f] as-is double num = 0.0, den = 0.0; - for (int j = j_lo; j <= j_hi; ++j) - for (int i = frame_start[j] + threadIdx.x; i < frame_start[j] + frame_count[j]; i += blockDim.x) { - const int a = perm ? perm[i] : i; - if (!sco_ok[a]) continue; - const double coeff = sco_coeff[a]; - const double w = inv_sigma[a]; - const double w2 = w * w; - num += w2 * coeff * double(I[a]); - den += w2 * coeff * coeff; - } + for (int i = lo + threadIdx.x; i < hi; i += blockDim.x) { + const int a = perm ? perm[i] : i; + if (!sco_ok[a]) continue; + const double coeff = sco_coeff[a]; + const double w = inv_sigma[a]; + const double w2 = w * w; + num += w2 * coeff * double(I[a]); + den += w2 * coeff * coeff; + } const double tnum = BlockReduceSum(num, sh); __syncthreads(); const double tden = BlockReduceSum(den, sh); if (threadIdx.x == 0) { const double G = tden > 0.0 ? tnum / tden : NAN; g[f] = isfinite(G) ? fmax(0.0, G) : 1.0; scaled[f] = 1; + info[f] = tden; } } @@ -873,7 +861,7 @@ struct RotationScaleMergeGPU::Impl { // per space group CudaDevicePtr group, group_perm, group_start, group_count; // scratch - CudaDevicePtr group_mean, g; + CudaDevicePtr group_mean, g, info; CudaDevicePtr scaled; CudaDevicePtr inv_sigma; // 1/sigma, hoisted out of the IRLS loop (sigma never changes) CudaDevicePtr sco_coeff; @@ -1007,6 +995,7 @@ void RotationScaleMergeGPU::SetPartialsLayout(int n_obs, int n_frames, d.clipped = d.Alloc(n); d.g = d.Alloc(n_frames); d.scaled = d.Alloc(n_frames); + d.info = d.Alloc(n_frames); d.inv_sigma = d.Alloc(n_obs); d.sco_coeff = d.Alloc(n_obs); d.sco_ok = d.Alloc(n_obs); @@ -1096,7 +1085,8 @@ void RotationScaleMergeGPU::ScalePartials(int iters, double min_partiality, bool d.sco_coeff.get(), d.sco_ok.get()); CudaCheck(cudaGetLastError(), "PrepScaleObsKernel launch"); FitPerFrameGKernel<<s()>>>(d.n_frames, d.frame_start.get(), d.frame_count.get(), - d.I.get(), d.inv_sigma.get(), d.sco_coeff.get(), d.sco_ok.get(), nullptr, /*pool=*/false, d.g.get(), d.scaled.get()); + d.I.get(), d.inv_sigma.get(), d.sco_coeff.get(), d.sco_ok.get(), nullptr, long(MIN_REFLECTIONS), d.g.get(), d.scaled.get(), + d.info.get()); CudaCheck(cudaGetLastError(), "FitPerFrameGKernel launch"); UpdateCorrKernel<<s()>>>(d.n_obs, d.frame.get(), d.prescaling_corr.get(), d.partiality.get(), d.g.get(), d.scaled.get(), d.corr.get()); @@ -1111,6 +1101,26 @@ void RotationScaleMergeGPU::GetCorr(float *corr_out) const { cudaMemcpyDeviceToHost, impl_->s(), "download corr"); } +void RotationScaleMergeGPU::SetUnitPartialScale() { + DeviceGuard guard(impl_->device, impl_->available); + auto &d = *impl_; + const std::vector one(d.n_frames, 1.0); + const std::vector all(d.n_frames, 1); + d.Upload(d.g, one.data(), d.n_frames); + d.Upload(d.scaled, all.data(), d.n_frames); + const int blocks = std::min(65535, (d.n_obs + BLK - 1) / BLK); + UpdateCorrKernel<<s()>>>(d.n_obs, d.frame.get(), d.prescaling_corr.get(), + d.partiality.get(), d.g.get(), d.scaled.get(), d.corr.get()); + CudaCheck(cudaGetLastError(), "unit partial scale launch"); + CudaCheck(cudaStreamSynchronize(impl_->s()), "unit partial scale sync"); +} + +void RotationScaleMergeGPU::GetInfo(double *info_out) const { + DeviceGuard guard(impl_->device, impl_->available); + CopyAndWait(info_out, impl_->info.get(), size_t(impl_->n_frames) * sizeof(double), + cudaMemcpyDeviceToHost, impl_->s(), "download info"); +} + void RotationScaleMergeGPU::GetG(double *g_out, uint8_t *scaled_out) const { DeviceGuard guard(impl_->device, impl_->available); CopyAndWait(g_out, impl_->g.get(), size_t(impl_->n_frames) * sizeof(double), @@ -1512,37 +1522,32 @@ void RotationScaleMergeGPU::ResetFullsScale() { CudaCheck(cudaStreamSynchronize(impl_->s()), "reset fulls scale sync"); } -void RotationScaleMergeGPU::ScaleFulls(int iters, double min_partiality, bool pool) { +void RotationScaleMergeGPU::FitFullsScale(double min_partiality) { DeviceGuard guard(impl_->device, impl_->available); auto &d = *impl_; const int nf = d.n_fulls; if (nf == 0) return; - const int obs_blocks = std::min(65535, (nf + BLK - 1) / BLK); const int grp_blocks = std::min(65535, (d.n_groups + BLK - 1) / BLK); // Reset per call, as ScalePartials: the host keeps the G of a frame across calls. CudaCheck(cudaMemsetAsync(d.scaled.get(), 0, size_t(d.n_frames) * sizeof(uint8_t), impl_->s()), "memset f scaled"); CudaCheck(cudaMemsetAsync(d.g.get(), 0, size_t(d.n_frames) * sizeof(double), impl_->s()), "memset f g"); - for (int it = 0; it < iters; ++it) { - ReduceGroupMeansKernel<<s()>>>(d.n_groups, min_partiality, - d.f_gperm.get(), d.f_gstart.get(), d.f_gcount.get(), - d.f_I.get(), d.f_sigma.get(), d.f_partiality.get(), d.f_corr.get(), d.group_mean.get()); - CudaCheck(cudaGetLastError(), "ReduceGroupMeansKernel launch"); - // Not grid-stride, so its grid has to cover every full - unlike the grid-stride kernels - // below, which the 65535 cap is there for. Capped, it would silently leave the tail of - // sco_coeff/sco_ok stale above 16.8M fulls. - PrepScaleObsKernel<<<(nf + BLK - 1) / BLK, BLK, 0, impl_->s()>>>(nf, min_partiality, d.f_group.get(), d.f_partiality.get(), - d.f_rlp.get(), d.f_zeta.get(), d.f_on_ice.get(), d.group_mean.get(), - d.f_sigma.get(), d.f_inv_sigma.get(), d.f_sco_coeff.get(), d.f_sco_ok.get()); - CudaCheck(cudaGetLastError(), "PrepScaleObsKernel launch"); - FitPerFrameGKernel<<s()>>>(d.n_frames, - d.f_frame_start.get(), d.f_frame_count.get(), d.f_I.get(), d.f_inv_sigma.get(), - d.f_sco_coeff.get(), d.f_sco_ok.get(), d.f_frame_perm.get(), pool, d.g.get(), d.scaled.get()); - CudaCheck(cudaGetLastError(), "FitPerFrameGKernel launch"); - UpdateCorrKernel<<s()>>>(nf, d.f_frame.get(), d.f_rlp.get(), d.f_partiality.get(), - d.g.get(), d.scaled.get(), d.f_corr.get()); - } + ReduceGroupMeansKernel<<s()>>>(d.n_groups, min_partiality, + d.f_gperm.get(), d.f_gstart.get(), d.f_gcount.get(), + d.f_I.get(), d.f_sigma.get(), d.f_partiality.get(), d.f_corr.get(), d.group_mean.get()); + CudaCheck(cudaGetLastError(), "ReduceGroupMeansKernel launch"); + // Not grid-stride, so its grid has to cover every full - unlike the grid-stride kernels + // below, which the 65535 cap is there for. Capped, it would silently leave the tail of + // sco_coeff/sco_ok stale above 16.8M fulls. + PrepScaleObsKernel<<<(nf + BLK - 1) / BLK, BLK, 0, impl_->s()>>>(nf, min_partiality, d.f_group.get(), d.f_partiality.get(), + d.f_rlp.get(), d.f_zeta.get(), d.f_on_ice.get(), d.group_mean.get(), + d.f_sigma.get(), d.f_inv_sigma.get(), d.f_sco_coeff.get(), d.f_sco_ok.get()); + CudaCheck(cudaGetLastError(), "PrepScaleObsKernel launch"); + FitPerFrameGKernel<<s()>>>(d.n_frames, + d.f_frame_start.get(), d.f_frame_count.get(), d.f_I.get(), d.f_inv_sigma.get(), + d.f_sco_coeff.get(), d.f_sco_ok.get(), d.f_frame_perm.get(), 1L, d.g.get(), d.scaled.get(), + d.info.get()); CudaCheck(cudaGetLastError(), "scale fulls launch"); CudaCheck(cudaStreamSynchronize(impl_->s()), "scale fulls sync"); } diff --git a/image_analysis/scale_merge/RotationScaleMergeGPU.h b/image_analysis/scale_merge/RotationScaleMergeGPU.h index 98aa1b18a..0307ccbbf 100644 --- a/image_analysis/scale_merge/RotationScaleMergeGPU.h +++ b/image_analysis/scale_merge/RotationScaleMergeGPU.h @@ -61,10 +61,15 @@ public: // in place on the resident corr. Rotation model (partiality folded via the stored partiality). void ScalePartials(int iters, double min_partiality, bool has_d_min); + // corr = prescaling_corr / partiality on every partial: the partials at a per-frame scale of one. + void SetUnitPartialScale(); + // Copy the updated corr back to the host (length n_obs), and the fitted per-frame G (length // n_frames, as double) plus the per-frame "was fitted" flag. void GetCorr(float *corr_out) const; void GetG(double *g_out, uint8_t *scaled_out) const; + // The information each frame's last fit had on G, sum w^2 c^2 (length n_frames, 0 = not fitted). + void GetInfo(double *info_out) const; // Apply the smooth-G correction to the resident corr: corr[i] *= ratio[frame[i]] for frames with // apply[f] (both length n_frames), matching CPU SmoothG. Keeps corr resident (no round-trip). @@ -146,7 +151,7 @@ public: double min_captured_fraction, float max_frame_gap); // Download the combined fulls SoA (length = Combine()'s return). The working corr is downloaded - // separately by GetFullsCorr (it is only meaningful after ScaleFulls; otherwise the caller sets it). + // separately by GetFullsCorr (it is only meaningful after the fulls scaling; otherwise the caller sets it). void GetFulls(int32_t *h, int32_t *k, int32_t *l, float *I, float *sigma, float *d, float *image_number, int32_t *frame, uint8_t *on_ice, uint8_t *clipped, int32_t *group) const; @@ -178,20 +183,19 @@ public: const int32_t *gstart, const int32_t *gcount); // Start the Unity scaling of the resident fulls: corr, partiality, prescaling and zeta all 1, no - // frame fitted yet. Once before the ScaleFulls iterations. + // frame fitted yet. Once before the FitFullsScale rounds. void ResetFullsScale(); - // Run `iters` of the Unity scaling loop on the resident fulls (reduce group means -> per-frame LS G - // -> update corr), in place on the fulls' working corr. Requires SetFullsFrameCSR + SetFullsGroups - // and ResetFullsScale. - // pool: a frame with fewer than MIN_REFLECTIONS fulls is fitted over its neighbours - // (RotationScaleMerge::PoolHalfWidth). - void ScaleFulls(int iters, double min_partiality, bool pool); + // One Unity fit of the resident fulls (reduce group means -> per-frame LS G and its information, + // GetG / GetInfo), leaving the fulls' working corr as it is: the host smooths the fitted scales + // and applies them with SmoothFullsCorr. Requires SetFullsFrameCSR + SetFullsGroups and + // ResetFullsScale. + void FitFullsScale(double min_partiality); // The fulls' counterpart of SmoothCorr: f_corr[i] *= ratio[f_frame[i]] where apply[f]. void SmoothFullsCorr(const uint8_t *apply, const double *ratio); - // Download the fulls' working corr (length = n_fulls), valid after ScaleFulls. + // Download the fulls' working corr (length = n_fulls), valid after the fulls scaling. void GetFullsCorr(float *corr) const; // --- correction-surface fit (RotationScaleMerge::ApplyCellSurface) --- diff --git a/tests/RotationScaleWalkTest.cpp b/tests/RotationScaleWalkTest.cpp index f01334b3f..248bc9ecd 100644 --- a/tests/RotationScaleWalkTest.cpp +++ b/tests/RotationScaleWalkTest.cpp @@ -94,14 +94,43 @@ TEST_CASE("WalkRotationScale_StoredAnglesStand", "[RotationScale]") { } } -TEST_CASE("PoolHalfWidth_FitsSparseFramesOverTheirNeighbours", "[RotationScale]") { - // 20 fulls of its own: fitted on its own. - CHECK(RotationScaleMerge::PoolHalfWidth({0, 20, 0}, 1) == 0); - // 5 per frame: two frames either side bring 25 >= 20. - const std::vector five(11, 5); - CHECK(RotationScaleMerge::PoolHalfWidth(five, 5) == 2); - // At the end of the sweep the window grows on the one side there is. - CHECK(RotationScaleMerge::PoolHalfWidth(five, 0) == 3); - // A sweep that never holds enough stops at its length. - CHECK(RotationScaleMerge::PoolHalfWidth({1, 1, 1}, 1) == 3); +TEST_CASE("SmoothLogScale_FollowsInformationBridgesGaps", "[RotationScale]") { + const int n = 60; + // A ramp is no curvature, so any amount of smoothing keeps it exactly. + std::vector ramp(n), J(n, 1.0); + for (int f = 0; f < n; ++f) ramp[f] = -0.1 * f; + auto x = RotationScaleMerge::SmoothLogScale(ramp, J, 1e6); + for (int f = 0; f < n; ++f) CHECK(x[f] == Catch::Approx(ramp[f]).margin(1e-6)); + // A stretch with no information is bridged by the straight line through its neighbours. + std::vector Jgap(J); + for (int f = 20; f < 40; ++f) Jgap[f] = 0.0; + x = RotationScaleMerge::SmoothLogScale(ramp, Jgap, 1.0); + CHECK(x[30] == Catch::Approx(-3.0).margin(1e-6)); + // A frame with far more information than its neighbours keeps its own value. + std::vector y(n, 0.0), Jone(n, 1.0); + y[30] = 1.0; Jone[30] = 1e6; + x = RotationScaleMerge::SmoothLogScale(y, Jone, 10.0); + CHECK(x[30] == Catch::Approx(1.0).margin(1e-3)); + // Fewer than two frames with information: nothing to smooth against. + std::vector Jsingle(n, 0.0); + Jsingle[5] = 1.0; + CHECK(RotationScaleMerge::SmoothLogScale(y, Jsingle, 1.0) == y); +} + +TEST_CASE("ChooseLogScaleSmoothing_SmoothsNoiseFollowsSignal", "[RotationScale]") { + const int n = 400; + std::vector J(n, 1.0), noisy(n), step(n); + // Deterministic noise about a flat scale: the chosen curve is close to flat. + for (int f = 0; f < n; ++f) noisy[f] = 0.3 * std::sin(12.9898 * f) * std::cos(78.233 * f); + const double l_noise = RotationScaleMerge::ChooseLogScaleSmoothing(noisy, J, 1); + const auto flat = RotationScaleMerge::SmoothLogScale(noisy, J, l_noise); + double rms = 0.0; + for (double v : flat) rms += v * v; + CHECK(std::sqrt(rms / n) < 0.05); + // A precise slow wave is followed. + for (int f = 0; f < n; ++f) step[f] = 2.0 * std::sin(f / 30.0); + const double l_wave = RotationScaleMerge::ChooseLogScaleSmoothing(step, J, 1); + const auto wave = RotationScaleMerge::SmoothLogScale(step, J, l_wave); + CHECK(wave[47] == Catch::Approx(step[47]).margin(0.01)); + CHECK(l_wave < l_noise); }