Jungfraujoch/tests/XtalOptimizerTest.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/geom_refinement/XtalOptimizer.h"
#include "../image_analysis/bragg_prediction/BraggPrediction.h"

TEST_CASE("XtalOptimizer") {
    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);

    CrystalLattice latt_i(40,40,80,90,90,90);

    BraggPredictionSettings prediction_settings{
    .high_res_A = 1.5,
    .ewald_dist_cutoff = 0.001};
    BraggPrediction prediction;
    auto count = prediction.Calc(exp_i, latt_i, prediction_settings);

    std::vector<SpotToSave> spots;

    for (int i = 0; i < count; ++i) {
        auto refl = prediction.GetReflections().at(i);
        spots.push_back(SpotToSave{refl.predicted_x, refl.predicted_y});
    }

    XtalOptimizerData xtal_opt{};
    xtal_opt.max_time = 30.0;
    xtal_opt.latt = CrystalLattice(40.2,39.4,80.2, 90,91, 89);
    xtal_opt.geom.BeamX_pxl(1010).BeamY_pxl(995).DetectorDistance_mm(200)
    .PoniRot1_rad(0.01).PoniRot2_rad(0.02);
    xtal_opt.crystal_system = gemmi::CrystalSystem::Triclinic;

    auto start = std::chrono::high_resolution_clock::now();
    REQUIRE(XtalOptimizer(xtal_opt, {spots}));
    auto end = std::chrono::high_resolution_clock::now();
    std::cout << "XtalOptimizer took " << std::chrono::duration_cast<std::chrono::microseconds>(end - start).count()
              << " microseconds" << std::endl;

    auto uc_i = latt_i.GetUnitCell();
    auto uc_o = xtal_opt.latt.GetUnitCell();

    std::cout << "Beam center: " << xtal_opt.geom.GetBeamX_pxl() << " " << xtal_opt.geom.GetBeamY_pxl() << std::endl;
    std::cout << "Unit cell: " << uc_o.a << " " << uc_o.b << " " << uc_o.c << std::endl;
    CHECK(fabsf(xtal_opt.geom.GetBeamX_pxl() - exp_i.GetBeamX_pxl()) < 0.05);
    CHECK(fabsf(xtal_opt.geom.GetBeamY_pxl() - exp_i.GetBeamY_pxl()) < 0.05);

    CHECK(fabsf(uc_i.a - uc_o.a) < 0.1);
    CHECK(fabsf(uc_i.b - uc_o.b) < 0.1);
    CHECK(fabsf(uc_i.c - uc_o.c) < 0.2);
    CHECK(fabsf(uc_i.alpha - uc_o.alpha) < 0.1);
    CHECK(fabsf(uc_i.beta - uc_o.beta) < 0.1);
    CHECK(fabsf(uc_i.gamma - uc_o.gamma) < 0.1);
}


TEST_CASE("XtalOptimizer_NoBeamCenter") {
    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);

    CrystalLattice latt_i(40,50,80,90,95,90);

    BraggPredictionSettings prediction_settings{
        .high_res_A = 1.5,
        .ewald_dist_cutoff = 0.001};
    BraggPrediction prediction;
    auto count = prediction.Calc(exp_i, latt_i, prediction_settings);

    std::vector<SpotToSave> spots;

    for (int i = 0; i < count; ++i) {
        auto refl = prediction.GetReflections().at(i);
        spots.push_back(SpotToSave{refl.predicted_x, refl.predicted_y});
    }

    XtalOptimizerData xtal_opt{};
    xtal_opt.latt = CrystalLattice(40.2,49.4,80.2, 90,94, 89);
    xtal_opt.geom.BeamX_pxl(999.8).BeamY_pxl(1000.2).DetectorDistance_mm(200)
    .PoniRot1_rad(0.01).PoniRot2_rad(0.02);
    xtal_opt.crystal_system = gemmi::CrystalSystem::Monoclinic;
    xtal_opt.refine_beam_center = false;
    xtal_opt.max_time = 30.0;

    auto start = std::chrono::high_resolution_clock::now();
    REQUIRE(XtalOptimizer(xtal_opt, {spots}));
    auto end = std::chrono::high_resolution_clock::now();
    std::cout << "XtalOptimizer took " << std::chrono::duration_cast<std::chrono::microseconds>(end - start).count()
              << " microseconds" << std::endl;

    auto uc_i = latt_i.GetUnitCell();
    auto uc_o = xtal_opt.latt.GetUnitCell();

    std::cout << "Beam center: " << xtal_opt.geom.GetBeamX_pxl() << " " << xtal_opt.geom.GetBeamY_pxl() << std::endl;
    std::cout << "Unit cell: " << uc_o.a << " " << uc_o.b << " " << uc_o.c << std::endl;
    CHECK(fabsf(xtal_opt.geom.GetBeamX_pxl() - 999.8) < 0.01);
    CHECK(fabsf(xtal_opt.geom.GetBeamY_pxl() - 1000.2) < 0.01);

    CHECK(fabsf(uc_i.a - uc_o.a) < 0.1);
    CHECK(fabsf(uc_i.b - uc_o.b) < 0.1);
    CHECK(fabsf(uc_i.c - uc_o.c) < 0.2);
    CHECK(fabsf(uc_i.alpha - uc_o.alpha) < 0.1);
    CHECK(fabsf(uc_i.beta - uc_o.beta) < 0.1);
    CHECK(fabsf(uc_i.gamma - uc_o.gamma) < 0.1);
}

