Stills merge: the error-model fit and its rounds with the means run in a fraction of the time

The stills error model (fit + means, iterated) had made the "error model and merge" phase 8-18x
slower than before it. Now:
- the fit puts its bin boundaries in place with nth_element on compact (key, position) pairs, the two
  halves of each split side by side, instead of sorting the whole pool by a five-field tuple; the
  per-group passes of its iteration run on the workers;
- each round's fit starts from the previous round's model, so it settles in a few iterations;
- the rounds stop once a round leaves the model where the previous one put it (to 1e-3);
- those rounds only carry the means to where the model settles, so on a large run they fit every
  stride-th observation (about four million); one last round fits them all and makes the final means;
- the observations are grouped by one parallel sort on (reflection, image order) into one flat
  array instead of a hash map of vectors, and the per-reflection passes (means, samples) run on the
  workers.
Every sum runs in a fixed order over fixed pieces, so the result is the same on any number of
threads (the fit is checked at 1 and 16 threads; the merge test checks 1 against 8). On the sets
tried the merged data are unchanged to within the last round's convergence.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01HVeDQqN9qRD57vnShzTC37

image_analysis/scale_merge/ErrorModel.cpp

@@ -7,7 +7,6 @@
 #include <cmath>
 #include <functional>
 #include <future>
-#include <tuple>
 
 #include <Eigen/Dense>
 
@@ -156,60 +155,93 @@ bool ErrorModelUpdate(const std::vector<double> &bs2, const std::vector<double>
     return !settled;
 }
 
