leonarski_fandClaude Fable 5
cdbd0bed5d
BraggPredictionRot halves settings.wedge_deg to the +/- half-wedge of the partiality erf pair, but IndexAndRefine already passed GetWedge_deg()/2, so the two /2 compounded to a half-wedge of increment/4 - half the correct Kabsch value (increment/2, which ScaleOnTheFly's RotationPartiality already used). Pass the full increment so prediction's partiality matches scaling. With prediction correct, ScaleOnTheFly now uses the stored r.partiality directly (the value the reflection was integrated with) rather than recomputing the erf pair per reflection - recomputing only when scaling overrides the geometry (-w wedge refinement, --mosaicity, or a scaling wedge override). Output-neutral on the /data/rotation_test battery (SG/cell/completeness identical, ISa/CC1/2 within run noise on the stable crystals). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
480 lines
18 KiB
C++
480 lines
18 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 (r.on_ice_ring) // ice-contaminated intensity would drag the per-image scale; keep it out of the fit
|
|
return false;
|
|
|
|
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 (r.on_ice_ring)
|
|
continue;
|
|
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;
|
|
}
|
|
|
|
// The partiality stored on each reflection was computed at prediction with this image's mosaicity and
|
|
// the (now correct) rotation wedge, so it is used directly. Recompute it only when scaling overrides
|
|
// that geometry: wedge refinement, a forced mosaicity, or a scaling wedge override.
|
|
const bool recompute_partiality = model == PartialityModel::Rotation
|
|
&& (refine_rot_wedge || s.GetForcedMosaicity().has_value() || s.GetRotationWedgeForScaling().has_value());
|
|
|
|
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 = recompute_partiality
|
|
? RotationPartiality(r.delta_phi_deg, r.zeta, result.mos, result.wedge)
|
|
: r.partiality;
|
|
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 (recompute_partiality && 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();
|
|
}
|
|
}
|