TEST_CASE("XtalOptimizer_orthorombic") {
    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);

    CrystalLattice latt_i(40,50,80,90,90,90);

    BraggPredictionSettings prediction_settings{
        .high_res_A = 1.5,
        .ewald_dist_cutoff = 0.001};
    BraggPrediction prediction;
    auto count = prediction.Calc(exp_i, latt_i, prediction_settings);

    std::vector<SpotToSave> spots;

    for (int i = 0; i < count; ++i) {
        auto refl = prediction.GetReflections().at(i);
        spots.push_back(SpotToSave{refl.predicted_x, refl.predicted_y});
    }
    XtalOptimizerData xtal_opt{};
    xtal_opt.latt = CrystalLattice(40.2,49.6,80.3, 90,91, 89);
    xtal_opt.geom.BeamX_pxl(1005).BeamY_pxl(997).DetectorDistance_mm(200)
            .PoniRot1_rad(0.01).PoniRot2_rad(0.02);
    xtal_opt.max_time = 30.0;
    xtal_opt.crystal_system = gemmi::CrystalSystem::Orthorhombic;
    auto start = std::chrono::high_resolution_clock::now();
    REQUIRE(XtalOptimizer(xtal_opt, {spots}));
    auto end = std::chrono::high_resolution_clock::now();
    std::cout << "XtalOptimizer took " << std::chrono::duration_cast<std::chrono::microseconds>(end - start).count()
              << " microseconds" << std::endl;

    auto uc_i = latt_i.GetUnitCell();
    auto uc_o = xtal_opt.latt.GetUnitCell();

    std::cout << "Beam center: " << xtal_opt.geom.GetBeamX_pxl() << " " << xtal_opt.geom.GetBeamY_pxl() << std::endl;
    std::cout << "Unit cell: " << uc_o.a << " " << uc_o.b << " " << uc_o.c << std::endl;
    CHECK(fabsf(xtal_opt.geom.GetBeamX_pxl() - exp_i.GetBeamX_pxl()) < 0.1);
    CHECK(fabsf(xtal_opt.geom.GetBeamY_pxl() - exp_i.GetBeamY_pxl()) < 0.1);

    CHECK(fabsf(uc_i.a - uc_o.a) < 0.1);
    CHECK(fabsf(uc_i.b - uc_o.b) < 0.1);
    CHECK(fabsf(uc_i.c - uc_o.c) < 0.2);
    CHECK(fabs(uc_o.alpha - 90) < 0.02);
    CHECK(fabs(uc_o.beta - 90) < 0.02);
    CHECK(fabs(uc_o.gamma - 90) < 0.02);
}

TEST_CASE("XtalOptimizer_triclinic") {
    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);

    CrystalLattice latt_i(40,55,120,95,97,100);

    BraggPredictionSettings prediction_settings{
        .high_res_A = 1.5,
        .ewald_dist_cutoff = 0.001,
    };
    BraggPrediction prediction;
    auto count = prediction.Calc(exp_i, latt_i, prediction_settings);

    std::vector<SpotToSave> spots;

    for (int i = 0; i < count; ++i) {
        auto refl = prediction.GetReflections().at(i);
        spots.push_back(SpotToSave{refl.predicted_x, refl.predicted_y});
    }
    XtalOptimizerData xtal_opt{};
    xtal_opt.latt = CrystalLattice(40.1,54.9,121, 95,97, 99.5);
    xtal_opt.geom.BeamX_pxl(997).BeamY_pxl(1005).DetectorDistance_mm(200)
            .PoniRot1_rad(0.01).PoniRot2_rad(0.02);
    xtal_opt.crystal_system = gemmi::CrystalSystem::Triclinic;
    xtal_opt.max_time = 36.0;
    auto start = std::chrono::high_resolution_clock::now();
    REQUIRE(XtalOptimizer(xtal_opt, {spots}));
    auto end = std::chrono::high_resolution_clock::now();
    std::cout << "XtalOptimizer took " << std::chrono::duration_cast<std::chrono::microseconds>(end - start).count()
              << " microseconds" << std::endl;

    auto uc_i = latt_i.GetUnitCell();
    auto uc_o = xtal_opt.latt.GetUnitCell();

    std::cout << "Beam center: " << xtal_opt.geom.GetBeamX_pxl() << " " << xtal_opt.geom.GetBeamY_pxl() << std::endl;
    std::cout << "Unit cell: " << uc_o.a << " " << uc_o.b << " " << uc_o.c << std::endl;
    CHECK(fabsf(xtal_opt.geom.GetBeamX_pxl() - exp_i.GetBeamX_pxl()) < 0.2);
    CHECK(fabsf(xtal_opt.geom.GetBeamY_pxl() - exp_i.GetBeamY_pxl()) < 0.2);

    CHECK(fabsf(uc_i.a - uc_o.a) < 0.1);
    CHECK(fabsf(uc_i.b - uc_o.b) < 0.1);
    CHECK(fabsf(uc_i.c - uc_o.c) < 0.5);
    CHECK(fabsf(uc_i.alpha - uc_o.alpha) < 0.1);
    CHECK(fabsf(uc_i.beta - uc_o.beta) < 0.1);
    CHECK(fabsf(uc_i.gamma - uc_o.gamma) < 0.1);
}


