Merge branch 'cpu-prediction-cap' into rc173 (CPU Bragg prediction grows its buffer like the GPU, no silent 20000 cap)

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01D1G8gJVAy6gp1K5Dz3NE5C

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
 

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,

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<Reflection> 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<OrderKey> order_keys;
     std::vector<Reflection> 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

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,

tests/CalcBraggPredictionTest.cpp

@@ -10,6 +10,7 @@
 #include "../image_analysis/bragg_prediction/BraggPredictionRot.h"
 #include "../image_analysis/bragg_prediction/RockingSlice.h"
 #include <map>
+#include <set>
 
 // 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<std::pair<BraggPrediction *, BraggPrediction *>> 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<Reflection> 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<std::tuple<int, int, int, float>> 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<std::pair<BraggPrediction *, BraggPrediction *>> 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<std::tuple<int, int, int, bool>, 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