Files
Jungfraujoch/image_analysis/geom_refinement/XtalOptimizer.cpp
T
leonarski_fandClaude Opus 4.8 c140b3567d rugnux: restrain the gauge-weak beam centre toward the header in geometry refinement
Single-axis rotation has a gauge degeneracy: rotating the whole experiment
about the spindle leaves every spot position unchanged, so the beam-centre
component parallel to the spindle is a null / under-determined direction.
Refining it freely lets it wander a few pixels and absorb centroid
systematics into a wrong beam that the co-refined orientation keeps
position-consistent.

Add a soft header prior on that one component in the shared XtalOptimizer:
residual = w*(beam[parallel] - header), with w set so the restraint behaves
like a sigma_px-pixel prior competing with the positional residuals. The
gauge direction has ~zero data sensitivity, so the prior pins it near the
LaB6-monitored header, while a real, well-supported drift can still overcome
it (the beam does drift). One restraint covers all three refinement sites at
once: the noisy accumulated-spot primary, the per-frame filter, and
post-refine.

Validated against independent XDS processing: neutral on well-determined
data (the gauge direction has no effect there), and a real low-resolution
accuracy gain where the free refinement had otherwise wandered off the
header - the merged intensities move back into agreement with XDS while the
internal CC1/2, which is blind to this, barely changes. No space-group
determinations change.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-07-18 17:31:45 +02:00

