Files
Jungfraujoch/image_analysis/geom_refinement/XtalOptimizer.cpp
T
leonarski_fandClaude Opus 5 164f15c903 geom_refinement: stop committing refinements that did not converge
Four of the seven ceres::Solve calls in image_analysis obtained a
Solver::Summary and never looked at it, so a solve that failed numerically
had its parameters written back and was reported as success.
StillsPartialityRefine and both PostRefine solves already gated on
IsSolutionUsable(); this brings the rest to the same contract.

IsSolutionUsable() is the right test rather than checking for CONVERGENCE:
it accepts a solve that ran out of iterations or wall-clock time but still
descended, which is exactly what the real-time callers depend on when they
set max_solver_time instead of max_num_iterations. Only FAILURE and
USER_FAILURE are rejected.

XtalOptimizer checks before the write-back, so a failed refinement now
leaves the caller's geom and latt untouched instead of half-updated.
GeometryRefiner folds it into result.ok, which previously reported success
from spot and frame counts alone. RingOptimizer returns a geometry by
value that both callers assign straight back over their input, so it hands
back the unchanged reference rather than a diverged beam centre.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
2026-07-31 15:02:11 +02:00

549 lines
24 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 <algorithm>
#include <numeric>
#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;
};
// Prior confidence weight per spot: how strong the spot is FOR ITS RESOLUTION. The frame's spots are
// ordered by resolution and cut into equal-count shells, and each intensity is divided by its shell
// median. Refinement needs the high-resolution spots (they carry the cell and distance information) and
// those are legitimately weaker, so a raw intensity weight would suppress exactly the wrong ones; the
// shell normalisation makes the weight resolution-neutral by construction.
//
// The weight enters as w^2 on the squared residual, w^2 = r/(1+r): the shell median contributes half,
// a 4x-median spot 0.8, a quarter-median spot 0.2. Weak spots still pull, they just do not drive. Unlike
// a robust loss this is a PRIOR - it never looks at the current residual, so it cannot mistake a genuine
// spot for an outlier when the starting geometry is far off and leave the fit unable to move.
static std::vector<double> SpotConfidenceWeights(const std::vector<SpotToSave> &spots) {
constexpr size_t spots_per_shell = 32;
std::vector<size_t> by_res(spots.size());
std::iota(by_res.begin(), by_res.end(), 0);
std::ranges::sort(by_res, {}, [&](size_t i) { return spots[i].d_A; });
const size_t nshells = std::max<size_t>(1, spots.size() / spots_per_shell);
std::vector<double> weight(spots.size());
std::vector<float> shell_intensity;
for (size_t s = 0; s < nshells; s++) {
const size_t begin = s * spots.size() / nshells;
const size_t end = (s + 1) * spots.size() / nshells;
shell_intensity.clear();
for (size_t i = begin; i < end; i++)
shell_intensity.push_back(spots[by_res[i]].intensity);
std::ranges::nth_element(shell_intensity, shell_intensity.begin() + shell_intensity.size() / 2);
const double median = std::max(1e-3f, shell_intensity[shell_intensity.size() / 2]);
for (size_t i = begin; i < end; i++) {
const double r = std::max(0.0f, spots[by_res[i]].intensity) / median;
weight[by_res[i]] = std::sqrt(r / (1.0 + r));
}
}
return 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;
// Sum of w^2 over the spots that entered - the beam prior below is scaled by it so that its
// strength relative to the data is the same weighted or not. Equals the residual block count
// when the spots are unweighted.
double effective_spots = 0.0;
for (int i = 0; i < spots.size(); i++) {
if (spots[i].empty())
continue;
std::vector<double> weight; // empty = unweighted
if (data.weight_spots_by_confidence)
weight = SpotConfidenceWeights(spots[i]);
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 (size_t j = 0; j < spots[i].size(); j++) {
const auto &pt = spots[i][j];
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;
const double weight_sq = weight.empty() ? 1.0 : weight[j] * weight[j];
effective_spots += weight_sq;
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)),
// Ceres has no per-residual weight; ScaledLoss(nullptr, a) multiplies the squared
// residual by the constant a, i.e. it applies a weight of sqrt(a) to the residual.
weight.empty()
? nullptr
: new ceres::ScaledLoss(nullptr, weight_sq, ceres::TAKE_OWNERSHIP),
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(effective_spots) / 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;
if (data.max_iterations > 0)
options.max_num_iterations = data.max_iterations;
else
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);
// Only a genuine numerical failure is rejected here: a solve that ran out of iterations or
// out of time but still descended counts as usable, which is what the real-time caller
// relies on when it sets max_solver_time. Checked before anything is written back, so a
// failed refinement leaves data untouched rather than committing half a fit.
if (!summary.IsSolutionUsable())
return false;
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;
if (data.max_iterations > 0)
options.max_num_iterations = data.max_iterations;
else
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);
if (!summary.IsSolutionUsable())
return false;
// 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;
}
}