diff --git a/image_analysis/scale_merge/RotationScaleMerge.cpp b/image_analysis/scale_merge/RotationScaleMerge.cpp index 9ee732ae..31dd013d 100644 --- a/image_analysis/scale_merge/RotationScaleMerge.cpp +++ b/image_analysis/scale_merge/RotationScaleMerge.cpp @@ -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 fit_scratch; auto fit_error_model = [&](const std::vector &pool) { - // Takes the samples by value (it sorts them in place). - auto fit_ab = [&](std::vector 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 &in) { + fit_scratch.assign(in.begin(), in.end()); + std::vector &smp = fit_scratch; if (smp.size() < static_cast(8 * n_bins)) return; - std::sort(smp.begin(), smp.end(), [](const Sample &a, const Sample &b) { return a.I2 < b.I2; }); - std::vector 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 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 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(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 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(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(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> 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 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 iv(start[n_groups]); + { + std::vector 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); }