leonarski_fandClaude Opus 5.5
ab6b6740f8
The 2mFo-DFc and mFo-DFc coefficients written to <prefix>_maps.mtz and <prefix>_refined_maps.mtz were on the observed scale, which keeps the anisotropic part of the overall scale (k_total). On anisotropic data every atom of the map was stretched along the weak direction. Each coefficient is now divided by exp(-h^T B*_aniso h / 4), B*_aniso the overall anisotropic B of the fit less its isotropic part, as phenix.refine does with its maps. FP and FC stay on the observed scale, and the maps the refinement and the validation read themselves (blob finder, solvent audit, mean density at the atoms) are unchanged, so the R-factors, the blobs and every REFINE_ key are the same. On a 0.97 A set with dB ~10 A^2 the median half-maximum axis ratio of well-ordered carbons rises from 0.67 to 0.75; on isotropic sets the maps hardly change. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01UqXYdL4qBGYRiSXSGXfCLu
599 lines
26 KiB
C++
599 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;
|
|
}
|
|
|
|
gemmi::SMat33<double> AnisotropicPart(const gemmi::SMat33<double> &b_star, const gemmi::UnitCell &cell) {
|
|
const gemmi::SMat33<double> b = b_star.transformed_by(cell.orth.mat);
|
|
const double b_iso = (b.u11 + b.u22 + b.u33) / 3.0;
|
|
// b_iso I in Cartesian is b_iso F F^T in the reciprocal basis
|
|
const gemmi::Mat33 f = cell.frac.mat.multiply(cell.frac.mat.transpose());
|
|
return {b_star.u11 - b_iso * f[0][0], b_star.u22 - b_iso * f[1][1], b_star.u33 - b_iso * f[2][2],
|
|
b_star.u12 - b_iso * f[0][1], b_star.u13 - b_iso * f[0][2], b_star.u23 - b_iso * f[1][2]};
|
|
}
|
|
|
|
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);
|
|
ParallelChunks(static_cast<int>(n), ThreadsForWork(n, nthreads), [&](int lo, int hi) {
|
|
for (int i = lo; i < hi; 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
|