diff --git a/image_analysis/geom_refinement/CMakeLists.txt b/image_analysis/geom_refinement/CMakeLists.txt index 51df98fa..dcf9936f 100644 --- a/image_analysis/geom_refinement/CMakeLists.txt +++ b/image_analysis/geom_refinement/CMakeLists.txt @@ -6,6 +6,9 @@ ADD_LIBRARY(JFJochGeomRefinement STATIC AssignSpotsToRings.h XtalOptimizer.cpp XtalOptimizer.h + XtalResidual.h + GeometryRefiner.cpp + GeometryRefiner.h LatticeReduction.cpp LatticeReduction.h ) diff --git a/image_analysis/geom_refinement/GeometryRefiner.cpp b/image_analysis/geom_refinement/GeometryRefiner.cpp new file mode 100644 index 00000000..b10ca118 --- /dev/null +++ b/image_analysis/geom_refinement/GeometryRefiner.cpp @@ -0,0 +1,311 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#include "GeometryRefiner.h" + +#include +#include +#include +#include + +#include "../../common/JFJochMath.h" // PI +#include "LatticeReduction.h" +#include "XtalResidual.h" +#include "ceres/ceres.h" + +namespace { + +// Trigonal is described with the hexagonal cell + B, same as the per-image XtalOptimizer. +gemmi::CrystalSystem MapSystem(gemmi::CrystalSystem s) { + return (s == gemmi::CrystalSystem::Trigonal) ? gemmi::CrystalSystem::Hexagonal : s; +} + +// The cell angles (radians) that XtalResidual bakes into its B matrix for a given system. Only +// monoclinic (beta) and triclinic (all three) read p2; the rest use fixed symmetry angles. We keep the +// angles constant in v1, so these are set once from the input cell. +void InitCellAngles(gemmi::CrystalSystem system, const UnitCell &cell, double ang[3]) { + const double d2r = PI / 180.0; + switch (system) { + case gemmi::CrystalSystem::Hexagonal: + ang[0] = PI / 2.0; ang[1] = PI / 2.0; ang[2] = 2.0 * PI / 3.0; + break; + case gemmi::CrystalSystem::Monoclinic: + // XtalResidual reads p2[0] as beta. + ang[0] = cell.beta * d2r; ang[1] = PI / 2.0; ang[2] = PI / 2.0; + break; + case gemmi::CrystalSystem::Triclinic: + ang[0] = cell.alpha * d2r; ang[1] = cell.beta * d2r; ang[2] = cell.gamma * d2r; + break; + default: // orthorhombic / tetragonal / cubic + ang[0] = PI / 2.0; ang[1] = PI / 2.0; ang[2] = PI / 2.0; + break; + } +} + +// Effective (lengths, alpha,beta,gamma) that reproduce the lattice XtalResidual builds from (len, ang), +// so a frame's lattice can be reconstructed with AngleAxisAndCellToLattice for HKL re-assignment. +void EffectiveCell(gemmi::CrystalSystem system, const double len[3], const double ang[3], + double lengths[3], double &alpha, double &beta, double &gamma) { + switch (system) { + case gemmi::CrystalSystem::Hexagonal: + lengths[0] = len[0]; lengths[1] = len[0]; lengths[2] = len[2]; + alpha = PI / 2.0; beta = PI / 2.0; gamma = 2.0 * PI / 3.0; + break; + case gemmi::CrystalSystem::Tetragonal: + lengths[0] = len[0]; lengths[1] = len[0]; lengths[2] = len[2]; + alpha = PI / 2.0; beta = PI / 2.0; gamma = PI / 2.0; + break; + case gemmi::CrystalSystem::Cubic: + lengths[0] = len[0]; lengths[1] = len[0]; lengths[2] = len[0]; + alpha = PI / 2.0; beta = PI / 2.0; gamma = PI / 2.0; + break; + case gemmi::CrystalSystem::Monoclinic: + lengths[0] = len[0]; lengths[1] = len[1]; lengths[2] = len[2]; + alpha = PI / 2.0; beta = ang[0]; gamma = PI / 2.0; + break; + case gemmi::CrystalSystem::Triclinic: + lengths[0] = len[0]; lengths[1] = len[1]; lengths[2] = len[2]; + alpha = ang[0]; beta = ang[1]; gamma = ang[2]; + break; + default: // orthorhombic + lengths[0] = len[0]; lengths[1] = len[1]; lengths[2] = len[2]; + alpha = PI / 2.0; beta = PI / 2.0; gamma = PI / 2.0; + break; + } +} + +// Orientation seed (angle-axis) for a frame's lattice, using the per-system decomposition the per-image +// optimizer uses so it matches XtalResidual's B convention. +void OrientationSeed(gemmi::CrystalSystem system, const CrystalLattice &latt, double rod[3]) { + double scratch[3]; + if (system == gemmi::CrystalSystem::Hexagonal) { + LatticeToRodriguesAndLengths_Hex(latt, rod, scratch); + } else if (system == gemmi::CrystalSystem::Monoclinic) { + double beta; + LatticeToRodriguesLengthsBeta_Mono(latt, rod, scratch, beta); + } else { + LatticeToRodriguesAndLengths_GS(latt, rod, scratch); + } +} + +UnitCell BuildUnitCell(gemmi::CrystalSystem system, const double len[3], const double ang[3]) { + const double r2d = 180.0 / PI; + UnitCell uc{}; + switch (system) { + case gemmi::CrystalSystem::Hexagonal: + uc = {(float)len[0], (float)len[0], (float)len[2], 90.0f, 90.0f, 120.0f}; + break; + case gemmi::CrystalSystem::Tetragonal: + uc = {(float)len[0], (float)len[0], (float)len[2], 90.0f, 90.0f, 90.0f}; + break; + case gemmi::CrystalSystem::Cubic: + uc = {(float)len[0], (float)len[0], (float)len[0], 90.0f, 90.0f, 90.0f}; + break; + case gemmi::CrystalSystem::Monoclinic: + uc = {(float)len[0], (float)len[1], (float)len[2], 90.0f, (float)(ang[0] * r2d), 90.0f}; + break; + case gemmi::CrystalSystem::Triclinic: + uc = {(float)len[0], (float)len[1], (float)len[2], + (float)(ang[0] * r2d), (float)(ang[1] * r2d), (float)(ang[2] * r2d)}; + break; + default: // orthorhombic + uc = {(float)len[0], (float)len[1], (float)len[2], 90.0f, 90.0f, 90.0f}; + break; + } + return uc; +} + +// Anchors the cell lengths to the input cell. Penalises each length's deviation from its input value +// with a large weight, breaking the low-resolution distance<->cell-scale degeneracy (position ~ +// distance / cell_scale) so distance is determined by the data with the cell pinned. +struct CellLengthRegularizer { + double w0, w1, w2; + double l0, l1, l2; + template + bool operator()(const T *const len, T *r) const { + r[0] = T(w0) * (len[0] - T(l0)); + r[1] = T(w1) * (len[1] - T(l1)); + r[2] = T(w2) * (len[2] - T(l2)); + return true; + } +}; + +} // namespace + +GeometryRefinerResult RefineGlobalGeometry(const DiffractionGeometry &nominal_geom, + const UnitCell &input_cell, + const std::vector &frames, + const GeometryRefinerSettings &settings) { + GeometryRefinerResult result; + try { + const gemmi::CrystalSystem system = MapSystem(settings.crystal_system); + + // Shared geometry parameter blocks (refined). Detector tilt and the rotation axis are constant. + double beam[2] = {nominal_geom.GetBeamX_pxl(), nominal_geom.GetBeamY_pxl()}; + double distance_mm = nominal_geom.GetDetectorDistance_mm(); + double detector_rot[2] = {nominal_geom.GetPoniRot1_rad(), nominal_geom.GetPoniRot2_rad()}; + double rot_vec[3] = {1.0, 0.0, 0.0}; + const double rot3 = nominal_geom.GetPoniRot3_rad(); + const double lambda = nominal_geom.GetWavelength_A(); + const double pixel = nominal_geom.GetPixelSize_mm(); + + // Shared cell blocks. Lengths refined (regularized); angles constant. + double cell_len[3] = {input_cell.a, input_cell.b, input_cell.c}; + const double cell_len0[3] = {input_cell.a, input_cell.b, input_cell.c}; + double cell_ang[3]; + InitCellAngles(system, input_cell, cell_ang); + + double eff_len[3], eff_alpha, eff_beta, eff_gamma; // recomputed each round from cell_len/ang + + // Per-frame orientation seeds. + std::vector> orient(frames.size()); + for (size_t i = 0; i < frames.size(); ++i) + OrientationSeed(system, frames[i].lattice, orient[i].data()); + + // Reciprocal length of one detector pixel at the nominal geometry, used to scale the robust + // loss (~3 px) into the reciprocal-space residual units. + const double recip_per_px = pixel / (distance_mm * lambda); + const double loss_scale = 3.0 * recip_per_px; + + // Tolerance schedule (fractional-hkl norm), tightening each round like the per-image optimizer. + auto round_tol = [&](int r) { + if (settings.rounds <= 1) return 0.2; + const double t = static_cast(r) / (settings.rounds - 1); + return 0.3 - (0.3 - 0.1) * t; + }; + + for (int round = 0; round < settings.rounds; ++round) { + EffectiveCell(system, cell_len, cell_ang, eff_len, eff_alpha, eff_beta, eff_gamma); + + DiffractionGeometry cur = nominal_geom; + cur.BeamX_pxl(beam[0]).BeamY_pxl(beam[1]).DetectorDistance_mm(distance_mm); + + const double tol = round_tol(round); + const double tol_sq = tol * tol; + + ceres::Problem problem; + auto ordering = std::make_shared(); + + int spots_used = 0, frames_used = 0; + for (size_t i = 0; i < frames.size(); ++i) { + const CrystalLattice latt = AngleAxisAndCellToLattice( + orient[i].data(), eff_len, eff_alpha, eff_beta, eff_gamma); + const Coord v0 = latt.Vec0(), v1 = latt.Vec1(), v2 = latt.Vec2(); + + int used_in_frame = 0; + for (const auto &s : frames[i].spots) { + const Coord recip = cur.DetectorToRecip(s.x, s.y); + const double hf = recip * v0, kf = recip * v1, lf = recip * v2; + const double h = std::round(hf), k = std::round(kf), l = std::round(lf); + const double dn = (hf - h) * (hf - h) + (kf - k) * (kf - k) + (lf - l) * (lf - l); + if (dn > tol_sq) + continue; + + problem.AddResidualBlock( + new ceres::AutoDiffCostFunction( + new XtalResidual(s.x, s.y, lambda, pixel, rot3, 0.0, h, k, l, system)), + new ceres::CauchyLoss(loss_scale), + beam, &distance_mm, detector_rot, rot_vec, orient[i].data(), cell_len, cell_ang); + ++used_in_frame; + ++spots_used; + } + if (used_in_frame > 0) { + ordering->AddElementToGroup(orient[i].data(), 0); // eliminate orientations first + ++frames_used; + } + } + + if (spots_used < settings.min_spots_total) { + if (round == 0) + return result; // ok stays false + break; // keep the previous round's solution + } + + // Cell-length regularizer (weight ~ coeff * sqrt(N_obs) / L0^2 -> effectively pins the cell). + const double sqrtN = std::sqrt(static_cast(spots_used)); + CellLengthRegularizer *reg = new CellLengthRegularizer{ + settings.cell_reg_coeff * sqrtN / (cell_len0[0] * cell_len0[0]), + settings.cell_reg_coeff * sqrtN / (cell_len0[1] * cell_len0[1]), + settings.cell_reg_coeff * sqrtN / (cell_len0[2] * cell_len0[2]), + cell_len0[0], cell_len0[1], cell_len0[2]}; + problem.AddResidualBlock( + new ceres::AutoDiffCostFunction(reg), nullptr, cell_len); + + // Shared blocks go in the second elimination group; constant ones are stripped by Ceres. + for (double *b : {beam, &distance_mm, detector_rot, rot_vec, cell_len, cell_ang}) + ordering->AddElementToGroup(b, 1); + + // Beam + distance bounds around the nominal value; cell lengths loosely bounded (the + // regularizer does the anchoring). + problem.SetParameterLowerBound(beam, 0, nominal_geom.GetBeamX_pxl() - settings.beam_range_px); + problem.SetParameterUpperBound(beam, 0, nominal_geom.GetBeamX_pxl() + settings.beam_range_px); + problem.SetParameterLowerBound(beam, 1, nominal_geom.GetBeamY_pxl() - settings.beam_range_px); + problem.SetParameterUpperBound(beam, 1, nominal_geom.GetBeamY_pxl() + settings.beam_range_px); + problem.SetParameterLowerBound(&distance_mm, 0, + std::max(1.0, nominal_geom.GetDetectorDistance_mm() - settings.distance_range_mm)); + problem.SetParameterUpperBound(&distance_mm, 0, + nominal_geom.GetDetectorDistance_mm() + settings.distance_range_mm); + for (int j = 0; j < 3; ++j) { + problem.SetParameterLowerBound(cell_len, j, 0.7 * cell_len0[j]); + problem.SetParameterUpperBound(cell_len, j, 1.3 * cell_len0[j]); + } + + problem.SetParameterBlockConstant(detector_rot); // tilt not refined + problem.SetParameterBlockConstant(rot_vec); + problem.SetParameterBlockConstant(cell_ang); // angles not refined in v1 + + ceres::Solver::Options options; + options.linear_solver_type = ceres::DENSE_SCHUR; + options.linear_solver_ordering = ordering; + options.num_threads = std::max(1, settings.num_threads); + options.max_num_iterations = 100; + options.max_solver_time_in_seconds = 60.0; + options.logging_type = ceres::LoggingType::SILENT; + options.minimizer_progress_to_stdout = false; + + ceres::Solver::Summary summary; + ceres::Solve(options, &problem, &summary); + + result.frames_used = frames_used; + result.spots_used = spots_used; + } + + // Final fit quality: median detector residual of the retained spots at the refined geometry. + EffectiveCell(system, cell_len, cell_ang, eff_len, eff_alpha, eff_beta, eff_gamma); + DiffractionGeometry cur = nominal_geom; + cur.BeamX_pxl(beam[0]).BeamY_pxl(beam[1]).DetectorDistance_mm(distance_mm); + std::vector residuals; + for (size_t i = 0; i < frames.size(); ++i) { + const CrystalLattice latt = AngleAxisAndCellToLattice( + orient[i].data(), eff_len, eff_alpha, eff_beta, eff_gamma); + const Coord a = latt.Astar(), b = latt.Bstar(), c = latt.Cstar(); + const Coord v0 = latt.Vec0(), v1 = latt.Vec1(), v2 = latt.Vec2(); + for (const auto &s : frames[i].spots) { + const Coord recip = cur.DetectorToRecip(s.x, s.y); + const double hf = recip * v0, kf = recip * v1, lf = recip * v2; + const double h = std::round(hf), k = std::round(kf), l = std::round(lf); + const double dn = (hf - h) * (hf - h) + (kf - k) * (kf - k) + (lf - l) * (lf - l); + if (dn > 0.1 * 0.1) + continue; + const Coord recip_pred = a * (float)h + b * (float)k + c * (float)l; + const auto [px, py] = cur.RecipToDetector(recip_pred); + if (std::isfinite(px) && std::isfinite(py)) + residuals.push_back(std::hypot(px - s.x, py - s.y)); + } + } + if (!residuals.empty()) { + std::nth_element(residuals.begin(), residuals.begin() + residuals.size() / 2, residuals.end()); + result.median_residual_px = residuals[residuals.size() / 2]; + } + + result.ok = result.spots_used >= settings.min_spots_total && result.frames_used > 0; + result.beam_x_px = beam[0]; + result.beam_y_px = beam[1]; + result.distance_mm = distance_mm; + result.cell = BuildUnitCell(system, cell_len, cell_ang); + return result; + } catch (...) { + result.ok = false; + return result; + } +} diff --git a/image_analysis/geom_refinement/GeometryRefiner.h b/image_analysis/geom_refinement/GeometryRefiner.h new file mode 100644 index 00000000..8bb3b787 --- /dev/null +++ b/image_analysis/geom_refinement/GeometryRefiner.h @@ -0,0 +1,59 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#pragma once + +#include +#include + +#include "../../common/CrystalLattice.h" +#include "../../common/DiffractionGeometry.h" +#include "../../common/UnitCell.h" +#include "gemmi/symmetry.hpp" + +// Offline global (multi-frame) geometry refinement for serial stills. Pass 1 indexes many frames +// independently; this joint bundle adjustment then determines the ONE shared detector geometry (beam +// centre, distance, cell scale) that per-image refinement cannot reach from a single sparse still, so +// re-indexing with it converts more frames. Mirrors the validated Python prototype: one Ceres problem +// with shared beam/distance/cell blocks and a per-frame orientation block, robust loss, and a cell +// regularizer that anchors the (known) cell to break the distance<->cell-scale degeneracy at low +// resolution. Detector tilt (rot1/rot2) is NOT refined (gauge-coupled with the beam; zero indexing gain). + +struct GeomRefineSpot { + float x = 0, y = 0; // observed (raw) spot position, pixels + int32_t h = 0, k = 0, l = 0; // integer Miller index assigned in pass 1 +}; + +struct GeomRefineFrame { + CrystalLattice lattice; // per-frame indexed lattice (real space), for the orientation seed + std::vector spots; // indexed spots on this frame +}; + +struct GeometryRefinerSettings { + gemmi::CrystalSystem crystal_system = gemmi::CrystalSystem::Triclinic; + // Cell-scale regularizer strength (anchors the cell lengths to the input cell so the otherwise + // degenerate distance is determined). Large => cell effectively fixed (the validated safe regime). + double cell_reg_coeff = 30.0; + int rounds = 3; // hkl-reassignment / tolerance-tightening rounds + double beam_range_px = 8.0; // beam bound around the nominal value + double distance_range_mm = 8.0; // distance bound around the nominal value + int min_spots_total = 100; // refuse to refine below this many indexed spots + int num_threads = 1; +}; + +struct GeometryRefinerResult { + bool ok = false; + double beam_x_px = 0, beam_y_px = 0; // refined beam centre (== direct beam, tilt not refined) + double distance_mm = 0; + UnitCell cell{}; // refined cell + int frames_used = 0; + int spots_used = 0; + double median_residual_px = 0; // final per-spot detector residual (fit quality) +}; + +// nominal_geom supplies the starting beam/distance/tilt/wavelength/pixel size; input_cell is anchored by +// the regularizer. Returns ok=false (and leaves the caller's geometry untouched) on any failure. +GeometryRefinerResult RefineGlobalGeometry(const DiffractionGeometry &nominal_geom, + const UnitCell &input_cell, + const std::vector &frames, + const GeometryRefinerSettings &settings); diff --git a/image_analysis/geom_refinement/XtalOptimizer.cpp b/image_analysis/geom_refinement/XtalOptimizer.cpp index aa7f4b95..acf4071f 100644 --- a/image_analysis/geom_refinement/XtalOptimizer.cpp +++ b/image_analysis/geom_refinement/XtalOptimizer.cpp @@ -5,203 +5,11 @@ #include #include "XtalOptimizer.h" +#include "XtalResidual.h" #include "ceres/ceres.h" #include "ceres/rotation.h" #include "LatticeReduction.h" -struct XtalResidual { - XtalResidual(double x, double y, - double lambda, - double pixel_size, - double rot3, - double angle_rad, - double exp_h, double exp_k, - double exp_l, - gemmi::CrystalSystem symmetry) - : obs_x(x), obs_y(y), - inv_lambda(1.0/lambda), - pixel_size(pixel_size), - rot3(rot3), - exp_h(exp_h), - exp_k(exp_k), - exp_l(exp_l), - angle_rad(angle_rad), - symmetry(symmetry) { - if (std::fabs(lambda) < 1e-6) - throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, - "Lambda cannot be close to zero"); - } - - template - bool operator()(const T *const beam, - 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, - T *residual) const { - // PyFAI convention: poni_rot = Rz(-rot3) * Rx(-rot2) * Ry(+rot1). - // detector_rot[0] = rot1, detector_rot[1] = rot2 are refined; rot3 is fixed - // (e.g. from a PONI import) and baked in here as a constant so that a non-zero - // rot3 is not silently dropped during refinement. - const T rot1 = detector_rot[0]; - const T rot2 = detector_rot[1]; - - // Ry(+rot1): rotation around Y-axis - const T c1 = ceres::cos(rot1); - const T s1 = ceres::sin(rot1); - - // Rx(-rot2): rotation around X-axis with inverted sign (PyFAI left-handed) - const T c2 = ceres::cos(rot2); - const T s2 = ceres::sin(rot2); - - // Rz(-rot3): rotation around Z (beam); constant, identity when rot3 == 0 - const T c3 = T(cos(rot3)); - const T s3 = T(sin(rot3)); - - // Detector coordinates in mm - const T det_x = (T(obs_x) - beam[0]) * T(pixel_size); - const T det_y = (T(obs_y) - beam[1]) * T(pixel_size); - const T det_z = T(distance_mm[0]); - - // Apply Ry(rot1) first: rotate around Y - const T t1_x = c1 * det_x + s1 * det_z; - const T t1_y = det_y; - const T t1_z = -s1 * det_x + c1 * det_z; - - // Then apply Rx(-rot2): rotate around X - const T t2_x = t1_x; - const T t2_y = c2 * t1_y + s2 * t1_z; - const T t2_z = -s2 * t1_y + c2 * t1_z; - - // Then apply Rz(-rot3): rotate around Z (beam) - const T x = c3 * t2_x + s3 * t2_y; - const T y = -s3 * t2_x + c3 * t2_y; - const T z = t2_z; - - // convert to recip space - const T lab_norm = ceres::sqrt(x * x + y * y + z * z); - const T inv_norm = T(1) / lab_norm; - - T recip_raw[3]; - recip_raw[0] = x * inv_norm * T(inv_lambda); - recip_raw[1] = y * inv_norm * T(inv_lambda); - recip_raw[2] = (z * inv_norm - T(1.0)) * T(inv_lambda); - - // Apply goniometer "back-to-start" rotation: - // brings observed reciprocal from image orientation into reference crystal frame - const T aa_back[3] = { - T(angle_rad) * rotation_axis[0], - T(angle_rad) * rotation_axis[1], - T(angle_rad) * rotation_axis[2] - }; - - T recip_obs[3]; - ceres::AngleAxisRotatePoint(aa_back, recip_raw, recip_obs); - - const Eigen::Matrix e_obs_recip(recip_obs[0], recip_obs[1], recip_obs[2]); - - // Build unit cell lengths and B (convention: columns are a, b, c prior to global rotation) - Eigen::Matrix e_uc_len = Eigen::Matrix::Zero(); - Eigen::Matrix B = Eigen::Matrix::Identity(); - - if (symmetry == gemmi::CrystalSystem::Hexagonal) { - e_uc_len << p1[0], p1[0], p1[2]; - B(0, 1) = T(-0.5); // cos(120) - B(1, 1) = T(sqrt(3.0) / 2.0); // sin(120) - } else if (symmetry == gemmi::CrystalSystem::Orthorhombic) { - e_uc_len << p1[0], p1[1], p1[2]; - } else if (symmetry == gemmi::CrystalSystem::Tetragonal) { - e_uc_len << p1[0], p1[0], p1[2]; - } else if (symmetry == gemmi::CrystalSystem::Cubic) { - e_uc_len << p1[0], p1[0], p1[0]; - } else if (symmetry == gemmi::CrystalSystem::Monoclinic) { - // Unique axis b: alpha = gamma = 90°, beta free (angle between a and c) - e_uc_len << p1[0], p1[1], p1[2]; - B(0, 2) = ceres::cos(p2[0]); - B(2, 2) = ceres::sin(p2[0]); - } else { - // Triclinic: p1 = (a,b,c), p2 = (alpha, beta, gamma) in radians - const T ca = ceres::cos(p2[0]); - const T cb = ceres::cos(p2[1]); - const T cg = ceres::cos(p2[2]); - const T sg = ceres::sin(p2[2]); - - e_uc_len << p1[0], p1[1], p1[2]; - - B(0, 0) = T(1); - B(1, 0) = T(0); - B(2, 0) = T(0); - B(0, 1) = cg; - B(1, 1) = sg; - B(2, 1) = T(0); - - // c vector components: - const T cx = cb; - const T cy = (ca - cb * cg) / sg; - const T v = T(1) - cx * cx - cy * cy; - const T cz = (v >= T(0)) ? ceres::sqrt(v) : T(0); - - B(0, 2) = cx; - B(1, 2) = cy; - B(2, 2) = cz; - } - - // Build unrotated direct lattice columns: (B * D), then rotate them by p0. - // This avoids AngleAxisToRotationMatrix + matrix multiplications. - const T L0 = e_uc_len[0]; - const T L1 = e_uc_len[1]; - const T L2 = e_uc_len[2]; - - T col0_unrot[3] = {B(0, 0) * L0, B(1, 0) * L0, B(2, 0) * L0}; - T col1_unrot[3] = {B(0, 1) * L1, B(1, 1) * L1, B(2, 1) * L1}; - T col2_unrot[3] = {B(0, 2) * L2, B(1, 2) * L2, B(2, 2) * L2}; - - T col0_rot[3], col1_rot[3], col2_rot[3]; - ceres::AngleAxisRotatePoint(p0, col0_unrot, col0_rot); - ceres::AngleAxisRotatePoint(p0, col1_unrot, col1_rot); - ceres::AngleAxisRotatePoint(p0, col2_unrot, col2_rot); - - const Eigen::Matrix A(col0_rot[0], col0_rot[1], col0_rot[2]); - const Eigen::Matrix Bv(col1_rot[0], col1_rot[1], col1_rot[2]); - const Eigen::Matrix C(col2_rot[0], col2_rot[1], col2_rot[2]); - - const Eigen::Matrix BxC = Bv.cross(C); - const Eigen::Matrix CxA = C.cross(A); - const Eigen::Matrix AxB = A.cross(Bv); - - 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 T h = T(exp_h); - const T k = T(exp_k); - const T l = T(exp_l); - - const Eigen::Matrix e_pred_recip = Astar * h + Bstar * k + Cstar * l; - - residual[0] = e_obs_recip[0] - e_pred_recip[0]; - residual[1] = e_obs_recip[1] - e_pred_recip[1]; - residual[2] = e_obs_recip[2] - e_pred_recip[2]; - - return true; - } - - const double obs_x, obs_y; - const double inv_lambda; - const double pixel_size; - const double rot3; - const double exp_h; - const double exp_k; - const double exp_l; - const double angle_rad; - gemmi::CrystalSystem symmetry; -}; - struct XtalResidualRotationOnlyPrecomp { XtalResidualRotationOnlyPrecomp(const Coord &recip_obs, const CrystalLattice &latt, diff --git a/image_analysis/geom_refinement/XtalResidual.h b/image_analysis/geom_refinement/XtalResidual.h new file mode 100644 index 00000000..6bdcbc3c --- /dev/null +++ b/image_analysis/geom_refinement/XtalResidual.h @@ -0,0 +1,214 @@ +// SPDX-FileCopyrightText: 2025 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#pragma once + +#include + +#include + +#include "ceres/ceres.h" +#include "ceres/rotation.h" +#include "gemmi/symmetry.hpp" + +#include "../../common/JFJochException.h" + +// Detector -> reciprocal geometry residual, shared by the per-image XtalOptimizer (one lattice, one +// frame) and the offline GeometryRefiner (shared beam/distance/cell blocks, one orientation block per +// frame). Parameter blocks: beam(2), distance_mm(1), detector_rot(2 = rot1,rot2), rotation_axis(3), +// p0 = lattice orientation angle-axis(3), p1 = unit-cell lengths(3), p2 = unit-cell angles(3, rad). +// The residual is the difference between the observed reciprocal vector (from the detector position via +// the current beam/distance/tilt) and the predicted one (h*a* + k*b* + l*c* for the current cell + +// orientation), in Angstrom^-1. rot3 and the goniometer angle_rad are baked in as constants. +struct XtalResidual { + XtalResidual(double x, double y, + double lambda, + double pixel_size, + double rot3, + double angle_rad, + double exp_h, double exp_k, + double exp_l, + gemmi::CrystalSystem symmetry) + : obs_x(x), obs_y(y), + inv_lambda(1.0/lambda), + pixel_size(pixel_size), + rot3(rot3), + exp_h(exp_h), + exp_k(exp_k), + exp_l(exp_l), + angle_rad(angle_rad), + symmetry(symmetry) { + if (std::fabs(lambda) < 1e-6) + throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, + "Lambda cannot be close to zero"); + } + + template + bool operator()(const T *const beam, + 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, + T *residual) const { + // PyFAI convention: poni_rot = Rz(-rot3) * Rx(-rot2) * Ry(+rot1). + // detector_rot[0] = rot1, detector_rot[1] = rot2 are refined; rot3 is fixed + // (e.g. from a PONI import) and baked in here as a constant so that a non-zero + // rot3 is not silently dropped during refinement. + const T rot1 = detector_rot[0]; + const T rot2 = detector_rot[1]; + + // Ry(+rot1): rotation around Y-axis + const T c1 = ceres::cos(rot1); + const T s1 = ceres::sin(rot1); + + // Rx(-rot2): rotation around X-axis with inverted sign (PyFAI left-handed) + const T c2 = ceres::cos(rot2); + const T s2 = ceres::sin(rot2); + + // Rz(-rot3): rotation around Z (beam); constant, identity when rot3 == 0 + const T c3 = T(cos(rot3)); + const T s3 = T(sin(rot3)); + + // Detector coordinates in mm + const T det_x = (T(obs_x) - beam[0]) * T(pixel_size); + const T det_y = (T(obs_y) - beam[1]) * T(pixel_size); + const T det_z = T(distance_mm[0]); + + // Apply Ry(rot1) first: rotate around Y + const T t1_x = c1 * det_x + s1 * det_z; + const T t1_y = det_y; + const T t1_z = -s1 * det_x + c1 * det_z; + + // Then apply Rx(-rot2): rotate around X + const T t2_x = t1_x; + const T t2_y = c2 * t1_y + s2 * t1_z; + const T t2_z = -s2 * t1_y + c2 * t1_z; + + // Then apply Rz(-rot3): rotate around Z (beam) + const T x = c3 * t2_x + s3 * t2_y; + const T y = -s3 * t2_x + c3 * t2_y; + const T z = t2_z; + + // convert to recip space + const T lab_norm = ceres::sqrt(x * x + y * y + z * z); + const T inv_norm = T(1) / lab_norm; + + T recip_raw[3]; + recip_raw[0] = x * inv_norm * T(inv_lambda); + recip_raw[1] = y * inv_norm * T(inv_lambda); + recip_raw[2] = (z * inv_norm - T(1.0)) * T(inv_lambda); + + // Apply goniometer "back-to-start" rotation: + // brings observed reciprocal from image orientation into reference crystal frame + const T aa_back[3] = { + T(angle_rad) * rotation_axis[0], + T(angle_rad) * rotation_axis[1], + T(angle_rad) * rotation_axis[2] + }; + + T recip_obs[3]; + ceres::AngleAxisRotatePoint(aa_back, recip_raw, recip_obs); + + const Eigen::Matrix e_obs_recip(recip_obs[0], recip_obs[1], recip_obs[2]); + + // Build unit cell lengths and B (convention: columns are a, b, c prior to global rotation) + Eigen::Matrix e_uc_len = Eigen::Matrix::Zero(); + Eigen::Matrix B = Eigen::Matrix::Identity(); + + if (symmetry == gemmi::CrystalSystem::Hexagonal) { + e_uc_len << p1[0], p1[0], p1[2]; + B(0, 1) = T(-0.5); // cos(120) + B(1, 1) = T(sqrt(3.0) / 2.0); // sin(120) + } else if (symmetry == gemmi::CrystalSystem::Orthorhombic) { + e_uc_len << p1[0], p1[1], p1[2]; + } else if (symmetry == gemmi::CrystalSystem::Tetragonal) { + e_uc_len << p1[0], p1[0], p1[2]; + } else if (symmetry == gemmi::CrystalSystem::Cubic) { + e_uc_len << p1[0], p1[0], p1[0]; + } else if (symmetry == gemmi::CrystalSystem::Monoclinic) { + // Unique axis b: alpha = gamma = 90°, beta free (angle between a and c) + e_uc_len << p1[0], p1[1], p1[2]; + B(0, 2) = ceres::cos(p2[0]); + B(2, 2) = ceres::sin(p2[0]); + } else { + // Triclinic: p1 = (a,b,c), p2 = (alpha, beta, gamma) in radians + const T ca = ceres::cos(p2[0]); + const T cb = ceres::cos(p2[1]); + const T cg = ceres::cos(p2[2]); + const T sg = ceres::sin(p2[2]); + + e_uc_len << p1[0], p1[1], p1[2]; + + B(0, 0) = T(1); + B(1, 0) = T(0); + B(2, 0) = T(0); + B(0, 1) = cg; + B(1, 1) = sg; + B(2, 1) = T(0); + + // c vector components: + const T cx = cb; + const T cy = (ca - cb * cg) / sg; + const T v = T(1) - cx * cx - cy * cy; + const T cz = (v >= T(0)) ? ceres::sqrt(v) : T(0); + + B(0, 2) = cx; + B(1, 2) = cy; + B(2, 2) = cz; + } + + // Build unrotated direct lattice columns: (B * D), then rotate them by p0. + // This avoids AngleAxisToRotationMatrix + matrix multiplications. + const T L0 = e_uc_len[0]; + const T L1 = e_uc_len[1]; + const T L2 = e_uc_len[2]; + + T col0_unrot[3] = {B(0, 0) * L0, B(1, 0) * L0, B(2, 0) * L0}; + T col1_unrot[3] = {B(0, 1) * L1, B(1, 1) * L1, B(2, 1) * L1}; + T col2_unrot[3] = {B(0, 2) * L2, B(1, 2) * L2, B(2, 2) * L2}; + + T col0_rot[3], col1_rot[3], col2_rot[3]; + ceres::AngleAxisRotatePoint(p0, col0_unrot, col0_rot); + ceres::AngleAxisRotatePoint(p0, col1_unrot, col1_rot); + ceres::AngleAxisRotatePoint(p0, col2_unrot, col2_rot); + + const Eigen::Matrix A(col0_rot[0], col0_rot[1], col0_rot[2]); + const Eigen::Matrix Bv(col1_rot[0], col1_rot[1], col1_rot[2]); + const Eigen::Matrix C(col2_rot[0], col2_rot[1], col2_rot[2]); + + const Eigen::Matrix BxC = Bv.cross(C); + const Eigen::Matrix CxA = C.cross(A); + const Eigen::Matrix AxB = A.cross(Bv); + + 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 T h = T(exp_h); + const T k = T(exp_k); + const T l = T(exp_l); + + const Eigen::Matrix e_pred_recip = Astar * h + Bstar * k + Cstar * l; + + residual[0] = e_obs_recip[0] - e_pred_recip[0]; + residual[1] = e_obs_recip[1] - e_pred_recip[1]; + residual[2] = e_obs_recip[2] - e_pred_recip[2]; + + return true; + } + + const double obs_x, obs_y; + const double inv_lambda; + const double pixel_size; + const double rot3; + const double exp_h; + const double exp_k; + const double exp_l; + const double angle_rad; + gemmi::CrystalSystem symmetry; +}; diff --git a/rugnux/Rugnux.cpp b/rugnux/Rugnux.cpp index 72406ddb..656e0d38 100644 --- a/rugnux/Rugnux.cpp +++ b/rugnux/Rugnux.cpp @@ -10,6 +10,7 @@ #include #include #include +#include #include #include #include @@ -23,6 +24,7 @@ #include "../writer/FileWriter.h" #include "../image_analysis/MXAnalysisWithoutFPGA.h" #include "../image_analysis/IndexAndRefine.h" +#include "../image_analysis/geom_refinement/GeometryRefiner.h" #include "../image_analysis/indexing/IndexerThreadPool.h" #include "../image_analysis/azint/AzIntEngineCPU.h" #include "../image_analysis/image_preprocessing/ImagePreprocessorCPU.h" @@ -74,6 +76,144 @@ Rugnux::Rugnux(JFJochHDF5Reader &reader, DiffractionExperiment experiment, : reader_(reader), experiment_(std::move(experiment)), pixel_mask_(std::move(pixel_mask)), config_(std::move(config)) {} +void Rugnux::RefineStillsGeometry(int start_image, int end_image, int images_to_process, + RugnuxObserver *observer) { + Logger logger("Rugnux"); + + // Rotation has its own two-pass; this is stills only. + if (experiment_.IsRotationIndexing()) { + logger.Warning("--refine-geometry is stills-only; rotation data uses two-pass indexing. Ignoring."); + return; + } + // The bundle anchors a known cell; without one there is nothing to refine against. + const auto cell = experiment_.GetUnitCell(); + if (!cell.has_value()) { + logger.Warning("--refine-geometry needs a reference unit cell (-C / -S / reference MTZ). Skipping."); + return; + } + + const int refine_frames = std::max(1, config_.refine_geometry.value()); + if (observer) + observer->OnPhase("Geometry refinement (first pass)"); + + const auto dataset = reader_.GetDataset(); + + // Index a spread sample large enough to yield ~refine_frames strong frames even at a low hit rate. + const int sample_budget = std::min(images_to_process, std::max(refine_frames * 50, 8000)); + const std::vector sample = select_equally_spaced_image_ordinals(images_to_process, sample_budget); + + // First-pass engines, built from the current (nominal) geometry: an index-only pass (no + // integration) to obtain each frame's spots + assigned HKL + orientation. + AzimuthalIntegrationMapping mapping(experiment_, pixel_mask_); + IndexerThreadPool pool(experiment_.GetIndexingSettings()); + IndexAndRefine indexer(experiment_, &pool, /*retain_outcomes=*/false); + + auto pass_settings = config_.spot_finding; + pass_settings.enable = true; + pass_settings.indexing = true; + pass_settings.quick_integration = false; + + std::vector frames; + std::mutex frames_mutex; + std::atomic next_i = 0; + + auto worker = [&]() { + pin_gpu(); // round-robin per worker thread; must precede engine construction + MXAnalysisWithoutFPGA analysis(experiment_, mapping, pixel_mask_, indexer); + AzimuthalIntegrationProfile profile(mapping); + + while (!cancelled_) { + const int idx = next_i.fetch_add(1); + if (idx >= static_cast(sample.size())) break; + const int ordinal = sample[idx]; + const int image_idx = start_image + ordinal * config_.stride; + + std::shared_ptr img; + try { + img = reader_.GetRawImage(image_idx); + } catch (const std::exception &e) { + logger.Warning("Geometry refinement: failed to load image {}: {}", image_idx, e.what()); + continue; + } + if (!img) continue; + + DataMessage msg{}; + msg.image = img->image; + msg.number = ordinal; + msg.original_number = image_idx; + if (dataset->efficiency.size() > image_idx) + msg.image_collection_efficiency = dataset->efficiency[image_idx]; + + try { + analysis.Analyze(msg, profile, pass_settings); + } catch (const std::exception &e) { + continue; + } + if (!msg.indexing_result.value_or(false) || !msg.indexing_lattice.has_value()) + continue; + + GeomRefineFrame f; + f.lattice = *msg.indexing_lattice; + for (const auto &s : msg.spots) + if (s.indexed && s.lattice == 0) + f.spots.push_back(GeomRefineSpot{s.x, s.y, + static_cast(s.h), static_cast(s.k), static_cast(s.l)}); + if (f.spots.size() >= 6) { + const std::unique_lock ul(frames_mutex); + frames.push_back(std::move(f)); + } + } + }; + + std::vector> futures; + futures.reserve(config_.nthreads); + for (int i = 0; i < config_.nthreads; ++i) + futures.push_back(std::async(std::launch::async, worker)); + for (auto &fut : futures) + fut.get(); + + if (cancelled_) + return; + + // Keep the strongest frames (most indexed spots) for the bundle - the true cell indexes many + // spots per frame, and orientation diversity comes for free from independent serial stills. + constexpr int MIN_STRONG_SPOTS = 20; + std::vector strong; + for (auto &f : frames) + if (static_cast(f.spots.size()) >= MIN_STRONG_SPOTS) + strong.push_back(std::move(f)); + std::sort(strong.begin(), strong.end(), + [](const GeomRefineFrame &a, const GeomRefineFrame &b) { return a.spots.size() > b.spots.size(); }); + if (static_cast(strong.size()) > refine_frames) + strong.resize(refine_frames); + + logger.Info("Geometry refinement: indexed {} of {} sampled frames, bundling {} strong frames (>= {} spots)", + frames.size(), sample.size(), strong.size(), MIN_STRONG_SPOTS); + + GeometryRefinerSettings gr_settings; + gr_settings.crystal_system = experiment_.GetCrystalSystem(); + gr_settings.num_threads = config_.nthreads; + + const GeometryRefinerResult r = RefineGlobalGeometry( + experiment_.GetDiffractionGeometry(), *cell, strong, gr_settings); + if (!r.ok) { + logger.Warning("Geometry refinement did not run (too few strong frames / did not converge); " + "keeping the input geometry"); + return; + } + + logger.Info("Geometry refinement: beam ({:.2f}, {:.2f}) -> ({:.2f}, {:.2f}) px, distance {:.4f} -> {:.4f} mm", + experiment_.GetBeamX_pxl(), experiment_.GetBeamY_pxl(), r.beam_x_px, r.beam_y_px, + experiment_.GetDetectorDistance_mm(), r.distance_mm); + logger.Info("Geometry refinement: cell a/b/c {:.3f}/{:.3f}/{:.3f} -> {:.3f}/{:.3f}/{:.3f} A " + "(median residual {:.3f} px, {} frames, {} spots)", + cell->a, cell->b, cell->c, r.cell.a, r.cell.b, r.cell.c, + r.median_residual_px, r.frames_used, r.spots_used); + + experiment_.BeamX_pxl(r.beam_x_px).BeamY_pxl(r.beam_y_px).DetectorDistance_mm(r.distance_mm); + experiment_.SetUnitCell(r.cell); +} + ProcessResult Rugnux::Run(RugnuxObserver *observer) { Logger logger("Rugnux"); ProcessResult result; @@ -130,6 +270,13 @@ ProcessResult Rugnux::Run(RugnuxObserver *observer) { experiment_.Goniometer(shifted); } + // Stills global geometry-refinement first pass: bundle-adjust the shared beam/distance/cell from a + // sample of strongly-indexed frames and apply it to experiment_, so everything below (azimuthal + // mapping, the main index+integrate pass, scaling) uses the refined geometry. Rotation is skipped + // inside (it has its own two-pass). + if (full && config_.refine_geometry.has_value()) + RefineStillsGeometry(start_image, end_image, images_to_process, observer); + AzimuthalIntegrationMapping mapping(experiment_, pixel_mask_); JFJochReceiverPlots plots; diff --git a/rugnux/Rugnux.h b/rugnux/Rugnux.h index da918b96..f8aa5dd5 100644 --- a/rugnux/Rugnux.h +++ b/rugnux/Rugnux.h @@ -61,6 +61,12 @@ struct ProcessConfig { int rotation_indexing_image_count = 100; std::optional forced_rotation_lattice; + // Stills global geometry-refinement two-pass (FullAnalysis, stills only). When set, an extra first + // pass indexes a spread sample of frames, a single joint bundle adjustment over the strongest ones + // determines the shared beam/distance/cell that per-image refinement cannot reach, and the main pass + // re-indexes with it. The value is the number of strong frames fed to the bundle (--refine-geometry). + std::optional refine_geometry; + // Scaling / merging (FullAnalysis). reference_data (from a reference MTZ) also drives the // per-image live scaling path when present. bool run_scaling = false; @@ -129,6 +135,13 @@ class Rugnux { ProcessConfig config_; std::atomic cancelled_{false}; + // Stills global geometry-refinement first pass (config_.refine_geometry): index a spread sample of + // frames, bundle-adjust the shared beam/distance/cell from the strongest ones, and apply the result + // to experiment_ so the main pass re-indexes with it. No-op (leaves experiment_ unchanged) if it + // cannot run. Stills only; rotation has its own two-pass. + void RefineStillsGeometry(int start_image, int end_image, int images_to_process, + RugnuxObserver *observer); + public: Rugnux(JFJochHDF5Reader &reader, DiffractionExperiment experiment, PixelMask pixel_mask, ProcessConfig config); diff --git a/rugnux/rugnux_cli.cpp b/rugnux/rugnux_cli.cpp index 371cc39c..a31bdaa8 100644 --- a/rugnux/rugnux_cli.cpp +++ b/rugnux/rugnux_cli.cpp @@ -76,6 +76,7 @@ void print_usage() { std::cout << " -S, --space-group Space group number - used for both indexing and scaling" << std::endl; std::cout << " -C, --unit-cell Fix reference unit cell: \"a,b,c,alpha,beta,gamma\"" << std::endl; std::cout << " -r, --refine Geometry refinement algorithm (none|orientation|beam_and_lattice|multi); multi tries all three per image and keeps whichever indexes the most spots" << std::endl; + std::cout << " --refine-geometry[=N] Stills: extra first pass that bundle-adjusts the shared beam/distance/cell from N strongly-indexed frames (default: 200), then re-indexes with it (lifts weak-stills indexing)" << std::endl; std::cout << std::endl; std::cout << " Scaling and merging (on by default)" << std::endl; @@ -146,6 +147,7 @@ enum { OPT_SINGLE_PASS_ROTATION, OPT_REDO_ROTATION_SPOTS, OPT_FORCE_ROTATION_LATTICE, + OPT_REFINE_GEOMETRY, OPT_BANDWIDTH, OPT_INTEGRATION_RADIUS, OPT_REJECT_OUTLIERS, @@ -226,6 +228,7 @@ static option long_options[] = { {"polarization", required_argument, nullptr, OPT_POLARIZATION}, {"redo-rotation-spots", no_argument, nullptr, OPT_REDO_ROTATION_SPOTS}, {"force-rotation-lattice", required_argument, nullptr, OPT_FORCE_ROTATION_LATTICE}, + {"refine-geometry", optional_argument, nullptr, OPT_REFINE_GEOMETRY}, {"spot-sigma", required_argument, nullptr, OPT_SPOT_SIGMA}, @@ -476,6 +479,7 @@ int main(int argc, char **argv) { double min_image_cc = 0.0; int64_t scaling_iter = 3; std::optional forced_rotation_lattice; + std::optional refine_geometry; // --refine-geometry[=N]: stills global geometry-refinement pass std::optional bandwidth_fwhm; // relative FWHM of dlambda/lambda @@ -550,6 +554,15 @@ int main(int argc, char **argv) { case OPT_REDO_ROTATION_SPOTS: reuse_rotation_spots = false; break; + case OPT_REFINE_GEOMETRY: { + int n = optarg ? atoi(optarg) : 200; + if (n <= 0) { + logger.Error("Invalid --refine-geometry frame count: {}", optarg ? optarg : ""); + exit(EXIT_FAILURE); + } + refine_geometry = n; + break; + } case OPT_FORCE_ROTATION_LATTICE: { if (rotation_indexing) { logger.Error("Rotation indexing already enabled"); @@ -1374,6 +1387,7 @@ int main(int argc, char **argv) { config.reuse_rotation_spots = reuse_rotation_spots; config.rotation_indexing_image_count = rotation_indexing_image_count; config.forced_rotation_lattice = forced_rotation_lattice; + config.refine_geometry = refine_geometry; config.run_scaling = run_scaling; config.scaling_iter = scaling_iter; config.reference_data = reference_data;