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

Bin the error model's samples without sorting them

Browse Source 
The (a, b) fit wants sixteen equal-count bins in I^2 and takes three medians out of each. It was
getting them by sorting the whole pool - millions of 32-byte samples - and it did that fourteen
times a run: the fit runs once per merge and twice where the resolution cutoff refits, the outlier
refit doubles it again, and there are five merges. Each call also took its pool BY VALUE, so every
one of those began by copying tens of megabytes, and each bin then built three more vectors by
push_back to hand to a median.

A bin only has to be the right SET. Put each boundary in place with nth_element instead, splitting
the boundaries down the middle so every level halves the range it works on - four levels of linear
work against n log n - and take the three medians straight off the bin's own span with the field
wanted, which is what median_of was doing anyway: it returns the lower median, exactly the element
nth_element leaves at that index. No copy is made at all, and the sixteen bins are disjoint so they
divide over the cores.

The comparator is now total. The sort it replaces was not stable, so which of two samples of equal
I^2 landed in which bin was decided by the order the pool happened to arrive in - and the refit is
handed a different order from the first fit. Ordering on the remaining fields, which are in the same
cache line, makes the bin a property of the samples instead. This is why the merged intensities are
not byte-identical to the previous release on about half a percent of reflections, at a median
difference of zero and a worst case of 1.2e-2: those are the ties, whose old resolution was
arbitrary. Every fitted (a, b, ISa, chi2) in the run agrees to four significant figures.

The per-group outlier median goes the same way. It was building a vector per ASU group to hold a
handful of floats - over a million allocations, their growth and their frees, five times a run -
where the counts were already to hand from the pass above. One flat array with a per-group span
gives the identical median, since a median does not care how the multiset was laid out.

