From 3c4d730f898a32001d19fdf5a4233db531a146d9 Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Wed, 14 Oct 2026 11:59:09 +0200 Subject: [PATCH] RESULT-CHANGING (owner decision): stills tilt post-refinement solved directly instead of by Ceres Each crystal's (psi_x, psi_y) tilt was refined by a Ceres problem rebuilt for every inner iteration: one auto-differentiated residual block and one Cauchy loss per reflection. With the scaling reference cleaned up it was 70% of the stills scale/merge. The same objective - Cauchy on the sigma-weighted intensity residuals plus the Gaussian prior on the tilt, the tilt bounded - is now minimised by a small Levenberg-Marquardt on the rho'-reweighted 2x2 normal equations, the tilt applied to all reflections through one rotation matrix and the Jacobian taken by central differences. It reaches the same answer to optimizer precision: on the new synthetic test (StillsPartialityRefine_RecoversTheTilt) partiality error 0.0559 -> 0.03019 with Ceres and 0.03016 with this, mean tilt 0.02065 vs 0.02067 deg; on two 25k-frame XFEL runs (--mode scale) CC1/2, I/sigma, R_pim and the CC against a deposited lysozyme model agree to 0.0005, mean |dpsi| 0.023 deg on both. Scale+merge 6.9 -> 3.0 s (--mode scale, loaded box). Co-Authored-By: Claude Opus 5.5 (1M context) --- .../scale_merge/StillsPartialityRefine.cpp | 176 +++++++++++------- tests/MergeScaleTest.cpp | 96 ++++++++++ 2 files changed, 205 insertions(+), 67 deletions(-) diff --git a/image_analysis/scale_merge/StillsPartialityRefine.cpp b/image_analysis/scale_merge/StillsPartialityRefine.cpp index 60ee3f7dd..21785a263 100644 --- a/image_analysis/scale_merge/StillsPartialityRefine.cpp +++ b/image_analysis/scale_merge/StillsPartialityRefine.cpp @@ -3,13 +3,14 @@ #include "StillsPartialityRefine.h" +#include +#include #include #include #include #include #include -#include #include #include "Merge.h" @@ -101,44 +102,114 @@ namespace { 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; + // The rotation by the angle-axis vector (psi_x, psi_y, 0) as a matrix (Rodrigues), so that one tilt + // is applied to all the reflections of a crystal by a matrix product each. + std::array TiltMatrix(double psi_x, double psi_y) { + const double theta = std::sqrt(psi_x * psi_x + psi_y * psi_y); + if (theta < 1e-15) + return {1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0}; + const double kx = psi_x / theta, ky = psi_y / theta; + const double c = std::cos(theta), s = std::sin(theta), t = 1.0 - c; + return {c + t * kx * kx, t * kx * ky, s * ky, + t * kx * ky, c + t * ky * ky, -s * kx, + -s * ky, s * kx, c}; + } + + // ComputeP with the tilt given as its rotation matrix. + struct TiltModel { double s0x, s0y, s0z, inv_lambda; double gamma0, gamma_e, bw; - double G, inv_corr, Iref, Iobs, weight; - template - 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_ang = T(gamma_e) * dstar; - const T sbw = T(bw) * ceres::abs(qr[2]); - const T sig2 = T(gamma0) * T(gamma0) + sig_ang * sig_ang + sbw * sbw; - const T p = ceres::exp(T(-0.5) * de * de / sig2); - residual[0] = T(weight) * (T(G) * p * T(inv_corr) * T(Iref) - T(Iobs)); - return true; + double P(const std::array &R, const FitObs &o) const { + const double qx = R[0] * o.qx + R[1] * o.qy + R[2] * o.qz; + const double qy = R[3] * o.qx + R[4] * o.qy + R[5] * o.qz; + const double qz = R[6] * o.qx + R[7] * o.qy + R[8] * o.qz; + const double de = VecLen(qx + s0x, qy + s0y, qz + s0z) - inv_lambda; + const double sig_ang = gamma_e * VecLen(qx, qy, qz); + const double sbw = bw * std::fabs(qz); + const double sig2 = gamma0 * gamma0 + sig_ang * sig_ang + sbw * sbw; + if (!(sig2 > 0.0)) + return 1.0; + return std::exp(-0.5 * de * de / sig2); } }; - // 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 - 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; + // Refine the orientation tilt (psi_x, psi_y) holding the scale G fixed: minimise + // sum_i rho(r_i^2) / 2 + |psi / sigma_prior|^2 / 2, r_i = weight_i (G p_i(psi) Iref_i / corr_i - Iobs_i), + // with the Cauchy rho(s) = k^2 log(1 + s / k^2) - the intensity-space objective of ScaleOnTheFly, and + // the Gaussian prior a proper MAP term since r_i is already in sigma units. Levenberg-Marquardt on the + // normal equations reweighted by rho', the Jacobian by central differences, every trial tilt held inside + // the bound. Returns false when the cost cannot be evaluated, leaving psi as it came in. + bool RefineTilt(const std::vector &obs, const TiltModel &model, double G, double robust_k, + double inv_sigma_prior, double max_tilt, double psi[2]) { + constexpr int MAX_ITERATIONS = 25; + constexpr double STEP = 1e-7; // rad, for the central differences + const double k2 = robust_k * robust_k; + const double prior2 = inv_sigma_prior * inv_sigma_prior; + + auto cost = [&](const double p[2]) { + const auto R = TiltMatrix(p[0], p[1]); + double c = 0.5 * prior2 * (p[0] * p[0] + p[1] * p[1]); + for (const auto &o: obs) { + const double r = o.weight * (G * model.P(R, o) * o.inv_corr * o.Iref - o.Iobs); + c += 0.5 * k2 * std::log1p(r * r / k2); + } + return c; + }; + + double current = cost(psi); + if (!std::isfinite(current)) + return false; + double lambda = 1e-4; + for (int iter = 0; iter < MAX_ITERATIONS; ++iter) { + const auto R = TiltMatrix(psi[0], psi[1]); + const auto Rxp = TiltMatrix(psi[0] + STEP, psi[1]), Rxm = TiltMatrix(psi[0] - STEP, psi[1]); + const auto Ryp = TiltMatrix(psi[0], psi[1] + STEP), Rym = TiltMatrix(psi[0], psi[1] - STEP); + double A00 = prior2, A01 = 0.0, A11 = prior2; + double b0 = prior2 * psi[0], b1 = prior2 * psi[1]; + for (const auto &o: obs) { + const double a = o.weight * G * o.inv_corr * o.Iref; + const double r = a * model.P(R, o) - o.weight * o.Iobs; + const double jx = a * (model.P(Rxp, o) - model.P(Rxm, o)) / (2.0 * STEP); + const double jy = a * (model.P(Ryp, o) - model.P(Rym, o)) / (2.0 * STEP); + const double w = 1.0 / (1.0 + r * r / k2); // rho' + A00 += w * jx * jx; + A01 += w * jx * jy; + A11 += w * jy * jy; + b0 += w * jx * r; + b1 += w * jy * r; + } + + bool accepted = false; + double trial[2] = {psi[0], psi[1]}; + double trial_cost = current; + while (lambda < 1e10) { + const double M00 = A00 * (1.0 + lambda), M11 = A11 * (1.0 + lambda); + const double det = M00 * M11 - A01 * A01; + if (!(det > 0.0)) + break; + trial[0] = std::clamp(psi[0] - (M11 * b0 - A01 * b1) / det, -max_tilt, max_tilt); + trial[1] = std::clamp(psi[1] - (M00 * b1 - A01 * b0) / det, -max_tilt, max_tilt); + trial_cost = cost(trial); + if (std::isfinite(trial_cost) && trial_cost < current) { + accepted = true; + lambda = std::max(lambda / 10.0, 1e-12); + break; + } + lambda *= 10.0; + } + if (!accepted) + break; + const double step = std::hypot(trial[0] - psi[0], trial[1] - psi[1]); + const double decrease = current - trial_cost; + psi[0] = trial[0]; + psi[1] = trial[1]; + current = trial_cost; + if (decrease <= 1e-6 * current || step <= 1e-10) + break; } - }; + return true; + } } StillsPartialityRefine::StillsPartialityRefine(const DiffractionExperiment &x) @@ -238,6 +309,8 @@ double StillsPartialityRefine::RefineOne(IntegrationOutcome &outcome, double psi[2] = {0.0, 0.0}; double G = 1.0; + const TiltModel model{.s0x = S0.x, .s0y = S0.y, .s0z = S0.z, .inv_lambda = inv_lambda, + .gamma0 = gamma0, .gamma_e = gamma_e, .bw = bw}; const bool refine_tilt = obs.size() >= settings_.min_reflections; const int inner = refine_tilt ? settings_.inner_iterations : 1; @@ -258,42 +331,11 @@ double StillsPartialityRefine::RefineOne(IntegrationOutcome &outcome, if (!refine_tilt) break; - // (2) Refine the tilt holding G fixed. - ceres::Problem problem; - for (const auto &o: obs) { - auto *cost = new ceres::AutoDiffCostFunction(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, .inv_corr = o.inv_corr, .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( - 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; - const double psi_before[2] = {psi[0], psi[1]}; - ceres::Solve(options, &problem, &summary); - // A solve that failed leaves whatever the minimiser last wrote in psi, and that tilt would go - // straight onto every reflection's partiality below. Keep the tilt this crystal came in with - // and stop refining it instead - G alone is still a usable model. - if (!summary.IsSolutionUsable()) { - psi[0] = psi_before[0]; - psi[1] = psi_before[1]; + // (2) Refine the tilt holding G fixed. A cost that cannot be evaluated keeps the tilt this crystal + // came in with and stops refining it - G alone is still a usable model. + const double inv_sigma_prior = settings_.prior_sigma_deg > 0.0 ? kRadToDeg / settings_.prior_sigma_deg : 0.0; + if (!RefineTilt(obs, model, G, settings_.robust_k, inv_sigma_prior, settings_.max_tilt_rad, psi)) break; - } } // Keep what the crystal came in with. This refinement is on by default, so a crystal the tilt model diff --git a/tests/MergeScaleTest.cpp b/tests/MergeScaleTest.cpp index ed548a294..d862fd213 100644 --- a/tests/MergeScaleTest.cpp +++ b/tests/MergeScaleTest.cpp @@ -10,6 +10,8 @@ #include "../image_analysis/scale_merge/HKLKey.h" #include "../image_analysis/scale_merge/Merge.h" #include "../image_analysis/scale_merge/ResolutionCutoff.h" +#include "../image_analysis/scale_merge/ScaleOnTheFly.h" +#include "../image_analysis/scale_merge/StillsPartialityRefine.h" #include "gemmi/reciproc.hpp" TEST_CASE("HKLKey_NoSG_noMergeFriedel") { @@ -650,3 +652,97 @@ TEST_CASE("MergeOnTheFly_AddImagesIsTheSerialMerge") { CHECK(serial_stats.overall.cc_half == parallel_stats.overall.cc_half); CHECK(serial_stats.overall.r_meas == parallel_stats.overall.r_meas); } + +// Stills partiality post-refinement on noiseless synthetic crystals: each crystal's true orientation is +// tilted from the stored one by up to 0.05 deg, its intensities are the true partials of a fixed set of +// full intensities, and its stored partialities those of the stored orientation. The refinement has to +// move the partialities towards the true ones. (The Ceres solve this replaced gave the same numbers: error +// 0.0559 -> 0.0302, mean tilt 0.0206 deg.) +TEST_CASE("StillsPartialityRefine_RecoversTheTilt") { + constexpr float CELL = 50.0f, WAVELENGTH = 1.0f, D_MIN = 2.0f; + constexpr double GAMMA0 = 0.001, ACCEPT = 1.7 * GAMMA0; // RLP width; |dist_ewald| accepted at prediction + DiffractionExperiment x; + x.SetSpaceGroup(gemmi::get_spacegroup_by_name("P 1")); + DiffractionGeometry geom; + geom.Wavelength_A(WAVELENGTH); + const Coord S0 = geom.GetScatteringVector(); + + std::mt19937 rng(11); + std::normal_distribution gauss(0.0, 1.0); + std::uniform_real_distribution unit(0.0, 1.0); + auto full_I = [](int h, int k, int l) { // a fixed, Friedel-symmetric pseudo-random full intensity + const uint64_t key = HKLKeyGenerator(true, gemmi::get_spacegroup_by_name("P 1"))(h, k, l).pack(); + return 10.0 + 990.0 * static_cast((key * 0x9e3779b97f4a7c15ULL) >> 40) / static_cast(1ULL << 24); + }; + auto rotate = [](const double R[9], const Coord &v) { + return Coord(static_cast(R[0] * v.x + R[1] * v.y + R[2] * v.z), + static_cast(R[3] * v.x + R[4] * v.y + R[5] * v.z), + static_cast(R[6] * v.x + R[7] * v.y + R[8] * v.z)); + }; + auto rotation = [](double ax, double ay, double az, double angle, double R[9]) { // Rodrigues + const double n = std::sqrt(ax * ax + ay * ay + az * az); + ax /= n; ay /= n; az /= n; + const double c = std::cos(angle), s = std::sin(angle), t = 1.0 - c; + const double M[9] = {c + t * ax * ax, t * ax * ay - s * az, t * ax * az + s * ay, + t * ax * ay + s * az, c + t * ay * ay, t * ay * az - s * ax, + t * ax * az - s * ay, t * ay * az + s * ax, c + t * az * az}; + std::copy(M, M + 9, R); + }; + auto partiality = [&](const Coord &q) { + const double de = (q + S0).Length() - 1.0 / WAVELENGTH; + return std::exp(-0.5 * de * de / (GAMMA0 * GAMMA0)); + }; + + std::vector outcomes(150); + std::vector> true_p(outcomes.size()); + const int hmax = static_cast(CELL / D_MIN); + for (size_t c = 0; c < outcomes.size(); ++c) { + double R[9], T[9]; + rotation(gauss(rng), gauss(rng), gauss(rng), 2.0 * M_PI * unit(rng), R); + const double tilt = (0.02 + 0.03 * unit(rng)) * M_PI / 180.0, phi = 2.0 * M_PI * unit(rng); + rotation(std::cos(phi), std::sin(phi), 0.0, tilt, T); // about an axis normal to the beam + auto &o = outcomes[c]; + o.geom = geom; + o.latt = CrystalLattice(rotate(R, Coord(CELL, 0, 0)), rotate(R, Coord(0, CELL, 0)), + rotate(R, Coord(0, 0, CELL))); + const double G = 0.5 + unit(rng); + for (int h = -hmax; h <= hmax; ++h) + for (int k = -hmax; k <= hmax; ++k) + for (int l = -hmax; l <= hmax; ++l) { + const Coord q = o.latt.Astar() * h + o.latt.Bstar() * k + o.latt.Cstar() * l; + const float dstar = q.Length(); + if (dstar == 0.0f || dstar > 1.0f / D_MIN) + continue; + if (std::fabs((q + S0).Length() - 1.0 / WAVELENGTH) > ACCEPT) + continue; + const double p_true = partiality(rotate(T, q)); + Reflection r{}; + r.h = h; r.k = k; r.l = l; + r.d = 1.0f / dstar; + r.I = static_cast(G * p_true * full_I(h, k, l)); + r.sigma = static_cast(1.0 + 0.01 * r.I); + r.prescaling_corr = 1.0f; + r.partiality = static_cast(partiality(q)); + r.image_scale_corr = 1.0f / r.partiality; + o.reflections.push_back(r); + true_p[c].push_back(p_true); + } + REQUIRE(o.reflections.size() >= 40); + } + + auto mean_abs_error = [&]() { + double sum = 0.0; + size_t n = 0; + for (size_t c = 0; c < outcomes.size(); ++c) + for (size_t i = 0; i < outcomes[c].reflections.size(); ++i, ++n) + sum += std::fabs(outcomes[c].reflections[i].partiality - true_p[c][i]); + return sum / static_cast(n); + }; + const double before = mean_abs_error(); + ScaleOnTheFly(x, MergeAll(x, outcomes, 4, std::nullopt)).Scale(outcomes, 4); + const double mean_tilt = StillsPartialityRefine(x).Run(outcomes, 4, std::nullopt); + const double after = mean_abs_error(); + CHECK(mean_tilt > 0.015); + CHECK(mean_tilt < 0.06); + CHECK(after < 0.6 * before); // measured 0.54: the prior holds the larger tilts back +}