From 34d5bb563138055c31c3cd9ceb482f4871eb3ab5 Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Wed, 14 Oct 2026 10:09:23 +0200 Subject: [PATCH 1/3] stills merge 2.3-2.5x faster, output unchanged: parallel observation passes, hash accumulator, no amplitudes on scaling references The stills scale/merge after the image loop was almost entirely serial: - every scaling-reference merge (MergeAll: the per-image scaling reference and the two partiality post-refinement references) ran French-Wilson single-threaded on its export, about 3 s each, and assigned free-R flags - neither is read by a reference, which uses I only. MergeAll now exports intensities alone (MergeOnTheFly::ExportIntensities). - the merge accumulated into a std::map; it is now an unordered_map exported in key order, which is the std::map order. - AddImage, the error-model collection and the R_meas pass of MergeStats did the per-observation work (systematic-absence test, ASU key, filters, error-model sigma, outlier test, merged-I lookup) in one serial loop. They now reduce a block of images in parallel and consume it in image order (ReduceImagesInOrder), so every floating-point sum, half-set choice and hash-map insertion is made in the serial order. Measured on a 25k-frame XFEL stills run (box loaded, load 9-14): merge P1 19.0 -> 7.1 s, final-group merge 27.7 -> 13.2 s; on a second run 23.4 -> 7.7 s and 29.7 -> 13.6 s. mtz/unmerged mtz data, hkl, cif and report identical on both runs, mx and --mode scale. New test MergeOnTheFly_AddImagesIsTheSerialMerge checks the parallel merge against the serial loop across a block boundary. Co-Authored-By: Claude Opus 5.5 (1M context) --- image_analysis/scale_merge/Merge.cpp | 253 ++++++++++++------ image_analysis/scale_merge/Merge.h | 28 +- .../scale_merge/StillsPartialityRefine.cpp | 2 +- rugnux/RugnuxScaleMerge.cpp | 11 +- rugnux/rugnux_cli.cpp | 6 +- tests/MergeScaleTest.cpp | 65 +++++ 6 files changed, 269 insertions(+), 96 deletions(-) diff --git a/image_analysis/scale_merge/Merge.cpp b/image_analysis/scale_merge/Merge.cpp index 54ec8d9eb..2c3bf37f1 100644 --- a/image_analysis/scale_merge/Merge.cpp +++ b/image_analysis/scale_merge/Merge.cpp @@ -15,6 +15,7 @@ #include "../../common/CorrelationCoefficient.h" #include "../../common/ResolutionShells.h" #include "../../common/Definitions.h" +#include "../../common/ParallelFor.h" #include "HKLKey.h" #include "RfreeFlags.h" #include "FrenchWilson.h" @@ -32,6 +33,27 @@ namespace { z = z ^ (z >> 31); return static_cast(z & 1ULL); } + + // Calls reduce(i, out) for the images in parallel, a block of them at a time, and then consume(i, out) + // for each image of the block in image order - the order that every floating-point sum, every + // half-set choice and every hash-map insertion of the merge depends on. The result is therefore the + // serial loop's, bit for bit, at any thread count. + template + void ReduceImagesInOrder(size_t n_images, size_t nthreads, Reduce reduce, Consume consume) { + constexpr size_t BLOCK_IMAGES = 4096; + std::vector> out(std::min(BLOCK_IMAGES, n_images)); + for (size_t first = 0; first < n_images; first += BLOCK_IMAGES) { + const int n = static_cast(std::min(BLOCK_IMAGES, n_images - first)); + ParallelChunks(n, nthreads, [&](int lo, int hi) { + for (int j = lo; j < hi; ++j) { + out[j].clear(); + reduce(first + j, out[j]); + } + }); + for (int j = 0; j < n; ++j) + consume(first + j, out[j]); + } + } } MergeOnTheFly::MergeOnTheFly(const DiffractionExperiment &x) @@ -57,14 +79,8 @@ MergeOnTheFly &MergeOnTheFly::ReferenceCell(const std::optional &cell) return *this; } -void MergeOnTheFly::AddImage(const IntegrationOutcome &outcome, int64_t image_id) { - std::unique_lock ul(merged_mutex); - - if (Mask(outcome)) - return; - - const int prefer = HalfForImage(image_id); - +size_t MergeOnTheFly::Contributions(const IntegrationOutcome &outcome, std::vector &out) const { + size_t rejected = 0; for (const auto &r: outcome.reflections) { if (generator.IsSystematicallyAbsent(r)) continue; @@ -94,24 +110,36 @@ void MergeOnTheFly::AddImage(const IntegrationOutcome &outcome, int64_t image_id const auto mit = reject_median_I.find(hkl_key); if (mit != reject_median_I.end() && std::fabs(I_corr - mit->second) > reject_nsigma * sigma_corr) { - ++reject_count; + ++rejected; continue; } } - auto it = accumulator.find(hkl_key); - if (it == accumulator.end()) - it = accumulator.emplace(hkl_key, MergeAccum{ - .h = hkl.plus ? hkl.h : -hkl.h, - .k = hkl.plus ? hkl.k : -hkl.k, - .l = hkl.plus ? hkl.l : -hkl.l, - }).first; - const float w = 1.0f / (sigma_corr * sigma_corr); - const float wI = w * I_corr; + out.push_back(Contribution{ + .key = hkl_key, + .h = hkl.plus ? hkl.h : -hkl.h, + .k = hkl.plus ? hkl.k : -hkl.k, + .l = hkl.plus ? hkl.l : -hkl.l, + .w = w, + .wI = w * I_corr, + .d = r.d, + .on_ice_ring = r.on_ice_ring + }); + } + return rejected; +} - it->second.sum_wI += wI; - it->second.sum_w += w; +void MergeOnTheFly::Accumulate(const std::vector &contributions, int64_t image_id) { + const int prefer = HalfForImage(image_id); + + for (const auto &c: contributions) { + auto it = accumulator.find(c.key); + if (it == accumulator.end()) + it = accumulator.emplace(c.key, MergeAccum{.h = c.h, .k = c.k, .l = c.l}).first; + + it->second.sum_wI += c.wI; + it->second.sum_w += c.w; // CC1/2 correlates the two half-set means, so a reflection whose observations all land in one // half has no second mean and drops out of the statistic entirely. Taking the image's own @@ -119,18 +147,44 @@ void MergeOnTheFly::AddImage(const IntegrationOutcome &outcome, int64_t image_id // reflection floor(n/2) against ceil(n/2), which is the split cctbx and XDS use. const int half = it->second.n_half[prefer] > it->second.n_half[1 - prefer] ? 1 - prefer : prefer; - it->second.sum_wI_half[half] += wI; - it->second.sum_w_half[half] += w; + it->second.sum_wI_half[half] += c.wI; + it->second.sum_w_half[half] += c.w; it->second.n_half[half]++; - if (r.on_ice_ring) + if (c.on_ice_ring) it->second.on_ice_ring = true; - if (!std::isfinite(it->second.d) && std::isfinite(r.d) && r.d > 0.0f) - it->second.d = r.d; + if (!std::isfinite(it->second.d) && std::isfinite(c.d) && c.d > 0.0f) + it->second.d = c.d; } } +void MergeOnTheFly::AddImage(const IntegrationOutcome &outcome, int64_t image_id) { + std::unique_lock ul(merged_mutex); + + if (Mask(outcome)) + return; + std::vector contributions; + reject_count += Contributions(outcome, contributions); + Accumulate(contributions, image_id); +} + +void MergeOnTheFly::AddImages(const std::vector &outcomes) { + std::unique_lock ul(merged_mutex); + + std::vector rejected(outcomes.size(), 0); + ReduceImagesInOrder( + outcomes.size(), nthreads, + [&](size_t i, std::vector &out) { + if (!Mask(outcomes[i])) + rejected[i] = Contributions(outcomes[i], out); + }, + [&](size_t i, const std::vector &out) { + reject_count += rejected[i]; + Accumulate(out, static_cast(i)); + }); +} + float MergeOnTheFly::CorrectedSigma(float I_corr, float sigma_corr, float image_scale_corr, float var_bkg, uint64_t hkl_key) const { @@ -188,29 +242,36 @@ void MergeOnTheFly::RefineErrorModel(const std::vector &outc // --- 1. Collect accepted, scaled observations grouped by symmetry-equivalent hkl, // applying exactly the filters AddImage uses. --- struct Obs { float I, sigma; }; + struct KeyedObs { uint64_t key; Obs obs; }; std::unordered_map> groups; - for (const auto &outcome: outcomes) { - if (Mask(outcome)) - continue; - for (const auto &r: outcome.reflections) { - if (generator.IsSystematicallyAbsent(r)) - continue; - if (r.image_scale_corr <= 0.0 || !std::isfinite(r.image_scale_corr)) - continue; - if (!AcceptReflection(r, high_resolution_limit, low_resolution_limit)) - continue; - if (exclude_ice_rings && r.on_ice_ring) - continue; - if (r.partiality < min_partiality) - continue; - const float I_corr = r.I * r.image_scale_corr; - const float sigma_corr = r.sigma * r.image_scale_corr; - if (!std::isfinite(I_corr) || !std::isfinite(sigma_corr) || sigma_corr <= 0.0f) - continue; - groups[generator(r).pack()].push_back({I_corr, sigma_corr}); - } - } + ReduceImagesInOrder( + outcomes.size(), nthreads, + [&](size_t i, std::vector &out) { + if (Mask(outcomes[i])) + return; + for (const auto &r: outcomes[i].reflections) { + if (generator.IsSystematicallyAbsent(r)) + continue; + if (r.image_scale_corr <= 0.0 || !std::isfinite(r.image_scale_corr)) + continue; + if (!AcceptReflection(r, high_resolution_limit, low_resolution_limit)) + continue; + if (exclude_ice_rings && r.on_ice_ring) + continue; + if (r.partiality < min_partiality) + continue; + const float I_corr = r.I * r.image_scale_corr; + const float sigma_corr = r.sigma * r.image_scale_corr; + if (!std::isfinite(I_corr) || !std::isfinite(sigma_corr) || sigma_corr <= 0.0f) + continue; + out.push_back({generator(r).pack(), {I_corr, sigma_corr}}); + } + }, + [&](size_t, const std::vector &out) { + for (const auto &o: out) + groups[o.key].push_back(o.obs); + }); // --- 2. One global pool of (sigma^2, ^2, bias-corrected squared deviation). For an // observation in a group of n, the residual from the inverse-variance mean has @@ -337,7 +398,7 @@ void MergeOnTheFly::RefineErrorModel(const std::vector &outc error_model_chi2 = chi2.empty() ? 0.0 : median(chi2) / CHI2_1_MEDIAN; } -bool MergeOnTheFly::Mask(const IntegrationOutcome &outcome) { +bool MergeOnTheFly::Mask(const IntegrationOutcome &outcome) const { if (reference_cell) { auto cell = outcome.latt.GetUnitCell(); if (!cell.is_close(*reference_cell, @@ -355,12 +416,20 @@ bool MergeOnTheFly::Mask(const IntegrationOutcome &outcome) { return false; } -std::vector MergeOnTheFly::ExportReflections() { +std::vector MergeOnTheFly::ExportIntensities() { std::unique_lock ul(merged_mutex); std::vector out; + std::vector> by_key; + by_key.reserve(accumulator.size()); + for (const auto &[key, accum]: accumulator) + by_key.emplace_back(key, &accum); + std::sort(by_key.begin(), by_key.end(), + [](const auto &a, const auto &b) { return a.first < b.first; }); + out.reserve(accumulator.size()); - for (const auto &accum: accumulator | std::views::values) { + for (const auto &[key, accum_ptr]: by_key) { + const MergeAccum &accum = *accum_ptr; if (accum.sum_w <= 0.0) continue; @@ -388,18 +457,22 @@ std::vector MergeOnTheFly::ExportReflections() { out.emplace_back(mr); } + return out; +} +std::vector MergeOnTheFly::ExportReflections() { + auto out = ExportIntensities(); AssignRfreeFlags(out, space_group, scaling_settings.GetRfreeFraction(), 500, reference_cell); ApplyFrenchWilson(out, space_group); return out; } std::vector MergeAll(const DiffractionExperiment &x, - const std::vector &integration_outcome) { + const std::vector &integration_outcome, + size_t nthreads) { MergeOnTheFly merge(x); - for (size_t i = 0; i < integration_outcome.size(); ++i) - merge.AddImage(integration_outcome[i], static_cast(i)); - return merge.ExportReflections(); + merge.Threads(nthreads).AddImages(integration_outcome); + return merge.ExportIntensities(); } struct ShellAccum { @@ -630,43 +703,51 @@ MergeStatistics MergeOnTheFly::MergeStats(const std::vector &m std::unordered_map rmeas_obs; rmeas_obs.reserve(merged.size()); - for (int i = 0; i < integration_outcome.size(); ++i) { - if (Mask(integration_outcome[i])) - continue; + struct ShellObs { uint64_t key; float I_corr; float merged_I; int shell; bool merged; }; + ReduceImagesInOrder( + integration_outcome.size(), nthreads, + [&](size_t i, std::vector &out) { + if (Mask(integration_outcome[i])) + return; + for (const auto &r: integration_outcome[i].reflections) { + if (generator.IsSystematicallyAbsent(r)) + continue; + if (r.image_scale_corr <= 0.0 || !std::isfinite(r.image_scale_corr)) + continue; + if (!AcceptReflection(r, d_min_limit_A, d_max_limit_A)) + continue; + if (r.partiality < min_partiality) + continue; - for (const auto &r: integration_outcome[i].reflections) { - if (generator.IsSystematicallyAbsent(r)) - continue; - if (r.image_scale_corr <= 0.0 || !std::isfinite(r.image_scale_corr)) - continue; - if (!AcceptReflection(r, d_min_limit_A, d_max_limit_A)) - continue; - if (r.partiality < min_partiality) - continue; + const float I_corr = r.I * r.image_scale_corr; + const float sigma_corr = r.sigma * r.image_scale_corr; + if (!std::isfinite(I_corr) || !std::isfinite(sigma_corr) || sigma_corr <= 0.0f) + continue; - const float I_corr = r.I * r.image_scale_corr; - const float sigma_corr = r.sigma * r.image_scale_corr; - if (!std::isfinite(I_corr) || !std::isfinite(sigma_corr) || sigma_corr <= 0.0f) - continue; - - const auto shell = shells.GetShell(r.d); - if (!shell.has_value()) - continue; - const int s = *shell; - if (s >= 0 && s < n_shells) { - acc[s].total_obs++; - const auto key = generator(r).pack(); - const auto mit = merged_I.find(key); - if (mit != merged_I.end()) { - auto &ra = rmeas_obs[key]; - ra.sum_abs_dev += std::abs(static_cast(I_corr) - mit->second); - ra.sum_I += I_corr; - ra.n++; - ra.shell = s; + const auto shell = shells.GetShell(r.d); + if (!shell.has_value()) + continue; + const int s = *shell; + if (s >= 0 && s < n_shells) { + const auto key = generator(r).pack(); + const auto mit = merged_I.find(key); + out.push_back({key, I_corr, mit != merged_I.end() ? mit->second : NAN, s, + mit != merged_I.end()}); } } - } - } + }, + [&](size_t, const std::vector &out) { + for (const auto &o: out) { + acc[o.shell].total_obs++; + if (o.merged) { + auto &ra = rmeas_obs[o.key]; + ra.sum_abs_dev += std::abs(static_cast(o.I_corr) - o.merged_I); + ra.sum_I += o.I_corr; + ra.n++; + ra.shell = o.shell; + } + } + }); // R_meas per shell: sum over reflections of sqrt(n/(n-1)) * sum_i|I_i - |, over sum of I_i. // R_merge and R_pim beside it, with the factor 1 and sqrt(1/(n-1)) - R_pim following Weiss (2001) diff --git a/image_analysis/scale_merge/Merge.h b/image_analysis/scale_merge/Merge.h index b8f2b797e..7cca243a3 100644 --- a/image_analysis/scale_merge/Merge.h +++ b/image_analysis/scale_merge/Merge.h @@ -283,10 +283,11 @@ class MergeOnTheFly { // merged intensities feed the space-group search and the error model - those model fits must not // see the ice-contaminated intensities. The final in-symmetry merge keeps them (for completeness). bool exclude_ice_rings = false; + size_t nthreads = 1; HKLKeyGenerator generator; - std::map accumulator; + std::unordered_map accumulator; // exported in key order // Global error model (XDS form): sigma_corr^2 = a*sigma^2 + (b*)^2. a rescales the // (under-estimated) counting variance; the (b*)^2 term adds the intensity- @@ -315,12 +316,26 @@ class MergeOnTheFly { std::unordered_map reject_median_I; size_t reject_count = 0; - bool Mask(const IntegrationOutcome &outcome); + bool Mask(const IntegrationOutcome &outcome) const; + + // One accepted observation, reduced to what the accumulator adds for it. + struct Contribution { + uint64_t key; + int32_t h, k, l; // the reflection's own index, as the accumulator records it + float w, wI, d; + bool on_ice_ring; + }; + // The observations of one image that reach the merge; returns how many the outlier test dropped. + size_t Contributions(const IntegrationOutcome &outcome, std::vector &out) const; + void Accumulate(const std::vector &contributions, int64_t image_id); public: MergeOnTheFly(const DiffractionExperiment &x); MergeOnTheFly& ReferenceCell(const std::optional &cell); MergeOnTheFly& ExcludeIceRings(bool input) { exclude_ice_rings = input; return *this; } MergeOnTheFly& FilterByImageCC(bool input) { filter_by_image_cc = input; return *this; } + // Workers for the passes over every observation (AddImages, RefineErrorModel, MergeStats); the + // result does not depend on the number. + MergeOnTheFly& Threads(size_t input) { nthreads = std::max(1, input); return *this; } // Fit the global error model from the spread of symmetry-equivalent observations. // Call once before merging; AddImage then applies it. @@ -341,6 +356,9 @@ public: // because it holds every observation at once; a streaming accumulator cannot, and there is no // second, device-side walk of the stills observations for this to disagree with. void AddImage(const IntegrationOutcome& outcome, int64_t image_id); + // AddImage for every outcome, the image id being its index: the same merge, bit for bit, with the + // per-observation work spread over the workers and the accumulation kept in image order. + void AddImages(const std::vector &outcomes); // d_min_override, when set, is the effective high-resolution limit for the shell table (used for // the automatic resolution cutoff computed by the caller); otherwise the manual @@ -355,10 +373,14 @@ public: std::optional declared_range = std::nullopt); std::vector ExportReflections(); + // The merged intensities alone - no free-R flags, no French-Wilson amplitudes - which is all a + // scaling reference reads. + std::vector ExportIntensities(); }; std::vector MergeAll(const DiffractionExperiment &x, - const std::vector &reflections); + const std::vector &reflections, + size_t nthreads); // Pearson CC between one image's corrected intensities (I * image_scale_corr) and a reference set of // full intensities, over the reflections that would enter the merge (non-ice, within the resolution diff --git a/image_analysis/scale_merge/StillsPartialityRefine.cpp b/image_analysis/scale_merge/StillsPartialityRefine.cpp index 44910f6cb..12265a510 100644 --- a/image_analysis/scale_merge/StillsPartialityRefine.cpp +++ b/image_analysis/scale_merge/StillsPartialityRefine.cpp @@ -358,7 +358,7 @@ double StillsPartialityRefine::Run(std::vector &outcomes, si for (int outer = 0; outer < settings_.outer_iterations; ++outer) { // Reference full intensities from the current corrections. - const std::vector merged = MergeAll(experiment_, outcomes); + const std::vector merged = MergeAll(experiment_, outcomes, nthreads); std::map reference; for (const auto &m: merged) reference[hkl_key_generator_(m)] = m.I; diff --git a/rugnux/RugnuxScaleMerge.cpp b/rugnux/RugnuxScaleMerge.cpp index a2b5f44a4..e42731bfa 100644 --- a/rugnux/RugnuxScaleMerge.cpp +++ b/rugnux/RugnuxScaleMerge.cpp @@ -1250,7 +1250,8 @@ bool Rugnux::ScaleMergeAndSymmetry(PipelineLocals &p) { // degrade weak stills (measured CC1/2 collapse at the default 3 iters). { phase("Scaling images (" + label + ")"); - auto merge_result = MergeAll(experiment_, indexer->GetIntegrationOutcome()); + TimingMark scale_mark(logger, "stills scaling images " + label); + auto merge_result = MergeAll(experiment_, indexer->GetIntegrationOutcome(), config_.nthreads); indexer->ScaleAllImages(merge_result); } // Physical partiality post-refinement (default on; --simple-stills disables): refine a per-crystal @@ -1258,6 +1259,7 @@ bool Rugnux::ScaleMergeAndSymmetry(PipelineLocals &p) { // correction (no re-integration). Never for the P1 search pass. if (!for_search && experiment_.GetScalingSettings().GetStillsPartialityRefine()) { phase("Partiality post-refine (" + label + ")"); + TimingMark refine_mark(logger, "stills partiality post-refine " + label); StillsPartialityRefine refiner(experiment_); const double mean_tilt = refiner.Run(indexer->GetIntegrationOutcome(), config_.nthreads); logger.Info("Stills partiality post-refine: mean |dpsi| = {:.3f} deg", mean_tilt); @@ -1272,7 +1274,9 @@ bool Rugnux::ScaleMergeAndSymmetry(PipelineLocals &p) { const std::vector &merge_input = indexer->GetIntegrationOutcome(); phase("Merging"); + TimingMark merge_mark(logger, "stills error model and merge " + label); MergeOnTheFly merge_engine(experiment_); + merge_engine.Threads(config_.nthreads); // For the de-novo P1 pass (for_search) drop the ice-ring reflections from the merged // intensities and the error model, so the space-group search sees clean data; the final // in-symmetry merge keeps them so completeness is not lost. @@ -1296,11 +1300,11 @@ bool Rugnux::ScaleMergeAndSymmetry(PipelineLocals &p) { logger.Info("Error model (XDS convention): a={:.3f} b={:.3e} ISa={:.1f} chi2={:.2f}", em.a, em.b, em.isa, merge_engine.ErrorModelChi2()); } - for (size_t i = 0; i < merge_input.size(); ++i) - merge_engine.AddImage(merge_input[i], static_cast(i)); + merge_engine.AddImages(merge_input); ScaleMergeResult out; out.merged = merge_engine.ExportReflections(); + merge_mark.Stop(); // Automatic high-resolution cutoff (post-merge): trim the written reflections + reported // shells to the CC1/2 fall-off. The merge, scaling and error model above ran over the full @@ -1313,6 +1317,7 @@ bool Rugnux::ScaleMergeAndSymmetry(PipelineLocals &p) { cut_ss.GetResolutionCCTarget(), for_search, logger, &result.resolution_fit_A); phase("Computing statistics"); + TimingMark stats_mark(logger, "stills statistics " + label); out.statistics = merge_engine.MergeStats(out.merged, merge_input, config_.reference_data, effective_d_min); // The same kept reflections binned over the reference range (--report-resolution); the // stills error model is fitted over the whole range already, so no ISa of its own. diff --git a/rugnux/rugnux_cli.cpp b/rugnux/rugnux_cli.cpp index 714bd73d1..5527130e9 100644 --- a/rugnux/rugnux_cli.cpp +++ b/rugnux/rugnux_cli.cpp @@ -2103,7 +2103,7 @@ static int RunRugnux(int argc, char **argv) { // against a foreign dataset injects cross-dataset systematics and is a worse reference than // the data's own merge. The per-image scale G is the exact one-pass solution, so one pass // (iterating a self-rebuilt reference only re-fits the freshly-scaled noise on weak stills). - ScaleOnTheFly(experiment, MergeAll(experiment, reflections)).Scale(reflections, nthreads); + ScaleOnTheFly(experiment, MergeAll(experiment, reflections, nthreads)).Scale(reflections, nthreads); // Physical partiality post-refinement (default on; --simple-stills disables): refine a per-crystal // orientation tilt against the merge and recompute each reflection's partiality + scale correction // (no re-integration), then merge with the improved corrections. @@ -2113,6 +2113,7 @@ static int RunRugnux(int argc, char **argv) { logger.Info("Stills partiality post-refine: mean |dpsi| = {:.3f} deg", mean_tilt); } MergeOnTheFly merge_engine(experiment); + merge_engine.Threads(nthreads); merge_engine.ReferenceCell(experiment.GetUnitCell()); // --min-image-cc has to hold for the merge itself, not only for the reported statistics. merge_engine.FilterByImageCC(experiment.GetScalingSettings().GetMinCCForImage() > 0.0); @@ -2128,8 +2129,7 @@ static int RunRugnux(int argc, char **argv) { logger.Info("Error model (XDS convention): a={:.3f} b={:.3e} ISa={:.1f} chi2={:.2f}", em.a, em.b, em.isa, merge_engine.ErrorModelChi2()); } - for (size_t i = 0; i < reflections.size(); ++i) - merge_engine.AddImage(reflections[i], static_cast(i)); + merge_engine.AddImages(reflections); merged_reflections = merge_engine.ExportReflections(); // Automatic high-resolution cutoff (post-merge), matching the full-analysis path: a manual diff --git a/tests/MergeScaleTest.cpp b/tests/MergeScaleTest.cpp index aab3d5e9d..ed548a294 100644 --- a/tests/MergeScaleTest.cpp +++ b/tests/MergeScaleTest.cpp @@ -3,6 +3,7 @@ #include +#include #include #include "../image_analysis/bragg_integration/CalcISigma.h" #include "../image_analysis/scale_merge/ErrorModel.h" @@ -585,3 +586,67 @@ TEST_CASE("Global Wilson B over a widened window, with no upper cap", "[wilson]" CHECK(1.0 / (3.9 * 3.9) - 1.0 / (low_res.d_low_A * low_res.d_low_A) == Catch::Approx(1.0 / 9.0 - 1.0 / 16.0).epsilon(0.01)); } + +// The parallel AddImages is the serial AddImage loop, bit for bit: the per-observation work is spread +// over the workers, but every sum and every half-set choice is still made in image order. More images +// than one of its blocks, so the block boundary is crossed. +TEST_CASE("MergeOnTheFly_AddImagesIsTheSerialMerge") { + DiffractionExperiment x; + x.SetSpaceGroup(TestSpaceGroup()); + + std::mt19937 rng(7); + std::uniform_int_distribution index(-6, 6); + std::uniform_real_distribution unit(0.0f, 1.0f); + std::vector outcomes(5000); + for (auto &o: outcomes) { + for (int j = 0; j < 20; ++j) { + Reflection r{}; + r.h = index(rng); + r.k = index(rng); + r.l = index(rng); + if (r.h == 0 && r.k == 0 && r.l == 0) + continue; + r.d = 20.0f / std::sqrt(static_cast(r.h * r.h + r.k * r.k + r.l * r.l)); + r.I = 100.0f * unit(rng) - 5.0f; + r.sigma = 1.0f + 5.0f * unit(rng); + r.var_bkg = r.sigma * r.sigma * unit(rng); + r.prescaling_corr = 1.0f; + r.partiality = 0.5f + 0.5f * unit(rng); + r.image_scale_corr = 0.5f + unit(rng); + r.on_ice_ring = unit(rng) < 0.05f; + o.reflections.push_back(r); + } + } + + MergeOnTheFly serial(x); + serial.RefineErrorModel(outcomes); + for (size_t i = 0; i < outcomes.size(); ++i) + serial.AddImage(outcomes[i], static_cast(i)); + const auto serial_merged = serial.ExportReflections(); + + MergeOnTheFly parallel(x); + parallel.Threads(8); + parallel.RefineErrorModel(outcomes); + parallel.AddImages(outcomes); + const auto parallel_merged = parallel.ExportReflections(); + + REQUIRE(serial.ErrorModelActive()); + CHECK(serial.ErrorModelA() == parallel.ErrorModelA()); + CHECK(serial.ErrorModelB() == parallel.ErrorModelB()); + REQUIRE(serial_merged.size() == parallel_merged.size()); + for (size_t i = 0; i < serial_merged.size(); ++i) { + CHECK(serial_merged[i].h == parallel_merged[i].h); + CHECK(serial_merged[i].k == parallel_merged[i].k); + CHECK(serial_merged[i].l == parallel_merged[i].l); + CHECK(serial_merged[i].I == parallel_merged[i].I); + CHECK(serial_merged[i].sigma == parallel_merged[i].sigma); + for (int half = 0; half < 2; ++half) // NaN where a half is empty, on both sides + CHECK(std::memcmp(&serial_merged[i].I_half[half], ¶llel_merged[i].I_half[half], sizeof(float)) == 0); + } + + const auto serial_stats = serial.MergeStats(serial_merged, outcomes); + const auto parallel_stats = parallel.MergeStats(parallel_merged, outcomes); + CHECK(serial_stats.overall.total_observations == parallel_stats.overall.total_observations); + CHECK(serial_stats.overall.cc_half == parallel_stats.overall.cc_half); + CHECK(serial_stats.overall.r_meas == parallel_stats.overall.r_meas); +} From 4936f9b408c042371131651ed4eaefedbf8820db Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Wed, 14 Oct 2026 11:32:51 +0200 Subject: [PATCH 2/3] RESULT-CHANGING (owner decision): stills scaling reference and tilt refinement only over the crystals the merge keeps The final stills merge drops every crystal whose cell is not close to the consensus cell (MergeOnTheFly::Mask), but the scaling reference (MergeAll) and the partiality post-refinement were made over all of them. On two 25k-frame XFEL runs of a tetragonal crystal those were 26-30% of the indexed frames and 60% of the integrated reflections: weak frames (median 8 indexed spots) indexed in a 2x/3x/4x supercell or an unrelated lattice, whose supercell-only reflections are pure noise and whose indices do not name the true reflections. They went into the reference every merged crystal was scaled against, and through the tilt refinement. MergeAll and StillsPartialityRefine::Run now take the reference cell and use only the crystals it accepts - the same test, with the same tolerances, as the merge. Measured (both runs, mx path): CC1/2 0.9599 -> 0.9598 and 0.9645 -> 0.9644, I/sigma +0.05 and +0.02, R_pim unchanged, space group and resolution unchanged; --mode scale against a deposited tetragonal lysozyme model CC 0.9452 -> 0.9477 and 0.9461 -> 0.9473. Scale+merge 2.1x faster on --mode scale (14.6 -> 6.9 s, loaded box). Co-Authored-By: Claude Opus 5.5 (1M context) --- image_analysis/scale_merge/Merge.cpp | 5 +++-- image_analysis/scale_merge/Merge.h | 5 ++++- .../scale_merge/StillsPartialityRefine.cpp | 13 ++++++++++--- image_analysis/scale_merge/StillsPartialityRefine.h | 5 ++++- rugnux/RugnuxScaleMerge.cpp | 5 +++-- rugnux/rugnux_cli.cpp | 4 ++-- 6 files changed, 26 insertions(+), 11 deletions(-) diff --git a/image_analysis/scale_merge/Merge.cpp b/image_analysis/scale_merge/Merge.cpp index 2c3bf37f1..4f72271c7 100644 --- a/image_analysis/scale_merge/Merge.cpp +++ b/image_analysis/scale_merge/Merge.cpp @@ -469,9 +469,10 @@ std::vector MergeOnTheFly::ExportReflections() { std::vector MergeAll(const DiffractionExperiment &x, const std::vector &integration_outcome, - size_t nthreads) { + size_t nthreads, + const std::optional &reference_cell) { MergeOnTheFly merge(x); - merge.Threads(nthreads).AddImages(integration_outcome); + merge.Threads(nthreads).ReferenceCell(reference_cell).AddImages(integration_outcome); return merge.ExportIntensities(); } diff --git a/image_analysis/scale_merge/Merge.h b/image_analysis/scale_merge/Merge.h index 7cca243a3..cfb5cd615 100644 --- a/image_analysis/scale_merge/Merge.h +++ b/image_analysis/scale_merge/Merge.h @@ -378,9 +378,12 @@ public: std::vector ExportIntensities(); }; +// The merged intensities of the images whose cell is close to reference_cell (all of them without one): +// the reference a per-image scale is fitted against. std::vector MergeAll(const DiffractionExperiment &x, const std::vector &reflections, - size_t nthreads); + size_t nthreads, + const std::optional &reference_cell); // Pearson CC between one image's corrected intensities (I * image_scale_corr) and a reference set of // full intensities, over the reflections that would enter the merge (non-ice, within the resolution diff --git a/image_analysis/scale_merge/StillsPartialityRefine.cpp b/image_analysis/scale_merge/StillsPartialityRefine.cpp index 12265a510..60ee3f7dd 100644 --- a/image_analysis/scale_merge/StillsPartialityRefine.cpp +++ b/image_analysis/scale_merge/StillsPartialityRefine.cpp @@ -349,7 +349,14 @@ double StillsPartialityRefine::RefineOne(IntegrationOutcome &outcome, return tilt_deg; } -double StillsPartialityRefine::Run(std::vector &outcomes, size_t nthreads) const { +double StillsPartialityRefine::Run(std::vector &outcomes, size_t nthreads, + const std::optional &reference_cell) const { + const auto &indexing = experiment_.GetIndexingSettings(); + auto merged_crystal = [&](const IntegrationOutcome &o) { + return !reference_cell || o.latt.GetUnitCell().is_close(*reference_cell, + indexing.GetUnitCellDistTolerance(), + indexing.GetUnitCellAngleTolerance_deg()); + }; if (nthreads == 0) nthreads = AvailableCpus(); nthreads = std::max(1, nthreads); @@ -358,7 +365,7 @@ double StillsPartialityRefine::Run(std::vector &outcomes, si for (int outer = 0; outer < settings_.outer_iterations; ++outer) { // Reference full intensities from the current corrections. - const std::vector merged = MergeAll(experiment_, outcomes, nthreads); + const std::vector merged = MergeAll(experiment_, outcomes, nthreads, reference_cell); std::map reference; for (const auto &m: merged) reference[hkl_key_generator_(m)] = m.I; @@ -370,7 +377,7 @@ double StillsPartialityRefine::Run(std::vector &outcomes, si auto worker = [&]() { size_t i = next.fetch_add(1); while (i < outcomes.size()) { - const double t = RefineOne(outcomes[i], reference); + const double t = merged_crystal(outcomes[i]) ? RefineOne(outcomes[i], reference) : 0.0; if (t > 0.0) { double prev = tilt_sum.load(); while (!tilt_sum.compare_exchange_weak(prev, prev + t)) {} diff --git a/image_analysis/scale_merge/StillsPartialityRefine.h b/image_analysis/scale_merge/StillsPartialityRefine.h index fbce8c1dd..f74568b65 100644 --- a/image_analysis/scale_merge/StillsPartialityRefine.h +++ b/image_analysis/scale_merge/StillsPartialityRefine.h @@ -43,7 +43,10 @@ public: explicit StillsPartialityRefine(const DiffractionExperiment &x); // Refine all crystals in place. Returns the mean |dpsi| applied (degrees), for diagnostics. - double Run(std::vector &outcomes, size_t nthreads = 0) const; + // Only the crystals whose cell is close to reference_cell (all without one) are refined and make the + // reference: the others are not merged. + double Run(std::vector &outcomes, size_t nthreads, + const std::optional &reference_cell) const; private: const DiffractionExperiment experiment_; diff --git a/rugnux/RugnuxScaleMerge.cpp b/rugnux/RugnuxScaleMerge.cpp index e42731bfa..f390ebc3f 100644 --- a/rugnux/RugnuxScaleMerge.cpp +++ b/rugnux/RugnuxScaleMerge.cpp @@ -1251,7 +1251,8 @@ bool Rugnux::ScaleMergeAndSymmetry(PipelineLocals &p) { { phase("Scaling images (" + label + ")"); TimingMark scale_mark(logger, "stills scaling images " + label); - auto merge_result = MergeAll(experiment_, indexer->GetIntegrationOutcome(), config_.nthreads); + auto merge_result = MergeAll(experiment_, indexer->GetIntegrationOutcome(), config_.nthreads, + result.consensus_cell); indexer->ScaleAllImages(merge_result); } // Physical partiality post-refinement (default on; --simple-stills disables): refine a per-crystal @@ -1261,7 +1262,7 @@ bool Rugnux::ScaleMergeAndSymmetry(PipelineLocals &p) { phase("Partiality post-refine (" + label + ")"); TimingMark refine_mark(logger, "stills partiality post-refine " + label); StillsPartialityRefine refiner(experiment_); - const double mean_tilt = refiner.Run(indexer->GetIntegrationOutcome(), config_.nthreads); + const double mean_tilt = refiner.Run(indexer->GetIntegrationOutcome(), config_.nthreads, result.consensus_cell); logger.Info("Stills partiality post-refine: mean |dpsi| = {:.3f} deg", mean_tilt); // The refine fits its OWN per-image scale, with no floor, and adopts it unconditionally // when the image has no prior CC - which is exactly the state the guard above leaves. So diff --git a/rugnux/rugnux_cli.cpp b/rugnux/rugnux_cli.cpp index 5527130e9..6b29580ca 100644 --- a/rugnux/rugnux_cli.cpp +++ b/rugnux/rugnux_cli.cpp @@ -2103,13 +2103,13 @@ static int RunRugnux(int argc, char **argv) { // against a foreign dataset injects cross-dataset systematics and is a worse reference than // the data's own merge. The per-image scale G is the exact one-pass solution, so one pass // (iterating a self-rebuilt reference only re-fits the freshly-scaled noise on weak stills). - ScaleOnTheFly(experiment, MergeAll(experiment, reflections, nthreads)).Scale(reflections, nthreads); + ScaleOnTheFly(experiment, MergeAll(experiment, reflections, nthreads, experiment.GetUnitCell())).Scale(reflections, nthreads); // Physical partiality post-refinement (default on; --simple-stills disables): refine a per-crystal // orientation tilt against the merge and recompute each reflection's partiality + scale correction // (no re-integration), then merge with the improved corrections. if (experiment.GetScalingSettings().GetStillsPartialityRefine()) { StillsPartialityRefine refiner(experiment); - const double mean_tilt = refiner.Run(reflections, nthreads); + const double mean_tilt = refiner.Run(reflections, nthreads, experiment.GetUnitCell()); logger.Info("Stills partiality post-refine: mean |dpsi| = {:.3f} deg", mean_tilt); } MergeOnTheFly merge_engine(experiment); From 3c4d730f898a32001d19fdf5a4233db531a146d9 Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Wed, 14 Oct 2026 11:59:09 +0200 Subject: [PATCH 3/3] RESULT-CHANGING (owner decision): stills tilt post-refinement solved directly instead of by Ceres Each crystal's (psi_x, psi_y) tilt was refined by a Ceres problem rebuilt for every inner iteration: one auto-differentiated residual block and one Cauchy loss per reflection. With the scaling reference cleaned up it was 70% of the stills scale/merge. The same objective - Cauchy on the sigma-weighted intensity residuals plus the Gaussian prior on the tilt, the tilt bounded - is now minimised by a small Levenberg-Marquardt on the rho'-reweighted 2x2 normal equations, the tilt applied to all reflections through one rotation matrix and the Jacobian taken by central differences. It reaches the same answer to optimizer precision: on the new synthetic test (StillsPartialityRefine_RecoversTheTilt) partiality error 0.0559 -> 0.03019 with Ceres and 0.03016 with this, mean tilt 0.02065 vs 0.02067 deg; on two 25k-frame XFEL runs (--mode scale) CC1/2, I/sigma, R_pim and the CC against a deposited lysozyme model agree to 0.0005, mean |dpsi| 0.023 deg on both. Scale+merge 6.9 -> 3.0 s (--mode scale, loaded box). Co-Authored-By: Claude Opus 5.5 (1M context) --- .../scale_merge/StillsPartialityRefine.cpp | 176 +++++++++++------- tests/MergeScaleTest.cpp | 96 ++++++++++ 2 files changed, 205 insertions(+), 67 deletions(-) diff --git a/image_analysis/scale_merge/StillsPartialityRefine.cpp b/image_analysis/scale_merge/StillsPartialityRefine.cpp index 60ee3f7dd..21785a263 100644 --- a/image_analysis/scale_merge/StillsPartialityRefine.cpp +++ b/image_analysis/scale_merge/StillsPartialityRefine.cpp @@ -3,13 +3,14 @@ #include "StillsPartialityRefine.h" +#include +#include #include #include #include #include #include -#include #include #include "Merge.h" @@ -101,44 +102,114 @@ namespace { return G; } - // Ceres residual: refine the orientation tilt (psi_x, psi_y) holding the scale G fixed. The tilt - // rotates the base reciprocal vector q; partiality follows analytically. Residual is the intensity - // mismatch weighted by 1/sigma, exactly matching ScaleOnTheFly's intensity-space objective. - struct PsiResidual { - double qx, qy, qz; + // The rotation by the angle-axis vector (psi_x, psi_y, 0) as a matrix (Rodrigues), so that one tilt + // is applied to all the reflections of a crystal by a matrix product each. + std::array TiltMatrix(double psi_x, double psi_y) { + const double theta = std::sqrt(psi_x * psi_x + psi_y * psi_y); + if (theta < 1e-15) + return {1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0}; + const double kx = psi_x / theta, ky = psi_y / theta; + const double c = std::cos(theta), s = std::sin(theta), t = 1.0 - c; + return {c + t * kx * kx, t * kx * ky, s * ky, + t * kx * ky, c + t * ky * ky, -s * kx, + -s * ky, s * kx, c}; + } + + // ComputeP with the tilt given as its rotation matrix. + struct TiltModel { double s0x, s0y, s0z, inv_lambda; double gamma0, gamma_e, bw; - double G, inv_corr, Iref, Iobs, weight; - template - bool operator()(const T *const psi, T *residual) const { - T q[3] = {T(qx), T(qy), T(qz)}; - T aa[3] = {psi[0], psi[1], T(0.0)}; - T qr[3]; - ceres::AngleAxisRotatePoint(aa, q, qr); - const T Sx = qr[0] + T(s0x), Sy = qr[1] + T(s0y), Sz = qr[2] + T(s0z); - const T de = ceres::sqrt(Sx * Sx + Sy * Sy + Sz * Sz) - T(inv_lambda); - const T dstar = ceres::sqrt(qr[0] * qr[0] + qr[1] * qr[1] + qr[2] * qr[2]); - const T sig_ang = T(gamma_e) * dstar; - const T sbw = T(bw) * ceres::abs(qr[2]); - const T sig2 = T(gamma0) * T(gamma0) + sig_ang * sig_ang + sbw * sbw; - const T p = ceres::exp(T(-0.5) * de * de / sig2); - residual[0] = T(weight) * (T(G) * p * T(inv_corr) * T(Iref) - T(Iobs)); - return true; + double P(const std::array &R, const FitObs &o) const { + const double qx = R[0] * o.qx + R[1] * o.qy + R[2] * o.qz; + const double qy = R[3] * o.qx + R[4] * o.qy + R[5] * o.qz; + const double qz = R[6] * o.qx + R[7] * o.qy + R[8] * o.qz; + const double de = VecLen(qx + s0x, qy + s0y, qz + s0z) - inv_lambda; + const double sig_ang = gamma_e * VecLen(qx, qy, qz); + const double sbw = bw * std::fabs(qz); + const double sig2 = gamma0 * gamma0 + sig_ang * sig_ang + sbw * sbw; + if (!(sig2 > 0.0)) + return 1.0; + return std::exp(-0.5 * de * de / sig2); } }; - // Gaussian prior N(0, sigma_prior^2) on the tilt. The data residuals above are (model-obs)/sigma, a - // proper chi^2, so the MAP prior residual is simply dpsi/sigma_prior - no scale calibration needed. - struct PsiPrior { - double inv_sigma; - template - bool operator()(const T *const psi, T *residual) const { - residual[0] = T(inv_sigma) * psi[0]; - residual[1] = T(inv_sigma) * psi[1]; - return true; + // Refine the orientation tilt (psi_x, psi_y) holding the scale G fixed: minimise + // sum_i rho(r_i^2) / 2 + |psi / sigma_prior|^2 / 2, r_i = weight_i (G p_i(psi) Iref_i / corr_i - Iobs_i), + // with the Cauchy rho(s) = k^2 log(1 + s / k^2) - the intensity-space objective of ScaleOnTheFly, and + // the Gaussian prior a proper MAP term since r_i is already in sigma units. Levenberg-Marquardt on the + // normal equations reweighted by rho', the Jacobian by central differences, every trial tilt held inside + // the bound. Returns false when the cost cannot be evaluated, leaving psi as it came in. + bool RefineTilt(const std::vector &obs, const TiltModel &model, double G, double robust_k, + double inv_sigma_prior, double max_tilt, double psi[2]) { + constexpr int MAX_ITERATIONS = 25; + constexpr double STEP = 1e-7; // rad, for the central differences + const double k2 = robust_k * robust_k; + const double prior2 = inv_sigma_prior * inv_sigma_prior; + + auto cost = [&](const double p[2]) { + const auto R = TiltMatrix(p[0], p[1]); + double c = 0.5 * prior2 * (p[0] * p[0] + p[1] * p[1]); + for (const auto &o: obs) { + const double r = o.weight * (G * model.P(R, o) * o.inv_corr * o.Iref - o.Iobs); + c += 0.5 * k2 * std::log1p(r * r / k2); + } + return c; + }; + + double current = cost(psi); + if (!std::isfinite(current)) + return false; + double lambda = 1e-4; + for (int iter = 0; iter < MAX_ITERATIONS; ++iter) { + const auto R = TiltMatrix(psi[0], psi[1]); + const auto Rxp = TiltMatrix(psi[0] + STEP, psi[1]), Rxm = TiltMatrix(psi[0] - STEP, psi[1]); + const auto Ryp = TiltMatrix(psi[0], psi[1] + STEP), Rym = TiltMatrix(psi[0], psi[1] - STEP); + double A00 = prior2, A01 = 0.0, A11 = prior2; + double b0 = prior2 * psi[0], b1 = prior2 * psi[1]; + for (const auto &o: obs) { + const double a = o.weight * G * o.inv_corr * o.Iref; + const double r = a * model.P(R, o) - o.weight * o.Iobs; + const double jx = a * (model.P(Rxp, o) - model.P(Rxm, o)) / (2.0 * STEP); + const double jy = a * (model.P(Ryp, o) - model.P(Rym, o)) / (2.0 * STEP); + const double w = 1.0 / (1.0 + r * r / k2); // rho' + A00 += w * jx * jx; + A01 += w * jx * jy; + A11 += w * jy * jy; + b0 += w * jx * r; + b1 += w * jy * r; + } + + bool accepted = false; + double trial[2] = {psi[0], psi[1]}; + double trial_cost = current; + while (lambda < 1e10) { + const double M00 = A00 * (1.0 + lambda), M11 = A11 * (1.0 + lambda); + const double det = M00 * M11 - A01 * A01; + if (!(det > 0.0)) + break; + trial[0] = std::clamp(psi[0] - (M11 * b0 - A01 * b1) / det, -max_tilt, max_tilt); + trial[1] = std::clamp(psi[1] - (M00 * b1 - A01 * b0) / det, -max_tilt, max_tilt); + trial_cost = cost(trial); + if (std::isfinite(trial_cost) && trial_cost < current) { + accepted = true; + lambda = std::max(lambda / 10.0, 1e-12); + break; + } + lambda *= 10.0; + } + if (!accepted) + break; + const double step = std::hypot(trial[0] - psi[0], trial[1] - psi[1]); + const double decrease = current - trial_cost; + psi[0] = trial[0]; + psi[1] = trial[1]; + current = trial_cost; + if (decrease <= 1e-6 * current || step <= 1e-10) + break; } - }; + return true; + } } StillsPartialityRefine::StillsPartialityRefine(const DiffractionExperiment &x) @@ -238,6 +309,8 @@ double StillsPartialityRefine::RefineOne(IntegrationOutcome &outcome, double psi[2] = {0.0, 0.0}; double G = 1.0; + const TiltModel model{.s0x = S0.x, .s0y = S0.y, .s0z = S0.z, .inv_lambda = inv_lambda, + .gamma0 = gamma0, .gamma_e = gamma_e, .bw = bw}; const bool refine_tilt = obs.size() >= settings_.min_reflections; const int inner = refine_tilt ? settings_.inner_iterations : 1; @@ -258,42 +331,11 @@ double StillsPartialityRefine::RefineOne(IntegrationOutcome &outcome, if (!refine_tilt) break; - // (2) Refine the tilt holding G fixed. - ceres::Problem problem; - for (const auto &o: obs) { - auto *cost = new ceres::AutoDiffCostFunction(new PsiResidual{ - .qx = o.qx, .qy = o.qy, .qz = o.qz, - .s0x = S0.x, .s0y = S0.y, .s0z = S0.z, .inv_lambda = inv_lambda, - .gamma0 = gamma0, .gamma_e = gamma_e, .bw = bw, - .G = G, .inv_corr = o.inv_corr, .Iref = o.Iref, .Iobs = o.Iobs, .weight = o.weight}); - problem.AddResidualBlock(cost, new ceres::CauchyLoss(settings_.robust_k), psi); - } - if (settings_.prior_sigma_deg > 0.0) { - const double inv_sigma = kRadToDeg / settings_.prior_sigma_deg; // 1 / sigma_prior (rad) - problem.AddResidualBlock(new ceres::AutoDiffCostFunction( - new PsiPrior{inv_sigma}), nullptr, psi); - } - problem.SetParameterLowerBound(psi, 0, -settings_.max_tilt_rad); - problem.SetParameterUpperBound(psi, 0, settings_.max_tilt_rad); - problem.SetParameterLowerBound(psi, 1, -settings_.max_tilt_rad); - problem.SetParameterUpperBound(psi, 1, settings_.max_tilt_rad); - - ceres::Solver::Options options; - options.linear_solver_type = ceres::DENSE_QR; - options.minimizer_progress_to_stdout = false; - options.num_threads = 1; - options.max_num_iterations = 25; - ceres::Solver::Summary summary; - const double psi_before[2] = {psi[0], psi[1]}; - ceres::Solve(options, &problem, &summary); - // A solve that failed leaves whatever the minimiser last wrote in psi, and that tilt would go - // straight onto every reflection's partiality below. Keep the tilt this crystal came in with - // and stop refining it instead - G alone is still a usable model. - if (!summary.IsSolutionUsable()) { - psi[0] = psi_before[0]; - psi[1] = psi_before[1]; + // (2) Refine the tilt holding G fixed. A cost that cannot be evaluated keeps the tilt this crystal + // came in with and stops refining it - G alone is still a usable model. + const double inv_sigma_prior = settings_.prior_sigma_deg > 0.0 ? kRadToDeg / settings_.prior_sigma_deg : 0.0; + if (!RefineTilt(obs, model, G, settings_.robust_k, inv_sigma_prior, settings_.max_tilt_rad, psi)) break; - } } // Keep what the crystal came in with. This refinement is on by default, so a crystal the tilt model diff --git a/tests/MergeScaleTest.cpp b/tests/MergeScaleTest.cpp index ed548a294..d862fd213 100644 --- a/tests/MergeScaleTest.cpp +++ b/tests/MergeScaleTest.cpp @@ -10,6 +10,8 @@ #include "../image_analysis/scale_merge/HKLKey.h" #include "../image_analysis/scale_merge/Merge.h" #include "../image_analysis/scale_merge/ResolutionCutoff.h" +#include "../image_analysis/scale_merge/ScaleOnTheFly.h" +#include "../image_analysis/scale_merge/StillsPartialityRefine.h" #include "gemmi/reciproc.hpp" TEST_CASE("HKLKey_NoSG_noMergeFriedel") { @@ -650,3 +652,97 @@ TEST_CASE("MergeOnTheFly_AddImagesIsTheSerialMerge") { CHECK(serial_stats.overall.cc_half == parallel_stats.overall.cc_half); CHECK(serial_stats.overall.r_meas == parallel_stats.overall.r_meas); } + +// Stills partiality post-refinement on noiseless synthetic crystals: each crystal's true orientation is +// tilted from the stored one by up to 0.05 deg, its intensities are the true partials of a fixed set of +// full intensities, and its stored partialities those of the stored orientation. The refinement has to +// move the partialities towards the true ones. (The Ceres solve this replaced gave the same numbers: error +// 0.0559 -> 0.0302, mean tilt 0.0206 deg.) +TEST_CASE("StillsPartialityRefine_RecoversTheTilt") { + constexpr float CELL = 50.0f, WAVELENGTH = 1.0f, D_MIN = 2.0f; + constexpr double GAMMA0 = 0.001, ACCEPT = 1.7 * GAMMA0; // RLP width; |dist_ewald| accepted at prediction + DiffractionExperiment x; + x.SetSpaceGroup(gemmi::get_spacegroup_by_name("P 1")); + DiffractionGeometry geom; + geom.Wavelength_A(WAVELENGTH); + const Coord S0 = geom.GetScatteringVector(); + + std::mt19937 rng(11); + std::normal_distribution gauss(0.0, 1.0); + std::uniform_real_distribution unit(0.0, 1.0); + auto full_I = [](int h, int k, int l) { // a fixed, Friedel-symmetric pseudo-random full intensity + const uint64_t key = HKLKeyGenerator(true, gemmi::get_spacegroup_by_name("P 1"))(h, k, l).pack(); + return 10.0 + 990.0 * static_cast((key * 0x9e3779b97f4a7c15ULL) >> 40) / static_cast(1ULL << 24); + }; + auto rotate = [](const double R[9], const Coord &v) { + return Coord(static_cast(R[0] * v.x + R[1] * v.y + R[2] * v.z), + static_cast(R[3] * v.x + R[4] * v.y + R[5] * v.z), + static_cast(R[6] * v.x + R[7] * v.y + R[8] * v.z)); + }; + auto rotation = [](double ax, double ay, double az, double angle, double R[9]) { // Rodrigues + const double n = std::sqrt(ax * ax + ay * ay + az * az); + ax /= n; ay /= n; az /= n; + const double c = std::cos(angle), s = std::sin(angle), t = 1.0 - c; + const double M[9] = {c + t * ax * ax, t * ax * ay - s * az, t * ax * az + s * ay, + t * ax * ay + s * az, c + t * ay * ay, t * ay * az - s * ax, + t * ax * az - s * ay, t * ay * az + s * ax, c + t * az * az}; + std::copy(M, M + 9, R); + }; + auto partiality = [&](const Coord &q) { + const double de = (q + S0).Length() - 1.0 / WAVELENGTH; + return std::exp(-0.5 * de * de / (GAMMA0 * GAMMA0)); + }; + + std::vector outcomes(150); + std::vector> true_p(outcomes.size()); + const int hmax = static_cast(CELL / D_MIN); + for (size_t c = 0; c < outcomes.size(); ++c) { + double R[9], T[9]; + rotation(gauss(rng), gauss(rng), gauss(rng), 2.0 * M_PI * unit(rng), R); + const double tilt = (0.02 + 0.03 * unit(rng)) * M_PI / 180.0, phi = 2.0 * M_PI * unit(rng); + rotation(std::cos(phi), std::sin(phi), 0.0, tilt, T); // about an axis normal to the beam + auto &o = outcomes[c]; + o.geom = geom; + o.latt = CrystalLattice(rotate(R, Coord(CELL, 0, 0)), rotate(R, Coord(0, CELL, 0)), + rotate(R, Coord(0, 0, CELL))); + const double G = 0.5 + unit(rng); + for (int h = -hmax; h <= hmax; ++h) + for (int k = -hmax; k <= hmax; ++k) + for (int l = -hmax; l <= hmax; ++l) { + const Coord q = o.latt.Astar() * h + o.latt.Bstar() * k + o.latt.Cstar() * l; + const float dstar = q.Length(); + if (dstar == 0.0f || dstar > 1.0f / D_MIN) + continue; + if (std::fabs((q + S0).Length() - 1.0 / WAVELENGTH) > ACCEPT) + continue; + const double p_true = partiality(rotate(T, q)); + Reflection r{}; + r.h = h; r.k = k; r.l = l; + r.d = 1.0f / dstar; + r.I = static_cast(G * p_true * full_I(h, k, l)); + r.sigma = static_cast(1.0 + 0.01 * r.I); + r.prescaling_corr = 1.0f; + r.partiality = static_cast(partiality(q)); + r.image_scale_corr = 1.0f / r.partiality; + o.reflections.push_back(r); + true_p[c].push_back(p_true); + } + REQUIRE(o.reflections.size() >= 40); + } + + auto mean_abs_error = [&]() { + double sum = 0.0; + size_t n = 0; + for (size_t c = 0; c < outcomes.size(); ++c) + for (size_t i = 0; i < outcomes[c].reflections.size(); ++i, ++n) + sum += std::fabs(outcomes[c].reflections[i].partiality - true_p[c][i]); + return sum / static_cast(n); + }; + const double before = mean_abs_error(); + ScaleOnTheFly(x, MergeAll(x, outcomes, 4, std::nullopt)).Scale(outcomes, 4); + const double mean_tilt = StillsPartialityRefine(x).Run(outcomes, 4, std::nullopt); + const double after = mean_abs_error(); + CHECK(mean_tilt > 0.015); + CHECK(mean_tilt < 0.06); + CHECK(after < 0.6 * before); // measured 0.54: the prior holds the larger tilts back +}