Measured together with the previous commit on a high-multiplicity rotation set: 38.3 s -> 35.1 s.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_011n8riB6X59oRjkrSHzNPAU
This commit is contained in:
jungfrau
2026-08-23 11:52:54 -04:00
co-authored by Claude Opus 5
parent c9fc46e6e2
commit 4a537dbfc2
1 changed files with 74 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files
image_analysis/scale_merge/RotationScaleMerge.cpp
+74 -22
View File
@@ -2551,25 +2551,63 @@ RotationScaleMerge::Result RotationScaleMerge::MergeAndStats(int n_groups, bool
constexpr int n_bins = 16;
// Fit (a, b) from the intensity-binned median deviations of a pool of samples, then refit on that
// pool's misfit-free subset. A lambda because the pool changes once the cutoff below is known.
std::vector<Sample> fit_scratch;
auto fit_error_model = [&](const std::vector<Sample> &pool) {
// Takes the samples by value (it sorts them in place).
auto fit_ab = [&](std::vector<Sample> smp) {
// Works on a scratch copy, and the caller's pool keeps its order. That matters: the refit
// below builds its own pool by walking this one, and the partition is not a total order, so
// handing the refit a reordered pool moves which of two samples of equal I2 lands in which
// bin. One reused buffer rather than a fresh vector per call - the copy is a memcpy and costs
// a few milliseconds against the hundreds the sort it replaces used to.
auto fit_ab = [&](const std::vector<Sample> &in) {
fit_scratch.assign(in.begin(), in.end());
std::vector<Sample> &smp = fit_scratch;
if (smp.size() < static_cast<size_t>(8 * n_bins))
return;
std::sort(smp.begin(), smp.end(), [](const Sample &a, const Sample &b) { return a.I2 < b.I2; });
std::vector<double> bs2, bI2, bd2;
const size_t per = smp.size() / n_bins;
for (int bin = 0; bin < n_bins; ++bin) {
const size_t lo = bin * per;
const size_t hi = (bin == n_bins - 1) ? smp.size() : lo + per;
std::vector<double> vs2, vI2, vd2;
for (size_t i = lo; i < hi; ++i) {
vs2.push_back(smp[i].s2); vI2.push_back(smp[i].I2); vd2.push_back(smp[i].dev2);
// The bins are equal COUNTS in I2 rank, and all that comes out of one is three medians - so
// the bin has to be the right SET, not a sorted one. Put each boundary in place with
// nth_element instead, splitting the boundaries down the middle so each level halves the
// range it works on: order n per level and four levels, against n log n for the sort. The
// medians are then taken off the bin's own span, by the field wanted, with no copy at all -
// the three vectors this used to build were the whole pool again, three times over, and
// median_of takes the lower median, which is what nth_element leaves at that index.
// Total, so which of two samples of equal I2 falls in which bin is a property of the
// samples rather than of the algorithm that partitioned them. The full sort this replaces
// was unstable, so it resolved those ties by whatever order it happened to receive - and
// the refit below is handed a different order from the first fit. Ordering on the other
// two fields costs nothing (they are already in the cache line) and settles it for good.
const auto by_I2 = [](const Sample &a, const Sample &b) {
if (a.I2 != b.I2) return a.I2 < b.I2;
if (a.s2 != b.s2) return a.s2 < b.s2;
if (a.dev2 != b.dev2) return a.dev2 < b.dev2;
return a.d < b.d;
};
const std::function<void(size_t, size_t, int, int)> split =
[&](size_t lo, size_t hi, int b0, int b1) {
if (b0 >= b1) return;
const int mid = (b0 + b1) / 2;
const size_t k = static_cast<size_t>(mid) * per;
std::nth_element(smp.begin() + lo, smp.begin() + k, smp.begin() + hi, by_I2);
split(lo, k, b0, mid);
split(k, hi, mid + 1, b1);
};
split(0, smp.size(), 1, n_bins);
std::vector<double> bs2(n_bins), bI2(n_bins), bd2(n_bins);
ParallelChunks(n_bins, ThreadsForWork(smp.size(), nthreads, 8 * 32768), [&](int blo, int bhi) {
for (int bin = blo; bin < bhi; ++bin) {
const size_t lo = static_cast<size_t>(bin) * per;
const size_t hi = (bin == n_bins - 1) ? smp.size() : lo + per;
const auto mid = smp.begin() + lo + (hi - lo) / 2;
std::nth_element(smp.begin() + lo, mid, smp.begin() + hi,
[](const Sample &a, const Sample &b) { return a.s2 < b.s2; });
bs2[bin] = mid->s2;
std::nth_element(smp.begin() + lo, mid, smp.begin() + hi, by_I2);
bI2[bin] = mid->I2;
std::nth_element(smp.begin() + lo, mid, smp.begin() + hi,
[](const Sample &a, const Sample &b) { return a.dev2 < b.dev2; });
bd2[bin] = mid->dev2 / CHI2_1_MEDIAN;
}
bs2.push_back(median_of(vs2));
bI2.push_back(median_of(vI2));
bd2.push_back(median_of(vd2) / CHI2_1_MEDIAN);
}
});
// `b` is identified ONLY by the spread of I^2/sigma^2 across the bins, and the bins hold equal
// COUNTS - so when fewer reflections are strong than one bin holds (1/16 of the pool), the top
// bin's median sits at an intensity where b cannot be measured, and the fit assigns it the
@@ -2641,7 +2679,7 @@ RotationScaleMerge::Result RotationScaleMerge::MergeAndStats(int n_groups, bool
if (v > 0.0 && s.dev2 <= ns2 * v) kept.push_back(s);
}
if (kept.size() >= static_cast<size_t>(8 * n_bins) && kept.size() < pool.size())
fit_ab(std::move(kept));
fit_ab(kept);
}
};
{
@@ -2690,15 +2728,29 @@ RotationScaleMerge::Result RotationScaleMerge::MergeAndStats(int n_groups, bool
// Per-group outlier-rejection median of I*corr (host both paths - a per-group median is awkward on
// the GPU; cheap here, cnt >= 2 filter from the em pass). Fed to the merge accumulate.
if (reject_outliers) {
std::vector<std::vector<float>> iv(n_groups);
// One flat array with a per-group span, not a vector per group: n_groups is over a million
// on a P1 pass, so a vector each is a million allocations, their growth, and a million frees
// - five times a run. A median is the k-th smallest of a multiset and does not care how the
// multiset was laid out, so this is the same number.
std::vector<int32_t> start(n_groups + 1, 0);
for (const auto &o : fulls)
if (usable_merge(o) && cnt[o.group] >= 2)
iv[o.group].push_back(o.I * o.corr);
for (int g = 0; g < n_groups; ++g)
if (!iv[g].empty()) {
std::nth_element(iv[g].begin(), iv[g].begin() + iv[g].size() / 2, iv[g].end());
reject_median[g] = iv[g][iv[g].size() / 2];
if (usable_merge(o) && cnt[o.group] >= 2) start[o.group + 1]++;
for (int g = 0; g < n_groups; ++g) start[g + 1] += start[g];
std::vector<float> iv(start[n_groups]);
{
std::vector<int32_t> fill(start.begin(), start.end() - 1);
for (const auto &o : fulls)
if (usable_merge(o) && cnt[o.group] >= 2) iv[fill[o.group]++] = o.I * o.corr;
}
ParallelChunks(n_groups, ThreadsForWork(iv.size(), nthreads), [&](int glo, int ghi) {
for (int g = glo; g < ghi; ++g) {
const int lo = start[g], hi = start[g + 1];
if (lo == hi) continue;
const auto mid = iv.begin() + lo + (hi - lo) / 2;
std::nth_element(iv.begin() + lo, mid, iv.begin() + hi);
reject_median[g] = *mid;
}
});
}
fit_error_model(samples);
}