From 09bb88d4d4d604033d6d140a64db8e9c3d50f320 Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Mon, 28 Sep 2026 18:15:50 +0200 Subject: [PATCH] Bragg prediction on the CPU grows its buffer instead of stopping at 20000 BraggPrediction::Calc and BraggPredictionRot::Calc stopped filling at max_reflections (kPredictionCapacity = 20000), silently, inside the h/k/l walk - so a frame that predicted more kept whichever reflections the walk reached first, while the GPU predictors grow their buffer (GrowCapacity) and keep them all. A large cell therefore merged a different reflection set on a CPU-only build (the macOS/Windows viewer path, or CUDA_VISIBLE_DEVICES=) than on a GPU one. Both CPU predictors now double the buffer when the walk fills it, so they return the same set as the GPU; output_limit and TruncateToOutput are unchanged and remain the one cap on what leaves Calc, on both paths. Tests: BraggPrediction_GrowsPastStartingCapacity ([portable], CPU only) - a large cell predicting 28k (still) / 54k (rotation) reflections on one frame gives exactly what a buffer large enough from the outset gives, and an output_limit keeps the best-recorded half. BraggPrediction_CPU_GPU_ consistency_large_cell - the same frame on CPU and GPU: 28384/28384 and 53671/53671 matched by hkl, positions within 0.1 px. Both fail on the old code (count == 20000). Validation (CPU build): four in-house sets md5-identical to the base. A dense rotation set from the private arm predicted 72.0M (20000 x 3600 frames) before and 235.9M now, against the GPU's 235.9M; it ingests 218.65M partials, as the GPU does, and now reaches the GPU's cell, space group and merging statistics, where the capped CPU run settled on a different lattice. Cost on the CPU: 72.5 GB peak RSS (17.5 GB before), 44.5 min wall on a shared machine (7.2 min before). Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01D1G8gJVAy6gp1K5Dz3NE5C --- docs/CHANGELOG.md | 1 + .../bragg_prediction/BraggPrediction.cpp | 6 +- .../bragg_prediction/BraggPrediction.h | 27 ++-- .../bragg_prediction/BraggPredictionRot.cpp | 8 +- tests/CalcBraggPredictionTest.cpp | 124 ++++++++++++++++++ 5 files changed, 143 insertions(+), 23 deletions(-) diff --git a/docs/CHANGELOG.md b/docs/CHANGELOG.md index 9fc33481c..67b416ebf 100644 --- a/docs/CHANGELOG.md +++ b/docs/CHANGELOG.md @@ -15,6 +15,7 @@ * Rugnux's `--model` rigid-body refinement computes the model's density and solvent mask on all threads, with results unchanged. * Rugnux's `--model` rigid-body refinement computes the model's structure factors from one copy of the model and derives the translation part of its Jacobian exactly, so it runs faster; placements can differ slightly. * `rugnux --model` also draws a null replicate again when its rigid-body placement ends near an orientation equivalent to the model's, not only when it starts near one. +* Rugnux without a GPU no longer stops Bragg prediction at 20,000 reflections per frame, so a large unit cell integrates the same reflections on a CPU-only build as on a GPU. ### 1.0.0-rc.173 diff --git a/image_analysis/bragg_prediction/BraggPrediction.cpp b/image_analysis/bragg_prediction/BraggPrediction.cpp index 6c4288ab9..53fa376c2 100644 --- a/image_analysis/bragg_prediction/BraggPrediction.cpp +++ b/image_analysis/bragg_prediction/BraggPrediction.cpp @@ -127,9 +127,6 @@ int BraggPrediction::Calc(const DiffractionExperiment &experiment, const Crystal if (systematic_absence(h, k, l, settings.centering)) continue; - if (i >= max_reflections) - continue; - float recip_x = AhBk_x + Cstar.x * l; float recip_y = AhBk_y + Cstar.y * l; float recip_z = AhBk_z + Cstar.z * l; @@ -206,6 +203,9 @@ int BraggPrediction::Calc(const DiffractionExperiment &experiment, const Crystal const float qe_corr = sensor_qe.Factor(cos_alpha); const float flight_corr = air.Factor(cos_alpha); + if (i == max_reflections) + GrowCapacity(2 * max_reflections); + float d = 1.0f / sqrtf(recip_sq); reflections[i] = Reflection{ .h = h, diff --git a/image_analysis/bragg_prediction/BraggPrediction.h b/image_analysis/bragg_prediction/BraggPrediction.h index 4d14a44b4..2248e716c 100644 --- a/image_analysis/bragg_prediction/BraggPrediction.h +++ b/image_analysis/bragg_prediction/BraggPrediction.h @@ -39,18 +39,17 @@ struct BraggPredictionSettings { class BraggPrediction { protected: - // Not const: on the GPU path the buffer grows to fit a frame that predicts more than it holds. + // Not const: the buffer grows to fit a frame that predicts more than it holds. int max_reflections; std::vector reflections; // 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. - // - // NOTE: only the GPU Calc overrides call this. BraggPrediction::Calc and BraggPredictionRot::Calc - // stop filling at max_reflections instead, silently - and because that cap is applied inside the - // h/k/l walk, before the resolution test, what survives is the low-|h| block rather than the - // reflections nearest the Ewald sphere. A cell large enough to overflow 20000 therefore yields - // different merged reflections on a CPU-only build than on a GPU one. + // Every Calc grows rather than stops: the CPU ones double the buffer when the walk fills it, the GPU + // ones re-run the kernel at the counted total. A CPU Calc that stopped at the capacity instead kept + // whichever reflections the h/k/l walk reached first, not the best-recorded ones, so a large cell + // merged a different reflection set on a CPU-only build than on a GPU one. The one cap on what + // leaves Calc is output_limit, applied by TruncateToOutput on both. virtual void GrowCapacity(int count); // Deterministically cap Calc's output at output_limit: if more were predicted, keep the best-recorded @@ -78,14 +77,12 @@ protected: std::vector order_keys; std::vector order_scratch; public: - // The prediction buffer holds up to kPredictionCapacity reflections so a strong lattice does not - // overflow it. Calc returns at most output_limit, the number that flows downstream and is serialized - kept low so the - // per-image reflection list stays within the frame transport headroom. - // Starting size. On the GPU path the buffer grows to whatever a frame actually predicts - // (GrowCapacity), so a large cell is not truncated there; the CPU path still caps at this value, - // see the note on GrowCapacity. It used to be a hard cap on both, and overflowing it was 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. + // Calc returns at most output_limit, the number that flows downstream and is serialized - kept low so + // the per-image reflection list stays within the frame transport headroom. + // kPredictionCapacity is the prediction buffer's starting size. The buffer grows to whatever a frame + // actually predicts (GrowCapacity), so a large cell is not truncated. It used to be a hard cap, and + // overflowing it was lossy - and on the GPU NON-DETERMINISTIC, since the kernels claim slots with an + // atomicAdd, so which reflections survived depended on block scheduling and changed between runs. static constexpr int kPredictionCapacity = 20000; // How many reflections may flow downstream per image, offline. 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 diff --git a/image_analysis/bragg_prediction/BraggPredictionRot.cpp b/image_analysis/bragg_prediction/BraggPredictionRot.cpp index 4a516262f..33b022a6c 100644 --- a/image_analysis/bragg_prediction/BraggPredictionRot.cpp +++ b/image_analysis/bragg_prediction/BraggPredictionRot.cpp @@ -109,8 +109,6 @@ int BraggPredictionRot::Calc(const DiffractionExperiment &experiment, const Crys if (systematic_absence(h, k, l, settings.centering)) continue; - if (i >= max_reflections) - continue; Coord p0 = Astar * h + Bstar * k + Cstar * l; float p0_sq = p0 * p0; @@ -144,9 +142,6 @@ int BraggPredictionRot::Calc(const DiffractionExperiment &experiment, const Crys } for (const auto& p_m1 : p_m1_opt) { - if (i >= max_reflections) - continue; - const float cosphi = (p_m1 * p0_m1 + p_m3 * p0_m3) / rho_sq; const float sinphi = (p_m1 * p0_m3 - p_m3 * p0_m1) / rho_sq; const double c = cosphi, s = sinphi; @@ -237,6 +232,9 @@ int BraggPredictionRot::Calc(const DiffractionExperiment &experiment, const Crys const float qe_corr = sensor_qe.Factor(cos_alpha); const float flight_corr = air.Factor(cos_alpha); + if (i == max_reflections) + GrowCapacity(2 * max_reflections); + float d = 1.0f / sqrtf(p0_sq); reflections[i] = Reflection{ .h = h, diff --git a/tests/CalcBraggPredictionTest.cpp b/tests/CalcBraggPredictionTest.cpp index 48f0c1661..6319a8efc 100644 --- a/tests/CalcBraggPredictionTest.cpp +++ b/tests/CalcBraggPredictionTest.cpp @@ -10,6 +10,7 @@ #include "../image_analysis/bragg_prediction/BraggPredictionRot.h" #include "../image_analysis/bragg_prediction/RockingSlice.h" #include +#include // The flight-path term is a zero-parameter prediction: a NIST attenuation coefficient, the stated // sample-to-detector distance, and Beer-Lambert. Nothing about it is fitted, so it can be checked @@ -359,6 +360,85 @@ TEST_CASE("BraggPredictionRot_PartialWalksAlongItsRing", "[rocking_slice]") { CHECK(tested > 20); } +// A large cell, and a frame that predicts more reflections than the prediction buffer starts with, on +// both the still and the rotation path. Shared by the CPU test below and the CPU/GPU parity test. +namespace { + DiffractionExperiment LargeCellExperiment() { + DiffractionExperiment experiment(DetJF4M()); + experiment.DetectorDistance_mm(100.0).BeamX_pxl(1500.0).BeamY_pxl(1000.0) + .IncidentEnergy_keV(12.0) + .Goniometer(GoniometerAxis("omega", 0.0f, 1.0f, Coord(-1, 0, 0), {})); + return experiment; + } + + CrystalLattice LargeCell() { + return CrystalLattice(Coord{150, 10, 0}, Coord{-10, 140, 5}, Coord{0, -5, 160}); + } + + BraggPredictionSettings LargeCellSettings() { + return BraggPredictionSettings{ + .high_res_A = 1.6, + .ewald_dist_cutoff = 0.004, + .max_h = 95, .max_k = 90, .max_l = 100, + .wedge_deg = 1.0f + }; + } +} + +TEST_CASE("BraggPrediction_GrowsPastStartingCapacity", "[portable]") { + // A frame predicting more than the buffer starts with must come back whole - the same reflections a + // buffer large enough from the outset returns - and not as whichever the h/k/l walk reached first. + const auto experiment = LargeCellExperiment(); + const auto lattice = LargeCell(); + const auto settings = LargeCellSettings(); + + BraggPrediction still_grown, still_large(1000000); + BraggPredictionRot rot_grown, rot_large(1000000); + const std::vector> pairs = { + {&still_grown, &still_large}, {&rot_grown, &rot_large} + }; + for (const auto &[grown, large] : pairs) { + const int n_grown = grown->Calc(experiment, lattice, settings); + const int n_large = large->Calc(experiment, lattice, settings); + REQUIRE(n_grown > BraggPrediction::kPredictionCapacity); + REQUIRE(n_grown == n_large); + for (int i = 0; i < n_grown; i++) { + const auto &a = grown->GetReflections()[i]; + const auto &b = large->GetReflections()[i]; + REQUIRE(a.h == b.h); + REQUIRE(a.k == b.k); + REQUIRE(a.l == b.l); + REQUIRE(a.predicted_x == b.predicted_x); + REQUIRE(a.predicted_y == b.predicted_y); + } + + // Where a limit applies, what is kept is the best-recorded part: no dropped reflection ranks above + // the worst kept one - by partiality on the rotation path, by excitation error on the still path. + const bool rotation = (grown == &rot_grown); + std::vector all(grown->GetReflections().begin(), grown->GetReflections().begin() + n_grown); + grown->output_limit = n_grown / 2; + const int n_kept = grown->Calc(experiment, lattice, settings); + REQUIRE(n_kept == n_grown / 2); + + std::set> kept; + float worst_partiality = 1.0f, worst_dist_ewald = 0.0f; + for (int i = 0; i < n_kept; i++) { + const auto &r = grown->GetReflections()[i]; + kept.insert({r.h, r.k, r.l, r.delta_phi_deg}); + worst_partiality = std::min(worst_partiality, r.partiality); + worst_dist_ewald = std::max(worst_dist_ewald, r.dist_ewald); + } + for (const auto &r : all) { + if (kept.contains({r.h, r.k, r.l, r.delta_phi_deg})) + continue; + if (rotation) + REQUIRE(r.partiality <= worst_partiality); + else + REQUIRE(r.dist_ewald >= worst_dist_ewald); + } + } +} + #ifdef JFJOCH_USE_CUDA #include "../image_analysis/bragg_prediction/BraggPredictionGPU.h" #include "../image_analysis/bragg_prediction/BraggPredictionRotGPU.h" @@ -652,4 +732,48 @@ TEST_CASE("BraggPredictionRot_CPU_GPU_consistency_tilted") { CHECK(min_corr < 0.99f); CHECK(max_flight > 1.0f); } +TEST_CASE("BraggPrediction_CPU_GPU_consistency_large_cell") { + // A frame predicting more than kPredictionCapacity: both paths grow their buffer, so both predict the + // same reflections. The CPU used to stop at the capacity and return the first 20000 of the walk. + const auto experiment = LargeCellExperiment(); + const auto lattice = LargeCell(); + const auto settings = LargeCellSettings(); + + BraggPrediction still_cpu; + BraggPredictionGPU still_gpu; + BraggPredictionRot rot_cpu; + BraggPredictionRotGPU rot_gpu; + const std::vector> pairs = { + {&still_cpu, &still_gpu}, {&rot_cpu, &rot_gpu} + }; + for (const auto &[cpu, gpu] : pairs) { + const int cpu_count = cpu->Calc(experiment, lattice, settings); + const int gpu_count = gpu->Calc(experiment, lattice, settings); + REQUIRE(cpu_count > BraggPrediction::kPredictionCapacity); + REQUIRE(gpu_count > BraggPrediction::kPredictionCapacity); + + // Keyed by hkl and, on the rotation path, by the side of the frame centre, which separates the + // two rocking solutions one hkl can have. A still has one solution per hkl. + const bool rotation = (cpu == &rot_cpu); + auto key = [rotation](const Reflection &r) { + return std::make_tuple(r.h, r.k, r.l, rotation && r.delta_phi_deg > 0); + }; + std::map, const Reflection *> gpu_map; + for (int i = 0; i < gpu_count; i++) + gpu_map[key(gpu->GetReflections()[i])] = &gpu->GetReflections()[i]; + int matched = 0; + for (int i = 0; i < cpu_count; i++) { + const auto &c = cpu->GetReflections()[i]; + auto it = gpu_map.find(key(c)); + if (it == gpu_map.end()) + continue; + CHECK(c.predicted_x == Catch::Approx(it->second->predicted_x).margin(0.1)); + CHECK(c.predicted_y == Catch::Approx(it->second->predicted_y).margin(0.1)); + matched++; + } + // The same set up to float rounding at the acceptance edges. + CHECK(matched >= cpu_count - cpu_count / 1000); + CHECK(matched >= gpu_count - gpu_count / 1000); + } +} #endif