diff --git a/image_analysis/scale_merge/ErrorModel.cpp b/image_analysis/scale_merge/ErrorModel.cpp index dc2868ad9..d6b8ed5e1 100644 --- a/image_analysis/scale_merge/ErrorModel.cpp +++ b/image_analysis/scale_merge/ErrorModel.cpp @@ -7,7 +7,6 @@ #include #include #include -#include #include @@ -156,60 +155,93 @@ bool ErrorModelUpdate(const std::vector &bs2, const std::vector return !settled; } -StillsErrorModelFit FitStillsErrorModel(std::vector &pool) { +StillsErrorModelFit FitStillsErrorModel(const std::vector &pool, size_t nthreads, + const StillsErrorModelFit &start) { constexpr int n_bins = ERROR_MODEL_BINS; constexpr int n_sub = 4; StillsErrorModelFit fit; if (pool.size() < static_cast(8 * n_bins * n_sub)) return fit; - // Bins of counting (I/sigma)^2, each split by the partiality term's share of the intensity term. Total - // orders, so which sample lands in which bin is a property of the samples alone. - auto snr = [](const StillsErrorModelSample &x) { return x.I2 / x.s2; }; - auto share = [](const StillsErrorModelSample &x) { return x.I2 > 0.0 ? x.P2 / x.I2 : 0.0; }; - std::sort(pool.begin(), pool.end(), [&](const auto &x, const auto &y) { - return std::make_tuple(snr(x), x.s2, x.I2, x.P2, x.dev2) < std::make_tuple(snr(y), y.s2, y.I2, y.P2, y.dev2); + // Bins of counting (I/sigma)^2, each split by the partiality term's share of the intensity term. The + // boundaries are put in place with nth_element on (key, position in the pool) - a total order, so which + // sample lands in which bin depends on the samples and the order the caller gave them in alone. The keys + // are computed once and carried in the comparison, never recomputed inside it. + struct Key { double v; size_t i; }; + const auto by_key = [](const Key &p, const Key &q) { return p.v < q.v || (p.v == q.v && p.i < q.i); }; + std::vector key(pool.size()); + ParallelChunks(static_cast(pool.size()), ThreadsForWork(pool.size(), nthreads), [&](int lo, int hi) { + for (int i = lo; i < hi; ++i) + key[i] = {pool[i].I2 / pool[i].s2, static_cast(i)}; }); const size_t per = pool.size() / n_bins; - std::vector> groups; - for (int bin = 0; bin < n_bins; ++bin) { - const size_t lo = bin * per, hi = bin == n_bins - 1 ? pool.size() : (bin + 1) * per; - std::sort(pool.begin() + lo, pool.begin() + hi, [&](const auto &x, const auto &y) { - return std::make_tuple(share(x), x.s2, x.I2, x.dev2) < std::make_tuple(share(y), y.s2, y.I2, y.dev2); - }); - const size_t sub = (hi - lo) / n_sub; - for (int q = 0; q < n_sub; ++q) - groups.emplace_back(lo + q * sub, q == n_sub - 1 ? hi : lo + (q + 1) * sub); - } + // The split runs down the middle, and the two halves of a split are disjoint ranges, so they run side + // by side while the range is still worth a thread (as in FitErrorModel). + constexpr size_t SPLIT_MIN_PARALLEL = 1u << 17; + const std::function split = [&](size_t lo, size_t hi, int b0, int b1, int depth) { + if (b0 >= b1) return; + const int mid = (b0 + b1) / 2; + const size_t k = static_cast(mid) * per; + std::nth_element(key.begin() + lo, key.begin() + k, key.begin() + hi, by_key); + if (depth > 0 && hi - lo > SPLIT_MIN_PARALLEL) { + auto left = std::async(std::launch::async, [&, lo, k, b0, mid, depth] { + RestoreThreadAffinity(); + split(lo, k, b0, mid, depth - 1); + }); + split(k, hi, mid + 1, b1, depth - 1); + left.get(); + } else { + split(lo, k, b0, mid, 0); + split(k, hi, mid + 1, b1, 0); + } + }; + split(0, key.size(), 1, n_bins, nthreads > 1 ? 3 : 0); // b only where the strongest bin reaches past the counting term, as in FitErrorModel. std::vector top; - for (size_t i = (n_bins - 1) * per; i < pool.size(); ++i) - top.push_back(snr(pool[i])); + for (size_t j = (n_bins - 1) * per; j < key.size(); ++j) + top.push_back(key[j].v); std::nth_element(top.begin(), top.begin() + top.size() / 2, top.end()); const bool b_measured = top[top.size() / 2] > ERROR_MODEL_B_LEVER_MIN; + // Each bin in four by the share, and the samples copied out in that order. + std::vector> groups(n_bins * n_sub); + std::vector x(pool.size()); + ParallelFor(n_bins, nthreads, [&](int bin) { + const size_t lo = bin * per, hi = bin == n_bins - 1 ? key.size() : (bin + 1) * per; + for (size_t j = lo; j < hi; ++j) { + const auto &o = pool[key[j].i]; + key[j].v = o.I2 > 0.0 ? o.P2 / o.I2 : 0.0; + } + const size_t sub = (hi - lo) / n_sub; + for (int q = 1; q < n_sub; ++q) + std::nth_element(key.begin() + lo + (q - 1) * sub, key.begin() + lo + q * sub, key.begin() + hi, by_key); + for (int q = 0; q < n_sub; ++q) + groups[bin * n_sub + q] = {lo + q * sub, q == n_sub - 1 ? hi : lo + (q + 1) * sub}; + for (size_t j = lo; j < hi; ++j) + x[j] = pool[key[j].i]; + }); + key = {}; const int ng = static_cast(groups.size()); std::vector z2(pool.size()); Eigen::MatrixXd X(ng, 3); Eigen::VectorXd y(ng); - double a = 1.0, b2 = 0.0, c2 = 0.0; + double a = start.a, b2 = start.b2, c2 = start.c2; for (int iter = 0; iter < 30; ++iter) { // Per group: the three terms divided by the current variance, and the median normalised deviation // expressed as the sum it stands for. - for (int g = 0; g < ng; ++g) { + ParallelFor(ng, ThreadsForWork(pool.size(), nthreads), [&](int g) { const auto [lo, hi] = groups[g]; double s = 0.0, I = 0.0, P = 0.0; for (size_t i = lo; i < hi; ++i) { - const auto &x = pool[i]; - const double v = a * x.s2 + b2 * x.I2 + c2 * x.P2; - s += x.s2 / v; I += x.I2 / v; P += x.P2 / v; - z2[i] = x.dev2 / v; + const double v = a * x[i].s2 + b2 * x[i].I2 + c2 * x[i].P2; + s += x[i].s2 / v; I += x[i].I2 / v; P += x[i].P2 / v; + z2[i] = x[i].dev2 / v; } const auto mid = z2.begin() + lo + (hi - lo) / 2; std::nth_element(z2.begin() + lo, mid, z2.begin() + hi); X.row(g) << s, I, P; y(g) = *mid / ERROR_MODEL_CHI2_1_MEDIAN * static_cast(hi - lo); - } + }); // Every group weighs the same in relative terms; a negative term is dropped and the rest refitted. const Eigen::VectorXd w = y.array().square().inverse(); std::vector terms = b_measured ? std::vector{0, 1, 2} : std::vector{0, 2}; diff --git a/image_analysis/scale_merge/ErrorModel.h b/image_analysis/scale_merge/ErrorModel.h index 11898071e..d76981479 100644 --- a/image_analysis/scale_merge/ErrorModel.h +++ b/image_analysis/scale_merge/ErrorModel.h @@ -82,4 +82,7 @@ struct StillsErrorModelFit { bool b_resolved = false; // as ErrorModelFit::b_resolved }; -StillsErrorModelFit FitStillsErrorModel(std::vector &pool); +// start: where the iteration begins - identity by default, or the fit of a previous round on a pool that +// has moved little since. +StillsErrorModelFit FitStillsErrorModel(const std::vector &pool, size_t nthreads = 1, + const StillsErrorModelFit &start = {}); diff --git a/image_analysis/scale_merge/Merge.cpp b/image_analysis/scale_merge/Merge.cpp index b00142b75..d6924587a 100644 --- a/image_analysis/scale_merge/Merge.cpp +++ b/image_analysis/scale_merge/Merge.cpp @@ -7,6 +7,7 @@ #include #include #include +#include #include #include @@ -298,9 +299,7 @@ void MergeOnTheFly::RefineErrorModel(const std::vector &outc // AddImage uses, down to the floor of the merge with an error model. struct Obs { float I, sigma, corr, var_bkg, p; }; struct KeyedObs { uint64_t key; Obs obs; }; - std::unordered_map> groups; - std::vector keys; // in order of first appearance, so every sum below is in a fixed order - + std::vector all; // in image order ReduceImagesInOrder( outcomes.size(), nthreads, [&](size_t i, std::vector &out) { @@ -325,86 +324,147 @@ void MergeOnTheFly::RefineErrorModel(const std::vector &outc } }, [&](size_t, const std::vector &out) { - for (const auto &o: out) { - auto &g = groups[o.key]; - if (g.empty()) - keys.push_back(o.key); - g.push_back(o.obs); - } + all.insert(all.end(), out.begin(), out.end()); }); + // The observations of each reflection side by side, in key order and within a reflection in image + // order - a sort on (key, place), a total order, so the same on any number of workers. Every sum + // below runs over them in that fixed order. + struct KeyPlace { uint64_t key; size_t place; }; + std::vector order(all.size()); + const size_t obs_threads = ThreadsForWork(all.size(), nthreads); + ParallelChunks(static_cast(all.size()), obs_threads, [&](int lo, int hi) { + for (int i = lo; i < hi; ++i) + order[i] = {all[i].key, static_cast(i)}; + }); + ParallelSort(order.begin(), order.end(), nthreads, [](const KeyPlace &x, const KeyPlace &y) { + return x.key < y.key || (x.key == y.key && x.place < y.place); + }); + std::vector flat(all.size()); + ParallelChunks(static_cast(all.size()), obs_threads, [&](int lo, int hi) { + for (int i = lo; i < hi; ++i) + flat[i] = all[order[i].place].obs; + }); + all = {}; + std::vector keys; + std::vector start; // reflection k is flat[start[k] .. start[k + 1]) + for (size_t i = 0; i < order.size(); ++i) + if (i == 0 || order[i].key != order[i - 1].key) { + keys.push_back(order[i].key); + start.push_back(i); + } + start.push_back(order.size()); + order = {}; + const auto group = [&](int k) { return std::span(flat.data() + start[k], start[k + 1] - start[k]); }; + const int n_keys = static_cast(keys.size()); + const size_t key_threads = ThreadsForWork(keys.size(), nthreads, 1024); + // --- 2. The starting mean of each reflection (counting-sigma weights) and, where there are three // observations for it to stand on, the median - the robust centre of the outlier test. - std::unordered_map mean; - for (const auto key: keys) { - const auto &obs = groups[key]; - double sum_w = 0.0, sum_wI = 0.0; - for (const auto &o: obs) { - const double w = 1.0 / (static_cast(o.sigma) * o.sigma); - sum_w += w; - sum_wI += w * o.I; + std::vector mean(n_keys); + std::vector median(n_keys, NAN); + ParallelChunks(n_keys, key_threads, [&](int lo, int hi) { + for (int k = lo; k < hi; ++k) { + const auto obs = group(k); + double sum_w = 0.0, sum_wI = 0.0; + for (const auto &o: obs) { + const double w = 1.0 / (static_cast(o.sigma) * o.sigma); + sum_w += w; + sum_wI += w * o.I; + } + mean[k] = sum_wI / sum_w; + if (reject_outliers && obs.size() >= 3) { + std::vector iv; + iv.reserve(obs.size()); + for (const auto &o: obs) + iv.push_back(o.I); + std::nth_element(iv.begin(), iv.begin() + iv.size() / 2, iv.end()); + median[k] = iv[iv.size() / 2]; + } } - mean[key] = sum_wI / sum_w; - if (reject_outliers && obs.size() >= 3) { - std::vector iv; - iv.reserve(obs.size()); - for (const auto &o: obs) - iv.push_back(o.I); - std::nth_element(iv.begin(), iv.begin() + iv.size() / 2, iv.end()); - reject_median_I[key] = iv[iv.size() / 2]; - } - } + }); + for (int k = 0; k < n_keys; ++k) + if (!std::isnan(median[k])) + reject_median_I[keys[k]] = median[k]; // --- 3. Fit the model (ErrorModel.h, FitStillsErrorModel) on the scatter about the reflection means, // the counting variance rebuilt at the mean, and refit with the means the fitted model weights: // the means and the model each depend on the other, and a single pass leaves the weights of // the strong reflections at the counting-sigma means the fit started from. Measured on serial // and rotation-as-serial sets, a few rounds raise the anomalous signal of the merge; they have - // settled after four. + // settled after four, most often already after two or three: a round that leaves the model where + // the previous one put it (to 1e-3) would only reproduce the same means, and ends the iteration. + // Those rounds only carry the means to where the model settles, so they fit every stride-th + // observation (some four million of them); one last round fits them all and makes the final means. + // Each observation's sample goes to its own slot, so the pool is in key order whatever the number + // of workers; the slots left empty (single observations, the leverage cut) are dropped after. constexpr int ROUNDS = 4; + const size_t stride = std::max(1, flat.size() >> 22); + bool last_round = false; std::vector samples; StillsErrorModelFit fit; double a = 1.0, b2 = 0.0, c2 = 0.0; auto variance = [&](const Obs &o, double m) { return a * CountingVariance(o.sigma, o.corr, o.var_bkg, o.p, m) + (b2 + c2 * PartialityTerm(o.p)) * m * m; }; - for (int round = 0; round < ROUNDS; ++round) { - samples.clear(); - for (const auto key: keys) { - const auto &obs = groups[key]; - if (obs.size() < 2) - continue; - const double m = mean[key]; - // The residual from the weighted mean has E[(I_i - )^2] = var_i (1 - h_i), h_i = w_i / sum_w its - // leverage, with the weights of the model that made the mean. - double sum_w = 0.0; - for (const auto &o: obs) - sum_w += 1.0 / variance(o, m); - for (const auto &o: obs) { - const double factor = 1.0 - 1.0 / (variance(o, m) * sum_w); - if (factor < 0.05) + for (int round = 0; ; ++round) { + const size_t step = last_round ? 1 : stride; + samples.assign((flat.size() + step - 1) / step, StillsErrorModelSample{NAN, NAN, NAN, NAN}); + ParallelChunks(n_keys, key_threads, [&](int lo, int hi) { + std::vector var; + for (int k = lo; k < hi; ++k) { + const auto obs = group(k); + if (obs.size() < 2) continue; - const double resid = static_cast(o.I) - m; - samples.push_back({CountingVariance(o.sigma, o.corr, o.var_bkg, o.p, m), m * m, - PartialityTerm(o.p) * m * m, resid * resid / factor}); + const double m = mean[k]; + // The residual from the weighted mean has E[(I_i - )^2] = var_i (1 - h_i), h_i = w_i / sum_w + // its leverage, with the weights of the model that made the mean. + var.resize(obs.size()); + double sum_w = 0.0; + for (size_t j = 0; j < obs.size(); ++j) { + var[j] = variance(obs[j], m); + sum_w += 1.0 / var[j]; + } + for (size_t j = 0; j < obs.size(); ++j) { + if ((start[k] + j) % step != 0) + continue; + const auto &o = obs[j]; + const double factor = 1.0 - 1.0 / (var[j] * sum_w); + if (factor < 0.05) + continue; + const double resid = static_cast(o.I) - m; + samples[(start[k] + j) / step] = {CountingVariance(o.sigma, o.corr, o.var_bkg, o.p, m), m * m, + PartialityTerm(o.p) * m * m, resid * resid / factor}; + } } - } - fit = FitStillsErrorModel(samples); + }); + samples.erase(std::remove_if(samples.begin(), samples.end(), + [](const StillsErrorModelSample &x) { return std::isnan(x.dev2); }), + samples.end()); + const StillsErrorModelFit last = fit; + fit = FitStillsErrorModel(samples, nthreads, last); // each round starts from the last one's model if (!fit.active) return; // too little multiplicity to fit -> leave identity a = fit.a; b2 = fit.b2; c2 = fit.c2; - for (const auto key: keys) { - double sum_w = 0.0, sum_wI = 0.0; - for (const auto &o: groups[key]) { - const double w = 1.0 / variance(o, mean[key]); - sum_w += w; - sum_wI += w * o.I; + const auto close = [](double x, double y) { return std::fabs(x - y) <= 1e-3 * std::fabs(y); }; + const bool settled = round > 0 && close(a, last.a) && close(b2, last.b2) && close(c2, last.c2); + ParallelChunks(n_keys, key_threads, [&](int lo, int hi) { + for (int k = lo; k < hi; ++k) { + double sum_w = 0.0, sum_wI = 0.0; + for (const auto &o: group(k)) { + const double w = 1.0 / variance(o, mean[k]); + sum_w += w; + sum_wI += w * o.I; + } + mean[k] = sum_wI / sum_w; } - mean[key] = sum_wI / sum_w; - } + }); + if (last_round || ((settled || round + 1 == ROUNDS) && step == 1)) + break; + last_round = settled || round + 1 == ROUNDS; } - for (const auto key: keys) - error_model_mean_I[key] = static_cast(mean[key]); + for (int k = 0; k < n_keys; ++k) + error_model_mean_I[keys[k]] = static_cast(mean[k]); error_model_a = a; error_model_b = std::sqrt(b2); diff --git a/tests/MergeScaleTest.cpp b/tests/MergeScaleTest.cpp index 74b9e4825..e9ae92b10 100644 --- a/tests/MergeScaleTest.cpp +++ b/tests/MergeScaleTest.cpp @@ -546,6 +546,10 @@ TEST_CASE("StillsErrorModel_RecoversThePartialityTerm") { } const auto fit = FitStillsErrorModel(pool); REQUIRE(fit.active); + const auto many = FitStillsErrorModel(pool, 16); // the same numbers on any number of threads + CHECK(many.a == fit.a); + CHECK(many.b2 == fit.b2); + CHECK(many.c2 == fit.c2); CHECK(fit.a == Catch::Approx(a).epsilon(0.05)); CHECK(std::sqrt(fit.b2) == Catch::Approx(b).epsilon(0.05)); CHECK(std::sqrt(fit.c2) == Catch::Approx(c).epsilon(0.05));