Jungfraujoch/tests/IndexingUnitTest.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 <optional>
#include <random>
#include "../writer/HDF5Objects.h"
#include "../image_analysis/indexing/IndexerFactory.h"
#include "../image_analysis/indexing/PostIndexingRefinement.h"
#include "../image_analysis/IndexAndRefine.h"
#include "../image_analysis/bragg_prediction/BraggPrediction.h"
#include "../common/Logger.h"

inline double round_err(double x) {
    return std::abs(x - std::round(x));
}

#ifdef JFJOCH_USE_CUDA

#include <Eigen/Dense>
#include <Eigen/Geometry>

namespace {
    Eigen::Matrix3f MakeRotation(float ax_deg, float ay_deg, float az_deg) {
        const float ax = ax_deg * static_cast<float>(M_PI) / 180.0f;
        const float ay = ay_deg * static_cast<float>(M_PI) / 180.0f;
        const float az = az_deg * static_cast<float>(M_PI) / 180.0f;

        return (Eigen::AngleAxisf(az, Eigen::Vector3f::UnitZ())
              * Eigen::AngleAxisf(ay, Eigen::Vector3f::UnitY())
              * Eigen::AngleAxisf(ax, Eigen::Vector3f::UnitX())).toRotationMatrix();
    }

    CrystalLattice RotateLattice(const CrystalLattice &lattice, const Eigen::Matrix3f &rot) {
        auto apply = [&](const Coord &v) {
            Eigen::Vector3f x(v.x, v.y, v.z);
            x = rot * x;
            return Coord(x.x(), x.y(), x.z());
        };

        return {
            apply(lattice.Vec0()),
            apply(lattice.Vec1()),
            apply(lattice.Vec2())
        };
    }

    std::vector<uint8_t> BuildIndexedMask(const std::vector<Coord> &spots,
                                          const CrystalLattice &lattice,
                                          float tolerance,
                                          int64_t &count) {
        const Coord a = lattice.Vec0();
        const Coord b = lattice.Vec1();
        const Coord c = lattice.Vec2();
        const float tol_sq = tolerance * tolerance;

        std::vector<uint8_t> mask(spots.size(), 0);
        count = 0;

        for (size_t i = 0; i < spots.size(); ++i) {
            const float h_fp = spots[i] * a;
            const float k_fp = spots[i] * b;
            const float l_fp = spots[i] * c;

            const float h_frac = h_fp - std::round(h_fp);
            const float k_frac = k_fp - std::round(k_fp);
            const float l_frac = l_fp - std::round(l_fp);

            const float norm_sq = h_frac * h_frac + k_frac * k_frac + l_frac * l_frac;
            if (norm_sq < tol_sq) {
                mask[i] = 1;
                ++count;
            }
        }

        return mask;
    }

    int64_t MaskOverlap(const std::vector<uint8_t> &a, const std::vector<uint8_t> &b) {
        int64_t overlap = 0;
        for (size_t i = 0; i < a.size(); ++i) {
            if (a[i] && b[i])
                ++overlap;
        }
        return overlap;
    }

    bool MatchesCellLengths(const UnitCell &lhs, const UnitCell &rhs, float rel_tol = 0.08f) {
        std::array<float, 3> a = {
            static_cast<float>(lhs.a),
            static_cast<float>(lhs.b),
            static_cast<float>(lhs.c)
        };
        std::array<float, 3> b = {
            static_cast<float>(rhs.a),
            static_cast<float>(rhs.b),
            static_cast<float>(rhs.c)
        };
        std::sort(a.begin(), a.end());
        std::sort(b.begin(), b.end());

        for (int i = 0; i < 3; ++i) {
            const float denom = std::max(b[i], 1e-6f);
            if (std::abs(a[i] - b[i]) / denom > rel_tol)
                return false;
        }
        return true;
    }

    std::vector<Coord> BuildPredictedReciprocalSpots(const DiffractionExperiment &experiment,
                                                     const CrystalLattice &lattice,
                                                     const BraggPredictionSettings &settings) {
        BraggPrediction prediction(20000);
        const int nref = prediction.Calc(experiment, lattice, settings);
        REQUIRE(nref > 0);

        const auto &refs = prediction.GetReflections();
        const Coord astar = lattice.Astar();
        const Coord bstar = lattice.Bstar();
        const Coord cstar = lattice.Cstar();

        std::vector<Coord> spots;
        spots.reserve(nref);

        for (int i = 0; i < nref; ++i) {
            const auto &r = refs[i];
            spots.emplace_back(static_cast<float>(r.h) * astar
                             + static_cast<float>(r.k) * bstar
                             + static_cast<float>(r.l) * cstar);
        }

        return spots;
    }
} // namespace

