From df57810eb470df7d13d505d376d426040432bfd9 Mon Sep 17 00:00:00 2001
From: Filip Leonarski <filip.leonarski@psi.ch>
Date: Mon, 1 Dec 2025 15:30:54 +0100
Subject: [PATCH] jfjoch_test: Add test for XtalOptimizer with rotation (WIP)

---
 tests/XtalOptimizerTest.cpp | 79 +++++++++++++++++++++++++++++++++++++
 1 file changed, 79 insertions(+)

diff --git a/tests/XtalOptimizerTest.cpp b/tests/XtalOptimizerTest.cpp
index 451fa615..40007100 100644
--- a/tests/XtalOptimizerTest.cpp
+++ b/tests/XtalOptimizerTest.cpp
@@ -547,3 +547,82 @@ TEST_CASE("LatticeToRodrigues_Hex") {
     CHECK(uc_o.beta == Catch::Approx(90.0));
     CHECK(uc_o.gamma == Catch::Approx(120.0));
 }
+
+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
+    };
+
+    std::vector<SpotToSave> spots;
+    BraggPrediction prediction;
+
+    // Predict reflections for images at 0-30 deg.
+    for (int img = 0; img < 10; ++img) {
+        // For a rotated image, per-image lattice is obtained as Multiply(rot.transpose())
+        const RotMatrix rot = axis.GetTransformation(img);
+        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.ice_ring = false;
+            s.indexed = true;
+            spots.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 = axis;
+    xtal_opt.min_spots = 200;
+    xtal_opt.refine_beam_center = true;
+    xtal_opt.refine_distance_mm = false;
+    xtal_opt.refine_detector_angles = false;
+
+    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);
+}