TEST_CASE("XtalOptimizer_tetragonal") {
    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);

    CrystalLattice latt_i(40,40,80,90,90,90);

    BraggPredictionSettings prediction_settings{
        .high_res_A = 1.5,
        .ewald_dist_cutoff = 0.001};
    BraggPrediction prediction;
    auto count = prediction.Calc(exp_i, latt_i, prediction_settings);

    std::vector<SpotToSave> spots;

    for (int i = 0; i < count; ++i) {
        auto refl = prediction.GetReflections().at(i);
        spots.push_back(SpotToSave{refl.predicted_x, refl.predicted_y});
    }

    XtalOptimizerData xtal_opt{};
    xtal_opt.latt = CrystalLattice(40.6,39.3,80.5, 90,91, 89);
    xtal_opt.geom.BeamX_pxl(1010).BeamY_pxl(995).DetectorDistance_mm(200)
            .PoniRot1_rad(0.01).PoniRot2_rad(0.02);
    xtal_opt.crystal_system = gemmi::CrystalSystem::Tetragonal;
    xtal_opt.max_time = 30.0;
    auto start = std::chrono::high_resolution_clock::now();
    REQUIRE(XtalOptimizer(xtal_opt, {spots}));
    auto end = std::chrono::high_resolution_clock::now();
    std::cout << "XtalOptimizer took " << std::chrono::duration_cast<std::chrono::microseconds>(end - start).count()
              << " microseconds" << std::endl;

    auto uc_i = latt_i.GetUnitCell();
    auto uc_o = xtal_opt.latt.GetUnitCell();

    std::cout << "Beam center: " << xtal_opt.geom.GetBeamX_pxl() << " " << xtal_opt.geom.GetBeamY_pxl() << std::endl;
    std::cout << "Unit cell: " << uc_o.a << " " << uc_o.b << " " << uc_o.c << std::endl;
    CHECK(fabsf(xtal_opt.geom.GetBeamX_pxl() - exp_i.GetBeamX_pxl()) < 0.1);
    CHECK(fabsf(xtal_opt.geom.GetBeamY_pxl() - exp_i.GetBeamY_pxl()) < 0.1);

    CHECK(fabsf(uc_i.a - uc_o.a) < 0.1);
    CHECK(fabsf(uc_i.b - uc_o.b) < 0.1);
    CHECK(fabsf(uc_i.c - uc_o.c) < 0.5);
    CHECK(fabs(uc_o.alpha - 90) < 0.02);
    CHECK(fabs(uc_o.beta - 90) < 0.02);
    CHECK(fabs(uc_o.gamma - 90) < 0.02);
}

TEST_CASE("XtalOptimizer_hexagonal") {
    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);

    CrystalLattice latt_i(40,40,70,90,90,120);

    BraggPredictionSettings prediction_settings{
        .high_res_A = 1.5,
        .ewald_dist_cutoff = 0.001};
    BraggPrediction prediction;
    auto count = prediction.Calc(exp_i, latt_i, prediction_settings);

    std::vector<SpotToSave> spots;

    for (int i = 0; i < count; ++i) {
        auto refl = prediction.GetReflections().at(i);
        spots.push_back(SpotToSave{refl.predicted_x, refl.predicted_y});
    }

    XtalOptimizerData xtal_opt{};
    xtal_opt.latt = CrystalLattice(39.5,39.8,70.1, 90,90, 119.5);
    xtal_opt.geom.BeamX_pxl(1007).BeamY_pxl(990).DetectorDistance_mm(200)
            .PoniRot1_rad(0.01).PoniRot2_rad(0.02);
    xtal_opt.crystal_system = gemmi::CrystalSystem::Hexagonal;
    xtal_opt.max_time = 60.0;
    auto start = std::chrono::high_resolution_clock::now();
    bool ret = XtalOptimizer(xtal_opt, {spots});
    auto end = std::chrono::high_resolution_clock::now();
    std::cout << "XtalOptimizer took " << std::chrono::duration_cast<std::chrono::microseconds>(end - start).count()
              << " microseconds" << std::endl;
    REQUIRE(ret);
    auto uc_i = latt_i.GetUnitCell();
    auto uc_o = xtal_opt.latt.GetUnitCell();

    std::cout << "Beam center: " << xtal_opt.geom.GetBeamX_pxl() << " " << xtal_opt.geom.GetBeamY_pxl() << std::endl;
    std::cout << "Unit cell: " << uc_o.a << " " << uc_o.b << " " << uc_o.c << " " << uc_o.alpha << " " << uc_o.beta
              << " " << uc_o.gamma << std::endl;

    CHECK(fabsf(xtal_opt.geom.GetBeamX_pxl() - exp_i.GetBeamX_pxl()) < 0.1);
    CHECK(fabsf(xtal_opt.geom.GetBeamY_pxl() - exp_i.GetBeamY_pxl()) < 0.1);

    CHECK(fabsf(uc_i.a - uc_o.a) < 0.1);
    CHECK(fabsf(uc_i.b - uc_o.b) < 0.1);
    CHECK(fabsf(uc_i.c - uc_o.c) < 0.2);
    CHECK(fabs(uc_o.alpha - 90) < 0.02);
    CHECK(fabs(uc_o.beta - 90) < 0.01);
    CHECK(fabs(uc_o.gamma - 120) < 0.01);
}

