From aa6ec193e9c999f1fb8648f3af553282904c5841 Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Mon, 28 Sep 2026 16:52:13 +0200 Subject: [PATCH] Unmerged MTZ: build the file on several workers, byte-identical The unmerged export (_unmerged.mtz) was serial apart from its final row sort: gathering the partials, sorting them by (h,k,l,image), summing the rocking events and filling the rows all ran on one thread. On a 17 M partial rotation set that was 2.8 s (sort 1.6 s, event sums 0.6 s, gather 0.25 s, rows 0.19 s); on a synthetic 122 M partial sweep 23 s. - The partials are gathered into preallocated slots per outcome, left uninitialised so the workers first-touch the pages, and sorted with ParallelSort. The sort key now ends with the part's input position (outcome, index), so the order is total and the same on any number of workers. It differs from the previous std::sort order only among parts with equal (h,k,l,image_number); measured on four in-house sweeps (464-6138 such positions each), every summed full is bitwise the same. - Events are summed in chunks whose bounds move forward to the next event start, so no event is split; the chunks are returned as pieces in order rather than concatenated. - Rows are filled by index into a preallocated table; absences and the batch range are still decided in one serial pass. The final row gather after the sort is parallel too. Measured: the 17 M set 2.82 s -> 0.86 s, the 8 M set 1.37 -> 0.49 s, all while the P1 cross-check merge shares the pool; a synthetic 122 M partial sweep 22.7 s -> 7.7 s at 32 workers. What remains is mostly ParallelSort's pairwise merge levels. The file write (0.1 s for 112 MB) is I/O and left as it is. p.hkl, p.mtz, p_P1.mtz and p_unmerged.mtz are md5-identical to the base on four in-house sets, on the GPU and the CPU build, and at -N 1. Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01D1G8gJVAy6gp1K5Dz3NE5C --- image_analysis/WriteReflections.cpp | 258 ++++++++++++++++------------ tests/WriteReflectionsTest.cpp | 48 ++++++ 2 files changed, 196 insertions(+), 110 deletions(-) diff --git a/image_analysis/WriteReflections.cpp b/image_analysis/WriteReflections.cpp index bd654a999..798467ec4 100644 --- a/image_analysis/WriteReflections.cpp +++ b/image_analysis/WriteReflections.cpp @@ -14,6 +14,7 @@ #include #include #include +#include #include #include #include @@ -648,90 +649,109 @@ float DetectorY(const Reflection &r) { return std::isfinite(r.observed_y) ? r.ob // not the number rugnux scales on (the merge recomputes every partiality from a frame-order-smoothed // mosaicity, RotationScaleMerge::SmoothMosaicity), so clamping it to 1 would hide the spread and buy // a reading program nothing. -std::vector SumRockingEvents(const std::vector &outcomes, - double min_partiality, double min_captured_fraction, - float wedge_deg) { +// +// The fulls come back in pieces, one per worker, which laid end to end are all of them in (h,k,l) +// order - joining them would be one more serial copy of the whole list. +std::vector> SumRockingEvents(const std::vector &outcomes, + double min_partiality, double min_captured_fraction, + float wedge_deg, size_t nthreads) { const float max_frame_gap = RockingEventFrameGap(wedge_deg); // == RotationScaleMerge's - // The sort key travels with the part instead of being read back through the pointer, the way the + // The sort key travels with the part instead of being read back from the reflection, the way the // merge's own ingest sort carries it (RotationScaleMerge's SortKey): there are millions of parts - // and an indirect compare is a cache miss on every one of them. The keys are the same values in - // the same order, so introsort makes the same comparisons and the same swaps and leaves the same - // order - which matters, because two parts can genuinely share (h,k,l) and image_number and the - // event sums below are floating point. + // and an indirect compare is a cache miss on every one of them. Two parts can genuinely share + // (h,k,l) and image_number, and the event sums below are floating point, so the part's position + // in the input - its outcome, then its place there - is the last key: the order is total, and the + // sort gives the same one on any number of workers. struct Part { int32_t h, k, l; float image_number; - const Reflection *r; + uint32_t o, n; // the outcome, and the part's place in it }; - std::vector parts; - size_t n_parts = 0; - for (const auto &outcome : outcomes) - n_parts += outcome.reflections.size(); - parts.reserve(n_parts); - for (const auto &outcome : outcomes) - for (const auto &r : outcome.reflections) - parts.push_back({r.h, r.k, r.l, r.image_number, &r}); - std::sort(parts.begin(), parts.end(), [](const Part &a, const Part &b) { - return std::tie(a.h, a.k, a.l, a.image_number) < std::tie(b.h, b.k, b.l, b.image_number); + std::vector first(outcomes.size() + 1, 0); + for (size_t o = 0; o < outcomes.size(); ++o) + first[o + 1] = first[o] + outcomes[o].reflections.size(); + // Left uninitialised, so that the pages are first touched by the workers filling them. + const size_t n_parts = first.back(); + const auto parts = std::make_unique_for_overwrite(n_parts); + ParallelFor(static_cast(outcomes.size()), nthreads, [&](int o) { + for (size_t m = 0; m < outcomes[o].reflections.size(); ++m) { + const Reflection &r = outcomes[o].reflections[m]; + parts[first[o] + m] = {r.h, r.k, r.l, r.image_number, static_cast(o), static_cast(m)}; + } + }); + ParallelSort(parts.get(), parts.get() + n_parts, nthreads, [](const Part &a, const Part &b) { + return std::tie(a.h, a.k, a.l, a.image_number, a.o, a.n) < std::tie(b.h, b.k, b.l, b.image_number, b.o, b.n); }); - std::vector fulls; - for (size_t i = 0; i < parts.size(); ) { - size_t j = i + 1; - while (j < parts.size() && parts[j].h == parts[i].h && parts[j].k == parts[i].k - && parts[j].l == parts[i].l - && parts[j].image_number - parts[j - 1].image_number <= max_frame_gap) - ++j; + // Part i continues the event of part i - 1. + const auto continues = [&](size_t i) { + return parts[i].h == parts[i - 1].h && parts[i].k == parts[i - 1].k && parts[i].l == parts[i - 1].l + && parts[i].image_number - parts[i - 1].image_number <= max_frame_gap; + }; - double sum_p = 0.0, sum_I = 0.0, sum_var = 0.0, sum_var_bkg = 0.0; - double p_corr = 0.0, p_qe = 0.0, p_flight = 0.0, p_frame = 0.0, p_x = 0.0, p_y = 0.0, - p_delta_phi = 0.0, p_zeta = 0.0, p_bkg = 0.0; - for (size_t m = i; m < j; ++m) { - const Reflection &r = *parts[m].r; - const double p = r.partiality; - const float pc = r.prescaling_corr * r.qe_corr * r.flight_corr; // LP x QE x flight path - sum_p += p; - sum_I += static_cast(r.I) * pc; - sum_var += static_cast(r.sigma) * r.sigma * pc * pc; - sum_var_bkg += static_cast(r.var_bkg) * pc * pc; - p_corr += p * pc; - p_qe += p * r.qe_corr; - p_flight += p * r.flight_corr; - p_frame += p * r.image_number; - p_x += p * DetectorX(r); - p_y += p * DetectorY(r); - p_delta_phi += p * r.delta_phi_deg; - p_zeta += p * r.zeta; - p_bkg += p * r.bkg; + // The events are summed in chunks, each moved on to the next event start, so no event is split. + const size_t nchunks = ThreadsForWork(n_parts, nthreads); + std::vector> chunk_fulls(nchunks); + ParallelFor(static_cast(nchunks), nthreads, [&](int c) { + size_t i = c * n_parts / nchunks; + size_t end = (c + 1) * n_parts / nchunks; + while (i > 0 && i < n_parts && continues(i)) ++i; + while (end < n_parts && continues(end)) ++end; + while (i < end) { + size_t j = i + 1; + while (j < n_parts && continues(j)) + ++j; + + double sum_p = 0.0, sum_I = 0.0, sum_var = 0.0, sum_var_bkg = 0.0; + double p_corr = 0.0, p_qe = 0.0, p_flight = 0.0, p_frame = 0.0, p_x = 0.0, p_y = 0.0, + p_delta_phi = 0.0, p_zeta = 0.0, p_bkg = 0.0; + for (size_t m = i; m < j; ++m) { + const Reflection &r = outcomes[parts[m].o].reflections[parts[m].n]; + const double p = r.partiality; + const float pc = r.prescaling_corr * r.qe_corr * r.flight_corr; // LP x QE x flight path + sum_p += p; + sum_I += static_cast(r.I) * pc; + sum_var += static_cast(r.sigma) * r.sigma * pc * pc; + sum_var_bkg += static_cast(r.var_bkg) * pc * pc; + p_corr += p * pc; + p_qe += p * r.qe_corr; + p_flight += p * r.flight_corr; + p_frame += p * r.image_number; + p_x += p * DetectorX(r); + p_y += p * DetectorY(r); + p_delta_phi += p * r.delta_phi_deg; + p_zeta += p * r.zeta; + p_bkg += p * r.bkg; + } + Reflection full = outcomes[parts[i].o].reflections[parts[i].n]; + i = j; + if (sum_p < min_partiality || sum_p < min_captured_fraction) + continue; + + // The prescaling factor is applied by the writer, which multiplies I by the whole correction, so + // divide the event's own factor back out of the sums here. It is the same geometry for every part + // of one event to a median 2e-4, so the file's I / LP * QE * FLIGHT is still the raw count sum. The + // event's mean is taken on the whole correction and on the two path terms; the + // Lorentz-polarization part is what is left when those are divided out of it. + const float mean_corr = static_cast(p_corr / sum_p); + full.qe_corr = static_cast(p_qe / sum_p); + full.flight_corr = static_cast(p_flight / sum_p); + full.prescaling_corr = mean_corr / (full.qe_corr * full.flight_corr); + full.I = static_cast(sum_I / mean_corr); + full.sigma = static_cast(std::sqrt(sum_var) / mean_corr); + full.var_bkg = static_cast(sum_var_bkg / (static_cast(mean_corr) * mean_corr)); + full.partiality = static_cast(sum_p); + full.image_number = static_cast(p_frame / sum_p); + full.observed_x = static_cast(p_x / sum_p); + full.observed_y = static_cast(p_y / sum_p); + full.delta_phi_deg = static_cast(p_delta_phi / sum_p); + full.zeta = static_cast(p_zeta / sum_p); + full.bkg = static_cast(p_bkg / sum_p); + chunk_fulls[c].push_back(full); } - Reflection full = *parts[i].r; - i = j; - if (sum_p < min_partiality || sum_p < min_captured_fraction) - continue; - - // The prescaling factor is applied by the writer, which multiplies I by the whole correction, so - // divide the event's own factor back out of the sums here. It is the same geometry for every part - // of one event to a median 2e-4, so the file's I / LP * QE * FLIGHT is still the raw count sum. The - // event's mean is taken on the whole correction and on the two path terms; the - // Lorentz-polarization part is what is left when those are divided out of it. - const float mean_corr = static_cast(p_corr / sum_p); - full.qe_corr = static_cast(p_qe / sum_p); - full.flight_corr = static_cast(p_flight / sum_p); - full.prescaling_corr = mean_corr / (full.qe_corr * full.flight_corr); - full.I = static_cast(sum_I / mean_corr); - full.sigma = static_cast(std::sqrt(sum_var) / mean_corr); - full.var_bkg = static_cast(sum_var_bkg / (static_cast(mean_corr) * mean_corr)); - full.partiality = static_cast(sum_p); - full.image_number = static_cast(p_frame / sum_p); - full.observed_x = static_cast(p_x / sum_p); - full.observed_y = static_cast(p_y / sum_p); - full.delta_phi_deg = static_cast(p_delta_phi / sum_p); - full.zeta = static_cast(p_zeta / sum_p); - full.bkg = static_cast(p_bkg / sum_p); - fulls.push_back(full); - } - return fulls; + }); + return chunk_fulls; } } // namespace @@ -830,53 +850,69 @@ gemmi::Mtz UnmergedMtz(const std::vector &outcomes, // removing them turns its test into an assumption. XDS and DIALS draw the line in the same place. const char centering = mtz.spacegroup ? mtz.spacegroup->hm[0] : 'P'; + std::vector> fulls; + if (scanning && sum_partials) + fulls = SumRockingEvents(outcomes, + experiment.GetScalingSettings().GetMinPartiality(), + experiment.GetScalingSettings().GetMinCapturedFraction(), + wedge_deg, nthreads); + std::vector rows; const auto add_row = [&](const Reflection &r) { if (systematic_absence(r.h, r.k, r.l, centering)) return; - std::array hkl{r.h, r.k, r.l}; - const int isym = hkl_mover.move_to_asu(hkl); // A summed full's image_number is its rocking-curve centroid, so this is the batch the // reflection is centred on - which is what a batch means for a full everywhere else. const int batch = 1 + static_cast(std::lround(r.image_number)); if (last_batch < first_batch) { first_batch = batch; last_batch = batch; } first_batch = std::min(first_batch, batch); last_batch = std::max(last_batch, batch); - mtz.data.push_back(static_cast(hkl[0])); - mtz.data.push_back(static_cast(hkl[1])); - mtz.data.push_back(static_cast(hkl[2])); - mtz.data.push_back(static_cast((partials ? 256 : 0) + isym)); - mtz.data.push_back(static_cast(batch)); - const float corr = r.prescaling_corr * r.qe_corr * r.flight_corr; - mtz.data.push_back(r.I * corr); - mtz.data.push_back(r.sigma * corr); - mtz.data.push_back(r.partiality); - mtz.data.push_back(DetectorX(r)); - mtz.data.push_back(DetectorY(r)); - mtz.data.push_back(phi_start_deg(r.image_number) + wedge_deg / 2.0f); - // prescaling_corr is Lorentz x polarization, which is what LP means; qe_corr is the sensor - // efficiency and flight_corr the flight path beside it, each written in the divide-by - // direction, so that I / LP * QE * FLIGHT is the raw count. FLIGHT is <= 1 where QE is >= 1: - // the sensor makes an oblique reflection read high and the medium makes it read low. - mtz.data.push_back(r.prescaling_corr); - mtz.data.push_back(1.0f / r.qe_corr); - mtz.data.push_back(1.0f / r.flight_corr); - mtz.data.push_back(0.0f); // FLAG: nothing here is a rejected observation - mtz.data.push_back(r.delta_phi_deg); - mtz.data.push_back(r.zeta); - mtz.data.push_back(r.bkg); - mtz.data.push_back(r.var_bkg); + rows.push_back(&r); }; if (scanning && sum_partials) { - for (const auto &r : SumRockingEvents(outcomes, - experiment.GetScalingSettings().GetMinPartiality(), - experiment.GetScalingSettings().GetMinCapturedFraction(), - wedge_deg)) - add_row(r); + for (const auto &piece : fulls) + for (const auto &r : piece) + add_row(r); } else { for (const auto &outcome : outcomes) for (const auto &r : outcome.reflections) add_row(r); } + + // Every row has its own slot, so they are filled on several workers. + const size_t w = mtz.columns.size(); + mtz.data.resize(rows.size() * w); + ParallelChunks(static_cast(rows.size()), ThreadsForWork(rows.size(), nthreads), [&](int lo, int hi) { + for (int i = lo; i < hi; ++i) { + const Reflection &r = *rows[i]; + float *row = &mtz.data[static_cast(i) * w]; + std::array hkl{r.h, r.k, r.l}; + const int isym = hkl_mover.move_to_asu(hkl); + row[0] = static_cast(hkl[0]); + row[1] = static_cast(hkl[1]); + row[2] = static_cast(hkl[2]); + row[3] = static_cast((partials ? 256 : 0) + isym); + row[4] = static_cast(1 + static_cast(std::lround(r.image_number))); // BATCH, as above + const float corr = r.prescaling_corr * r.qe_corr * r.flight_corr; + row[5] = r.I * corr; + row[6] = r.sigma * corr; + row[7] = r.partiality; + row[8] = DetectorX(r); + row[9] = DetectorY(r); + row[10] = phi_start_deg(r.image_number) + wedge_deg / 2.0f; + // prescaling_corr is Lorentz x polarization, which is what LP means; qe_corr is the sensor + // efficiency and flight_corr the flight path beside it, each written in the divide-by + // direction, so that I / LP * QE * FLIGHT is the raw count. FLIGHT is <= 1 where QE is >= 1: + // the sensor makes an oblique reflection read high and the medium makes it read low. + row[11] = r.prescaling_corr; + row[12] = 1.0f / r.qe_corr; + row[13] = 1.0f / r.flight_corr; + row[14] = 0.0f; // FLAG: nothing here is a rejected observation + row[15] = r.delta_phi_deg; + row[16] = r.zeta; + row[17] = r.bkg; + row[18] = r.var_bkg; + } + }); mtz.nreflections = static_cast(mtz.data.size() / mtz.columns.size()); // The batch header's orientation matrix is the crystal at rotation angle zero - each batch's own @@ -955,10 +991,9 @@ gemmi::Mtz UnmergedMtz(const std::vector &outcomes, if (!mtz.has_data()) mtz.sort(5); // no rows: its own "No data." failure, as before { - const size_t w = mtz.columns.size(); - std::vector rows(mtz.nreflections); - std::iota(rows.begin(), rows.end(), 0); - ParallelSort(rows.begin(), rows.end(), nthreads, [&](int i, int j) { + std::vector order(mtz.nreflections); + std::iota(order.begin(), order.end(), 0); + ParallelSort(order.begin(), order.end(), nthreads, [&](int i, int j) { const float *a = &mtz.data[static_cast(i) * w], *b = &mtz.data[static_cast(j) * w]; for (int n = 0; n < 5; ++n) if (a[n] != b[n]) @@ -966,8 +1001,11 @@ gemmi::Mtz UnmergedMtz(const std::vector &outcomes, return i < j; }); std::vector sorted(mtz.data.size()); - for (size_t i = 0; i < rows.size(); ++i) - std::memcpy(&sorted[i * w], &mtz.data[static_cast(rows[i]) * w], w * sizeof(float)); + ParallelChunks(static_cast(order.size()), ThreadsForWork(order.size(), nthreads), [&](int lo, int hi) { + for (int i = lo; i < hi; ++i) + std::memcpy(&sorted[static_cast(i) * w], &mtz.data[static_cast(order[i]) * w], + w * sizeof(float)); + }); mtz.data.swap(sorted); mtz.sort_order = {{1, 2, 3, 4, 5}}; } diff --git a/tests/WriteReflectionsTest.cpp b/tests/WriteReflectionsTest.cpp index b8e6d06c4..24a532447 100644 --- a/tests/WriteReflectionsTest.cpp +++ b/tests/WriteReflectionsTest.cpp @@ -5,6 +5,7 @@ #include #include +#include #include @@ -137,3 +138,50 @@ TEST_CASE("Unmerged MTZ: LP is Lorentz-polarization, QE the sensor efficiency, F CHECK(I == Catch::Approx(raw[i] * lp[i] * qe[i] * fl[i]).epsilon(1e-5)); } } + +TEST_CASE("Unmerged MTZ: built on several workers, the same rows as on one", "[write_reflections][portable]") { + // A sweep large enough that the sort, the event sums and the row fill are all split into + // pieces, with events of several parts, events cut short by the sweep's ends, and parts that + // share (h,k,l) and image_number - the tie the sort has to break the same way every time. + auto experiment = TestExperiment(); + experiment.Goniometer(GoniometerAxis("omega", 0.0f, 0.1f, Coord(-1, 0, 0), {})); + + uint32_t seed = 12345; + const auto next = [&seed] { seed = seed * 1664525u + 1013904223u; return (seed >> 8) / 16777216.0f; }; + + std::vector outcomes(200); + for (int f = 0; f < 200; ++f) { + for (int q = 0; q < 8000; ++q) { + if ((f + q) % 7 >= 3) + continue; // each reflection is seen on three consecutive images out of seven + Reflection r{}; + r.h = q % 20 - 10; r.k = (q / 20) % 20 - 10; r.l = q / 400 + 1; + r.image_number = static_cast(f); + r.I = 1000.0f * next(); + r.sigma = 1.0f + 10.0f * next(); + r.bkg = next(); + r.var_bkg = next(); + r.prescaling_corr = 1.0f + next(); + r.partiality = 0.2f + 0.2f * next(); + r.predicted_x = 1000.0f * next(); r.predicted_y = 1000.0f * next(); + r.observed_x = NAN; r.observed_y = NAN; + outcomes[f].reflections.push_back(r); + if (q % 97 == 0) { + r.I = 1000.0f * next(); + outcomes[f].reflections.push_back(r); + } + } + } + + for (bool sum_partials : {true, false}) { + INFO("sum_partials " << sum_partials); + const gemmi::Mtz serial = UnmergedMtz(outcomes, TETRAGONAL_CELL, experiment, sum_partials, 1); + const gemmi::Mtz parallel = UnmergedMtz(outcomes, TETRAGONAL_CELL, experiment, sum_partials, 8); + REQUIRE(serial.nreflections > 100000); + REQUIRE(parallel.nreflections == serial.nreflections); + CHECK(std::memcmp(parallel.data.data(), serial.data.data(), serial.data.size() * sizeof(float)) == 0); + REQUIRE(parallel.batches.size() == serial.batches.size()); + CHECK(parallel.batches.front().number == serial.batches.front().number); + CHECK(parallel.batches.back().number == serial.batches.back().number); + } +}