Files
Jungfraujoch/image_analysis/geom_refinement/PostRefine.cpp
T
jungfrauandClaude Opus 5 8ae53b7fdf Say what actually makes the post-refine bucket sort reproducible
The bucketing commit justified itself with "the partials order became total in an earlier commit",
which is true of the scale/merge ingest and not of this sort: part_less ends at the image number,
so two partials of one reflection on one image tie, exactly as they did before.

Nothing is wrong with the result. The counting-sort prefix lays each bucket out chunk by chunk,
and a chunk is a contiguous span of the gathered order, so every bucket arrives at std::sort in
global gather order no matter how many threads scattered it - the order is reproducible run to run
and identical across -N. Giving Partial a rank field to make the comparator total would settle
those ties by index instead, at eight more bytes on an array that reaches tens of millions of
elements, and would change nothing anyone can observe.

So state the invariant where the comparator is, rather than leaving the next reader to trust a
claim that does not hold for this half.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_011n8riB6X59oRjkrSHzNPAU
2026-08-23 08:23:57 -04:00

642 lines
40 KiB
C++

// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only
#include "../../common/ParallelFor.h"
#include "PostRefine.h"
#include <algorithm>
#include <array>
#include <cmath>
#include <limits>
#include <numeric>
#include "../../common/JFJochMath.h" // PI
#include "XtalResidual.h" // XtalResidual (the positional detector<->reciprocal residual, step B)
#include "LatticeReduction.h"
#include "ceres/ceres.h"
#include "ceres/rotation.h"
namespace {
// One integrated partial, flattened across all images.
struct Partial {
int h, k, l;
float img;
double I, sigma;
double angle_rad; // frame mid-exposure goniometer angle
double obs_x, obs_y; // observed spot centroid (pixels); NAN if the box sum found no centroid
};
// A rocking event and its precomputed reference reciprocal vector (phi=0 frame, from the indexed lattice).
struct Event {
double phi_obs; // rad, intensity-weighted rocking centroid
double weight; // sqrt(sum I / sum sigma)
double e_ref[3]; // h*a* + k*b* + l*c* at the reference (unrefined) cell/orientation
int h, k, l;
};
// Distance-INDEPENDENT Ewald excitation residual for a uniform cell-scale parameter s and a refined
// goniometer axis (3-vector). The header-distance miscalibration leaves a uniform cell scale; the axis is
// the other phi_obs lever. Both are phi_obs-constrained (distance-independent). e_ref is the reference
// reciprocal (h*a* + k*b* + l*c* at the indexed cell). On the Ewald sphere <=> |p|^2 + 2 p_z/lambda == 0.
struct ScaleAxisExcitationResidual {
ScaleAxisExcitationResidual(double lambda, double angle_rad, double weight, const double e_ref[3])
: inv_lambda(1.0 / lambda), angle_rad(angle_rad), weight(weight),
ex(e_ref[0]), ey(e_ref[1]), ez(e_ref[2]) {}
template<typename T>
bool operator()(const T *const s, const T *const axis, T *residual) const {
const T inv_s = T(1) / s[0];
const T p_ref[3] = {T(ex) * inv_s, T(ey) * inv_s, T(ez) * inv_s};
const T aa[3] = {T(-angle_rad) * axis[0], T(-angle_rad) * axis[1], T(-angle_rad) * axis[2]};
T p_lab[3];
ceres::AngleAxisRotatePoint(aa, p_ref, p_lab);
const T zeta = p_lab[0] * p_lab[0] + p_lab[1] * p_lab[1] + p_lab[2] * p_lab[2]
+ T(2.0) * p_lab[2] * T(inv_lambda);
residual[0] = T(weight) * zeta * T(0.5) / T(inv_lambda);
return true;
}
const double inv_lambda, angle_rad, weight, ex, ey, ez;
};
// GONIOMETER ROTATION SCALE k: the same Ewald excitation residual, but with the cell scale and the axis
// DIRECTION already committed by step A, so the single free quantity is how far the stage actually turned
// per unit of commanded angle. Two differences from step A matter:
// * the angle is measured from the CENTRE of the sweep, not from the goniometer's zero. The reference
// orientation is the one rotation indexing fitted against the commanded angles, so it has already
// absorbed the MEAN angle error; only the part that varies across the sweep is left to fit. Scaling the
// absolute angle instead - which is what reading k off the length of step A's axis vector does - asks
// the fit to also produce a constant offset it has no parameter for, and the least-squares compromise
// shrinks k towards 1 by var(phi) / (var(phi) + phi_centre^2): exactly a factor of four for the common
// case of a sweep starting at zero.
// * e_mid is the reference reciprocal vector already turned to the sweep centre and divided by the
// committed cell scale, so nothing but k is free.
struct RotationScaleResidual {
RotationScaleResidual(double lambda, double dangle_rad, const double u[3], const double e_mid[3])
: inv_lambda(1.0 / lambda), dangle_rad(dangle_rad),
ux(u[0]), uy(u[1]), uz(u[2]), ex(e_mid[0]), ey(e_mid[1]), ez(e_mid[2]) {}
template<typename T>
bool operator()(const T *const k, T *residual) const {
const T a = T(-dangle_rad) * k[0];
const T aa[3] = {a * T(ux), a * T(uy), a * T(uz)};
const T p_ref[3] = {T(ex), T(ey), T(ez)};
T p_lab[3];
ceres::AngleAxisRotatePoint(aa, p_ref, p_lab);
const T zeta = p_lab[0] * p_lab[0] + p_lab[1] * p_lab[1] + p_lab[2] * p_lab[2]
+ T(2.0) * p_lab[2] * T(inv_lambda);
residual[0] = zeta * T(0.5) / T(inv_lambda);
return true;
}
const double inv_lambda, dangle_rad, ux, uy, uz, ex, ey, ez;
};
} // namespace
PostRefineResult PostRefineRotationGeometry(const std::vector<IntegrationOutcome> &outcomes,
const GoniometerAxis &axis,
const DiffractionGeometry &nominal_geom,
const CrystalLattice &reference_latt,
const PostRefineSettings &settings,
Logger &logger) {
PostRefineResult result;
result.geom = nominal_geom;
result.cell = reference_latt.GetUnitCell();
result.distance_before_mm = nominal_geom.GetDetectorDistance_mm();
result.distance_after_mm = nominal_geom.GetDetectorDistance_mm();
try {
const double wedge_half = axis.GetWedge_deg() / 2.0;
const double lambda = nominal_geom.GetWavelength_A();
const Coord ax = axis.GetAxis();
const Coord Astar = reference_latt.Astar(), Bstar = reference_latt.Bstar(), Cstar = reference_latt.Cstar();
// Count first, then fill. Growing one vector by push_back over tens of millions of
// reflections copies the whole thing every time it doubles - several gigabytes of pure
// copying - and the counts are cheap to take. Each outcome then owns a slice, so the fill
// runs on all threads and lands in the order the serial loop produced.
const size_t nthreads = std::max(1, settings.num_threads);
const int n_out = static_cast<int>(outcomes.size());
std::vector<size_t> pts_offset(n_out + 1, 0);
ParallelChunks(n_out, nthreads, [&](int lo, int hi) {
for (int o = lo; o < hi; o++) {
size_t keep = 0;
for (const auto &r : outcomes[o].reflections)
if (std::isfinite(r.I) && std::isfinite(r.sigma) && r.sigma > 0.0f)
keep++;
pts_offset[o + 1] = keep;
}
});
for (int o = 0; o < n_out; o++)
pts_offset[o + 1] += pts_offset[o];
std::vector<Partial> pts(pts_offset[n_out]);
ParallelChunks(n_out, nthreads, [&](int lo, int hi) {
for (int o = lo; o < hi; o++) {
size_t at = pts_offset[o];
for (const auto &r : outcomes[o].reflections) {
if (!std::isfinite(r.I) || !std::isfinite(r.sigma) || r.sigma <= 0.0f) continue;
const double mid_deg = axis.GetAngle_deg(r.image_number) + wedge_half;
const double ox = std::isfinite(r.observed_x) ? r.observed_x : NAN;
const double oy = std::isfinite(r.observed_y) ? r.observed_y : NAN;
pts[at++] = Partial{r.h, r.k, r.l, r.image_number, r.I, r.sigma,
mid_deg * PI / 180.0, ox, oy};
}
}
});
logger.Info("Post-refine: {} partials gathered", pts.size());
if (pts.size() < static_cast<size_t>(settings.min_events)) return result;
// Bucket by h, then sort the buckets. h is the leading key, so the sorted array is the
// buckets laid end to end, and each bucket sorts on its own thread. Sorting the whole thing
// in one pass moved 56 bytes per element through every level of a comparison sort, on one
// thread, over tens of millions of reflections.
{
// Not a total order: two partials of one reflection on one image still tie, as they did
// before this was bucketed. What makes the result reproducible is the scatter below
// rather than the comparator - the prefix lays each bucket out chunk by chunk and a
// chunk is a contiguous span of the gathered order, so a bucket reaches std::sort in
// global gather order whatever the thread count. Ties therefore resolve the same way on
// every run and at every -N; they are simply not resolved by rank.
const auto part_less = [](const Partial &a, const Partial &b) {
if (a.h != b.h) return a.h < b.h;
if (a.k != b.k) return a.k < b.k;
if (a.l != b.l) return a.l < b.l;
return a.img < b.img;
};
const int n = static_cast<int>(pts.size());
int h_lo = std::numeric_limits<int>::max(), h_hi = std::numeric_limits<int>::min();
for (const auto &q : pts) { h_lo = std::min(h_lo, q.h); h_hi = std::max(h_hi, q.h); }
const int H = (h_lo <= h_hi) ? (h_hi - h_lo + 1) : 1;
const int nt = static_cast<int>(std::clamp<size_t>(nthreads, 1, std::max(1, n)));
const int chunk = (n + nt - 1) / nt;
std::vector<std::vector<int32_t>> hist(nt, std::vector<int32_t>(H, 0));
ParallelChunks(nt, nthreads, [&](int tlo, int thi) {
for (int t = tlo; t < thi; ++t) {
const int lo = t * chunk, hi = std::min(n, lo + chunk);
for (int i = lo; i < hi; ++i) hist[t][pts[i].h - h_lo]++;
}
});
std::vector<int32_t> bstart(H + 1, 0);
int32_t acc = 0;
for (int b = 0; b < H; ++b) {
bstart[b] = acc;
for (int t = 0; t < nt; ++t) { const int32_t c = hist[t][b]; hist[t][b] = acc; acc += c; }
}
bstart[H] = acc;
std::vector<Partial> sorted(pts.size());
ParallelChunks(nt, nthreads, [&](int tlo, int thi) {
for (int t = tlo; t < thi; ++t) {
std::vector<int32_t> fill = hist[t];
const int lo = t * chunk, hi = std::min(n, lo + chunk);
for (int i = lo; i < hi; ++i) sorted[fill[pts[i].h - h_lo]++] = pts[i];
}
});
std::vector<int> order(H);
std::iota(order.begin(), order.end(), 0);
std::sort(order.begin(), order.end(),
[&](int a, int b) { return (bstart[a + 1] - bstart[a]) > (bstart[b + 1] - bstart[b]); });
ParallelFor(H, nthreads, [&](int oi) {
const int b = order[oi];
std::sort(sorted.begin() + bstart[b], sorted.begin() + bstart[b + 1], part_less);
});
pts.swap(sorted);
}
// Split into rocking events (same raw hkl, adjacent frames). Only >=2-frame events carry an
// unbiased phi_obs (a single-frame centroid is just the frame centre); precompute e_ref per event.
constexpr float MAX_FRAME_GAP = 2.0f;
std::vector<Event> events;
size_t i = 0, event_frames = 0;
while (i < pts.size()) {
size_t j = i + 1;
while (j < pts.size() && pts[j].h == pts[i].h && pts[j].k == pts[i].k && pts[j].l == pts[i].l
&& pts[j].img - pts[j - 1].img <= MAX_FRAME_GAP)
++j;
if (j - i >= 2) {
double sumI = 0, sumIphi = 0, sumSig = 0;
for (size_t m = i; m < j; ++m) {
const double Ipos = std::max(0.0, pts[m].I);
sumI += Ipos; sumIphi += Ipos * pts[m].angle_rad; sumSig += pts[m].sigma;
}
if (sumI > 0.0 && sumSig > 0.0) {
const double phi = sumIphi / sumI;
event_frames += j - i;
const Coord e = Astar * static_cast<float>(pts[i].h) + Bstar * static_cast<float>(pts[i].k)
+ Cstar * static_cast<float>(pts[i].l);
events.push_back(Event{phi, std::sqrt(sumI / sumSig), {e.x, e.y, e.z},
pts[i].h, pts[i].k, pts[i].l});
}
}
i = j;
}
// Frames per event is the phi_obs sampling: near 2 the reflections barely rock, so the angle
// this refinement is fitted to is under-determined. It is a geometry count, so unlike an
// intensity-weighted width it cannot be inflated by noise.
logger.Info("Post-refine: {} multi-frame rocking events ({:.1f} frames per event)", events.size(),
events.empty() ? 0.0 : static_cast<double>(event_frames) / events.size());
if (static_cast<int>(events.size()) < settings.min_events) return result;
// The rotation-scale fit further down is a single scalar whose whole point is how the residual
// varies ALONG the sweep, so it keeps every event. The cap below ranks by I/sigma, and on the
// crystals that have a stage fault the strong events sit in the middle of the sweep - the part
// that still indexes - so a capped set would leave the ends unrepresented in exactly the fit that
// has to see them.
const std::vector<Event> scale_events = events;
constexpr size_t MAX_EVENTS = 20000;
if (events.size() > MAX_EVENTS) {
std::nth_element(events.begin(), events.begin() + MAX_EVENTS, events.end(),
[](const Event &a, const Event &b) { return a.weight > b.weight; });
events.resize(MAX_EVENTS);
}
// ---- GEOMETRY REFINEMENT: the XtalOptimizer-equivalent, done as TWO SEPARATE
// cross-validated steps rather than one joint fit (the same lesson as integration: refining the
// profile width and the scale jointly fails, refining them separately works). Each step is committed
// only if it lowers a HELD-OUT (deterministic split-half) residual - otherwise that part of the
// geometry is left at nominal ("quit when things go wrong"):
// Step A: cell scale + rotation axis from phi_obs (distance-independent excitation residual).
// Step B: detector distance + beam centre from the observed spot positions, with the cell FIXED at
// step A (so the positional residual is no longer degenerate with the cell scale).
// Detector tilt is held fixed (gauge-coupled to orientation on a single crystal). ----
if (settings.refine_geometry) {
const gemmi::CrystalSystem sys =
(settings.crystal_system == gemmi::CrystalSystem::Trigonal) ? gemmi::CrystalSystem::Hexagonal
: settings.crystal_system;
const double ax0[3] = {ax.x, ax.y, ax.z};
const double lambda_l = lambda;
const double rot3 = nominal_geom.GetPoniRot3_rad();
const double pixel_mm = nominal_geom.GetPixelSize_mm();
const double det_rot[2] = {nominal_geom.GetPoniRot1_rad(), nominal_geom.GetPoniRot2_rad()};
const UnitCell r0 = reference_latt.GetUnitCell();
// Deterministic split of the reflections into a fit half and a held-out half. Avalanche-mix the
// hkl hash so the split bit is decorrelated from the LSB - a plain h+k+l parity collides with the
// lattice centering condition (e.g. an I-centred lattice has h+k+l even for EVERY present
// reflection, so a parity split would leave the validation half empty).
auto is_val = [](int h, int k, int l) {
unsigned u = static_cast<unsigned>(h) * 2654435761u + static_cast<unsigned>(k) * 2246822519u
+ static_cast<unsigned>(l) * 3266489917u;
u ^= u >> 15; u *= 2246822519u; u ^= u >> 13;
return (u & 1u) != 0u;
};
enum Subset { FIT, VAL, ALL };
auto in = [&](int h, int k, int l, Subset s) {
return s == ALL || (is_val(h, k, l) == (s == VAL)); };
// ===== Step A: cell scale s + rotation axis from phi_obs =====
auto excit_cost = [&](Subset s, double sc, const double axv[3]) {
double c = 0.0; int n = 0;
for (const auto &ev : events) {
if (!in(ev.h, ev.k, ev.l, s)) continue;
ScaleAxisExcitationResidual r(lambda_l, ev.phi_obs, 1.0, ev.e_ref);
double sd = sc, av[3] = {axv[0], axv[1], axv[2]}, resid = 0.0;
r(&sd, av, &resid); c += resid * resid; ++n;
}
return n ? c / n : 0.0;
};
auto solve_scale_axis = [&](Subset s, double &s_out, double ax_out[3]) {
double sc = 1.0, axv[3] = {ax0[0], ax0[1], ax0[2]};
ceres::Problem p;
for (const auto &ev : events) {
if (!in(ev.h, ev.k, ev.l, s)) continue;
p.AddResidualBlock(new ceres::AutoDiffCostFunction<ScaleAxisExcitationResidual, 1, 1, 3>(
new ScaleAxisExcitationResidual(lambda_l, ev.phi_obs, settings.excitation_weight, ev.e_ref)),
new ceres::CauchyLoss(0.02), &sc, axv);
}
p.SetParameterLowerBound(&sc, 0, 0.9); p.SetParameterUpperBound(&sc, 0, 1.1);
for (int j = 0; j < 3; ++j) { p.SetParameterLowerBound(axv, j, ax0[j] - 0.05);
p.SetParameterUpperBound(axv, j, ax0[j] + 0.05); }
ceres::Solver::Options o; o.linear_solver_type = ceres::DENSE_QR; o.max_num_iterations = 50;
o.num_threads = std::max(1, settings.num_threads); o.logging_type = ceres::LoggingType::SILENT;
ceres::Solver::Summary sum; ceres::Solve(o, &p, &sum);
s_out = sc; ax_out[0] = axv[0]; ax_out[1] = axv[1]; ax_out[2] = axv[2];
return sum.IsSolutionUsable();
};
double s_fit = 1.0, ax_fit[3];
const bool convA = solve_scale_axis(FIT, s_fit, ax_fit);
const double cvA_nom = excit_cost(VAL, 1.0, ax0);
const double cvA_ref = excit_cost(VAL, s_fit, ax_fit);
// Commit the cell scale only for a small, credible move: a well-calibrated header needs < ~0.6 %,
// so a > 1 % scale is a red flag (on multi-lattice / noisy data the excitation fit is biased the
// same way in every cross-validation fold, so the relative-improvement gate cannot catch it).
result.cell_refined = convA && cvA_ref < 0.98 * cvA_nom && std::fabs(s_fit - 1.0) < 0.01;
double s = 1.0, axv[3] = {ax0[0], ax0[1], ax0[2]};
if (result.cell_refined) solve_scale_axis(ALL, s, axv); // commit: re-fit on all data
const double axlen = std::sqrt(axv[0]*axv[0] + axv[1]*axv[1] + axv[2]*axv[2]);
const double axdev = std::acos(std::clamp((axv[0]*ax0[0]+axv[1]*ax0[1]+axv[2]*ax0[2])
/ std::max(1e-9, axlen), -1.0, 1.0)) * 180.0 / PI;
logger.Info("Post-refine GEOM step A (cell/axis): s = {:.5f}, rot-axis {:.3f} deg, held-out excit "
"{:.3e} -> {:.3e} => {}", s, axdev, cvA_nom, cvA_ref,
result.cell_refined ? "COMMIT" : "reject (kept nominal cell)");
// ===== Goniometer rotation SCALE k, its own one-parameter fit on the same rocking events =====
// The angles stored in the file are the COMMANDED ones, so a stage that turned k times as far
// is invisible in the header. Nothing else here can represent it: the cell scale, the axis
// direction, the distance and the beam are all orthogonal to a rotation MAGNITUDE error. Fitted
// after step A so the cell scale and the axis direction are fixed at their committed values and
// k is the only free quantity.
const double u[3] = {axv[0] / axlen, axv[1] / axlen, axv[2] / axlen};
double phi_c = 0.0, phi_lo = scale_events[0].phi_obs, phi_hi = scale_events[0].phi_obs;
for (const auto &ev : scale_events) {
phi_c += ev.phi_obs;
phi_lo = std::min(phi_lo, ev.phi_obs);
phi_hi = std::max(phi_hi, ev.phi_obs);
}
phi_c /= static_cast<double>(scale_events.size());
const double sweep_deg = (phi_hi - phi_lo) * 180.0 / PI;
// The reference reciprocal vector turned to the sweep centre, at the committed cell scale. The
// angle then enters the fit measured FROM that centre. A constant crystal missetting about the
// spindle is k with a slope in phi, so measuring the angle from the goniometer's zero instead
// lets a missetting leak into k with gain <phi>/<phi^2> - which depends only on where the sweep
// happens to sit. On a short sweep starting near zero that gain is enormous: a 0.14 deg
// missetting on a 10 deg wedge fakes 1.4 % of k. Referred to the sweep centre the leak is
// identically zero at any width, and no parameter has to be added to get it.
// The residual is closed-form in k, so this is a one-parameter minimisation rather than a
// solver problem. A rotation preserves length, so |p_lab| = |e_mid| whatever k is, and only
// the z component moves; Rodrigues gives it exactly:
//
// r(k) = C + A cos(a k) - B sin(a k) = C + R cos(a k + psi)
// C = lambda |e|^2 / 2 + u_z (u.e), A = e_z - u_z (u.e), B = (u x e)_z, a = phi_obs - phi_c
//
// which is the same function the residual functor computes, to the last bit. Handing 8 million
// one-parameter residual blocks to Ceres instead cost tens of millions of allocations and a
// dense factorisation per iteration, for a fit that a scan over a bounded interval settles.
// Coefficients are computed in double and stored narrowed: their rounding perturbs the
// minimiser by ~1e-10, and k is carried downstream as a float.
struct ScaleTerm { float a, C, R, psi; };
std::vector<ScaleTerm> terms(scale_events.size());
std::vector<int> fifth_of(scale_events.size());
const double aa_c[3] = {-phi_c * u[0], -phi_c * u[1], -phi_c * u[2]};
ParallelChunks(static_cast<int>(scale_events.size()), nthreads, [&](int lo, int hi) {
for (int e = lo; e < hi; ++e) {
const double p[3] = {scale_events[e].e_ref[0] / s, scale_events[e].e_ref[1] / s,
scale_events[e].e_ref[2] / s};
double em[3];
ceres::AngleAxisRotatePoint(aa_c, p, em);
const double ue = u[0] * em[0] + u[1] * em[1] + u[2] * em[2];
const double e2 = em[0] * em[0] + em[1] * em[1] + em[2] * em[2];
const double C = 0.5 * lambda_l * e2 + u[2] * ue;
const double A = em[2] - u[2] * ue;
const double B = u[0] * em[1] - u[1] * em[0];
terms[e] = ScaleTerm{static_cast<float>(scale_events[e].phi_obs - phi_c),
static_cast<float>(C), static_cast<float>(std::hypot(A, B)),
static_cast<float>(std::atan2(B, A))};
fifth_of[e] = std::clamp(static_cast<int>(
5.0 * (scale_events[e].phi_obs - phi_lo) / std::max(1e-9, phi_hi - phi_lo)), 0, 4);
}
});
// Robust-loss scale from the scatter the events actually have: it varies by more than a decade
// between datasets, so any fixed constant is either inert or throws away real data. Taken once,
// over every event, so the all-data fit and every jackknife fold share it.
const auto residual_at = [&](const ScaleTerm &t, double k) {
return static_cast<double>(t.C)
+ static_cast<double>(t.R) * std::cos(static_cast<double>(t.a) * k + t.psi);
};
double rms = 0.0;
for (const auto &t : terms) {
const double r = residual_at(t, 1.0);
rms += r * r;
}
rms = std::sqrt(rms / static_cast<double>(terms.size()));
const double huber_delta = std::max(1e-12, 2.0 * rms);
const double huber_d2 = huber_delta * huber_delta;
// Ceres minimises half the sum of the loss applied to the SQUARED residual, so that is what is
// reproduced here. One pass yields the five per-fifth partial sums, which serve the all-data
// fit and all five leave-a-fifth-out folds together.
// Each chunk folds into its own slot and the slots are summed in chunk order, so the
// sum does not depend on which worker finishes first: the same events always add up in
// the same sequence, and the fit is reproducible run to run.
const int n_terms = static_cast<int>(terms.size());
const int cost_nt = static_cast<int>(std::max<size_t>(1, std::min(nthreads,
static_cast<size_t>(n_terms))));
const int cost_chunk = (n_terms + cost_nt - 1) / cost_nt;
const auto cost_by_fifth = [&](double k) {
std::vector<std::array<double, 5>> per_chunk(cost_nt);
ParallelChunks(n_terms, nthreads, [&](int lo, int hi) {
std::array<double, 5> acc{};
for (int e = lo; e < hi; ++e) {
const double r = residual_at(terms[e], k);
const double s2 = r * r;
acc[fifth_of[e]] += (s2 <= huber_d2) ? s2
: (2.0 * huber_delta * std::sqrt(s2) - huber_d2);
}
per_chunk[lo / cost_chunk] = acc;
});
std::array<double, 5> total{};
for (const auto &acc : per_chunk)
for (int j = 0; j < 5; ++j) total[j] += acc[j];
return total;
};
// Scan the interval Ceres was bounded to, then close in. No event's phase can move by more than
// a fraction of a period over an interval this narrow, so the objective has no structure the
// grid could step over; the refinement is only there to place the minimum precisely.
constexpr int SCALE_GRID = 101;
constexpr double SCALE_K_LO = 0.95, SCALE_K_HI = 1.05;
std::vector<std::array<double, 5>> grid(SCALE_GRID);
for (int g = 0; g < SCALE_GRID; ++g)
grid[g] = cost_by_fifth(SCALE_K_LO + (SCALE_K_HI - SCALE_K_LO) * g / (SCALE_GRID - 1));
auto solve_scale = [&](int drop_fifth) {
const auto total = [&](const std::array<double, 5> &f) {
double t = 0.0;
for (int j = 0; j < 5; ++j)
if (j != drop_fifth) t += f[j];
return t;
};
int best = 0;
for (int g = 1; g < SCALE_GRID; ++g)
if (total(grid[g]) < total(grid[best])) best = g;
const double step = (SCALE_K_HI - SCALE_K_LO) / (SCALE_GRID - 1);
double a = std::max(SCALE_K_LO, SCALE_K_LO + step * (best - 1));
double b = std::min(SCALE_K_HI, SCALE_K_LO + step * (best + 1));
// Golden section: the objective is smooth but its curvature jumps wherever an event
// crosses the Huber knee, which a derivative method would have to cope with.
constexpr double INV_PHI = 0.6180339887498949;
double c = b - INV_PHI * (b - a), d = a + INV_PHI * (b - a);
double fc = total(cost_by_fifth(c)), fd = total(cost_by_fifth(d));
// The fit is narrowed to a float before it is applied (Rugnux.h prepass_rotation_scale_),
// so bracketing it below that type's epsilon, 6e-8, only buys about ten more full
// passes over the events for a digit that cannot survive being stored.
while (b - a > 1e-7) {
if (fc < fd) { b = d; d = c; fd = fc; c = b - INV_PHI * (b - a); fc = total(cost_by_fifth(c)); }
else { a = c; c = d; fc = fd; d = a + INV_PHI * (b - a); fd = total(cost_by_fifth(d)); }
}
return 0.5 * (a + b);
};
const double k_fit = solve_scale(-1);
result.rotation_scale = k_fit;
// ----- Whether to COMMIT it. A stage fault is rare - 36 of 37 rotation datasets sit at 1.0000
// on a direct scan - and a 1 % angle correction applied to a healthy dataset would damage it
// silently, so every test below has to pass.
// Preconditions: below these the fit is reported but never acted on. Under ~30 deg of sweep k
// entangles with the axis direction and 10-20 deg truncations of a perfect dataset wander by
// +-0.6 %; a screening wedge must not trigger a correction.
constexpr int MIN_SCALE_EVENTS = 5000;
constexpr double MIN_SCALE_SWEEP_DEG = 30.0;
// T1 significance: 0.5 % is 18 sigma on the between-dataset scatter of healthy stages
// (robust sd 2.8e-4) and still 3.5x below the one measured fault.
constexpr double ROTATION_SCALE_TOL = 0.005;
// T2 relevance: the misorientation the error produces at each end of the sweep. A large k over
// a short sweep moves nothing and is not worth correcting.
constexpr double MIN_SCALE_END_ERROR_DEG = 0.5;
// T3 uniformity: a stage error is a ramp present in EVERY part of the sweep, so dropping any
// fifth of it must leave the same k. A second lattice that dominates ONE END of the sweep -
// exactly what happens where the primary stops indexing - fakes a k indistinguishable from a
// real fault on T1 and T2, and is the reason this test is not optional. It replaces the
// hkl-hash split used elsewhere here, which cannot see it: both halves of that split sit at
// the same angles, so anything structured in phi survives in both folds.
constexpr double MIN_SCALE_JACKKNIFE_FRAC = 0.5;
const double end_error_deg = std::fabs(k_fit - 1.0) * sweep_deg / 2.0;
const bool enough_data = static_cast<int>(scale_events.size()) >= MIN_SCALE_EVENTS
&& sweep_deg >= MIN_SCALE_SWEEP_DEG;
const bool big_enough = enough_data && std::fabs(k_fit - 1.0) >= ROTATION_SCALE_TOL
&& end_error_deg >= MIN_SCALE_END_ERROR_DEG;
double jackknife = 1.0;
if (big_enough)
for (int f = 0; f < 5; ++f)
jackknife = std::min(jackknife, (solve_scale(f) - 1.0) / (k_fit - 1.0));
result.rotation_scale_suspect = big_enough && jackknife >= MIN_SCALE_JACKKNIFE_FRAC;
logger.Info("Post-refine rotation SCALE: k = {:.5f} over {:.0f} deg of sweep centred on {:.1f} "
"deg ({} events): end error {:.2f} deg, leave-a-fifth-out {:.2f} => {}",
k_fit, sweep_deg, phi_c * 180.0 / PI, scale_events.size(), end_error_deg, jackknife,
result.rotation_scale_suspect ? "COMMIT"
: !enough_data ? "report only (too little sweep or too few events)"
: "reject (kept the stored angles)");
if (result.rotation_scale_suspect)
logger.Warning("Goniometer rotation scale looks off by {:+.2f} % (fitted {:.5f}): the stage "
"appears to have turned {} than the angles stored in the file, which are the "
"COMMANDED values. This is a hardware calibration fault, not a data problem - "
"left uncorrected it inflates mosaicity, biases the cell and loses "
"high-resolution reflections",
100.0 * (k_fit - 1.0), k_fit, k_fit > 1.0 ? "further" : "less far");
// Cell (scale s, shape fixed) as the XtalResidual parameter blocks p0/p1/p2, held CONSTANT in step B.
double p0[3] = {0, 0, 0}, p1[3] = {0, 0, 0}, p2[3] = {0, 0, 0};
double beta = r0.beta;
switch (sys) {
case gemmi::CrystalSystem::Tetragonal:
LatticeToRodriguesAndLengths_GS(reference_latt, p0, p1);
p1[0] = (p1[0] + p1[1]) / 2.0; break;
case gemmi::CrystalSystem::Cubic:
LatticeToRodriguesAndLengths_GS(reference_latt, p0, p1);
p1[0] = (p1[0] + p1[1] + p1[2]) / 3.0; break;
case gemmi::CrystalSystem::Hexagonal:
LatticeToRodriguesAndLengths_Hex(reference_latt, p0, p1); break;
case gemmi::CrystalSystem::Monoclinic:
LatticeToRodriguesLengthsBeta_Mono(reference_latt, p0, p1, beta);
p2[0] = beta; break;
case gemmi::CrystalSystem::Orthorhombic:
LatticeToRodriguesAndLengths_GS(reference_latt, p0, p1); break;
default:
LatticeToRodriguesAndLengths_GS(reference_latt, p0, p1);
p2[0] = r0.alpha * PI / 180.0; p2[1] = r0.beta * PI / 180.0; p2[2] = r0.gamma * PI / 180.0; break;
}
for (int j = 0; j < 3; ++j) p1[j] *= s; // apply the committed cell scale
double rot_vec[3] = {axv[0], axv[1], axv[2]}; // committed (or nominal) axis
// ===== Step B: detector distance + beam from the observed positions, cell fixed =====
std::vector<const Partial *> obs;
for (const auto &pp : pts)
if (std::isfinite(pp.obs_x) && std::isfinite(pp.obs_y)) obs.push_back(&pp);
constexpr size_t MAX_OBS = 20000;
if (obs.size() > MAX_OBS) {
std::nth_element(obs.begin(), obs.begin() + MAX_OBS, obs.end(),
[](const Partial *a, const Partial *b) {
return a->I / std::max(1e-9, a->sigma) > b->I / std::max(1e-9, b->sigma); });
obs.resize(MAX_OBS);
}
result.obs_used = static_cast<int>(obs.size());
const double beam_x0 = nominal_geom.GetBeamX_pxl(), beam_y0 = nominal_geom.GetBeamY_pxl();
const double dist0 = nominal_geom.GetDetectorDistance_mm();
auto pos_cost = [&](Subset s, const double beam[2], const double dist[1]) {
double c = 0.0; int n = 0;
for (const Partial *pp : obs) {
if (!in(pp->h, pp->k, pp->l, s)) continue;
XtalResidual r(pp->obs_x, pp->obs_y, lambda_l, pixel_mm, rot3, pp->angle_rad,
pp->h, pp->k, pp->l, sys);
double resid[3] = {0, 0, 0};
r(beam, dist, det_rot, rot_vec, p0, p1, p2, resid);
c += resid[0]*resid[0] + resid[1]*resid[1] + resid[2]*resid[2]; ++n;
}
return n ? c / n : 0.0;
};
auto solve_detector = [&](Subset s, double beam_out[2], double &dist_out) {
double beam[2] = {beam_x0, beam_y0}, dist[1] = {dist0};
ceres::Problem p;
for (const Partial *pp : obs) {
if (!in(pp->h, pp->k, pp->l, s)) continue;
p.AddResidualBlock(new ceres::AutoDiffCostFunction<XtalResidual, 3, 2, 1, 2, 3, 3, 3, 3>(
new XtalResidual(pp->obs_x, pp->obs_y, lambda_l, pixel_mm, rot3, pp->angle_rad,
pp->h, pp->k, pp->l, sys)),
new ceres::CauchyLoss(0.02), beam, dist,
const_cast<double *>(det_rot), rot_vec, p0, p1, p2);
p.SetParameterBlockConstant(const_cast<double *>(det_rot));
p.SetParameterBlockConstant(rot_vec);
p.SetParameterBlockConstant(p0); p.SetParameterBlockConstant(p1); p.SetParameterBlockConstant(p2);
}
if (p.NumResidualBlocks() == 0) { beam_out[0] = beam_x0; beam_out[1] = beam_y0; dist_out = dist0; return false; }
p.SetParameterLowerBound(dist, 0, dist0 * 0.95); p.SetParameterUpperBound(dist, 0, dist0 * 1.05);
for (int j = 0; j < 2; ++j) { p.SetParameterLowerBound(beam, j, beam[j] - 15.0);
p.SetParameterUpperBound(beam, j, beam[j] + 15.0); }
ceres::Solver::Options o; o.linear_solver_type = ceres::DENSE_QR; o.max_num_iterations = 60;
o.num_threads = std::max(1, settings.num_threads); o.logging_type = ceres::LoggingType::SILENT;
ceres::Solver::Summary sum; ceres::Solve(o, &p, &sum);
beam_out[0] = beam[0]; beam_out[1] = beam[1]; dist_out = dist[0];
return sum.IsSolutionUsable();
};
double beam[2] = {beam_x0, beam_y0}, dist = dist0;
if (obs.size() >= static_cast<size_t>(settings.min_events)) {
double beam_fit[2], dist_fit;
const bool convB = solve_detector(FIT, beam_fit, dist_fit);
const double b_nom[2] = {beam_x0, beam_y0}, d_nom[1] = {dist0};
const double b_ref[2] = {beam_fit[0], beam_fit[1]}, d_ref[1] = {dist_fit};
const double cvB_nom = pos_cost(VAL, b_nom, d_nom);
const double cvB_ref = pos_cost(VAL, b_ref, d_ref);
// Commit the detector geometry only for a small, credible move: distance < 1 % (a calibrated
// header needs < ~0.6 %). A larger move is the red flag for an unreliable fit - typically a
// second lattice whose spots bias every cross-validation fold identically, so the relative
// "it improved" gate is blind to it and pulls a spurious distance<->cell pair (the radial
// degeneracy) far off. The absolute size of the move discriminates a genuine header correction
// from that failure far better than the absolute residual, which real marginal (noisy / iced)
// data shares with the multi-lattice case.
const bool in_bounds = std::fabs(dist_fit - dist0) < 0.01 * dist0
&& std::hypot(beam_fit[0] - beam_x0, beam_fit[1] - beam_y0) < 15.0;
result.detector_refined = convB && cvB_ref < 0.98 * cvB_nom && in_bounds;
if (result.detector_refined) { double bo[2]; solve_detector(ALL, bo, dist); beam[0] = bo[0]; beam[1] = bo[1]; }
logger.Info("Post-refine GEOM step B (distance/beam): dist {:.3f} -> {:.3f} mm, beam "
"({:.2f},{:.2f}) -> ({:.2f},{:.2f}), held-out pos {:.3e} -> {:.3e} => {}",
dist0, result.detector_refined ? dist : dist0, beam_x0, beam_y0,
result.detector_refined ? beam[0] : beam_x0, result.detector_refined ? beam[1] : beam_y0,
cvB_nom, cvB_ref, result.detector_refined ? "COMMIT" : "reject (kept nominal detector)");
} else {
logger.Info("Post-refine GEOM step B: only {} positional observations - skipped", obs.size());
}
// Assemble the committed geometry.
UnitCell cellA = r0;
if (result.cell_refined) { cellA.a = static_cast<float>(r0.a * s); cellA.b = static_cast<float>(r0.b * s);
cellA.c = static_cast<float>(r0.c * s); }
result.cell = cellA;
result.distance_after_mm = dist;
result.beam_x_before_px = beam_x0; result.beam_x_after_px = beam[0];
result.beam_y_before_px = beam_y0; result.beam_y_after_px = beam[1];
result.events_used = static_cast<int>(events.size());
result.ok = result.cell_refined || result.detector_refined;
if (!result.ok)
logger.Info("Post-refine GEOM: neither step passed cross-validation - geometry left at nominal");
return result;
}
return result; // refine_geometry is the only supported mode; nothing refined otherwise
} catch (...) {
result.ok = false;
return result;
}
}