Jungfraujoch/tests/CalcBraggPredictionTest.cpp

// SPDX-FileCopyrightText: 2024 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only

#include <catch2/catch_all.hpp>

#include "../image_analysis/bragg_prediction/BraggPrediction.h"
#include "../common/JFJochMath.h"
#include <iostream>
#include "../image_analysis/SensorAbsorption.h"
#include "../image_analysis/bragg_prediction/BraggPredictionRot.h"
#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
// against the tables it is computed from rather than against any measurement.
TEST_CASE("BraggPrediction_FlightPathAttenuation", "[air_path]") {
    using sensor_absorption::FlightPathAttenuation;
    constexpr double D_mm = 160.0;
    auto lambda_of = [](double keV) { return 12.39842 / keV; };
    auto air = [&](double keV) {
        return FlightPathAttenuation::Build(FlightPathMedium::Air, D_mm, lambda_of(keV));
    };

    // Attenuation length of dry air, as D/L at a known distance. These follow from the NIST dry-air
    // mass attenuation coefficients and 1.205e-3 g/cm^3, and span four decades over the corpus's
    // energy range: L is 8.2 m at 18 keV but only 8.9 cm at 3.8 keV.
    CHECK(air(18.0).d_over_L == Catch::Approx(0.01959).epsilon(0.01));
    CHECK(air(12.4).d_over_L == Catch::Approx(0.05351).epsilon(0.01));
    CHECK(air(3.76).d_over_L == Catch::Approx(1.7972).epsilon(0.01));

    // The argon K edge at 3.2029 keV is in the table and is a real step: argon is 1.28% of air by
    // mass but dominates its absorption here, so interpolating straight across it would understate
    // the correction in exactly the regime where the correction is largest.
    CHECK(air(3.21).d_over_L > air(3.19).d_over_L * 1.05f);

    const auto a18 = air(18.0);
    // Normalised at normal incidence: head-on is untouched, by construction rather than by rounding.
    CHECK(a18.Factor(1.0f) == 1.0f);
    // Always >= 1 and monotone in the obliquity, because a reflection that arrives at a larger angle
    // crossed strictly more of the medium.
    float prev = 1.0f;
    for (double alpha_deg : {10.0, 20.0, 30.0, 40.0, 50.0, 55.0}) {
        const float f = a18.Factor(static_cast<float>(std::cos(alpha_deg * PI / 180.0)));
        CHECK(f > prev);
        prev = f;
    }
    const float cos55 = static_cast<float>(std::cos(55.0 * PI / 180.0));
    // 18 keV over 160 mm at 55 degrees: +1.5%.
    CHECK(a18.Factor(cos55) == Catch::Approx(1.0147).epsilon(0.002));

    // It runs the other way to the sensor term at the same angle: the sensor makes an oblique
    // reflection read high, the medium makes it read low, and neither is the other's undoing.
    const auto qe = sensor_absorption::SensorQE::Build("Si", 450.0, lambda_of(18.0));
    CHECK(qe.Factor(cos55) < 1.0f);
    CHECK(a18.Factor(cos55) > 1.0f);

    // HELIUM is why a long-wavelength station is usable at all: at 3.8 keV it attenuates some three
    // orders of magnitude less than air, so the same geometry that costs a factor of several in air
    // costs a fraction of a per cent in helium. It is NOT vacuum, and is not modelled as one.
    const auto he = FlightPathAttenuation::Build(FlightPathMedium::Helium, D_mm, lambda_of(3.76));
    CHECK(he.d_over_L > 0.0f);
    CHECK(he.d_over_L < air(3.76).d_over_L / 100.0f);
    CHECK(he.Factor(cos55) > 1.0f);
    CHECK(he.Factor(cos55) < 1.01f);

    // VACUUM leaves every intensity exactly untouched, at every angle - bit-identical to applying
    // no correction at all.
    const auto vac = FlightPathAttenuation::Build(FlightPathMedium::Vacuum, D_mm, lambda_of(3.76));
    CHECK(vac.d_over_L == 0.0f);
    CHECK(vac.Factor(1.0f) == 1.0f);
    CHECK(vac.Factor(0.5f) == 1.0f);
    CHECK(vac.Factor(0.1f) == 1.0f);
    // Same for a file that states no distance or no wavelength.
    CHECK(FlightPathAttenuation::Build(FlightPathMedium::Air, 0.0, lambda_of(12.4)).d_over_L == 0.0f);
    CHECK(FlightPathAttenuation::Build(FlightPathMedium::Air, D_mm, 0.0).d_over_L == 0.0f);

    // What the report quotes as the worth of the assumption. On an untilted detector the correction
    // is a pure function of resolution, so its whole effect on merged data is a Wilson-B shift; the
    // sign is negative because lifting the high-angle data flattens the fall-off. Measured on real
    // data at three geometries, this predicts the observed shift to within 8%.
    const double dB_hard = air(13.0).WilsonBShift_A2(1.28, 50.0, lambda_of(13.0));
    const double dB_soft = air(3.76).WilsonBShift_A2(3.02, 50.0, lambda_of(3.76));
    CHECK(dB_hard < 0.0);
    CHECK(std::fabs(dB_hard) < 1.0);          // hard X-rays: nothing a user could read off the data
    CHECK(dB_soft < -5.0);                    // long wavelength: the dominant correction in the run
    CHECK(std::fabs(dB_soft) > std::fabs(dB_hard) * 10.0);
    // Vacuum is worth exactly nothing, which is what makes the report line meaningful.
    CHECK(vac.WilsonBShift_A2(3.02, 50.0, lambda_of(3.76)) == 0.0);
}

