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); + } +}