TEST_CASE("XtalOptimizer_hexagonal_unconstrained") {
    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);

    CrystalLattice latt_i(40,40,70,90,90,120);

    BraggPredictionSettings prediction_settings{
        .high_res_A = 1.5,
        .ewald_dist_cutoff = 0.001};
    BraggPrediction prediction;
    auto count = prediction.Calc(exp_i, latt_i, prediction_settings);

    std::vector<SpotToSave> spots;

    for (int i = 0; i < count; ++i) {
        auto refl = prediction.GetReflections().at(i);
        spots.push_back(SpotToSave{refl.predicted_x, refl.predicted_y});
    }

    XtalOptimizerData xtal_opt{};
    xtal_opt.latt = CrystalLattice(39.9,39.8,70.1, 90,90, 120);
    xtal_opt.geom.BeamX_pxl(1002).BeamY_pxl(998).DetectorDistance_mm(200)
            .PoniRot1_rad(0.01).PoniRot2_rad(0.02);
    xtal_opt.crystal_system = gemmi::CrystalSystem::Triclinic;
    xtal_opt.max_time = 30.0;
    auto start = std::chrono::high_resolution_clock::now();
    REQUIRE(XtalOptimizer(xtal_opt, {spots}));
    auto end = std::chrono::high_resolution_clock::now();
    std::cout << "XtalOptimizer took " << std::chrono::duration_cast<std::chrono::microseconds>(end - start).count()
              << " microseconds" << std::endl;

    auto uc_i = latt_i.GetUnitCell();
    auto uc_o = xtal_opt.latt.GetUnitCell();

    std::cout << "Beam center: " << xtal_opt.geom.GetBeamX_pxl() << " " << xtal_opt.geom.GetBeamY_pxl() << std::endl;
    std::cout << "Unit cell: " << uc_o.a << " " << uc_o.b << " " << uc_o.c << " " << uc_o.alpha << " " << uc_o.beta
              << " " << uc_o.gamma << std::endl;

    CHECK(fabsf(xtal_opt.geom.GetBeamX_pxl() - exp_i.GetBeamX_pxl()) < 0.3);
    CHECK(fabsf(xtal_opt.geom.GetBeamY_pxl() - exp_i.GetBeamY_pxl()) < 0.3);

    CHECK(fabsf(uc_i.a - uc_o.a) < 0.1);
    CHECK(fabsf(uc_i.b - uc_o.b) < 0.1);
    CHECK(fabsf(uc_i.c - uc_o.c) < 0.2);
    CHECK(fabs(uc_o.alpha - 90) < 0.1);
    CHECK(fabs(uc_o.beta - 90) < 0.1);
    CHECK(fabs(uc_o.gamma - 120) < 0.1);
}