TEST_CASE("FastFeedbackIndexer","[Indexing]") {
    std::vector<Coord> hkl;
    for (int i = 1; i < 7; i++)
        for (int j = 1; j<6; j++)
            for (int k = 1; k < 4; k++)
                hkl.emplace_back(i,j,k);

    std::vector<UnitCell> cells;
    cells.emplace_back(30,40,50,90,90,90);
    cells.emplace_back(80,80,90,90,90,120);
    cells.emplace_back(40,45,80,90,82.5,90);

    DiffractionExperiment experiment;
    experiment.SetUnitCell(cells[0]);
    experiment.IndexingAlgorithm(IndexingAlgorithmEnum::FFBIDX);

    REQUIRE(experiment.GetIndexingAlgorithm() == IndexingAlgorithmEnum::FFBIDX);

    std::unique_ptr<Indexer> indexer = CreateIndexer(experiment);

    for (auto &c: cells) {
        CrystalLattice l(c);

        Eigen::Matrix3f m;
        m << l.Vec0().x, l.Vec0().y, l.Vec0().z,
                l.Vec1().x, l.Vec1().y, l.Vec1().z,
                l.Vec2().x, l.Vec2().y, l.Vec2().z;
        auto m1 = m.transpose().inverse();

        CrystalLattice recip_l(Coord(m1(0,0), m1(0,1), m1(0,2)),
                               Coord(m1(1,0), m1(1,1), m1(1,2)),
                               Coord(m1(2,0), m1(2,1), m1(2,2)));

        std::vector<Coord> recip;
        recip.reserve(hkl.size());
        for (const auto &i: hkl)
            recip.emplace_back(i.x * recip_l.Vec0() + i.y * recip_l.Vec1() + i.z * recip_l.Vec2());

        experiment.SetUnitCell(c);

        indexer->Setup(experiment);
        auto ret = indexer->Run(recip);
        REQUIRE(!ret.lattice.empty());

        double err[3] = {0.0, 0.0, 0.0};
        for (const auto &iter: recip) {
            err[0] += round_err(ret.lattice[0].Vec0() * iter);
            err[1] += round_err(ret.lattice[0].Vec1() * iter);
            err[2] += round_err(ret.lattice[0].Vec2() * iter);
        }
        REQUIRE (err[0] < 0.001 * recip.size());
        REQUIRE (err[1] < 0.001 * recip.size());
        REQUIRE (err[2] < 0.001 * recip.size());
    }
}

TEST_CASE("FFTIndexer","[Indexing]") {
    Logger logger("FFTIndexer");

    UnitCell uc(39,45,78,90,90,90);
    CrystalLattice cl(uc);

    DiffractionExperiment experiment;
    IndexingSettings settings;
    settings.Algorithm(IndexingAlgorithmEnum::FFT)
            .FFT_MaxUnitCell_A(250.0).FFT_HighResolution_A(2 * M_PI / 3.0);
    experiment.ImportIndexingSettings(settings).SetUnitCell(uc);

    REQUIRE(experiment.GetIndexingAlgorithm() == IndexingAlgorithmEnum::FFT);
    REQUIRE(experiment.GetIndexingSettings().GetTolerance() == Catch::Approx(0.1f));
    std::unique_ptr<Indexer> indexer = CreateIndexer(experiment);
    REQUIRE(indexer);

    std::vector<Coord> vec;
    for (int h = -2; h < 10; h++) {
        for (int k = -5; k < 10; k++) {
            for (int l = -3; l < 10; l++) {
                vec.push_back(h * cl.Astar() + k * cl.Bstar() + l * cl.Cstar());
            }
        }
    }
    logger.Info("Spots {}", vec.size());

    auto start = std::chrono::high_resolution_clock::now();
    auto result = indexer->Run(vec);
    auto end = std::chrono::high_resolution_clock::now();

    REQUIRE(result.lattice.size() == 1);

    auto uc_out = result.lattice[0].GetUnitCell();

    // Collect and sort both sets of lengths to compare order-independently
    std::array<float, 3> out_lengths = {
        static_cast<float>(uc_out.a),
        static_cast<float>(uc_out.b),
        static_cast<float>(uc_out.c)
    };
    std::array<float, 3> ref_lengths = {
        static_cast<float>(uc.a),
        static_cast<float>(uc.b),
        static_cast<float>(uc.c)
    };
    std::sort(out_lengths.begin(), out_lengths.end());
    std::sort(ref_lengths.begin(), ref_lengths.end());

    CHECK(out_lengths[0] == Catch::Approx(ref_lengths[0]));
    CHECK(out_lengths[1] == Catch::Approx(ref_lengths[1]));
    CHECK(out_lengths[2] == Catch::Approx(ref_lengths[2]));

    CHECK(uc_out.alpha == Catch::Approx(uc.alpha));
    CHECK(uc_out.beta == Catch::Approx(uc.beta));
    CHECK(uc_out.gamma == Catch::Approx(uc.gamma));

    logger.Info("Time: {} ms", std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count());
}

