From 1ca5a69054eaaa2906ba92dc8ab357ba9bc876ab Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Thu, 2 Jul 2026 13:55:04 +0200 Subject: [PATCH] ScaleOnTheFly: solve the per-image scale by IRLS instead of Ceres In the default rotation (rot3d) path only G is refined - B is fixed, mosaicity is pinned and the wedge is not refined - so the predicted intensity G*coeff is linear in G and the robust (Cauchy) per-image scale is a 1-D M-estimate. Solve it directly by iteratively reweighted least squares (a few closed-form weighted ratios) instead of building a Ceres problem per image. Ceres is kept for the cases that are genuinely nonlinear: refining the B-factor or the rotation wedge. Same Cauchy objective as the Ceres path, but ~4x faster at scaling and ~30% faster overall on the /data/rotation_test battery, with space group, cell, ISa, completeness and CC1/2 matching across all 18 crystals (the two that look different, EP_cs_01-17 and EcwtAL500, are run-to-run unstable for both solvers). lyso_ref scaling 25.2->4.3s, cytC_2 15.2->2.6s, battery total 468->316s. Also drop the per-image G/B regularizers (gated by GetScalingRegularize, which nothing enables) and the now-unused RegularizationResidual. Co-Authored-By: Claude Fable 5 --- image_analysis/scale_merge/ScaleOnTheFly.cpp | 254 +++++++++++-------- 1 file changed, 154 insertions(+), 100 deletions(-) diff --git a/image_analysis/scale_merge/ScaleOnTheFly.cpp b/image_analysis/scale_merge/ScaleOnTheFly.cpp index fda12d97..36a1e8d8 100644 --- a/image_analysis/scale_merge/ScaleOnTheFly.cpp +++ b/image_analysis/scale_merge/ScaleOnTheFly.cpp @@ -3,7 +3,10 @@ #include "ScaleOnTheFly.h" +#include +#include #include +#include #include #include @@ -32,23 +35,51 @@ namespace { } - struct RegularizationResidual { - RegularizationResidual(const double weight, const double expected) - : weight(weight), - expected(expected) { - } - - template - bool operator()(const T *parameter, T *residual) const { - residual[0] = T(weight) * (parameter[0] - T(expected)); - return true; - } - - private: - const double weight; - const double expected; + // One reflection reduced to the 1-D scale fit: predicted intensity is G * coeff (coeff is constant + // while B, mosaicity and wedge are fixed), measured is Iobs, weighted by 1/sigma. + struct ScaleObs { + double coeff; + double Iobs; + double weight; }; + // Robust per-image scale: minimise sum_i Cauchy_k( weight_i (G*coeff_i - Iobs_i) ) over G >= 0. The + // model is linear in G, so this M-estimate is a few reweighted-least-squares steps (each a closed-form + // weighted ratio) - the same objective the Ceres path solves, without a per-image problem/autodiff/ + // trust-region. Seeded from the plain weighted-LS solution; Cauchy weight is 1/(1 + (res/k)^2). + double SolveScaleIRLS(const std::vector &obs, double robust_k) { + auto weighted_scale = [&obs](auto robust_weight) { + double num = 0.0, den = 0.0; + for (const auto &o: obs) { + const double rw = robust_weight(o); + const double w2 = o.weight * o.weight; + num += rw * w2 * o.coeff * o.Iobs; + den += rw * w2 * o.coeff * o.coeff; + } + return den > 0.0 ? num / den : NAN; + }; + + double G = weighted_scale([](const ScaleObs &) { return 1.0; }); + if (!std::isfinite(G)) + return 1.0; + G = std::max(0.0, G); + + const double k2 = robust_k * robust_k; + for (int iter = 0; iter < 30; ++iter) { + const double G_prev = G; + const double G_next = weighted_scale([&](const ScaleObs &o) { + const double res = o.weight * (G * o.coeff - o.Iobs); + return 1.0 / (1.0 + res * res / k2); + }); + if (!std::isfinite(G_next)) + break; + G = std::max(0.0, G_next); + if (std::abs(G - G_prev) <= 1e-7 * std::max(G, 1.0)) + break; + } + return G; + } + class ScalingResidual { protected: const double Iobs; @@ -212,8 +243,6 @@ void ScaleOnTheFly::Scale(IntegrationOutcome &integration_outcome) const { auto start = std::chrono::steady_clock::now(); - ceres::Problem problem; - ScaleOnTheFlyResult result{ .B = 0.0, .G = 1.0, @@ -235,103 +264,128 @@ void ScaleOnTheFly::Scale(IntegrationOutcome &integration_outcome) const { result.wedge = NAN; } - size_t n_reflections = 0; - for (const auto &r: integration_outcome.reflections) { - if (!Accept(r)) - continue; - - const HKLKey key = hkl_key_generator(r); - - if (!reference_data.contains(key)) - continue; - - ++n_reflections; - - const double Itrue = reference_data.at(key); - const double sigma = r.sigma; - - switch (model) { - case PartialityModel::Fixed: { - auto *cost = new ceres::AutoDiffCostFunction( - new IntensityFixedResidual(r, Itrue, sigma, r.partiality)); - problem.AddResidualBlock(cost, new ceres::CauchyLoss(SCALE_ROBUST_K), &result.G, &result.B); - } - break; - case PartialityModel::Unity: { - auto *cost = new ceres::AutoDiffCostFunction( - new IntensityFixedResidual(r, Itrue, sigma, 1.0)); - problem.AddResidualBlock(cost, new ceres::CauchyLoss(SCALE_ROBUST_K), &result.G, &result.B); - } - break; - case PartialityModel::Rotation: { - auto *cost = new ceres::AutoDiffCostFunction( - new ScalingRotationResidual(r, Itrue, sigma)); - problem.AddResidualBlock(cost, new ceres::CauchyLoss(SCALE_ROBUST_K), &result.G, &result.B, &result.mos, - &result.wedge); - } - break; - default: - throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, - "Not supported partiality model"); - } - } - - if (n_reflections < MIN_REFLECTIONS) { + auto clear_scale = [&]() { integration_outcome.image_scale_cc.reset(); integration_outcome.image_scale_cc_n.reset(); integration_outcome.image_scale_g.reset(); integration_outcome.image_scale_b_factor_Ang2.reset(); - return; - } + }; - constexpr double SIGMA_SCALE_FACTOR = 2; // We assume scale factor is +/- 2.0 - const double scale_factor_regularization_weight = std::sqrt(static_cast(n_reflections) / SIGMA_SCALE_FACTOR); + // With B, mosaicity and wedge all fixed the predicted intensity G * coeff is linear in G, so the + // robust per-image scale is a 1-D M-estimate solved directly (IRLS) instead of building a Ceres + // problem per image. Ceres is kept only for the cases that make it genuinely nonlinear: refining the + // B-factor (exp(-B/...)) or the rotation wedge (which moves the partiality). + const bool linear_in_g = !s.GetRefineB() + && (model != PartialityModel::Rotation || !refine_rot_wedge); - if (s.GetScalingRegularize()) { - auto *scale_reg_cost = new ceres::AutoDiffCostFunction( - new RegularizationResidual(scale_factor_regularization_weight, 1.0)); - problem.AddResidualBlock(scale_reg_cost, nullptr, &result.G); - } - - problem.SetParameterLowerBound(&result.G, 0, 0.0); - - if (s.GetRefineB()) { - constexpr double SIGMA_B = 10.0; // in A^2 - const double b_regularization_weight = std::sqrt(static_cast(n_reflections) / SIGMA_B); - - if (s.GetScalingRegularize()) { - auto *b_reg_cost = new ceres::AutoDiffCostFunction( - new RegularizationResidual(b_regularization_weight, 0.0)); - problem.AddResidualBlock(b_reg_cost, nullptr, &result.B); + if (linear_in_g) { + std::vector obs; + obs.reserve(integration_outcome.reflections.size()); + for (const auto &r: integration_outcome.reflections) { + if (!Accept(r)) + continue; + const auto it = reference_data.find(hkl_key_generator(r)); + if (it == reference_data.end()) + continue; + double partiality; + switch (model) { + case PartialityModel::Unity: + partiality = 1.0; + break; + case PartialityModel::Rotation: + partiality = RotationPartiality(r.delta_phi_deg, r.zeta, result.mos, result.wedge); + break; + default: + partiality = r.partiality; + break; + } + const double B_term = std::exp(result.B * -SafeInv(4.0 * r.d * r.d, 0.0)); + const double coeff = partiality * B_term * SafeInv(r.rlp, 1.0) * it->second; + obs.push_back({coeff, static_cast(r.I), SafeInv(r.sigma, 1.0)}); } - problem.SetParameterLowerBound(&result.B, 0, s.GetMinB()); - problem.SetParameterUpperBound(&result.B, 0, s.GetMaxB()); - } else { - problem.SetParameterBlockConstant(&result.B); - } + if (obs.size() < MIN_REFLECTIONS) { + clear_scale(); + return; + } - if (model == PartialityModel::Rotation) { - if (refine_rot_wedge) { + result.G = SolveScaleIRLS(obs, SCALE_ROBUST_K); + } else { + ceres::Problem problem; + + size_t n_reflections = 0; + for (const auto &r: integration_outcome.reflections) { + if (!Accept(r)) + continue; + + const HKLKey key = hkl_key_generator(r); + + if (!reference_data.contains(key)) + continue; + + ++n_reflections; + + const double Itrue = reference_data.at(key); + const double sigma = r.sigma; + + switch (model) { + case PartialityModel::Fixed: { + auto *cost = new ceres::AutoDiffCostFunction( + new IntensityFixedResidual(r, Itrue, sigma, r.partiality)); + problem.AddResidualBlock(cost, new ceres::CauchyLoss(SCALE_ROBUST_K), &result.G, &result.B); + } + break; + case PartialityModel::Unity: { + auto *cost = new ceres::AutoDiffCostFunction( + new IntensityFixedResidual(r, Itrue, sigma, 1.0)); + problem.AddResidualBlock(cost, new ceres::CauchyLoss(SCALE_ROBUST_K), &result.G, &result.B); + } + break; + case PartialityModel::Rotation: { + auto *cost = new ceres::AutoDiffCostFunction( + new ScalingRotationResidual(r, Itrue, sigma)); + problem.AddResidualBlock(cost, new ceres::CauchyLoss(SCALE_ROBUST_K), &result.G, &result.B, + &result.mos, &result.wedge); + } + break; + default: + throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, + "Not supported partiality model"); + } + } + + if (n_reflections < MIN_REFLECTIONS) { + clear_scale(); + return; + } + + problem.SetParameterLowerBound(&result.G, 0, 0.0); + + if (s.GetRefineB()) { + problem.SetParameterLowerBound(&result.B, 0, s.GetMinB()); + problem.SetParameterUpperBound(&result.B, 0, s.GetMaxB()); + } else { + problem.SetParameterBlockConstant(&result.B); + } + + // Only when the wedge is refined are mos/wedge parameter blocks in the problem. Bound the wedge + // and keep mosaicity fixed: the per-image fit is degenerate between G and mosaicity and collapses + // it toward its floor (3x too small), which corrupts the partiality. + if (model == PartialityModel::Rotation && refine_rot_wedge) { problem.SetParameterLowerBound(&result.wedge, 0, s.GetMinWedge()); problem.SetParameterUpperBound(&result.wedge, 0, s.GetMaxWedge()); - } else { - problem.SetParameterBlockConstant(&result.wedge); + problem.SetParameterBlockConstant(&result.mos); } - // Trust the (now correctly estimated) indexing mosaicity; do NOT re-refine it here. The - // per-image scaling residual fit is degenerate between G and the mosaicity, and it collapses - // the mosaicity toward its floor (3x too small) which corrupts the partiality. Keep it fixed. - problem.SetParameterBlockConstant(&result.mos); + + ceres::Solver::Options options; + options.linear_solver_type = ceres::DENSE_QR; + options.minimizer_progress_to_stdout = false; + options.num_threads = 1; + + ceres::Solver::Summary summary; + ceres::Solve(options, &problem, &summary); } - ceres::Solver::Options options; - options.linear_solver_type = ceres::DENSE_QR; - options.minimizer_progress_to_stdout = false; - options.num_threads = 1; - - ceres::Solver::Summary summary; - ceres::Solve(options, &problem, &summary); - for (auto &r: integration_outcome.reflections) { const double B_term = exp(result.B * -SafeInv(4.0 * r.d * r.d, 0.0));