Jungfraujoch/tests/RotationIndexerTest.cpp

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

#include <catch2/catch_all.hpp>
#include <iostream>

#include "../image_analysis/RotationIndexer.h"
#include "../image_analysis/bragg_prediction/BraggPrediction.h"

TEST_CASE("RotationIndexer") {
    DiffractionExperiment exp_i;
    exp_i.IncidentEnergy_keV(WVL_1A_IN_KEV)
         .BeamX_pxl(1000)
         .BeamY_pxl(1000)
         .PoniRot1_rad(0.01)
         .PoniRot2_rad(0.02)
         .DetectorDistance_mm(200)
    .ImagesPerTrigger(50);

    IndexingSettings settings;
#ifdef JFJOCH_USE_CUDA
    settings.Algorithm(IndexingAlgorithmEnum::FFT);
#elif JFJOCH_USE_FFTW
    settings.Algorithm(IndexingAlgorithmEnum::FFTW);
#else
    return;
#endif

    settings.RotationIndexing(true).RotationIndexingAngularStride_deg(1.0).RotationIndexingMinAngularRange_deg(30.0);
    exp_i.ImportIndexingSettings(settings);

    // Base lattice (non-pathological)
    CrystalLattice latt_base(40, 50, 80, 90, 90, 90);
    latt_base = latt_base.Multiply(RotMatrix(2.0, Coord(sqrt(3)/3,sqrt(3)/3,sqrt(3)/3)));

    // Rotation axis: around X with 1 deg per image
    GoniometerAxis axis("omega", 0.0f, 1.0f, Coord(1,0,0), std::nullopt);
    exp_i.Goniometer(axis);

    BraggPredictionSettings prediction_settings{
        .high_res_A = 1.3,
        .ewald_dist_cutoff = 0.002
    };

    IndexerThreadPool indexer_thread_pool(exp_i.GetIndexingSettings());
    RotationIndexer indexer(exp_i, indexer_thread_pool);

    BraggPrediction prediction;

    int cnt = 0;

    // Predict reflections for images at 0-30 deg.
    for (int img = 0; img < 50; ++img) {
        std::vector<SpotToSave> spots;
        // For a rotated image, per-image lattice is obtained as Multiply(rot.transpose())
        const float angle_deg = axis.GetAngle_deg(img) + axis.GetWedge_deg() / 2.0f;
        const RotMatrix rot = axis.GetTransformationAngle(angle_deg);
        const CrystalLattice latt_img = latt_base.Multiply(rot.transpose());

        const auto n = prediction.Calc(exp_i, latt_img, prediction_settings);
        for (int i = 0; i < n; ++i) {
            const auto& r = prediction.GetReflections().at(i);
            SpotToSave s{};
            s.x = r.predicted_x;
            s.y = r.predicted_y;
            s.image = img;       // provide image index for rotation-aware refinement
            s.intensity = 1.0f;  // minimal positive value
            s.phi = angle_deg;
            s.ice_ring = false;
            s.indexed = true;
            spots.push_back(s);
        }
        if (indexer.ProcessImage(img, spots).has_value())
            cnt++;
    }

    CHECK(cnt == 20);

    auto ret = indexer.GetLattice();
    REQUIRE(ret.has_value());
    auto uc = ret->lattice.GetUnitCell();
    auto uc_ref = latt_base.GetUnitCell();
    REQUIRE(std::fabs(uc.a - uc_ref.a) < 0.02 );
    REQUIRE(std::fabs(uc.b - uc_ref.b) < 0.02 );
    REQUIRE(std::fabs(uc.c - uc_ref.c) < 0.02 );
    REQUIRE(std::fabs(uc.alpha - uc_ref.alpha) < 0.1);
    REQUIRE(std::fabs(uc.beta - uc_ref.beta) < 0.1);
    REQUIRE(std::fabs(uc.gamma - uc_ref.gamma) < 0.1);
    CHECK(ret->search_result.centering == 'P');
    CHECK(ret->search_result.system == gemmi::CrystalSystem::Orthorhombic);
}