TEST_CASE("XtalOptimizer_cubic") {
    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);

    CrystalLattice latt_i(Coord(40,0,0),
                          Coord(0, 40 / sqrt(2), -40 / sqrt(2)),
                          Coord(0, 40 / sqrt(2), 40 / sqrt(2)));
    auto uc_i = latt_i.GetUnitCell();

    BraggPredictionSettings prediction_settings{
        .high_res_A = 1.5,
        .ewald_dist_cutoff = 0.001};
    BraggPrediction prediction;
    auto count = prediction.Calc(exp_i, latt_i, prediction_settings);

    std::vector<SpotToSave> spots;

    for (int i = 0; i < count; ++i) {
        auto refl = prediction.GetReflections().at(i);
        spots.push_back(SpotToSave{refl.predicted_x, refl.predicted_y});
    }

    XtalOptimizerData xtal_opt{};
    xtal_opt.latt = CrystalLattice(Coord(39,0,0),
                                   Coord(0, 39.5 / sqrt(2), -40.5 / sqrt(2)),
                                   Coord(0, 39.2 / sqrt(2), 39.7 / sqrt(2)));
    xtal_opt.geom.BeamX_pxl(1007).BeamY_pxl(990).DetectorDistance_mm(200)
            .PoniRot1_rad(0.01).PoniRot2_rad(0.02);
    xtal_opt.crystal_system = gemmi::CrystalSystem::Cubic;
    xtal_opt.max_time = 30.0;
    auto start = std::chrono::high_resolution_clock::now();
    REQUIRE(XtalOptimizer(xtal_opt, {spots}));
    auto end = std::chrono::high_resolution_clock::now();
    std::cout << "XtalOptimizer took " << std::chrono::duration_cast<std::chrono::microseconds>(end - start).count()
              << " microseconds" << std::endl;

    auto uc_o = xtal_opt.latt.GetUnitCell();

    std::cout << "Beam center: " << xtal_opt.geom.GetBeamX_pxl() << " " << xtal_opt.geom.GetBeamY_pxl() << std::endl;
    std::cout << "Unit cell: " << uc_o.a << " " << uc_o.b << " " << uc_o.c << std::endl;
    CHECK(fabsf(xtal_opt.geom.GetBeamX_pxl() - exp_i.GetBeamX_pxl()) < 0.1);
    CHECK(fabsf(xtal_opt.geom.GetBeamY_pxl() - exp_i.GetBeamY_pxl()) < 0.1);

    CHECK(fabsf(uc_i.a - uc_o.a) < 0.1);
    CHECK(fabsf(uc_i.b - uc_o.b) < 0.1);
    CHECK(fabsf(uc_i.c - uc_o.c) < 0.2);
    CHECK(fabs(uc_o.alpha - 90) < 0.02);
    CHECK(fabs(uc_o.beta - 90) < 0.02);
    CHECK(fabs(uc_o.gamma - 90) < 0.02);
}

TEST_CASE("XtalOptimizer_monoclinic") {
    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);

    CrystalLattice latt_i(50,60,70,90,115,90);

    auto uc_i = latt_i.GetUnitCell();

    BraggPredictionSettings prediction_settings{
        .high_res_A = 1.5,
        .ewald_dist_cutoff = 0.001
    };
    BraggPrediction prediction;
    auto count = prediction.Calc(exp_i, latt_i, prediction_settings);

    std::vector<SpotToSave> spots;

    for (int i = 0; i < count; ++i) {
        auto refl = prediction.GetReflections().at(i);
        spots.push_back(SpotToSave{refl.predicted_x, refl.predicted_y});
    }

    XtalOptimizerData xtal_opt{};
    xtal_opt.latt = CrystalLattice(49.5, 60.5, 69.8, 90, 113.5, 90);
    xtal_opt.geom.BeamX_pxl(1007).BeamY_pxl(990).DetectorDistance_mm(200)
            .PoniRot1_rad(0.01).PoniRot2_rad(0.02);
    xtal_opt.crystal_system = gemmi::CrystalSystem::Monoclinic;
    xtal_opt.max_time = 30.0;
    auto start = std::chrono::high_resolution_clock::now();
    REQUIRE(XtalOptimizer(xtal_opt, {spots}));
    auto end = std::chrono::high_resolution_clock::now();
    std::cout << "XtalOptimizer took " << std::chrono::duration_cast<std::chrono::microseconds>(end - start).count()
              << " microseconds" << std::endl;

    auto uc_o = xtal_opt.latt.GetUnitCell();

    std::cout << "Beam center: " << xtal_opt.geom.GetBeamX_pxl() << " " << xtal_opt.geom.GetBeamY_pxl() << std::endl;
    std::cout << "Unit cell: " << uc_o.a << " " << uc_o.b << " " << uc_o.c << " " << uc_o.alpha << " " << uc_o.beta << " " << uc_o.gamma << std::endl;
    CHECK(fabsf(xtal_opt.geom.GetBeamX_pxl() - exp_i.GetBeamX_pxl()) < 0.2);
    CHECK(fabsf(xtal_opt.geom.GetBeamY_pxl() - exp_i.GetBeamY_pxl()) < 0.2);

    CHECK(fabsf(uc_i.a - uc_o.a) < 0.1);
    CHECK(fabsf(uc_i.b - uc_o.b) < 0.1);
    CHECK(fabsf(uc_i.c - uc_o.c) < 0.2);
    CHECK(fabs(uc_o.alpha - 90) < 0.05);
    CHECK(fabs(uc_o.beta - uc_i.beta) < 0.05);
    CHECK(fabs(uc_o.gamma - 90) < 0.05);
}

