diff --git a/docs/ACKNOWLEDGEMENT.md b/docs/ACKNOWLEDGEMENT.md index ede0116af..8f41d08a4 100644 --- a/docs/ACKNOWLEDGEMENT.md +++ b/docs/ACKNOWLEDGEMENT.md @@ -384,6 +384,15 @@ J. Beilsten-Edmands, G. Winter, R. Gildea et al., "Scaling diffraction data in t package: algorithms and new approaches for multi-crystal scaling" (2020), Acta Cryst. D76, 385-399 [doi:10.1107/S2059798320003198](https://doi.org/10.1107/S2059798320003198). +**Several sweeps merged as one** - each sweep's partials combined into fulls on their own, and the +fulls of all the sweeps then scaled together, every sweep with its own frame scales and a relative B +against the joint merge, is how XSCALE merges several data sets (W. Kabsch, "Integration, scaling, +space-group assignment and post-refinement" (2010), above). Reading the radiation dose a sweep was taken +at from its relative B, which orders the sweeps by dose from the data rather than from the order they +were given in, follows K. Diederichs, "Some aspects of quantitative analysis and correction of radiation +damage" (2006), Acta Cryst. D62, 96-101 +[doi:10.1107/S0907444905031537](https://doi.org/10.1107/S0907444905031537). + **Uncertainty conventions** follow the IUCr Commission on Crystallographic Nomenclature: D. Schwarzenbach, S. C. Abrahams, H. D. Flack et al., "Statistical descriptors in crystallography: Report of the IUCr Subcommittee on Statistical Descriptors" (1989), Acta Cryst. A45, 63-75 diff --git a/image_analysis/scale_merge/RotationScaleMerge.cpp b/image_analysis/scale_merge/RotationScaleMerge.cpp index c3071337e..694ca83d4 100644 --- a/image_analysis/scale_merge/RotationScaleMerge.cpp +++ b/image_analysis/scale_merge/RotationScaleMerge.cpp @@ -531,11 +531,28 @@ void RotationScaleMerge::IngestHost() { if (frame_limit >= 0) n_frames = std::min(n_frames, frame_limit); partials_released = false; resident_ingest = false; + frame_segment.clear(); + if (!segments.empty()) { + frame_segment.assign(n_frames, -1); + int32_t next = 0; + for (size_t k = 0; k < segments.size(); ++k) { + const SweepSegment &seg = segments[k]; + if (seg.first_frame != next || seg.n_frames <= 0 || seg.first_frame + seg.n_frames > n_frames) + throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, + "RotationScaleMerge: the sweep segments do not cover the frames in order"); + for (int32_t f = seg.first_frame; f < seg.first_frame + seg.n_frames; ++f) + frame_segment[f] = static_cast(k); + next = seg.first_frame + seg.n_frames; + } + if (next != n_frames) + throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, + "RotationScaleMerge: the sweep segments do not cover the frames in order"); + } #ifdef JFJOCH_USE_CUDA // Probed before the observations are built: whether the build lays down the full Obs array or // only the ingest arrays depends on whether the device pipeline will be active. gpu_ = std::make_unique(gpu_device); - gpu_active_ = gpu_->Available() && !scale_on_cpu_; + gpu_active_ = gpu_->Available() && !scale_on_cpu_ && segments.empty(); resident_ingest = gpu_active_ && !observation_dump; #endif @@ -605,7 +622,12 @@ void RotationScaleMerge::IngestHost() { } }); - MeasureIncidentFlux(mean_bkg); + // Fulls of several sweeps carry each sweep's own background: not one flux meter, so the per-frame + // scale alone takes the flux. + if (segments.empty()) + MeasureIncidentFlux(mean_bkg); + else + frame_flux.assign(n_frames, 1.0); // The occupied range of h, which is the buckets the sort works on. int b_lo = H_ALL, b_hi = -1; @@ -728,12 +750,19 @@ void RotationScaleMerge::IngestHost() { } }); } - BuildInRangeObservations(keys, static_cast(total), SmoothedFrameGeometry()); + // Pre-combined fulls have no rocking geometry left to smooth or to recompute a partiality from. + BuildInRangeObservations(keys, static_cast(total), + segments.empty() ? SmoothedFrameGeometry() : SmoothedGeometry{}); rawrun_group.assign(rawrun_start.size(), -1); logger->Info("RotationScaleMerge: ingested {} partial observations from {} frames ({} distinct hkl)", n_partials_obs, n_frames, rawrun_start.size()); - SmoothMosaicityAndPartiality(); + if (segments.empty()) { + SmoothMosaicityAndPartiality(); + } else { + ingest_zeta.reset(); + ingest_d.reset(); + } ParallelChunks(static_cast(partials.size()), ThreadsForWork(partials.size(), nthreads), [&](int lo, int hi) { for (int i = lo; i < hi; ++i) @@ -2174,6 +2203,7 @@ RotationScaleMerge::ScalingLoopOutcome RotationScaleMerge::RunFullsScalingLoop( has_fulls[f] = frame_obs_count[f] > 0 ? 1 : 0; } double best_step = INFINITY, lambda = 0.0; + std::vector segment_lambda(segments.size(), 0.0); bool lambda_kept = false; std::vector steps, g_best; AndersonMixing anderson(SCALING_ANDERSON_DEPTH); @@ -2188,9 +2218,16 @@ RotationScaleMerge::ScalingLoopOutcome RotationScaleMerge::RunFullsScalingLoop( if (fitted) y[f] += g_fit[f] > 0.0 ? std::max(std::log(g_fit[f] / g[f]), -1.0) : -1.0; } - if (!lambda_kept) - lambda = ChooseLogScaleSmoothing(y, J, std::max(1, rocking_event_frames_at_start), nthreads); - const std::vector x = SmoothLogScale(y, J, lambda); + std::vector x; + if (segments.empty()) { + if (!lambda_kept) + lambda = ChooseLogScaleSmoothing(y, J, std::max(1, rocking_event_frames_at_start), nthreads); + x = SmoothLogScale(y, J, lambda); + } else { + x = SmoothLogScaleBySegment(y, J, segment_lambda, !lambda_kept); + std::vector l = segment_lambda; + lambda = median_of(l); + } for (int f = 0; f < n_frames; ++f) g_new[f] = std::exp(x[f]); const double g_ref = TypicalFrameScale(g_new, n_obs); @@ -2342,8 +2379,14 @@ void RotationScaleMerge::RefineDecay(int n_groups) { return o.group >= 0 && o.corr > 0.0f && std::isfinite(o.corr) && (o.d > 0.0f) && o.partiality >= min_partiality; }; + // Where an observation sits on the dose axis: its frame, or with several sweeps its place in the dose + // order the data put them in (MeasureSegmentDose). + auto dose = [&](const auto &o) { + return segments.empty() ? static_cast(o.image_number) + : frame_dose[o.frame] + (o.image_number - static_cast(o.frame)); + }; auto decay_factor = [&](const auto &o, double slope) { - return std::exp(slope * (o.image_number - fcenter) * s2_of(o.d)); + return std::exp(slope * (dose(o) - fcenter) * s2_of(o.d)); }; // The fields the seven passes below read (three slope fits, four held-out scores), pulled out of the @@ -2358,6 +2401,8 @@ void RotationScaleMerge::RefineDecay(int n_groups) { BuildDecayTerms(n_groups, t); decay_terms_kept = {}; decay_terms_kept_valid = false; + if (!segments.empty()) + MeasureSegmentDose(t, n_groups); const int nt = static_cast(ThreadsForWork(t.term.size(), nthreads)); std::vector sw(n_groups), swI(n_groups); @@ -2385,7 +2430,7 @@ void RotationScaleMerge::RefineDecay(int n_groups) { const double Is = static_cast(o.I) * o.corr, sc = static_cast(o.sigma) * o.corr; if (!std::isfinite(Iref) || Iref <= 0.0 || !(Is > 0.0)) return; const double w = (Is / sc) * (Is / sc); - const double x = (o.image_number - fcenter) * s2_of(o.d), y = std::log(Iref / Is); + const double x = (dose(o) - fcenter) * s2_of(o.d), y = std::log(Iref / Is); a[0] += w; a[1] += w * x; a[2] += w * y; a[3] += w * x * x; a[4] += w * x * y; }, [](Lin &a, const Lin &b) { for (int j = 0; j < 5; ++j) a[j] += b[j]; }); const double Sw = S[0], Sx = S[1], Sy = S[2], Sxx = S[3], Sxy = S[4]; @@ -2460,6 +2505,81 @@ void RotationScaleMerge::RefineDecay(int n_groups) { 0.5 * slope * n_frames, 0.5 * slope); } +void RotationScaleMerge::MeasureSegmentDose(const DecayTerms &t, int n_groups) { + // A relative B per data set against the joint merge, as XSCALE fits it: Kabsch (2010) Acta Cryst. + // D66, 133-144; the dose read from it: Diederichs (2006) Acta Cryst. D62, 96-101. + // Radiation damage raises the relative B of a sweep collected later, so the B factors of the + // segments against their joint merge give the order the crystal took the dose in - read from the + // data, not from the order the sweeps were given in or the times their headers state. Per segment, + // the inverse-variance-of-log weighted slope of ln(I_ref / I_obs) against s^2 = 1/(4 d^2) is 2 B, + // with an intercept of its own (the scale is the per-frame G's). The least damaged segment is B = 0, + // and its frames start the dose axis. + const size_t ns = segments.size(); + std::vector sw(n_groups, 0.0), swI(n_groups, 0.0); + for (const DecayTerm &o : t.term) { + const double sc = static_cast(o.sigma) * o.corr, w = 1.0 / (sc * sc); + sw[o.group] += w; + swI[o.group] += w * static_cast(o.I) * o.corr; + } + std::vector> S(ns, std::array{}); // Sw, Sx, Sy, Sxx, Sxy + for (const DecayTerm &o : t.term) { + if (sw[o.group] <= 0.0) continue; + const double Iref = swI[o.group] / sw[o.group]; + const double Is = static_cast(o.I) * o.corr, sc = static_cast(o.sigma) * o.corr; + if (!std::isfinite(Iref) || Iref <= 0.0 || !(Is > 0.0)) continue; + const double w = (Is / sc) * (Is / sc); + const double x = 1.0 / (4.0 * static_cast(o.d) * o.d), y = std::log(Iref / Is); + auto &a = S[frame_segment[o.frame]]; + a[0] += w; a[1] += w * x; a[2] += w * y; a[3] += w * x * x; a[4] += w * x * y; + } + std::vector b(ns, NAN); + for (size_t k = 0; k < ns; ++k) { + const auto &a = S[k]; + const double var = a[0] > 0.0 ? a[3] - a[1] * a[1] / a[0] : 0.0; + if (var > 0.0) + b[k] = 0.5 * (a[4] - a[1] * a[2] / a[0]) / var; + } + // A segment with nothing to measure keeps its place in the input order, after the measured ones. + std::vector order(ns); + std::iota(order.begin(), order.end(), 0); + std::stable_sort(order.begin(), order.end(), [&](int p, int q) { + if (std::isfinite(b[p]) != std::isfinite(b[q])) return std::isfinite(b[p]); + return std::isfinite(b[p]) && b[p] < b[q]; + }); + const double b0 = std::isfinite(b[order[0]]) ? b[order[0]] : 0.0; + frame_dose.assign(n_frames, 0.0); + double offset = 0.0; + std::string text; + for (size_t rank = 0; rank < ns; ++rank) { + const int k = order[rank]; + const SweepSegment &seg = segments[k]; + for (int32_t f = seg.first_frame; f < seg.first_frame + seg.n_frames; ++f) + frame_dose[f] = offset + (f - seg.first_frame); + offset += seg.n_frames; + segment_results[k].relative_b = b[k] - b0; + segment_results[k].dose_rank = static_cast(rank); + text += fmt::format("{}segment {} (B {:+.2f} A^2)", rank == 0 ? "" : ", ", k + 1, b[k] - b0); + } + logger->Info("Dose order of the {} sweep segments, from their relative B against the joint merge: {}", + ns, text); +} + +std::vector RotationScaleMerge::SmoothLogScaleBySegment(const std::vector &y, + const std::vector &J, + std::vector &lambda, bool choose) const { + std::vector x(y.size()); + for (size_t k = 0; k < segments.size(); ++k) { + const auto lo = y.begin() + segments[k].first_frame, hi = lo + segments[k].n_frames; + const auto jlo = J.begin() + segments[k].first_frame; + const std::vector ys(lo, hi), Js(jlo, jlo + segments[k].n_frames); + if (choose) + lambda[k] = ChooseLogScaleSmoothing(ys, Js, std::max(1, rocking_event_frames_at_start), nthreads); + const std::vector xs = SmoothLogScale(ys, Js, lambda[k]); + std::copy(xs.begin(), xs.end(), x.begin() + segments[k].first_frame); + } + return x; +} + // A^2, the bound the Gauss-Seidel solve clamps each batch's relative-B to. It is a numerical guard against // divergence, never an answer: a run of identical values in the solved curve means the solve reached the // bound, not that the crystal did anything. The CORRECTION keeps the tight bound it was tuned with; the @@ -4716,12 +4836,15 @@ void RotationScaleMerge::Combine() { size_t i = 0; while (i < ev.size()) { size_t kk = i + 1; - float last_frame = partials[ev[i]].r->image_number; - while (kk < ev.size()) { - const float frame = partials[ev[kk]].r->image_number; - if (frame - last_frame > max_frame_gap) break; - last_frame = frame; - ++kk; + // A pre-combined full is a whole event already (SweepSegment). + if (segments.empty()) { + float last_frame = partials[ev[i]].r->image_number; + while (kk < ev.size()) { + const float frame = partials[ev[kk]].r->image_number; + if (frame - last_frame > max_frame_gap) break; + last_frame = frame; + ++kk; + } } double pooled_bkg = 0.0; @@ -6692,6 +6815,12 @@ RotationScaleMerge::Result RotationScaleMerge::Run(bool for_search, bool full_st const gemmi::SpaceGroup *group) { const ScopedGpuPin pin(gpu_device); // see IngestHost run_group = group; + // Pre-combined fulls: there are no partials to fit a scale on, and every event is one frame. + if (!segments.empty()) { + scale_partials_choice = false; + rocking_event_frames_at_start = 1; + rocking_events_at_start = n_partials_obs; + } struct ClearGroup { const gemmi::SpaceGroup *&g; ~ClearGroup() { g = nullptr; } } clear_group{run_group}; // Whether a frame's scale can be fitted on its partials depends on whether the partials tell a change // of scale from an error of the partiality model. The count of rocking events per frame is the prior @@ -7042,7 +7171,7 @@ RotationScaleMerge::Result RotationScaleMerge::RunArm(bool scale_partials, bool // kernels on it, and the restore re-uploads. // --- 2a. On the de-novo search pass only, drop observations whose Lorentz geometry is poor. --- - if (for_search && search_min_zeta > 0.0) { + if (for_search && search_min_zeta > 0.0 && segments.empty()) { // a full has no rocking geometry int64_t n_dropped = 0; bool zeta_on_gpu = false; #ifdef JFJOCH_USE_CUDA @@ -7352,20 +7481,26 @@ RotationScaleMerge::Result RotationScaleMerge::RunArm(bool scale_partials, bool // is a correction to the merged INTENSITIES, and a pass whose intensities are thrown away has // nowhere to put one - the fitted gain lives on the fulls, which the next pass rebuilds. --- const bool refine_surfaces = !for_search && full_stats; + // What runs along the frames of ONE sweep - the damage monitor, the sweep ledger, the per-batch B and + // the surfaces, each read in one goniometer and one detector setting - has no meaning across the + // segments of several, and is left out there. The decay is fitted along the dose order instead. + const bool along_frames = segments.empty(); const bool corrections = - refine_surfaces && (refine_decay_b || absorption_iter > 0 || modulation_iter > 0 || relative_b_deg > 0.0); + refine_surfaces && (refine_decay_b || (along_frames && (absorption_iter > 0 || modulation_iter > 0 + || relative_b_deg > 0.0))); + segment_results.assign(segments.size(), Result::Segment{}); // Radiation-damage monitor: measure the per-batch relative-B on the scaled fulls BEFORE any correction // (report-only; captures the full damage signature, not a residual). Skipped on the P1 search pass. std::unique_ptr sub_mark; auto sub = [&](const char *n) { sub_mark.reset(); sub_mark = std::make_unique(*logger, std::string("rsm/s ") + n); }; sub("damage B"); - if (refine_surfaces) + if (refine_surfaces && along_frames) MeasureRadiationDamageB(n_groups); // Sweep-quality diagnostic, on the per-frame scale the partial scaling just fitted (the flux is // already out of it) and the per-frame CC computed above. It drops nothing itself; the ranges it // segments are the ledger 4c below decides on. sub("sweep quality"); - if (refine_surfaces) + if (refine_surfaces && along_frames) MeasureSweepQuality(partial_scaled, cc, cc_n); // CC1/2 of these same fulls with the surfaces not yet folded in, for the caller's cross-pass quality // guard: the pass this one is judged against fits no surfaces of its own, so it is the uncorrected @@ -7391,7 +7526,7 @@ RotationScaleMerge::Result RotationScaleMerge::RunArm(bool scale_partials, bool decay_terms_kept = {}; decay_terms_kept_valid = false; sub("relative B"); - if (refine_surfaces && relative_b_deg > 0.0) + if (refine_surfaces && along_frames && relative_b_deg > 0.0) RefineRelativeB(n_groups); // per-batch relative-B on top of the single decay slope // The three surfaces overlap, so the one fitted first gets the first claim and the others are tested // on what it left - the order decides what is adopted. Detector first: every frame of the sweep @@ -7407,16 +7542,16 @@ RotationScaleMerge::Result RotationScaleMerge::RunArm(bool scale_partials, bool // the grid is then tested on what the smooth surface left, and where that surface is refused the // grid sees exactly what it saw before. sub("modulation"); - if (refine_surfaces && modulation_iter > 0) + if (refine_surfaces && along_frames && modulation_iter > 0) RefineModulation(modulation_iter, n_groups); sub("absorption time"); - if (refine_surfaces && absorption_iter > 0) + if (refine_surfaces && along_frames && absorption_iter > 0) RefineAbsorptionTime(absorption_iter, n_groups); sub("absorption SH"); - if (refine_surfaces && absorption_iter > 0) + if (refine_surfaces && along_frames && absorption_iter > 0) RefineAbsorptionSH(absorption_iter, n_groups); sub("absorption gonio"); - if (refine_surfaces && absorption_iter > 0) + if (refine_surfaces && along_frames && absorption_iter > 0) RefineAbsorption(absorption_iter, n_groups); sub("disposition"); @@ -7457,6 +7592,16 @@ RotationScaleMerge::Result RotationScaleMerge::RunArm(bool scale_partials, bool // its CC1/2 is the answer without a second merge. A caller that did not ask keeps the NaN. if (measure_cc_before_corrections) r.cc_half_before_corrections = cc_half_before_corrections.value_or(r.statistics.overall.cc_half); + // Each segment's scale: the median of its frames' scales, at the run's gauge. + for (size_t k = 0; k < segments.size(); ++k) { + std::vector g; + for (int32_t f = segments[k].first_frame; f < segments[k].first_frame + segments[k].n_frames; ++f) + if (fulls_frame_count[f] > 0 && std::isfinite(g_full[f]) && g_full[f] > 0.0) + g.push_back(g_full[f]); + if (!g.empty()) + segment_results[k].scale = median_of(g); + } + r.segments = segment_results; // The fulls and the merge's working arrays are rebuilt by every arm before they are read, so they go // with it: kept, they sat beside the next arm's combine and through everything the caller did between // merges - on a large cell in P1, the fulls and about 330 bytes per ASU group, gigabytes. Faulting them diff --git a/image_analysis/scale_merge/RotationScaleMerge.h b/image_analysis/scale_merge/RotationScaleMerge.h index 6291bdf1e..6b34a5129 100644 --- a/image_analysis/scale_merge/RotationScaleMerge.h +++ b/image_analysis/scale_merge/RotationScaleMerge.h @@ -80,6 +80,17 @@ std::pair TrimBlankSweepEnds(std::vector &outcomes const std::vector &tally); std::pair TrimBlankSweepEnds(std::vector &outcomes, size_t nthreads); +// One sweep of several whose frames are laid end to end in one engine (the joint merge of several +// sweeps): frames [first_frame, first_frame + n_frames) of the outcomes. Every observation of a segment's +// frames is a full already combined in its own sweep (partiality 1, filed under the frame of its peak), +// so nothing is combined again, nothing is recomputed from the rocking geometry, and nothing that runs +// along the frames - the per-frame scale's smoothness above all - reaches across from one sweep into the +// next. +struct SweepSegment { + int32_t first_frame = 0; + int32_t n_frames = 0; +}; + class RotationScaleMerge { public: // The log scale x minimising sum_f J_f (x_f - y_f)^2 + lambda * sum_f (x_f - 2 x_{f+1} + x_{f+2})^2: @@ -139,6 +150,13 @@ public: // The fulls the merge kept, unmerged (see ScaledFull): what passed the outlier tests, inside the // resolution cut. Filled only when SetExportScaledFulls asked for it, and never on a search merge. std::vector scaled_fulls; + // Per sweep segment (SetSweepSegments), in segment order; empty without segments. + struct Segment { + double scale = NAN; // median per-frame scale of the segment's fulls (run gauge) + double relative_b = NAN; // A^2, against the joint merge, the least damaged segment at 0 + int dose_rank = -1; // the segment's place in the dose order those B factors decide + }; + std::vector segments; }; // experiment: read live (its space group is changed by the caller between Run() calls). @@ -288,6 +306,12 @@ public: // Merge only the first n frames of the outcomes (by default all of them). Set before Ingest(). void SetFrameLimit(int n) { frame_limit = n; } + // The outcomes are several sweeps of pre-combined fulls laid end to end (SweepSegment). Set before + // Ingest(); the segments must cover the frames in order. Such an engine runs on the host: its fulls + // are a few per frame where a sweep's partials are millions, and the device pipeline is the partials'. + // Without segments (the default) nothing of this runs. + void SetSweepSegments(std::vector s) { segments = std::move(s); } + // The GPU the engine runs on (device 0 by default). Set before Ingest(). The device only says where // the merge runs: on cards of one model the result is the same on any of them. void SetGpuDevice(int device) { gpu_device = device; } @@ -724,6 +748,15 @@ private: std::vector last_cc_n; std::vector last_frame_scaled; int frame_limit = -1; // see SetFrameLimit + std::vector segments; // see SetSweepSegments + std::vector frame_segment; // each frame's segment (Ingest; empty without segments) + std::vector segment_results; // the last arm's, copied into its Result + // Each frame's place on the dose axis the decay is fitted along: the frames of the segments laid end + // to end in the order the segments' relative B factors put them (MeasureSegmentDose). + std::vector frame_dose; + // SmoothLogScale / ChooseLogScaleSmoothing within each segment on its own; lambda per segment. + std::vector SmoothLogScaleBySegment(const std::vector &y, const std::vector &J, + std::vector &lambda, bool choose) const; int gpu_device = 0; // see SetGpuDevice bool export_scaled_fulls = false; // see SetExportScaledFulls bool french_wilson = true; // see SetFrenchWilson @@ -958,6 +991,9 @@ private: std::vector g_start; }; void BuildDecayTerms(int n_groups, DecayTerms &t) const; + // Per segment: the relative B against the merged reference, and the dose order it puts the segments + // in (segment_results, frame_dose). + void MeasureSegmentDose(const DecayTerms &t, int n_groups); // The terms the radiation-damage monitor built, handed on to RefineDecay: nothing between the two // touches corr, so they are the terms RefineDecay would build. Empty outside that stretch of RunArm. DecayTerms decay_terms_kept; diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 63599b77e..d2075443c 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -139,6 +139,7 @@ ADD_EXECUTABLE(jfjoch_test PostRefineTest.cpp CellEsdTest.cpp RotationScaleWalkTest.cpp + RotationScaleMergeSegmentsTest.cpp AndersonMixingTest.cpp RugnuxLargeTest.cpp TestData.h diff --git a/tests/RotationScaleMergeSegmentsTest.cpp b/tests/RotationScaleMergeSegmentsTest.cpp new file mode 100644 index 000000000..1b33d0454 --- /dev/null +++ b/tests/RotationScaleMergeSegmentsTest.cpp @@ -0,0 +1,212 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#include + +#include +#include +#include +#include +#include +#include + +#include "../common/CUDAWrapper.h" +#include "../common/Logger.h" +#include "../image_analysis/scale_merge/RotationScaleMerge.h" + +namespace { + +// A synthetic crystal measured as pre-combined fulls on several sweeps: a cubic 20 A cell, Wilson +// intensities at B = 15 A^2 to 1.6 A, each sweep at its own scale and relative B, Poisson-like noise. +struct SweepSpec { + int n_frames; + double scale; + double relative_b; +}; + +struct Synthetic { + std::vector frames; + std::map, double> I_true; // by the Friedel-merged index +}; + +constexpr double CELL = 20.0; + +std::array FriedelKey(int h, int k, int l) { + const std::array a{h, k, l}, b{-h, -k, -l}; + return std::max(a, b); +} + +Synthetic MakeSweeps(const std::vector &sweeps, uint32_t seed) { + std::mt19937 rng(seed); + std::vector> hkl; + for (int h = -12; h <= 12; ++h) + for (int k = -12; k <= 12; ++k) + for (int l = -12; l <= 12; ++l) { + const double d = CELL / std::sqrt(static_cast(h * h + k * k + l * l)); + if ((h || k || l) && d >= 1.6) hkl.push_back({h, k, l}); + } + Synthetic out; + std::exponential_distribution wilson(1.0); + for (const auto &m: hkl) { + const auto key = FriedelKey(m[0], m[1], m[2]); + if (!out.I_true.contains(key)) { + const double s2 = (m[0] * m[0] + m[1] * m[1] + m[2] * m[2]) / (4.0 * CELL * CELL); + out.I_true[key] = 2000.0 * std::exp(-2.0 * 15.0 * s2) * (0.05 + wilson(rng)); + } + } + std::uniform_int_distribution pick(0, hkl.size() - 1); + std::normal_distribution noise(0.0, 1.0); + int frame = 0; + for (const auto &sweep: sweeps) + for (int f = 0; f < sweep.n_frames; ++f, ++frame) { + IntegrationOutcome o; + o.latt = CrystalLattice(UnitCell{static_cast(CELL), static_cast(CELL), + static_cast(CELL), 90, 90, 90}); + for (int j = 0; j < 400; ++j) { + const auto &m = hkl[pick(rng)]; + const double s2 = (m[0] * m[0] + m[1] * m[1] + m[2] * m[2]) / (4.0 * CELL * CELL); + const double expected = out.I_true.at(FriedelKey(m[0], m[1], m[2])) * sweep.scale + * std::exp(-2.0 * sweep.relative_b * s2); + const double sigma = std::sqrt(expected + 20.0); + Reflection r{}; + r.h = m[0]; r.k = m[1]; r.l = m[2]; + r.image_number = static_cast(frame); + r.d = static_cast(1.0 / std::sqrt(4.0 * s2)); + r.I = static_cast(expected + sigma * noise(rng)); + r.sigma = static_cast(sigma); + // The combine rebuilds a full's variance as var_bkg + I: what is not the intensity's own + r.var_bkg = static_cast(sigma * sigma - std::max(0.0, static_cast(r.I))); + r.bkg = 20.0f; + r.prescaling_corr = 1.0f; + r.partiality = 1.0f; + r.zeta = 1.0f; + r.image_scale_corr = 1.0f; + r.predicted_x = 500.0f; + r.predicted_y = 500.0f; + o.reflections.push_back(r); + } + out.frames.push_back(std::move(o)); + } + return out; +} + +DiffractionExperiment SyntheticExperiment(bool correction_surfaces) { + DiffractionExperiment x; + x.IncidentEnergy_keV(12.4f); + x.Goniometer(GoniometerAxis("omega", 0.0f, 0.1f, Coord(1, 0, 0), {})); + x.SetUnitCell(UnitCell{static_cast(CELL), static_cast(CELL), static_cast(CELL), 90, 90, 90}); + ScalingSettings s; + s.ScaleFulls(true).CorrectionSurfaces(correction_surfaces).ResolutionCutoff(ResolutionCutoffMethod::Off); + x.ImportScalingSettings(s); + return x; +} + +std::vector Segments(const std::vector &sweeps) { + std::vector out; + int32_t first = 0; + for (const auto &s: sweeps) { + out.push_back({first, s.n_frames}); + first += s.n_frames; + } + return out; +} + +RotationScaleMerge::Result Merge(Synthetic &data, const DiffractionExperiment &x, + const std::vector &segments, bool export_fulls = false, + bool *on_gpu = nullptr) { + Logger logger("RotationScaleMergeSegmentsTest"); + const UnitCell cell{static_cast(CELL), static_cast(CELL), static_cast(CELL), 90, 90, 90}; + RotationScaleMerge engine(x, data.frames, cell, 20, 4, logger); + if (!segments.empty()) + engine.SetSweepSegments(segments); + engine.SetExportScaledFulls(export_fulls); + engine.Ingest(); + if (on_gpu) + *on_gpu = engine.OnGPU(); + return engine.Run(false, true, false); +} + +// The median over a frame range of the scaled fulls' ratio to the true intensity (the strong ones). +double MedianRatio(const RotationScaleMerge::Result &r, const Synthetic &data, int f0, int f1) { + std::vector ratio; + for (const auto &o: r.scaled_fulls) + if (o.image_number >= static_cast(f0) && o.image_number < static_cast(f1)) { + const double t = data.I_true.at(FriedelKey(o.h, o.k, o.l)); + if (t > 500.0) + ratio.push_back(o.I / t); + } + REQUIRE(ratio.size() > 20); + std::nth_element(ratio.begin(), ratio.begin() + ratio.size() / 2, ratio.end()); + return ratio[ratio.size() / 2]; +} + +} // namespace + +TEST_CASE("RotationScaleMergeSegments_PerSweepScaleRecovered", "[RotationScale]") { + const std::vector sweeps{{60, 1.0, 0.0}, {60, 0.5, 0.0}}; + auto data = MakeSweeps(sweeps, 11); + const auto r = Merge(data, SyntheticExperiment(false), Segments(sweeps), true); + REQUIRE(r.segments.size() == 2); + CHECK(r.segments[1].scale / r.segments[0].scale == Catch::Approx(0.5).epsilon(0.03)); + CHECK(r.statistics.overall.cc_half > 0.99); + + // No coupling across the boundary: the frames on either side of it are put on the common scale as + // well as the middle of their sweeps are. + const double middle = MedianRatio(r, data, 20, 40); + CHECK(MedianRatio(r, data, 57, 60) / middle == Catch::Approx(1.0).epsilon(0.03)); + CHECK(MedianRatio(r, data, 60, 63) / middle == Catch::Approx(1.0).epsilon(0.03)); + CHECK(MedianRatio(r, data, 80, 100) / middle == Catch::Approx(1.0).epsilon(0.03)); +} + +TEST_CASE("RotationScaleMergeSegments_DoseOrderFromTheRelativeB", "[RotationScale]") { + // Given first, the sweep taken second: its B is 8 A^2 higher. The order comes from the data. + const std::vector sweeps{{60, 1.0, 8.0}, {60, 1.0, 0.0}, {60, 1.0, 4.0}}; + auto data = MakeSweeps(sweeps, 12); + const auto r = Merge(data, SyntheticExperiment(true), Segments(sweeps)); + REQUIRE(r.segments.size() == 3); + CHECK(r.segments[1].dose_rank == 0); + CHECK(r.segments[2].dose_rank == 1); + CHECK(r.segments[0].dose_rank == 2); + CHECK(r.segments[1].relative_b == Catch::Approx(0.0).margin(1e-9)); + CHECK(r.segments[0].relative_b == Catch::Approx(8.0).margin(1.0)); + CHECK(r.segments[2].relative_b == Catch::Approx(4.0).margin(1.0)); +} + +TEST_CASE("RotationScaleMergeSegments_SegmentsMustCoverTheFrames", "[RotationScale]") { + const std::vector sweeps{{10, 1.0, 0.0}, {10, 1.0, 0.0}}; + auto data = MakeSweeps(sweeps, 13); + const auto x = SyntheticExperiment(false); + CHECK_THROWS_AS(Merge(data, x, {{0, 10}, {11, 9}}), JFJochException); + CHECK_THROWS_AS(Merge(data, x, {{0, 10}}), JFJochException); +} + +TEST_CASE("RotationScaleMergeSegments_NoSegmentsIsTheEngineAsBefore", "[RotationScale]") { + // The same synthetic sweep merged without segments gives, to the bit, what the engine gave before + // segments existed: FNV-1a over the merged I and sigma, recorded on the engine of rc181 e56b19abb - + // on the CPU, and on the GPU for the card model it was recorded on (the device path matches the host + // path to the bit only where every reduction does; this one carries its own, so it is checked against + // its own record). + const std::vector sweeps{{40, 1.0, 0.0}, {40, 0.7, 3.0}}; + auto data = MakeSweeps(sweeps, 14); + bool on_gpu = false; + const auto r = Merge(data, SyntheticExperiment(true), {}, false, &on_gpu); + uint64_t hash = 1469598103934665603ULL; + auto mix = [&hash](float v) { + uint32_t u; + std::memcpy(&u, &v, sizeof u); + for (int b = 0; b < 4; ++b) { + hash ^= (u >> (8 * b)) & 0xffu; + hash *= 1099511628211ULL; + } + }; + for (const auto &m: r.merged) { + mix(m.I); + mix(m.sigma); + } + CHECK(r.segments.empty()); + CHECK(r.merged.size() == 4106); + if (!on_gpu) + CHECK(hash == 3376995616949734726ULL); + else if (get_gpu_description().find("RTX 5080") != std::string::npos) + CHECK(hash == 6594155396455850264ULL); +}