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

Spread the scaling corrections and the space-group search over the cores

Browse Source 
Two thirds of a rotation run is one thread. The image loop is not the problem -
on the heaviest crystal of the battery it is 1.8 s of 40 - and neither GPU nor
CPU is saturated, because while the corrections and the space-group search run
there is one core working and 47 idle. Mean occupancy over the whole run: 3.9 of
48.

In the correction surfaces (absorption in the goniometer frame, detector-plane
modulation, absorption against time and detector position - all one function):
the per-cell accumulation, the score reduction and the final apply are now
chunked, as are the three loops that assign a full to its cell, one of which
spends a sine and a cosine per full de-rotating it into the crystal frame. Two
full sorts of four million floats went with them: only the nine bin edges are
wanted, so they are selected instead, each selection starting where the last one
left off.

The per-group pass is deliberately left serial. The terms of one group are
spread all over the list, so the only way to give a thread groups of its own is
to walk in group order, and that trades a near-sequential read of the fulls for
a random one over a few hundred megabytes - the trade that already lost once in
the combine kernel.

The space-group search scores each candidate rotation by correlating I(h)
against I(Rh) over the whole merge. Every operator it can ask about comes from a
fixed list and none of them depend on each other, so they are scored up front,
in parallel, and the search reads the cache. The scratch that stops a pair being
counted twice is now per worker rather than shared.

Worker counts are gated on how much work there is, not on how many cores the
machine has (ThreadsForWork). Both parallel helpers start a thread per chunk, so
a small dataset on a large node would otherwise pay for 48 thread starts to sum
a few thousand terms - and this runs on 8-core laptops as well as on this node.