TEST_CASE("XtalOptimizer_rotation") {
    // Geometry
    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);

    // Base lattice (non-pathological)
    CrystalLattice latt_base(40, 50, 80, 90, 95, 90);
    auto uc_ref = latt_base.GetUnitCell();

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

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

    size_t nimages = 10;

    std::vector<std::vector<SpotToSave>> spots(nimages);
    BraggPrediction prediction;

    // Predict reflections for images at 0-30 deg.
    for (int img = 0; img < nimages; ++img) {
        // 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.phi = angle_deg;
            s.intensity = 1.0f;  // minimal positive value
            s.ice_ring = false;
            s.indexed = true;
            spots[img].push_back(s);
        }
    }

    // Seed slightly perturbed geometry and lattice; provide rotation axis for refinement.
    // The beam-centre component PARALLEL to the spindle (here X, the rotation axis) is a gauge-weak
    // direction that XtalOptimizer now deliberately restrains toward the header rather than refining
    // freely (see the BeamComponentPrior in XtalOptimizer.cpp). So only the perpendicular component
    // (Y) is seeded off-truth to exercise beam-centre recovery; X is seeded at its true value.
    // The distance is seeded at its true value too: XtalOptimizer does not refine it - the rotation
    // post-refinement fits it globally instead, where the cell is held fixed and the two are no longer
    // degenerate - so a distance seeded off-truth here would simply be absorbed by the cell.
    XtalOptimizerData xtal_opt{};
    xtal_opt.latt = CrystalLattice(39.7f, 50.6f, 79.6f, 90.0f, 94.5f, 90.5f);
    xtal_opt.geom.BeamX_pxl(1000).BeamY_pxl(997).DetectorDistance_mm(200.0)
            .PoniRot1_rad(0.01).PoniRot2_rad(0.02);
    xtal_opt.crystal_system = gemmi::CrystalSystem::Monoclinic;
    xtal_opt.axis = axis;
    xtal_opt.min_spots = 200;
    xtal_opt.refine_beam_center = true;
    xtal_opt.refine_detector_angles = false;
    xtal_opt.max_time = 30.0;

    auto t0 = std::chrono::high_resolution_clock::now();
    REQUIRE(XtalOptimizer(xtal_opt, spots));
    auto t1 = std::chrono::high_resolution_clock::now();
    std::cout << "XtalOptimizer (rotation 4 images) took "
              << std::chrono::duration_cast<std::chrono::microseconds>(t1 - t0).count()
              << " microseconds" << std::endl;

    const auto uc_out = xtal_opt.latt.GetUnitCell();

    // Geometry checks
    CHECK(fabsf(xtal_opt.geom.GetBeamX_pxl() - exp_i.GetBeamX_pxl()) < 0.2f);
    CHECK(fabsf(xtal_opt.geom.GetBeamY_pxl() - exp_i.GetBeamY_pxl()) < 0.2f);

    // Lattice checks
    CHECK(fabsf(uc_ref.a - uc_out.a) < 0.2f);
    CHECK(fabsf(uc_ref.b - uc_out.b) < 0.2f);
    CHECK(fabsf(uc_ref.c - uc_out.c) < 0.4f);
    CHECK(fabsf(uc_ref.alpha - uc_out.alpha) < 0.2f);
    CHECK(fabsf(uc_ref.beta - uc_out.beta) < 0.2f);
    CHECK(fabsf(uc_ref.gamma - uc_out.gamma) < 0.2f);
}

TEST_CASE("XtalOptimizer_refine_rotation_axis") {
    // Geometry
    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);

    // Base lattice (non-pathological)
    CrystalLattice latt_base(40, 50, 80, 90, 95, 90);
    auto uc_ref = latt_base.GetUnitCell();

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

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


    BraggPrediction prediction;

    const size_t nimages = 10;
    std::vector<std::vector<SpotToSave>> spots(nimages);

    // Predict reflections for images at 0-30 deg.
    for (int img = 0; img < nimages; ++img) {
        // 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.at(img).push_back(s);
        }
    }

    // Seed slightly perturbed geometry and lattice; provide rotation axis for refinement
    XtalOptimizerData xtal_opt{};
    xtal_opt.latt = CrystalLattice(39.7f, 50.6f, 79.6f, 90.0f, 94.5f, 90.5f);
    xtal_opt.geom.BeamX_pxl(1003).BeamY_pxl(997).DetectorDistance_mm(200.0)
            .PoniRot1_rad(0.01).PoniRot2_rad(0.02);
    xtal_opt.crystal_system = gemmi::CrystalSystem::Monoclinic;
    xtal_opt.axis = GoniometerAxis("omega", 0.0f, 3.0f,
        Coord(0.8, 0.05, 0.05).Normalize(),
        std::nullopt);
    xtal_opt.min_spots = 200;
    xtal_opt.refine_beam_center = true;
    xtal_opt.refine_detector_angles = false;
    xtal_opt.refine_rotation_axis = true;
    xtal_opt.max_time = 30.0;

    auto t0 = std::chrono::high_resolution_clock::now();
    REQUIRE(XtalOptimizer(xtal_opt, spots));
    auto t1 = std::chrono::high_resolution_clock::now();
    std::cout << "XtalOptimizer (rotation 4 images) took "
              << std::chrono::duration_cast<std::chrono::microseconds>(t1 - t0).count()
              << " microseconds" << std::endl;

    const auto uc_out = xtal_opt.latt.GetUnitCell();

    // Geometry checks
    CHECK(fabsf(xtal_opt.geom.GetBeamX_pxl() - exp_i.GetBeamX_pxl()) < 0.2f);
    CHECK(fabsf(xtal_opt.geom.GetBeamY_pxl() - exp_i.GetBeamY_pxl()) < 0.2f);

    // Lattice checks
    CHECK(fabsf(uc_ref.a - uc_out.a) < 0.2f);
    CHECK(fabsf(uc_ref.b - uc_out.b) < 0.2f);
    CHECK(fabsf(uc_ref.c - uc_out.c) < 0.3f);
    CHECK(fabsf(uc_ref.alpha - uc_out.alpha) < 0.2f);
    CHECK(fabsf(uc_ref.beta - uc_out.beta) < 0.2f);
    CHECK(fabsf(uc_ref.gamma - uc_out.gamma) < 0.2f);
    CHECK(fabsf(xtal_opt.axis->GetAxis().x - 1.0) < 0.01f);
    CHECK(fabsf(xtal_opt.axis->GetAxis().y) < 0.01f);
    CHECK(fabsf(xtal_opt.axis->GetAxis().z) < 0.01f);
}

