BraggPredictionRot: walk only the l of each column that can pass the phi test

Rotation prediction enumerated the whole (2h+1)(2k+1)(2l+1) box every frame and kept the ~1% whose
rotation solution lands near the frame. A solution within phi_limit of the frame needs
|f| = ||p0|^2 + 2 S0.p0| <= 2 |S0_perp| |p0| phi_limit, and along an (h, k) column f is a quadratic
in l, so the l that can pass are at most two intervals. Those are walked - widened by 1% in f and by
a whole index at each end - in the same order and with the same arithmetic; the rest the phi test
rejected anyway. Without a phi limit (min_zeta 0) the whole column is walked as before. The GPU
kernel is unchanged.

A new test predicts 1000 random lattices, geometries, axes and rocking widths both ways and
requires the same reflections in the same order, every value bit for bit (30000 cases passed once
locally). CPU build, two interleaved warm runs: 8a1a output loop 241.5 -> 216.3 s, pre-pass loop
90.7 -> 73.8 s, whole run 427.0 -> 384.2 s; 8qaw (with the gzip change before this) loops
116.6 / 251.3 -> 94.4 / 206.8 s, user CPU 14700 -> 12718 s. GPU build within noise.
p.mtz byte-identical on myob, cytc, thau, kdp, 8a1a and 8qaw, GPU and CPU builds.

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

image_analysis/bragg_prediction/BraggPredictionRot.cpp

@@ -1,6 +1,9 @@
 // SPDX-FileCopyrightText: 2025 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
 // SPDX-License-Identifier: GPL-3.0-only
 
+#include <algorithm>
+#include <cmath>
+
 #include "../../common/JFJochMath.h"
 #include "BraggPredictionRot.h"
 #include "../SensorAbsorption.h"
