From 479e4e7261fd7026d0f81f6f9b1e1b1af035e22e Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Tue, 27 Jan 2026 16:29:44 +0100 Subject: [PATCH] XtalOptimizer: Refine rotation axis --- common/GoniometerAxis.cpp | 19 ++-- common/GoniometerAxis.h | 1 + .../geom_refinement/XtalOptimizer.cpp | 37 +++++--- .../geom_refinement/XtalOptimizer.h | 1 + tests/XtalOptimizerTest.cpp | 86 +++++++++++++++++++ 5 files changed, 127 insertions(+), 17 deletions(-) diff --git a/common/GoniometerAxis.cpp b/common/GoniometerAxis.cpp index b88d8e51..163316d9 100644 --- a/common/GoniometerAxis.cpp +++ b/common/GoniometerAxis.cpp @@ -31,6 +31,20 @@ GoniometerAxis::GoniometerAxis(const std::string& in_name, helical_step = in_helical_step; } +GoniometerAxis &GoniometerAxis::ScreeningWedge(const std::optional &input) { + screening_wedge = input; + return *this; +} + +GoniometerAxis &GoniometerAxis::Axis(const Coord &input) { + if (input.Length() == 0.0f) + throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, + "Rotation axis cannot have 0 length"); + + axis = input.Normalize(); + return *this; +} + std::string GoniometerAxis::GetName() const { return name; } @@ -97,11 +111,6 @@ std::vector GoniometerAxis::GetAngleContainer(int64_t max_image_number) return angle_container; } -GoniometerAxis &GoniometerAxis::ScreeningWedge(const std::optional &input) { - screening_wedge = input; - return *this; -} - std::optional GoniometerAxis::GetScreeningWedge() const { return screening_wedge; } diff --git a/common/GoniometerAxis.h b/common/GoniometerAxis.h index f644a40f..ecda91df 100644 --- a/common/GoniometerAxis.h +++ b/common/GoniometerAxis.h @@ -23,6 +23,7 @@ public: const Coord &axis, const std::optional &helical_step); + GoniometerAxis& Axis(const Coord &input); GoniometerAxis& ScreeningWedge(const std::optional& input); [[nodiscard]] std::string GetName() const; diff --git a/image_analysis/geom_refinement/XtalOptimizer.cpp b/image_analysis/geom_refinement/XtalOptimizer.cpp index 3ffa8b23..5e000cbd 100644 --- a/image_analysis/geom_refinement/XtalOptimizer.cpp +++ b/image_analysis/geom_refinement/XtalOptimizer.cpp @@ -9,7 +9,6 @@ struct XtalResidual { XtalResidual(double x, double y, double lambda, double pixel_size, - Coord rot_axis, double angle_rad, double exp_h, double exp_k, double exp_l, @@ -20,7 +19,6 @@ struct XtalResidual { exp_h(exp_h), exp_k(exp_k), exp_l(exp_l), - rot_axis(rot_axis), angle_rad(angle_rad), symmetry(symmetry) { } @@ -30,6 +28,7 @@ struct XtalResidual { const T *const beam_y, const T *const distance_mm, const T *const detector_rot, + const T *const rotation_axis, const T *const p0, const T *const p1, const T *const p2, @@ -61,13 +60,11 @@ struct XtalResidual { // Apply goniometer "back-to-start" rotation: // brings observed reciprocal from image orientation into reference crystal frame - T a[3]; - a[0] = T(angle_rad*rot_axis.x); - a[1] = T(angle_rad*rot_axis.y); - a[2] = T(angle_rad*rot_axis.z); - - T rot_arr[9]; - ceres::AngleAxisToRotationMatrix(a, rot_arr); + T v[3], rot_arr[9]; + v[0] = T(angle_rad) * rotation_axis[0]; + v[1] = T(angle_rad) * rotation_axis[1]; + v[2] = T(angle_rad) * rotation_axis[2]; + ceres::AngleAxisToRotationMatrix(v, rot_arr); Eigen::Matrix R_gonio_back(rot_arr); Eigen::Matrix e_obs_recip = R_gonio_back * e_obs_recip_raw; @@ -149,7 +146,6 @@ struct XtalResidual { const double exp_h; const double exp_k; const double exp_l; - const Coord rot_axis; const double angle_rad; gemmi::CrystalSystem symmetry; }; @@ -393,6 +389,8 @@ bool XtalOptimizerInternal(XtalOptimizerData &data, double latt_vec0[3], latt_vec1[3], latt_vec2[3]; + double rot_vec[3] = {1, 0, 0}; + switch (data.crystal_system) { case gemmi::CrystalSystem::Orthorhombic: LatticeToRodriguesAndLengths_GS(data.latt, latt_vec0, latt_vec1); @@ -424,6 +422,12 @@ bool XtalOptimizerInternal(XtalOptimizerData &data, break; } + if (data.axis) { + rot_vec[0] = data.axis->GetAxis().x; + rot_vec[1] = data.axis->GetAxis().y; + rot_vec[2] = data.axis->GetAxis().z; + } + const float tolerance_sq = tolerance * tolerance; // Add residuals for each point @@ -461,11 +465,10 @@ bool XtalOptimizerInternal(XtalOptimizerData &data, continue; problem.AddResidualBlock( - new ceres::AutoDiffCostFunction( + new ceres::AutoDiffCostFunction( new XtalResidual(pt.x, pt.y, data.geom.GetWavelength_A(), data.geom.GetPixelSize_mm(), - axis, angle_rad, h, k, l, data.crystal_system)), @@ -474,6 +477,7 @@ bool XtalOptimizerInternal(XtalOptimizerData &data, &beam_y, &distance_mm, detector_rot, + rot_vec, latt_vec0, latt_vec1, latt_vec2 @@ -506,6 +510,10 @@ bool XtalOptimizerInternal(XtalOptimizerData &data, } } + if (!data.refine_rotation_axis) { + problem.SetParameterBlockConstant(rot_vec); + } + if (!data.refine_unit_cell) { problem.SetParameterBlockConstant(latt_vec1); problem.SetParameterBlockConstant(latt_vec2); @@ -564,6 +572,11 @@ bool XtalOptimizerInternal(XtalOptimizerData &data, std::cout << "X " << rot_vector.Length() * 180.0 / M_PI << " " << rot_vector_norm << std::endl; } + if (data.axis && data.refine_rotation_axis) { + data.axis.value().Axis(Coord(rot_vec[0], rot_vec[1], rot_vec[2])); + std::cout << "Refined rotation axis " << rot_vec[0] << " " << rot_vec[1] << " " << rot_vec[2] << std::endl; + } + if (data.crystal_system == gemmi::CrystalSystem::Orthorhombic) data.latt = AngleAxisAndLengthsToLattice(latt_vec0, latt_vec1, false); else if (data.crystal_system == gemmi::CrystalSystem::Tetragonal) { diff --git a/image_analysis/geom_refinement/XtalOptimizer.h b/image_analysis/geom_refinement/XtalOptimizer.h index 1854459c..c78ff96f 100644 --- a/image_analysis/geom_refinement/XtalOptimizer.h +++ b/image_analysis/geom_refinement/XtalOptimizer.h @@ -27,6 +27,7 @@ struct XtalOptimizerData { bool refine_distance_mm = false; bool refine_detector_angles = false; bool refine_unit_cell = true; // This refines unit cell size + angles - orientation is always refined + bool refine_rotation_axis = false; bool index_ice_rings = true; diff --git a/tests/XtalOptimizerTest.cpp b/tests/XtalOptimizerTest.cpp index 8503c8b0..96bc42ee 100644 --- a/tests/XtalOptimizerTest.cpp +++ b/tests/XtalOptimizerTest.cpp @@ -627,3 +627,89 @@ TEST_CASE("XtalOptimizer_rotation") { 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 + }; + + std::vector 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 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.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 = 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_distance_mm = false; + xtal_opt.refine_detector_angles = false; + xtal_opt.refine_rotation_axis = true; + + 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(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); +} \ No newline at end of file