mx/Jungfraujoch
2
1
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 68 Wiki Activity

Fit the goniometer rotation scale in closed form

Browse Source 
The fit has ONE parameter, and it was handed to Ceres as one residual block
per rocking event - 8 million of them on a large crystal. Each block is a
functor, an auto-diff cost function and a loss object on the heap, and the
solver then factorises an 8-million-by-one Jacobian on every iteration. It cost
13.7 s.

The residual is closed-form in k. A rotation preserves length, so |p_lab| is
|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

with a the event's angle from the sweep centre. That is the same function the
functor computes - Ceres uses the exact Rodrigues form here, so there is no
small-angle branch to disagree with - and it reduces the fit to minimising a
smooth function of one variable over the interval the solver was bounded to.
It is scanned on a grid and then closed in by golden section; the objective's
curvature jumps wherever an event crosses the Huber knee, which is why this is
not a Newton iteration.

The coefficients are computed in double and stored narrowed. Their rounding
moves the minimiser by ~1e-10, and k is carried downstream as a float, so the
committed value is the same to far more digits than anything reads.

One pass over the events yields the five per-fifth partial sums, so the
all-data fit and the five leave-a-fifth-out folds share it. That matters
because the jackknife only runs when the fit is big enough to act on, and on a
crystal that trips it the old code paid for six full solves.

The partials gather ahead of it counted first and then filled instead of
growing one vector by push_back tens of millions of times, which copied the
whole thing on every doubling.

Measured: unchanged verdict and k to five decimals on the regression crystals.
Full 24-crystal battery: same space group on all 24, none failed, 15m32s ->
13m35s together with the scale/merge changes.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
jungfrau
2026-08-15 18:56:09 -04:00
co-authored by Claude Opus 5
parent 6368c00173
commit e11c2a2b20
1 changed files with 134 additions and 40 deletions
Download Patch File Download Diff File Expand all files Collapse all files
image_analysis/geom_refinement/PostRefine.cpp
+134 -40
View File
@@ -1,11 +1,13 @@
// 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 <mutex>
#include "../../common/JFJochMath.h" // PI
#include "XtalResidual.h" // XtalResidual (the positional detector<->reciprocal residual, step B)
@@ -105,16 +107,39 @@ PostRefineResult PostRefineRotationGeometry(const std::vector<IntegrationOutcome
const Coord ax = axis.GetAxis();
const Coord Astar = reference_latt.Astar(), Bstar = reference_latt.Bstar(), Cstar = reference_latt.Cstar();
std::vector<Partial> pts;
for (const auto &o : outcomes)
for (const auto &r : 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.push_back(Partial{r.h, r.k, r.l, r.image_number, r.I, r.sigma,
mid_deg * PI / 180.0, ox, oy});
// 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;
@@ -276,39 +301,108 @@ PostRefineResult PostRefineRotationGeometry(const std::vector<IntegrationOutcome
// 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.
std::vector<std::array<double, 3>> e_mid(scale_events.size());
for (size_t e = 0; e < scale_events.size(); ++e) {
const double aa[3] = {-phi_c * u[0], -phi_c * u[1], -phi_c * u[2]};
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};
ceres::AngleAxisRotatePoint(aa, p, e_mid[e].data());
}
// 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.
double rms = 0.0;
for (size_t e = 0; e < scale_events.size(); ++e) {
RotationScaleResidual r(lambda_l, scale_events[e].phi_obs - phi_c, u, e_mid[e].data());
double one = 1.0, resid = 0.0;
r(&one, &resid); rms += resid * resid;
}
rms = std::sqrt(rms / static_cast<double>(scale_events.size()));
// Fit k over the events whose phi lies outside the given fifth of the sweep (-1 = all of it).
auto solve_scale = [&](int drop_fifth) {
double kv = 1.0;
ceres::Problem p;
for (size_t e = 0; e < scale_events.size(); ++e) {
const int fifth = std::clamp(static_cast<int>(
// 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);
if (fifth == drop_fifth) continue;
p.AddResidualBlock(new ceres::AutoDiffCostFunction<RotationScaleResidual, 1, 1>(
new RotationScaleResidual(lambda_l, scale_events[e].phi_obs - phi_c, u, e_mid[e].data())),
new ceres::HuberLoss(std::max(1e-12, 2.0 * rms)), &kv);
}
p.SetParameterLowerBound(&kv, 0, 0.95); p.SetParameterUpperBound(&kv, 0, 1.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);
return sum.IsSolutionUsable() ? kv : 1.0;
});
// 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()), 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;