TEST_CASE("FFTIndexer_AxisLongerThanTheBound","[Indexing]") {
    // The projected histogram is sampled at 1/(2*max length), so an axis between the bound and twice it
    // folds back to 2*bound - L: a 360 A axis searched to 250 A peaks at 140 A. The fold must be undone,
    // not built into a cell.
    UnitCell uc(39, 45, 360, 90, 90, 90);
    CrystalLattice cl(uc);

    DiffractionExperiment experiment;
    IndexingSettings settings;
    settings.Algorithm(IndexingAlgorithmEnum::FFT)
            .FFT_MaxUnitCell_A(250.0).FFT_HighResolution_A(2 * M_PI / 3.0);
    experiment.ImportIndexingSettings(settings);
    std::unique_ptr<Indexer> indexer = CreateIndexer(experiment);
    REQUIRE(indexer);

    std::vector<Coord> vec;
    for (int h = -2; h < 8; h++)
        for (int k = -4; k < 8; k++)
            for (int l = -3; l < 60; l++)
                vec.push_back(h * cl.Astar() + k * cl.Bstar() + l * cl.Cstar());

    auto result = indexer->Run(vec);
    REQUIRE(!result.lattice.empty());
    const auto uc_out = result.lattice[0].GetUnitCell();
    std::array<float, 3> out = {uc_out.a, uc_out.b, uc_out.c};
    std::sort(out.begin(), out.end());
    CHECK(out[0] == Catch::Approx(39.0f).epsilon(0.01));
    CHECK(out[1] == Catch::Approx(45.0f).epsilon(0.01));
    CHECK(out[2] == Catch::Approx(360.0f).epsilon(0.01));
}

TEST_CASE("FFTIndexer_CentredReferenceCell","[Indexing]") {
    // A reference cell is quoted the way a deposit states it, which for a centred lattice is the
    // conventional cell - and the indexer's candidates are Niggli-reduced primitive cells, whose
    // edges that description does not have. The two descriptions of one lattice must therefore
    // select the same candidate, or the user who quotes the deposit gets no lattice at all.
    const UnitCell conventional{70.0f, 70.0f, 70.0f, 90.0f, 90.0f, 90.0f};
    const CrystalLattice primitive = CrystalLattice(conventional).ToPrimitive('F');
    const UnitCell primitive_uc = primitive.GetUnitCell();

    std::vector<Coord> vec;
    for (int h = -3; h < 8; h++)
        for (int k = -3; k < 8; k++)
            for (int l = -3; l < 8; l++)
                vec.push_back(h * primitive.Astar() + k * primitive.Bstar() + l * primitive.Cstar());

    auto index_with = [&](const UnitCell &reference) {
        DiffractionExperiment experiment;
        IndexingSettings settings;
        settings.Algorithm(IndexingAlgorithmEnum::FFT)
                .FFT_MaxUnitCell_A(250.0).FFT_HighResolution_A(2 * M_PI / 3.0);
        experiment.ImportIndexingSettings(settings).SetUnitCell(reference);
        std::unique_ptr<Indexer> indexer = CreateIndexer(experiment);
        REQUIRE(indexer);
        return indexer->Run(vec).lattice;
    };

    const auto from_primitive = index_with(primitive_uc);
    const auto from_conventional = index_with(conventional);

    REQUIRE(from_primitive.size() == 1);
    REQUIRE(from_conventional.size() == 1);

    for (const auto &found : {from_primitive[0], from_conventional[0]}) {
        const auto uc_out = found.GetUnitCell();
        CHECK(uc_out.a == Catch::Approx(primitive_uc.a).epsilon(0.01));
        CHECK(uc_out.b == Catch::Approx(primitive_uc.b).epsilon(0.01));
        CHECK(uc_out.c == Catch::Approx(primitive_uc.c).epsilon(0.01));
        CHECK(uc_out.alpha == Catch::Approx(primitive_uc.alpha).epsilon(0.01));
        CHECK(uc_out.beta == Catch::Approx(primitive_uc.beta).epsilon(0.01));
        CHECK(uc_out.gamma == Catch::Approx(primitive_uc.gamma).epsilon(0.01));
    }

}

TEST_CASE("FFTIndexer_SpindleSeverity", "[Indexing][Spindle]") {
    // End to end through the real indexer: a crystal whose shortest row lies on the spindle must
    // come back with a severity of 1 from a normal run, from a severity-only run - which must not
    // index anything - and not at all when the frame is below the spot floor.
    UnitCell uc{39, 45, 78, 90, 90, 90};
    CrystalLattice cl(uc);

    DiffractionExperiment experiment(DetJF4M());
    experiment.DetectorDistance_mm(75).BeamY_pxl(1136).BeamX_pxl(1090).IncidentEnergy_keV(12.4);
    // The 39 A axis of the cell lies along x, and so does the spindle.
    experiment.Goniometer(GoniometerAxis("omega", 0, 0.1f, Coord(1, 0, 0), {}));

    IndexingSettings settings;
    settings.Algorithm(IndexingAlgorithmEnum::FFT)
            .FFT_MaxUnitCell_A(250.0).FFT_HighResolution_A(2 * M_PI / 3.0);
    experiment.ImportIndexingSettings(settings).SetUnitCell(uc);

    std::unique_ptr<Indexer> indexer = CreateIndexer(experiment);
    REQUIRE(indexer);
    indexer->Setup(experiment);

    std::vector<Coord> vec;
    for (int h = -2; h < 10; h++)
        for (int k = -5; k < 10; k++)
            for (int l = -3; l < 10; l++)
                vec.push_back(h * cl.Astar() + k * cl.Bstar() + l * cl.Cstar());

    auto full = indexer->Run(vec);
    REQUIRE(full.spindle_blind_fraction.has_value());
    CHECK_THAT(*full.spindle_blind_fraction, Catch::Matchers::WithinAbs(1.0f, 1e-4));

    auto severity_only = indexer->Run(vec, /*severity_only=*/true);
    CHECK(severity_only.lattice.empty());
    CHECK_FALSE(severity_only.executed);
    REQUIRE(severity_only.spindle_blind_fraction.has_value());
    CHECK_THAT(*severity_only.spindle_blind_fraction, Catch::Matchers::WithinAbs(1.0f, 1e-4));

    // Below the spot floor there is no value - the CANNOT-SAY state - not a middling one.
    const std::vector<Coord> few(vec.begin(), vec.begin() + 40);
    auto starved = indexer->Run(few, /*severity_only=*/true);
    CHECK_FALSE(starved.spindle_blind_fraction.has_value());
}