-StillsErrorModelFit FitStillsErrorModel(std::vector<StillsErrorModelSample> &pool) {
+StillsErrorModelFit FitStillsErrorModel(const std::vector<StillsErrorModelSample> &pool, size_t nthreads,
+                                        const StillsErrorModelFit &start) {
     constexpr int n_bins = ERROR_MODEL_BINS;
     constexpr int n_sub = 4;
     StillsErrorModelFit fit;
     if (pool.size() < static_cast<size_t>(8 * n_bins * n_sub))
         return fit;
 
-    // Bins of counting (I/sigma)^2, each split by the partiality term's share of the intensity term. Total
-    // orders, so which sample lands in which bin is a property of the samples alone.
-    auto snr = [](const StillsErrorModelSample &x) { return x.I2 / x.s2; };
-    auto share = [](const StillsErrorModelSample &x) { return x.I2 > 0.0 ? x.P2 / x.I2 : 0.0; };
-    std::sort(pool.begin(), pool.end(), [&](const auto &x, const auto &y) {
-        return std::make_tuple(snr(x), x.s2, x.I2, x.P2, x.dev2) < std::make_tuple(snr(y), y.s2, y.I2, y.P2, y.dev2);
+    // Bins of counting (I/sigma)^2, each split by the partiality term's share of the intensity term. The
+    // boundaries are put in place with nth_element on (key, position in the pool) - a total order, so which
+    // sample lands in which bin depends on the samples and the order the caller gave them in alone. The keys
+    // are computed once and carried in the comparison, never recomputed inside it.
+    struct Key { double v; size_t i; };
+    const auto by_key = [](const Key &p, const Key &q) { return p.v < q.v || (p.v == q.v && p.i < q.i); };
+    std::vector<Key> key(pool.size());
+    ParallelChunks(static_cast<int>(pool.size()), ThreadsForWork(pool.size(), nthreads), [&](int lo, int hi) {
+        for (int i = lo; i < hi; ++i)
+            key[i] = {pool[i].I2 / pool[i].s2, static_cast<size_t>(i)};
     });
     const size_t per = pool.size() / n_bins;
-    std::vector<std::pair<size_t, size_t>> groups;
-    for (int bin = 0; bin < n_bins; ++bin) {
-        const size_t lo = bin * per, hi = bin == n_bins - 1 ? pool.size() : (bin + 1) * per;
-        std::sort(pool.begin() + lo, pool.begin() + hi, [&](const auto &x, const auto &y) {
-            return std::make_tuple(share(x), x.s2, x.I2, x.dev2) < std::make_tuple(share(y), y.s2, y.I2, y.dev2);
-        });
-        const size_t sub = (hi - lo) / n_sub;
-        for (int q = 0; q < n_sub; ++q)
-            groups.emplace_back(lo + q * sub, q == n_sub - 1 ? hi : lo + (q + 1) * sub);
-    }
+    // The split runs down the middle, and the two halves of a split are disjoint ranges, so they run side
+    // by side while the range is still worth a thread (as in FitErrorModel).
+    constexpr size_t SPLIT_MIN_PARALLEL = 1u << 17;
+    const std::function<void(size_t, size_t, int, int, int)> split = [&](size_t lo, size_t hi, int b0, int b1, int depth) {
+        if (b0 >= b1) return;
+        const int mid = (b0 + b1) / 2;
+        const size_t k = static_cast<size_t>(mid) * per;
+        std::nth_element(key.begin() + lo, key.begin() + k, key.begin() + hi, by_key);
+        if (depth > 0 && hi - lo > SPLIT_MIN_PARALLEL) {
+            auto left = std::async(std::launch::async, [&, lo, k, b0, mid, depth] {
+                RestoreThreadAffinity();
+                split(lo, k, b0, mid, depth - 1);
+            });
+            split(k, hi, mid + 1, b1, depth - 1);
+            left.get();
+        } else {
+            split(lo, k, b0, mid, 0);
+            split(k, hi, mid + 1, b1, 0);
+        }
+    };
+    split(0, key.size(), 1, n_bins, nthreads > 1 ? 3 : 0);
     // b only where the strongest bin reaches past the counting term, as in FitErrorModel.
     std::vector<double> top;
-    for (size_t i = (n_bins - 1) * per; i < pool.size(); ++i)
-        top.push_back(snr(pool[i]));
+    for (size_t j = (n_bins - 1) * per; j < key.size(); ++j)
+        top.push_back(key[j].v);
     std::nth_element(top.begin(), top.begin() + top.size() / 2, top.end());
     const bool b_measured = top[top.size() / 2] > ERROR_MODEL_B_LEVER_MIN;
+    // Each bin in four by the share, and the samples copied out in that order.
+    std::vector<std::pair<size_t, size_t>> groups(n_bins * n_sub);
+    std::vector<StillsErrorModelSample> x(pool.size());
+    ParallelFor(n_bins, nthreads, [&](int bin) {
+        const size_t lo = bin * per, hi = bin == n_bins - 1 ? key.size() : (bin + 1) * per;
+        for (size_t j = lo; j < hi; ++j) {
+            const auto &o = pool[key[j].i];
+            key[j].v = o.I2 > 0.0 ? o.P2 / o.I2 : 0.0;
+        }
+        const size_t sub = (hi - lo) / n_sub;
+        for (int q = 1; q < n_sub; ++q)
+            std::nth_element(key.begin() + lo + (q - 1) * sub, key.begin() + lo + q * sub, key.begin() + hi, by_key);
+        for (int q = 0; q < n_sub; ++q)
+            groups[bin * n_sub + q] = {lo + q * sub, q == n_sub - 1 ? hi : lo + (q + 1) * sub};
+        for (size_t j = lo; j < hi; ++j)
+            x[j] = pool[key[j].i];
+    });
+    key = {};
 
     const int ng = static_cast<int>(groups.size());
     std::vector<double> z2(pool.size());
     Eigen::MatrixXd X(ng, 3);
     Eigen::VectorXd y(ng);
-    double a = 1.0, b2 = 0.0, c2 = 0.0;
+    double a = start.a, b2 = start.b2, c2 = start.c2;
     for (int iter = 0; iter < 30; ++iter) {
         // Per group: the three terms divided by the current variance, and the median normalised deviation
         // expressed as the sum it stands for.
-        for (int g = 0; g < ng; ++g) {
+        ParallelFor(ng, ThreadsForWork(pool.size(), nthreads), [&](int g) {
             const auto [lo, hi] = groups[g];
             double s = 0.0, I = 0.0, P = 0.0;
             for (size_t i = lo; i < hi; ++i) {
-                const auto &x = pool[i];
-                const double v = a * x.s2 + b2 * x.I2 + c2 * x.P2;
-                s += x.s2 / v; I += x.I2 / v; P += x.P2 / v;
-                z2[i] = x.dev2 / v;
+                const double v = a * x[i].s2 + b2 * x[i].I2 + c2 * x[i].P2;
+                s += x[i].s2 / v; I += x[i].I2 / v; P += x[i].P2 / v;
+                z2[i] = x[i].dev2 / v;
             }
             const auto mid = z2.begin() + lo + (hi - lo) / 2;
             std::nth_element(z2.begin() + lo, mid, z2.begin() + hi);
             X.row(g) << s, I, P;
             y(g) = *mid / ERROR_MODEL_CHI2_1_MEDIAN * static_cast<double>(hi - lo);
-        }
+        });
         // Every group weighs the same in relative terms; a negative term is dropped and the rest refitted.
         const Eigen::VectorXd w = y.array().square().inverse();
         std::vector<int> terms = b_measured ? std::vector<int>{0, 1, 2} : std::vector<int>{0, 2};

image_analysis/scale_merge/ErrorModel.h

@@ -82,4 +82,7 @@ struct StillsErrorModelFit {
     bool b_resolved = false;   // as ErrorModelFit::b_resolved
 };
 
-StillsErrorModelFit FitStillsErrorModel(std::vector<StillsErrorModelSample> &pool);
+// start: where the iteration begins - identity by default, or the fit of a previous round on a pool that
+// has moved little since.
+StillsErrorModelFit FitStillsErrorModel(const std::vector<StillsErrorModelSample> &pool, size_t nthreads = 1,
+                                        const StillsErrorModelFit &start = {});

image_analysis/scale_merge/Merge.cpp

@@ -7,6 +7,7 @@
 #include <cmath>
 #include <limits>
 #include <map>
+#include <span>
 #include <tuple>
 #include <unordered_map>
 
@@ -298,9 +299,7 @@ void MergeOnTheFly::RefineErrorModel(const std::vector<IntegrationOutcome> &outc
     //        AddImage uses, down to the floor of the merge with an error model.
     struct Obs { float I, sigma, corr, var_bkg, p; };
     struct KeyedObs { uint64_t key; Obs obs; };
-    std::unordered_map<uint64_t, std::vector<Obs>> groups;
-    std::vector<uint64_t> keys;   // in order of first appearance, so every sum below is in a fixed order
-
+    std::vector<KeyedObs> all;   // in image order
     ReduceImagesInOrder<KeyedObs>(
         outcomes.size(), nthreads,
         [&](size_t i, std::vector<KeyedObs> &out) {
@@ -325,86 +324,147 @@ void MergeOnTheFly::RefineErrorModel(const std::vector<IntegrationOutcome> &outc
             }
         },
         [&](size_t, const std::vector<KeyedObs> &out) {
-            for (const auto &o: out) {
-                auto &g = groups[o.key];
-                if (g.empty())
-                    keys.push_back(o.key);
-                g.push_back(o.obs);
-            }
+            all.insert(all.end(), out.begin(), out.end());
         });
 
+    // The observations of each reflection side by side, in key order and within a reflection in image
+    // order - a sort on (key, place), a total order, so the same on any number of workers. Every sum
+    // below runs over them in that fixed order.
+    struct KeyPlace { uint64_t key; size_t place; };
+    std::vector<KeyPlace> order(all.size());
+    const size_t obs_threads = ThreadsForWork(all.size(), nthreads);
+    ParallelChunks(static_cast<int>(all.size()), obs_threads, [&](int lo, int hi) {
+        for (int i = lo; i < hi; ++i)
+            order[i] = {all[i].key, static_cast<size_t>(i)};
+    });
+    ParallelSort(order.begin(), order.end(), nthreads, [](const KeyPlace &x, const KeyPlace &y) {
+        return x.key < y.key || (x.key == y.key && x.place < y.place);
+    });
+    std::vector<Obs> flat(all.size());
+    ParallelChunks(static_cast<int>(all.size()), obs_threads, [&](int lo, int hi) {
+        for (int i = lo; i < hi; ++i)
+            flat[i] = all[order[i].place].obs;
+    });
+    all = {};
+    std::vector<uint64_t> keys;
+    std::vector<size_t> start;   // reflection k is flat[start[k] .. start[k + 1])
+    for (size_t i = 0; i < order.size(); ++i)
+        if (i == 0 || order[i].key != order[i - 1].key) {
+            keys.push_back(order[i].key);
+            start.push_back(i);
+        }
+    start.push_back(order.size());
+    order = {};
+    const auto group = [&](int k) { return std::span<const Obs>(flat.data() + start[k], start[k + 1] - start[k]); };
+    const int n_keys = static_cast<int>(keys.size());
+    const size_t key_threads = ThreadsForWork(keys.size(), nthreads, 1024);
+
     // --- 2. The starting mean of each reflection (counting-sigma weights) and, where there are three
     //        observations for it to stand on, the median - the robust centre of the outlier test.
-    std::unordered_map<uint64_t, double> mean;
-    for (const auto key: keys) {
-        const auto &obs = groups[key];
-        double sum_w = 0.0, sum_wI = 0.0;
-        for (const auto &o: obs) {
-            const double w = 1.0 / (static_cast<double>(o.sigma) * o.sigma);
-            sum_w += w;
-            sum_wI += w * o.I;
+    std::vector<double> mean(n_keys);
+    std::vector<float> median(n_keys, NAN);
+    ParallelChunks(n_keys, key_threads, [&](int lo, int hi) {
+        for (int k = lo; k < hi; ++k) {
+            const auto obs = group(k);
+            double sum_w = 0.0, sum_wI = 0.0;
+            for (const auto &o: obs) {
+                const double w = 1.0 / (static_cast<double>(o.sigma) * o.sigma);
+                sum_w += w;
+                sum_wI += w * o.I;
+            }
+            mean[k] = sum_wI / sum_w;
+            if (reject_outliers && obs.size() >= 3) {
+                std::vector<float> iv;
+                iv.reserve(obs.size());
+                for (const auto &o: obs)
+                    iv.push_back(o.I);
+                std::nth_element(iv.begin(), iv.begin() + iv.size() / 2, iv.end());
+                median[k] = iv[iv.size() / 2];
+            }
         }
-        mean[key] = sum_wI / sum_w;
-        if (reject_outliers && obs.size() >= 3) {
-            std::vector<float> iv;
-            iv.reserve(obs.size());
-            for (const auto &o: obs)
-                iv.push_back(o.I);
-            std::nth_element(iv.begin(), iv.begin() + iv.size() / 2, iv.end());
-            reject_median_I[key] = iv[iv.size() / 2];
-        }
-    }
+    });
+    for (int k = 0; k < n_keys; ++k)
+        if (!std::isnan(median[k]))
+            reject_median_I[keys[k]] = median[k];
 
     // --- 3. Fit the model (ErrorModel.h, FitStillsErrorModel) on the scatter about the reflection means,
     //        the counting variance rebuilt at the mean, and refit with the means the fitted model weights:
     //        the means and the model each depend on the other, and a single pass leaves the weights of
     //        the strong reflections at the counting-sigma means the fit started from. Measured on serial
     //        and rotation-as-serial sets, a few rounds raise the anomalous signal of the merge; they have
-    //        settled after four.
+    //        settled after four, most often already after two or three: a round that leaves the model where
+    //        the previous one put it (to 1e-3) would only reproduce the same means, and ends the iteration.
+    //        Those rounds only carry the means to where the model settles, so they fit every stride-th
+    //        observation (some four million of them); one last round fits them all and makes the final means.
+    //        Each observation's sample goes to its own slot, so the pool is in key order whatever the number
+    //        of workers; the slots left empty (single observations, the leverage cut) are dropped after.
     constexpr int ROUNDS = 4;
+    const size_t stride = std::max<size_t>(1, flat.size() >> 22);
+    bool last_round = false;
     std::vector<StillsErrorModelSample> samples;
     StillsErrorModelFit fit;
     double a = 1.0, b2 = 0.0, c2 = 0.0;
     auto variance = [&](const Obs &o, double m) {
         return a * CountingVariance(o.sigma, o.corr, o.var_bkg, o.p, m) + (b2 + c2 * PartialityTerm(o.p)) * m * m;
     };
-    for (int round = 0; round < ROUNDS; ++round) {
-        samples.clear();
-        for (const auto key: keys) {
-            const auto &obs = groups[key];
-            if (obs.size() < 2)
-                continue;
-            const double m = mean[key];
-            // The residual from the weighted mean has E[(I_i - <I>)^2] = var_i (1 - h_i), h_i = w_i / sum_w its
-            // leverage, with the weights of the model that made the mean.
-            double sum_w = 0.0;
-            for (const auto &o: obs)
-                sum_w += 1.0 / variance(o, m);
-            for (const auto &o: obs) {
-                const double factor = 1.0 - 1.0 / (variance(o, m) * sum_w);
-                if (factor < 0.05)
+    for (int round = 0; ; ++round) {
+        const size_t step = last_round ? 1 : stride;
+        samples.assign((flat.size() + step - 1) / step, StillsErrorModelSample{NAN, NAN, NAN, NAN});
+        ParallelChunks(n_keys, key_threads, [&](int lo, int hi) {
+            std::vector<double> var;
+            for (int k = lo; k < hi; ++k) {
+                const auto obs = group(k);
+                if (obs.size() < 2)
                     continue;
-                const double resid = static_cast<double>(o.I) - m;
-                samples.push_back({CountingVariance(o.sigma, o.corr, o.var_bkg, o.p, m), m * m,
-                                   PartialityTerm(o.p) * m * m, resid * resid / factor});
+                const double m = mean[k];
+                // The residual from the weighted mean has E[(I_i - <I>)^2] = var_i (1 - h_i), h_i = w_i / sum_w
+                // its leverage, with the weights of the model that made the mean.
+                var.resize(obs.size());
+                double sum_w = 0.0;
+                for (size_t j = 0; j < obs.size(); ++j) {
+                    var[j] = variance(obs[j], m);
+                    sum_w += 1.0 / var[j];
+                }
+                for (size_t j = 0; j < obs.size(); ++j) {
+                    if ((start[k] + j) % step != 0)
+                        continue;
+                    const auto &o = obs[j];
+                    const double factor = 1.0 - 1.0 / (var[j] * sum_w);
+                    if (factor < 0.05)
+                        continue;
+                    const double resid = static_cast<double>(o.I) - m;
+                    samples[(start[k] + j) / step] = {CountingVariance(o.sigma, o.corr, o.var_bkg, o.p, m), m * m,
+                                                     PartialityTerm(o.p) * m * m, resid * resid / factor};
+                }
             }
-        }
-        fit = FitStillsErrorModel(samples);
+        });
+        samples.erase(std::remove_if(samples.begin(), samples.end(),
+                                     [](const StillsErrorModelSample &x) { return std::isnan(x.dev2); }),
+                      samples.end());
+        const StillsErrorModelFit last = fit;
+        fit = FitStillsErrorModel(samples, nthreads, last);   // each round starts from the last one's model
         if (!fit.active)
             return;   // too little multiplicity to fit -> leave identity
         a = fit.a; b2 = fit.b2; c2 = fit.c2;
-        for (const auto key: keys) {
-            double sum_w = 0.0, sum_wI = 0.0;
-            for (const auto &o: groups[key]) {
-                const double w = 1.0 / variance(o, mean[key]);
-                sum_w += w;
-                sum_wI += w * o.I;
+        const auto close = [](double x, double y) { return std::fabs(x - y) <= 1e-3 * std::fabs(y); };
+        const bool settled = round > 0 && close(a, last.a) && close(b2, last.b2) && close(c2, last.c2);
+        ParallelChunks(n_keys, key_threads, [&](int lo, int hi) {
+            for (int k = lo; k < hi; ++k) {
+                double sum_w = 0.0, sum_wI = 0.0;
+                for (const auto &o: group(k)) {
+                    const double w = 1.0 / variance(o, mean[k]);
+                    sum_w += w;
+                    sum_wI += w * o.I;
+                }
+                mean[k] = sum_wI / sum_w;
             }
-            mean[key] = sum_wI / sum_w;
-        }
+        });
+        if (last_round || ((settled || round + 1 == ROUNDS) && step == 1))
+            break;
+        last_round = settled || round + 1 == ROUNDS;
     }
-    for (const auto key: keys)
-        error_model_mean_I[key] = static_cast<float>(mean[key]);
+    for (int k = 0; k < n_keys; ++k)
+        error_model_mean_I[keys[k]] = static_cast<float>(mean[k]);
 
     error_model_a = a;
     error_model_b = std::sqrt(b2);

tests/MergeScaleTest.cpp

@@ -546,6 +546,10 @@ TEST_CASE("StillsErrorModel_RecoversThePartialityTerm") {
     }
     const auto fit = FitStillsErrorModel(pool);
     REQUIRE(fit.active);
+    const auto many = FitStillsErrorModel(pool, 16);   // the same numbers on any number of threads
+    CHECK(many.a == fit.a);
+    CHECK(many.b2 == fit.b2);
+    CHECK(many.c2 == fit.c2);
     CHECK(fit.a == Catch::Approx(a).epsilon(0.05));
     CHECK(std::sqrt(fit.b2) == Catch::Approx(b).epsilon(0.05));
     CHECK(std::sqrt(fit.c2) == Catch::Approx(c).epsilon(0.05));