mx/Jungfraujoch
3
2
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 81 Wiki Activity

ModelScaling: FitModelScale does the same fits with less work

Browse Source 
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) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01SVmAWnzCmRKAXVUCdc4iNi
This commit is contained in:
leonarski_f
2026-10-10 10:47:27 +02:00
co-authored by Claude Opus 5.5
parent 03cd78322a
commit d271a55a74
1 changed files with 202 additions and 35 deletions
Download Patch File Download Diff File Expand all files Collapse all files
image_analysis/structure_refinement/ModelScaling.cpp
+202 -35
View File
@@ -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 <algorithm>
#include <cmath>
#include <complex>
#include <vector>
#include "../../common/ParallelFor.h"
namespace {
// gemmi::Scaling<float> 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<float> 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<float> &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<float> &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<float> &s, double k_sol, double b_sol,
double k_start, const gemmi::SMat33<double> &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<float> 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<float>(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<double>{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<double> &dy_da) const {
// For NA parameters: the B components are NA - 1 constraint rows.
template <int NA>
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 <int NA>
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<NA>(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<NA>() - the same
// iterations to the same answer, bit for bit.
template <int NA>
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<double> initial_a = target.get_parameters();
std::vector<double> best_a = initial_a;
double lambda = lambda_start;
double alpha[NA * NA], beta[NA], temp_alpha[NA * NA];
std::vector<double> temp_beta(NA);
const double initial_wssr = ComputeLMMatrices<NA>(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<NA>(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<float> &scaling, ModelScaleBox box, size_t nthreads) {
ModelScaleReport report;
report.n_points = static_cast<int>(scaling.points.size());
@@ -151,31 +307,41 @@ ModelScaleReport FitModelScale(gemmi::Scaling<float> &scaling, ModelScaleBox box
gemmi::SMat33<double> best_b_star{0, 0, 0, 0, 0, 0};
struct PointFit { double k_sol, b_sol, r = NAN, k_overall = 1.0; gemmi::SMat33<double> b_star{0, 0, 0, 0, 0, 0}; };
auto fit_point = [](gemmi::Scaling<float> &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<double> b_last = scaling.b_star;
auto fit_point = [&scaling](FixedSolventFit &fit, PointFit &pf, double k_start,
const gemmi::SMat33<double> &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<PointFit> fitted; // every point fitted so far, in grid order
auto try_points = [&](std::vector<PointFit> &pts) {
if (independent)
ParallelChunks(static_cast<int>(pts.size()), nthreads, [&](int lo, int hi) {
gemmi::Scaling<float> local = scaling;
FixedSolventFit fit(local);
for (int i = lo; i < hi; ++i)
fit_point(local, fit, pts[i]);
if (independent) {
std::vector<int> todo;
for (int i = 0; i < static_cast<int>(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<int>(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<float> &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.