@@ -84,6 +87,7 @@ int BraggPredictionRot::Calc(const DiffractionExperiment &experiment, const Crys
     // angle. phi_limit carries a 1e-3 rad margin, far above float rounding, so every solution rejected
     // here is rejected by that test too. -1 = nothing is rejected here.
     double cos_phi_limit = -1.0;
+    double phi_limit = -1.0;
     if (settings.min_zeta > 0.0f) {
         double mos_max_rad = mos_angle_rad;
         if (bandwidth_sigma > 0.0f) {
@@ -93,13 +97,33 @@ int BraggPredictionRot::Calc(const DiffractionExperiment &experiment, const Crys
                                    : INFINITY;
             mos_max_rad = std::sqrt(mos_max_rad * mos_max_rad + dphi_bw * dphi_bw);
         }
-        const double phi_limit = half_wedge_angle_rad
-                                 + (settings.mosaicity_multiplier * mos_max_rad + 1e-5) / settings.min_zeta
-                                 + 1e-3;
+        phi_limit = half_wedge_angle_rad
+                    + (settings.mosaicity_multiplier * mos_max_rad + 1e-5) / settings.min_zeta
+                    + 1e-3;
         if (phi_limit < PI)
             cos_phi_limit = std::cos(phi_limit);
+        else
+            phi_limit = -1.0;
     }
 
+    // Which l of an (h, k) column can pass that test at all, so the walk does not visit the rest of
+    // the box. f(p0) = |p0|^2 + 2 S0.p0 is zero where p0 lies on the Ewald sphere, and rotating p0 by
+    // phi about the axis moves S0.p0 by at most |S0_perp| |p0_perp| |phi|. So a reflection whose
+    // solution lies within phi_limit has |f| <= 2 |S0_perp| |p0| phi_limit <= f_max, and along a column
+    // f is a quadratic in l, so the l that can pass are at most two intervals. Both are widened - f_max
+    // by 1%, each end by a whole index - so every l the tests below keep is still walked, in the same
+    // order and with the same arithmetic; the rest the phi test would have rejected anyway.
+    auto dot = [](const Coord &a, const Coord &b) {
+        return static_cast<double>(a.x) * b.x + static_cast<double>(a.y) * b.y + static_cast<double>(a.z) * b.z;
+    };
+    const double S0_perp = std::sqrt(std::max(0.0, dot(S0, S0) - static_cast<double>(m2_S0) * m2_S0));
+    const double f_max = 1.01 * 2.0 * S0_perp * one_over_dmax * phi_limit;
+    const double C_C = dot(Cstar, Cstar), S0_C = dot(S0, Cstar);
+    // A root of the quadratic as an index, held inside the box first so the conversion cannot overflow.
+    auto to_l = [&](double v) {
+        return static_cast<int>(std::clamp(v, -settings.max_l - 2.0, settings.max_l + 2.0));
+    };
+
     // p0 = A* h + B* k + C* l, evaluated as ((A* h) + (B* k)) + (C* l) exactly as before, with the
     // terms that do not change in the inner loops taken out of them.
     std::vector<Coord> Cstar_l(2 * settings.max_l + 1);
@@ -112,7 +136,32 @@ int BraggPredictionRot::Calc(const DiffractionExperiment &experiment, const Crys
         for (int k = -settings.max_k; k <= settings.max_k; k++) {
             const Coord AB = Astar_h + Bstar * k;
 
-            for (int l = -settings.max_l; l <= settings.max_l; l++) {
+            // The l of this column that can pass (above): f(l) = C_C l^2 + b l + c <= f_max on
+            // [l_lo, l_hi], and f(l) < -f_max strictly inside (skip_lo, skip_hi).
+            int l_lo = -settings.max_l, l_hi = settings.max_l;
+            int skip_lo = settings.max_l + 1, skip_hi = settings.max_l + 1;
+            if (phi_limit > 0.0 && !enumerate_full_box) {
+                const double b = 2.0 * (dot(AB, Cstar) + S0_C);
+                const double c = dot(AB, AB) + 2.0 * dot(S0, AB);
+                const double disc_hi = b * b - 4.0 * C_C * (c - f_max);
+                if (!(disc_hi >= 0.0))
+                    continue;
+                const double r_hi = std::sqrt(disc_hi);
+                l_lo = std::max(l_lo, to_l(std::ceil((-b - r_hi) / (2.0 * C_C))) - 1);
+                l_hi = std::min(l_hi, to_l(std::floor((-b + r_hi) / (2.0 * C_C))) + 1);
+                const double disc_lo = b * b - 4.0 * C_C * (c + f_max);
+                if (disc_lo > 0.0) {
+                    const double r_lo = std::sqrt(disc_lo);
+                    skip_lo = to_l(std::floor((-b - r_lo) / (2.0 * C_C))) + 2;
+                    skip_hi = to_l(std::ceil((-b + r_lo) / (2.0 * C_C))) - 2;
+                }
+            }
+
+            for (int l = l_lo; l <= l_hi; l++) {
+                if (l >= skip_lo && l <= skip_hi)
+                    l = skip_hi + 1;
+                if (l > l_hi)
+                    break;
                 if (systematic_absence(h, k, l, settings.centering))
                     continue;
 

image_analysis/bragg_prediction/BraggPredictionRot.h

@@ -12,4 +12,8 @@ public:
     int Calc(const DiffractionExperiment &experiment,
         const CrystalLattice &lattice,
         const BraggPredictionSettings &settings) override;
+protected:
+    // Walk every l of every column rather than only those that can pass the phi test - for the test
+    // that checks the two give the same reflections.
+    bool enumerate_full_box = false;
 };

tests/CalcBraggPredictionTest.cpp

@@ -11,6 +11,8 @@
 #include "../image_analysis/bragg_prediction/RockingSlice.h"
 #include <map>
 #include <set>
+#include <random>
+#include <algorithm>
 
 // 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
@@ -439,6 +441,89 @@ TEST_CASE("BraggPrediction_GrowsPastStartingCapacity", "[portable]") {
     }
 }
 
+namespace {
+    // The rotation predictor walking every l of every (h, k) column, which is what it did before it
+    // learned to skip the l that cannot pass its phi test.
+    class BraggPredictionRotFullBox : public BraggPredictionRot {
+    public:
+        BraggPredictionRotFullBox() { enumerate_full_box = true; }
+    };
+}
+
+TEST_CASE("BraggPredictionRot_ColumnRangeKeepsEveryReflection", "[portable]") {
+    // The walk visits only the l of each column whose phi solution can lie near the frame. On random
+    // lattices, geometries, axes and rocking widths it must predict exactly what the whole box does:
+    // the same reflections, in the same order, with every value bit for bit.
+    std::mt19937 rng(20261008);
+    std::uniform_real_distribution<float> u(0.0f, 1.0f);
+    auto unit = [&] {
+        while (true) {
+            const Coord v(2 * u(rng) - 1, 2 * u(rng) - 1, 2 * u(rng) - 1);
+            if (v.Length() > 0.1f && v.Length() < 1.0f) return v.Normalize();
+        }
+    };
+    const char centerings[] = {'P', 'I', 'C', 'F', 'R', 'A'};
+    int cases = 0, predicted = 0;
+    while (cases < 1000) {
+        const Coord a = unit() * (5.0f + 195.0f * u(rng));
+        const Coord b = unit() * (5.0f + 195.0f * u(rng));
+        const Coord c = unit() * (5.0f + 195.0f * u(rng));
+        if (std::fabs(a * (b % c)) < 0.3f * a.Length() * b.Length() * c.Length())
+            continue;   // too flat a cell
+        const CrystalLattice lattice(a, b, c);
+
+        DiffractionExperiment experiment(DetJF4M());
+        experiment.DetectorDistance_mm(50.0f + 350.0f * u(rng))
+                .BeamX_pxl(2068.0f * u(rng)).BeamY_pxl(2164.0f * u(rng))
+                .PoniRot1_rad(0.3f * (u(rng) - 0.5f)).PoniRot2_rad(0.3f * (u(rng) - 0.5f))
+                .IncidentEnergy_keV(6.0f + 19.0f * u(rng));
+        const float wedge_deg = 0.01f + 2.0f * u(rng);
+        experiment.Goniometer(GoniometerAxis("omega", 360.0f * u(rng), wedge_deg, unit(), {}));
+
+        // A box of at most 60 per index keeps the whole-box walk affordable.
+        const float longest = std::max({a.Length(), b.Length(), c.Length()});
+        BraggPredictionSettings settings{
+            .high_res_A = std::max(0.7f + 3.0f * u(rng), longest / 60.0f),
+            .max_h = 0, .max_k = 0, .max_l = 0,
+            .centering = centerings[rng() % 6],
+            .wedge_deg = wedge_deg,
+            .mosaicity_deg = 0.01f + u(rng),
+            .min_zeta = u(rng) < 0.2f ? 0.0f : 0.01f + 0.2f * u(rng),
+            .mosaicity_multiplier = 2.0f + 4.0f * u(rng),
+            .bandwidth_sigma = u(rng) < 0.5f ? 0.0f : 0.01f * u(rng)
+        };
+        settings.max_h = static_cast<int>(std::ceil(a.Length() / settings.high_res_A));
+        settings.max_k = static_cast<int>(std::ceil(b.Length() / settings.high_res_A));
+        settings.max_l = static_cast<int>(std::ceil(c.Length() / settings.high_res_A));
+
+        BraggPredictionRot pred;
+        BraggPredictionRotFullBox full;
+        pred.output_limit = full.output_limit = 1000000;
+        const int n = pred.Calc(experiment, lattice, settings);
+        const int n_full = full.Calc(experiment, lattice, settings);
+        INFO("case " << cases);
+        REQUIRE(n == n_full);
+        for (int i = 0; i < n; i++) {
+            const auto &x = pred.GetReflections()[i];
+            const auto &y = full.GetReflections()[i];
+            REQUIRE(x.h == y.h);
+            REQUIRE(x.k == y.k);
+            REQUIRE(x.l == y.l);
+            REQUIRE(x.delta_phi_deg == y.delta_phi_deg);
+            REQUIRE(x.predicted_x == y.predicted_x);
+            REQUIRE(x.predicted_y == y.predicted_y);
+            REQUIRE(x.d == y.d);
+            REQUIRE(x.dist_ewald == y.dist_ewald);
+            REQUIRE(x.partiality == y.partiality);
+            REQUIRE(x.zeta == y.zeta);
+            REQUIRE(x.image_scale_corr == y.image_scale_corr);
+        }
+        predicted += n;
+        ++cases;
+    }
+    CHECK(predicted > 100000);
+}
+
 #ifdef JFJOCH_USE_CUDA
 #include "../image_analysis/bragg_prediction/BraggPredictionGPU.h"
 #include "../image_analysis/bragg_prediction/BraggPredictionRotGPU.h"