TEST_CASE("PostIndexingRefinement_MultiLattice_TwoCrystals_BraggPrediction","[Indexing]") {
    Logger logger("PostIndexingRefinement_MultiLattice_TwoCrystals_BraggPrediction");

    UnitCell tetragonal_uc{36.9, 78.95, 78.95, 90.0, 90.0, 90.0};
    CrystalLattice tetragonal_base(tetragonal_uc);

    CrystalLattice crystal_rot_1 = RotateLattice(tetragonal_base, MakeRotation(10.0f, 18.0f, 27.0f));
    CrystalLattice crystal_rot_2 = RotateLattice(tetragonal_base, MakeRotation(66.0f, -14.0f, 101.0f));

    DiffractionExperiment experiment(DetJF4M());
    experiment.DetectorDistance_mm(75)
              .BeamY_pxl(1136)
              .BeamX_pxl(1090)
              .IncidentEnergy_keV(12.4);

    BraggPredictionSettings pred_settings{
        .high_res_A = 2.0f,
        .ewald_dist_cutoff = 0.0010f,
        .max_h = 20, .max_k = 20, .max_l = 20,
        .centering = 'P',
        .wedge_deg = 0.1f,
        .mosaicity_deg = 0.2f,
        .min_zeta = 0.05f,
        .mosaicity_multiplier = 4.0f
    };

    const auto spots_1 = BuildPredictedReciprocalSpots(experiment, crystal_rot_1, pred_settings);
    const auto spots_2 = BuildPredictedReciprocalSpots(experiment, crystal_rot_2, pred_settings);

    logger.Info("Predicted spots lattice 1: {}", spots_1.size());
    logger.Info("Predicted spots lattice 2: {}", spots_2.size());

    std::vector<Coord> spots;
    spots.reserve(spots_1.size() + spots_2.size());
    spots.insert(spots.end(), spots_1.begin(), spots_1.end());
    spots.insert(spots.end(), spots_2.begin(), spots_2.end());

    Eigen::MatrixX3<float> oCell(9, 3);
    Eigen::VectorX<float> scores(3);

    auto put_lattice = [&](int idx, const CrystalLattice &lattice) {
        oCell(idx * 3 + 0, 0) = lattice.Vec0().x;
        oCell(idx * 3 + 0, 1) = lattice.Vec0().y;
        oCell(idx * 3 + 0, 2) = lattice.Vec0().z;

        oCell(idx * 3 + 1, 0) = lattice.Vec1().x;
        oCell(idx * 3 + 1, 1) = lattice.Vec1().y;
        oCell(idx * 3 + 1, 2) = lattice.Vec1().z;

        oCell(idx * 3 + 2, 0) = lattice.Vec2().x;
        oCell(idx * 3 + 2, 1) = lattice.Vec2().y;
        oCell(idx * 3 + 2, 2) = lattice.Vec2().z;
    };

    put_lattice(0, crystal_rot_1);
    put_lattice(1, crystal_rot_2);
    put_lattice(2, crystal_rot_1); // duplicate to verify overlap rejection

    // Keep bootstrap scores tiny to disable candidate drift in iterative re-fitting.
    scores(0) = 1e-6f;
    scores(1) = 1.1e-6f;
    scores(2) = 2e-6f;

    RefineParameters params{
        .viable_cell_min_spots = 12,
        .dist_tolerance_vs_reference = 0.05f,
        .reference_unit_cell = std::nullopt,
        .min_length_A = 20.0f,
        .max_length_A = 120.0f,
        .min_angle_deg = 60.0f,
        .max_angle_deg = 120.0f,
        .indexing_tolerance = 0.05f
    };

    auto refined = Refine(spots, spots.size(), oCell, scores, params);

    REQUIRE(refined.size() >= 2);

    int crystal_count = 0;
    for (const auto &lattice : refined) {
        if (MatchesCellLengths(lattice.GetUnitCell(), tetragonal_uc))
            ++crystal_count;
    }
    CHECK(crystal_count >= 2);

    int64_t count_0 = 0;
    int64_t count_1 = 0;
    auto mask_0 = BuildIndexedMask(spots, refined[0], params.indexing_tolerance, count_0);
    auto mask_1 = BuildIndexedMask(spots, refined[1], params.indexing_tolerance, count_1);
    const int64_t overlap = MaskOverlap(mask_0, mask_1);
    const int64_t max_set = std::max(count_0, count_1);

    logger.Info("Returned lattice 0 indexes {} spots", count_0);
    logger.Info("Returned lattice 1 indexes {} spots", count_1);
    logger.Info("Overlap between returned lattices: {} / {}", overlap, max_set);

    CHECK(overlap <= static_cast<int64_t>(0.2f * static_cast<float>(max_set)));
}

