leonarski_fandClaude Opus 4.8
1a2b0181a5
Replace the frozen scalar-sigma stills partiality with a physical, refined model. Per crystal, refine an orientation tilt (dpsi_x, dpsi_y) against the running merge and recompute each reflection's partiality analytically from the refined geometry (angular Ewald-proximity model, sigma(d*) = gamma_e*d*), with the per-crystal scale G profiled out by the existing robust IRLS - no re-integration. A soft Gaussian prior on dpsi tames weak-data overfit while staying inert on strong data. The merge <-> refine loop iterates a few times. This is now the stills default via ScalingSettings::stills_partiality_refine (on). A single opt-out flag `--simple-stills` reverts to treating every reflection as a full (p=1, single pass). Retires the experimental `--still-partiality` flag. The viewer gains a "Partiality post-refinement (stills)" checkbox in Scaling settings. Validated (integrate-once / --scale): CC1/2 and R_meas both improve on three monochromatic serial-stills datasets (+2.8 / -10, +5.6 / -3.4, +2.1 / -4); neutral on a pink-beam DMM set (already-full reflections); R-free/R-work down vs a fixed model; competitive with CrystFEL partialator on matched frames. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
330 lines
14 KiB
C++
330 lines
14 KiB
C++
// SPDX-FileCopyrightText: 2025 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
|
|
// SPDX-License-Identifier: GPL-3.0-only
|
|
|
|
#include "StillsPartialityRefine.h"
|
|
|
|
#include <atomic>
|
|
#include <cmath>
|
|
#include <future>
|
|
#include <thread>
|
|
#include <vector>
|
|
|
|
#include <ceres/ceres.h>
|
|
#include <ceres/rotation.h>
|
|
|
|
#include "Merge.h"
|
|
|
|
namespace {
|
|
constexpr size_t MIN_FIT_REFLECTIONS = 20;
|
|
constexpr double kRadToDeg = 180.0 / 3.14159265358979323846;
|
|
|
|
double SafeInv(double x, double fallback) {
|
|
if (!std::isfinite(x) || x == 0.0)
|
|
return fallback;
|
|
return 1.0 / x;
|
|
}
|
|
|
|
// One accepted reflection reduced to the physical partiality fit: the base reciprocal vector q (crystal
|
|
// frame at the stored per-image orientation), the reference full intensity, the measured intensity, the
|
|
// Lorentz factor and the weight. dist_ewald / partiality are recomputed from q under the refined tilt.
|
|
struct FitObs {
|
|
double qx, qy, qz;
|
|
double Iref;
|
|
double Iobs;
|
|
double lp; // 1 / rlp
|
|
double weight; // 1 / sigma
|
|
};
|
|
|
|
double VecLen(double x, double y, double z) { return std::sqrt(x * x + y * y + z * z); }
|
|
|
|
// Analytic partiality for a reflection whose base reciprocal vector is q, tilted by (psi_x, psi_y).
|
|
// Mirrors BraggPrediction: dist_ewald = |S| - 1/lambda with S = q_rot + S0, and
|
|
// p = exp(-dist_ewald^2 / 2 sigma^2), sigma^2 = (gamma0 + gamma_e*d*)^2 + (bw*|q_z|)^2.
|
|
double ComputeP(double qx, double qy, double qz,
|
|
double psi_x, double psi_y,
|
|
double s0x, double s0y, double s0z, double inv_lambda,
|
|
double gamma0, double gamma_e, double bw) {
|
|
double aa[3] = {psi_x, psi_y, 0.0};
|
|
double q[3] = {qx, qy, qz};
|
|
double qr[3];
|
|
ceres::AngleAxisRotatePoint(aa, q, qr);
|
|
const double Sx = qr[0] + s0x, Sy = qr[1] + s0y, Sz = qr[2] + s0z;
|
|
const double de = std::sqrt(Sx * Sx + Sy * Sy + Sz * Sz) - inv_lambda;
|
|
const double dstar = std::sqrt(qr[0] * qr[0] + qr[1] * qr[1] + qr[2] * qr[2]);
|
|
const double sig = gamma0 + gamma_e * dstar;
|
|
const double sbw = bw * std::fabs(qr[2]);
|
|
const double sig2 = sig * sig + sbw * sbw;
|
|
if (!(sig2 > 0.0))
|
|
return 1.0;
|
|
return std::exp(-0.5 * de * de / sig2);
|
|
}
|
|
|
|
// Robust per-crystal scale G (linear in G given the model coefficients), identical objective to
|
|
// ScaleOnTheFly::SolveScaleIRLS: minimise sum Cauchy_k( w (G*coeff - Iobs) ) over G >= 0.
|
|
double SolveScaleIRLS(const std::vector<double> &coeff, const std::vector<double> &Iobs,
|
|
const std::vector<double> &weight, double robust_k) {
|
|
auto weighted_scale = [&](auto robust_weight) {
|
|
double num = 0.0, den = 0.0;
|
|
for (size_t i = 0; i < coeff.size(); ++i) {
|
|
const double rw = robust_weight(i);
|
|
const double w2 = weight[i] * weight[i];
|
|
num += rw * w2 * coeff[i] * Iobs[i];
|
|
den += rw * w2 * coeff[i] * coeff[i];
|
|
}
|
|
return den > 0.0 ? num / den : NAN;
|
|
};
|
|
|
|
double G = weighted_scale([](size_t) { 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([&](size_t i) {
|
|
const double res = weight[i] * (G * coeff[i] - Iobs[i]);
|
|
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;
|
|
}
|
|
|
|
// Ceres residual: refine the orientation tilt (psi_x, psi_y) holding the scale G fixed. The tilt
|
|
// rotates the base reciprocal vector q; partiality follows analytically. Residual is the intensity
|
|
// mismatch weighted by 1/sigma, exactly matching ScaleOnTheFly's intensity-space objective.
|
|
struct PsiResidual {
|
|
double qx, qy, qz;
|
|
double s0x, s0y, s0z, inv_lambda;
|
|
double gamma0, gamma_e, bw;
|
|
double G, lp, Iref, Iobs, weight;
|
|
|
|
template<typename T>
|
|
bool operator()(const T *const psi, T *residual) const {
|
|
T q[3] = {T(qx), T(qy), T(qz)};
|
|
T aa[3] = {psi[0], psi[1], T(0.0)};
|
|
T qr[3];
|
|
ceres::AngleAxisRotatePoint(aa, q, qr);
|
|
const T Sx = qr[0] + T(s0x), Sy = qr[1] + T(s0y), Sz = qr[2] + T(s0z);
|
|
const T de = ceres::sqrt(Sx * Sx + Sy * Sy + Sz * Sz) - T(inv_lambda);
|
|
const T dstar = ceres::sqrt(qr[0] * qr[0] + qr[1] * qr[1] + qr[2] * qr[2]);
|
|
const T sig = T(gamma0) + T(gamma_e) * dstar;
|
|
const T sbw = T(bw) * ceres::abs(qr[2]);
|
|
const T sig2 = sig * sig + sbw * sbw;
|
|
const T p = ceres::exp(T(-0.5) * de * de / sig2);
|
|
residual[0] = T(weight) * (T(G) * p * T(lp) * T(Iref) - T(Iobs));
|
|
return true;
|
|
}
|
|
};
|
|
|
|
// Gaussian prior N(0, sigma_prior^2) on the tilt. The data residuals above are (model-obs)/sigma, a
|
|
// proper chi^2, so the MAP prior residual is simply dpsi/sigma_prior - no scale calibration needed.
|
|
struct PsiPrior {
|
|
double inv_sigma;
|
|
template<typename T>
|
|
bool operator()(const T *const psi, T *residual) const {
|
|
residual[0] = T(inv_sigma) * psi[0];
|
|
residual[1] = T(inv_sigma) * psi[1];
|
|
return true;
|
|
}
|
|
};
|
|
}
|
|
|
|
StillsPartialityRefine::StillsPartialityRefine(const DiffractionExperiment &x)
|
|
: StillsPartialityRefine(x, Settings{}) {}
|
|
|
|
StillsPartialityRefine::StillsPartialityRefine(const DiffractionExperiment &x, Settings settings)
|
|
: experiment_(x),
|
|
settings_(settings),
|
|
hkl_key_generator_(x.GetScalingSettings().GetMergeFriedel(), x.GetSpaceGroupNumber().value_or(1)),
|
|
d_min_limit_(x.GetScalingSettings().GetHighResolutionLimit_A()),
|
|
bandwidth_sigma_(x.GetBandwidthFWHM().value_or(0.0f) / 2.3548f) {}
|
|
|
|
double StillsPartialityRefine::RefineOne(IntegrationOutcome &outcome,
|
|
const std::map<HKLKey, double> &reference) const {
|
|
if (outcome.reflections.empty())
|
|
return 0.0;
|
|
|
|
const Coord Astar = outcome.latt.Astar();
|
|
const Coord Bstar = outcome.latt.Bstar();
|
|
const Coord Cstar = outcome.latt.Cstar();
|
|
const Coord S0 = outcome.geom.GetScatteringVector();
|
|
const double inv_lambda = 1.0 / outcome.geom.GetWavelength_A();
|
|
const double bw = bandwidth_sigma_;
|
|
const double gamma0 = 0.0; // width is purely angular: sigma(d*) = gamma_e * d* (set per crystal below)
|
|
|
|
auto base_q = [&](const Reflection &r) {
|
|
return Astar * static_cast<float>(r.h) + Bstar * static_cast<float>(r.k)
|
|
+ Cstar * static_cast<float>(r.l);
|
|
};
|
|
|
|
// Collect the reflections that constrain the fit (accepted, non-ice, finite, present in the reference).
|
|
std::vector<FitObs> obs;
|
|
obs.reserve(outcome.reflections.size());
|
|
double sum_ang2 = 0.0; // RMS angular excitation error (dist_ewald / d*) -> per-crystal mosaic width
|
|
size_t n_de = 0;
|
|
for (const Reflection &r: outcome.reflections) {
|
|
if (r.on_ice_ring || !AcceptReflection(r, d_min_limit_))
|
|
continue;
|
|
if (!std::isfinite(r.I) || !std::isfinite(r.sigma) || r.sigma <= 0.0f)
|
|
continue;
|
|
const auto it = reference.find(hkl_key_generator_(r));
|
|
if (it == reference.end() || !std::isfinite(it->second))
|
|
continue;
|
|
|
|
const Coord q = base_q(r);
|
|
obs.push_back(FitObs{
|
|
.qx = q.x, .qy = q.y, .qz = q.z,
|
|
.Iref = it->second,
|
|
.Iobs = static_cast<double>(r.I),
|
|
.lp = SafeInv(r.rlp, 1.0),
|
|
.weight = SafeInv(r.sigma, 1.0),
|
|
});
|
|
|
|
// Angular excitation error delta_psi = dist_ewald / d* at the stored orientation (psi = 0). Using
|
|
// the ANGULAR distance (not the linear reciprocal-space distance) makes the partiality width
|
|
// resolution-clean: a fixed mosaic angle smears high-resolution rlps more in reciprocal space, so a
|
|
// constant linear width computes p too small at high resolution and over-divides those shells.
|
|
const double dstar = VecLen(q.x, q.y, q.z);
|
|
const double de0 = VecLen(q.x + S0.x, q.y + S0.y, q.z + S0.z) - inv_lambda;
|
|
if (dstar > 1e-9) {
|
|
const double dpsi = de0 / dstar;
|
|
sum_ang2 += dpsi * dpsi;
|
|
++n_de;
|
|
}
|
|
}
|
|
|
|
if (obs.size() < MIN_FIT_REFLECTIONS || n_de == 0)
|
|
return 0.0;
|
|
|
|
// Per-crystal angular mosaic width from the RMS angular excitation error. sigma(d*) = gamma_e * d*
|
|
// (gamma0 = 0), i.e. p = exp(-0.5 (delta_psi / gamma_e)^2) is a Gaussian in the angular distance from
|
|
// the Ewald sphere - the physical mosaic/divergence model, independent of resolution. A positive
|
|
// settings_.gamma_e overrides the per-crystal estimate with a shared (pooled) width.
|
|
const double gamma_e_ang = std::max(std::sqrt(sum_ang2 / static_cast<double>(n_de)), 1e-9);
|
|
const double gamma_e = settings_.gamma_e > 0.0 ? settings_.gamma_e : gamma_e_ang;
|
|
|
|
double psi[2] = {0.0, 0.0};
|
|
double G = 1.0;
|
|
|
|
const bool refine_tilt = obs.size() >= settings_.min_reflections;
|
|
const int inner = refine_tilt ? settings_.inner_iterations : 1;
|
|
|
|
for (int it = 0; it < inner; ++it) {
|
|
// (1) Solve G given the current partialities.
|
|
std::vector<double> coeff(obs.size()), Iobs(obs.size()), weight(obs.size());
|
|
for (size_t j = 0; j < obs.size(); ++j) {
|
|
const double p = ComputeP(obs[j].qx, obs[j].qy, obs[j].qz, psi[0], psi[1],
|
|
S0.x, S0.y, S0.z, inv_lambda, gamma0, gamma_e, bw);
|
|
coeff[j] = p * obs[j].lp * obs[j].Iref;
|
|
Iobs[j] = obs[j].Iobs;
|
|
weight[j] = obs[j].weight;
|
|
}
|
|
G = SolveScaleIRLS(coeff, Iobs, weight, settings_.robust_k);
|
|
if (!(G > 0.0) || !std::isfinite(G))
|
|
return 0.0;
|
|
|
|
if (!refine_tilt)
|
|
break;
|
|
|
|
// (2) Refine the tilt holding G fixed.
|
|
ceres::Problem problem;
|
|
for (const auto &o: obs) {
|
|
auto *cost = new ceres::AutoDiffCostFunction<PsiResidual, 1, 2>(new PsiResidual{
|
|
.qx = o.qx, .qy = o.qy, .qz = o.qz,
|
|
.s0x = S0.x, .s0y = S0.y, .s0z = S0.z, .inv_lambda = inv_lambda,
|
|
.gamma0 = gamma0, .gamma_e = gamma_e, .bw = bw,
|
|
.G = G, .lp = o.lp, .Iref = o.Iref, .Iobs = o.Iobs, .weight = o.weight});
|
|
problem.AddResidualBlock(cost, new ceres::CauchyLoss(settings_.robust_k), psi);
|
|
}
|
|
if (settings_.prior_sigma_deg > 0.0) {
|
|
const double inv_sigma = kRadToDeg / settings_.prior_sigma_deg; // 1 / sigma_prior (rad)
|
|
problem.AddResidualBlock(new ceres::AutoDiffCostFunction<PsiPrior, 2, 2>(
|
|
new PsiPrior{inv_sigma}), nullptr, psi);
|
|
}
|
|
problem.SetParameterLowerBound(psi, 0, -settings_.max_tilt_rad);
|
|
problem.SetParameterUpperBound(psi, 0, settings_.max_tilt_rad);
|
|
problem.SetParameterLowerBound(psi, 1, -settings_.max_tilt_rad);
|
|
problem.SetParameterUpperBound(psi, 1, settings_.max_tilt_rad);
|
|
|
|
ceres::Solver::Options options;
|
|
options.linear_solver_type = ceres::DENSE_QR;
|
|
options.minimizer_progress_to_stdout = false;
|
|
options.num_threads = 1;
|
|
options.max_num_iterations = 25;
|
|
ceres::Solver::Summary summary;
|
|
ceres::Solve(options, &problem, &summary);
|
|
}
|
|
|
|
// Write the refined partiality + scale correction onto every reflection of the crystal (not only the
|
|
// fit subset), so the merge sees a consistent model. image_scale_corr = rlp / (partiality * G).
|
|
for (auto &r: outcome.reflections) {
|
|
const Coord q = base_q(r);
|
|
const double p = ComputeP(q.x, q.y, q.z, psi[0], psi[1], S0.x, S0.y, S0.z, inv_lambda,
|
|
gamma0, gamma_e, bw);
|
|
r.partiality = static_cast<float>(p);
|
|
const double denom = p * G;
|
|
r.image_scale_corr = (std::isfinite(r.rlp) && denom > 0.0)
|
|
? static_cast<float>(r.rlp / denom)
|
|
: NAN;
|
|
}
|
|
outcome.image_scale_g = static_cast<float>(G);
|
|
|
|
const double tilt_deg = std::sqrt(psi[0] * psi[0] + psi[1] * psi[1]) * kRadToDeg;
|
|
return tilt_deg;
|
|
}
|
|
|
|
double StillsPartialityRefine::Run(std::vector<IntegrationOutcome> &outcomes, size_t nthreads) const {
|
|
if (nthreads == 0)
|
|
nthreads = std::thread::hardware_concurrency();
|
|
nthreads = std::max<size_t>(1, nthreads);
|
|
|
|
double last_mean_tilt = 0.0;
|
|
|
|
for (int outer = 0; outer < settings_.outer_iterations; ++outer) {
|
|
// Reference full intensities from the current corrections.
|
|
const std::vector<MergedReflection> merged = MergeAll(experiment_, outcomes, false);
|
|
std::map<HKLKey, double> reference;
|
|
for (const auto &m: merged)
|
|
reference[hkl_key_generator_(m)] = m.I;
|
|
|
|
std::atomic<double> tilt_sum{0.0};
|
|
std::atomic<size_t> tilt_n{0};
|
|
std::atomic<size_t> next{0};
|
|
|
|
auto worker = [&]() {
|
|
size_t i = next.fetch_add(1);
|
|
while (i < outcomes.size()) {
|
|
const double t = RefineOne(outcomes[i], reference);
|
|
if (t > 0.0) {
|
|
double prev = tilt_sum.load();
|
|
while (!tilt_sum.compare_exchange_weak(prev, prev + t)) {}
|
|
tilt_n.fetch_add(1);
|
|
}
|
|
i = next.fetch_add(1);
|
|
}
|
|
};
|
|
|
|
const size_t nt = std::min(nthreads, std::max<size_t>(1, outcomes.size()));
|
|
if (nt <= 1) {
|
|
worker();
|
|
} else {
|
|
std::vector<std::future<void>> futures;
|
|
futures.reserve(nt);
|
|
for (size_t t = 0; t < nt; ++t)
|
|
futures.emplace_back(std::async(std::launch::async, worker));
|
|
for (auto &f: futures)
|
|
f.get();
|
|
}
|
|
|
|
last_mean_tilt = tilt_n > 0 ? tilt_sum.load() / static_cast<double>(tilt_n.load()) : 0.0;
|
|
}
|
|
|
|
return last_mean_tilt;
|
|
}
|