TEST_CASE("RotationIndexer short scan indexes only at the end") {
    DiffractionExperiment exp_i;
    exp_i.IncidentEnergy_keV(WVL_1A_IN_KEV)
         .BeamX_pxl(1000)
         .BeamY_pxl(1000)
         .PoniRot1_rad(0.01)
         .PoniRot2_rad(0.02)
         .DetectorDistance_mm(200)
         .ImagesPerTrigger(6);

    IndexingSettings settings;
#ifdef JFJOCH_USE_CUDA
    settings.Algorithm(IndexingAlgorithmEnum::FFT);
#elif JFJOCH_USE_FFTW
    settings.Algorithm(IndexingAlgorithmEnum::FFTW);
#else
    return;
#endif

    settings.RotationIndexing(true).RotationIndexingAngularStride_deg(1.0).RotationIndexingMinAngularRange_deg(30.0);
    exp_i.ImportIndexingSettings(settings);

    CrystalLattice latt_base(40, 50, 80, 90, 90, 90);
    latt_base = latt_base.Multiply(RotMatrix(2.0, Coord(sqrt(3)/3,sqrt(3)/3,sqrt(3)/3)));

    GoniometerAxis axis("omega", 0.0f, 30.0f, Coord(1,0,0), std::nullopt);
    exp_i.Goniometer(axis);

    BraggPredictionSettings prediction_settings{
        .high_res_A = 1.3,
        .ewald_dist_cutoff = 0.002
    };

    IndexerThreadPool indexer_thread_pool(exp_i.GetIndexingSettings());
    RotationIndexer indexer(exp_i, indexer_thread_pool);

    BraggPrediction prediction;

    for (int img = 0; img < 5; ++img) {
        std::vector<SpotToSave> spots;
        const float angle_deg = axis.GetAngle_deg(img) + axis.GetWedge_deg() / 2.0f;
        const RotMatrix rot = axis.GetTransformationAngle(angle_deg);
        const CrystalLattice latt_img = latt_base.Multiply(rot.transpose());

        const auto n = prediction.Calc(exp_i, latt_img, prediction_settings);
        for (int i = 0; i < n; ++i) {
            const auto& r = prediction.GetReflections().at(i);
            SpotToSave s{};
            s.x = r.predicted_x;
            s.y = r.predicted_y;
            s.image = img;
            s.intensity = 1.0f;
            s.phi = angle_deg;
            s.ice_ring = false;
            s.indexed = true;
            spots.push_back(s);
        }

        CHECK_FALSE(indexer.ProcessImage(img, spots).has_value());
    }

    std::vector<SpotToSave> last_spots;
    const int last_img = 5;
    const float angle_deg = axis.GetAngle_deg(last_img) + axis.GetWedge_deg() / 2.0f;
    const RotMatrix rot = axis.GetTransformationAngle(angle_deg);
    const CrystalLattice latt_img = latt_base.Multiply(rot.transpose());

    const auto n = prediction.Calc(exp_i, latt_img, prediction_settings);
    for (int i = 0; i < n; ++i) {
        const auto& r = prediction.GetReflections().at(i);
        SpotToSave s{};
        s.x = r.predicted_x;
        s.y = r.predicted_y;
        s.image = last_img;
        s.intensity = 1.0f;
        s.phi = angle_deg;
        s.ice_ring = false;
        s.indexed = true;
        last_spots.push_back(s);
    }

    auto ret_last = indexer.ProcessImage(last_img, last_spots);
    REQUIRE(ret_last.has_value());

    auto ret = indexer.GetLattice();
    REQUIRE(ret.has_value());
    auto uc = ret->lattice.GetUnitCell();
    auto uc_ref = latt_base.GetUnitCell();
    REQUIRE(std::fabs(uc.a - uc_ref.a) < 0.02 );
    REQUIRE(std::fabs(uc.b - uc_ref.b) < 0.02 );
    REQUIRE(std::fabs(uc.c - uc_ref.c) < 0.02 );
    REQUIRE(std::fabs(uc.alpha - uc_ref.alpha) < 0.1);
    REQUIRE(std::fabs(uc.beta - uc_ref.beta) < 0.1);
    REQUIRE(std::fabs(uc.gamma - uc_ref.gamma) < 0.1);
    CHECK(ret->search_result.centering == 'P');
    CHECK(ret->search_result.system == gemmi::CrystalSystem::Orthorhombic);
}