/*
TEST_CASE("FFTIndexer_MultiLattice_TwoCrystals_BraggPrediction","[Indexing]") {
    Logger logger("FFTIndexer_MultiLattice_TwoCrystals_BraggPrediction");

    UnitCell tetragonal_uc{36.9, 78.95, 78.95, 90.0, 90.0, 90.0};
    CrystalLattice tetragonal_base(tetragonal_uc);

    CrystalLattice crystal_rot_1 = RotateLattice(tetragonal_base, MakeRotation(10.0f, 18.0f, 27.0f));
    CrystalLattice crystal_rot_2 = RotateLattice(tetragonal_base, MakeRotation(66.0f, -14.0f, 101.0f));

    DiffractionExperiment experiment(DetJF4M());
    experiment.DetectorDistance_mm(75)
              .BeamY_pxl(1136)
              .BeamX_pxl(1090)
              .IncidentEnergy_keV(12.4);

    BraggPredictionSettings pred_settings{
        .high_res_A = 2.0f,
        .ewald_dist_cutoff = 0.0010f,
        .max_h = 20, .max_k = 20, .max_l = 20,
        .centering = 'P',
        .wedge_deg = 0.1f,
        .mosaicity_deg = 0.2f,
        .min_zeta = 0.05f,
        .mosaicity_multiplier = 4.0f
    };

    auto spots_1 = BuildPredictedReciprocalSpots(experiment, crystal_rot_1, pred_settings);
    auto spots_2 = BuildPredictedReciprocalSpots(experiment, crystal_rot_2, pred_settings);

    logger.Info("Predicted spots lattice 1: {}", spots_1.size());
    logger.Info("Predicted spots lattice 2: {}", spots_2.size());

    std::vector<Coord> spots;
    spots.reserve(spots_1.size() + spots_2.size());
    spots.insert(spots.end(), spots_1.begin(), spots_1.end());
    spots.insert(spots.end(), spots_2.begin(), spots_2.end());

    IndexingSettings settings;
    settings.Algorithm(IndexingAlgorithmEnum::FFT)
            .FFT_MaxUnitCell_A(120.0)
            .FFT_HighResolution_A(2.0f)
            .FFT_NumVectors(1024);
    experiment.ImportIndexingSettings(settings)
              .SetUnitCell(tetragonal_uc);

    REQUIRE(experiment.GetIndexingAlgorithm() == IndexingAlgorithmEnum::FFT);

    std::unique_ptr<Indexer> indexer = CreateIndexer(experiment);
    REQUIRE(indexer);
    indexer->Setup(experiment);

    auto result = indexer->Run(spots);

    logger.Info("FFT returned {} lattices", result.lattice.size());
    REQUIRE(result.lattice.size() >= 2);

    const float tolerance = experiment.GetIndexingSettings().GetTolerance();

    int crystal_count = 0;
    for (size_t i = 0; i < result.lattice.size(); ++i) {
        auto uc = result.lattice[i].GetUnitCell();
        int64_t indexed_count = 0;
        BuildIndexedMask(spots, result.lattice[i], tolerance, indexed_count);
        logger.Info("Lattice {} cell ({:.1f} {:.1f} {:.1f}) indexes {} spots",
                    i, uc.a, uc.b, uc.c, indexed_count);
        if (MatchesCellLengths(uc, tetragonal_uc))
            ++crystal_count;
    }

    CHECK(crystal_count >= 2);

    // Verify the two best crystal lattices are distinct (low overlap)
    if (result.lattice.size() >= 2) {
        int64_t count_0 = 0, count_1 = 0;
        auto mask_0 = BuildIndexedMask(spots, result.lattice[0], tolerance, count_0);
        auto mask_1 = BuildIndexedMask(spots, result.lattice[1], tolerance, count_1);
        const int64_t overlap = MaskOverlap(mask_0, mask_1);
        const int64_t max_set = std::max(count_0, count_1);
        logger.Info("Top-2 overlap: {} / {}", overlap, max_set);
        CHECK(overlap <= static_cast<int64_t>(0.5f * static_cast<float>(max_set)));
    }
}

TEST_CASE("FFBIDXIndexer_MultiLattice_TwoCrystals_BraggPrediction","[Indexing]") {
    Logger logger("FFBIDXIndexer_MultiLattice_TwoCrystals_BraggPrediction");

    UnitCell tetragonal_uc{36.9, 78.95, 78.95, 90.0, 90.0, 90.0};
    CrystalLattice tetragonal_base(tetragonal_uc);

    CrystalLattice crystal_rot_1 = RotateLattice(tetragonal_base, MakeRotation(10.0f, 18.0f, 27.0f));
    CrystalLattice crystal_rot_2 = RotateLattice(tetragonal_base, MakeRotation(66.0f, -14.0f, 101.0f));

    DiffractionExperiment experiment(DetJF4M());
    experiment.DetectorDistance_mm(75)
              .BeamY_pxl(1136)
              .BeamX_pxl(1090)
              .IncidentEnergy_keV(12.4);

    BraggPredictionSettings pred_settings{
        .high_res_A = 2.0f,
        .ewald_dist_cutoff = 0.0010f,
        .max_h = 20, .max_k = 20, .max_l = 20,
        .centering = 'P',
        .wedge_deg = 0.1f,
        .mosaicity_deg = 0.2f,
        .min_zeta = 0.05f,
        .mosaicity_multiplier = 4.0f
    };

    auto spots_1 = BuildPredictedReciprocalSpots(experiment, crystal_rot_1, pred_settings);
    auto spots_2 = BuildPredictedReciprocalSpots(experiment, crystal_rot_2, pred_settings);

    logger.Info("Predicted spots lattice 1: {}", spots_1.size());
    logger.Info("Predicted spots lattice 2: {}", spots_2.size());

    std::vector<Coord> spots;
    spots.reserve(spots_1.size() + spots_2.size());
    spots.insert(spots.end(), spots_1.begin(), spots_1.end());
    spots.insert(spots.end(), spots_2.begin(), spots_2.end());

    experiment.SetUnitCell(tetragonal_uc);
    experiment.IndexingAlgorithm(IndexingAlgorithmEnum::FFBIDX);

    REQUIRE(experiment.GetIndexingAlgorithm() == IndexingAlgorithmEnum::FFBIDX);

    std::unique_ptr<Indexer> indexer = CreateIndexer(experiment);
    REQUIRE(indexer);
    indexer->Setup(experiment);

    auto result = indexer->Run(spots);

    logger.Info("FFBIDX returned {} lattices", result.lattice.size());
    REQUIRE(result.lattice.size() >= 2);

    const float tolerance = experiment.GetIndexingSettings().GetTolerance();

    int crystal_count = 0;
    for (size_t i = 0; i < result.lattice.size(); ++i) {
        auto uc = result.lattice[i].GetUnitCell();
        int64_t indexed_count = 0;
        BuildIndexedMask(spots, result.lattice[i], tolerance, indexed_count);
        logger.Info("Lattice {} cell ({:.1f} {:.1f} {:.1f}) indexes {} spots",
                    i, uc.a, uc.b, uc.c, indexed_count);
        if (MatchesCellLengths(uc, tetragonal_uc))
            ++crystal_count;
    }

    CHECK(crystal_count >= 2);

    // Verify the two best crystal lattices are distinct (low overlap)
    if (result.lattice.size() >= 2) {
        int64_t count_0 = 0, count_1 = 0;
        auto mask_0 = BuildIndexedMask(spots, result.lattice[0], tolerance, count_0);
        auto mask_1 = BuildIndexedMask(spots, result.lattice[1], tolerance, count_1);
        const int64_t overlap = MaskOverlap(mask_0, mask_1);
        const int64_t max_set = std::max(count_0, count_1);
        logger.Info("Top-2 overlap: {} / {}", overlap, max_set);
        CHECK(overlap <= static_cast<int64_t>(0.5f * static_cast<float>(max_set)));
    }
} */