TEST_CASE("BraggPrediction_11keV", "[portable]") {
    DiffractionExperiment experiment(DetJF4M());
    experiment.DetectorDistance_mm(100.0).BeamX_pxl(1500.0).BeamY_pxl(1000.0)
    .IncidentEnergy_keV(11.0);

    DiffractionGeometry geom = experiment.GetDiffractionGeometry();

    CrystalLattice lattice(Coord{20, 10, 0}, Coord{-20, 40, 0},
                           Coord{0, 0, 200});

    BraggPredictionSettings settings{
        .high_res_A = 2.0,
        .ewald_dist_cutoff = 0.001,
        .max_h = 40, .max_k = 40, .max_l = 40
    };
    BraggPrediction prediction;

    int count = prediction.Calc(experiment, lattice, settings);
    REQUIRE(count > 0);

    for (int i = 0; i < count; i++) {
        auto r = prediction.GetReflections().at(i);
        auto recip = r.h * lattice.Astar() + r.k * lattice.Bstar() + r.l * lattice.Cstar();

        REQUIRE(std::abs(r.h) < settings.max_h );
        REQUIRE(std::abs(r.k) < settings.max_k );
        REQUIRE(std::abs(r.l) < settings.max_l );
        REQUIRE(r.d >= settings.high_res_A);
        REQUIRE(r.d == Catch::Approx(1/std::sqrt(recip * recip)).margin(0.01f));
        REQUIRE(r.dist_ewald == Catch::Approx(std::abs(geom.DistFromEwaldSphere(recip))).epsilon(1e-4));
        auto [x,y] = geom.RecipToDetector(recip);
        REQUIRE(r.predicted_x == Catch::Approx(x).margin(0.01));
        REQUIRE(r.predicted_y == Catch::Approx(y).margin(0.01));
    }
}

TEST_CASE("BraggPrediction_15keV") {
    DiffractionExperiment experiment(DetJF4M());
    experiment.DetectorDistance_mm(100.0).BeamX_pxl(1500.0).BeamY_pxl(1000.0)
    .IncidentEnergy_keV(15.0);

    DiffractionGeometry geom = experiment.GetDiffractionGeometry();

    CrystalLattice lattice(Coord{20, 10, 0}, Coord{-20, 40, 0},
                           Coord{0, 0, 200});

    BraggPredictionSettings settings{
        .high_res_A = 2.0,
        .ewald_dist_cutoff = 0.001,
        .max_h = 40, .max_k = 40, .max_l = 40
    };
    BraggPrediction prediction;

    int count = prediction.Calc(experiment, lattice, settings);
    REQUIRE(count > 0);

    for (int i = 0; i < count; i++) {
        auto r = prediction.GetReflections().at(i);
        auto recip = r.h * lattice.Astar() + r.k * lattice.Bstar() + r.l * lattice.Cstar();

        REQUIRE(std::abs(r.h) < settings.max_h );
        REQUIRE(std::abs(r.k) < settings.max_k );
        REQUIRE(std::abs(r.l) < settings.max_l );
        REQUIRE(r.d >= settings.high_res_A);
        REQUIRE(r.d == Catch::Approx(1/std::sqrt(recip * recip)).margin(0.01f));
        REQUIRE(r.dist_ewald == Catch::Approx(std::abs(geom.DistFromEwaldSphere(recip))).epsilon(1e-3));
        auto [x,y] = geom.RecipToDetector(recip);
        REQUIRE(r.predicted_x == Catch::Approx(x).margin(0.01));
        REQUIRE(r.predicted_y == Catch::Approx(y).margin(0.01));
    }
}

