leonarski_f
46ad2c10ed
The score's robust fine pass (pseudo-Huber IRLS of ModelScaleBox::robust) on the device, so that refinement on a GPU no longer fits its scale on the host every cycle - with the trials side by side that was a robust grid search on one thread per slot (13-133 s per refinement on the seven timing sets, against 1.7-8 s before the score commit). - ModelScaleGPU: the weights multiply the WSSR and J^T J / J^T r in scale_sums (a weight row per fit, reserved at the first robust fit); robust_residuals / robust_weights kernels; the median of each fit's first residuals taken on the host as FitModelScale takes it (nth_element at n/2, on the fits' threads); LevMarBatch split out of FitBatch so that the rounds restart from where a fit stands. FitSolvent(..., robust) refits the coarse winner in the fine pass, as FitModelScale's robust fit does; the plain fit is unchanged. - ModelScaleGPUPool::Fit(..., robust); its streams at the highest priority: a fit is a chain of small launches the host waits on, and on a card shared with the trials' long kernels each wait queued behind them (refit.scale 58-87 ms -> 14-45 ms per cycle, the result the same to the bit). - Refinement: the main context and every trial slot fit on their pool (robust); the CPU engine keeps FitModelScale's robust fit. The card's fit is FitModelScale's to rounding (float residuals and weights, the steps' anisotropic factor in float), deterministic run to run and whichever engine a fit takes: R-work/R-free of the 14 parity sets on the GPU equal the CPU engine's to 4 digits (largest difference 1.4e-6 on the 3.6 A set). Test: ModelScaleGPUPool_RobustMatchesFitModelScale (one reflection in fifty a 4x outlier, five cells). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_015CeEGgpKdfkCTJb3YNMJcu (cherry picked from commit d15f048cbd7ed2d6b430f9f9cb5018ca284da8f0)
582 lines
25 KiB
C++
582 lines
25 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"
|
|
|
|
#ifdef JFJOCH_USE_CUDA
|
|
#include "../../common/CUDAWrapper.h"
|
|
#include "../indexing/CUDAMemHelpers.h"
|
|
#include "ModelScaleGPU.h"
|
|
#include "RigidBodyGPUEngine.h" // ModelScaleGPUTooFewReflections, RigidBodyGPUEngine::CurrentDevice
|
|
#endif
|
|
|
|
namespace {
|
|
|
|
constexpr int ROBUST_ROUNDS = 4; // reweightings of the robust fit (ModelScaleBox::robust)
|
|
|
|
// 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 weight = 1.0; // 1, but for the robust fit's reweighting (RobustReweight)
|
|
double get_y() const { return fobs; }
|
|
double get_weight() const { return weight; }
|
|
};
|
|
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), 1.0};
|
|
if (p.fobs < 1 || p.fobs < p.sigma) // skip weak reflections
|
|
continue;
|
|
double fcalc_abs = points[i].fcalc_abs;
|
|
if (fcalc_abs == 0) // local change in gemmi_gph: no log ratio
|
|
continue;
|
|
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;
|
|
}
|
|
};
|
|
|
|
// The robust fit's weights at the current parameters: each reflection's residual r = |Fo| - |F_model| weighted by
|
|
// 1 / sqrt(1 + (r / c)^2), with c the median |r| of the unweighted fit. Iterated, the weighted least squares then
|
|
// minimises sum c^2 2 (sqrt(1 + (r / c)^2) - 1) - quadratic for the bulk of the reflections and linear for the few
|
|
// a model fits badly (a beam-stop shadow, an ice ring, a missing part of the model), which least squares would
|
|
// otherwise let set the scale of all the rest.
|
|
// Following Huber (1964) Ann. Math. Statist. 35, 73-101 for the M-estimator, as iteratively reweighted least squares
|
|
void RobustReweight(FixedSolventFit &fit, double &c) {
|
|
std::vector<double> r(fit.points.size());
|
|
for (size_t i = 0; i < fit.points.size(); i++)
|
|
r[i] = std::fabs(fit.points[i].fobs - fit.compute_value(fit.points[i]));
|
|
if (c <= 0) {
|
|
std::vector<double> sorted = r;
|
|
std::nth_element(sorted.begin(), sorted.begin() + sorted.size() / 2, sorted.end());
|
|
c = sorted[sorted.size() / 2];
|
|
}
|
|
if (c <= 0)
|
|
return;
|
|
for (size_t i = 0; i < fit.points.size(); i++)
|
|
fit.points[i].weight = 1.0 / std::sqrt(1.0 + gemmi::sq(r[i] / c));
|
|
}
|
|
|
|
// 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;
|
|
const double w = p.get_weight();
|
|
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] * w;
|
|
alpha[NA * j + k] += used ? a : 0.0;
|
|
}
|
|
const double b = dy_sig * dy_da[j] * w;
|
|
beta[j] += used ? b : 0.0;
|
|
}
|
|
wssr += gemmi::sq(dy_sig) * w;
|
|
}
|
|
// 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()
|
|
wssr += gemmi::sq(p.get_y() - p.fcalc_abs * (float) (fit.k_overall * kaniso[i])) * p.get_weight();
|
|
}
|
|
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;
|
|
// Not in gemmi: a fit that has walked to where k_overall exp(-B s^2) |F| underflows at every
|
|
// reflection has no step left to take - the squares in alpha are all 0 while beta keeps a
|
|
// residue near 1e-210, and jordan_solve() throws on that. It ends here, at its best so far.
|
|
if (alpha[0] == 0)
|
|
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;
|
|
bool robust_pass = false; // the fine pass of a robust fit (ModelScaleBox::robust)
|
|
auto fit_point = [&scaling, &robust_pass](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
|
|
if (robust_pass) {
|
|
double c = 0;
|
|
for (int round = 0; round < ROBUST_ROUNDS; round++) {
|
|
RobustReweight(fit, c);
|
|
LevMarFit(fit);
|
|
}
|
|
}
|
|
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. b_sol is stepped by
|
|
// 10 A^2, or in eight values where the box is wider than the default, so a raised ceiling costs
|
|
// no more grid points.
|
|
const double b_step = std::max(10.0, (box.b_hi - box.b_lo) / 7.0);
|
|
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 += b_step)
|
|
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 + b_step);
|
|
std::vector<PointFit> fine;
|
|
// A robust fit takes the fine pass robustly: each of its points, and the coarse winner again, by robust least
|
|
// squares - the coarse pass only finds where to look.
|
|
if (box.robust) {
|
|
robust_pass = true;
|
|
fitted.clear();
|
|
best_r = -1;
|
|
}
|
|
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 - b_step); bs <= b_hi2 + 1e-9; bs += b_step / 2.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;
|
|
}
|
|
|
|
double SolventBMax(const gemmi::Model &model) {
|
|
double sum = 0;
|
|
size_t n = 0;
|
|
for (const gemmi::Chain &chain : model.chains)
|
|
for (const gemmi::Residue &residue : chain.residues)
|
|
for (const gemmi::Atom &atom : residue.atoms) {
|
|
sum += atom.b_iso;
|
|
++n;
|
|
}
|
|
const double b_max = ModelScaleBox{}.b_hi;
|
|
return n > 0 ? std::max(b_max, 2.0 * sum / static_cast<double>(n)) : b_max;
|
|
}
|
|
|
|
#ifdef JFJOCH_USE_CUDA
|
|
// A fit is a chain of small launches, each waited for by the host's Levenberg-Marquardt: on a card shared with
|
|
// long kernels (refinement's trials side by side) every wait would queue behind them, so the stream goes first.
|
|
int HighestStreamPriority() {
|
|
int least = 0, greatest = 0;
|
|
cudaDeviceGetStreamPriorityRange(&least, &greatest);
|
|
return greatest;
|
|
}
|
|
|
|
struct ModelScaleGPUPool::Engine {
|
|
CudaStream stream;
|
|
ModelScaleGPU scale;
|
|
CudaDevicePtr<float2> fcalc, fmask;
|
|
explicit Engine(size_t max_points)
|
|
: stream(cudaStreamNonBlocking, HighestStreamPriority()), scale(stream, max_points), fcalc(max_points),
|
|
fmask(max_points) {}
|
|
};
|
|
|
|
size_t ModelScaleGPUPool::DeviceBytes(size_t max_points, size_t engines, bool robust) {
|
|
return engines * (ModelScaleGPU::DeviceBytes(max_points) + 2 * max_points * sizeof(float2) +
|
|
(robust ? ModelScaleGPU::RobustBytes(max_points) : 0));
|
|
}
|
|
|
|
ModelScaleGPUPool::ModelScaleGPUPool(int device, size_t max_points, size_t engines)
|
|
: device_(device), max_points_(max_points) {
|
|
const int current = RigidBodyGPUEngine::CurrentDevice();
|
|
set_gpu(device_);
|
|
for (size_t i = 0; i < engines; i++) {
|
|
engines_.push_back(std::make_unique<Engine>(max_points));
|
|
idle_.push_back(engines_.back().get());
|
|
}
|
|
set_gpu(current);
|
|
}
|
|
|
|
ModelScaleGPUPool::~ModelScaleGPUPool() {
|
|
const int current = RigidBodyGPUEngine::CurrentDevice();
|
|
set_gpu(device_);
|
|
engines_.clear();
|
|
set_gpu(current);
|
|
}
|
|
|
|
ModelScaleReport ModelScaleGPUPool::Fit(gemmi::Scaling<float> &scaling, double b_hi, size_t nthreads, bool robust) {
|
|
const size_t n = scaling.points.size();
|
|
std::vector<std::array<int, 3>> hkl(n);
|
|
std::vector<double> stol2(n);
|
|
std::vector<float> fobs(n), sigma(n);
|
|
std::vector<float2> fc(n), fm(n);
|
|
for (size_t i = 0; i < n; i++) {
|
|
const auto &p = scaling.points[i];
|
|
hkl[i] = {p.hkl[0], p.hkl[1], p.hkl[2]};
|
|
stol2[i] = p.stol2;
|
|
fobs[i] = p.fobs;
|
|
sigma[i] = p.sigma;
|
|
fc[i] = {p.fcmol.real(), p.fcmol.imag()};
|
|
fm[i] = {p.fmask.real(), p.fmask.imag()};
|
|
}
|
|
std::vector<std::array<double, 6>> constraints;
|
|
for (const gemmi::Vec6 &c : scaling.constraint_matrix)
|
|
constraints.push_back(c);
|
|
double frac[9];
|
|
for (int i = 0; i < 3; i++)
|
|
for (int j = 0; j < 3; j++)
|
|
frac[3 * i + j] = scaling.cell.frac.mat[i][j];
|
|
|
|
Engine *e;
|
|
{
|
|
std::unique_lock lock(m_);
|
|
cv_.wait(lock, [&] { return !idle_.empty(); });
|
|
e = idle_.back();
|
|
idle_.pop_back();
|
|
}
|
|
auto release = [&] {
|
|
std::lock_guard lock(m_);
|
|
idle_.push_back(e);
|
|
cv_.notify_one();
|
|
};
|
|
const int current = RigidBodyGPUEngine::CurrentDevice();
|
|
ModelSolventFit fit;
|
|
try {
|
|
set_gpu(device_);
|
|
e->scale.SetPoints(hkl, stol2, fobs, sigma, constraints, frac);
|
|
if (cudaMemcpyAsync(e->fcalc, fc.data(), n * sizeof(float2), cudaMemcpyHostToDevice, e->stream) != cudaSuccess
|
|
|| cudaMemcpyAsync(e->fmask, fm.data(), n * sizeof(float2), cudaMemcpyHostToDevice, e->stream) != cudaSuccess)
|
|
throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "ModelScaleGPUPool: the upload failed");
|
|
fit = e->scale.FitSolvent(e->fcalc, e->fmask, b_hi, robust, nthreads);
|
|
} catch (const ModelScaleGPUTooFewReflections &) {
|
|
set_gpu(current);
|
|
release();
|
|
return FitModelScale(scaling, {.b_hi = b_hi, .robust = robust}, nthreads);
|
|
} catch (...) {
|
|
set_gpu(current);
|
|
release();
|
|
throw;
|
|
}
|
|
set_gpu(current);
|
|
release();
|
|
|
|
scaling.k_sol = fit.k_sol;
|
|
scaling.b_sol = fit.b_sol;
|
|
scaling.k_overall = fit.scale.k_overall;
|
|
const double *b = fit.scale.b_star;
|
|
scaling.b_star = {b[0], b[1], b[2], b[3], b[4], b[5]};
|
|
ModelScaleReport report;
|
|
report.n_points = static_cast<int>(n);
|
|
report.n_grid = fit.n_grid;
|
|
report.r_work_fit = fit.r;
|
|
return report;
|
|
}
|
|
#endif
|