ScaleAndMerge: Fix partiality model
All checks were successful
Build Packages / build:rpm (ubuntu2404_nocuda) (push) Successful in 12m46s
Build Packages / build:rpm (ubuntu2204_nocuda) (push) Successful in 13m31s
Build Packages / build:rpm (rocky8_nocuda) (push) Successful in 13m48s
Build Packages / Generate python client (push) Successful in 17s
Build Packages / build:rpm (ubuntu2204) (push) Successful in 14m8s
Build Packages / Create release (push) Has been skipped
Build Packages / build:rpm (rocky8_sls9) (push) Successful in 14m14s
Build Packages / build:rpm (rocky8) (push) Successful in 14m13s
Build Packages / Build documentation (push) Successful in 37s
Build Packages / build:rpm (rocky9_nocuda) (push) Successful in 14m52s
Build Packages / build:rpm (rocky9) (push) Successful in 14m58s
Build Packages / build:rpm (ubuntu2404) (push) Successful in 7m9s
Build Packages / Unit tests (push) Successful in 50m39s

This commit is contained in:
2026-02-09 12:29:48 +01:00
parent c40a0332e7
commit 684c3df204
5 changed files with 68 additions and 106 deletions

View File

@@ -47,8 +47,7 @@ int BraggPredictionRot::Calc(const DiffractionExperiment &experiment, const Crys
int i = 0;
const float mos_angle_rad = settings.mosaicity_deg * static_cast<float>(M_PI) / 180.f;
const float wedge_angle_rad = settings.wedge_deg * static_cast<float>(M_PI) / 180.f;
const float max_angle_rad = wedge_angle_rad / 2 + mos_angle_rad;
const float half_wedge_angle_rad = settings.wedge_deg * static_cast<float>(M_PI) / 180.f / 2.0f ;
for (int h = -settings.max_hkl; h <= settings.max_hkl; h++) {
// Precompute A* h contribution
@@ -109,10 +108,10 @@ int BraggPredictionRot::Calc(const DiffractionExperiment &experiment, const Crys
continue;
const float lorentz_reciprocal = std::fabs(m2 * (S % S0))/(S*S0);
const float c1 = std::sqrt(2.0f) * mos_angle_rad / zeta_abs;
const float c1 = zeta_abs / (std::sqrt(2.0f) * mos_angle_rad);
const float partiality = (std::erf((phi + wedge_angle_rad / 2.0f) * c1)
- std::erf((phi - wedge_angle_rad / 2.0f) * c1)) / 2.0f;
const float partiality = (std::erf((phi + half_wedge_angle_rad) * c1)
- std::erf((phi - half_wedge_angle_rad) * c1)) / 2.0f;
// Inlined RecipToDector with rot1 and rot2 (rot3 = 0)
// Apply rotation matrix transpose
float S_rot_x = rot[0] * S.x + rot[1] * S.y + rot[2] * S.z;

View File

@@ -125,7 +125,7 @@ namespace {
// Partiality calculation (Kabsch formulation)
// c1 = sqrt(2) * sigma / zeta, where sigma = mosaicity
float c1 = sqrtf(2.0f) * C.mos_angle_rad / zeta_abs;
float c1 = zeta_abs / (sqrtf(2.0f) * C.mos_angle_rad);
float half_wedge = C.wedge_angle_rad / 2.0f;
float partiality = (erff((phi + half_wedge) * c1)
- erff((phi - half_wedge) * c1)) / 2.0f;

View File

@@ -114,39 +114,30 @@ inline HKLKey CanonicalizeHKLKey(const Reflection& r, const ScaleMergeOptions& o
}
struct IntensityResidual {
IntensityResidual(const Reflection& r, double sigma_obs, double wedge_deg, double s2, bool refine_b,
bool partiality_model)
IntensityResidual(const Reflection& r, double sigma_obs, double wedge_deg, bool refine_partiality)
: Iobs_(static_cast<double>(r.I)),
inv_sigma_(SafeInv(sigma_obs, 1.0)),
delta_phi_rad_(static_cast<double>(r.delta_phi) * M_PI / 180.0),
lp_(SafeInv(static_cast<double>(r.rlp), 1.0)),
half_wedge_rad_(wedge_deg * M_PI / 180.0 / 2.0),
c1_(std::sqrt(2.0) * SafeInv(static_cast<double>(r.zeta), 1.0)),
s2_(s2),
refine_b_(refine_b),
partiality_model_(partiality_model) {}
c1_(r.zeta / std::sqrt(2.0)),
partiality_(r.partiality),
refine_partiality_(refine_partiality) {}
template<typename T>
bool operator()(const T* const log_k,
const T* const b,
bool operator()(const T* const k,
const T* const mosaicity_rad,
const T* const log_Itrue,
const T* const Itrue,
T* residual) const {
const T k = ceres::exp(log_k[0]);
const T Itrue = ceres::exp(log_Itrue[0]);
T partiality = T(1.0);
if (partiality_model_) {
const T arg_plus = T(delta_phi_rad_ + half_wedge_rad_) * (T(c1_) * mosaicity_rad[0]);
const T arg_minus = T(delta_phi_rad_ - half_wedge_rad_) * (T(c1_) * mosaicity_rad[0]);
T partiality;
if (refine_partiality_ && mosaicity_rad[0] != 0.0) {
const T arg_plus = T(delta_phi_rad_ + half_wedge_rad_) * T(c1_) / mosaicity_rad[0];
const T arg_minus = T(delta_phi_rad_ - half_wedge_rad_) * T(c1_) / mosaicity_rad[0];
partiality = (ceres::erf(arg_plus) - ceres::erf(arg_minus)) / T(2.0);
}
T wilson = T(1.0);
if (refine_b_) {
wilson = ceres::exp(-b[0] * T(s2_));
}
} else
partiality = T(partiality_);
const T Ipred = k * wilson * partiality * T(lp_) * Itrue;
const T Ipred = k[0] * partiality * T(lp_) * Itrue[0];
residual[0] = (Ipred - T(Iobs_)) * T(inv_sigma_);
return true;
}
@@ -157,9 +148,8 @@ struct IntensityResidual {
double lp_;
double half_wedge_rad_;
double c1_;
double s2_;
bool refine_b_;
bool partiality_model_;
double partiality_;
bool refine_partiality_;
};
struct ScaleRegularizationResidual {
@@ -167,9 +157,8 @@ struct ScaleRegularizationResidual {
: inv_sigma_(SafeInv(sigma_k, 1.0)) {}
template <typename T>
bool operator()(const T* const log_k, T* residual) const {
const T k = ceres::exp(log_k[0]);
residual[0] = (k - T(1.0)) * T(inv_sigma_);
bool operator()(const T* const k, T* residual) const {
residual[0] = (k[0] - T(1.0)) * T(inv_sigma_);
return true;
}
@@ -187,7 +176,6 @@ ScaleMergeResult ScaleAndMergeReflectionsCeres(const std::vector<Reflection>& ob
int img_id = 0;
int img_slot = -1;
int hkl_slot = -1;
double s2 = 0.0;
double sigma = 0.0;
};
@@ -213,7 +201,6 @@ ScaleMergeResult ScaleAndMergeReflectionsCeres(const std::vector<Reflection>& ob
continue;
const double sigma = SafeSigma(static_cast<double>(r.sigma), opt.min_sigma);
const double s2 = 1.0 / (d * d);
const int img_id = RoundImageId(r.image_number, opt.image_number_rounding);
@@ -247,7 +234,6 @@ ScaleMergeResult ScaleAndMergeReflectionsCeres(const std::vector<Reflection>& ob
o.img_id = img_id;
o.img_slot = img_slot;
o.hkl_slot = hkl_slot;
o.s2 = s2;
o.sigma = sigma;
obs.push_back(o);
}
@@ -256,16 +242,14 @@ ScaleMergeResult ScaleAndMergeReflectionsCeres(const std::vector<Reflection>& ob
const int nhkl = static_cast<int>(hklToSlot.size());
out.image_scale_k.assign(nimg, 1.0);
out.image_b_factor.assign(nimg, 0.0);
out.image_ids.assign(nimg, 0);
for (const auto& kv : imgIdToSlot) {
out.image_ids[kv.second] = kv.first;
}
std::vector<double> log_k(nimg, 0.0);
std::vector<double> b(nimg, 0.0);
std::vector<double> log_Itrue(nhkl, 0.0);
std::vector<double> k(nimg, 1.0);
std::vector<double> Itrue(nhkl, 0.0);
auto deg2rad = [](double deg) { return deg * M_PI / 180.0; };
@@ -281,31 +265,30 @@ ScaleMergeResult ScaleAndMergeReflectionsCeres(const std::vector<Reflection>& ob
for (int h = 0; h < nhkl; ++h) {
auto& v = per_hkl_I[h];
if (v.empty()) {
log_Itrue[h] = std::log(std::max(opt.min_sigma, 1e-6));
Itrue[h] = std::max(opt.min_sigma, 1e-6);
continue;
}
std::nth_element(v.begin(), v.begin() + static_cast<long>(v.size() / 2), v.end());
double med = v[v.size() / 2];
if (!std::isfinite(med) || med <= opt.min_sigma)
med = opt.min_sigma;
log_Itrue[h] = std::log(med);
Itrue[h] = med;
}
}
ceres::Problem problem;
const bool partiality_model = opt.wedge_deg > 0.0;
const bool refine_partiality = opt.wedge_deg > 0.0;
for (const auto& o : obs) {
auto* cost = new ceres::AutoDiffCostFunction<IntensityResidual, 1, 1, 1, 1, 1>(
new IntensityResidual(*o.r, o.sigma, opt.wedge_deg, o.s2, opt.refine_b_factor, partiality_model));
auto* cost = new ceres::AutoDiffCostFunction<IntensityResidual, 1, 1, 1, 1>(
new IntensityResidual(*o.r, o.sigma, opt.wedge_deg, refine_partiality));
problem.AddResidualBlock(cost,
nullptr, // no loss function
&log_k[o.img_slot],
&b[o.img_slot],
&k[o.img_slot],
&mosaicity_rad[o.img_slot],
&log_Itrue[o.hkl_slot]);
&Itrue[o.hkl_slot]);
}
// Optional Kabsch-like regularization for k
@@ -313,28 +296,14 @@ ScaleMergeResult ScaleAndMergeReflectionsCeres(const std::vector<Reflection>& ob
for (int i = 0; i < nimg; ++i) {
auto* rcost = new ceres::AutoDiffCostFunction<ScaleRegularizationResidual, 1, 1>(
new ScaleRegularizationResidual(opt.scale_regularization_sigma));
problem.AddResidualBlock(rcost, nullptr, &log_k[i]);
problem.AddResidualBlock(rcost, nullptr, &k[i]);
}
}
// Fix gauge freedom
if (opt.fix_first_image_scale && nimg > 0) {
log_k[0] = 0.0;
problem.SetParameterBlockConstant(&log_k[0]);
}
if (!opt.refine_b_factor) {
for (int i = 0; i < nimg; ++i) {
b[i] = 0.0;
problem.SetParameterBlockConstant(&b[i]);
}
} else {
for (int i = 0; i < nimg; ++i) {
if (opt.b_min)
problem.SetParameterLowerBound(&b[i], 0, *opt.b_min);
if (opt.b_max)
problem.SetParameterUpperBound(&b[i], 0, *opt.b_max);
}
k[0] = 1.0;
problem.SetParameterBlockConstant(&k[0]);
}
// Mosaicity refinement + bounds
@@ -342,14 +311,22 @@ ScaleMergeResult ScaleAndMergeReflectionsCeres(const std::vector<Reflection>& ob
for (int i = 0; i < nimg; ++i)
problem.SetParameterBlockConstant(&mosaicity_rad[i]);
} else {
const std::optional<double> min_rad = opt.mosaicity_min_deg
? std::optional<double>(deg2rad(*opt.mosaicity_min_deg)) : std::nullopt;
const std::optional<double> max_rad = opt.mosaicity_max_deg
? std::optional<double>(deg2rad(*opt.mosaicity_max_deg)) : std::nullopt;
const double min_rad = deg2rad(opt.mosaicity_min_deg);
const double max_rad = deg2rad(opt.mosaicity_max_deg);
for (int i = 0; i < nimg; ++i) {
if (min_rad) problem.SetParameterLowerBound(&mosaicity_rad[i], 0, *min_rad);
if (max_rad) problem.SetParameterUpperBound(&mosaicity_rad[i], 0, *max_rad);
problem.SetParameterLowerBound(&mosaicity_rad[i], 0, min_rad);
problem.SetParameterUpperBound(&mosaicity_rad[i], 0, max_rad);
}
}
// Force k[i] > 0 for all images (scale factors must be positive)
{
constexpr double k_floor = 1e-8;
for (int i = 0; i < nimg; ++i) {
if (!(opt.fix_first_image_scale && i == 0)) {
problem.SetParameterLowerBound(&k[i], 0, k_floor);
}
}
}
@@ -364,10 +341,8 @@ ScaleMergeResult ScaleAndMergeReflectionsCeres(const std::vector<Reflection>& ob
ceres::Solve(options, &problem, &summary);
// --- Export per-image results ---
for (int i = 0; i < nimg; ++i) {
out.image_scale_k[i] = std::exp(log_k[i]);
out.image_b_factor[i] = opt.refine_b_factor ? b[i] : 0.0;
}
for (int i = 0; i < nimg; ++i)
out.image_scale_k[i] = k[i];
out.mosaicity_deg.resize(nimg);
for (int i = 0; i < nimg; ++i)
@@ -375,13 +350,11 @@ ScaleMergeResult ScaleAndMergeReflectionsCeres(const std::vector<Reflection>& ob
// --- Compute goodness-of-fit (reduced chi-squared) ---
const int n_obs = static_cast<int>(obs.size());
// Count free parameters: nhkl log_Itrue + per-image (log_k + b + mosaicity) minus fixed ones
// Count free parameters: nhkl Itrue + per-image (k + mosaicity) minus fixed ones
int n_params = nhkl;
for (int i = 0; i < nimg; ++i) {
if (!(opt.fix_first_image_scale && i == 0))
n_params += 1; // log_k
if (opt.refine_b_factor)
n_params += 1; // b
n_params += 1; // k
if (opt.refine_mosaicity)
n_params += 1; // mosaicity
}
@@ -393,22 +366,20 @@ ScaleMergeResult ScaleAndMergeReflectionsCeres(const std::vector<Reflection>& ob
const int i = o.img_slot;
const int h = o.hkl_slot;
const double ki = std::exp(log_k[i]);
const double atten = opt.refine_b_factor ? std::exp(-b[i] * o.s2) : 1.0;
const double Itrue = std::exp(log_Itrue[h]);
const double ki = k[i];
const double zeta = static_cast<double>(o.r->zeta);
const double mosa_i = mosaicity_rad[i];
const double c1 = std::sqrt(2.0) * (mosa_i / zeta);
const double c2 = o.r->zeta / (std::sqrt(2.0) * mosa_i);
const double delta_phi_rad = static_cast<double>(o.r->delta_phi) * M_PI / 180.0;
const double partiality = partiality_model
? (std::erf((delta_phi_rad + half_wedge_rad) * c1) - std::erf((delta_phi_rad - half_wedge_rad) * c1)) / 2.0
: 1.0;
const double partiality = refine_partiality
? (std::erf((delta_phi_rad + half_wedge_rad) * c2)
- std::erf((delta_phi_rad - half_wedge_rad) * c2)) / 2.0
: o.r->partiality;
const double lp = SafeInv(static_cast<double>(o.r->rlp), 1.0);
const double lp = SafeInv(o.r->rlp, 1.0);
const double Ipred = ki * atten * partiality * lp * Itrue;
const double Ipred = ki * partiality * lp * Itrue[h];
const double resid = (Ipred - static_cast<double>(o.r->I)) / o.sigma;
sum_r2 += resid * resid;
}
@@ -427,7 +398,7 @@ ScaleMergeResult ScaleAndMergeReflectionsCeres(const std::vector<Reflection>& ob
out.merged[h].h = slotToHKL[h].h;
out.merged[h].k = slotToHKL[h].k;
out.merged[h].l = slotToHKL[h].l;
out.merged[h].I = std::exp(log_Itrue[h]);
out.merged[h].I = Itrue[h];
out.merged[h].sigma = std::numeric_limits<double>::quiet_NaN();
out.merged[h].d = 0.0;
}
@@ -457,17 +428,16 @@ ScaleMergeResult ScaleAndMergeReflectionsCeres(const std::vector<Reflection>& ob
std::vector<std::pair<const double*, const double*>> covariance_blocks;
covariance_blocks.reserve(nhkl);
for (int h = 0; h < nhkl; ++h) {
covariance_blocks.emplace_back(&log_Itrue[h], &log_Itrue[h]);
covariance_blocks.emplace_back(&Itrue[h], &Itrue[h]);
}
if (covariance.Compute(covariance_blocks, &problem)) {
for (int h = 0; h < nhkl; ++h) {
double var_log_I = 0.0;
covariance.GetCovarianceBlock(&log_Itrue[h], &log_Itrue[h], &var_log_I);
double var_I = 0.0;
covariance.GetCovarianceBlock(&Itrue[h], &Itrue[h], &var_I);
// σ(I) = I * sqrt( var(log I) * GoF² )
const double Itrue = std::exp(log_Itrue[h]);
out.merged[h].sigma = Itrue * std::sqrt(var_log_I * gof2);
// σ(I) = I * sqrt( var(I) * GoF² )
out.merged[h].sigma = std::sqrt(var_I * gof2);
}
}

View File

@@ -11,8 +11,6 @@
#include "gemmi/symmetry.hpp"
struct ScaleMergeOptions {
bool refine_b_factor = true;
int max_num_iterations = 100;
double max_solver_time_s = 1.0;
@@ -21,9 +19,6 @@ struct ScaleMergeOptions {
bool fix_first_image_scale = true;
std::optional<double> b_min = 0.0;
std::optional<double> b_max = 200.0;
// Symmetry canonicalization of HKL prior to merging/scaling.
// If not set, the routine uses raw HKL as-is.
std::optional<gemmi::SpaceGroup> space_group;
@@ -41,8 +36,8 @@ struct ScaleMergeOptions {
bool refine_mosaicity = true;
double mosaicity_init_deg = 0.17; // ~0.003 rad
std::optional<double> mosaicity_min_deg = 1e-3;
std::optional<double> mosaicity_max_deg = 2.0;
double mosaicity_min_deg = 1e-3;
double mosaicity_max_deg = 2.0;
// --- Optional: regularize per-image scale k towards 1 (Kabsch-like) ---
bool regularize_scale_to_one = false;
@@ -61,7 +56,6 @@ struct MergedReflection {
struct ScaleMergeResult {
std::vector<MergedReflection> merged;
std::vector<double> image_scale_k;
std::vector<double> image_b_factor;
std::vector<int> image_ids;
// One mosaicity value per image (degrees).

View File

@@ -392,7 +392,6 @@ int main(int argc, char **argv) {
logger.Info("Running scaling (mosaicity refinement) ...");
ScaleMergeOptions scale_opts;
scale_opts.refine_b_factor = false; // B-factor refinement doesn't make sense for rotation
scale_opts.refine_mosaicity = true;
scale_opts.max_num_iterations = 500;
scale_opts.max_solver_time_s = 240.0; // generous cutoff for now