TEST_CASE("BraggPrediction_Rot1_Rot2", "[portable]") {
    DiffractionExperiment experiment(DetJF4M());
    experiment.DetectorDistance_mm(100.0).BeamX_pxl(1500.0).BeamY_pxl(1000.0)
    .PoniRot1_rad(2.0/180.0 * PI).PoniRot2_rad(3.0/180.0 * PI)
    .IncidentEnergy_keV(11.0);

    DiffractionGeometry geom = experiment.GetDiffractionGeometry();

    CrystalLattice lattice(Coord{20, 10, 0}, Coord{-20, 40, 0},
                           Coord{0, 0, 200});

    BraggPredictionSettings settings{
        .high_res_A = 2.0,
        .ewald_dist_cutoff = 0.001,
        .max_h = 40, .max_k = 40, .max_l = 40
    };
    BraggPrediction prediction;

    int count = prediction.Calc(experiment, lattice, settings);
    REQUIRE(count > 0);

    for (int i = 0; i < count; i++) {
        auto r = prediction.GetReflections().at(i);
        auto recip = r.h * lattice.Astar() + r.k * lattice.Bstar() + r.l * lattice.Cstar();

        REQUIRE(std::abs(r.h) < settings.max_h );
        REQUIRE(std::abs(r.k) < settings.max_k );
        REQUIRE(std::abs(r.l) < settings.max_l );
        REQUIRE(r.d >= settings.high_res_A);
        REQUIRE(r.d == Catch::Approx(1/std::sqrt(recip * recip)).margin(0.01f));
        REQUIRE(r.dist_ewald == Catch::Approx(std::abs(geom.DistFromEwaldSphere(recip))).epsilon(1e-4));
        auto [x,y] = geom.RecipToDetector(recip);
        REQUIRE(r.predicted_x == Catch::Approx(x).margin(0.01));
        REQUIRE(r.predicted_y == Catch::Approx(y).margin(0.01));
    }
}

TEST_CASE("BraggPrediction_backscattering") {
    DiffractionExperiment experiment(DetJF9M());
    experiment.DetectorDistance_mm(120.0).BeamX_pxl(1500.0).BeamY_pxl(1000.0)
        .IncidentEnergy_keV(3.0/WVL_1A_IN_KEV);

    // Orthogonal basis, sufficient to test parity rules
    CrystalLattice lattice(
        Coord{40, 0, 0},
        Coord{0, 50, 0},
        Coord{0, 0, 60}
    );

    BraggPredictionSettings settings{
        .high_res_A = 3.0f,
        .ewald_dist_cutoff = 0.1f, // Very large cutoff, to be able to see as many reflections as possible
        .max_h = 50, .max_k = 50, .max_l = 50
    };

    BraggPrediction pred;
    int count = pred.Calc(experiment, lattice, settings);
    REQUIRE(count > 0);
    for (const auto& r : pred.GetReflections()) {
        if (r.d == 0) break;
        REQUIRE(r.d > 3.0 / sqrt(2.0));
    }
}

TEST_CASE("BraggPrediction_systematic_absences", "[portable]") {
    DiffractionExperiment experiment(DetJF4M());
    experiment.DetectorDistance_mm(120.0).BeamX_pxl(1500.0).BeamY_pxl(1000.0)
        .IncidentEnergy_keV(12.0);

    // Orthogonal basis, sufficient to test parity rules
    CrystalLattice lattice(
        Coord{40, 0, 0},
        Coord{0, 50, 0},
        Coord{0, 0, 60}
    );

    BraggPredictionSettings settings{
        .high_res_A = 3.0f,
        .ewald_dist_cutoff = 0.1f, // Very large cutoff, to be able to see as many reflections as possible
        .max_h = 50, .max_k = 50, .max_l = 50
    };

    BraggPrediction pred;

    SECTION("I centering") {
        // 1) Body-centered I: reflections with h+k+l odd must be absent
        settings.centering = 'I';
        int count_I = pred.Calc(experiment, lattice, settings);
        REQUIRE(count_I > 0);
        for (const auto& r : pred.GetReflections()) {
            if (r.d == 0) break; // ignore unfilled tail if any
            REQUIRE(((r.h + r.k + r.l) % 2) == 0);
        }
    }

    SECTION ("F centering") {
        // 2) Face-centered F: h,k,l all even or all odd
        settings.centering = 'F';

        int count_F = pred.Calc(experiment, lattice, settings);
        REQUIRE(count_F > 0);
        for (const auto& r : pred.GetReflections()) {
            if (r.d == 0) break;
            const bool he = (r.h & 1) == 0, ke = (r.k & 1) == 0, le = (r.l & 1) == 0;
            const bool all_even = he && ke && le;
            const bool all_odd = (!he) && (!ke) && (!le);
            REQUIRE((all_even || all_odd));
        }
    }

    SECTION("R centering") {
        settings.centering = 'R';
        int count_R = pred.Calc(experiment, lattice, settings);
        REQUIRE(count_R > 0);

        // R (hexagonal setting): -h + k + l = 3n
        for (const auto& r : pred.GetReflections()) {
            if (r.d == 0) break;
            int cond = (-r.h + r.k + r.l) % 3;
            if (cond < 0) cond += 3;
            REQUIRE(cond == 0);
        }
    }
}

