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