Measured on the heaviest crystal, idle machine, two runs each, summed over both
passes: those phases go 7.88 s -> 5.19 s. Whole-run wall time is the wrong ruler
for it - it moves +-4 s between identical runs. Battery 9m45s -> 9m23s, space
group 21/24, no failures; 16 of 24 crystals bit-identical to the previous run
and the rest inside the noise floor of running one binary twice.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
jungfrau
2026-08-23 08:25:43 -04:00
co-authored by Claude Opus 5
parent 01b619cfd0
commit ca3ca7170e
5 changed files with 175 additions and 71 deletions
Download Patch File Download Diff File Expand all files Collapse all files
common/ParallelFor.h
+11
View File
@@ -12,6 +12,17 @@
// worker count from the caller rather than asking the hardware, so a run that was told how many
// threads to use keeps to it.
// How many workers a pass over `n` cheap items should use: enough that each gets at least
// `min_per_thread` of them, and never more than the caller was given. Both helpers below start a
// thread per chunk, which costs tens of microseconds, so a pass that is repeated - or one whose data
// is small - can easily spend more on splitting the work than on doing it. That is not a large
// machine's problem: it is what makes the same code right on an 8-core laptop, a 16-core desktop and
// a two-socket node, none of which should be handed 48 chunks of a few thousand items.
inline size_t ThreadsForWork(size_t n, size_t nthreads, size_t min_per_thread = 32768) {
if (nthreads <= 1 || n == 0) return 1;
return std::clamp<size_t>(n / min_per_thread, 1, nthreads);
}
// Chunked: each worker gets one contiguous [lo, hi) range and there is no per-item synchronisation.
// Right for millions of cheap uniform items - the CPU stand-in for a flat CUDA grid-stride kernel.
// The split is fixed and deterministic, so a pass whose per-element work is independent gives the
image_analysis/scale_merge/RotationScaleMerge.cpp
+123 -66
View File
@@ -1707,16 +1707,20 @@ void RotationScaleMerge::RefineAbsorption(int n_iter, int n_groups) {
constexpr int NB = 8;
std::vector<int32_t> cell(fulls.size(), -1);
for (size_t i = 0; i < fulls.size(); ++i) {
const Obs &o = fulls[i];
if (!std::isfinite(o.px) || !std::isfinite(o.py))
continue;
const Coord s1 = Coord((o.px - beam_x) / F, (o.py - beam_y) / F, 1.0f).Normalize();
const Coord u = gon.GetTransformationAngle(gon.GetAngle_deg(o.image_number)).transpose() * s1;
const int ix = std::clamp(static_cast<int>((u.x + 1.0f) * 0.5f * NB), 0, NB - 1);
const int iy = std::clamp(static_cast<int>((u.y + 1.0f) * 0.5f * NB), 0, NB - 1);
cell[i] = ix * NB + iy;
}
// A de-rotation and a normalisation per full, each independent of the others - and a spindle
// rotation matrix costs a sine and a cosine, which is why this shows up at all.
ParallelChunks(static_cast<int>(fulls.size()), ThreadsForWork(fulls.size(), nthreads), [&](int lo, int hi) {
for (int i = lo; i < hi; ++i) {
const Obs &o = fulls[i];
if (!std::isfinite(o.px) || !std::isfinite(o.py))
continue;
const Coord s1 = Coord((o.px - beam_x) / F, (o.py - beam_y) / F, 1.0f).Normalize();
const Coord u = gon.GetTransformationAngle(gon.GetAngle_deg(o.image_number)).transpose() * s1;
const int ix = std::clamp(static_cast<int>((u.x + 1.0f) * 0.5f * NB), 0, NB - 1);
const int iy = std::clamp(static_cast<int>((u.y + 1.0f) * 0.5f * NB), 0, NB - 1);
cell[i] = ix * NB + iy;
}
});
ApplyCellSurface(cell, NB * NB, n_iter, n_groups, "Absorption (goniometer-frame 8x8)");
}
@@ -1755,23 +1759,37 @@ void RotationScaleMerge::RefineAbsorptionTime(int n_iter, int n_groups) {
if (usable(o)) { vx.push_back(o.px); vy.push_back(o.py); }
if (static_cast<int>(vx.size()) < NT * ND * ND)
return;
std::sort(vx.begin(), vx.end());
std::sort(vy.begin(), vy.end());
// Only the ND-1 bin edges are wanted, not the order of the millions of values between them, so
// select them instead of sorting. Taken in increasing order each selection starts where the last
// one left off, on a range already partitioned below that edge, so this is a couple of passes over
// the values rather than the n log n a full sort spends.
std::vector<float> ex(ND - 1), ey(ND - 1);
size_t prev = 0;
for (int i = 1; i < ND; ++i) {
ex[i - 1] = vx[vx.size() * i / ND];
ey[i - 1] = vy[vy.size() * i / ND];
const size_t pos = vx.size() * i / ND;
std::nth_element(vx.begin() + prev, vx.begin() + pos, vx.end());
ex[i - 1] = vx[pos];
prev = pos;
}
prev = 0;
for (int i = 1; i < ND; ++i) {
const size_t pos = vy.size() * i / ND;
std::nth_element(vy.begin() + prev, vy.begin() + pos, vy.end());
ey[i - 1] = vy[pos];
prev = pos;
}
const int frames_per_bin = std::max(1, n_frames / NT);
std::vector<int32_t> cell(fulls.size(), -1);
for (size_t i = 0; i < fulls.size(); ++i) {
const Obs &o = fulls[i];
if (!usable(o)) continue;
const int it = std::min(NT - 1, o.frame / frames_per_bin);
const int ix = static_cast<int>(std::upper_bound(ex.begin(), ex.end(), o.px) - ex.begin());
const int iy = static_cast<int>(std::upper_bound(ey.begin(), ey.end(), o.py) - ey.begin());
cell[i] = (it * ND + ix) * ND + iy;
}
ParallelChunks(static_cast<int>(fulls.size()), ThreadsForWork(fulls.size(), nthreads), [&](int lo, int hi) {
for (int i = lo; i < hi; ++i) {
const Obs &o = fulls[i];
if (!usable(o)) continue;
const int it = std::min(NT - 1, o.frame / frames_per_bin);
const int ix = static_cast<int>(std::upper_bound(ex.begin(), ex.end(), o.px) - ex.begin());
const int iy = static_cast<int>(std::upper_bound(ey.begin(), ey.end(), o.py) - ey.begin());
cell[i] = (it * ND + ix) * ND + iy;
}
});
ApplyCellSurface(cell, NT * ND * ND, n_iter, n_groups, "Absorption (time x detector 12x10x10)");
}
@@ -1797,13 +1815,15 @@ void RotationScaleMerge::RefineModulation(int n_iter, int n_groups) {
constexpr int NB = 24;
const float sx = NB / (pxmax - pxmin), sy = NB / (pymax - pymin);
std::vector<int32_t> cell(fulls.size(), -1);
for (size_t i = 0; i < fulls.size(); ++i) {
const Obs &o = fulls[i];
if (!std::isfinite(o.px) || !std::isfinite(o.py)) continue;
const int ix = std::clamp(static_cast<int>((o.px - pxmin) * sx), 0, NB - 1);
const int iy = std::clamp(static_cast<int>((o.py - pymin) * sy), 0, NB - 1);
cell[i] = ix * NB + iy;
}
ParallelChunks(static_cast<int>(fulls.size()), ThreadsForWork(fulls.size(), nthreads), [&](int lo, int hi) {
for (int i = lo; i < hi; ++i) {
const Obs &o = fulls[i];
if (!std::isfinite(o.px) || !std::isfinite(o.py)) continue;
const int ix = std::clamp(static_cast<int>((o.px - pxmin) * sx), 0, NB - 1);
const int iy = std::clamp(static_cast<int>((o.py - pymin) * sy), 0, NB - 1);
cell[i] = ix * NB + iy;
}
});
ApplyCellSurface(cell, NB * NB, n_iter, n_groups, "Modulation (detector-frame 16x16)");
}
@@ -1826,9 +1846,33 @@ void RotationScaleMerge::ApplyCellSurface(const std::vector<int32_t> &cell, int
auto subset = [&](int parity) -> const std::vector<int32_t> & {
return parity < 0 ? idx_all : (parity ? idx_odd : idx_even);
};
// Per-group reference accumulators, reused by every pass (n_groups is large, so re-allocating them
// ~22x per surface is ~22x the page faults for nothing).
// Every pass below runs over one of these lists, ~22 times per surface, so the worker count is
// set by how much data there is rather than by how many threads the machine has: a small dataset
// on a large node would otherwise pay for 48 thread starts per pass to sum a few thousand terms.
const int nt = static_cast<int>(ThreadsForWork(idx_all.size(), nthreads));
// Where thread t starts in a list of n items. The boundaries depend on nothing but n and the
// thread count, so a sum split this way gives the same answer on every run.
auto part = [nt](size_t n, int t) { return n * static_cast<size_t>(t) / static_cast<size_t>(nt); };
// The per-group reference of one subset with surface A applied - the first pass of both the fit
// and the score. It stays serial: the terms of one group are spread all over the list, so the only
// way to give a thread groups of its own is to walk in group order, which trades this pass's
// near-sequential read of `fulls` for a random one over a few hundred megabytes. That is the same
// trade that made length-bucketed run ordering a loss in the combine kernel. Worth revisiting with
// per-worker accumulators (n_groups doubles each, so it costs memory rather than locality) if this
// pass ever dominates.
std::vector<double> sw(n_groups), swI(n_groups);
auto reference = [&](int parity, const std::vector<double> &A) {
const std::vector<int32_t> &sel = subset(parity);
std::fill(sw.begin(), sw.end(), 0.0);
std::fill(swI.begin(), swI.end(), 0.0);
for (const int32_t i : sel) {
const Obs &o = fulls[i];
const double a = A[cell[i]], sc = static_cast<double>(o.sigma) * o.corr * a, w = 1.0 / (sc * sc);
sw[o.group] += w; swI[o.group] += w * static_cast<double>(o.I) * o.corr * a;
}
};
// Fit the per-cell factor over the subset {frame&1 == parity} (parity < 0 = all fulls), n_iter
// alternating rounds against that subset's own reference (Tikhonov pull to 1, gauge-fixed to a
@@ -1836,14 +1880,10 @@ void RotationScaleMerge::ApplyCellSurface(const std::vector<int32_t> &cell, int
auto fit_surface = [&](int parity) -> std::vector<double> {
const std::vector<int32_t> &sel = subset(parity);
std::vector<double> A(ncell, 1.0);
// Per-thread cell accumulators, allocated once for the whole fit rather than per round.
std::vector<std::vector<double>> tcross(nt, std::vector<double>(ncell)), tref2(nt, std::vector<double>(ncell));
for (int it = 0; it < n_iter; ++it) {
std::fill(sw.begin(), sw.end(), 0.0);
std::fill(swI.begin(), swI.end(), 0.0);
for (const int32_t i : sel) {
const Obs &o = fulls[i];
const double a = A[cell[i]], sc = static_cast<double>(o.sigma) * o.corr * a, w = 1.0 / (sc * sc);
sw[o.group] += w; swI[o.group] += w * static_cast<double>(o.I) * o.corr * a;
}
reference(parity, A);
// The cell's factor is the scale that carries the observation onto the reference, and it is
// fitted by regressing the OBSERVATION on the reference: A = sum w Iref^2 / sum w Is Iref,
// the inverse of the slope of Is against Iref. Not the other way round - a least-squares
@@ -1854,16 +1894,29 @@ void RotationScaleMerge::ApplyCellSurface(const std::vector<int32_t> &cell, int
// the gauge fix spreads over the whole surface and the n_iter rounds compound, and which no
// even/odd cross-validation can see because both halves carry it equally. The per-frame
// scale (FitPerFrameG) already regresses this way round.
// Chunked over the subset, each thread summing into its own cells: unlike the pass above
// there is no ordering that keeps threads off each other's bins, and there are only ncell
// of them, so per-thread copies are cheap and the fixed chunking keeps it reproducible.
std::vector<double> cross(ncell, 0.0), ref2(ncell, 0.0);
for (const int32_t i : sel) {
const Obs &o = fulls[i];
if (sw[o.group] <= 0.0) continue;
const double Iref = swI[o.group] / sw[o.group], a = A[cell[i]];
const double Is = static_cast<double>(o.I) * o.corr * a, sc = static_cast<double>(o.sigma) * o.corr * a;
if (!std::isfinite(Iref) || Iref <= 0.0 || !(Is > 0.0) || !(sc > 0.0)) continue;
const double w = 1.0 / (sc * sc);
cross[cell[i]] += w * Is * Iref; ref2[cell[i]] += w * Iref * Iref;
}
ParallelChunks(nt, nthreads, [&](int tlo, int thi) {
for (int t = tlo; t < thi; ++t) {
std::vector<double> &xcross = tcross[t], &xref2 = tref2[t];
std::fill(xcross.begin(), xcross.end(), 0.0);
std::fill(xref2.begin(), xref2.end(), 0.0);
for (size_t k = part(sel.size(), t); k < part(sel.size(), t + 1); ++k) {
const int32_t i = sel[k];
const Obs &o = fulls[i];
if (sw[o.group] <= 0.0) continue;
const double Iref = swI[o.group] / sw[o.group], a = A[cell[i]];
const double Is = static_cast<double>(o.I) * o.corr * a, sc = static_cast<double>(o.sigma) * o.corr * a;
if (!std::isfinite(Iref) || Iref <= 0.0 || !(Is > 0.0) || !(sc > 0.0)) continue;
const double w = 1.0 / (sc * sc);
xcross[cell[i]] += w * Is * Iref; xref2[cell[i]] += w * Iref * Iref;
}
}
});
for (int t = 0; t < nt; ++t)
for (int c = 0; c < ncell; ++c) { cross[c] += tcross[t][c]; ref2[c] += tref2[t][c]; }
std::vector<double> dsorted = cross;
std::nth_element(dsorted.begin(), dsorted.begin() + dsorted.size() / 2, dsorted.end());
const double lambda = 0.1 * std::max(1e-30, dsorted[dsorted.size() / 2]);
@@ -1880,27 +1933,29 @@ void RotationScaleMerge::ApplyCellSurface(const std::vector<int32_t> &cell, int
// scored against that subset's own reference (no leakage). Lower = tighter equivalents.
auto score = [&](int parity, const std::vector<double> &A) -> double {
const std::vector<int32_t> &sel = subset(parity);
std::fill(sw.begin(), sw.end(), 0.0);
std::fill(swI.begin(), swI.end(), 0.0);
for (const int32_t i : sel) {
const Obs &o = fulls[i];
const double a = A[cell[i]], Is = static_cast<double>(o.I) * o.corr * a;
const double sc = static_cast<double>(o.sigma) * o.corr * a, w = 1.0 / (sc * sc);
sw[o.group] += w; swI[o.group] += w * Is;
}
reference(parity, A);
// Rmeas-like (sigma-INDEPENDENT) agreement of the held-out equivalents: sum|Is - Iref| / sum|Iref|.
// Scoring on the studentized deviation instead lets a surface "improve" the held-out chi^2 by
// reshaping sigma (via corr) without tightening the actual intensities - which on mis-indexed / bad
// data passes cross-validation yet worsens Rmeas. A fractional metric cannot be gamed that way.
std::vector<double> tnum(nt, 0.0), tden(nt, 0.0);
ParallelChunks(nt, nthreads, [&](int tlo, int thi) {
for (int t = tlo; t < thi; ++t) {
double num = 0.0, den = 0.0;
for (size_t k = part(sel.size(), t); k < part(sel.size(), t + 1); ++k) {
const int32_t i = sel[k];
const Obs &o = fulls[i];
if (sw[o.group] <= 0.0) continue;
const double a = A[cell[i]], Is = static_cast<double>(o.I) * o.corr * a;
const double Iref = swI[o.group] / sw[o.group];
if (!std::isfinite(Iref) || Iref <= 0.0) continue;
num += std::abs(Is - Iref); den += Iref;
}
tnum[t] = num; tden[t] = den;
}
});
double num = 0.0, den = 0.0;
for (const int32_t i : sel) {
const Obs &o = fulls[i];
if (sw[o.group] <= 0.0) continue;
const double a = A[cell[i]], Is = static_cast<double>(o.I) * o.corr * a;
const double Iref = swI[o.group] / sw[o.group];
if (!std::isfinite(Iref) || Iref <= 0.0) continue;
num += std::abs(Is - Iref); den += Iref;
}
for (int t = 0; t < nt; ++t) { num += tnum[t]; den += tden[t]; }
return den > 0.0 ? num / den : 0.0;
};
@@ -1917,9 +1972,11 @@ void RotationScaleMerge::ApplyCellSurface(const std::vector<int32_t> &cell, int
return;
}
const std::vector<double> A = fit_surface(-1);
for (size_t i = 0; i < fulls.size(); ++i)
if (cell[i] >= 0)
fulls[i].corr = static_cast<float>(fulls[i].corr * A[cell[i]]);
ParallelChunks(static_cast<int>(fulls.size()), ThreadsForWork(fulls.size(), nthreads), [&](int lo, int hi) {
for (int i = lo; i < hi; ++i)
if (cell[i] >= 0)
fulls[i].corr = static_cast<float>(fulls[i].corr * A[cell[i]]);
});
logger.Info("{} correction: cross-validated, held-out gain {:.1f}%",
name, 100.0 * gain / std::max(base, 1e-30));
}
image_analysis/scale_merge/SearchSpaceGroup.cpp
+36 -5
View File
@@ -2,6 +2,7 @@
// SPDX-License-Identifier: GPL-3.0-only
#include "SearchSpaceGroup.h"
#include "../../common/ParallelFor.h"
#include <algorithm>
#include <array>
@@ -315,10 +316,10 @@ SearchSpaceGroupResult SearchSpaceGroup(
key_to_index.emplace(key[i], static_cast<int>(i));
// --- Stage A: score each distinct rotation operator once ---
std::vector<uint32_t> visited(n, 0);
uint32_t epoch = 0;
auto score_operator = [&](const gemmi::Op& op) -> SpaceGroupOperatorScore {
// `visited` and `epoch` are scratch, taken as arguments rather than captured so that several
// operators can be scored at once - each worker below keeps its own pair.
auto score_operator = [&](const gemmi::Op& op, std::vector<uint32_t>& visited,
uint32_t& epoch) -> SpaceGroupOperatorScore {
++epoch;
std::vector<double> x, y;
for (size_t i = 0; i < n; ++i) {
@@ -362,12 +363,14 @@ SearchSpaceGroupResult SearchSpaceGroup(
};
std::map<std::array<int, 9>, SpaceGroupOperatorScore> op_cache;
std::vector<uint32_t> visited(n, 0);
uint32_t epoch = 0;
auto operator_score = [&](const gemmi::Op& op) -> const SpaceGroupOperatorScore& {
const auto rk = RotKey(op);
auto it = op_cache.find(rk);
if (it != op_cache.end())
return it->second;
return op_cache.emplace(rk, score_operator(op)).first->second;
return op_cache.emplace(rk, score_operator(op, visited, epoch)).first->second;
};
// Conjugate rotations (symmetry-equivalent within the point group) relate symmetry-equivalent
@@ -423,6 +426,34 @@ SearchSpaceGroupResult SearchSpaceGroup(
holohedry = HolohedryRotationSet(opt.lattice_system.value());
const auto point_groups = EnumeratePointGroups(holohedry);
// Every operator the search can ask about comes from this list, and scoring one is a pass over the
// whole merge with a hash lookup per reflection - the most expensive thing in here. They do not
// depend on each other, so score the distinct ones now and let the search below read the cache.
// One `visited` per worker, not per operator: it is as long as the merge, so allocating it per
// operator would cost more than the scoring.
{
std::vector<gemmi::Op> distinct;
std::vector<std::array<int, 9>> keys;
for (const auto& pg : point_groups)
for (const auto& rot : pg.rotations) {
const auto rk = RotKey(rot);
if (std::find(keys.begin(), keys.end(), rk) == keys.end()) {
keys.push_back(rk);
distinct.push_back(rot);
}
}
std::vector<SpaceGroupOperatorScore> scored(distinct.size());
ParallelChunks(static_cast<int>(distinct.size()),
std::min(opt.nthreads, distinct.size()), [&](int lo, int hi) {
std::vector<uint32_t> scratch(n, 0);
uint32_t ep = 0;
for (int i = lo; i < hi; ++i)
scored[i] = score_operator(distinct[i], scratch, ep);
});
for (size_t i = 0; i < distinct.size(); ++i)
op_cache.emplace(keys[i], scored[i]);
}
// Mapping every observation onto its symmetry representative under a candidate's rotations - one
// apply_to_hkl + Canonicalize per observation per operator - is the expensive half of BOTH
// quantities below, and both need exactly the same mapping. Build it once per point group.
image_analysis/scale_merge/SearchSpaceGroup.h
+4
View File
@@ -274,6 +274,10 @@ struct SearchSpaceGroupOptions {
// in it, at p <= 2e-9 - loose enough that three well-measured dead axial reflections clear it,
// tight enough that two do not.
double min_screw_absence_evidence = 20.0;
// Workers for the operator-correlation stage, which is the bulk of the search: one pass over the
// whole merge per candidate rotation, and the rotations are independent of each other. 1 = serial.
size_t nthreads = 1;
};
struct SearchSpaceGroupResult {
rugnux/Rugnux.cpp
+1
View File
@@ -2193,6 +2193,7 @@ ProcessResult Rugnux::RunPipeline(RugnuxObserver *observer, bool write_output, b
bool promoted_point_group = prepass_merge_sg_.has_value() && prepass_promoted_point_group_;
if (!experiment_.GetGemmiSpaceGroup().has_value()) {
SearchSpaceGroupOptions sg_opts;
sg_opts.nthreads = static_cast<size_t>(std::max(1, config_.nthreads));
sg_opts.merge_friedel = experiment_.GetScalingSettings().GetMergeFriedel();
sg_opts.d_min_limit_A = std::max<double>(
d_min_search, experiment_.GetScalingSettings().GetHighResolutionLimit_A().value_or(0.0));