TEST_CASE("RockingSliceCentroid_TruncatedNormalMean", "[rocking_slice]") {
    // The closed form against a direct quadrature of the slice, on a frame that holds the centre, one on
    // the curve's flank and one on its far tail.
    const float sigma = 0.004f, half_wedge = 0.0015f, c1 = 1.0f / (std::sqrt(2.0f) * sigma);
    for (float phi : {0.0f, 0.001f, -0.006f, 0.012f}) {
        double sw = 0.0, stw = 0.0;
        for (int i = 0; i <= 20000; ++i) {
            const double t = phi - half_wedge + 2.0 * half_wedge * i / 20000.0;
            const double w = std::exp(-t * t / (2.0 * sigma * sigma));
            sw += w; stw += t * w;
        }
        const float partiality = (std::erf((phi + half_wedge) * c1) - std::erf((phi - half_wedge) * c1)) / 2.0f;
        INFO("phi " << phi);
        CHECK(RockingSliceCentroid_rad(phi, half_wedge, c1, partiality) == Catch::Approx(stw / sw).margin(2e-6));
    }
    CHECK(RockingSliceCentroid_rad(0.0f, half_wedge, c1, 0.3f) == 0.0f);
}

TEST_CASE("BraggPredictionRot_PartialWalksAlongItsRing", "[rocking_slice]") {
    // Each rotation frame predicts a partial where the frame's slice of its rocking curve puts it: the
    // exact-condition position walked along the Debye ring. Over the frames of one reflection the
    // predicted positions therefore stay at one distance from the beam and move monotonically along
    // the ring; a prediction at the exact condition would be the same point on every frame.
    DiffractionExperiment experiment(DetJF4M());
    experiment.DetectorDistance_mm(100.0).BeamX_pxl(1100.0).BeamY_pxl(1000.0).IncidentEnergy_keV(12.4);
    const GoniometerAxis axis("omega", 0.0f, 0.1f, Coord(-1, 0, 0), {});
    experiment.Goniometer(axis);
    const CrystalLattice lattice(Coord{40, 0, 0}, Coord{0, 50, 0}, Coord{0, 0, 60});
    BraggPredictionSettings settings{.high_res_A = 2.5, .ewald_dist_cutoff = 0.0015,
                                     .max_h = 20, .max_k = 25, .max_l = 30,
                                     .wedge_deg = 0.1f, .mosaicity_deg = 0.1f};
    const float bx = 1100.0f, by = 1000.0f;

    std::map<std::tuple<int, int, int>, std::vector<std::pair<float, float>>> track;   // (radius, azimuth)
    BraggPredictionRot pred;
    for (int frame = 0; frame < 300; ++frame) {
        const auto latt = lattice.Multiply(axis.GetTransformationAngle(frame * 0.1f));
        const int n = pred.Calc(experiment, latt, settings);
        for (int i = 0; i < n; ++i) {
            const auto &r = pred.GetReflections().at(i);
            if (r.zeta > 0.3f) continue;
            const float dx = r.predicted_x - bx, dy = r.predicted_y - by;
            track[{r.h, r.k, r.l}].emplace_back(std::hypot(dx, dy), std::atan2(dy, dx));
        }
    }
    int tested = 0;
    for (const auto &[hkl, pts] : track) {
        if (pts.size() < 20) continue;
        float rmin = pts[0].first, rmax = pts[0].first;
        float along_first = 0.0f, along_last = 0.0f;
        int sign_changes = 0;
        float prev_step = 0.0f;
        for (size_t j = 0; j < pts.size(); ++j) {
            rmin = std::min(rmin, pts[j].first);
            rmax = std::max(rmax, pts[j].first);
            const float along = (pts[j].second - pts[0].second) * pts[0].first;   // px along the ring
            if (j == 0) along_first = along;
            along_last = along;
            if (j > 0) {
                const float step = (pts[j].second - pts[j - 1].second);
                if (step * prev_step < 0.0f) ++sign_changes;
                if (step != 0.0f) prev_step = step;
            }
        }
        INFO("hkl " << std::get<0>(hkl) << " " << std::get<1>(hkl) << " " << std::get<2>(hkl) << " frames " << pts.size());
        CHECK(rmax - rmin < 0.05f);                              // stays on its ring
        CHECK(std::fabs(along_last - along_first) > 0.1f);       // but walks along it
        CHECK(sign_changes == 0);                                // in one direction
        ++tested;
    }
    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);
        }
    }
}

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 "../common/CUDAWrapper.h"
#include "../image_analysis/bragg_prediction/BraggPredictionGPU.h"
#include "../image_analysis/bragg_prediction/BraggPredictionRotGPU.h"
#include "../image_analysis/bragg_prediction/BraggPredictionRot.h"
#include <map>
#include <algorithm>

