mx/Jungfraujoch
2
1
Fork 0
Code Issues Pull Requests Actions 3 Packages Projects Releases 70 Wiki Activity
Files
0fed94d75beca70ee9f5fea68bf519953cf816c6
Jungfraujoch/image_analysis/geom_refinement/PostRefine.cpp
T
jungfrauandClaude Opus 5 0fed94d75b Gather post-refinement's partials without touching four gigabytes twice 
Post-refinement fits eight numbers, and it selects the twenty thousand best-recorded events to fit
them from. Before it can select, it copies every integrated partial into an array of its own and
sorts it. On a large cell that is 63 million of them, and the phase took 8.5 s of a 56 s run.

Almost none of that was the sort. `std::vector<Partial> pts(n)` value-initialises: one thread writes
3.5 GB of zeroes, page by page, before the parallel fill overwrites every byte of it - and being the
first touch, it also decides where the pages live, so the whole array lands on one NUMA node and
every later pass over it runs at one node's bandwidth. The same again for the sorted copy. Allocate
the storage without initialising it and let the parallel fill be the first touch.

The record itself carried more than the sort reads. `angle_rad` is a function of the image number
that the goniometer can give back on demand, and the two observed positions are wanted only by the
distance step, and only for the twenty thousand it keeps. Storing what is read - and as the floats
the fields already were, since widening a float to a double is exact - takes the record from 56
bytes to 32, which is a third off the fill and half off the sort's element moves.

Then three passes that walked the whole array to no purpose. The h range is now taken in the count
pass, which reads the same reflections anyway; the bucket histogram in the fill pass, which already
has h in hand. The event split walked serially and grew its output by doubling - about a gigabyte of
pure copying - although h is the leading sort key, so a rocking event never crosses an h bucket:
count per bucket, prefix, fill in parallel, and the events come out in the order the serial walk
produced them. And the copy of the whole event list, made only so that nth_element could destroy the
original, is now an index array.

Every one of these is the same arithmetic in the same order. Measured on a large-cell rotation set,
with the two commits that follow: 52.6 s -> 46.1 s, and the merged .hkl, .mtz and .cif are
byte-identical. `part_less` is deliberately left as it was, not a total order: what makes it
reproducible is that each bucket reaches the sort in gather order, and the new chunking is still a
contiguous span of that order.

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

704 lines
44 KiB
C++
Raw Blame History

