From 2e5c0df0a58ca9dea29d4d501346eb7e439adb92 Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Wed, 9 Sep 2026 16:17:46 +0200 Subject: [PATCH] scaling: the decay / relative-B passes read a compact array, and their per-group references run on all threads The three step-4b passes - the global decay slope, the per-batch relative-B and the radiation-damage monitor - were plain serial walks over the eighty-byte observation records, a dozen of them, each reading twenty bytes of each record and each allocating and zeroing a pair of per-group vectors. The decay fit alone runs seven of those passes (three slope fits, four held-out scores) on one thread while the workers and the GPU idle. Give them the treatment the correction surfaces below them already have: pull the fields they read out of the records once into a compact array in fulls order, and keep a second copy of the same terms in ASU-group order (stable counting sort) so a thread can own whole groups. The per-group inverse-variance references then run on all threads and still add each group's terms in fulls order, and the remaining scalar reductions keep their serial order over the compact array - so every sum is formed from the same terms in the same order as before, and the numbers are unchanged rather than merely close. The build is shared by the three curve fits and the three held-out scores of the relative-B pass instead of being repeated per call. Co-Authored-By: Claude Opus 5 (1M context) --- .../scale_merge/RotationScaleMerge.cpp | 224 ++++++++++++------ .../scale_merge/RotationScaleMerge.h | 16 +- 2 files changed, 173 insertions(+), 67 deletions(-) diff --git a/image_analysis/scale_merge/RotationScaleMerge.cpp b/image_analysis/scale_merge/RotationScaleMerge.cpp index fecb4c1cc..f18e20aa8 100644 --- a/image_analysis/scale_merge/RotationScaleMerge.cpp +++ b/image_analysis/scale_merge/RotationScaleMerge.cpp @@ -1214,6 +1214,46 @@ void RotationScaleMerge::UpdateCorr(std::vector &obs, const std::vector= 0 && o.corr > 0.0f && std::isfinite(o.corr) && (o.d > 0.0f) + && o.partiality >= min_partiality; + }; + const size_t nf = fulls.size(); + const int nch = static_cast(ThreadsForWork(nf, nthreads)); + auto cpart = [nf, nch](int c) { return nf * static_cast(c) / static_cast(nch); }; + std::vector off(nch + 1, 0); + ParallelChunks(nch, nthreads, [&](int clo, int chi) { + for (int c = clo; c < chi; ++c) { + size_t keep = 0; + for (size_t i = cpart(c); i < cpart(c + 1); ++i) + if (usable(fulls[i])) ++keep; + off[c + 1] = keep; + } + }); + for (int c = 0; c < nch; ++c) off[c + 1] += off[c]; + t.term.resize(off[nch]); + ParallelChunks(nch, nthreads, [&](int clo, int chi) { + for (int c = clo; c < chi; ++c) { + size_t w = off[c]; + for (size_t i = cpart(c); i < cpart(c + 1); ++i) { + const Obs &o = fulls[i]; + if (!usable(o)) continue; + t.term[w++] = {o.I, o.sigma, o.corr, o.d, o.image_number, o.frame, o.group}; + } + } + }); + t.g_start.assign(n_groups + 1, 0); + for (const DecayTerm &d : t.term) ++t.g_start[d.group + 1]; + for (int g = 0; g < n_groups; ++g) t.g_start[g + 1] += t.g_start[g]; + t.gterm.resize(t.term.size()); + std::vector fill(t.g_start.begin(), t.g_start.end() - 1); + for (const DecayTerm &d : t.term) t.gterm[fill[d.group]++] = d; +} + void RotationScaleMerge::RefineDecay(int n_groups) { // Radiation damage weakens later frames more at higher resolution - a resolution x time (Debye-Waller) // systematic the resolution-flat per-frame G leaves in. Model it as a single global relative-B rate: @@ -1236,22 +1276,37 @@ void RotationScaleMerge::RefineDecay(int n_groups) { return o.group >= 0 && o.corr > 0.0f && std::isfinite(o.corr) && (o.d > 0.0f) && o.partiality >= min_partiality; }; - auto decay_factor = [&](const Obs &o, double slope) { + auto decay_factor = [&](const auto &o, double slope) { return std::exp(slope * (o.image_number - fcenter) * s2_of(o.d)); }; + // The fields the seven passes below read (three slope fits, four held-out scores), pulled out of the + // fat records once, with the group-ordered copy that lets the per-group references run on all + // threads. See BuildDecayTerms: the compacted array is walked in fulls order and a group's terms are + // still summed in fulls order, so every sum here is formed exactly as the serial walks formed it. + DecayTerms t; + BuildDecayTerms(n_groups, t); + const int nt = static_cast(ThreadsForWork(t.term.size(), nthreads)); + std::vector sw(n_groups), swI(n_groups); + // Fit the global slope over the subset {frame&1 == parity} (parity < 0 = all fulls), using an // inverse-variance reference built from that same subset. slope*factor applied to the current corr. auto fit_slope = [&](int parity) -> double { - std::vector sw(n_groups, 0.0), swI(n_groups, 0.0); - for (const auto &o : fulls) { - if (!usable(o) || (parity >= 0 && (o.frame & 1) != parity)) continue; - const double sc = static_cast(o.sigma) * o.corr, w = 1.0 / (sc * sc); - sw[o.group] += w; swI[o.group] += w * static_cast(o.I) * o.corr; - } + ParallelChunks(n_groups, nt, [&](int glo, int ghi) { + for (int g = glo; g < ghi; ++g) { + double s_w = 0.0, s_wI = 0.0; + for (int k = t.g_start[g]; k < t.g_start[g + 1]; ++k) { + const DecayTerm &o = t.gterm[k]; + if (parity >= 0 && (o.frame & 1) != parity) continue; + const double sc = static_cast(o.sigma) * o.corr, w = 1.0 / (sc * sc); + s_w += w; s_wI += w * static_cast(o.I) * o.corr; + } + sw[g] = s_w; swI[g] = s_wI; + } + }); double Sw = 0, Sx = 0, Sy = 0, Sxx = 0, Sxy = 0; - for (const auto &o : fulls) { - if (!usable(o) || (parity >= 0 && (o.frame & 1) != parity) || sw[o.group] <= 0.0) continue; + for (const DecayTerm &o : t.term) { + if ((parity >= 0 && (o.frame & 1) != parity) || sw[o.group] <= 0.0) continue; const double Iref = swI[o.group] / sw[o.group]; const double Is = static_cast(o.I) * o.corr, sc = static_cast(o.sigma) * o.corr; if (!std::isfinite(Iref) || Iref <= 0.0 || !(Is > 0.0)) continue; @@ -1272,17 +1327,23 @@ void RotationScaleMerge::RefineDecay(int n_groups) { // absorption / modulation gate in ApplyCellSurface (which was switched off studentized chi^2 for the // same reason); the reference is still the inverse-variance mean. auto subset_disagreement = [&](int parity, double slope) -> double { - std::vector sw(n_groups, 0.0), swI(n_groups, 0.0); - for (const auto &o : fulls) { - if (!usable(o) || (o.frame & 1) != parity) continue; - const double Is = static_cast(o.I) * o.corr * decay_factor(o, slope); - const double sc = static_cast(o.sigma) * o.corr * decay_factor(o, slope); - const double w = 1.0 / (sc * sc); - sw[o.group] += w; swI[o.group] += w * Is; - } + ParallelChunks(n_groups, nt, [&](int glo, int ghi) { + for (int g = glo; g < ghi; ++g) { + double s_w = 0.0, s_wI = 0.0; + for (int k = t.g_start[g]; k < t.g_start[g + 1]; ++k) { + const DecayTerm &o = t.gterm[k]; + if ((o.frame & 1) != parity) continue; + const double Is = static_cast(o.I) * o.corr * decay_factor(o, slope); + const double sc = static_cast(o.sigma) * o.corr * decay_factor(o, slope); + const double w = 1.0 / (sc * sc); + s_w += w; s_wI += w * Is; + } + sw[g] = s_w; swI[g] = s_wI; + } + }); double num = 0.0, den = 0.0; - for (const auto &o : fulls) { - if (!usable(o) || (o.frame & 1) != parity || sw[o.group] <= 0.0) continue; + for (const DecayTerm &o : t.term) { + if ((o.frame & 1) != parity || sw[o.group] <= 0.0) continue; const double Iref = swI[o.group] / sw[o.group]; const double Is = static_cast(o.I) * o.corr * decay_factor(o, slope); if (!std::isfinite(Iref) || Iref <= 0.0) continue; @@ -1312,9 +1373,11 @@ void RotationScaleMerge::RefineDecay(int n_groups) { total_delta_B, 100.0 * gain / std::max(base, 1e-30)); return; } - for (auto &o : fulls) - if (usable(o)) - o.corr = static_cast(o.corr * decay_factor(o, slope)); + ParallelChunks(static_cast(fulls.size()), ThreadsForWork(fulls.size(), nthreads), [&](int lo, int hi) { + for (int i = lo; i < hi; ++i) + if (usable(fulls[i])) + fulls[i].corr = static_cast(fulls[i].corr * decay_factor(fulls[i], slope)); + }); logger.Info("Decay correction: total relative-B = {:.2f} A^2 over run (dB/dframe = {:.2e}, cross-validated)", 0.5 * slope * n_frames, 0.5 * slope); } @@ -1366,8 +1429,8 @@ std::vector RotationScaleMerge::SolveCurvatureSmoothedB(const std::vecto return b; } -std::vector RotationScaleMerge::FitRelativeBCurve(int n_groups, int n_batch, int frames_per_batch, - int gparity) const { +std::vector RotationScaleMerge::FitRelativeBCurve(const DecayTerms &t, int n_groups, int n_batch, + int frames_per_batch, int gparity) const { // Fit one relative-B per batch over the ASU-group subset {group&1 == gparity} (gparity < 0 = all fulls), // against an inverse-variance reference built from that same subset (no leakage). Per batch the model is // y = ln(I_ref / I_obs) = b * s^2 through the origin (the resolution-flat intercept is the per-frame G, @@ -1376,21 +1439,24 @@ std::vector RotationScaleMerge::FitRelativeBCurve(int n_groups, int n_ba // with overall scale, and cannot change equivalent agreement (within one ASU group every equivalent sits // at the same s^2, so only b VARYING between batches moves the merge). Drives the per-batch correction. auto s2_of = [](float d) { return d > 0.0f ? 1.0 / (4.0 * static_cast(d) * d) : 0.0; }; - auto usable = [&](const Obs &o) { - return o.group >= 0 && o.corr > 0.0f && std::isfinite(o.corr) && (o.d > 0.0f) - && o.partiality >= min_partiality; - }; - auto batch_of = [&](const Obs &o) { return std::min(n_batch - 1, o.frame / frames_per_batch); }; + auto batch_of = [&](const DecayTerm &o) { return std::min(n_batch - 1, o.frame / frames_per_batch); }; - std::vector sw(n_groups, 0.0), swI(n_groups, 0.0); - for (const auto &o : fulls) { - if (!usable(o) || (gparity >= 0 && (o.group & 1) != gparity)) continue; - const double sc = static_cast(o.sigma) * o.corr, w = 1.0 / (sc * sc); - sw[o.group] += w; swI[o.group] += w * static_cast(o.I) * o.corr; - } + std::vector sw(n_groups), swI(n_groups); + ParallelChunks(n_groups, static_cast(ThreadsForWork(t.term.size(), nthreads)), [&](int glo, int ghi) { + for (int g = glo; g < ghi; ++g) { + double s_w = 0.0, s_wI = 0.0; + if (gparity < 0 || (g & 1) == gparity) + for (int k = t.g_start[g]; k < t.g_start[g + 1]; ++k) { + const DecayTerm &o = t.gterm[k]; + const double sc = static_cast(o.sigma) * o.corr, w = 1.0 / (sc * sc); + s_w += w; s_wI += w * static_cast(o.I) * o.corr; + } + sw[g] = s_w; swI[g] = s_wI; + } + }); std::vector num(n_batch, 0.0), den(n_batch, 0.0); - for (const auto &o : fulls) { - if (!usable(o) || (gparity >= 0 && (o.group & 1) != gparity) || sw[o.group] <= 0.0) continue; + for (const DecayTerm &o : t.term) { + if ((gparity >= 0 && (o.group & 1) != gparity) || sw[o.group] <= 0.0) continue; const double Iref = swI[o.group] / sw[o.group]; const double Is = static_cast(o.I) * o.corr, sc = static_cast(o.sigma) * o.corr; if (!std::isfinite(Iref) || Iref <= 0.0 || !(Is > 0.0)) continue; @@ -1436,16 +1502,29 @@ void RotationScaleMerge::MeasureRadiationDamageB(int n_groups) { return o.group >= 0 && o.corr > 0.0f && std::isfinite(o.corr) && (o.d > 0.0f) && o.partiality >= min_partiality; }; - auto batch_of = [&](const Obs &o) { return std::min(n_batch - 1, o.frame / frames_per_batch); }; + auto batch_of = [&](const DecayTerm &o) { return std::min(n_batch - 1, o.frame / frames_per_batch); }; + + // The fields the three walks below read, pulled out of the fat records once, with the group-ordered + // copy the low-dose reference runs over (BuildDecayTerms). + DecayTerms t; + BuildDecayTerms(n_groups, t); + const int nt = static_cast(ThreadsForWork(t.term.size(), nthreads)); // Low-dose reference: inverse-variance group means over the first ~10% of frames (>= 1 batch). const int ref_frames = std::max(frames_per_batch, n_frames / 10); - std::vector sw0(n_groups, 0.0), swI0(n_groups, 0.0); - for (const auto &o : fulls) { - if (!usable(o) || o.frame >= ref_frames) continue; - const double sc = static_cast(o.sigma) * o.corr, w = 1.0 / (sc * sc); - sw0[o.group] += w; swI0[o.group] += w * static_cast(o.I) * o.corr; - } + std::vector sw0(n_groups), swI0(n_groups); + ParallelChunks(n_groups, nt, [&](int glo, int ghi) { + for (int g = glo; g < ghi; ++g) { + double s_w = 0.0, s_wI = 0.0; + for (int k = t.g_start[g]; k < t.g_start[g + 1]; ++k) { + const DecayTerm &o = t.gterm[k]; + if (o.frame >= ref_frames) continue; + const double sc = static_cast(o.sigma) * o.corr, w = 1.0 / (sc * sc); + s_w += w; s_wI += w * static_cast(o.I) * o.corr; + } + sw0[g] = s_w; swI0[g] = s_wI; + } + }); // Resolution shells of equal occupancy in s^2, from a strided sample of the fulls - the per-batch fit // regresses on the shell means, and an equal-WIDTH grid leaves the low-resolution shells (the ones that // still carry signal in a weak batch) with almost no reflections. @@ -1472,8 +1551,7 @@ void RotationScaleMerge::MeasureRadiationDamageB(int n_groups) { // all. Pooled over the whole run a shell either has signal or it has none, so walk out to the first // that has none and lay the shells inside that. std::vector rw(MONITOR_SHELLS, 0.0), rwI(MONITOR_SHELLS, 0.0); - for (const auto &o : fulls) { - if (!usable(o)) continue; + for (const DecayTerm &o : t.term) { const double sc = static_cast(o.sigma) * o.corr, w = 1.0 / (sc * sc); const int k = shell_of(s2_of(o.d)); rw[k] += w; rwI[k] += w * static_cast(o.I) * o.corr; @@ -1498,8 +1576,8 @@ void RotationScaleMerge::MeasureRadiationDamageB(int n_groups) { // its own I/sigma - which is what says whether the batch can be measured at all. const int n_cell = n_batch * MONITOR_SHELLS; std::vector sw(n_cell, 0.0), swI(n_cell, 0.0), swR(n_cell, 0.0), sws2(n_cell, 0.0); - for (const auto &o : fulls) { - if (!usable(o) || sw0[o.group] <= 0.0) continue; + for (const DecayTerm &o : t.term) { + if (sw0[o.group] <= 0.0) continue; const double Iref = swI0[o.group] / sw0[o.group]; if (!std::isfinite(Iref)) continue; const double s2 = s2_of(o.d), sc = static_cast(o.sigma) * o.corr, w = 1.0 / (sc * sc); @@ -1848,25 +1926,37 @@ void RotationScaleMerge::RefineRelativeB(int n_groups) { return o.group >= 0 && o.corr > 0.0f && std::isfinite(o.corr) && (o.d > 0.0f) && o.partiality >= min_partiality; }; - auto batch_of = [&](const Obs &o) { return std::min(n_batch - 1, o.frame / frames_per_batch); }; - auto b_factor = [&](const Obs &o, const std::vector &b) { + auto batch_of = [&](const auto &o) { return std::min(n_batch - 1, o.frame / frames_per_batch); }; + auto b_factor = [&](const auto &o, const std::vector &b) { return std::exp(b[batch_of(o)] * s2_of(o.d)); }; + // The fields the three curve fits and the three held-out scores below read, pulled out of the fat + // records once, with the group-ordered copy their per-group references run over (BuildDecayTerms). + DecayTerms t; + BuildDecayTerms(n_groups, t); + const int nt = static_cast(ThreadsForWork(t.term.size(), nthreads)); + std::vector sw(n_groups), swI(n_groups); // Held-out disagreement over the group subset {group&1 == gparity} with `b` applied: the same // sigma-INDEPENDENT Rmeas-like metric as RefineDecay (sum|Is-Iref|/sum|Iref|), so a relative-B cannot // "win" by reshaping sigma. Reference built from the scored subset (no leakage). auto subset_disagreement = [&](int gparity, const std::vector &b) -> double { - std::vector sw(n_groups, 0.0), swI(n_groups, 0.0); - for (const auto &o : fulls) { - if (!usable(o) || (o.group & 1) != gparity) continue; - const double f = b_factor(o, b); - const double Is = static_cast(o.I) * o.corr * f, sc = static_cast(o.sigma) * o.corr * f; - const double w = 1.0 / (sc * sc); - sw[o.group] += w; swI[o.group] += w * Is; - } + ParallelChunks(n_groups, nt, [&](int glo, int ghi) { + for (int g = glo; g < ghi; ++g) { + double s_w = 0.0, s_wI = 0.0; + if ((g & 1) == gparity) + for (int k = t.g_start[g]; k < t.g_start[g + 1]; ++k) { + const DecayTerm &o = t.gterm[k]; + const double f = b_factor(o, b); + const double Is = static_cast(o.I) * o.corr * f, sc = static_cast(o.sigma) * o.corr * f; + const double w = 1.0 / (sc * sc); + s_w += w; s_wI += w * Is; + } + sw[g] = s_w; swI[g] = s_wI; + } + }); double num = 0.0, den = 0.0; - for (const auto &o : fulls) { - if (!usable(o) || (o.group & 1) != gparity || sw[o.group] <= 0.0) continue; + for (const DecayTerm &o : t.term) { + if ((o.group & 1) != gparity || sw[o.group] <= 0.0) continue; const double Iref = swI[o.group] / sw[o.group]; const double Is = static_cast(o.I) * o.corr * b_factor(o, b); if (!std::isfinite(Iref) || Iref <= 0.0) continue; @@ -1879,7 +1969,7 @@ void RotationScaleMerge::RefineRelativeB(int n_groups) { // equivalents (same s^2, different batch) into extra scatter - net harmful, exactly like RefineDecay's // sub-floor slope. (RELATIVE_B_MIN_SPREAD ~ DECAY_MIN_DELTA_B.) constexpr double RELATIVE_B_MIN_SPREAD = 2.0; // A^2, minimum peak-to-peak relative-B to engage - const std::vector b_all = FitRelativeBCurve(n_groups, n_batch, frames_per_batch, -1); + const std::vector b_all = FitRelativeBCurve(t, n_groups, n_batch, frames_per_batch, -1); const double bmin = *std::min_element(b_all.begin(), b_all.end()); const double bmax = *std::max_element(b_all.begin(), b_all.end()); if (bmax - bmin < RELATIVE_B_MIN_SPREAD) { @@ -1891,17 +1981,19 @@ void RotationScaleMerge::RefineRelativeB(int n_groups) { // per-batch structure generalises across the equivalent split; over-fit per-batch noise does not. const std::vector zero(n_batch, 0.0); const double base = subset_disagreement(0, zero) + subset_disagreement(1, zero); - const std::vector b_even = FitRelativeBCurve(n_groups, n_batch, frames_per_batch, 0); - const std::vector b_odd = FitRelativeBCurve(n_groups, n_batch, frames_per_batch, 1); + const std::vector b_even = FitRelativeBCurve(t, n_groups, n_batch, frames_per_batch, 0); + const std::vector b_odd = FitRelativeBCurve(t, n_groups, n_batch, frames_per_batch, 1); const double gain = base - (subset_disagreement(1, b_even) + subset_disagreement(0, b_odd)); if (!(gain > CV_MIN_RELATIVE_GAIN * base)) { logger.Info("Relative-B: not cross-validated ({} batches, peak-to-peak {:.2f} A^2, held-out gain {:.1f}%, skipped)", n_batch, bmax - bmin, 100.0 * gain / std::max(base, 1e-30)); return; } - for (auto &o : fulls) - if (usable(o)) - o.corr = static_cast(o.corr * b_factor(o, b_all)); + ParallelChunks(static_cast(fulls.size()), ThreadsForWork(fulls.size(), nthreads), [&](int lo, int hi) { + for (int i = lo; i < hi; ++i) + if (usable(fulls[i])) + fulls[i].corr = static_cast(fulls[i].corr * b_factor(fulls[i], b_all)); + }); logger.Info("Relative-B: {} batches of ~{:.1f} deg, peak-to-peak {:.2f} A^2, held-out gain {:.1f}% (cross-validated)", n_batch, relative_b_deg, bmax - bmin, 100.0 * gain / std::max(base, 1e-30)); } diff --git a/image_analysis/scale_merge/RotationScaleMerge.h b/image_analysis/scale_merge/RotationScaleMerge.h index ecfd6cbde..1d7879617 100644 --- a/image_analysis/scale_merge/RotationScaleMerge.h +++ b/image_analysis/scale_merge/RotationScaleMerge.h @@ -392,6 +392,19 @@ private: // physical total-dB floor). RefineAbsorption fits a smooth factor over the diffracted-beam direction in // the goniometer frame (path-length / absorption; negligible at hard X-rays, matters at low energy). void RefineDecay(int n_groups); + // The observation fields the decay / relative-B passes read, pulled out of `fulls` once. Those + // passes run a dozen walks over the eighty-byte records for twenty bytes of each, and their + // per-group references had to stay serial because they scatter into a per-group array. `term` keeps + // fulls order, so a sum formed over it is formed from the same terms in the same order the serial + // walk used; `gterm` holds the same terms in ASU-group order, by a stable counting sort over `term`, + // so a group's terms are still added in fulls order - which lets a thread own whole groups and gives + // the same per-group sums. `g_start` is that array's CSR. + struct DecayTerm { float I, sigma, corr, d, image_number; int32_t frame, group; }; + struct DecayTerms { + std::vector term, gterm; + std::vector g_start; + }; + void BuildDecayTerms(int n_groups, DecayTerms &t) const; // Solve a smooth per-batch relative-B from the per-batch normal equations for b (num_c, den_c): // data-fidelity + a second-difference (curvature) penalty, by Gauss-Seidel, each batch clamped to // +-b_max. Returns the un-anchored curve; the caller sets the gauge and the clamp it can live with. @@ -401,7 +414,8 @@ private: // Fit a smoothed per-batch relative-B curve (A^2 per batch) on the fulls over the ASU-group subset // {group&1==gparity} (gparity<0 = all): the weighted s^2 slope of ln(Iref/Iobs) per batch against a // subset-global reference, smoothed and zero-mean-anchored. Drives the per-batch correction. - std::vector FitRelativeBCurve(int n_groups, int n_batch, int frames_per_batch, int gparity) const; + std::vector FitRelativeBCurve(const DecayTerms &t, int n_groups, int n_batch, + int frames_per_batch, int gparity) const; // Radiation-damage MONITOR (report-only): measure the per-batch relative-B on the scaled fulls before // any decay correction and store the first->last relative-B change + the per-batch curve on this object // (copied into the result statistics by MergeAndStats, then printed / logged / written to the mmCIF).