leonarski_fandClaude Opus 5.5
210664f454
Two ways the scale fit's normal matrix went singular and gemmi's jordan_solve threw "Trying to reverse singular matrix. Column 0 is zeroed.", ending the run: 1. The isotropic start (fit_isotropic_b_approximately) takes log(Fobs/|Fcalc|) over the strong reflections. At a phase-restricted reflection whose Fcalc and Fmask share a phase, some solvent pair of the k_sol/b_sol grid can cancel the model exactly in float, |Fcalc + k_sol exp(-b_sol s^2) Fmask| = 0. log(Fobs/0) = inf made that grid point's k_overall and B NaN, every sum of the Levenberg-Marquardt NaN, and the solve failed. This is what stopped three open-arm sets in the output pass after MR on the GPU (the scale fit after model validation's rigid body): one strong reflection, one grid point. The CPU path has the same exposure; which grid point hits an exact zero depends on the last bits of Fcalc, so the --no-gpu run of the same set went through. Such a reflection has no log ratio and is now left out of the start - in the GPU kernel, in FixedSolventFit::Start and, as a marked local change, in gemmi's own fit_isotropic_b_approximately, so the three stay equal and the held-solvent fits are covered too. 2. In FitModelScale's CPU Levenberg-Marquardt (ModelValidation_CCModelUnrelatedModelConfirmsNothing in a CUDA=OFF build): on 66 reflections of an unrelated model the fit walked to k_overall -424 and a B where k exp(-B s^2) |F| underflows at every reflection; the squares in J^T J were all 0 while J^T r kept a residue near 1e-210, which jordan_solve rejects. The fit now ends there, at its best so far, as it would at the lambda limit; that grid point then loses on R. The GPU fit takes the derivatives in float and their squares in double, so it cannot reach this state. Both changes only touch fits that threw before. Byte identity checked against an rc176 build on the standard and inhouse timing benchmarks and on 4nwv, 8tyy, 6oel (p.hkl, p_model.pdb, MTZ data, MR_* report keys). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01PXc1zD1yxHqaR1DGdCTHL4
531 lines
23 KiB
C++
531 lines
23 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 {
|
|
|
|
// 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;
|
|
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;
|
|
}
|
|
};
|
|
|
|
// 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;
|
|
// 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;
|
|
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. 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;
|
|
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
|
|
struct ModelScaleGPUPool::Engine {
|
|
CudaStream stream;
|
|
ModelScaleGPU scale;
|
|
CudaDevicePtr<float2> fcalc, fmask;
|
|
explicit Engine(size_t max_points) : scale(stream, max_points), fcalc(max_points), fmask(max_points) {}
|
|
};
|
|
|
|
size_t ModelScaleGPUPool::DeviceBytes(size_t max_points, size_t engines) {
|
|
return engines * (ModelScaleGPU::DeviceBytes(max_points) + 2 * max_points * sizeof(float2));
|
|
}
|
|
|
|
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) {
|
|
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);
|
|
} catch (const ModelScaleGPUTooFewReflections &) {
|
|
set_gpu(current);
|
|
release();
|
|
return FitModelScale(scaling, {.b_hi = b_hi}, 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
|