diff --git a/image_analysis/IndexAndRefine.cpp b/image_analysis/IndexAndRefine.cpp index c9d16b67..5a9357d5 100644 --- a/image_analysis/IndexAndRefine.cpp +++ b/image_analysis/IndexAndRefine.cpp @@ -103,20 +103,20 @@ void IndexAndRefine::RefineGeometryIfNeeded(DataMessage &msg, IndexAndRefine::In .max_time = 0.04 // 40 ms is max allowed time for the operation }; - if (experiment.IsRotationIndexing()) { - data.refine_beam_center = false; - data.refine_rotation_axis = false; - data.refine_unit_cell = false; - } if (outcome.symmetry.crystal_system == gemmi::CrystalSystem::Trigonal) data.crystal_system = gemmi::CrystalSystem::Hexagonal; + switch (experiment.GetIndexingSettings().GetGeomRefinementAlgorithm()) { case GeomRefinementAlgorithmEnum::None: break; case GeomRefinementAlgorithmEnum::BeamCenter: - if (XtalOptimizer(data, msg.spots)) { + if (experiment.IsRotationIndexing()) { + XtalOptimizerRotationOnly(data, msg.spots, 0.2); + XtalOptimizerRotationOnly(data, msg.spots, 0.1); + XtalOptimizerRotationOnly(data, msg.spots, 0.05); + } else if (XtalOptimizer(data, msg.spots)) { outcome.experiment.BeamX_pxl(data.geom.GetBeamX_pxl()) .BeamY_pxl(data.geom.GetBeamY_pxl()); outcome.beam_center_updated = true; diff --git a/image_analysis/geom_refinement/XtalOptimizer.cpp b/image_analysis/geom_refinement/XtalOptimizer.cpp index 0a4685df..aec455c9 100644 --- a/image_analysis/geom_refinement/XtalOptimizer.cpp +++ b/image_analysis/geom_refinement/XtalOptimizer.cpp @@ -187,6 +187,60 @@ struct XtalResidual { gemmi::CrystalSystem symmetry; }; +struct XtalResidualRotationOnlyPrecomp { + XtalResidualRotationOnlyPrecomp(const Coord &recip_obs, + const Coord &latt_a, + const Coord &latt_b, + const Coord &latt_c, + double h, double k, double l) + : s_obs(recip_obs), + a0(latt_a), b0(latt_b), c0(latt_c), + h(h), k(k), l(l) { + } + + template + bool operator()(const T *const rot_aa, T *residual) const { + // Rotate the CURRENT lattice vectors by rot_aa (proper SO(3) rotation) + T a_in[3] = {T(a0.x), T(a0.y), T(a0.z)}; + T b_in[3] = {T(b0.x), T(b0.y), T(b0.z)}; + T c_in[3] = {T(c0.x), T(c0.y), T(c0.z)}; + + T a_rot[3], b_rot[3], c_rot[3]; + ceres::AngleAxisRotatePoint(rot_aa, a_in, a_rot); + ceres::AngleAxisRotatePoint(rot_aa, b_in, b_rot); + ceres::AngleAxisRotatePoint(rot_aa, c_in, c_rot); + + const Eigen::Matrix A(a_rot[0], a_rot[1], a_rot[2]); + const Eigen::Matrix B(b_rot[0], b_rot[1], b_rot[2]); + const Eigen::Matrix C(c_rot[0], c_rot[1], c_rot[2]); + + // Reciprocal basis from rotated direct lattice + const Eigen::Matrix BxC = B.cross(C); + const Eigen::Matrix CxA = C.cross(A); + const Eigen::Matrix AxB = A.cross(B); + + const T V = A.dot(BxC); + const T invV = T(1) / V; + + const Eigen::Matrix Astar = BxC * invV; + const Eigen::Matrix Bstar = CxA * invV; + const Eigen::Matrix Cstar = AxB * invV; + + const Eigen::Matrix s_pred = + Astar * T(h) + Bstar * T(k) + Cstar * T(l); + + // Residual in reciprocal space + residual[0] = T(s_obs.x) - s_pred[0]; + residual[1] = T(s_obs.y) - s_pred[1]; + residual[2] = T(s_obs.z) - s_pred[2]; + return true; + } + + Coord s_obs; + Coord a0, b0, c0; + double h, k, l; +}; + inline void LatticeToRodriguesAndLengths_GS(const CrystalLattice &latt, double rod[3], double lengths[3]) { @@ -647,3 +701,110 @@ bool XtalOptimizer(XtalOptimizerData &data, const std::vector &spots return XtalOptimizerInternal(data, spots, 0.1); } +bool XtalOptimizerRotationOnly(XtalOptimizerData &data, + const std::vector &spots, + const float tolerance) { + try { + // Parameter: angle-axis for the extra rotation. Identity == {0,0,0}. + double rot_aa[3] = {0.0, 0.0, 0.0}; + + const Coord a0 = data.latt.Vec0(); + const Coord b0 = data.latt.Vec1(); + const Coord c0 = data.latt.Vec2(); + + // Spot selection by current indexing (same approach as XtalOptimizerInternal) + const Coord vec0 = data.latt.Vec0(); + const Coord vec1 = data.latt.Vec1(); + const Coord vec2 = data.latt.Vec2(); + + const float tol_sq = tolerance * tolerance; + + ceres::Problem problem; + + for (const auto &pt : spots) { + if (!data.index_ice_rings && pt.ice_ring) + continue; + + // Compute fractional HKL using the CURRENT lattice + Coord recip_index = pt.ReciprocalCoord(data.geom); + if (data.axis.has_value()) + recip_index = data.axis->GetTransformationAngle(pt.phi) * recip_index; + + const double h_fp = static_cast(recip_index * vec0); + const double k_fp = static_cast(recip_index * vec1); + const double l_fp = static_cast(recip_index * vec2); + + const double h = std::round(h_fp); + const double k = std::round(k_fp); + const double l = std::round(l_fp); + + const double norm_sq = + (h - h_fp) * (h - h_fp) + + (k - k_fp) * (k - k_fp) + + (l - l_fp) * (l - l_fp); + + if (norm_sq > static_cast(tol_sq)) + continue; + + const Coord s_obs = data.geom.DetectorToRecip(pt.x, pt.y); + + auto *cost = + new ceres::AutoDiffCostFunction( + new XtalResidualRotationOnlyPrecomp(s_obs, a0, b0, c0, h, k, l) + ); + + problem.AddResidualBlock(cost, nullptr, rot_aa); + } + + if (problem.NumResidualBlocks() < data.min_spots) + return false; + + ceres::Solver::Options options; + options.linear_solver_type = ceres::DENSE_QR; + options.minimizer_progress_to_stdout = false; + options.max_solver_time_in_seconds = data.max_time; + options.logging_type = ceres::LoggingType::SILENT; + options.num_threads = 1; + + ceres::Solver::Summary summary; + ceres::Solve(options, &problem, &summary); + + // Apply rotation to lattice (L' = R * L, acting on column vectors) + double R_raw[9]; + ceres::AngleAxisToRotationMatrix(rot_aa, R_raw); // row-major 3x3 + + Eigen::Matrix3d R; + R << R_raw[0], R_raw[1], R_raw[2], + R_raw[3], R_raw[4], R_raw[5], + R_raw[6], R_raw[7], R_raw[8]; + + const Eigen::Vector3d A(a0.x, a0.y, a0.z); + const Eigen::Vector3d B(b0.x, b0.y, b0.z); + const Eigen::Vector3d C(c0.x, c0.y, c0.z); + + const Eigen::Vector3d A2 = R * A; + const Eigen::Vector3d B2 = R * B; + const Eigen::Vector3d C2 = R * C; + + data.latt = CrystalLattice( + Coord(static_cast(A2.x()), static_cast(A2.y()), static_cast(A2.z())), + Coord(static_cast(B2.x()), static_cast(B2.y()), static_cast(B2.z())), + Coord(static_cast(C2.x()), static_cast(C2.y()), static_cast(C2.z())) + ); + + double theta = std::sqrt(rot_aa[0] * rot_aa[0] + rot_aa[1] * rot_aa[1] + rot_aa[2] * rot_aa[2]); + data.angle_corr = theta; + if (theta > 1e-6) { + Coord rot; + rot.x = rot_aa[0] / theta; + rot.y = rot_aa[1] / theta; + rot.z = rot_aa[2] / theta; + data.angle_axis = rot; + } else + data.angle_axis.reset(); + + return true; + } catch (...) { + return false; + } +} diff --git a/image_analysis/geom_refinement/XtalOptimizer.h b/image_analysis/geom_refinement/XtalOptimizer.h index 79422e6e..9ca306dd 100644 --- a/image_analysis/geom_refinement/XtalOptimizer.h +++ b/image_analysis/geom_refinement/XtalOptimizer.h @@ -38,6 +38,10 @@ struct XtalOptimizerData { // output std::optional beam_corr_x; std::optional beam_corr_y; + + // For rotation only optimizer + std::optional angle_corr; + std::optional angle_axis; }; void LatticeToRodriguesAndLengths_GS(const CrystalLattice &latt, double rod[3], double lengths[3]); @@ -54,5 +58,6 @@ CrystalLattice AngleAxisAndCellToLattice(const double rod[3], double gamma_rad); bool XtalOptimizer(XtalOptimizerData &data, const std::vector &spots); +bool XtalOptimizerRotationOnly(XtalOptimizerData &data, const std::vector &spots, float tolerance); #endif //JFJOCH_XTALOPTIMIZER_H