TEST_CASE("BraggPredictionGPU") {
    if (get_gpu_count() == 0)
        SKIP("No CUDA GPU present");
    DiffractionExperiment experiment(DetJF4M());
    experiment.DetectorDistance_mm(100.0).BeamX_pxl(1500.0).BeamY_pxl(1000.0)
    .IncidentEnergy_keV(13.0);

    DiffractionGeometry geom = experiment.GetDiffractionGeometry();

    CrystalLattice lattice(Coord{20, 10, 0}, Coord{-20, 40, 0},
                           Coord{0, 0, 200});

    BraggPredictionSettings settings{
        .high_res_A = 2.0,
        .ewald_dist_cutoff = 0.001,
        .max_h = 40, .max_k = 40, .max_l = 40
    };
    BraggPredictionGPU prediction;

    int count = prediction.Calc(experiment, lattice, settings);
    REQUIRE(count > 0);

    for (int i = 0; i < count; i++) {
        auto r = prediction.GetReflections().at(i);
        auto recip = r.h * lattice.Astar() + r.k * lattice.Bstar() + r.l * lattice.Cstar();

        REQUIRE(std::abs(r.h) < settings.max_h );
        REQUIRE(std::abs(r.k) < settings.max_k );
        REQUIRE(std::abs(r.l) < settings.max_l );
        REQUIRE(r.d >= settings.high_res_A);
        REQUIRE(r.d == Catch::Approx(1/std::sqrt(recip * recip)).margin(0.01f));
        REQUIRE(r.dist_ewald == Catch::Approx(std::abs(geom.DistFromEwaldSphere(recip))).epsilon(2e-2));
        auto [x,y] = geom.RecipToDetector(recip);
        REQUIRE(r.predicted_x == Catch::Approx(x).margin(0.05));
        REQUIRE(r.predicted_y == Catch::Approx(y).margin(0.05));
    }
}

TEST_CASE("BraggPredictionGPU_Rot1_Rot2") {
    if (get_gpu_count() == 0)
        SKIP("No CUDA GPU present");
    DiffractionExperiment experiment(DetJF4M());
    experiment.DetectorDistance_mm(100.0).BeamX_pxl(1500.0).BeamY_pxl(1000.0)
    .PoniRot1_rad(2.0/180.0 * PI).PoniRot2_rad(3.0/180.0 * PI)
    .IncidentEnergy_keV(11.0);

    DiffractionGeometry geom = experiment.GetDiffractionGeometry();

    CrystalLattice lattice(Coord{20, 10, 0}, Coord{-20, 40, 0},
                           Coord{0, 0, 200});

    BraggPredictionSettings settings{
        .high_res_A = 2.0,
        .ewald_dist_cutoff = 0.001,
        .max_h = 40, .max_k = 40, .max_l = 40
    };
    BraggPredictionGPU prediction;

    int count = prediction.Calc(experiment, lattice, settings);
    REQUIRE(count > 0);

    for (int i = 0; i < count; i++) {
        auto r = prediction.GetReflections().at(i);
        auto recip = r.h * lattice.Astar() + r.k * lattice.Bstar() + r.l * lattice.Cstar();

        REQUIRE(std::abs(r.h) < settings.max_h );
        REQUIRE(std::abs(r.k) < settings.max_k );
        REQUIRE(std::abs(r.l) < settings.max_l );
        REQUIRE(r.d >= settings.high_res_A);
        REQUIRE(r.d == Catch::Approx(1/std::sqrt(recip * recip)).margin(0.01f));
        REQUIRE(r.dist_ewald == Catch::Approx(std::abs(geom.DistFromEwaldSphere(recip))).epsilon(1e-3));
        auto [x,y] = geom.RecipToDetector(recip);
        REQUIRE(r.predicted_x == Catch::Approx(x).margin(0.05));
        REQUIRE(r.predicted_y == Catch::Approx(y).margin(0.05));
    }
}