#endif

namespace {
    // A frame of pure noise: reciprocal-space vectors of random direction and length between 50 A
    // and 2.5 A, with no periodicity of any kind in them. The seed is fixed, so the frame is the
    // same on every run and on every machine - a crash that needs a particular junk frame is no use
    // as a regression test if the frame is drawn afresh each time. With a plane normal given, the
    // cloud is flattened onto the plane through the origin, which is what a degenerate net looks
    // like to the indexer.
    std::vector<Coord> NoiseCloud(size_t count, uint32_t seed,
                                  const std::optional<Coord> &plane_normal = {}) {
        // The engine alone, scaled by hand: std::uniform_real_distribution is not specified to give
        // the same numbers in every standard library, and a fixed frame is the whole point here.
        std::mt19937 rng(seed);
        const auto draw = [&rng]() {
            return 2.0f * static_cast<float>(rng()) / static_cast<float>(std::mt19937::max()) - 1.0f;
        };

        const Coord n = plane_normal ? plane_normal->Normalize() : Coord(0, 0, 0);

        std::vector<Coord> vec;
        while (vec.size() < count) {
            Coord q(draw(), draw(), draw());
            if (plane_normal)
                q = q - n * (q * n);
            const float len = q.Length();
            if ((len < 0.05f) || (len > 1.0f))   // a ball, not the cube the three draws fill
                continue;
            vec.push_back(q * 0.4f);             // |q| in [0.02, 0.4] 1/A
        }
        return vec;
    }

    std::vector<IndexingAlgorithmEnum> FFTAlgorithms() {
        std::vector<IndexingAlgorithmEnum> ret;
#ifdef JFJOCH_USE_CUDA
        ret.push_back(IndexingAlgorithmEnum::FFT);
#endif
#ifdef JFJOCH_USE_FFTW
        ret.push_back(IndexingAlgorithmEnum::FFTW);
#endif
        return ret;
    }

    std::unique_ptr<Indexer> MakeFFTIndexer(IndexingAlgorithmEnum algorithm) {
        DiffractionExperiment experiment;
        IndexingSettings settings;
        settings.Algorithm(algorithm).FFT_MaxUnitCell_A(250.0).FFT_HighResolution_A(2 * M_PI / 3.0);
        experiment.ImportIndexingSettings(settings);
        return CreateIndexer(experiment);
    }

