Files
Jungfraujoch/image_analysis/scale_merge/ScaleOnTheFly.cpp
T

533 lines
20 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 <future>
#include <ceres/ceres.h>
#include <ceres/rotation.h>
namespace {
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);
}
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;
};
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(T(B[0]) * b_resolution_coeff);
residual[0] = (G[0] * T(partiality) * B_term * T(lp) * Itrue - T(Iobs)) * T(weight);
return true;
}
double partiality;
};
struct ScalingPostRefResidual : public ScalingResidual {
// In this algorithm at the point of post-refinement we don't anymore care for where maximum of
// the reflection was located and if it fits the observed position.
// This reflections was already integrated and we cannot integrate it better at this point.
// But, we could adjust partiality to indicate that this reflection was wrongly predicted.
// I.e., integrated position was far away from true reflection, so partiality must be low.
// This is an empiric model and need to see if this will work in practice at all.
// I hope it will allow the model to find that reflections were misindexed.
// We assume at this point that initial indexing was done properly and integration was generally OK => most low resolution reflections fit correctly
// Yet we know, that small errors in indexing are inducing misalignment at high resolution - sometimes it is visible that high-resolution reflections
// are away from shoe-boxes observed in the image, if we can catch this at post-refinement/scaling step, this would be great.
// Next logical step is to do this pixel-wise - for each pixel refine partiality and merge pixels
// This should work for per-image scaling, or even, maybe, for full rotation datasets (3600 images)
// Then we could properly take into account misalignment of shoe-box center vs. partiality and also remove most pixels
// in the shoe-box that don't really contribute to the reflection.
// But...it could also drift to downweighting partiality for all high resolution reflections to make loss function "fake happy".
// We assume rot3 == 0. Rot3 is not really helping much in crystallography (other than fixing polarization correction)
ScalingPostRefResidual(const Reflection &r, double Itrue, double sigma,
const DiffractionGeometry &geometry,
const CrystalLattice &lattice,
double exp_h, double exp_k,
double exp_l)
: ScalingResidual(r, Itrue, sigma),
integration_center_x(r.predicted_x),
integration_center_y(r.predicted_y),
inv_lambda(SafeInv(geometry.GetWavelength_A(), 0.0)),
pixel_size(geometry.GetPixelSize_mm()),
det_dist_mm(geometry.GetDetectorDistance_mm()),
beam_x(geometry.GetBeamX_pxl()),
beam_y(geometry.GetBeamY_pxl()),
exp_h(exp_h),
exp_k(exp_k),
exp_l(exp_l),
Astar(lattice.Astar()),
Bstar(lattice.Bstar()),
Cstar(lattice.Cstar()),
c1(std::cos(geometry.GetPoniRot1_rad())),
s1(std::sin(geometry.GetPoniRot1_rad())),
c2(std::cos(geometry.GetPoniRot2_rad())),
s2(std::sin(geometry.GetPoniRot2_rad())) {
}
template <typename T>
T CalcPartiality(const T *const R_radial,
const T *const R_tangential,
const T *const beam_corr,
const T *const p0) const {
// Detector coordinates in mm
const T det_x = (T(integration_center_x) - beam_x - beam_corr[0]) * T(pixel_size);
const T det_y = (T(integration_center_y) - beam_y - beam_corr[1]) * T(pixel_size);
const T det_z = T(det_dist_mm);
// Apply Ry(rot1) first: rotate around Y
const T t1_x = T(c1) * det_x + T(s1) * det_z;
const T t1_y = det_y;
const T t1_z = T(-s1) * det_x + T(c1) * det_z;
// Then apply Rx(-rot2): rotate around X
const T x = t1_x;
const T y = T(c2) * t1_y + T(s2) * t1_z;
const T z = - T(s2) * t1_y + T(c2) * t1_z;
// convert to recip space
const T lab_norm = ceres::sqrt(x * x + y * y + z * z);
const T inv_norm = T(1) / lab_norm;
T recip_obs[3];
recip_obs[0] = x * inv_norm * inv_lambda;
recip_obs[1] = y * inv_norm * inv_lambda;
recip_obs[2] = (z * inv_norm - T(1.0)) * inv_lambda;
const Eigen::Matrix<T, 3, 1> e_obs_recip(recip_obs[0], recip_obs[1], recip_obs[2]);
const T astar_unrot[3] = {T(Astar.x), T(Astar.y), T(Astar.z)};
const T bstar_unrot[3] = {T(Bstar.x), T(Bstar.y), T(Bstar.z)};
const T cstar_unrot[3] = {T(Cstar.x), T(Cstar.y), T(Cstar.z)};
T astar_rot[3], bstar_rot[3], cstar_rot[3];
ceres::AngleAxisRotatePoint(p0, astar_unrot, astar_rot);
ceres::AngleAxisRotatePoint(p0, bstar_unrot, bstar_rot);
ceres::AngleAxisRotatePoint(p0, cstar_unrot, cstar_rot);
const Eigen::Matrix<T, 3, 1> e_pred_recip(T(exp_h) * astar_rot[0] + T(exp_k) * bstar_rot[0] + T(exp_l) * cstar_rot[0],
T(exp_h) * astar_rot[1] + T(exp_k) * bstar_rot[1] + T(exp_l) * cstar_rot[1],
T(exp_h) * astar_rot[2] + T(exp_k) * bstar_rot[2] + T(exp_l) * cstar_rot[2]
);
// Ewald sphere centre is at -k_i = (0, 0, -inv_lambda)
// Radial direction: outward normal at g_hkl
const Eigen::Matrix<T, 3, 1> S_pred(
e_pred_recip[0],
e_pred_recip[1],
e_pred_recip[2] + T(inv_lambda) // g_hkl + k_i
);
const T S_pred_norm = S_pred.norm();
if (S_pred_norm < T(1e-10))
return T(0);
const Eigen::Matrix<T, 3, 1> n_radial = S_pred / S_pred_norm;
const Eigen::Matrix<T, 3, 1> delta_q = e_obs_recip - e_pred_recip;
const T eps_radial = delta_q.dot(n_radial);
const Eigen::Matrix<T, 3, 1> dq_tang = delta_q - eps_radial * n_radial;
const T eps_tangential_sq = dq_tang.squaredNorm(); // guaranteed ≥ 0
// ─────────────────────────────────────────────────────────────
return ceres::exp(
- eps_radial * eps_radial / (R_radial[0] * R_radial[0])
- eps_tangential_sq / (R_tangential[0] * R_tangential[0])
);
}
template<typename T>
bool operator()(const T *const G,
const T *const B,
const T *const R_radial,
const T *const R_tangential,
const T *const beam_corr,
const T *const p0,
T *residual) const {
if (R_radial[0] < T(1e-10) || R_tangential[0] < T(1e-10))
return false;
const T B_term = ceres::exp(B[0] * T(b_resolution_coeff));
const T partiality = CalcPartiality(R_radial, R_tangential, beam_corr, p0);
residual[0] = (G[0] * partiality * B_term * T(lp) * T(Itrue)
- T(Iobs)) * T(weight);
return true;
}
const double integration_center_x, integration_center_y;
const double inv_lambda;
const double pixel_size;
const double det_dist_mm;
const double beam_x, beam_y;
const double exp_h;
const double exp_k;
const double exp_l;
const Coord Astar, Bstar, Cstar;
const double c1,s1,c2,s2;
};
}
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();
ceres::Problem problem;
ScaleOnTheFlyResult result{
.B = 0.0,
.G = 1.0
};
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();
result.wedge = rot_wedge_deg.value_or(0.0);
} else {
result.mos = NAN;
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, nullptr, &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, nullptr, &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, nullptr, &result.G, &result.B, &result.mos,
&result.wedge);
}
break;
default:
throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
"Not supported partiality model");
}
}
if (n_reflections < MIN_REFLECTIONS) {
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);
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);
}
problem.SetParameterLowerBound(&result.B, 0, s.GetMinB());
problem.SetParameterUpperBound(&result.B, 0, s.GetMaxB());
} else {
problem.SetParameterBlockConstant(&result.B);
}
if (model == PartialityModel::Rotation) {
if (refine_rot_wedge) {
problem.SetParameterLowerBound(&result.wedge, 0, s.GetMinWedge());
problem.SetParameterUpperBound(&result.wedge, 0, s.GetMaxWedge());
} else {
problem.SetParameterBlockConstant(&result.wedge);
}
problem.SetParameterLowerBound(&result.mos, 0, s.GetMinMosaicity());
problem.SetParameterUpperBound(&result.mos, 0, s.GetMaxMosaicity());
}
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();
}
}