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

ScaleOnTheFly: solve the per-image scale by IRLS instead of Ceres

Browse Source 
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 <noreply@anthropic.com>
This commit is contained in:
leonarski_f
2026-07-02 13:55:04 +02:00
co-authored by Claude Fable 5
parent 5c9175688a
commit 1ca5a69054
1 changed files with 154 additions and 100 deletions
Download Patch File Download Diff File Expand all files Collapse all files
image_analysis/scale_merge/ScaleOnTheFly.cpp
+154 -100
View File
@@ -3,7 +3,10 @@
#include "ScaleOnTheFly.h"
#include <algorithm>
#include <cmath>
#include <future>
#include <vector>
#include <ceres/ceres.h>
#include <ceres/rotation.h>
@@ -32,23 +35,51 @@ namespace {
}
struct RegularizationResidual {
RegularizationResidual(const double weight, const double expected)
: weight(weight),
expected(expected) {
}
template<typename T>
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<ScaleObs> &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<IntensityFixedResidual, 1, 1, 1>(
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<IntensityFixedResidual, 1, 1, 1>(
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<ScalingRotationResidual, 1, 1, 1, 1, 1>(
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<double>(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<RegularizationResidual, 1, 1>(
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<double>(n_reflections) / SIGMA_B);
if (s.GetScalingRegularize()) {
auto *b_reg_cost = new ceres::AutoDiffCostFunction<RegularizationResidual, 1, 1>(
new RegularizationResidual(b_regularization_weight, 0.0));
problem.AddResidualBlock(b_reg_cost, nullptr, &result.B);
if (linear_in_g) {
std::vector<ScaleObs> 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<double>(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<IntensityFixedResidual, 1, 1, 1>(
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<IntensityFixedResidual, 1, 1, 1>(
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<ScalingRotationResidual, 1, 1, 1, 1, 1>(
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));