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 <noreply@anthropic.com>

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<ScaleObs> &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;
+        return ScaleFromSums(num, den, info);
-        info = den;
-        return std::isfinite(G) ? std::max(0.0, G) : 1.0;
     }
 
     double median_of(std::vector<double> &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<double>(sigma_corr) * sigma_corr);
-}
-
 void RotationScaleMerge::ReduceFullGroupMeans(int n_groups, std::vector<double> &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<doubl
 void RotationScaleMerge::ReduceScalingGroupMeans(int n_groups, const std::vector<ScalingRef> &refs,
                                                  const std::vector<float> &corr,
                                                  std::vector<double> &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,
+                const double w = ReferenceWeight(g, corr[p], r.partiality, r.I, r.sigma, min_partiality, I_corr);
-                                                 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<O> &obs, const std::vect
     frame_info_scratch = std::move(info);
 }
 
+void RotationScaleMerge::FitScalingFrames(const std::vector<ScalingFit> &fit, const std::vector<int32_t> &fit_start,
+                                          const std::vector<double> &group_mean_in, std::vector<double> &g) {
+    const int n_fr = static_cast<int>(fit_start.size()) - 1;
+    std::vector<uint8_t> scaled(n_fr, 0);
+    std::vector<double> 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<double> &g_restrained,
                                             const std::vector<uint8_t> &fitted,
                                             const std::vector<int32_t> &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
+            // The loop works on its own compact copy of the partials (ScalingRef and ScalingFit in the
-            // corr back to them when it is done.
+            // 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<int>(partials.size());
             const size_t partial_threads = ThreadsForWork(partials.size(), nthreads);
-            std::vector<ScalingObs> scaling_obs(n_partials);
+            const int n_grouped = partial_group_start.empty()
+                                  ? 0 : partial_group_start.back() + partial_group_count.back();
+            std::vector<int32_t> 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<ScalingRef> scaling_refs(n_partials);
             std::vector<float> scaling_corr(n_partials);
             ParallelChunks(n_partials, partial_threads, [&](int lo, int hi) {
-                for (int i = lo; i < hi; ++i) {
+                for (int q = lo; q < hi; ++q) {
-                    const Obs &o = partials[i];
+                    const Obs &o = partials[scaling_index[q]];
-                    scaling_obs[i] = {o.frame, o.group, o.I, o.sigma, o.partiality, o.prescaling_corr, o.zeta,
+                    scaling_refs[q] = {o.frame, o.I, o.sigma, o.partiality, o.prescaling_corr};
-                                      o.on_ice};
+                    scaling_corr[q] = o.corr;
-                    scaling_corr[i] = o.corr;
                 }
             });
-            std::vector<ScalingRef> scaling_refs(partial_group_start.empty()
+
-                                                 ? 0 : partial_group_start.back() + partial_group_count.back());
+            // FitPerFrameG's filters, all but the reference mean's, frame by frame.
-            ParallelChunks(n_groups, partial_threads, [&](int glo, int ghi) {
+            auto fitted_obs = [&](const Obs &o) {
-                for (int g = glo; g < ghi; ++g)
+                return o.group >= 0 && !o.on_ice && o.partiality >= min_partiality
-                    for (int p = partial_group_start[g]; p < partial_group_start[g] + partial_group_count[g]; ++p) {
+                       && std::isfinite(o.zeta) && o.zeta > 0.0f;
-                        const ScalingObs &o = scaling_obs[partial_group_perm[p]];
+            };
-                        scaling_refs[p] = {o.I, o.sigma, o.partiality};
+            const int n_fit_frames = static_cast<int>(frame_start.size());
-                    }
+            std::vector<int32_t> fit_start(n_fit_frames + 1, 0);
+            ParallelFor(n_fit_frames, nthreads, [&](int f) {
+                fit_start[f + 1] = static_cast<int32_t>(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<ScalingFit> 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<double>(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<uint8_t> &apply, const std::vector<double> &ratio) {
                 ParallelChunks(n_partials, partial_threads, [&](int lo, int hi) {
-                    for (int i = lo; i < hi; ++i) {
+                    for (int q = lo; q < hi; ++q) {
-                        const ScalingObs &o = scaling_obs[i];
+                        const ScalingRef &o = scaling_refs[q];
-                        float &corr = scaling_corr[i];
+                        float &corr = scaling_corr[q];
                         if (frame_fitted[o.frame]) {   // UpdateCorr
                             const double denom = static_cast<double>(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)

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
+    // eighty-byte Obs. ScalingRef holds what the group means and the corr update read, in the CSR's
-    // update read; ScalingRef holds, in the CSR's group order, what the group means read besides corr,
+    // group order with the partials in no group after them, and corr is kept beside it in the same
-    // so that only corr is gathered. The same values, so the same arithmetic in the same order.
+    // order - so neither pass gathers. ScalingFit holds, frame by frame, the observations
-    struct ScalingObs {
+    // FitScalingFrames fits: FitPerFrameG's filters, which do not change over the loop, applied once,
-        int32_t frame, group;
+    // and the factors of its sums that do not depend on the reference worked out once. The same
-        float I, sigma, partiality, prescaling_corr, zeta;
+    // values, so the same arithmetic in the same order.
-        uint8_t on_ice;
-    };
     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<double>(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<double> &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<double> &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<ScalingRef> &refs,
                                  const std::vector<float> &corr, std::vector<double> &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 <class O>
     void FitPerFrameG(const std::vector<O> &obs, const std::vector<int32_t> &fstart,
                       const std::vector<int32_t> &fcount, const std::vector<double> &group_mean_in,
                       bool unity, std::vector<double> &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<ScalingFit> &fit, const std::vector<int32_t> &fit_start,
+                          const std::vector<double> &group_mean_in, std::vector<double> &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),