// The rotation residual applies `angle_rad * |rot_vec|`, so the goniometer axis block must be
// constrained to unit length or the solve also fits a rotation SCALE that GoniometerAxis::Axis()
// then normalises away. Generate spots whose true rotation is k times the angles the optimizer is
// handed - k = 1 for a healthy goniometer, k = 1.01 for a stage that turned 1 % further than it was
// commanded - and check the axis DIRECTION and the cell come back either way. (Before the axis block
// was put on the unit sphere, the free length absorbed part of that 1 %: on this noise-free data it
// reached only 1.00538 of the 1.01 it would have had to reach to be a measurement.)
TEST_CASE("XtalOptimizer_rotation_axis_direction_only") {
    DiffractionExperiment exp_i;
    exp_i.IncidentEnergy_keV(WVL_1A_IN_KEV)
         .BeamX_pxl(1000).BeamY_pxl(1000)
         .DetectorDistance_mm(200);

    const CrystalLattice latt_base(40, 50, 80, 90, 95, 90);
    const auto uc_ref = latt_base.GetUnitCell();

    const double k_true = GENERATE(1.000, 1.010);

    // The angles the FILE records. The stage actually turned k_true times as far.
    GoniometerAxis nominal("omega", 0.0f, 3.0f, Coord(1, 0, 0), std::nullopt);

    BraggPredictionSettings prediction_settings{.high_res_A = 1.5, .ewald_dist_cutoff = 0.002};
    BraggPrediction prediction;

    const size_t nimages = 40;
    std::vector<std::vector<SpotToSave>> spots(nimages);
    for (size_t img = 0; img < nimages; ++img) {
        const float nominal_deg = nominal.GetAngle_deg(img) + nominal.GetWedge_deg() / 2.0f;
        const RotMatrix rot = nominal.GetTransformationAngle(static_cast<float>(nominal_deg * k_true));
        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 = static_cast<int32_t>(img);
            s.intensity = 1.0f;
            s.phi = nominal_deg;
            s.indexed = true;
            spots.at(img).push_back(s);
        }
    }

    XtalOptimizerData xtal_opt{};
    xtal_opt.latt = latt_base;
    xtal_opt.geom.BeamX_pxl(1000).BeamY_pxl(1000).DetectorDistance_mm(200.0);
    xtal_opt.crystal_system = gemmi::CrystalSystem::Monoclinic;
    xtal_opt.axis = GoniometerAxis("omega", 0.0f, 3.0f, Coord(0.999, 0.02, 0.02).Normalize(),
                                   std::nullopt);
    xtal_opt.min_spots = 200;
    xtal_opt.refine_beam_center = true;
    xtal_opt.refine_detector_angles = true;
    xtal_opt.refine_rotation_axis = true;
    xtal_opt.max_iterations = 200;

    REQUIRE(XtalOptimizer(xtal_opt, spots));

    const Coord axis_out = xtal_opt.axis->GetAxis();
    CHECK(std::fabs(axis_out.Length() - 1.0f) < 1e-5f);
    CHECK(std::fabs(axis_out.x - 1.0f) < 0.01f);
    CHECK(std::fabs(axis_out.y) < 0.01f);
    CHECK(std::fabs(axis_out.z) < 0.01f);

    const auto uc_out = xtal_opt.latt.GetUnitCell();
    CHECK(std::fabs(uc_ref.a - uc_out.a) < 0.2f);
    CHECK(std::fabs(uc_ref.b - uc_out.b) < 0.2f);
    CHECK(std::fabs(uc_ref.c - uc_out.c) < 0.3f);
}

// --- helpers for lattice sanity tests ---
#include <Eigen/Dense>