475 lines
20 KiB
C++
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
// SPDX-FileCopyrightText: 2025 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only
#include "../../common/JFJochMath.h"
#include <Eigen/Dense>
#include "XtalOptimizer.h"
#include "XtalResidual.h"
#include "ceres/ceres.h"
#include "ceres/rotation.h"
#include "LatticeReduction.h"
// Soft header prior on ONE beam-centre component (the spindle-parallel, gauge-weak one). Residual = w*(b - b0);
// the caller sets w so the prior behaves like a sigma-pixel restraint that competes with the (unit-weight)
// positional residuals - strong enough to pin the gauge direction, negligible in the well-constrained one.
struct BeamComponentPrior {
BeamComponentPrior(int component, double b0, double weight)
: component(component), b0(b0), weight(weight) {}
template<typename T>
bool operator()(const T *const beam, T *residual) const {
residual[0] = T(weight) * (beam[component] - T(b0));
return true;
}
int component;
double b0, weight;
};
struct XtalResidualRotationOnlyPrecomp {
XtalResidualRotationOnlyPrecomp(const Coord &recip_obs,
const CrystalLattice &latt,
double h, double k, double l)
: s_obs(recip_obs),
astar(latt.Astar()), bstar(latt.Bstar()), cstar(latt.Cstar()),
h(h), k(k), l(l) {
}
template<typename T>
bool operator()(const T *const rot_aa, T *residual) const {
const T astar_unrot[3] = {T(astar.x), T(astar.y), T(astar.z)};
const T bstar_unrot[3] = {T(bstar.x), T(bstar.y), T(bstar.z)};
const T cstar_unrot[3] = {T(cstar.x), T(cstar.y), T(cstar.z)};
T astar_rot[3], bstar_rot[3], cstar_rot[3];
ceres::AngleAxisRotatePoint(rot_aa, astar_unrot, astar_rot);
ceres::AngleAxisRotatePoint(rot_aa, bstar_unrot, bstar_rot);
ceres::AngleAxisRotatePoint(rot_aa, cstar_unrot, cstar_rot);
const Eigen::Matrix<T, 3, 1> s_pred(T(h) * astar_rot[0] + T(k) * bstar_rot[0] + T(l) * cstar_rot[0],
T(h) * astar_rot[1] + T(k) * bstar_rot[1] + T(l) * cstar_rot[1],
T(h) * astar_rot[2] + T(k) * bstar_rot[2] + T(l) * cstar_rot[2]
);
// 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;
}
const Coord s_obs;
const Coord astar, bstar, cstar;
const double h, k, l;
};
// Regularizer: penalises ||rot_aa|| to prefer the smallest rotation that
// explains the data. Weight should be chosen in the same units as the
// reciprocal-space residuals (Å⁻¹ per radian). A value of ~0.010.1 is
// typically enough to break degeneracy without biasing the solution.
struct RotationNormRegularizer {
explicit RotationNormRegularizer(double weight) : weight(weight) {}
template<typename T>
bool operator()(const T *const rot_aa, T *residual) const {
residual[0] = T(weight) * rot_aa[0];
residual[1] = T(weight) * rot_aa[1];
residual[2] = T(weight) * rot_aa[2];
return true;
}
const double weight;
};
bool XtalOptimizerInternal(XtalOptimizerData &data,
const std::vector<std::vector<SpotToSave>> &spots,
const float tolerance,
const int num_threads) {
try {
Coord vec0 = data.latt.Vec0();
Coord vec1 = data.latt.Vec1();
Coord vec2 = data.latt.Vec2();
double beta = data.latt.GetUnitCell().beta;
// Initial guess for the parameters
double beam[2] = {data.geom.GetBeamX_pxl(), data.geom.GetBeamY_pxl()};
double distance_mm = data.geom.GetDetectorDistance_mm();
double detector_rot[2] = {data.geom.GetPoniRot1_rad(), data.geom.GetPoniRot2_rad()};
ceres::Problem problem;
double latt_vec0[3] = {0.0, 0.0, 0.0};
double latt_vec1[3] = {0.0, 0.0, 0.0};
double latt_vec2[3] = {0.0, 0.0, 0.0};
double rot_vec[3] = {1, 0, 0};
switch (data.crystal_system) {
case gemmi::CrystalSystem::Orthorhombic:
LatticeToRodriguesAndLengths_GS(data.latt, latt_vec0, latt_vec1);
break;
case gemmi::CrystalSystem::Tetragonal:
LatticeToRodriguesAndLengths_GS(data.latt, latt_vec0, latt_vec1);
latt_vec1[0] = (latt_vec1[0] + latt_vec1[1]) / 2.0;
break;
case gemmi::CrystalSystem::Cubic:
LatticeToRodriguesAndLengths_GS(data.latt, latt_vec0, latt_vec1);
latt_vec1[0] = (latt_vec1[0] + latt_vec1[1] + latt_vec1[2]) / 3.0;
break;
case gemmi::CrystalSystem::Hexagonal:
LatticeToRodriguesAndLengths_Hex(data.latt, latt_vec0, latt_vec1);
break;
case gemmi::CrystalSystem::Monoclinic:
LatticeToRodriguesLengthsBeta_Mono(data.latt, latt_vec0, latt_vec1, beta);
latt_vec2[0] = beta;
latt_vec2[1] = 0.0;
latt_vec2[2] = 0.0;
break;
default:
// Triclinic: initialize a,b,c and α,β,γ from current unit cell
LatticeToRodriguesAndLengths_GS(data.latt, latt_vec0, latt_vec1);
auto uc = data.latt.GetUnitCell();
latt_vec2[0] = uc.alpha * PI / 180.0;
latt_vec2[1] = uc.beta * PI / 180.0;
latt_vec2[2] = uc.gamma * PI / 180.0;
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;
for (int i = 0; i < spots.size(); i++) {
if (spots[i].empty())
continue;
double angle_rad = 0.0;
std::optional<RotMatrix> rot_matr;
if (data.axis) {
const float angle_deg = data.axis->GetAngle_deg(i) + data.axis->GetWedge_deg() / 2.0;
angle_rad = angle_deg * PI / 180.0;
rot_matr = data.axis->GetTransformationAngle(angle_deg);
}
// Add residuals for each point
for (const auto &pt: spots[i]) {
if (!data.index_ice_rings && pt.ice_ring)
continue;
Coord recip = pt.ReciprocalCoord(data.geom);
if (rot_matr)
recip = rot_matr.value() * recip;
double h_fp = recip * vec0;
double k_fp = recip * vec1;
double l_fp = recip * vec2;
double h = std::round(h_fp);
double k = std::round(k_fp);
double l = std::round(l_fp);
double norm_sq = (h - h_fp) * (h - h_fp) + (k - k_fp) * (k - k_fp) + (l - l_fp) * (l - l_fp);
if (norm_sq > tolerance_sq)
continue;
problem.AddResidualBlock(
new ceres::AutoDiffCostFunction<XtalResidual, 3, 2, 1, 2, 3, 3, 3, 3>(
new XtalResidual(pt.x, pt.y,
data.geom.GetWavelength_A(),
data.geom.GetPixelSize_mm(),
data.geom.GetPoniRot3_rad(),
angle_rad,
h, k, l,
data.crystal_system)),
nullptr,
beam,
&distance_mm,
detector_rot,
rot_vec,
latt_vec0,
latt_vec1,
latt_vec2
);
}
}
if (problem.NumResidualBlocks() < data.min_spots)
return false;
if (!data.refine_distance_mm)
problem.SetParameterBlockConstant(&distance_mm);
else {
const double dist_range = 0.1;
problem.SetParameterLowerBound(&distance_mm, 0, distance_mm * (1.0 - dist_range));
problem.SetParameterUpperBound(&distance_mm, 0, distance_mm * (1.0 + dist_range));
}
if (!data.refine_beam_center)
problem.SetParameterBlockConstant(beam);
else if (data.axis) {
// Gauge handling (single-axis rotation): rotating the whole experiment about the spindle leaves every
// spot position unchanged, so the beam-centre component PARALLEL to the spindle is a null/gauge-weak
// direction. Refining it freely lets it wander (~+3 px) and absorb centroid systematics into a wrong
// beam that the co-refined orientation keeps position-consistent. Rather than freeze it (the beam
// does drift - it is only LaB6-monitored to ~a few px), RESTRAIN it toward the header with a soft
// prior: the gauge direction has ~zero data sensitivity so the prior pins it near the header, while a
// real, well-supported drift can still overcome it. The spindle is along a detector axis in standard
// geometry, so restrain the dominant of X / Y.
const Coord spindle = data.axis->GetAxis();
const int parallel = (std::fabs(spindle.x) >= std::fabs(spindle.y)) ? 0 : 1;
// Weight so the prior is a sigma_px-pixel restraint that competes with the positional residuals.
// k = d|recip|/d(beam_px) ~ pixel/(distance*lambda) [A^-1/px]; scaling by sqrt(#residuals) makes the
// prior's curvature ~ (1/9) of the well-constrained-data curvature at sigma_px=3, i.e. data wins the
// perpendicular direction, the prior wins the gauge one.
constexpr double sigma_px = 3.0;
const double k = data.geom.GetPixelSize_mm() / (distance_mm * data.geom.GetWavelength_A());
const double w = k * std::sqrt(static_cast<double>(problem.NumResidualBlocks())) / sigma_px;
problem.AddResidualBlock(
new ceres::AutoDiffCostFunction<BeamComponentPrior, 1, 2>(
new BeamComponentPrior(parallel, beam[parallel], w)),
nullptr, beam);
}
if (!data.refine_detector_angles) {
problem.SetParameterBlockConstant(detector_rot);
} else {
const double rot_range = 3.0 / 180.0 * PI;
for (int i = 0; i < 2; ++i) {
problem.SetParameterLowerBound(detector_rot, i, detector_rot[i] - rot_range);
problem.SetParameterUpperBound(detector_rot, i, detector_rot[i] + rot_range);
}
}
if (!data.refine_rotation_axis) {
problem.SetParameterBlockConstant(rot_vec);
}
if (!data.refine_unit_cell) {
problem.SetParameterBlockConstant(latt_vec1);
problem.SetParameterBlockConstant(latt_vec2);
} else {
// Parameter bounds
// Lengths
for (int i = 0; i < 3; ++i) {
problem.SetParameterLowerBound(latt_vec1, i, data.min_length_A);
problem.SetParameterUpperBound(latt_vec1, i, data.max_length_A);
}
if (data.crystal_system == gemmi::CrystalSystem::Monoclinic) {
const double beta_lo = std::max(1e-6, PI * (data.min_angle_deg / 180.0));
const double beta_hi = std::min(PI - 1e-6, PI * (data.max_angle_deg / 180.0));
problem.SetParameterLowerBound(latt_vec2, 0, beta_lo);
problem.SetParameterUpperBound(latt_vec2, 0, beta_hi);
} else if (data.crystal_system == gemmi::CrystalSystem::Triclinic) {
// α, β, γ bounds (radians)
const double alo = PI * (data.min_angle_deg / 180.0);
const double ahi = PI * (data.max_angle_deg / 180.0);
for (int i = 0; i < 3; ++i) {
problem.SetParameterLowerBound(latt_vec2, i, alo);
problem.SetParameterUpperBound(latt_vec2, i, ahi);
}
} else {
// Orthorhombic / Tetragonal / Cubic / Hexagonal:
// latt_vec2 has no meaning for these systems — always freeze it.
problem.SetParameterBlockConstant(latt_vec2);
}
}
// Configure solver
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 = num_threads; // usually 1 (called from many threads); caller may raise it
ceres::Solver::Summary summary;
// Run optimization
ceres::Solve(options, &problem, &summary);
if (data.refine_beam_center) {
data.beam_corr_x = data.geom.GetBeamX_pxl() - beam[0];
data.beam_corr_y = data.geom.GetBeamY_pxl() - beam[1];
data.geom.BeamX_pxl(beam[0]).BeamY_pxl(beam[1]);
}
if (data.refine_distance_mm)
data.geom.DetectorDistance_mm(distance_mm);
if (data.refine_detector_angles)
data.geom.PoniRot1_rad(detector_rot[0]).PoniRot2_rad(detector_rot[1]);
if (data.axis && data.refine_rotation_axis)
data.axis.value().Axis(Coord(rot_vec[0], rot_vec[1], rot_vec[2]));
if (data.crystal_system == gemmi::CrystalSystem::Orthorhombic)
data.latt = AngleAxisAndCellToLattice(latt_vec0, latt_vec1, PI / 2.0, PI / 2.0, PI / 2.0);
else if (data.crystal_system == gemmi::CrystalSystem::Tetragonal) {
latt_vec1[1] = latt_vec1[0];
data.latt = AngleAxisAndCellToLattice(latt_vec0, latt_vec1, PI / 2.0, PI / 2.0, PI / 2.0);
} else if (data.crystal_system == gemmi::CrystalSystem::Cubic) {
latt_vec1[1] = latt_vec1[0];
latt_vec1[2] = latt_vec1[0];
data.latt = AngleAxisAndCellToLattice(latt_vec0, latt_vec1, PI / 2.0, PI / 2.0, PI / 2.0);
} else if (data.crystal_system == gemmi::CrystalSystem::Hexagonal) {
latt_vec1[1] = latt_vec1[0];
data.latt = AngleAxisAndCellToLattice(latt_vec0, latt_vec1,PI / 2.0, PI / 2.0, 2.0 * PI / 3.0);
} else if (data.crystal_system == gemmi::CrystalSystem::Monoclinic) {
data.latt = AngleAxisAndCellToLattice(latt_vec0, latt_vec1, PI / 2.0, latt_vec2[0], PI / 2.0);
} else {
// Triclinic via the same generic builder
data.latt = AngleAxisAndCellToLattice(latt_vec0, latt_vec1, latt_vec2[0], latt_vec2[1], latt_vec2[2]);
}
return true;
} catch (...) {
// Convergence problems, likely not updated
return false;
}
}
bool XtalOptimizer(XtalOptimizerData &data, const std::vector<std::vector<SpotToSave>> &spots,
int num_threads) {
if (!XtalOptimizerInternal(data, spots, 0.3, num_threads))
return false;
XtalOptimizerInternal(data, spots, 0.2, num_threads);
return XtalOptimizerInternal(data, spots, 0.1, num_threads);
}
bool XtalOptimizerRotationOnly(XtalOptimizerData &data,
const std::vector<SpotToSave> &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};
// Spot selection by current indexing (same approach as XtalOptimizerInternal)
const Coord a0 = data.latt.Vec0();
const Coord b0 = data.latt.Vec1();
const Coord c0 = 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<double>(recip_index * a0);
const double k_fp = static_cast<double>(recip_index * b0);
const double l_fp = static_cast<double>(recip_index * c0);
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<double>(tol_sq))
continue;
// s_obs must be in the same reference frame as the
// predicted reciprocal vector (h·a* + k·b* + l·c*), which is the
// phi=0 crystal frame. Apply the same goniometer back-rotation
// that was used above for the HKL assignment.
Coord s_obs = data.geom.DetectorToRecip(pt.x, pt.y);
if (data.axis.has_value())
s_obs = data.axis->GetTransformationAngle(pt.phi) * s_obs;
auto *cost =
new ceres::AutoDiffCostFunction<XtalResidualRotationOnlyPrecomp, 3, 3>(
new XtalResidualRotationOnlyPrecomp(s_obs, data.latt, h, k, l)
);
problem.AddResidualBlock(cost, nullptr, rot_aa);
}
if (problem.NumResidualBlocks() < data.min_spots)
return false;
// Regularization: prefer the smallest rotation correction that fits the
// data. This is essential when spots are nearly coplanar in reciprocal
// space (e.g. still images), where the rotation component perpendicular
// to the scattering plane is otherwise underdetermined.
// The weight is in Å⁻¹ rad⁻¹; tune relative to your typical residual.
{
const double reg_weight = 0.05; // e.g. 0.05
problem.AddResidualBlock(
new ceres::AutoDiffCostFunction<RotationNormRegularizer, 3, 3>(
new RotationNormRegularizer(reg_weight)),
nullptr, rot_aa);
}
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 direct-lattice vectors.
// ceres::AngleAxisToRotationMatrix writes a **row-major** 3×3 matrix,
// and Eigen's << operator also fills row-by-row, so the assignment
// below is correct without any transposing.
//
// Note: for a pure orthogonal rotation R, R⁻ᵀ = R, so rotating the
// direct-lattice vectors (A, B, C) by R is exactly equivalent to
// rotating the reciprocal vectors (a*, b*, c*) by the same R. No
// transpose or inversion of R is needed here.
double R_raw[9];
ceres::AngleAxisToRotationMatrix(rot_aa, R_raw); // row-major 3x3
Eigen::Matrix3d R;
R << R_raw[0], R_raw[3], R_raw[6],
R_raw[1], R_raw[4], R_raw[7],
R_raw[2], R_raw[5], 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<float>(A2.x()), static_cast<float>(A2.y()), static_cast<float>(A2.z())),
Coord(static_cast<float>(B2.x()), static_cast<float>(B2.y()), static_cast<float>(B2.z())),
Coord(static_cast<float>(C2.x()), static_cast<float>(C2.y()), static_cast<float>(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;
}
}