leonarski_f
e406eb17ea
- New docs/CPU_DATA_ANALYSIS_REFINEMENT.md (§15): the ML amplitude target and its FFT gradient, scale/robust fit/fine solvent mask, D and the variance (work/free), the embedded Monomer Library subset and topology, geometry/metal/KL-B restraints, the minimiser, the weight trials, refined maps, blobs, GPU/CPU, and a "Not yet" list; with its own References. Linked from TECHNICAL.md and CPU_DATA_ANALYSIS.md; §14 no longer says the structure is never refined. - RUGNUX_REPORT.md: a Refinement (REFINE_*) section - every REFINE_ and REFINE_SCORE_ key, the --mode refine input keys and the _blobs.json fields. - RUGNUX_MODEL_AND_REFERENCE.md: "What it does not do yet"; RUGNUX_ADVANCED.md: --model mentions the refinement, --no-refine as in the usage, a --mode refine options table; GPU_DATA_ANALYSIS.md: the refinement's CUDA engines and the missing Metal path. - ACKNOWLEDGEMENT.md: Brunger, Karplus & Petsko 1989 (gradient-balanced starting weight), Nocedal 1980 (two-loop recursion), Charbonnier et al. 1997 (pseudo-Huber), Jiang & Brunger 1994 (flat solvent mask), Kullback & Leibler 1951; all DOIs checked against Crossref. - In-source credit lines: the weight balancing (replacing a "Brunger (1992)-style" line that named the R-free paper), Nocedal 1980, Charbonnier 1997, Jiang & Brunger 1994, Emsley et al. 2010 at FindDensityBlobs, Read 1986 at the refined maps. Comments only. - docs/THIRD_PARTY_NOTICES.md regenerated from the root file with update_version.sh's own step (it lacked the CCP4 Monomer Library row). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_015CeEGgpKdfkCTJb3YNMJcu (cherry picked from commit 09375dba351bae0cc73e9a6f0f79455dc92b3a10)
588 lines
26 KiB
C++
588 lines
26 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
|
|
// The smooth (pseudo-Huber) loss of Charbonnier et al. (1997) IEEE Trans. Image Process. 6, 298-311
|
|
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);
|
|
double k0 = box.k_start, b0 = box.b_start;
|
|
if (k0 < 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);
|
|
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,
|
|
double k_start, double b_start) {
|
|
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, k_start, b_start);
|
|
} catch (const ModelScaleGPUTooFewReflections &) {
|
|
set_gpu(current);
|
|
release();
|
|
return FitModelScale(scaling, {.b_hi = b_hi, .robust = robust, .k_start = k_start, .b_start = b_start}, 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
|