TEST_CASE("XtalOptimizer_prior_shrinks_toward_its_centre") {
    DiffractionExperiment exp_i;
    exp_i.IncidentEnergy_keV(WVL_1A_IN_KEV).BeamX_pxl(1000).BeamY_pxl(1000).DetectorDistance_mm(200);

    const CrystalLattice true_latt(40, 50, 80, 90, 90, 90);
    BraggPredictionSettings prediction_settings{.high_res_A = 1.5, .ewald_dist_cutoff = 0.001};
    BraggPrediction prediction;
    const auto count = prediction.Calc(exp_i, true_latt, prediction_settings);
    std::vector<SpotToSave> spots;
    for (int i = 0; i < count; ++i) {
        const auto &refl = prediction.GetReflections().at(i);
        spots.push_back(SpotToSave{refl.predicted_x, refl.predicted_y});
    }

    // The centre the refinement starts from and is restrained to: the crystal turned by 0.2 deg and the
    // beam centre one pixel off, with the data all pointing back at the truth.
    const CrystalLattice centre = true_latt.Multiply(RotMatrix(0.2f * static_cast<float>(M_PI) / 180.0f, Coord(0, 1, 0)));
    const auto refined = [&](double orientation_weight, double beam_weight) {
        XtalOptimizerData data{};
        data.max_iterations = 50;
        data.refine_unit_cell = false;
        data.crystal_system = gemmi::CrystalSystem::Orthorhombic;
        data.latt = centre;
        data.geom.BeamX_pxl(1001).BeamY_pxl(1000).DetectorDistance_mm(200);
        data.prior_lattice = centre;
        data.prior_beam_x = 1001;
        data.prior_beam_y = 1000;
        data.orientation_prior_weight = orientation_weight;
        data.beam_prior_weight = beam_weight;
        REQUIRE(XtalOptimizer(data, spots));
        return data;
    };
    const auto moved_deg = [&](const XtalOptimizerData &d) {
        const auto r = d.latt.Vec0() - centre.Vec0();
        return r.Length() / centre.Vec0().Length() * 180.0 / M_PI;
    };

    // No weight: the data alone, which reach the truth.
    const auto free_fit = refined(0.0, 0.0);
    CHECK(moved_deg(free_fit) == Catch::Approx(0.2).margin(0.02));
    CHECK(free_fit.geom.GetBeamX_pxl() == Catch::Approx(1000.0).margin(0.05));

    // A weight far above the data's: the centre.
    const auto held = refined(1e3, 1e3);
    CHECK(moved_deg(held) < 0.005);
    CHECK(held.geom.GetBeamX_pxl() == Catch::Approx(1001.0).margin(0.01));

    // A weight of the data's own order: in between.
    const auto between = refined(0.3, 3e-3);
    CHECK(moved_deg(between) > moved_deg(held));
    CHECK(moved_deg(between) < moved_deg(free_fit));
}

TEST_CASE("XtalOptimizer_prior_same_whatever_the_basis_is_turned_by") {
    // The same crystal written with a and b negated - a basis turned by 180 deg about c - is the same
    // physics and must be restrained the same: the restraint acts on the rotation, not on the angle-axis
    // parameters, which across the axis change by 1/0.64 as much for the same rotation at 180 deg.
    DiffractionExperiment exp_i;
    exp_i.IncidentEnergy_keV(WVL_1A_IN_KEV).BeamX_pxl(1000).BeamY_pxl(1000).DetectorDistance_mm(200);

    const CrystalLattice true_latt(40, 50, 80, 90, 90, 90);
    BraggPredictionSettings prediction_settings{.high_res_A = 1.5, .ewald_dist_cutoff = 0.001};
    BraggPrediction prediction;
    const auto count = prediction.Calc(exp_i, true_latt, prediction_settings);
    std::vector<SpotToSave> spots;
    for (int i = 0; i < count; ++i) {
        const auto &refl = prediction.GetReflections().at(i);
        spots.push_back(SpotToSave{refl.predicted_x, refl.predicted_y});
    }

    const CrystalLattice centre = true_latt.Multiply(RotMatrix(0.2f * static_cast<float>(M_PI) / 180.0f, Coord(1, 1, 0).Normalize()));
    const CrystalLattice flipped(-centre.Vec0(), -centre.Vec1(), centre.Vec2());
    const auto moved_deg = [&](const CrystalLattice &start) {
        XtalOptimizerData data{};
        data.max_iterations = 50;
        data.refine_unit_cell = false;
        data.refine_beam_center = false;
        data.crystal_system = gemmi::CrystalSystem::Orthorhombic;
        data.latt = start;
        data.geom.BeamX_pxl(1000).BeamY_pxl(1000).DetectorDistance_mm(200);
        data.prior_lattice = start;
        data.orientation_prior_weight = 3.0;
        REQUIRE(XtalOptimizer(data, spots));
        const auto r = data.latt.Vec2() - centre.Vec2();
        return r.Length() / centre.Vec2().Length() * 180.0 / M_PI;
    };
    const double a = moved_deg(centre), b = moved_deg(flipped);
    CHECK(a > 0.01);
    CHECK(b == Catch::Approx(a).epsilon(0.02));
}