// 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 <memory>
#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. Kept as narrow as the sort and the event
// split allow: on a large cell this array is gigabytes, and the scatter and every level of the
// per-bucket sort move all of it. The goniometer angle is not stored - it is a function of the
// image number alone, and is rebuilt from it where it is needed.
struct Partial {
int h, k, l;
float img;
float I, sigma;
float 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();
// The goniometer angle is a function of the image number alone, so it is recomputed where it
// is used rather than carried through the array below: eight bytes per partial cost more in
// the fill, the scatter and every level of the sort than the multiply-add that rebuilds them.
const auto angle_rad = [&](float img) { return (axis.GetAngle_deg(img) + wedge_half) * PI / 180.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. The h range comes out
// of the same sweep: the sort below buckets by h and needs to know how many buckets that is,
// and this pass already reads every reflection.
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);
std::vector<int> h_lo_of(n_out), h_hi_of(n_out);
ParallelChunks(n_out, nthreads, [&](int lo, int hi) {
for (int o = lo; o < hi; o++) {
size_t keep = 0;
int lmin = std::numeric_limits<int>::max(), lmax = std::numeric_limits<int>::min();
for (const auto &r : outcomes[o].reflections)
if (std::isfinite(r.I) && std::isfinite(r.sigma) && r.sigma > 0.0f) {
keep++;
lmin = std::min(lmin, r.h);
lmax = std::max(lmax, r.h);
}
pts_offset[o + 1] = keep;
h_lo_of[o] = lmin;
h_hi_of[o] = lmax;
}
});
int h_lo = std::numeric_limits<int>::max(), h_hi = std::numeric_limits<int>::min();
for (int o = 0; o < n_out; o++) {
pts_offset[o + 1] += pts_offset[o];
h_lo = std::min(h_lo, h_lo_of[o]);
h_hi = std::max(h_hi, h_hi_of[o]);
}
const size_t n_pts = pts_offset[n_out];
const int H = (h_lo <= h_hi) ? (h_hi - h_lo + 1) : 1;
// A vector of n partials VALUE-initialises them: on a large cell that is gigabytes of zeroing
// on one thread, and it is that one thread which first touches every page - which on a
// multi-socket machine leaves the whole array on its node, so every pass that follows runs at
// one node's memory bandwidth. new[] leaves the partials untouched, so the parallel fill is
// the first touch and each page lands on the node of the thread that filled it.
std::unique_ptr<Partial[]> pts(new Partial[n_pts]);
// The bucket histogram the sort needs is taken here rather than in a pass of its own, since
// the fill already has h in hand. Its chunks are the outcome chunks ParallelChunks makes, so
// the scatter below has to be split the same way.
const int nt = static_cast<int>(std::clamp<size_t>(nthreads, 1, std::max(1, n_out)));
const int chunk = (n_out + nt - 1) / nt;
std::vector<std::vector<int32_t>> hist(nt, std::vector<int32_t>(H, 0));
ParallelChunks(n_out, nthreads, [&](int lo, int hi) {
std::vector<int32_t> &h_count = hist[lo / chunk];
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 float ox = std::isfinite(r.observed_x) ? r.observed_x : NAN;
const float 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, ox, oy};
h_count[r.h - h_lo]++;
}
}
});
logger.Info("Post-refine: {} partials gathered", n_pts);
if (n_pts < static_cast<size_t>(settings.min_events)) return result;
// Where each bucket starts, and the buckets largest first: the sort lays the array out this
// way and the event split below walks the same buckets.
std::vector<int32_t> bstart(H + 1, 0);
std::vector<int> order(H);
// 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 every partial 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;
};
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::unique_ptr<Partial[]> sorted(new Partial[n_pts]);
ParallelChunks(n_out, nthreads, [&](int lo, int hi) {
std::vector<int32_t> fill = hist[lo / chunk];
for (size_t i = pts_offset[lo]; i < pts_offset[hi]; ++i)
sorted[fill[pts[i].h - h_lo]++] = pts[i];
});
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.get() + bstart[b], sorted.get() + 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;
const auto run_end = [&](size_t i, size_t end) {
size_t j = i + 1;
while (j < end && 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;
return j;
};
// The event the partials [i, j) make, or false where their intensities cannot place a
// centroid. Counting the events and writing them both walk the buckets, and both build the
// event this way.
const auto make_event = [&](size_t i, size_t j, Event &out) {
double sumI = 0, sumIphi = 0, sumSig = 0;
for (size_t m = i; m < j; ++m) {
const double Ipos = std::max(0.0, static_cast<double>(pts[m].I));
sumI += Ipos; sumIphi += Ipos * angle_rad(pts[m].img); sumSig += pts[m].sigma;
}
if (!(sumI > 0.0 && sumSig > 0.0)) return false;
const Coord e = Astar * static_cast<float>(pts[i].h) + Bstar * static_cast<float>(pts[i].k)
+ Cstar * static_cast<float>(pts[i].l);
out = Event{sumIphi / sumI, std::sqrt(sumI / sumSig), {e.x, e.y, e.z},
pts[i].h, pts[i].k, pts[i].l};
return true;
};
// An event never crosses an h boundary - h is the leading sort key - so the buckets can be
// walked independently, and laying their events out in bucket order gives exactly the order
// the serial walk produced. Counting first also sizes the array in one go, in place of a
// push_back that grew a gigabyte by doubling.
std::vector<int32_t> ev_count(H, 0);
std::vector<size_t> ev_frames(H, 0);
ParallelFor(H, nthreads, [&](int oi) {
const int b = order[oi];
const size_t end = bstart[b + 1];
int c = 0;
Event ev;
for (size_t i = bstart[b]; i < end; ) {
const size_t j = run_end(i, end);
if (j - i >= 2 && make_event(i, j, ev)) ++c;
i = j;
}
ev_count[b] = c;
});
std::vector<int32_t> ev_start(H + 1, 0);
for (int b = 0; b < H; ++b) ev_start[b + 1] = ev_start[b] + ev_count[b];
const size_t n_events = ev_start[H];
std::unique_ptr<Event[]> events(new Event[n_events]);
ParallelFor(H, nthreads, [&](int oi) {
const int b = order[oi];
const size_t end = bstart[b + 1];
int at = ev_start[b];
size_t frames = 0;
for (size_t i = bstart[b]; i < end; ) {
const size_t j = run_end(i, end);
if (j - i >= 2 && make_event(i, j, events[at])) { frames += j - i; ++at; }
i = j;
}
ev_frames[b] = frames;
});
size_t event_frames = 0;
for (int b = 0; b < H; ++b) event_frames += ev_frames[b];
// 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)", n_events,
n_events == 0 ? 0.0 : static_cast<double>(event_frames) / n_events);
if (static_cast<int>(n_events) < 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. Both sets come out of the one array by selecting on indices instead: with
// the weights in the same places nth_element takes the same decisions it would take on the
// events themselves, so the selection is the same one in the same order and the whole list no
// longer has to be duplicated to survive it.
constexpr size_t MAX_EVENTS = 20000;
std::vector<int32_t> selected(n_events);
std::iota(selected.begin(), selected.end(), 0);
if (selected.size() > MAX_EVENTS) {
std::nth_element(selected.begin(), selected.begin() + MAX_EVENTS, selected.end(),
[&](int32_t a, int32_t b) { return events[a].weight > events[b].weight; });
selected.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 int32_t i : selected) {
const Event &ev = events[i];
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 int32_t i : selected) {
const Event &ev = events[i];
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 = events[0].phi_obs, phi_hi = events[0].phi_obs;
for (size_t e = 0; e < n_events; ++e) {
phi_c += events[e].phi_obs;
phi_lo = std::min(phi_lo, events[e].phi_obs);
phi_hi = std::max(phi_hi, events[e].phi_obs);
}
phi_c /= static_cast<double>(n_events);
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(n_events);
std::vector<int> fifth_of(n_events);
const double aa_c[3] = {-phi_c * u[0], -phi_c * u[1], -phi_c * u[2]};
ParallelChunks(static_cast<int>(n_events), nthreads, [&](int lo, int hi) {
for (int e = lo; e < hi; ++e) {
const double p[3] = {events[e].e_ref[0] / s, events[e].e_ref[1] / s,
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>(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 * (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>(n_events) >= 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, n_events, 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 (size_t i = 0; i < n_pts; ++i)
if (std::isfinite(pts[i].obs_x) && std::isfinite(pts[i].obs_y)) obs.push_back(&pts[i]);
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, static_cast<double>(a->sigma))
> b->I / std::max(1e-9, static_cast<double>(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, angle_rad(pp->img),
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, angle_rad(pp->img),
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>(selected.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;
}
}
Reference in New Issue View Git Blame Copy Permalink