From 4e177bdb2190743a0ceb40bdf2c6e1412a4cbee0 Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Thu, 15 Oct 2026 07:50:16 +0200 Subject: [PATCH] RotationScaleMerge: CPU partial scaling loop reads its working copy in order, output unchanged The loop's corr now lives in the group CSR's order next to the references (the partials in no group after them), so the group means read both sequentially instead of gathering corr through the permutation, and the corr update walks the same arrays. FitPerFrameG's filters, which do not change over the loop, are applied once into a per-frame ScalingFit list with the reference-independent factors precomputed (FitScalingFrames). ReferenceWeight is inline, and the group means are no longer pre-filled with NaN serially. Same values, same operations in the same order: p.hkl and the merged and unmerged MTZ data are byte-identical on lysozyme, cytochrome C, myoglobin, a small-molecule set and 4nwv. On 4nwv (Apple M6, 12 threads) the partial scaling iterations take 3-4x less time, ~18 s -> ~6 s over 168 iterations. Co-Authored-By: Claude Opus 5.5 --- .../scale_merge/RotationScaleMerge.cpp | 121 ++++++++++++------ .../scale_merge/RotationScaleMerge.h | 40 ++++-- 2 files changed, 109 insertions(+), 52 deletions(-) diff --git a/image_analysis/scale_merge/RotationScaleMerge.cpp b/image_analysis/scale_merge/RotationScaleMerge.cpp index 6840ec09e..08bbd1b51 100644 --- a/image_analysis/scale_merge/RotationScaleMerge.cpp +++ b/image_analysis/scale_merge/RotationScaleMerge.cpp @@ -275,6 +275,12 @@ 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 ScaleFromSums(double num, double den, double &info) { + const double G = den > 0.0 ? num / den : NAN; + info = den; + return std::isfinite(G) ? std::max(0.0, G) : 1.0; + } + double SolveScale(const std::vector &obs, double &info) { double num = 0.0, den = 0.0; for (const auto &o : obs) { @@ -282,9 +288,7 @@ namespace { num += w2 * o.coeff * o.Iobs; 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; + return ScaleFromSums(num, den, info); } double median_of(std::vector &v) { @@ -1512,18 +1516,6 @@ int RotationScaleMerge::ComputeAsuGroups(const HKLKeyGenerator &keygen) { return n_groups; } -double RotationScaleMerge::ReferenceWeight(int32_t group, float corr, float partiality, float I, float sigma, - double min_partiality, float &I_corr) { - if (group < 0) return 0.0; - if (!(corr > 0.0f) || !std::isfinite(corr)) return 0.0; - if (partiality < min_partiality) return 0.0; - const float Ic = I * corr; - const float sigma_corr = sigma * corr; - if (!std::isfinite(Ic) || !std::isfinite(sigma_corr) || sigma_corr <= 0.0f) return 0.0; - I_corr = Ic; - return 1.0 / (static_cast(sigma_corr) * sigma_corr); -} - void RotationScaleMerge::ReduceFullGroupMeans(int n_groups, std::vector &out_mean) const { // Inverse-variance per-group mean of I*corr = the merge reference (no cell mask: the scaling // reference is built without a reference cell - only the final merge applies it). A group per task @@ -1565,7 +1557,7 @@ void RotationScaleMerge::ReducePartialGroupMeans(int n_groups, std::vector &refs, const std::vector &corr, std::vector &out_mean) const { - out_mean.assign(n_groups, NAN); + out_mean.resize(n_groups); ParallelChunks(n_groups, ThreadsForWork(corr.size(), nthreads), [&](int glo, int ghi) { for (int g = glo; g < ghi; ++g) { double sw = 0.0, swI = 0.0; @@ -1573,13 +1565,12 @@ void RotationScaleMerge::ReduceScalingGroupMeans(int n_groups, const std::vector for (int p = lo; p < hi; ++p) { const ScalingRef &r = refs[p]; float I_corr = 0.0f; - const double w = ReferenceWeight(g, corr[partial_group_perm[p]], r.partiality, r.I, r.sigma, - min_partiality, I_corr); + const double w = ReferenceWeight(g, corr[p], r.partiality, r.I, r.sigma, min_partiality, I_corr); if (w <= 0.0) continue; sw += w; swI += w * I_corr; } - if (sw > 0.0) out_mean[g] = swI / sw; + out_mean[g] = sw > 0.0 ? swI / sw : NAN; } }); } @@ -1709,6 +1700,31 @@ void RotationScaleMerge::FitPerFrameG(const std::vector &obs, const std::vect frame_info_scratch = std::move(info); } +void RotationScaleMerge::FitScalingFrames(const std::vector &fit, const std::vector &fit_start, + const std::vector &group_mean_in, std::vector &g) { + const int n_fr = static_cast(fit_start.size()) - 1; + std::vector scaled(n_fr, 0); + std::vector info(n_fr, 0.0); + ParallelFor(n_fr, nthreads, [&](int f) { + double num = 0.0, den = 0.0; + size_t n = 0; + for (int i = fit_start[f]; i < fit_start[f + 1]; ++i) { + const ScalingFit &o = fit[i]; + const double mean = group_mean_in[o.group]; + if (!std::isfinite(mean)) continue; + const double coeff = o.coeff_per_mean * mean; + num += o.w2 * coeff * o.Iobs; + den += o.w2 * coeff * coeff; + ++n; + } + if (n < MIN_REFLECTIONS) return; + g[f] = ScaleFromSums(num, den, info[f]); + scaled[f] = 1; + }); + frame_scaled_scratch = std::move(scaled); + frame_info_scratch = std::move(info); +} + void RotationScaleMerge::ShrinkToRestrained(const char *what, const std::vector &g_restrained, const std::vector &fitted, const std::vector &frame_obs_count, @@ -6405,28 +6421,49 @@ RotationScaleMerge::Result RotationScaleMerge::RunArm(bool scale_partials, bool logger->Info("Per-frame scaling (partials): the same grouping and settings as an earlier merge, " "so its {} iterations stand", partial_loop.iterations); } else { - // The loop works on its own compact copy of the partials (ScalingObs in the header) and hands - // corr back to them when it is done. + // The loop works on its own compact copy of the partials (ScalingRef and ScalingFit in the + // header) and hands corr back to them when it is done. scaling_index[q] is the partial at + // position q of the copy: the CSR's group order, then the partials in no group. const int n_partials = static_cast(partials.size()); const size_t partial_threads = ThreadsForWork(partials.size(), nthreads); - std::vector scaling_obs(n_partials); + const int n_grouped = partial_group_start.empty() + ? 0 : partial_group_start.back() + partial_group_count.back(); + std::vector scaling_index(partial_group_perm.get(), partial_group_perm.get() + n_grouped); + for (int i = 0; i < n_partials; ++i) + if (partials[i].group < 0) scaling_index.push_back(i); + std::vector scaling_refs(n_partials); std::vector scaling_corr(n_partials); ParallelChunks(n_partials, partial_threads, [&](int lo, int hi) { - for (int i = lo; i < hi; ++i) { - const Obs &o = partials[i]; - scaling_obs[i] = {o.frame, o.group, o.I, o.sigma, o.partiality, o.prescaling_corr, o.zeta, - o.on_ice}; - scaling_corr[i] = o.corr; + for (int q = lo; q < hi; ++q) { + const Obs &o = partials[scaling_index[q]]; + scaling_refs[q] = {o.frame, o.I, o.sigma, o.partiality, o.prescaling_corr}; + scaling_corr[q] = o.corr; } }); - std::vector scaling_refs(partial_group_start.empty() - ? 0 : partial_group_start.back() + partial_group_count.back()); - ParallelChunks(n_groups, partial_threads, [&](int glo, int ghi) { - for (int g = glo; g < ghi; ++g) - for (int p = partial_group_start[g]; p < partial_group_start[g] + partial_group_count[g]; ++p) { - const ScalingObs &o = scaling_obs[partial_group_perm[p]]; - scaling_refs[p] = {o.I, o.sigma, o.partiality}; - } + + // FitPerFrameG's filters, all but the reference mean's, frame by frame. + auto fitted_obs = [&](const Obs &o) { + return o.group >= 0 && !o.on_ice && o.partiality >= min_partiality + && std::isfinite(o.zeta) && o.zeta > 0.0f; + }; + const int n_fit_frames = static_cast(frame_start.size()); + std::vector fit_start(n_fit_frames + 1, 0); + ParallelFor(n_fit_frames, nthreads, [&](int f) { + fit_start[f + 1] = static_cast(std::count_if( + partials.begin() + frame_start[f], partials.begin() + frame_start[f] + frame_count[f], + fitted_obs)); + }); + for (int f = 0; f < n_fit_frames; ++f) fit_start[f + 1] += fit_start[f]; + std::vector scaling_fit(fit_start.back()); + ParallelFor(n_fit_frames, nthreads, [&](int f) { + int32_t k = fit_start[f]; + for (int i = frame_start[f]; i < frame_start[f] + frame_count[f]; ++i) { + const Obs &o = partials[i]; + if (!fitted_obs(o)) continue; + const double weight = SafeInv(o.sigma, 1.0); + scaling_fit[k++] = {o.group, o.partiality * SafeInv(o.prescaling_corr, 1.0), + static_cast(o.I), weight * weight}; + } }); // RunScalingLoop rescales after every iteration before anything reads corr, so the iteration's @@ -6437,14 +6474,14 @@ RotationScaleMerge::Result RotationScaleMerge::RunArm(bool scale_partials, bool partial_loop = RunScalingLoop("partials", g_partial, frame_count, smooth_window, [&] { ReduceScalingGroupMeans(n_groups, scaling_refs, scaling_corr, partial_mean); - FitPerFrameG(scaling_obs, frame_start, frame_count, partial_mean, /*unity=*/false, g_partial); + FitScalingFrames(scaling_fit, fit_start, partial_mean, g_partial); g_fitted = g_partial; frame_fitted = frame_scaled_scratch; }, [&](const std::vector &apply, const std::vector &ratio) { ParallelChunks(n_partials, partial_threads, [&](int lo, int hi) { - for (int i = lo; i < hi; ++i) { - const ScalingObs &o = scaling_obs[i]; - float &corr = scaling_corr[i]; + for (int q = lo; q < hi; ++q) { + const ScalingRef &o = scaling_refs[q]; + float &corr = scaling_corr[q]; if (frame_fitted[o.frame]) { // UpdateCorr const double denom = static_cast(o.partiality) * g_fitted[o.frame]; if (std::isfinite(o.prescaling_corr) && std::isfinite(denom) && denom > 0.0) @@ -6459,7 +6496,11 @@ RotationScaleMerge::Result RotationScaleMerge::RunArm(bool scale_partials, bool }); ParallelChunks(n_partials, partial_threads, [&](int lo, int hi) { - for (int i = lo; i < hi; ++i) partials[i].corr = scaling_corr[i]; + for (int q = lo; q < hi; ++q) partials[scaling_index[q]].corr = scaling_corr[q]; + }); + // The memo keeps corr in the partials' own order. + ParallelChunks(n_partials, partial_threads, [&](int lo, int hi) { + for (int i = lo; i < hi; ++i) scaling_corr[i] = partials[i].corr; }); PartialLoopMemo m{key, partial_loop, frame_scaled_scratch, g_partial, std::move(scaling_corr)}; if (partial_loop_memos.size() == MAX_PARTIAL_LOOP_MEMOS) diff --git a/image_analysis/scale_merge/RotationScaleMerge.h b/image_analysis/scale_merge/RotationScaleMerge.h index 34baebcd9..cb77e956b 100644 --- a/image_analysis/scale_merge/RotationScaleMerge.h +++ b/image_analysis/scale_merge/RotationScaleMerge.h @@ -616,22 +616,35 @@ private: // The CPU partial scaling loop's working copy of the partials. It passes over them several times // per iteration and reads a handful of fields each time, so it works on these rather than on the - // eighty-byte Obs: ScalingObs holds, in the partials' own order, what FitPerFrameG and the corr - // update read; ScalingRef holds, in the CSR's group order, what the group means read besides corr, - // so that only corr is gathered. The same values, so the same arithmetic in the same order. - struct ScalingObs { - int32_t frame, group; - float I, sigma, partiality, prescaling_corr, zeta; - uint8_t on_ice; - }; + // eighty-byte Obs. ScalingRef holds what the group means and the corr update read, in the CSR's + // group order with the partials in no group after them, and corr is kept beside it in the same + // order - so neither pass gathers. ScalingFit holds, frame by frame, the observations + // FitScalingFrames fits: FitPerFrameG's filters, which do not change over the loop, applied once, + // and the factors of its sums that do not depend on the reference worked out once. The same + // values, so the same arithmetic in the same order. struct ScalingRef { - float I, sigma, partiality; + int32_t frame; + float I, sigma, partiality, prescaling_corr; + }; + struct ScalingFit { + int32_t group; + double coeff_per_mean; // partiality / prescaling_corr: FitPerFrameG's coeff for a mean of 1 + double Iobs, w2; }; // The observation's inverse-variance weight in the scaling reference, and its scaled intensity; // 0 when the observation is not in the reference at all. static double ReferenceWeight(int32_t group, float corr, float partiality, float I, float sigma, - double min_partiality, float &I_corr); + double min_partiality, float &I_corr) { + if (group < 0) return 0.0; + if (!(corr > 0.0f) || !std::isfinite(corr)) return 0.0; + if (partiality < min_partiality) return 0.0; + const float Ic = I * corr; + const float sigma_corr = sigma * corr; + if (!std::isfinite(Ic) || !std::isfinite(sigma_corr) || sigma_corr <= 0.0f) return 0.0; + I_corr = Ic; + return 1.0 / (static_cast(sigma_corr) * sigma_corr); + } static double ReferenceWeight(const Obs &o, double min_partiality, float &I_corr) { return ReferenceWeight(o.group, o.corr, o.partiality, o.I, o.sigma, min_partiality, I_corr); } @@ -640,17 +653,20 @@ private: void ReduceFullGroupMeans(int n_groups, std::vector &out_mean) const; // The same over the partials, using the CSR above - the same sums in the same order. void ReducePartialGroupMeans(int n_groups, std::vector &out_mean) const; - // The same again over the scaling loop's copy: refs in group order, corr in the partials' order. + // The same again over the scaling loop's copy, refs and corr both in group order. void ReduceScalingGroupMeans(int n_groups, const std::vector &refs, const std::vector &corr, std::vector &out_mean) const; // Robust per-frame G fit (IRLS, Cauchy k=3), unity=false uses the rotation partiality, unity=true the // scale-fulls (partiality already folded in). Reads out_mean[group] as the reference intensity. - // O is Obs or ScalingObs. template void FitPerFrameG(const std::vector &obs, const std::vector &fstart, const std::vector &fcount, const std::vector &group_mean_in, bool unity, std::vector &g); + // FitPerFrameG(unity = false) over the scaling loop's ScalingFit, frame f's at + // [fit_start[f], fit_start[f + 1]). + void FitScalingFrames(const std::vector &fit, const std::vector &fit_start, + const std::vector &group_mean_in, std::vector &g); // The alternating per-frame scale fit run to convergence: `iterate` fits every frame once against // the current reference (leaving G in `g` and the frames it fitted in frame_scaled_scratch),