leonarski_fandClaude Fable 5
1ca5a69054
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>
466 lines
17 KiB
C++
466 lines
17 KiB
C++
// SPDX-FileCopyrightText: 2025 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
|
|
// SPDX-License-Identifier: GPL-3.0-only
|
|
|
|
#include "ScaleOnTheFly.h"
|
|
|
|
#include <algorithm>
|
|
#include <cmath>
|
|
#include <future>
|
|
#include <vector>
|
|
#include <ceres/ceres.h>
|
|
#include <ceres/rotation.h>
|
|
|
|
namespace {
|
|
// Robust loss scale (in sigma units) for the per-image scale fit: a few outlier reflections
|
|
// (zingers, overlaps, a mis-predicted spot) must not drag a frame's G/B into a bad optimum -
|
|
// that is the stochastic per-frame mis-scaling that elevates R-meas and collapses CC1/2 at low
|
|
// symmetry. Cauchy down-weights residuals beyond ~this many sigma without a hard cut.
|
|
constexpr double SCALE_ROBUST_K = 3.0;
|
|
|
|
double SafeInv(double x, double fallback) {
|
|
if (!std::isfinite(x) || x == 0.0)
|
|
return fallback;
|
|
return 1.0 / x;
|
|
}
|
|
|
|
float RotationPartiality( double delta_phi_deg,
|
|
double zeta,
|
|
double mosaicity_deg,
|
|
double wedge_deg) {
|
|
const double half_wedge = wedge_deg / 2.0;
|
|
const double c1 = zeta / std::sqrt(2.0);
|
|
const double arg_plus = (delta_phi_deg + half_wedge) * c1 / mosaicity_deg;
|
|
const double arg_minus = (delta_phi_deg - half_wedge) * c1 / mosaicity_deg;
|
|
return static_cast<float>((std::erf(arg_plus) - std::erf(arg_minus)) / 2.0);
|
|
}
|
|
|
|
|
|
// 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;
|
|
const double Itrue;
|
|
const double weight;
|
|
const double lp;
|
|
const double b_resolution_coeff;
|
|
|
|
ScalingResidual(const Reflection &r, double Itrue, double sigma)
|
|
: Iobs(r.I),
|
|
Itrue(Itrue),
|
|
weight(SafeInv(sigma, 1.0)),
|
|
lp(SafeInv(r.rlp, 1.0)),
|
|
b_resolution_coeff(-SafeInv(4.0 * r.d * r.d, 0.0)) {
|
|
}
|
|
};
|
|
|
|
struct ScalingRotationResidual : public ScalingResidual {
|
|
ScalingRotationResidual(const Reflection &r, double Itrue, double sigma)
|
|
: ScalingResidual(r, Itrue, sigma),
|
|
delta_phi_deg(r.delta_phi_deg),
|
|
c1(r.zeta / std::sqrt(2.0)) {
|
|
}
|
|
|
|
template<typename T>
|
|
bool operator()(const T *const G,
|
|
const T *const B,
|
|
const T *const mosaicity,
|
|
const T *const wedge,
|
|
T *residual) const {
|
|
if (mosaicity[0] < 1e-6)
|
|
return false;
|
|
|
|
const T half_wedge = wedge[0] / T(2.0);
|
|
const T arg_plus = (T(delta_phi_deg) + half_wedge) * T(c1) / mosaicity[0];
|
|
const T arg_minus = (T(delta_phi_deg) - half_wedge) * T(c1) / mosaicity[0];
|
|
const T partiality = (ceres::erf(arg_plus) - ceres::erf(arg_minus)) / T(2.0);
|
|
const T B_term = ceres::exp(B[0] * T(b_resolution_coeff));
|
|
residual[0] = (G[0] * partiality * B_term * T(lp) * T(Itrue) - T(Iobs)) * T(weight);
|
|
return true;
|
|
}
|
|
|
|
double delta_phi_deg;
|
|
double c1;
|
|
};
|
|
|
|
struct IntensityFixedResidual : public ScalingResidual {
|
|
IntensityFixedResidual(const Reflection &r, double Itrue, double sigma, double partiality)
|
|
: ScalingResidual(r, Itrue, sigma),
|
|
partiality(partiality) {
|
|
}
|
|
|
|
template<typename T>
|
|
bool operator()(const T *const G, const T *const B, T *residual) const {
|
|
const T B_term = ceres::exp(B[0] * T(b_resolution_coeff));
|
|
residual[0] = (G[0] * T(partiality) * B_term * T(lp) * Itrue - T(Iobs)) * T(weight);
|
|
return true;
|
|
}
|
|
|
|
double partiality;
|
|
};
|
|
|
|
struct RotationNormRegularizer {
|
|
explicit RotationNormRegularizer(double weight) : weight(weight) {}
|
|
|
|
template<typename T>
|
|
bool operator()(const T *const rot_aa, T *residual) const {
|
|
residual[0] = T(weight) * rot_aa[0];
|
|
residual[1] = T(weight) * rot_aa[1];
|
|
residual[2] = T(weight) * rot_aa[2];
|
|
return true;
|
|
}
|
|
|
|
const double weight;
|
|
};
|
|
|
|
}
|
|
|
|
ScaleOnTheFly::ScaleOnTheFly(const DiffractionExperiment &x, const std::vector<MergedReflection> &ref)
|
|
: sg(x.GetGemmiSpaceGroup()),
|
|
model(x.GetPartialityModel()),
|
|
s(x.GetScalingSettings()),
|
|
rot_wedge_deg(x.GetRotationWedgeForScaling()),
|
|
refine_rot_wedge(x.GetRefineRotationWedgeInScaling()),
|
|
hkl_key_generator(s.GetMergeFriedel(), x.GetSpaceGroupNumber().value_or(1)) {
|
|
for (const auto &r: ref) {
|
|
const auto key = hkl_key_generator(r);
|
|
reference_data[key] = r.I;
|
|
}
|
|
}
|
|
|
|
bool ScaleOnTheFly::Accept(const Reflection &r) const {
|
|
if (!AcceptReflection(r, s.GetHighResolutionLimit_A()))
|
|
return false;
|
|
|
|
switch (model) {
|
|
case PartialityModel::Rotation:
|
|
return std::isfinite(r.zeta) && r.zeta > 0.0f;
|
|
case PartialityModel::Fixed:
|
|
case PartialityModel::Unity:
|
|
return true;
|
|
}
|
|
|
|
return true;
|
|
}
|
|
|
|
std::pair<double, size_t> ScaleOnTheFly::CalculateGlobalCC(const std::vector<Reflection> &reflections) const {
|
|
double sum_x = 0.0;
|
|
double sum_y = 0.0;
|
|
double sum_x2 = 0.0;
|
|
double sum_y2 = 0.0;
|
|
double sum_xy = 0.0;
|
|
size_t n = 0;
|
|
|
|
for (const auto &r: reflections) {
|
|
if (!AcceptReflection(r, s.GetHighResolutionLimit_A()))
|
|
continue;
|
|
if (r.partiality < s.GetMinPartiality())
|
|
continue;
|
|
if (!std::isfinite(r.I) || !std::isfinite(r.image_scale_corr) || r.image_scale_corr <= 0.0f)
|
|
continue;
|
|
if (!std::isfinite(r.sigma) || r.sigma <= 0.0f)
|
|
continue;
|
|
|
|
const HKLKey key = hkl_key_generator(r);
|
|
const auto it = reference_data.find(key);
|
|
if (it == reference_data.end())
|
|
continue;
|
|
|
|
const double image_i = static_cast<double>(r.I) * static_cast<double>(r.image_scale_corr);
|
|
const double ref_i = it->second;
|
|
|
|
if (!std::isfinite(image_i) || !std::isfinite(ref_i))
|
|
continue;
|
|
|
|
sum_x += image_i;
|
|
sum_y += ref_i;
|
|
sum_x2 += image_i * image_i;
|
|
sum_y2 += ref_i * ref_i;
|
|
sum_xy += image_i * ref_i;
|
|
++n;
|
|
}
|
|
|
|
if (n < MIN_REFLECTIONS)
|
|
return {NAN, n};
|
|
|
|
const double nd = static_cast<double>(n);
|
|
const double cov = sum_xy - sum_x * sum_y / nd;
|
|
const double var_x = sum_x2 - sum_x * sum_x / nd;
|
|
const double var_y = sum_y2 - sum_y * sum_y / nd;
|
|
|
|
if (!(var_x > 0.0 && var_y > 0.0))
|
|
return {NAN, n};
|
|
|
|
return {cov / std::sqrt(var_x * var_y), n};
|
|
}
|
|
|
|
void ScaleOnTheFly::Scale(IntegrationOutcome &integration_outcome) const {
|
|
if (integration_outcome.reflections.empty())
|
|
return;
|
|
|
|
auto start = std::chrono::steady_clock::now();
|
|
|
|
ScaleOnTheFlyResult result{
|
|
.B = 0.0,
|
|
.G = 1.0,
|
|
.R = {0.005, 0.005}
|
|
};
|
|
|
|
if (model == PartialityModel::Rotation) {
|
|
if (integration_outcome.mosaicity_deg
|
|
&& std::isfinite(*integration_outcome.mosaicity_deg)
|
|
&& *integration_outcome.mosaicity_deg > 0.0)
|
|
result.mos = *integration_outcome.mosaicity_deg;
|
|
else
|
|
result.mos = s.GetDefaultMosaicity();
|
|
if (const auto forced = s.GetForcedMosaicity(); forced.has_value())
|
|
result.mos = *forced;
|
|
result.wedge = rot_wedge_deg.value_or(0.0);
|
|
} else {
|
|
result.mos = NAN;
|
|
result.wedge = NAN;
|
|
}
|
|
|
|
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();
|
|
};
|
|
|
|
// 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 (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)});
|
|
}
|
|
|
|
if (obs.size() < MIN_REFLECTIONS) {
|
|
clear_scale();
|
|
return;
|
|
}
|
|
|
|
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());
|
|
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);
|
|
}
|
|
|
|
for (auto &r: integration_outcome.reflections) {
|
|
const double B_term = exp(result.B * -SafeInv(4.0 * r.d * r.d, 0.0));
|
|
|
|
switch (model) {
|
|
case PartialityModel::Unity:
|
|
r.partiality = 1.0;
|
|
break;
|
|
case PartialityModel::Rotation: {
|
|
if (std::isfinite(r.delta_phi_deg) && std::isfinite(r.zeta) && result.mos > 1e-6)
|
|
r.partiality = RotationPartiality(r.delta_phi_deg, r.zeta, result.mos, result.wedge);
|
|
break;
|
|
}
|
|
default:
|
|
// For fixed partiality there is no need to change anything
|
|
break;
|
|
}
|
|
const double denom = B_term * r.partiality * result.G;
|
|
if (std::isfinite(r.rlp) && std::isfinite(denom) && denom > 0.0) {
|
|
r.image_scale_corr = static_cast<float>(r.rlp / denom);
|
|
} else {
|
|
r.image_scale_corr = NAN;
|
|
}
|
|
}
|
|
|
|
const auto [cc, cc_n] = CalculateGlobalCC(integration_outcome.reflections);
|
|
result.cc = cc;
|
|
result.cc_n = cc_n;
|
|
|
|
auto end = std::chrono::steady_clock::now();
|
|
result.time_s = std::chrono::duration<float>(end - start).count();
|
|
|
|
integration_outcome.image_scale_cc = cc;
|
|
integration_outcome.image_scale_cc_n = cc_n;
|
|
integration_outcome.image_scale_g = result.G;
|
|
|
|
if (s.GetRefineB())
|
|
integration_outcome.image_scale_b_factor_Ang2 = result.B;
|
|
else
|
|
integration_outcome.image_scale_b_factor_Ang2.reset();
|
|
if (model == PartialityModel::Rotation) {
|
|
integration_outcome.mosaicity_deg = result.mos;
|
|
if (refine_rot_wedge)
|
|
integration_outcome.image_scale_wedge_deg = result.wedge;
|
|
else
|
|
integration_outcome.image_scale_wedge_deg.reset();
|
|
} else {
|
|
integration_outcome.mosaicity_deg.reset();
|
|
integration_outcome.image_scale_wedge_deg.reset();
|
|
}
|
|
}
|
|
|
|
void ScaleOnTheFly::Scale(std::vector<IntegrationOutcome> &integration, size_t nthreads) const {
|
|
if (nthreads == 0)
|
|
nthreads = std::thread::hardware_concurrency();
|
|
|
|
if (nthreads <= 1) {
|
|
for (auto & i : integration)
|
|
Scale(i);
|
|
} else {
|
|
auto local_nthreads = std::min(nthreads, integration.size());
|
|
std::vector<std::future<void>> futures;
|
|
futures.reserve(local_nthreads);
|
|
std::atomic<size_t> curr_image = 0;
|
|
|
|
for (size_t t = 0; t < local_nthreads; ++t)
|
|
futures.emplace_back(std::async(std::launch::async, [&] {
|
|
size_t i = curr_image.fetch_add(1);
|
|
while (i < integration.size()) {
|
|
Scale(integration[i]);
|
|
i = curr_image.fetch_add(1);
|
|
}
|
|
}));
|
|
|
|
for (auto &f: futures)
|
|
f.get();
|
|
}
|
|
}
|