TEST_CASE("BraggPredictionGPU_systematic_absences") {
    if (get_gpu_count() == 0)
        SKIP("No CUDA GPU present");
    DiffractionExperiment experiment(DetJF4M());
    experiment.DetectorDistance_mm(120.0).BeamX_pxl(1500.0).BeamY_pxl(1000.0)
        .IncidentEnergy_keV(12.0);

    CrystalLattice lattice(
        Coord{40, 0, 0},
        Coord{0, 50, 0},
        Coord{0, 0, 60}
    );

    BraggPredictionSettings settings{
        .high_res_A = 3.0f,
        .ewald_dist_cutoff = 0.1f,
        .max_h = 50, .max_k = 50, .max_l = 50
    };

    BraggPredictionGPU pred;
    SECTION ("I centering") {
        // 1) Body-centered I
        settings.centering = 'I';

        int count_I = pred.Calc(experiment, lattice, settings);
        REQUIRE(count_I > 0);
        for (const auto& r : pred.GetReflections()) {
            if (r.d == 0) break;
            REQUIRE(((r.h + r.k + r.l) % 2) == 0);
        }
    }

    SECTION ("F centering") {
        // 2) Face-centered F
        settings.centering = 'F';
        int count_F = pred.Calc(experiment, lattice, settings);
        REQUIRE(count_F > 0);
        for (const auto& r : pred.GetReflections()) {
            if (r.d == 0) break;
            const bool he = (r.h & 1) == 0, ke = (r.k & 1) == 0, le = (r.l & 1) == 0;
            const bool all_even = he && ke && le;
            const bool all_odd = (!he) && (!ke) && (!le);
            REQUIRE((all_even || all_odd));
        }
    }

    SECTION("R centering") {
        settings.centering = 'R';
        int count_R = pred.Calc(experiment, lattice, settings);
        REQUIRE(count_R > 0);

        // R (hexagonal setting): -h + k + l = 3n
        for (const auto& r : pred.GetReflections()) {
            if (r.d == 0) break;
            int cond = (-r.h + r.k + r.l) % 3;
            if (cond < 0) cond += 3;
            REQUIRE(cond == 0);
        }
    }
}

TEST_CASE("BraggPredictionGPU_backscattering") {
    if (get_gpu_count() == 0)
        SKIP("No CUDA GPU present");
    DiffractionExperiment experiment(DetJF9M());
    experiment.DetectorDistance_mm(120.0).BeamX_pxl(1500.0).BeamY_pxl(1000.0)
        .IncidentEnergy_keV(3.0/WVL_1A_IN_KEV);

    // Orthogonal basis, sufficient to test parity rules
    CrystalLattice lattice(
        Coord{40, 0, 0},
        Coord{0, 50, 0},
        Coord{0, 0, 60}
    );

    BraggPredictionSettings settings{
        .high_res_A = 3.0f,
        .ewald_dist_cutoff = 0.1f, // Very large cutoff, to be able to see as many reflections as possible
        .max_h = 50, .max_k = 50, .max_l = 50
    };

    BraggPredictionGPU pred;
    int count = pred.Calc(experiment, lattice, settings);
    REQUIRE(count > 0);
    for (const auto& r : pred.GetReflections()) {
        if (r.d == 0) break;
        REQUIRE(r.d > 3.0 / sqrt(2.0));
    }
}

