From cfe202e9a1ec97d5c4c0b63e84a028ecfaa4cca2 Mon Sep 17 00:00:00 2001
From: Filip Leonarski <filip.leonarski@psi.ch>
Date: Sat, 26 Sep 2026 22:52:09 +0200
Subject: [PATCH] ScaledObservations: sample each image on its own thread

The ASU-key sample over every partial ran on one thread (0.2-1 s of the merge tail). Each image's
parts are now taken on the workers and joined in image order, so the sort sees the same parts in the
same order.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01D1G8gJVAy6gp1K5Dz3NE5C
---
 .../scale_merge/AnisotropyAnalysis.cpp        | 25 +++++++++++++------
 .../scale_merge/AnisotropyAnalysis.h          |  3 ++-
 rugnux/Rugnux.cpp                             |  3 ++-
 3 files changed, 21 insertions(+), 10 deletions(-)

diff --git a/image_analysis/scale_merge/AnisotropyAnalysis.cpp b/image_analysis/scale_merge/AnisotropyAnalysis.cpp
index 77c4b2758..18d4adfd7 100644
--- a/image_analysis/scale_merge/AnisotropyAnalysis.cpp
+++ b/image_analysis/scale_merge/AnisotropyAnalysis.cpp
@@ -2,6 +2,7 @@
 // SPDX-License-Identifier: GPL-3.0-only
 
 #include "AnisotropyAnalysis.h"
+#include "../../common/ParallelFor.h"
 
 #include <algorithm>
 #include <array>
@@ -1029,7 +1030,8 @@ namespace {
 
 std::vector<AnisotropyObservation> ScaledObservations(const std::vector<IntegrationOutcome> &outcomes,
                                                      bool rotation, const gemmi::SpaceGroup *space_group,
-                                                     float wedge_deg, double min_partiality) {
+                                                     float wedge_deg, double min_partiality,
+                                                     size_t nthreads) {
     // Per-image scale, indexed the way the outcomes are.
     std::vector<double> g(outcomes.size(), 0.0);
     for (size_t i = 0; i < outcomes.size(); ++i)
@@ -1105,20 +1107,27 @@ std::vector<AnisotropyObservation> ScaledObservations(const std::vector<Integrat
                                   space_group ? *space_group : *gemmi::find_spacegroup_by_number(1));
     const uint64_t sample_mask =
             std::bit_ceil((n_parts + MAX_SAMPLE_PARTS - 1) / MAX_SAMPLE_PARTS) - 1;
-    std::vector<Part> parts;
-    // Reserve what the sample holds rather than letting the vector grow: the growth chain copies a few
-    // hundred megabytes on a large sweep.
-    parts.reserve(std::min(n_parts, MAX_SAMPLE_PARTS));
-    for (size_t i = 0; i < outcomes.size(); ++i) {
+    // Each image's sample on its own thread, then joined in image order - the same parts in the same
+    // order as one loop over the images.
+    std::vector<std::vector<Part>> image_parts(outcomes.size());
+    ParallelFor(static_cast<int>(outcomes.size()), nthreads, [&](int i) {
         if (!(g[i] > 0.0))
-            continue;
+            return;
         for (const auto &r : outcomes[i].reflections) {
             if (!usable(r))
                 continue;
             if ((MixKey(asu_key(r.h, r.k, r.l).pack()) & sample_mask) != 0)
                 continue;
-            parts.push_back({r.h, r.k, r.l, r.image_number, &r, i});
+            image_parts[i].push_back({r.h, r.k, r.l, r.image_number, &r, static_cast<size_t>(i)});
         }
+    });
+    std::vector<Part> parts;
+    // Reserve what the sample holds rather than letting the vector grow: the growth chain copies a few
+    // hundred megabytes on a large sweep.
+    parts.reserve(std::min(n_parts, MAX_SAMPLE_PARTS));
+    for (auto &v : image_parts) {
+        parts.insert(parts.end(), v.begin(), v.end());
+        std::vector<Part>().swap(v);
     }
     std::sort(parts.begin(), parts.end(), [](const Part &a, const Part &b) {
         return std::tie(a.h, a.k, a.l, a.image_number) < std::tie(b.h, b.k, b.l, b.image_number);
diff --git a/image_analysis/scale_merge/AnisotropyAnalysis.h b/image_analysis/scale_merge/AnisotropyAnalysis.h
index 742202d8c..61f5573ca 100644
--- a/image_analysis/scale_merge/AnisotropyAnalysis.h
+++ b/image_analysis/scale_merge/AnisotropyAnalysis.h
@@ -150,7 +150,8 @@ std::vector<AnisotropyObservation> ScaledObservations(const std::vector<Integrat
                                                      bool rotation,
                                                      const gemmi::SpaceGroup *space_group = nullptr,
                                                      float wedge_deg = 0.0f,
-                                                     double min_partiality = 0.5);
+                                                     double min_partiality = 0.5,
+                                                     size_t nthreads = 1);
 
 // What the caller knows about the run and the merge that the reflections alone do not say.
 struct AnisotropyRunInfo {
diff --git a/rugnux/Rugnux.cpp b/rugnux/Rugnux.cpp
index ba9b26063..2fb91b89b 100644
--- a/rugnux/Rugnux.cpp
+++ b/rugnux/Rugnux.cpp
@@ -8359,7 +8359,8 @@ ProcessResult Rugnux::RunPipeline(RugnuxObserver *observer, bool write_output, b
                         ScaledObservations(indexer->GetIntegrationOutcome(),
                                            experiment_.IsRotationIndexing(), twin_sg,
                                            experiment_.GetGoniometer()
-                                               ? experiment_.GetGoniometer()->GetWedge_deg() : 0.0f),
+                                               ? experiment_.GetGoniometer()->GetWedge_deg() : 0.0f,
+                                           0.5, config_.nthreads),
                         *result.consensus_cell, twin_sg, aniso_run);
                 stats_text << AnisotropyToText(sm.statistics.anisotropy) << "\n";
             }