leonarski_fandClaude Opus 5.5
7a6b194e09
LevMarFit takes compute_wssr() in a copy (ComputeWSSR) that keeps each reflection's exp(-0.25 * b_star.r_u_r(h)); when the step is accepted, ComputeLMMatrices is evaluated at exactly those parameters and takes the same values from there instead of calling exp again. Same expressions, same parameters, same numbers. Bench against the base FitModelScale (single thread, 48k reflections, P2_1): 1.65 -> 1.10 s for this and the previous commit together; P1, P2_12_12_1, P6_122, I23 and the dependent n<=5 path identical to the bit and 30-40% faster. Validation outputs (p.hkl md5, p.mtz / p_maps.mtz data, maps, p_model.pdb, log) identical to the base: CPU build --mode scale on 5reo, myob_x10sa, 8sa8, 8xtg, 5ebi at -N 1, default and -N 8; full CPU runs of 9ea5, 6iu5, 7dkp; the GPU build with no visible device on the five --mode scale sets. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01SVmAWnzCmRKAXVUCdc4iNi
402 lines
18 KiB
C++
402 lines
18 KiB
C++
// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
|
|
// 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), 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 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 {
|
|
struct Point {
|
|
gemmi::Miller hkl;
|
|
float fobs;
|
|
float fcalc_abs; // std::abs(Scaling::get_fcalc(p))
|
|
double get_y() const { return fobs; }
|
|
double get_weight() const { return 1.0; }
|
|
};
|
|
const std::vector<gemmi::Vec6> &constraint_matrix;
|
|
double k_overall = 1.0;
|
|
gemmi::SMat33<double> b_star{0, 0, 0, 0, 0, 0};
|
|
std::vector<Point> points;
|
|
|
|
explicit FixedSolventFit(const gemmi::Scaling<float> &s) : constraint_matrix(s.constraint_matrix) {}
|
|
|
|
// 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());
|
|
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
|
|
// compute_value_and_derivatives(), with k_sol and b_sol fixed.
|
|
std::vector<double> get_parameters() const {
|
|
std::vector<double> ret;
|
|
ret.push_back(k_overall);
|
|
for (const gemmi::Vec6 &v : constraint_matrix)
|
|
ret.push_back(gemmi::vec6_dot(v, b_star));
|
|
return ret;
|
|
}
|
|
void set_parameters(const double *p) {
|
|
k_overall = p[0];
|
|
int n = 0;
|
|
b_star = {0, 0, 0, 0, 0, 0};
|
|
for (const gemmi::Vec6 &row : constraint_matrix) {
|
|
double d = p[++n];
|
|
b_star.u11 += row[0] * d;
|
|
b_star.u22 += row[1] * d;
|
|
b_star.u33 += row[2] * d;
|
|
b_star.u12 += row[3] * d;
|
|
b_star.u13 += row[4] * d;
|
|
b_star.u23 += row[5] * d;
|
|
}
|
|
}
|
|
void set_parameters(const std::vector<double> &p) { set_parameters(p.data()); }
|
|
double get_overall_scale_factor(const gemmi::Miller &hkl) const {
|
|
return k_overall * std::exp(-0.25 * b_star.r_u_r(hkl));
|
|
}
|
|
double compute_value(const Point &p) const {
|
|
return p.fcalc_abs * (float) get_overall_scale_factor(p.hkl);
|
|
}
|
|
// For NA parameters: the B components are NA - 1 constraint rows. `kaniso` is
|
|
// exp(-0.25 * b_star.r_u_r(h)) at the current parameters, as compute_value() takes it too.
|
|
template <int NA>
|
|
double compute_value_and_derivatives(const Point &p, double kaniso, double *dy_da) const {
|
|
gemmi::Vec3 h(p.hkl);
|
|
double fcalc_abs = p.fcalc_abs;
|
|
int n = 1;
|
|
double fe = fcalc_abs * kaniso;
|
|
double y = k_overall * fe;
|
|
dy_da[0] = fe;
|
|
gemmi::SMat33<double> du = {
|
|
-0.25 * y * (h.x * h.x),
|
|
-0.25 * y * (h.y * h.y),
|
|
-0.25 * y * (h.z * h.z),
|
|
-0.5 * y * (h.x * h.y),
|
|
-0.5 * y * (h.x * h.z),
|
|
-0.5 * y * (h.y * h.z),
|
|
};
|
|
for (int j = 0; j < NA - 1; ++j)
|
|
dy_da[n + j] = gemmi::vec6_dot(constraint_matrix[j], du);
|
|
return y;
|
|
}
|
|
};
|
|
|
|
// R-factor of the current parameters over the fitted reflections. This is what the grid is
|
|
// selected on, and it is the quantity the scale exists to make small.
|
|
double RFactor(const FixedSolventFit &fit) {
|
|
double num = 0, den = 0;
|
|
for (const auto &p : fit.points) {
|
|
num += std::fabs(p.fobs - fit.compute_value(p));
|
|
den += p.fobs;
|
|
}
|
|
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.
|
|
//
|
|
// `kaniso`: each reflection's exp(-0.25 * b_star.r_u_r(h)) at the current parameters where ComputeWSSR()
|
|
// has just taken them there (nullptr: computed here) - the same value, which need not be taken twice.
|
|
template <int NA>
|
|
double ComputeLMMatrices(const FixedSolventFit &fit, const double *kaniso, 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 (size_t i = 0; i < fit.points.size(); ++i) {
|
|
const auto &p = fit.points[i];
|
|
const double k = kaniso ? kaniso[i] : std::exp(-0.25 * fit.b_star.r_u_r(gemmi::Vec3(p.hkl)));
|
|
double y = fit.compute_value_and_derivatives<NA>(p, k, 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 compute_wssr() (levmar.hpp) on a FixedSolventFit, keeping each reflection's
|
|
// exp(-0.25 * b_star.r_u_r(h)) in `kaniso` for ComputeLMMatrices() at the same parameters.
|
|
double ComputeWSSR(const FixedSolventFit &fit, std::vector<double> &kaniso) {
|
|
long double wssr = 0; // long double here notably increases the accuracy
|
|
for (size_t i = 0; i < fit.points.size(); ++i) {
|
|
const auto &p = fit.points[i];
|
|
kaniso[i] = std::exp(-0.25 * fit.b_star.r_u_r(p.hkl));
|
|
// FixedSolventFit::compute_value(), weight 1
|
|
wssr += gemmi::sq(p.get_y() - p.fcalc_abs * (float) (fit.k_overall * kaniso[i]));
|
|
}
|
|
return (double) wssr;
|
|
}
|
|
|
|
// gemmi's LevMar::fit() (levmar.hpp) with its default settings, on ComputeLMMatrices<NA>() and
|
|
// ComputeWSSR() - 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);
|
|
|
|
std::vector<double> kaniso(target.points.size());
|
|
const double initial_wssr = ComputeLMMatrices<NA>(target, nullptr, 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 = ComputeWSSR(target, kaniso);
|
|
++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, kaniso.data(), alpha, beta); // at the parameters just tried
|
|
++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
|
|
// the overall scale and the anisotropic B refitted at every grid point - Afonine, Grosse-Kunstleve
|
|
// & Adams, Acta Cryst. D61, 850-855, 2005, which searches b_sol over 10-80 A^2 in steps of 5.
|
|
// The fit is unweighted, as in both phenix and Refmac (Murshudov, Skubak, Lebedev, Pannu, Steiner,
|
|
// Nicholls, Winn, Long & Vagin, Acta Cryst. D67, 355-367, 2011, eq. 11). The physical range and the
|
|
// starting values are those of Fokine & Urzhumtsev, Acta Cryst. D58, 1387-1392, 2002.
|
|
//
|
|
// The point of the grid is that k_sol and b_sol cannot leave the physical box: gemmi's own
|
|
// fit_parameters() is an unbounded Levenberg-Marquardt, and on this corpus it reached b_sol of
|
|
// 1707 A^2 - a solvent term switched off in all but the lowest-resolution shell. Here the solvent
|
|
// pair is held fixed at each grid point and only the overall scale and the symmetry-constrained
|
|
// anisotropic B are refined, which is the well-conditioned half of the problem and is left to
|
|
// gemmi's solver rather than reimplemented.
|
|
//
|
|
// 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, 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());
|
|
if (scaling.points.size() < 20)
|
|
return report;
|
|
|
|
const bool had_solvent = scaling.use_solvent;
|
|
scaling.fix_k_sol = true; // the grid owns the solvent pair; the solver never sees it
|
|
scaling.fix_b_sol = true;
|
|
|
|
// The reflections fit_isotropic_b_approximately() fits on (its own filter).
|
|
int n_isotropic = 0;
|
|
for (const auto &p : scaling.points)
|
|
if (!(p.fobs < 1 || p.fobs < p.sigma))
|
|
++n_isotropic;
|
|
const bool independent = n_isotropic > 5;
|
|
|
|
double best_r = -1, best_k_sol = 0.35, best_b_sol = 46.0, best_k_overall = 1.0;
|
|
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}; };
|
|
// 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) {
|
|
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 {
|
|
FixedSolventFit fit(scaling);
|
|
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;
|
|
// A diverged fit gives r = NaN; latched as best_r it wins every later r < best_r.
|
|
if (std::isfinite(pf.r) && (best_r < 0 || pf.r < best_r)) {
|
|
best_r = pf.r;
|
|
best_k_sol = pf.k_sol;
|
|
best_b_sol = pf.b_sol;
|
|
best_k_overall = pf.k_overall;
|
|
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.
|
|
std::vector<PointFit> coarse;
|
|
for (double ks = box.k_lo; ks <= box.k_hi + 1e-9; ks += 0.05)
|
|
for (double bs = box.b_lo; bs <= box.b_hi + 1e-9; bs += 10.0)
|
|
coarse.push_back(PointFit{ks, bs});
|
|
try_points(coarse);
|
|
const double k0 = best_k_sol, b0 = best_b_sol;
|
|
const double k_hi2 = std::min(box.k_hi, k0 + 0.05);
|
|
const double b_hi2 = std::min(box.b_hi, b0 + 10.0);
|
|
std::vector<PointFit> fine;
|
|
for (double ks = std::max(box.k_lo, k0 - 0.05); ks <= k_hi2 + 1e-9; ks += 0.025)
|
|
for (double bs = std::max(box.b_lo, b0 - 10.0); bs <= b_hi2 + 1e-9; bs += 5.0)
|
|
fine.push_back(PointFit{ks, bs});
|
|
try_points(fine);
|
|
|
|
scaling.k_sol = best_k_sol;
|
|
scaling.b_sol = best_b_sol;
|
|
scaling.k_overall = best_k_overall;
|
|
scaling.b_star = best_b_star;
|
|
scaling.use_solvent = had_solvent;
|
|
report.r_work_fit = best_r;
|
|
return report;
|
|
}
|