leonarski_fandClaude Opus 5
641f890a40
The per-image geometry refinement is the largest stage of the image loop, and a third of it was arithmetic on numbers that never change. The residual derives the detector angles' sines and cosines, the goniometer's back-rotation - a three-argument hypot, a sine, a cosine and a division - and the reciprocal basis of the cell on every evaluation. On the rotation path the detector angles and the axis are held fixed and stored as plain doubles, so all of it is constant, not merely constant per block: there is one frame per image and one cell. Three solves an image, fifty iterations a solve and a thousand spots make it tens of thousands of repetitions of the same result. The frame's constants are now built once and handed in. The body they feed is the same body, split out rather than copied, so no expression is reassociated - in particular the reciprocal vector is still formed as the basis times the inverse volume, with the volume not folded into the basis. The spot confidence weights depend only on each spot's resolution and intensity, which no solver touches, and were recomputed identically for each of the three passes. They are computed once. The sort behind them ordered indices through a projection that chased a random eighty-byte-strided element per comparison; it now sorts a packed resolution and index, which makes the same comparisons in the same sequence and therefore the same permutation. The spot list itself was copied per image through an initializer list whose elements are const; it is passed as a view. The integration engine was the last one in the loop copying through pageable host memory - three transfers in and eight out per image, twenty-six bytes a reflection, while every other engine already page-locks its staging. A driver copy from pageable memory stages through its own pinned buffer on the calling thread, which is why an asynchronous copy was averaging a hundred and thirteen microseconds. Page-locked, the same seventeen thousand calls cost four hundred and thirty-two milliseconds instead of one and a half seconds, and the wait moves to the synchronisation point where it belongs. Two smaller ones: the reflections were copied into the per-image message for a process file that a merging run does not write, so the copy is made where a writer exists; and the intensity statistics and the Wilson estimate walked the same eighty-byte array twice to read twelve bytes, which is now one pass with each accumulation in its own order. Every reflection file is byte-identical on four crystals; the process file's reflections match dataset for dataset, and its azimuthal arrays differ no more between this build and the last than the last differs from itself. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01EGpGdgmJ8MyY9pCGWjktyi
607 lines
28 KiB
C++
607 lines
28 KiB
C++
// 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 <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];
|
||
|
||
const AngleAxisRotator<T> rot(rot_aa);
|
||
rot.Rotate(astar_unrot, astar_rot);
|
||
rot.Rotate(bstar_unrot, bstar_rot);
|
||
rot.Rotate(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.01–0.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;
|
||
|
||
// Resolution order. Sorting a packed (resolution, index) array rather than an index vector with a
|
||
// projection into the spots keeps the comparisons off the 80-byte records - the same keys in the
|
||
// same order, so introsort makes the same comparisons and the same swaps, and the order it leaves
|
||
// is the same.
|
||
struct SpotByRes {
|
||
float d_A;
|
||
uint32_t index;
|
||
};
|
||
std::vector<SpotByRes> by_res(spots.size());
|
||
for (size_t i = 0; i < spots.size(); i++)
|
||
by_res[i] = {spots[i].d_A, static_cast<uint32_t>(i)};
|
||
std::ranges::sort(by_res, {}, &SpotByRes::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].index].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].index].intensity) / median;
|
||
weight[by_res[i].index] = std::sqrt(r / (1.0 + r));
|
||
}
|
||
}
|
||
return weight;
|
||
}
|
||
|
||
bool XtalOptimizerInternal(XtalOptimizerData &data,
|
||
std::span<const std::vector<SpotToSave>> spots,
|
||
const std::vector<std::vector<double>> &weights,
|
||
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()};
|
||
|
||
// The per-frame constants of the reduced residual (see XtalFrameConstants), one entry per frame
|
||
// that contributes. Reserved up front and never grown past that, so the residual blocks' pointers
|
||
// into it stay valid, and declared before the problem so that it outlives it.
|
||
std::vector<XtalFrameConstants> frame_const;
|
||
frame_const.reserve(spots.size());
|
||
|
||
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;
|
||
|
||
// The same for every spot of every frame, so taken once here rather than per residual.
|
||
const double cos_rot3 = std::cos(data.geom.GetPoniRot3_rad());
|
||
const double sin_rot3 = std::sin(data.geom.GetPoniRot3_rad());
|
||
|
||
// Per-image rotation refinement frees only the beam and the orientation and holds the other five
|
||
// blocks constant, so the seven-block residual makes Ceres differentiate 17 parameters to use 5.
|
||
// Where that is the configuration, use the reduced residual instead - identical fit, Jet<5>
|
||
// autodiff. Any other combination (stills also free the cell, the offline refiner frees distance
|
||
// and detector angles) keeps the general form below.
|
||
const bool beam_and_orientation_only = data.refine_beam_center
|
||
&& !data.refine_detector_angles
|
||
&& !data.refine_rotation_axis
|
||
&& !data.refine_unit_cell;
|
||
|
||
// 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;
|
||
|
||
const std::vector<double> &weight = weights[i]; // empty = unweighted
|
||
|
||
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);
|
||
}
|
||
|
||
if (beam_and_orientation_only)
|
||
frame_const.emplace_back(detector_rot, rot_vec, angle_rad, latt_vec1, latt_vec2,
|
||
data.crystal_system);
|
||
|
||
// 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;
|
||
|
||
const XtalResidual residual(pt.x, pt.y,
|
||
data.geom.GetWavelength_A(),
|
||
data.geom.GetPixelSize_mm(),
|
||
cos_rot3, sin_rot3,
|
||
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.
|
||
ceres::LossFunction *loss = weight.empty()
|
||
? nullptr
|
||
: new ceres::ScaledLoss(nullptr, weight_sq,
|
||
ceres::TAKE_OWNERSHIP);
|
||
|
||
if (beam_and_orientation_only)
|
||
problem.AddResidualBlock(
|
||
new ceres::AutoDiffCostFunction<XtalResidualBeamOrientation, 3, 2, 3>(
|
||
new XtalResidualBeamOrientation(residual, distance_mm, frame_const.back())),
|
||
loss,
|
||
beam,
|
||
latt_vec0
|
||
);
|
||
else
|
||
problem.AddResidualBlock(
|
||
new ceres::AutoDiffCostFunction<XtalResidualFixedDistance, 3, 2, 2, 3, 3, 3, 3>(
|
||
new XtalResidualFixedDistance(residual, distance_mm)),
|
||
loss,
|
||
beam,
|
||
detector_rot,
|
||
rot_vec,
|
||
latt_vec0,
|
||
latt_vec1,
|
||
latt_vec2
|
||
);
|
||
}
|
||
}
|
||
|
||
if (problem.NumResidualBlocks() < data.min_spots)
|
||
return false;
|
||
|
||
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);
|
||
}
|
||
|
||
// Distance, detector angles, rotation axis and cell are parameter blocks only in the general
|
||
// seven-block residual; the reduced one bakes them in, so there is nothing left to configure.
|
||
if (!beam_and_orientation_only) {
|
||
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;
|
||
// Normal equations, not QR. The problem is very tall and thin - thousands of spots against at
|
||
// most 17 parameters - and that is the shape DENSE_QR handles worst: it copies the Jacobian out
|
||
// of Ceres' row-major storage into a column-major buffer on every solve, and Eigen's blocked
|
||
// Householder then degenerates to the unblocked path because its block size is min(48, columns).
|
||
// Accumulating J^T J reads the Jacobian once instead. Both solve the same damped system, so the
|
||
// step is the same to round-off; the column scaling Ceres applies by default and the LM diagonal
|
||
// keep the squared condition number in hand.
|
||
options.linear_solver_type = ceres::DENSE_NORMAL_CHOLESKY;
|
||
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_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, std::span<const std::vector<SpotToSave>> spots,
|
||
int num_threads) {
|
||
// A spot's confidence weight is set by its resolution and its intensity, neither of which the solver
|
||
// touches, so the three passes below all get the same weights: take them once.
|
||
std::vector<std::vector<double>> weights(spots.size());
|
||
if (data.weight_spots_by_confidence)
|
||
for (size_t i = 0; i < spots.size(); i++)
|
||
if (!spots[i].empty())
|
||
weights[i] = SpotConfidenceWeights(spots[i]);
|
||
|
||
if (!XtalOptimizerInternal(data, spots, weights, 0.3, num_threads))
|
||
return false;
|
||
XtalOptimizerInternal(data, spots, weights, 0.2, num_threads);
|
||
return XtalOptimizerInternal(data, spots, weights, 0.1, num_threads);
|
||
}
|
||
|
||
bool XtalOptimizer(XtalOptimizerData &data, const std::vector<SpotToSave> &spots, int num_threads) {
|
||
return XtalOptimizer(data, std::span(&spots, 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_NORMAL_CHOLESKY; // tall and thin, as above
|
||
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;
|
||
}
|
||
} |