Files
Jungfraujoch/image_analysis/geom_refinement/PostRefine.cpp
T
jungfrauandClaude Opus 5 65eb86ab66 Size the post-refine gather by its reflections, not by its frames
Both passes of the gather are chunked over OUTCOMES, one per frame, but their cost is
the tens of thousands of reflections inside each. ThreadsForWork was being asked about
the frames: a rotation dataset has on the order of a thousand of them against a floor
of 32768 items per thread, so it answered 1, and the pass that walks every reflection
in the run ran on a single thread on every dataset in the test set.

RotationScaleMerge::Ingest does the identical thing and asks with the observation
count, which is what makes this an inconsistency rather than a tuning choice. The
reflection total is already to hand a few lines up.

The split is unchanged - only the worker count moves - so the output is identical.
Measured on the heaviest crystal: the phase drops 33 %, from 8.6 s to 5.5 s, and its
mean occupancy goes from 4.1 threads to 12.3.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
2026-08-17 18:05:19 -04:00

656 lines
41 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 <mutex>
#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();
// This refinement fits five numbers - a uniform cell scale, the rotation-axis direction, the
// detector distance and the beam centre - and the two solver steps below already cap themselves at
// 20 000 rocking events and 20 000 spot positions. Gathering, sorting and event-splitting every
// partial of a large dataset to feed those caps is work nothing reads. So thin the input by RAW
// HKL: a whole rocking event is then kept or dropped together, so no phi_obs centroid is biased;
// every part of the sweep and every part of the detector face is thinned by the same factor, which
// is what the distance/beam step and the rotation-scale scan are sensitive to; and the subset
// depends on the indices alone, so it is the same one on every run.
const size_t nthreads = std::max(1, settings.num_threads);
const int n_out = static_cast<int>(outcomes.size());
size_t n_integrated = 0;
for (const auto &o : outcomes) n_integrated += o.reflections.size();
const size_t stride = std::max<size_t>(1, n_integrated / std::max<size_t>(1, settings.max_partials));
// Mixed differently from the cross-validation split further down (is_val), so which hkls are
// sampled and which fold they land in are independent.
const auto in_subset = [stride](int h, int k, int l) {
if (stride == 1) return true;
unsigned u = static_cast<unsigned>(h) * 2246822519u + static_cast<unsigned>(k) * 3266489917u
+ static_cast<unsigned>(l) * 668265263u;
u ^= u >> 15; u *= 2654435761u; u ^= u >> 16;
return u % stride == 0;
};
// 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.
// Both passes below are chunked over OUTCOMES but their cost is in the reflections inside them,
// tens of thousands per frame, so the worker count has to come from the reflection count. Asked
// about the outcomes instead, ThreadsForWork sees a frame count against its 32768-per-thread
// floor and answers 1 - which ran the whole gather on one thread on every dataset in the test
// set. RotationScaleMerge::Ingest already asks the same question the right way round.
const size_t gather_threads = ThreadsForWork(n_integrated, nthreads);
std::vector<size_t> pts_offset(n_out + 1, 0);
ParallelChunks(n_out, gather_threads, [&](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
&& in_subset(r.h, r.k, r.l))
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, gather_threads, [&](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;
if (!in_subset(r.h, r.k, r.l)) 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 (1 raw hkl in {} of {} integrated)", pts.size(),
stride, n_integrated);
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.
{
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()), ThreadsForWork(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.
const auto cost_by_fifth = [&](double k) {
std::array<double, 5> total{};
std::mutex mx;
ParallelChunks(static_cast<int>(terms.size()), ThreadsForWork(terms.size(), 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);
}
std::unique_lock ul(mx);
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));
while (b - a > 1e-9) {
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;
}
}