Files
Jungfraujoch/image_analysis/geom_refinement/XtalOptimizer.cpp
T
leonarski_fandClaude Opus 5 b161da05c1 Geometry refinement: bound offline reprocessing by iterations, not the clock
The per-image refinement stopped on a wall-clock budget (40 ms, and 20 ms for
the rotation-only extra pass). Online that is exactly right - the budget is real
and an image that overruns it costs the acquisition. Offline it means the same
file refines to a different lattice depending on what else the machine was doing
at the time, which is not a property reprocessing should have.

Bound it by iteration count instead when the caller is offline. IndexAndRefine
takes the workflow as a constructor argument: the receiver asks for the
wall-clock bound, rugnux and the viewer get the reproducible one. 50 iterations
is Ceres' own default; the per-image problem converges well inside it, so it
bounds the pathological case rather than the normal one - measured on five
battery crystals, every number is unchanged from the timed version.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
2026-07-30 20:42:04 +02:00

539 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);
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);
// 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;
}
}