    // Fraction of the cloud that sits on the lattice, the same test the indexer's own scoring uses.
    float IndexedFraction(const std::vector<Coord> &spots, const CrystalLattice &lattice, float tolerance) {
        const float tol_sq = tolerance * tolerance;
        size_t indexed = 0;
        for (const auto &q: spots) {
            const float h = q * lattice.Vec0() - std::round(q * lattice.Vec0());
            const float k = q * lattice.Vec1() - std::round(q * lattice.Vec1());
            const float l = q * lattice.Vec2() - std::round(q * lattice.Vec2());
            if (h * h + k * k + l * l < tol_sq)
                indexed++;
        }
        return static_cast<float>(indexed) / static_cast<float>(spots.size());
    }

    // The lattice of the FFTIndexer test above, as spots.
    std::vector<Coord> LatticeSpots(const CrystalLattice &cl) {
        std::vector<Coord> vec;
        for (int h = -2; h < 10; h++)
            for (int k = -5; k < 10; k++)
                for (int l = -3; l < 10; l++)
                    vec.push_back(h * cl.Astar() + k * cl.Bstar() + l * cl.Cstar());
        return vec;
    }
}

TEST_CASE("FFTIndexer_NoiseFrame", "[Indexing]") {
    // The frame that works the indexer hardest is not a clean crystal but junk. A real lattice
    // prunes the FFT shortlist to four or five distinct directions; noise leaves dozens, which is
    // what drives the length sort and the degenerate-plane fallback - where a production broker
    // segfaulted. The run has to come back, and what it brings back has to be recognisably nothing.
    const auto vec = NoiseCloud(1500, 20260917);

    Logger logger("FFTIndexer_NoiseFrame");

    for (const auto algorithm: FFTAlgorithms()) {
        INFO("algorithm " << static_cast<int>(algorithm));
        auto indexer = MakeFFTIndexer(algorithm);
        REQUIRE(indexer);

        const auto result = indexer->Run(vec);
        logger.Info("algorithm {}: {} candidate lattices from noise", static_cast<int>(algorithm),
                    result.lattice.size());

        // Candidates off a junk frame are spurious by construction, and the indexer is allowed to
        // offer them - the caller scores them. What must not happen is one that looks like a real
        // crystal: a true lattice takes most of its frame, and these take about 1%.
        for (const auto &lattice: result.lattice)
            CHECK(IndexedFraction(vec, lattice, 0.1f) < 0.05f);
    }
}

TEST_CASE("FFTIndexer_CoplanarNoiseFrameLeavesTheIndexerUsable", "[Indexing]") {
    // Noise confined to one plane: a shortlist that is coplanar cannot close a cell whatever is done
    // with it, so the indexer goes looking for the missing row inside a 3 deg cap (SearchCap), which
    // swaps the whole search grid out and back. The indexer is not thrown away afterwards - the
    // broker's pool hands the same object the next image - so the frame after a junk one must index
    // exactly as it would have on its own.
    const auto noise = NoiseCloud(1500, 20260918, Coord(0.3f, -0.5f, 0.81f));

    const UnitCell uc(39, 45, 78, 90, 90, 90);
    const auto lattice_spots = LatticeSpots(CrystalLattice(uc));
    Logger logger("FFTIndexer_CoplanarNoiseFrame");

    for (const auto algorithm: FFTAlgorithms()) {
        INFO("algorithm " << static_cast<int>(algorithm));
        auto indexer = MakeFFTIndexer(algorithm);
        REQUIRE(indexer);

        const auto noise_result = indexer->Run(noise);
        logger.Info("algorithm {}: {} candidate lattices from coplanar noise",
                    static_cast<int>(algorithm), noise_result.lattice.size());
        for (const auto &lattice: noise_result.lattice)
            CHECK(IndexedFraction(noise, lattice, 0.1f) < 0.05f);

        const auto result = indexer->Run(lattice_spots);
        REQUIRE(result.lattice.size() == 1);
        const auto found = result.lattice[0].GetUnitCell();
        std::array<float, 3> lengths = {static_cast<float>(found.a), static_cast<float>(found.b),
                                        static_cast<float>(found.c)};
        std::sort(lengths.begin(), lengths.end());
        CHECK(lengths[0] == Catch::Approx(uc.a).epsilon(0.01));
        CHECK(lengths[1] == Catch::Approx(uc.b).epsilon(0.01));
        CHECK(lengths[2] == Catch::Approx(uc.c).epsilon(0.01));
    }
}

