mx/Jungfraujoch
2
1
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases 68 Wiki Activity

Sort the partials and the post-refine events in buckets, in parallel

Browse Source 
Both were one std::sort on one thread over tens of millions of elements, and
together they were a third of a crowded crystal's run.

Bucketing by h first makes them parallel. h is the comparator's leading key, so
the sorted array is exactly the buckets laid end to end, and each bucket sorts
on its own thread. In Ingest the keys are built straight into their bucket slot,
so this replaces the build pass rather than adding one and the packed-key array
is never duplicated; the extra memory is a few hundred kilobytes of histograms.
Buckets are taken largest first, because the tail of the phase is whichever
bucket finishes last.

The run split falls out of the same structure for free: a run of equal (h,k,l)
never crosses an h boundary, so each bucket counts its own runs, a scan over the
buckets gives the offsets, and the arrays are sized exactly - which also removes
the repeated growth the push_backs were paying for. The h range comes from the
finiteness pass, which already reads every observation.

The partials order became total in an earlier commit, when the observation index
was added as the last key. That is what makes this safe rather than merely fast:
the permutation is uniquely determined, so a bucket sort produces the same one a
single sort would.

Measured on a crystal with 66 million partial observations: Ingest 15.2 s and
14.3 s -> 8.3 s and 7.4 s, the post-refine event sort out of the top ten gaps
entirely, the whole crystal 2m22s -> 1m24s. Battery 15m32s -> 10m05s. Same space
group on all 24 crystals, none failed, and no crystal's R_meas moved by more
than 0.3 points.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
jungfrau
2026-08-15 22:31:05 -04:00
co-authored by Claude Opus 5
parent 621019140e
commit f36349060a
2 changed files with 194 additions and 40 deletions
Download Patch File Download Diff File Expand all files Collapse all files
image_analysis/geom_refinement/PostRefine.cpp
+53 -6
View File
@@ -7,6 +7,8 @@
#include <algorithm>
#include <array>
#include <cmath>
#include <limits>
#include <numeric>
#include <mutex>
#include "../../common/JFJochMath.h" // PI
@@ -143,12 +145,57 @@ PostRefineResult PostRefineRotationGeometry(const std::vector<IntegrationOutcome
logger.Info("Post-refine: {} partials gathered", pts.size());
if (pts.size() < static_cast<size_t>(settings.min_events)) return result;
std::sort(pts.begin(), pts.end(), [](const Partial &a, const Partial &b) {
if (a.h != b.h) return a.h < b.h;
if (a.k != b.k) return a.k < b.k;
if (a.l != b.l) return a.l < b.l;
return a.img < b.img;
});
// Bucket by h, then sort the buckets. h is the leading key, so the sorted array is the
// buckets laid end to end, and each bucket sorts on its own thread. Sorting the whole thing
// in one pass moved 56 bytes per element through every level of a comparison sort, on one
// thread, over tens of millions of reflections.
{
const auto part_less = [](const Partial &a, const Partial &b) {
if (a.h != b.h) return a.h < b.h;
if (a.k != b.k) return a.k < b.k;
if (a.l != b.l) return a.l < b.l;
return a.img < b.img;
};
const int n = static_cast<int>(pts.size());
int h_lo = std::numeric_limits<int>::max(), h_hi = std::numeric_limits<int>::min();
for (const auto &q : pts) { h_lo = std::min(h_lo, q.h); h_hi = std::max(h_hi, q.h); }
const int H = (h_lo <= h_hi) ? (h_hi - h_lo + 1) : 1;
const int nt = static_cast<int>(std::clamp<size_t>(nthreads, 1, std::max(1, n)));
const int chunk = (n + nt - 1) / nt;
std::vector<std::vector<int32_t>> hist(nt, std::vector<int32_t>(H, 0));
ParallelChunks(nt, nthreads, [&](int tlo, int thi) {
for (int t = tlo; t < thi; ++t) {
const int lo = t * chunk, hi = std::min(n, lo + chunk);
for (int i = lo; i < hi; ++i) hist[t][pts[i].h - h_lo]++;
}
});
std::vector<int32_t> bstart(H + 1, 0);
int32_t acc = 0;
for (int b = 0; b < H; ++b) {
bstart[b] = acc;
for (int t = 0; t < nt; ++t) { const int32_t c = hist[t][b]; hist[t][b] = acc; acc += c; }
}
bstart[H] = acc;
std::vector<Partial> sorted(pts.size());
ParallelChunks(nt, nthreads, [&](int tlo, int thi) {
for (int t = tlo; t < thi; ++t) {
std::vector<int32_t> fill = hist[t];
const int lo = t * chunk, hi = std::min(n, lo + chunk);
for (int i = lo; i < hi; ++i) sorted[fill[pts[i].h - h_lo]++] = pts[i];
}
});
std::vector<int> order(H);
std::iota(order.begin(), order.end(), 0);
std::sort(order.begin(), order.end(),
[&](int a, int b) { return (bstart[a + 1] - bstart[a]) > (bstart[b + 1] - bstart[b]); });
ParallelFor(H, nthreads, [&](int oi) {
const int b = order[oi];
std::sort(sorted.begin() + bstart[b], sorted.begin() + bstart[b + 1], part_less);
});
pts.swap(sorted);
}
// Split into rocking events (same raw hkl, adjacent frames). Only >=2-frame events carry an
// unbiased phi_obs (a single-frame centroid is just the frame centre); precompute e_ref per event.
image_analysis/scale_merge/RotationScaleMerge.cpp
+141 -34
View File
@@ -12,6 +12,7 @@
#include <fstream>
#include <future>
#include <limits>
#include <numeric>
#include <random>
#include <unordered_map>
@@ -270,10 +271,36 @@ void RotationScaleMerge::Ingest() {
// Per-obs AcceptReflection finiteness (immutable) - lets ComputeAsuGroups stamp the ASU-group id per
// obs from a flat 1-byte array instead of re-reading the fat Obs struct for every space group.
finite_ok.resize(partials.size());
for (size_t i = 0; i < partials.size(); ++i) {
const auto &o = partials[i];
finite_ok[i] = (std::isfinite(o.I) && std::isfinite(o.rlp) && o.rlp != 0.0f
&& std::isfinite(o.sigma) && o.sigma > 0.0f) ? 1 : 0;
// The h range comes out of the same sweep - the sort below buckets by h and needs to know how
// many buckets that is, and this pass already reads every observation.
int h_min = 0, h_max = 0;
{
const int n_obs = static_cast<int>(partials.size());
const int nt = static_cast<int>(std::clamp<size_t>(nthreads, 1, std::max(1, n_obs)));
constexpr int INT_LO = std::numeric_limits<int>::min();
constexpr int INT_HI = std::numeric_limits<int>::max();
std::vector<int> lo_of(nt, INT_HI), hi_of(nt, INT_LO);
const int chunk = (n_obs + nt - 1) / nt;
ParallelChunks(nt, nthreads, [&](int tlo, int thi) {
for (int t = tlo; t < thi; ++t) {
int lmin = INT_HI, lmax = INT_LO;
const int lo = t * chunk, hi = std::min(n_obs, lo + chunk);
for (int i = lo; i < hi; ++i) {
const auto &o = partials[i];
finite_ok[i] = (std::isfinite(o.I) && std::isfinite(o.rlp) && o.rlp != 0.0f
&& std::isfinite(o.sigma) && o.sigma > 0.0f) ? 1 : 0;
lmin = std::min(lmin, o.h);
lmax = std::max(lmax, o.h);
}
lo_of[t] = lmin; hi_of[t] = lmax;
}
});
int gmin = INT_HI, gmax = INT_LO;
for (int t = 0; t < nt; ++t) {
gmin = std::min(gmin, lo_of[t]);
gmax = std::max(gmax, hi_of[t]);
}
if (gmin <= gmax) { h_min = gmin; h_max = gmax; } // otherwise there are no observations
}
// Sort ONCE by (raw h,k,l, image_number) and split into raw-hkl runs. This is the one expensive sort;
@@ -290,41 +317,121 @@ void RotationScaleMerge::Ingest() {
// the result independent of the sort algorithm, which is what lets the order be reproduced by a
// faster one.
struct SortKey { int32_t h, k, l; float image_number; int32_t idx; };
const auto key_less = [](const SortKey &a, const SortKey &b) {
if (a.h != b.h) return a.h < b.h;
if (a.k != b.k) return a.k < b.k;
if (a.l != b.l) return a.l < b.l;
if (a.image_number != b.image_number) return a.image_number < b.image_number;
return a.idx < b.idx;
};
perm.resize(partials.size());
rawrun_start.clear(); rawrun_count.clear();
rawrun_h.clear(); rawrun_k.clear(); rawrun_l.clear(); rawrun_d.clear();
{
std::vector<SortKey> keys(partials.size());
for (size_t i = 0; i < partials.size(); ++i) {
const auto &o = partials[i];
keys[i] = SortKey{o.h, o.k, o.l, o.image_number, static_cast<int32_t>(i)};
}
std::sort(keys.begin(), keys.end(), [](const SortKey &a, const SortKey &b) {
if (a.h != b.h) return a.h < b.h;
if (a.k != b.k) return a.k < b.k;
if (a.l != b.l) return a.l < b.l;
if (a.image_number != b.image_number) return a.image_number < b.image_number;
return a.idx < b.idx;
// Bucket by h first, then sort each bucket. h is the comparator's leading key, so the sorted
// array is exactly the buckets laid end to end - and because the order is total (see above)
// the permutation is the one a single std::sort produces, whichever algorithm gets there.
// That is what makes this checkable: the output is byte-identical, not merely equivalent.
//
// The keys are built straight into their bucket slot, so this replaces the build pass rather
// than adding to it and the 20-bytes-per-observation array is never duplicated.
const int n = static_cast<int>(partials.size());
const int nt = static_cast<int>(std::clamp<size_t>(nthreads, 1, std::max(1, n)));
const int chunk = (n + nt - 1) / nt;
const int H = h_max - h_min + 1;
std::vector<std::vector<int32_t>> hist(nt, std::vector<int32_t>(H, 0));
ParallelChunks(nt, nthreads, [&](int tlo, int thi) {
for (int t = tlo; t < thi; ++t) {
const int lo = t * chunk, hi = std::min(n, lo + chunk);
for (int i = lo; i < hi; ++i) hist[t][partials[i].h - h_min]++;
}
});
for (size_t i = 0; i < keys.size(); ++i) perm[i] = keys[i].idx;
for (int i = 0; i < static_cast<int>(keys.size()); ) {
const SortKey k0 = keys[i];
int j = i;
float d = NAN;
while (j < static_cast<int>(keys.size())) {
if (keys[j].h != k0.h || keys[j].k != k0.k || keys[j].l != k0.l) break;
if (!std::isfinite(d)) {
const float dj = partials[keys[j].idx].d; // resolution: only until one is usable
if (std::isfinite(dj) && dj > 0.0f) d = dj;
}
++j;
// Where each chunk writes each bucket: buckets in order, and within a bucket the chunks in
// index order, so the scatter below is stable and lands in original order inside a bucket.
std::vector<int32_t> bstart(H + 1, 0);
int32_t acc = 0;
for (int b = 0; b < H; ++b) {
bstart[b] = acc;
for (int t = 0; t < nt; ++t) {
const int32_t c = hist[t][b];
hist[t][b] = acc;
acc += c;
}
rawrun_start.push_back(i);
rawrun_count.push_back(j - i);
rawrun_h.push_back(k0.h); rawrun_k.push_back(k0.k); rawrun_l.push_back(k0.l);
rawrun_d.push_back(d);
i = j;
}
bstart[H] = acc;
std::vector<SortKey> keys(partials.size());
ParallelChunks(nt, nthreads, [&](int tlo, int thi) {
for (int t = tlo; t < thi; ++t) {
std::vector<int32_t> fill = hist[t];
const int lo = t * chunk, hi = std::min(n, lo + chunk);
for (int i = lo; i < hi; ++i) {
const auto &o = partials[i];
keys[fill[o.h - h_min]++] =
SortKey{o.h, o.k, o.l, o.image_number, static_cast<int32_t>(i)};
}
}
});
// Largest bucket first: the tail of this is one bucket, so it should be the big one.
std::vector<int> order(H);
std::iota(order.begin(), order.end(), 0);
std::sort(order.begin(), order.end(), [&](int a, int b) {
return (bstart[a + 1] - bstart[a]) > (bstart[b + 1] - bstart[b]);
});
ParallelFor(H, nthreads, [&](int oi) {
const int b = order[oi];
std::sort(keys.begin() + bstart[b], keys.begin() + bstart[b + 1], key_less);
});
ParallelChunks(n, nthreads, [&](int lo, int hi) {
for (int i = lo; i < hi; ++i) perm[i] = keys[i].idx;
});
// The runs never cross an h boundary, so each bucket owns its own and they can be counted,
// scanned and written without touching each other. Sizing the arrays from the count also
// removes the repeated growth the push_backs paid for.
std::vector<int32_t> nruns(H, 0);
ParallelFor(H, nthreads, [&](int b) {
int c = 0;
for (int i = bstart[b]; i < bstart[b + 1]; ) {
int j = i;
while (j < bstart[b + 1] && keys[j].h == keys[i].h && keys[j].k == keys[i].k
&& keys[j].l == keys[i].l) ++j;
++c;
i = j;
}
nruns[b] = c;
});
std::vector<int32_t> run_base(H + 1, 0);
for (int b = 0; b < H; ++b) run_base[b + 1] = run_base[b] + nruns[b];
const int n_runs = run_base[H];
rawrun_start.resize(n_runs); rawrun_count.resize(n_runs);
rawrun_h.resize(n_runs); rawrun_k.resize(n_runs); rawrun_l.resize(n_runs);
rawrun_d.resize(n_runs);
ParallelFor(H, nthreads, [&](int b) {
int at = run_base[b];
for (int i = bstart[b]; i < bstart[b + 1]; ) {
const SortKey k0 = keys[i];
int j = i;
float d = NAN;
while (j < bstart[b + 1]) {
if (keys[j].h != k0.h || keys[j].k != k0.k || keys[j].l != k0.l) break;
if (!std::isfinite(d)) {
const float dj = partials[keys[j].idx].d; // resolution: only until one is usable
if (std::isfinite(dj) && dj > 0.0f) d = dj;
}
++j;
}
rawrun_start[at] = i;
rawrun_count[at] = j - i;
rawrun_h[at] = k0.h; rawrun_k[at] = k0.k; rawrun_l[at] = k0.l;
rawrun_d[at] = d;
++at;
i = j;
}
});
}
rawrun_group.assign(rawrun_start.size(), -1);
logger.Info("RotationScaleMerge: ingested {} partial observations from {} frames ({} distinct hkl)",