TEST_CASE("BraggPrediction_CPU_GPU_consistency_tilted") {
    if (get_gpu_count() == 0)
        SKIP("No CUDA GPU present");
    // Verify CPU and GPU implementations produce identical results with tilted detector
    DiffractionExperiment experiment(DetJF4M());
    experiment.DetectorDistance_mm(100.0).BeamX_pxl(1500.0).BeamY_pxl(1000.0)
        .PoniRot1_rad(0.04).PoniRot2_rad(-0.025)
        .IncidentEnergy_keV(12.0);

    CrystalLattice lattice(Coord{30, 10, 0}, Coord{-15, 45, 0}, Coord{0, 0, 150});

    BraggPredictionSettings settings{
        .high_res_A = 2.0,
        .ewald_dist_cutoff = 0.0015,
        .max_h = 30, .max_k = 30, .max_l = 30
    };

    BraggPrediction cpu_pred;
    BraggPredictionGPU gpu_pred;

    int cpu_count = cpu_pred.Calc(experiment, lattice, settings);
    int gpu_count = gpu_pred.Calc(experiment, lattice, settings);

    REQUIRE(cpu_count > 0);
    REQUIRE(gpu_count > 0);

    // Build map of GPU reflections by hkl
    std::map<std::tuple<int,int,int>, const Reflection*> gpu_refl_map;
    for (int i = 0; i < gpu_count; ++i) {
        const auto& r = gpu_pred.GetReflections().at(i);
        gpu_refl_map[{r.h, r.k, r.l}] = &r;
    }

    // Check that each CPU reflection has a matching GPU reflection
    int matched = 0;
    float min_corr = 1.0f;
    float max_flight = 1.0f;
    for (int i = 0; i < cpu_count; ++i) {
        const auto& cpu_r = cpu_pred.GetReflections().at(i);
        auto key = std::make_tuple(cpu_r.h, cpu_r.k, cpu_r.l);
        auto it = gpu_refl_map.find(key);
        if (it != gpu_refl_map.end()) {
            const auto& gpu_r = *it->second;
            CHECK(cpu_r.predicted_x == Catch::Approx(gpu_r.predicted_x).margin(0.1));
            CHECK(cpu_r.predicted_y == Catch::Approx(gpu_r.predicted_y).margin(0.1));
            CHECK(cpu_r.d == Catch::Approx(gpu_r.d).margin(0.01));
            // Both halves of the correction are part of the prediction, not a downstream product:
            // the sensor-efficiency term lived on the CPU path alone for a while because nothing here
            // compared it, and it is now a field of its own - so it is compared as a field of its own.
            CHECK(cpu_r.prescaling_corr == Catch::Approx(gpu_r.prescaling_corr).epsilon(1e-4));
            CHECK(cpu_r.qe_corr == Catch::Approx(gpu_r.qe_corr).epsilon(1e-4));
            CHECK(cpu_r.flight_corr == Catch::Approx(gpu_r.flight_corr).epsilon(1e-4));
            CHECK(cpu_r.image_scale_corr == Catch::Approx(gpu_r.image_scale_corr).epsilon(1e-4));
            min_corr = std::min(min_corr, cpu_r.qe_corr);
            max_flight = std::max(max_flight, cpu_r.flight_corr);
            matched++;
        }
    }

    // Most reflections should match (allow for some numerical differences at boundaries)
    CHECK(matched > cpu_count * 0.95);
    // ... and the comparison above must not be vacuous: on this geometry (320 um Si at 12 keV,
    // reflections out to 2 A) the sensor correction is several per cent, so a qe_corr that stayed
    // at 1 on both sides would mean the correction had been dropped from BOTH paths. It is checked
    // on qe_corr and not on prescaling_corr, which no longer carries it.
    CHECK(min_corr < 0.99f);
    // The same for the air term, which runs the other way: at 12 keV over this distance the air
    // correction reaches a few tenths of a per cent at the detector corner, so an flight_corr pinned at
    // 1 on both sides would mean it had been dropped from BOTH paths rather than agreeing.
    CHECK(max_flight > 1.0f);
}

