From e11c2a2b200d006d2ea82f2d97021135f8286810 Mon Sep 17 00:00:00 2001 From: jungfrau Date: Sat, 15 Aug 2026 18:56:09 -0400 Subject: [PATCH] Fit the goniometer rotation scale in closed form 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) --- image_analysis/geom_refinement/PostRefine.cpp | 174 ++++++++++++++---- 1 file changed, 134 insertions(+), 40 deletions(-) diff --git a/image_analysis/geom_refinement/PostRefine.cpp b/image_analysis/geom_refinement/PostRefine.cpp index fae7cf9f..6ae2df1b 100644 --- a/image_analysis/geom_refinement/PostRefine.cpp +++ b/image_analysis/geom_refinement/PostRefine.cpp @@ -1,11 +1,13 @@ // SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute // SPDX-License-Identifier: GPL-3.0-only +#include "../../common/ParallelFor.h" #include "PostRefine.h" #include #include #include +#include #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 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(outcomes.size()); + std::vector 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 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(settings.min_events)) return result; @@ -276,39 +301,108 @@ PostRefineResult PostRefineRotationGeometry(const std::vector> 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(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( + // 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 terms(scale_events.size()); + std::vector 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(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(scale_events[e].phi_obs - phi_c), + static_cast(C), static_cast(std::hypot(A, B)), + static_cast(std::atan2(B, A))}; + fifth_of[e] = std::clamp(static_cast( 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( - 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(t.C) + + static_cast(t.R) * std::cos(static_cast(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(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 total{}; + std::mutex mx; + ParallelChunks(static_cast(terms.size()), nthreads, [&](int lo, int hi) { + std::array 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> 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 &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;