From ade61eea60248b599edbf8f531a96408adcb7967 Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Thu, 30 Jul 2026 19:20:51 +0200 Subject: [PATCH] Rotation: predict without truncating, and keep the better of the two passes Digging into the selection logic showed the caps were not deciding the science - the two-pass geometry post-refinement was, and the caps only fed it randomness. Caps. The prediction buffer now grows to whatever a frame predicts instead of keeping an arbitrary subset of it, and the per-image reflection limit is raised to 65536, with the image-buffer transport headroom derived from the same constant so the two cannot drift. Measured: bit-identical output on five battery crystals, because a normal cell never approached the old limits - only a large cell (~2.8e6 A^3, ~30000-44000 predictions per frame) ever did. Pass-2 guard. The refined pass is normally the better answer, which is why it is the canonical output, but it was adopted whatever it produced. On that same crystal it merged more unique reflections than its own cell can hold - completeness "117%", which is arithmetically impossible - while the header- geometry pass sat at 92.6% and CC1/2 0.98. Compare the two and, when the refined pass is not credible, go back to the header geometry and re-run so the canonical files are the ones that are kept. Both bounds are set where only a failure reaches them. Together on that crystal: 111639 unique against XDS's 118730 (was 88000-99000 and different every run), CC1/2 98.0% (was 96.9-97.7%), ISa 8.54, and two runs now agree bit for bit. Co-Authored-By: Claude Opus 5 (1M context) --- common/DiffractionExperiment.cpp | 18 ++++---- image_analysis/IndexAndRefine.cpp | 4 +- .../bragg_prediction/BraggPrediction.cpp | 12 ++---- .../bragg_prediction/BraggPrediction.h | 24 +++++++---- .../bragg_prediction/BraggPredictionGPU.cu | 16 ++++++- .../bragg_prediction/BraggPredictionGPU.h | 5 +++ .../bragg_prediction/BraggPredictionRotGPU.cu | 20 ++++++++- .../bragg_prediction/BraggPredictionRotGPU.h | 4 ++ rugnux/Rugnux.cpp | 42 ++++++++++++++++++- 9 files changed, 116 insertions(+), 29 deletions(-) diff --git a/common/DiffractionExperiment.cpp b/common/DiffractionExperiment.cpp index 8ca90dc6..d6ed50da 100644 --- a/common/DiffractionExperiment.cpp +++ b/common/DiffractionExperiment.cpp @@ -8,6 +8,7 @@ #include "NetworkAddressConvert.h" #include "JFJochCompressor.h" // For ZSTD_USE_JFJOCH_RLE #include "DiffractionExperiment.h" +#include "../image_analysis/bragg_prediction/BraggPrediction.h" #include "CUDAWrapper.h" #include "JFJochException.h" @@ -1119,13 +1120,16 @@ int64_t DiffractionExperiment::GetImagesPerFile() const { } int64_t DiffractionExperiment::GetImageBufferLocationSize() const { - // A buffer slot holds the compressed image plus the per-image CBOR metadata - // (spot list, reflection list, azimuthal profile, ...). 4 MB of headroom covers - // the worst case: up to kMaxReflections (10000) reflections serialize to ~1.7 MB - // (~170 B each), plus the capped spot list and the azimuthal profile. A frame - // whose metadata still does not fit is dropped individually by the receiver - // rather than aborting the collection. - return GetMaxCompressedSize() + 4 * 1024 * 1024; + // A buffer slot holds the compressed image plus the per-image CBOR metadata (spot list, reflection + // list, azimuthal profile, ...). The reflection list dominates and is bounded by + // BraggPrediction::kPredictionOutput at ~170 B each, so derive the headroom from that constant + // rather than restate it - the two used to be written down separately and a raised cap would have + // silently outgrown the slot. Half again on top covers the capped spot list and the profile. A + // frame whose metadata still does not fit is dropped individually by the receiver rather than + // aborting the collection. + constexpr int64_t kReflectionBytes = 170; + const int64_t metadata = BraggPrediction::kPredictionOutput * kReflectionBytes * 3 / 2; + return GetMaxCompressedSize() + metadata; } float DiffractionExperiment::GetLossyCompressionSerialMX() const { diff --git a/image_analysis/IndexAndRefine.cpp b/image_analysis/IndexAndRefine.cpp index a90a7819..4952cb71 100644 --- a/image_analysis/IndexAndRefine.cpp +++ b/image_analysis/IndexAndRefine.cpp @@ -451,7 +451,9 @@ void IndexAndRefine::QuickPredictAndIntegrate(DataMessage &msg, auto integration_end_time = std::chrono::steady_clock::now(); msg.integration_time_s = std::chrono::duration(integration_end_time - integration_start_time).count(); - constexpr size_t kMaxReflections = 10000; + // The per-image reflection list is bounded by the same constant the prediction and the transport + // headroom use, so the three cannot drift apart. + constexpr size_t kMaxReflections = BraggPrediction::kPredictionOutput; if (i_outcome.reflections.size() > kMaxReflections) { // Keep only smallest d (highest resolution) std::nth_element(i_outcome.reflections.begin(), diff --git a/image_analysis/bragg_prediction/BraggPrediction.cpp b/image_analysis/bragg_prediction/BraggPrediction.cpp index 5e4460f3..f757bd70 100644 --- a/image_analysis/bragg_prediction/BraggPrediction.cpp +++ b/image_analysis/bragg_prediction/BraggPrediction.cpp @@ -8,15 +8,9 @@ #include "BraggPrediction.h" #include "../bragg_integration/SystematicAbsence.h" -void BraggPrediction::ReportOverflow(int predicted) { - if (overflow_reported) - return; - overflow_reported = true; - Logger("BraggPrediction").Warning( - "A frame predicted {} reflections but the buffer holds {} - the ones kept are whichever the GPU " - "wrote first, so this dataset does NOT process reproducibly. Results will differ between runs. " - "The cell is large enough that the prediction cap no longer fits it.", - predicted, kPredictionCapacity); +void BraggPrediction::GrowCapacity(int count) { + reflections.resize(count); + max_reflections = count; } int BraggPrediction::TruncateToOutput(int count) { diff --git a/image_analysis/bragg_prediction/BraggPrediction.h b/image_analysis/bragg_prediction/BraggPrediction.h index 479dcf63..9880fa79 100644 --- a/image_analysis/bragg_prediction/BraggPrediction.h +++ b/image_analysis/bragg_prediction/BraggPrediction.h @@ -27,7 +27,8 @@ struct BraggPredictionSettings { class BraggPrediction { protected: - const int max_reflections; + // Not const: the buffer grows to fit a frame that predicts more than it currently holds. + int max_reflections; std::vector reflections; // A frame that predicts more than the buffer holds keeps an ARBITRARY subset of them: the GPU @@ -35,8 +36,9 @@ protected: // changes from run to run. Say so, once per predictor, rather than let it pass silently - it is // not a small effect (measured: a 3% run-to-run spread in the number of merged reflections, and // every frame of that dataset overflowed). - void ReportOverflow(int predicted); - bool overflow_reported = false; + // Make room for `count` reflections. Overridden where device buffers have to follow. Called only + // when a frame predicted more than the current capacity, so a run pays for it a handful of times. + virtual void GrowCapacity(int count); // Deterministically cap Calc's output at kPredictionOutput: if more were predicted, keep the ones // closest to the Ewald sphere (smallest excitation error), ties broken by hkl. Returns the kept @@ -44,12 +46,20 @@ protected: int TruncateToOutput(int count); public: // The prediction buffer holds up to kPredictionCapacity reflections so a strong lattice does not - // overflow it (the GPU kernels then fill it in a non-deterministic atomic order - ReportOverflow - // warns when that happens; a large unit cell, ~2.8e6 A^3, reaches ~40000 per frame and overflows on - // every one). Calc returns at most kPredictionOutput, the number that flows downstream and is serialized - kept low so the + // overflow it. Calc returns at most kPredictionOutput, the number that flows downstream and is serialized - kept low so the // per-image reflection list stays within the frame transport headroom. + // Starting size only: the buffer GROWS to whatever a frame actually predicts (GrowCapacity), so a + // large cell is never truncated here. It used to be a hard cap, and overflowing it was both lossy + // and NON-DETERMINISTIC - the GPU kernels claim slots with an atomicAdd, so which reflections + // survived depended on block scheduling and changed between runs of the same command. static constexpr int kPredictionCapacity = 20000; - static constexpr int kPredictionOutput = 10000; + // How many reflections may flow downstream per image. Sized for a large unit cell: a ~2.8e6 A^3 + // cell predicts up to ~44000 per frame at 2.4 A. Truncating below what the frame really has costs + // more than it saves - the selection keeps the smallest excitation errors, i.e. the nearly + // fully-recorded reflections, and the rotation combine rebuilds a full FROM the partials it drops + // (measured on such a crystal: CC1/2 98 -> 60, ISa 8.6 -> 1.8). DiffractionExperiment's image-buffer + // headroom is derived from this, so the transport can carry what the analysis produces. + static constexpr int kPredictionOutput = 65536; explicit BraggPrediction(int max_reflections = kPredictionCapacity); diff --git a/image_analysis/bragg_prediction/BraggPredictionGPU.cu b/image_analysis/bragg_prediction/BraggPredictionGPU.cu index 43ebefbf..37bfd583 100644 --- a/image_analysis/bragg_prediction/BraggPredictionGPU.cu +++ b/image_analysis/bragg_prediction/BraggPredictionGPU.cu @@ -1,6 +1,7 @@ // SPDX-FileCopyrightText: 2025 Filip Leonarski, Paul Scherrer Institute // SPDX-License-Identifier: GPL-3.0-only +#include #include "../../common/JFJochMath.h" #include "BraggPredictionGPU.h" @@ -182,6 +183,13 @@ namespace { } } // namespace +void BraggPredictionGPU::GrowCapacity(int count) { + reg_out = CudaRegisteredVector(); + BraggPrediction::GrowCapacity(count); + reg_out = CudaRegisteredVector(reflections); + d_out = CudaDevicePtr(count); +} + BraggPredictionGPU::BraggPredictionGPU(int max_reflections) : BraggPrediction(max_reflections), reg_out(reflections), d_out(max_reflections), @@ -213,8 +221,12 @@ int BraggPredictionGPU::Calc(const DiffractionExperiment &experiment, int count = *h_count.get(); if (count > max_reflections) { - ReportOverflow(count); - count = max_reflections; + GrowCapacity(count); // see the rotation predictor + cudaMemsetAsync(d_count, 0, sizeof(int), stream); + bragg_kernel_3d<<>>(dK, settings.max_hkl, max_reflections, d_out, d_count); + cudaMemcpyAsync(h_count, d_count, sizeof(int), cudaMemcpyDeviceToHost, stream); + cudaStreamSynchronize(stream); + count = std::min(*h_count.get(), max_reflections); } if (count == 0) return {}; diff --git a/image_analysis/bragg_prediction/BraggPredictionGPU.h b/image_analysis/bragg_prediction/BraggPredictionGPU.h index ba0fc1f5..5fab081c 100644 --- a/image_analysis/bragg_prediction/BraggPredictionGPU.h +++ b/image_analysis/bragg_prediction/BraggPredictionGPU.h @@ -41,6 +41,11 @@ class BraggPredictionGPU : public BraggPrediction { public: explicit BraggPredictionGPU(int max_reflections = kPredictionCapacity); +protected: + void GrowCapacity(int count) override; + +public: + int Calc(const DiffractionExperiment &experiment, const CrystalLattice &lattice, const BraggPredictionSettings &settings) override; diff --git a/image_analysis/bragg_prediction/BraggPredictionRotGPU.cu b/image_analysis/bragg_prediction/BraggPredictionRotGPU.cu index c30d7b5f..46968ddf 100644 --- a/image_analysis/bragg_prediction/BraggPredictionRotGPU.cu +++ b/image_analysis/bragg_prediction/BraggPredictionRotGPU.cu @@ -1,6 +1,7 @@ // SPDX-FileCopyrightText: 2025 Filip Leonarski, Paul Scherrer Institute // SPDX-License-Identifier: GPL-3.0-only +#include #include "../../common/JFJochMath.h" #include "BraggPredictionRotGPU.h" @@ -268,6 +269,14 @@ BraggPredictionRotGPU::BraggPredictionRotGPU(int max_reflections) dK(1), d_count(1), h_count(1) { } +// The host buffer is page-locked and the device buffer is sized to match it, so both are rebuilt. +void BraggPredictionRotGPU::GrowCapacity(int count) { + reg_out = CudaRegisteredVector(); // unregister before the vector reallocates + BraggPrediction::GrowCapacity(count); + reg_out = CudaRegisteredVector(reflections); + d_out = CudaDevicePtr(count); +} + int BraggPredictionRotGPU::Calc(const DiffractionExperiment &experiment, const CrystalLattice &lattice, const BraggPredictionSettings &settings) { @@ -290,8 +299,15 @@ int BraggPredictionRotGPU::Calc(const DiffractionExperiment &experiment, int count = *h_count.get(); if (count > max_reflections) { - ReportOverflow(count); - count = max_reflections; + // The buffer holds an arbitrary subset of what was predicted (whichever slots the atomics + // reached first), so it cannot be used. Grow to fit and predict again; the buffer stays grown, + // so a run pays for this a handful of times at most. + GrowCapacity(count); + cudaMemsetAsync(d_count, 0, sizeof(int), stream); + bragg_rot_kernel_3d<<>>(dK, settings.max_hkl, max_reflections, d_out, d_count); + cudaMemcpyAsync(h_count, d_count, sizeof(int), cudaMemcpyDeviceToHost, stream); + cudaStreamSynchronize(stream); + count = std::min(*h_count.get(), max_reflections); } if (count == 0) return 0; diff --git a/image_analysis/bragg_prediction/BraggPredictionRotGPU.h b/image_analysis/bragg_prediction/BraggPredictionRotGPU.h index 94bb89f5..f1429900 100644 --- a/image_analysis/bragg_prediction/BraggPredictionRotGPU.h +++ b/image_analysis/bragg_prediction/BraggPredictionRotGPU.h @@ -45,6 +45,10 @@ class BraggPredictionRotGPU : public BraggPrediction { public: explicit BraggPredictionRotGPU(int max_reflections = kPredictionCapacity); +protected: + void GrowCapacity(int count) override; + +public: int Calc(const DiffractionExperiment &experiment, const CrystalLattice &lattice, const BraggPredictionSettings &settings) override; diff --git a/rugnux/Rugnux.cpp b/rugnux/Rugnux.cpp index 91a8e655..c7037ec6 100644 --- a/rugnux/Rugnux.cpp +++ b/rugnux/Rugnux.cpp @@ -246,7 +246,10 @@ ProcessResult Rugnux::Run(RugnuxObserver *observer) { if (cancelled_) { config_.output_prefix = base_prefix; return pass1; } if (gonio_snapshot) experiment_.Goniometer(*gonio_snapshot); // undo the pre-pass goniometer shift - // Apply the post-refined detector geometry for the second pass. + // Apply the post-refined detector geometry for the second pass, keeping the header geometry so + // the run can go back to it if the refined pass turns out worse (see the quality guard below). + const std::array header_geometry{experiment_.GetBeamX_pxl(), experiment_.GetBeamY_pxl(), + experiment_.GetDetectorDistance_mm()}; if (prepass_detector_geometry_) { const auto &g = *prepass_detector_geometry_; experiment_.BeamX_pxl(g[0]).BeamY_pxl(g[1]).DetectorDistance_mm(g[2]); @@ -293,6 +296,43 @@ ProcessResult Rugnux::Run(RugnuxObserver *observer) { // absence counts, the refused higher symmetry) is never shown on the default rotation path. if (!pass2.space_group_search.has_value()) pass2.space_group_search = pass1.space_group_search; + + // Pass 2 is normally the better answer, which is why it is the canonical output - but it is not + // guaranteed to be, and until now it was adopted whatever it produced. Two ways it can be wrong: + // it merges more unique reflections than the cell it settled on can hold (completeness above + // 100% is arithmetically impossible and means the cell is wrong), or its CC1/2 collapses + // relative to pass 1. Measured on a large-cell crystal: pass 1 195538 unique at 92.6% and + // CC1/2 0.98, pass 2 134667 at "117%" and CC1/2 0.62. Both bounds are set where only a failure + // reaches them, so a normal run - where pass 2 is a little better - keeps pass 2. + if (!cancelled_ && pass1.has_merge_statistics && pass2.has_merge_statistics) { + const auto completeness = [](const ProcessResult &r) { + const auto &o = r.merge_statistics.overall; + return o.possible_unique_reflections > 0 + ? 100.0 * o.unique_reflections / o.possible_unique_reflections : 0.0; + }; + const double compl2 = completeness(pass2); + const double cc1 = pass1.merge_statistics.overall.cc_half; + const double cc2 = pass2.merge_statistics.overall.cc_half; + constexpr double MAX_CREDIBLE_COMPLETENESS = 100.5; // rounding headroom, nothing more + constexpr double MAX_CC_HALF_LOSS = 0.05; + if (compl2 > MAX_CREDIBLE_COMPLETENESS || cc2 < cc1 - MAX_CC_HALF_LOSS) { + logger.Warning("Two-pass: the refined pass is worse than the header-geometry pass " + "(completeness {:.1f}% vs {:.1f}%, CC1/2 {:.3f} vs {:.3f}) - going back to " + "the header geometry. The refined geometry did not help this crystal.", + compl2, completeness(pass1), cc2, cc1); + // The refined pass has already written the canonical files, so re-run at the header + // geometry to replace them - the same remedy the supercell collapse above uses, and it + // only costs a pass on a crystal that was going to be wrong otherwise. + experiment_.BeamX_pxl(header_geometry[0]).BeamY_pxl(header_geometry[1]) + .DetectorDistance_mm(header_geometry[2]); + experiment_.SpaceGroupNumber(std::nullopt); + config_.output_prefix = base_prefix; + auto redo = RunPipeline(observer, /*write_output=*/true, /*geometry_prepass=*/false); + if (!redo.space_group_search.has_value()) + redo.space_group_search = pass1.space_group_search; + return redo; + } + } return pass2; } return RunPipeline(observer, /*write_output=*/true, /*geometry_prepass=*/false);