TEST_CASE("BraggPredictionRot_CPU_GPU_consistency_tilted") {
    if (get_gpu_count() == 0)
        SKIP("No CUDA GPU present");
    // The rotation counterpart of the test above. Same purpose: every per-reflection quantity the
    // two implementations both produce has to agree - the Lorentz-polarization factor and the sensor
    // efficiency each in their own field, so neither can hide behind the other in a product.
    DiffractionExperiment experiment(DetJF4M());
    experiment.DetectorDistance_mm(100.0).BeamX_pxl(1500.0).BeamY_pxl(1000.0)
        .PoniRot1_rad(0.04).PoniRot2_rad(-0.025)
        .IncidentEnergy_keV(12.0)
        .Goniometer(GoniometerAxis("omega", 0.0f, 0.1f, Coord(-1, 0, 0), {}));

    CrystalLattice lattice(Coord{30, 10, 0}, Coord{-15, 45, 0}, Coord{0, 0, 150});

    BraggPredictionSettings settings{
        .high_res_A = 2.0,
        .ewald_dist_cutoff = 0.0015,
        .max_h = 30, .max_k = 30, .max_l = 30
    };

    BraggPredictionRot cpu_pred;
    BraggPredictionRotGPU gpu_pred;

    int cpu_count = cpu_pred.Calc(experiment, lattice, settings);
    int gpu_count = gpu_pred.Calc(experiment, lattice, settings);

    REQUIRE(cpu_count > 0);
    REQUIRE(gpu_count > 0);

    std::map<std::tuple<int,int,int>, const Reflection*> gpu_refl_map;
    for (int i = 0; i < gpu_count; ++i) {
        const auto& r = gpu_pred.GetReflections().at(i);
        gpu_refl_map[{r.h, r.k, r.l}] = &r;
    }

    int matched = 0;
    float min_corr = 1.0f;
    float max_flight = 1.0f;
    for (int i = 0; i < cpu_count; ++i) {
        const auto& cpu_r = cpu_pred.GetReflections().at(i);
        auto it = gpu_refl_map.find(std::make_tuple(cpu_r.h, cpu_r.k, cpu_r.l));
        if (it != gpu_refl_map.end()) {
            const auto& gpu_r = *it->second;
            CHECK(cpu_r.predicted_x == Catch::Approx(gpu_r.predicted_x).margin(0.1));
            CHECK(cpu_r.predicted_y == Catch::Approx(gpu_r.predicted_y).margin(0.1));
            CHECK(cpu_r.d == Catch::Approx(gpu_r.d).margin(0.01));
            CHECK(cpu_r.prescaling_corr == Catch::Approx(gpu_r.prescaling_corr).epsilon(1e-3));
            CHECK(cpu_r.qe_corr == Catch::Approx(gpu_r.qe_corr).epsilon(1e-3));
            CHECK(cpu_r.flight_corr == Catch::Approx(gpu_r.flight_corr).epsilon(1e-3));
            min_corr = std::min(min_corr, cpu_r.qe_corr);
            max_flight = std::max(max_flight, cpu_r.flight_corr);
            matched++;
        }
    }

    CHECK(matched > cpu_count * 0.95);
    CHECK(min_corr < 0.99f);
    CHECK(max_flight > 1.0f);
}
TEST_CASE("BraggPrediction_CPU_GPU_consistency_large_cell") {
    if (get_gpu_count() == 0)
        SKIP("No CUDA GPU present");
    // 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);
    }
}

TEST_CASE("BraggPredictionRotGPU_DeviceOrderMatchesHost") {
    if (get_gpu_count() == 0)
        SKIP("No CUDA GPU present");
    // The GPU predictor orders and selects its reflections on the device. What it returns has to be what
    // OrderOutput and TruncateToOutput make of the same predictions on the host: the full list in hkl,
    // delta_phi order, and under a cap the best-recorded ones in TruncateToOutput's order.
    const auto experiment = LargeCellExperiment();
    const auto lattice = LargeCell();
    const auto settings = LargeCellSettings();

    BraggPredictionRotGPU gpu;
    gpu.output_limit = 1 << 30;
    const int n = gpu.Calc(experiment, lattice, settings);
    REQUIRE(n > BraggPrediction::kPredictionCapacity);
    const std::vector<Reflection> all(gpu.GetReflections().begin(), gpu.GetReflections().begin() + n);
    CHECK(std::is_sorted(all.begin(), all.end(), [](const Reflection &a, const Reflection &b) {
        if (a.h != b.h) return a.h < b.h;
        if (a.k != b.k) return a.k < b.k;
        if (a.l != b.l) return a.l < b.l;
        return a.delta_phi_deg < b.delta_phi_deg;
    }));

    for (const int limit : {n - 1, n / 2, n / 7, 1}) {
        std::vector<Reflection> host = all;
        std::partial_sort(host.begin(), host.begin() + limit, host.end(),
                          [](const Reflection &a, const Reflection &b) {
                              if (a.partiality != b.partiality) return a.partiality > b.partiality;
                              if (a.dist_ewald != b.dist_ewald) return a.dist_ewald < b.dist_ewald;
                              if (a.h != b.h) return a.h < b.h;
                              if (a.k != b.k) return a.k < b.k;
                              return a.l < b.l;
                          });
        gpu.output_limit = limit;
        REQUIRE(gpu.Calc(experiment, lattice, settings) == limit);
        INFO("limit " << limit);
        // Field by field: the kernel leaves the intensity fields unset, so the bytes of two runs differ there.
        auto same = [](const Reflection &a, const Reflection &b) {
            return a.h == b.h && a.k == b.k && a.l == b.l && a.delta_phi_deg == b.delta_phi_deg
                   && a.predicted_x == b.predicted_x && a.predicted_y == b.predicted_y && a.d == b.d
                   && a.dist_ewald == b.dist_ewald && a.partiality == b.partiality && a.zeta == b.zeta
                   && a.image_scale_corr == b.image_scale_corr;
        };
        size_t ndiff = 0;
        for (int i = 0; i < limit; i++)
            ndiff += !same(host[i], gpu.GetReflections()[i]);
        CHECK(ndiff == 0);
    }
}
#endif