From d271a55a74ac0a8e41f08a386f70b2f5ece082e0 Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Sat, 10 Oct 2026 10:47:27 +0200 Subject: [PATCH] ModelScaling: FitModelScale does the same fits with less work All exact to the bit (bench against the previous FitModelScale on P1, P2_1, P2_12_12_1, P4_32_12, P6_122, P2_13, I23 and the dependent n<=5 path; the ModelScaling cases; validation outputs identical on real data): - the start of a grid point (fit_isotropic_b_approximately) is computed from the same |Fcalc| the fit takes, once instead of twice, and without copying the Scaling's points per chunk; - gemmi's Levenberg-Marquardt is followed in a copy (LevMarFit) whose compute_lm_matrices is templated on the parameter count (2..7), so alpha/beta stay in registers; a zero derivative row adds +0 instead of being skipped (alpha/beta are never -0), and the full square is summed (the lower half is gemmi's, the upper is its mirror) - branch-free and vectorised; - grid points are scheduled one per task (ParallelFor) instead of fixed chunks; - a fine-pass pair bit-identical to a coarse-pass pair reuses that fit. Bench (single thread, 48k reflections, P2_1): 1.58 -> 1.33 s; ~15-25% on the other groups. Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01SVmAWnzCmRKAXVUCdc4iNi --- .../structure_refinement/ModelScaling.cpp | 237 +++++++++++++++--- 1 file changed, 202 insertions(+), 35 deletions(-) diff --git a/image_analysis/structure_refinement/ModelScaling.cpp b/image_analysis/structure_refinement/ModelScaling.cpp index 2c6cf5cd7..8ce43a4e2 100644 --- a/image_analysis/structure_refinement/ModelScaling.cpp +++ b/image_analysis/structure_refinement/ModelScaling.cpp @@ -2,22 +2,25 @@ // SPDX-License-Identifier: GPL-3.0-only // FixedSolventFit follows the scaling target of GEMMI's scaling.hpp -// (https://github.com/project-gemmi/gemmi/blob/master/include/gemmi/scaling.hpp) +// (https://github.com/project-gemmi/gemmi/blob/master/include/gemmi/scaling.hpp), and +// ComputeLMMatrices / LevMarFit the Levenberg-Marquardt of GEMMI's levmar.hpp +// (https://github.com/project-gemmi/gemmi/blob/master/include/gemmi/levmar.hpp) // (c) Global Phasing Ltd., Mozilla Public License Version 2.0 #include "ModelScaling.h" #include #include +#include #include #include "../../common/ParallelFor.h" namespace { -// gemmi::Scaling as a grid point fits it: the same parameters, model values, derivatives -// and R. The solvent pair is fixed there, so |Fcalc + k_sol exp(-b_sol s^2) Fmask| of a reflection -// is the same at every evaluation of the solver; it is taken once per grid point here rather than +// gemmi::Scaling as a grid point fits it: the same starting point, parameters, model values, +// derivatives and R. The solvent pair is fixed there, so |Fcalc + k_sol exp(-b_sol s^2) Fmask| of a +// reflection is the same at every evaluation of the solver; it is taken once per grid point here rather than // at every evaluation, where it was most of the cost of the fit. Every expression is the one // scaling.hpp evaluates, in the same types, so every number is the same to the bit. struct FixedSolventFit { @@ -35,13 +38,41 @@ struct FixedSolventFit { explicit FixedSolventFit(const gemmi::Scaling &s) : constraint_matrix(s.constraint_matrix) {} - // The reflections at the solvent pair of `s`, and its scale as the start of the fit. - void Start(const gemmi::Scaling &s) { - k_overall = s.k_overall; - b_star = s.b_star; + // The reflections of `s` at the solvent pair (k_sol, b_sol), and the start of the fit there: what + // Scaling::fit_isotropic_b_approximately() sets, from the same |Fcalc|, taken once for both. Where + // that finds five or fewer reflections to fit on it sets nothing, and the fit starts from + // (k_start, b_start) - the scale gemmi's would have been left at. + void Start(const gemmi::Scaling &s, double k_sol, double b_sol, + double k_start, const gemmi::SMat33 &b_start) { + k_overall = k_start; + b_star = b_start; points.resize(s.points.size()); - for (size_t i = 0; i < points.size(); ++i) - points[i] = {s.points[i].hkl, s.points[i].fobs, std::abs(s.get_fcalc(s.points[i]))}; + double sx = 0, sy = 0, sxx = 0, sxy = 0; + int n = 0; + for (size_t i = 0; i < points.size(); ++i) { + const auto &p = s.points[i]; + // Scaling::get_fcalc() at (k_sol, b_sol) + const std::complex fcalc = + s.use_solvent ? p.fcmol + (float) (k_sol * std::exp(-b_sol * p.stol2)) * p.fmask : p.fcmol; + points[i] = {p.hkl, p.fobs, std::abs(fcalc)}; + if (p.fobs < 1 || p.fobs < p.sigma) // skip weak reflections + continue; + double fcalc_abs = points[i].fcalc_abs; + double x = p.stol2; + double y = std::log(static_cast(p.fobs / fcalc_abs)); + sx += x; + sy += y; + sxx += x * x; + sxy += x * y; + n += 1; + } + if (n <= 5) + return; + double slope = (n * sxy - sx * sy) / (n * sxx - sx * sx); + double intercept = (sy - slope * sx) / n; + double b_iso = -slope; + k_overall = std::exp(intercept); + b_star = gemmi::SMat33{b_iso, b_iso, b_iso, 0, 0, 0}.transformed_by(s.cell.frac.mat); } // Scaling::get_parameters(), set_parameters(), get_overall_scale_factor(), compute_value() and @@ -74,7 +105,9 @@ struct FixedSolventFit { double compute_value(const Point &p) const { return p.fcalc_abs * (float) get_overall_scale_factor(p.hkl); } - double compute_value_and_derivatives(const Point &p, std::vector &dy_da) const { + // For NA parameters: the B components are NA - 1 constraint rows. + template + double compute_value_and_derivatives(const Point &p, double *dy_da) const { gemmi::Vec3 h(p.hkl); double kaniso = std::exp(-0.25 * b_star.r_u_r(h)); double fcalc_abs = p.fcalc_abs; @@ -90,7 +123,7 @@ struct FixedSolventFit { -0.5 * y * (h.x * h.z), -0.5 * y * (h.y * h.z), }; - for (size_t j = 0; j < constraint_matrix.size(); ++j) + for (int j = 0; j < NA - 1; ++j) dy_da[n + j] = gemmi::vec6_dot(constraint_matrix[j], du); return y; } @@ -107,6 +140,127 @@ double RFactor(const FixedSolventFit &fit) { return den > 0 ? num / den : 1.0; } +// gemmi's compute_lm_matrices() (levmar.hpp) on a FixedSolventFit of NA parameters. With the count +// known to the compiler the matrices are held in registers rather than written to memory at every +// reflection; the operations are gemmi's, in gemmi's order, so the matrices are gemmi's to the bit. +// A reflection's weight is 1 (FixedSolventFit::Point::get_weight), which gemmi multiplies by. +// gemmi skips the row of a derivative that is exactly 0; here the row gets +0 instead, which keeps the +// loop free of branches and changes nothing: an element of alpha or beta is never -0, the one value +// adding +0 would change, since it starts at +0 and a sum is only -0 when both its terms are. And +// where gemmi fills the lower half of alpha, the whole square is summed here, a rectangle the compiler +// can vectorise; the lower half is gemmi's, and the upper is replaced by it at the end, as in gemmi. +template +double ComputeLMMatrices(const FixedSolventFit &fit, double *alpha_out, double *beta_out) { + long double wssr = 0; // long double here notably increases the accuracy + double alpha[NA * NA] = {}, beta[NA] = {}; + double dy_da[NA]; + for (const auto &p : fit.points) { + double y = fit.compute_value_and_derivatives(p, dy_da); + double dy_sig = p.get_y() - y; + for (int j = 0; j != NA; ++j) { + const bool used = dy_da[j] != 0; + for (int k = 0; k < NA; ++k) { + const double a = dy_da[j] * dy_da[k]; + alpha[NA * j + k] += used ? a : 0.0; + } + const double b = dy_sig * dy_da[j]; + beta[j] += used ? b : 0.0; + } + wssr += gemmi::sq(dy_sig); + } + // The upper half of alpha is the lower half's mirror, as in gemmi. + for (int j = 1; j < NA; j++) + for (int k = 0; k < j; k++) + alpha[NA * k + j] = alpha[NA * j + k]; + std::copy(alpha, alpha + NA * NA, alpha_out); + std::copy(beta, beta + NA, beta_out); + return (double) wssr; +} + +// gemmi's LevMar::fit() (levmar.hpp) with its default settings, on ComputeLMMatrices() - the same +// iterations to the same answer, bit for bit. +template +void LevMarFit(FixedSolventFit &target) { + const int eval_limit = 100; + const double lambda_limit = 1e+15; + const double stop_rel_change = 1e-5; + const double lambda_up_factor = 10; + const double lambda_down_factor = 0.1; + const double lambda_start = 0.001; + + std::vector initial_a = target.get_parameters(); + std::vector best_a = initial_a; + double lambda = lambda_start; + double alpha[NA * NA], beta[NA], temp_alpha[NA * NA]; + std::vector temp_beta(NA); + + const double initial_wssr = ComputeLMMatrices(target, alpha, beta); + double wssr = initial_wssr; + + int small_change_counter = 0; + int eval_count = 1; // number of function evaluations so far + for (;;) { + if (eval_limit > 0 && eval_count >= eval_limit) + break; + + // prepare next parameters -> temp_beta + std::copy(alpha, alpha + NA * NA, temp_alpha); + // Using '*=' not '+=' below applies the dampling factor as: + // J^T J + lambda * diag(J^T J); not ... + lambda * I. + for (int j = 0; j < NA; j++) + temp_alpha[NA * j + j] *= (1.0 + lambda); + std::copy(beta, beta + NA, temp_beta.begin()); + + // Matrix solution (Ax=b) temp_alpha * da == temp_beta + gemmi::jordan_solve(temp_alpha, temp_beta.data(), NA); + + for (int i = 0; i < NA; i++) + // put new a[] into temp_beta[] + temp_beta[i] += best_a[i]; + + target.set_parameters(temp_beta); + double new_wssr = gemmi::compute_wssr(target); + ++eval_count; + if (new_wssr < wssr) { + double rel_change = (wssr - new_wssr) / wssr; + wssr = new_wssr; + best_a = temp_beta; + + if (wssr == 0) + break; + // termination criterion: negligible change of wssr + if (rel_change < stop_rel_change) { + if (++small_change_counter >= 2) + break; + } else { + small_change_counter = 0; + } + ComputeLMMatrices(target, alpha, beta); + ++eval_count; + lambda *= lambda_down_factor; + } else { // worse fitting + if (lambda > lambda_limit) // termination criterion: large lambda + break; + lambda *= lambda_up_factor; + } + } + + target.set_parameters(wssr < initial_wssr ? best_a : initial_a); +} + +// LevMarFit() at the fit's parameter count: k_overall and one per symmetry-allowed component of B, +// which is one (cubic) to six (triclinic). +void LevMarFit(FixedSolventFit &fit) { + switch (fit.constraint_matrix.size()) { + case 1: LevMarFit<2>(fit); break; + case 2: LevMarFit<3>(fit); break; + case 3: LevMarFit<4>(fit); break; + case 4: LevMarFit<5>(fit); break; + case 5: LevMarFit<6>(fit); break; + default: LevMarFit<7>(fit); break; + } +} + } // namespace // Following the phenix bulk-solvent and scaling procedure: k_sol and b_sol by a grid search, with @@ -125,11 +279,13 @@ double RFactor(const FixedSolventFit &fit) { // // Every grid point is its own fit: fit_isotropic_b_approximately() sets k_overall and b_star from the // data and the point's solvent pair alone, so a point does not depend on the one fitted before it, and -// the points run in parallel, each chunk on its own copy of `scaling`. The winner is then read off in -// grid order with the serial rule (lowest finite R, the first on a tie), so the answer is the serial -// loop's bit for bit. The one exception is fit_isotropic_b_approximately() finding five or fewer -// reflections to fit on - it then returns without setting anything and a point WOULD start from where -// the previous one ended - so there the grid is walked in order, on `scaling` itself, as it always was. +// the points run in parallel, one point per task, each on a FixedSolventFit of its own. The winner is +// then read off in grid order with the serial rule (lowest finite R, the first on a tie), so the answer +// is the serial loop's bit for bit. For the same reason a pair the refinement pass shares with the +// coarse one - the same two numbers, to the bit - is the same fit, and is taken from there rather than +// made again. The one exception is fit_isotropic_b_approximately() finding five or fewer reflections to +// fit on - it then returns without setting anything and a point WOULD start from where the previous one +// ended - so there the grid is walked in order, each point from where the last ended, as it always was. ModelScaleReport FitModelScale(gemmi::Scaling &scaling, ModelScaleBox box, size_t nthreads) { ModelScaleReport report; report.n_points = static_cast(scaling.points.size()); @@ -151,31 +307,41 @@ ModelScaleReport FitModelScale(gemmi::Scaling &scaling, ModelScaleBox box gemmi::SMat33 best_b_star{0, 0, 0, 0, 0, 0}; struct PointFit { double k_sol, b_sol, r = NAN, k_overall = 1.0; gemmi::SMat33 b_star{0, 0, 0, 0, 0, 0}; }; - auto fit_point = [](gemmi::Scaling &s, FixedSolventFit &fit, PointFit &pf) { - s.k_sol = pf.k_sol; - s.b_sol = pf.b_sol; - s.fit_isotropic_b_approximately(); // a fresh starting point for this solvent pair - fit.Start(s); - gemmi::LevMar levmar; - levmar.fit(fit); // k_overall + anisotropic B only - s.k_overall = fit.k_overall; - s.b_star = fit.b_star; + // The scale the walk stands at, which a point starts from only where the data give it no start of its own. + double k_last = scaling.k_overall; + gemmi::SMat33 b_last = scaling.b_star; + auto fit_point = [&scaling](FixedSolventFit &fit, PointFit &pf, double k_start, + const gemmi::SMat33 &b_start) { + fit.Start(scaling, pf.k_sol, pf.b_sol, k_start, b_start); + LevMarFit(fit); // k_overall + anisotropic B only pf.r = RFactor(fit); pf.k_overall = fit.k_overall; pf.b_star = fit.b_star; }; + std::vector fitted; // every point fitted so far, in grid order auto try_points = [&](std::vector &pts) { - if (independent) - ParallelChunks(static_cast(pts.size()), nthreads, [&](int lo, int hi) { - gemmi::Scaling local = scaling; - FixedSolventFit fit(local); - for (int i = lo; i < hi; ++i) - fit_point(local, fit, pts[i]); + if (independent) { + std::vector todo; + for (int i = 0; i < static_cast(pts.size()); ++i) { + const auto same = std::find_if(fitted.begin(), fitted.end(), [&](const PointFit &f) { + return f.k_sol == pts[i].k_sol && f.b_sol == pts[i].b_sol; + }); + if (same != fitted.end()) + pts[i] = *same; + else + todo.push_back(i); + } + ParallelFor(static_cast(todo.size()), nthreads, [&](int j) { + FixedSolventFit fit(scaling); + fit_point(fit, pts[todo[j]], k_last, b_last); }); - else { + } else { FixedSolventFit fit(scaling); - for (auto &pf : pts) - fit_point(scaling, fit, pf); + for (auto &pf : pts) { + fit_point(fit, pf, k_last, b_last); + k_last = pf.k_overall; + b_last = pf.b_star; + } } for (const auto &pf : pts) { ++report.n_grid; @@ -188,6 +354,7 @@ ModelScaleReport FitModelScale(gemmi::Scaling &scaling, ModelScaleBox box best_b_star = pf.b_star; } } + fitted.insert(fitted.end(), pts.begin(), pts.end()); }; // Coarse pass over the whole box, then one refinement pass around the winner.