TEST_CASE("FFTIndexer_ManyNoiseFrames", "[Indexing]") {
    // A grid scan: frame after frame of junk, sparse or dense, through one pooled indexer. Noise
    // leaves dozens of shortlist directions whose lengths fall in the same FFT bins, so the length
    // sort sees many exact ties. When that sort recomputed Length() inside its comparator, an LTO
    // build contracted it into FMAs differently at different inlined sites, a tie flipped by one ulp
    // between comparisons, and std::sort ran off its index array - the broker SEGV. Only about one
    // noise frame in a thousand did it; frame 1632 crashed the rc.172 LTO build
    // (-march=x86-64-v3 -flto=auto) on the GPU FFT path. A non-LTO build never crashed: there Length()
    // stays out of line and rounds the same way at every call.
    constexpr int FIRST_FRAME = 1628, LAST_FRAME = 1636;

    for (const auto algorithm: FFTAlgorithms()) {
        INFO("algorithm " << static_cast<int>(algorithm));
        auto indexer = MakeFFTIndexer(algorithm);
        REQUIRE(indexer);

        size_t lattices = 0;
        for (int frame = FIRST_FRAME; frame <= LAST_FRAME; frame++)
            lattices += indexer->Run(NoiseCloud(20 + (frame * 37) % 1500, 20260923 + frame)).lattice.size();
        SUCCEED(lattices << " candidate lattices from noise frames");
    }
}

TEST_CASE("FitSupercellProbe", "[SupercellProbe]") {
    // A class whose intensity is 3 at any rocking plus 40 times its partiality: the fit separates the
    // two, and the part that rocks is known to the noise the scatter about the line allows.
    SupercellProbeClass c;
    std::mt19937 rng(1);
    std::normal_distribution<double> noise(0.0, 1.0);
    for (int i = 0; i < 20000; i++) {
        const double p = (i % 100) / 100.0;
        const double I = 3.0 + 40.0 * p + noise(rng);
        c.n++;
        c.sum_i += I;
        c.sum_i_over_sigma += I;
        c.sum_p += p;
        c.sum_pp += p * p;
        c.sum_pi += p * I;
        c.sum_ii += I * I;
    }
    const SupercellProbeFit f = FitSupercellProbe(c);
    CHECK(f.b == Catch::Approx(40.0).margin(0.1));
    CHECK(f.a == Catch::Approx(3.0).margin(0.05));
    CHECK(f.b_se == Catch::Approx(1.0 / std::sqrt(20000.0 * (0.9999 / 12.0))).epsilon(0.05));
    CHECK(f.mean_p == Catch::Approx(0.495));

    // Too few reflections to fit a line: nothing.
    CHECK(FitSupercellProbe(SupercellProbeClass{}).b == 0.0);
}

TEST_CASE("DoubledLattice", "[SupercellProbe]") {
    const CrystalLattice prim(Coord(40, 0, 0), Coord(0, 50, 0), Coord(0, 0, 60));
    // Class 0 0 1: the c axis doubled, the others kept.
    const CrystalLattice c2 = DoubledLattice(prim, {0, 0, 1});
    CHECK(c2.Vec0().x == Catch::Approx(40.0));
    CHECK(c2.Vec1().y == Catch::Approx(50.0));
    CHECK(c2.Vec2().z == Catch::Approx(120.0));
    // Class 1 1 0: a centred doubling of the ab plane - twice a, and b plus a - so every vector of the
    // supercell has h + k even, at twice the volume.
    const CrystalLattice ab = DoubledLattice(prim, {1, 1, 0});
    CHECK(std::fabs(ab.CalcVolume()) == Catch::Approx(2.0 * 40 * 50 * 60));
    CHECK(ab.Vec0().x == Catch::Approx(80.0));
    CHECK(ab.Vec1().x == Catch::Approx(40.0));
    CHECK(ab.Vec1().y == Catch::Approx(50.0));
    CHECK(ab.Vec2().z == Catch::Approx(60.0));
}

TEST_CASE("PostIndexingRefinement_CoplanarCandidateIsRejected", "[Indexing]") {
    // A candidate whose three rows lie in one plane passes every length and angle test - three rows
    // 120 deg apart in a plane meet the 30-150 deg window - and every spot along the plane normal
    // indexes as (0,0,0), so the spot count passes too. Such a cell has no reciprocal basis; handed
    // on, it made the online analysis throw "Crystal lattice is coplanar" for the frame instead of
    // leaving the frame unindexed.
    const Coord a(50, 0, 0);
    const Coord b(-25, 25.0f * std::sqrt(3.0f), 0);
    const Coord c(-25, -25.0f * std::sqrt(3.0f), 0);
    REQUIRE(CrystalLattice(a, b, c).VolumeFraction() < MIN_BASIS_VOLUME_FRACTION);

    std::vector<Coord> spots;
    for (int i = 1; i <= 40; i++)
        spots.emplace_back(0.0f, 0.0f, 0.01f * static_cast<float>(i));

    Eigen::MatrixX3<float> cells(3, 3);
    cells << a.x, a.y, a.z,
             b.x, b.y, b.z,
             c.x, c.y, c.z;
    // A score below the refinement's distance floor, so the candidate reaches the filters as given.
    Eigen::VectorX<float> scores(1);
    scores << -10.99f;

    RefineParameters parameters{
        .viable_cell_min_spots = 9,
        .dist_tolerance_vs_reference = 0.05,
        .reference_unit_cell = std::nullopt,
        .min_length_A = 5,
        .max_length_A = 500,
        .min_angle_deg = 30,
        .max_angle_deg = 150,
        .indexing_tolerance = 0.1
    };

    const auto lattices = Refine(spots, spots.size(), cells, scores, parameters);
    CHECK(lattices.empty());
    for (const auto &l: lattices)
        CHECK_NOTHROW(l.GetUBMatrix());
}