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)
1127 lines
51 KiB
C++
1127 lines
51 KiB
C++
// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
|
|
// SPDX-License-Identifier: GPL-3.0-only
|
|
|
|
#include "Refinement.h"
|
|
|
|
#include <algorithm>
|
|
#include <chrono>
|
|
#include <map>
|
|
#include <optional>
|
|
#include <cmath>
|
|
#include <random>
|
|
#include <thread>
|
|
#include <unordered_map>
|
|
|
|
#include <ceres/gradient_problem.h>
|
|
#include <ceres/gradient_problem_solver.h>
|
|
#include <ceres/jet.h>
|
|
#include <ceres/rotation.h>
|
|
|
|
#include "gemmi/fourier.hpp" // get_size_for_hkl, get_f_phi_on_grid
|
|
#include "gemmi/modify.hpp" // remove_hydrogens
|
|
#include "gemmi/scaling.hpp"
|
|
#include "gemmi/stats.hpp" // calculate_data_statistics
|
|
|
|
#include "../../common/JFJochException.h"
|
|
#include "../../common/JFJochMath.h"
|
|
#include "../../common/Logger.h"
|
|
#include "ModelFFT.h"
|
|
#include "ModelScaling.h"
|
|
#include "RefinementMath.h"
|
|
#include "RefinementXray.h"
|
|
#include "SigmaA.h"
|
|
#ifdef JFJOCH_USE_CUDA
|
|
#include "../../common/CUDAWrapper.h"
|
|
#include "RefinementGPU.h"
|
|
#endif
|
|
|
|
namespace {
|
|
|
|
constexpr double B_FLOOR = 2.0; // A^2: no atom's B goes below it
|
|
// nor above twice the model's largest B, within these bounds (the GPU's gather is reserved for it)
|
|
constexpr double B_CEILING_MIN = 100.0, B_CEILING_MAX = 300.0;
|
|
// How far below the smallest B of the model a macro-cycle's B refinement may take an atom; the blur of the
|
|
// cycle is chosen for that.
|
|
constexpr double B_STEP_DOWN = 10.0;
|
|
constexpr double REPULSION_MARGIN = 0.5; // A: pairs this much beyond their contact distance are listed too
|
|
// The geometry the weight is chosen within (Afonine et al. 2012 select among weights by R-free with the
|
|
// geometry held to a target in the same way).
|
|
constexpr double MAX_RMSD_BONDS = 0.02, MAX_RMSD_ANGLES = 2.0;
|
|
// R-free differences smaller than this do not decide between two weights: the tighter geometry is taken.
|
|
constexpr double R_FREE_TIE = 0.001;
|
|
|
|
std::string Joined(const std::vector<std::string> &names) {
|
|
std::string out;
|
|
for (const std::string &n : names)
|
|
out += (out.empty() ? "" : " ") + n;
|
|
return out;
|
|
}
|
|
|
|
std::vector<gemmi::Atom *> AtomList(gemmi::Model &model) {
|
|
std::vector<gemmi::Atom *> atoms;
|
|
for (gemmi::Chain &ch : model.chains)
|
|
for (gemmi::Residue &r : ch.residues)
|
|
for (gemmi::Atom &a : r.atoms)
|
|
atoms.push_back(&a);
|
|
return atoms;
|
|
}
|
|
|
|
// The model as refinement takes it: the first model only, without hydrogens and atoms of unknown element,
|
|
// every B isotropic (an anisotropic one becomes its equivalent isotropic B) and inside the allowed range.
|
|
gemmi::Structure PrepareModel(const gemmi::Structure &placed, const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg) {
|
|
gemmi::Structure st = placed;
|
|
st.models.resize(1);
|
|
st.cell = cell;
|
|
st.spacegroup_hm = sg.xhm();
|
|
gemmi::remove_hydrogens(st.models[0]);
|
|
for (gemmi::Chain &ch : st.models[0].chains)
|
|
for (gemmi::Residue &r : ch.residues) {
|
|
std::erase_if(r.atoms, [](const gemmi::Atom &a) { return a.element == gemmi::El::X; });
|
|
for (gemmi::Atom &a : r.atoms) {
|
|
if (a.aniso.nonzero()) {
|
|
a.b_iso = static_cast<float>(gemmi::u_to_b() * a.aniso.trace() / 3.0);
|
|
a.aniso = {0, 0, 0, 0, 0, 0};
|
|
}
|
|
a.b_iso = static_cast<float>(std::clamp<double>(a.b_iso, B_FLOOR, B_CEILING_MAX));
|
|
}
|
|
}
|
|
gemmi::remove_empty_children(st);
|
|
st.setup_cell_images();
|
|
return st;
|
|
}
|
|
|
|
// What every stage works on: the model, the reflections, the engine and the row model of the cycle.
|
|
struct Context {
|
|
gemmi::Structure st;
|
|
std::vector<gemmi::Atom *> atoms;
|
|
gemmi::UnitCell cell;
|
|
const gemmi::SpaceGroup *sg = nullptr;
|
|
double d_min = 0;
|
|
RefinementReflections refl;
|
|
std::unique_ptr<RefinementEngine> engine;
|
|
std::unique_ptr<BulkSolventMask> mask;
|
|
#ifdef JFJOCH_USE_CUDA
|
|
std::unique_ptr<ModelScaleGPUPool> scale_gpu;
|
|
#endif
|
|
size_t nthreads = 1;
|
|
int device = -1; // the engine's GPU
|
|
size_t engine_bytes = 0; // and what it holds there
|
|
XrayRowModel rows;
|
|
std::vector<std::complex<float>> fcalc, fmask;
|
|
double k_sol = 0, b_sol = 0;
|
|
double k_overall = 0;
|
|
gemmi::SMat33<double> b_star{0, 0, 0, 0, 0, 0};
|
|
double b_floor = B_FLOOR; // the cycle's lowest allowed B
|
|
double b_ceiling = B_CEILING_MAX;
|
|
// Riding hydrogens: per atom the parent it rides on (-1: none), and where it sits from the parent.
|
|
std::vector<int> rider_parent;
|
|
std::vector<gemmi::Position> rider_offset;
|
|
int evaluations = 0;
|
|
|
|
gemmi::Model &Model() { return st.models[0]; }
|
|
|
|
// The hydrogens moved with their parents, at the offsets they were last placed at, and given their B.
|
|
void MoveRiders() {
|
|
for (size_t i = 0; i < rider_parent.size(); i++)
|
|
if (rider_parent[i] >= 0) {
|
|
atoms[i]->pos = atoms[rider_parent[i]]->pos + rider_offset[i];
|
|
atoms[i]->b_iso = atoms[rider_parent[i]]->b_iso;
|
|
}
|
|
}
|
|
// What the target says of a hydrogen is said of its parent.
|
|
void FoldRiders(std::vector<std::array<double, 4>> &g) const {
|
|
for (size_t i = 0; i < rider_parent.size(); i++)
|
|
if (rider_parent[i] >= 0)
|
|
for (int c = 0; c < 4; c++)
|
|
g[rider_parent[i]][c] += g[i][c];
|
|
}
|
|
|
|
std::vector<std::array<double, 4>> XyzB() const {
|
|
std::vector<std::array<double, 4>> out(atoms.size());
|
|
for (size_t i = 0; i < atoms.size(); i++)
|
|
out[i] = {atoms[i]->pos.x, atoms[i]->pos.y, atoms[i]->pos.z, atoms[i]->b_iso};
|
|
return out;
|
|
}
|
|
std::vector<gemmi::Position> Positions() const {
|
|
std::vector<gemmi::Position> out;
|
|
for (const gemmi::Atom *a : atoms)
|
|
out.push_back(a->pos);
|
|
return out;
|
|
}
|
|
void SetPositions(const std::vector<gemmi::Position> &p) {
|
|
for (size_t i = 0; i < atoms.size(); i++)
|
|
atoms[i]->pos = p[i];
|
|
}
|
|
std::vector<double> B() const {
|
|
std::vector<double> out;
|
|
for (const gemmi::Atom *a : atoms)
|
|
out.push_back(a->b_iso);
|
|
return out;
|
|
}
|
|
void SetB(const std::vector<double> &b) {
|
|
for (size_t i = 0; i < atoms.size(); i++)
|
|
atoms[i]->b_iso = static_cast<float>(b[i]);
|
|
}
|
|
|
|
// A new macro-cycle: the blur for the B the cycle may reach, the bulk-solvent mask of the model as it is,
|
|
// the scale and the solvent fitted to the working set, and the likelihood's parameters from the free set.
|
|
// Returns the R-factors of the model as it stands.
|
|
RFactors Refit() {
|
|
const RefinementStage stage("refit");
|
|
double b_min = b_ceiling;
|
|
for (const gemmi::Atom *a : atoms)
|
|
b_min = std::min<double>(b_min, a->b_iso);
|
|
b_floor = std::max(B_FLOOR, b_min - B_STEP_DOWN);
|
|
engine->SetBlur(RefinementBlur(cell, *sg, d_min, b_floor));
|
|
{
|
|
const RefinementStage s1("refit.fcalc");
|
|
engine->Fcalc(Model(), fcalc);
|
|
}
|
|
{
|
|
const RefinementStage s2("refit.fmask");
|
|
mask->Fmask(Model(), refl.hkl, fmask);
|
|
}
|
|
std::optional<RefinementStage> s3("refit.scale");
|
|
|
|
gemmi::AsuData<std::complex<float>> fc, fm;
|
|
gemmi::AsuData<gemmi::ValueSigma<float>> fo;
|
|
fc.unit_cell_ = fm.unit_cell_ = fo.unit_cell_ = cell;
|
|
fc.spacegroup_ = fm.spacegroup_ = fo.spacegroup_ = sg;
|
|
for (size_t m = 0; m < refl.Size(); m++) {
|
|
fc.v.push_back({refl.hkl[m], fcalc[m]});
|
|
fm.v.push_back({refl.hkl[m], fmask[m]});
|
|
if (!refl.free[m])
|
|
fo.v.push_back({refl.hkl[m], {refl.fo[m], 1.0f}});
|
|
}
|
|
gemmi::Scaling<float> scaling(cell, sg);
|
|
scaling.use_solvent = true;
|
|
scaling.prepare_points(fc, fo, &fm);
|
|
// the scale of the score: by robust least squares (ModelScaleBox::robust)
|
|
#ifdef JFJOCH_USE_CUDA
|
|
if (scale_gpu)
|
|
scale_gpu->Fit(scaling, SolventBMax(Model()), nthreads, true);
|
|
else
|
|
#endif
|
|
FitModelScale(scaling, {.b_hi = SolventBMax(Model()), .robust = true}, nthreads);
|
|
s3.reset();
|
|
k_sol = scaling.k_sol;
|
|
b_sol = scaling.b_sol;
|
|
k_overall = scaling.k_overall;
|
|
b_star = scaling.b_star;
|
|
rows.k_total.resize(refl.Size());
|
|
rows.fmask.resize(refl.Size());
|
|
std::vector<std::complex<float>> fmodel(refl.Size());
|
|
for (size_t m = 0; m < refl.Size(); m++) {
|
|
rows.k_total[m] = static_cast<float>(scaling.get_overall_scale_factor(refl.hkl[m]));
|
|
rows.fmask[m] = static_cast<float>(scaling.get_solvent_scale(cell.calculate_stol_sq(refl.hkl[m]))) * fmask[m];
|
|
fmodel[m] = rows.k_total[m] * (fcalc[m] + rows.fmask[m]);
|
|
}
|
|
{
|
|
const RefinementStage s4("refit.ml");
|
|
EstimateLikelihoodParameters(refl, fmodel, rows, nthreads);
|
|
}
|
|
engine->SetRowModel(rows);
|
|
return ComputeRFactors(refl, fmodel);
|
|
}
|
|
|
|
// Every term of the score of the atoms as they are, at the scale, the solvent and the likelihood parameters
|
|
// of the last Refit().
|
|
RefinementScore Score() {
|
|
const std::vector<std::complex<float>> fmodel = FModel();
|
|
RefinementScore sc;
|
|
sc.k_overall = k_overall;
|
|
sc.k_sol = k_sol;
|
|
sc.b_sol = b_sol;
|
|
sc.mask_spacing = mask->Spacing();
|
|
// b_star is in the reciprocal basis (exp(-h^T b_star h / 4)); in Cartesian B = O b_star O^T
|
|
const gemmi::Mat33 b = cell.orth.mat.multiply(b_star.as_mat33()).multiply(cell.orth.mat.transpose());
|
|
sc.b_iso = (b[0][0] + b[1][1] + b[2][2]) / 3.0;
|
|
sc.b_aniso = {b[0][0] - sc.b_iso, b[1][1] - sc.b_iso, b[2][2] - sc.b_iso, b[0][1], b[0][2], b[1][2]};
|
|
const RFactors r = ComputeRFactors(refl, fmodel);
|
|
sc.r_work = r.r_work;
|
|
sc.r_free = r.r_free;
|
|
sc.n_work = r.n_work;
|
|
sc.n_free = r.n_free;
|
|
sc.nll_work = sc.nll_free = 0;
|
|
double fom_sum = 0;
|
|
for (size_t m = 0; m < refl.Size(); m++) {
|
|
double dfc = 0;
|
|
const double fc = std::abs(fmodel[m]);
|
|
const double nll = refine_math::AmplitudeNegLogLikelihood<double>(refl.fo[m], fc, rows.d[m], rows.var[m],
|
|
refl.centric[m], dfc);
|
|
(refl.free[m] ? sc.nll_free : sc.nll_work) += nll;
|
|
// the figure of merit of the same likelihood: <cos(dphi)> = I1(X)/I0(X), X = 2 Fo D Fc / var, for an
|
|
// acentric reflection; tanh(X / 2) for a centric one
|
|
// Following Murshudov, Vagin & Dodson (1997) Acta Cryst. D53, 240-255
|
|
const double x = 2.0 * refl.fo[m] * rows.d[m] * fc / rows.var[m];
|
|
double log_i0 = 0, ratio = 0;
|
|
refine_math::BesselI0I1(x, log_i0, ratio);
|
|
fom_sum += refl.centric[m] ? std::tanh(0.5 * x) : ratio;
|
|
}
|
|
sc.fom = fom_sum / static_cast<double>(refl.Size());
|
|
for (int p : rider_parent)
|
|
sc.hydrogens += p >= 0;
|
|
constexpr int SHELLS = 20;
|
|
double s_lo = refl.inv_d2[0], s_hi = refl.inv_d2[0];
|
|
for (double s2 : refl.inv_d2) {
|
|
s_lo = std::min(s_lo, s2);
|
|
s_hi = std::max(s_hi, s2);
|
|
}
|
|
std::vector<std::array<double, 4>> sums(SHELLS, {0, 0, 0, 0});
|
|
sc.shells.resize(SHELLS);
|
|
for (int k = 0; k < SHELLS; k++) {
|
|
sc.shells[k].d_max = 1.0 / std::sqrt(s_lo + (s_hi - s_lo) * k / SHELLS);
|
|
sc.shells[k].d_min = 1.0 / std::sqrt(s_lo + (s_hi - s_lo) * (k + 1) / SHELLS);
|
|
}
|
|
for (size_t m = 0; m < refl.Size(); m++) {
|
|
const int k = std::min(SHELLS - 1, static_cast<int>((refl.inv_d2[m] - s_lo) / (s_hi - s_lo) * SHELLS));
|
|
const int f = refl.free[m] ? 2 : 0;
|
|
sums[k][f] += std::fabs(refl.fo[m] - std::abs(fmodel[m]));
|
|
sums[k][f + 1] += refl.fo[m];
|
|
(refl.free[m] ? sc.shells[k].n_free : sc.shells[k].n_work)++;
|
|
}
|
|
for (int k = 0; k < SHELLS; k++) {
|
|
if (sums[k][1] > 0)
|
|
sc.shells[k].r_work = sums[k][0] / sums[k][1];
|
|
if (sums[k][3] > 0)
|
|
sc.shells[k].r_free = sums[k][2] / sums[k][3];
|
|
}
|
|
return sc;
|
|
}
|
|
|
|
// F_model with the cycle's scale and solvent, for the atoms as they are now.
|
|
std::vector<std::complex<float>> FModel() {
|
|
engine->Fcalc(Model(), fcalc);
|
|
std::vector<std::complex<float>> fmodel(refl.Size());
|
|
for (size_t m = 0; m < refl.Size(); m++)
|
|
fmodel[m] = rows.k_total[m] * (fcalc[m] + rows.fmask[m]);
|
|
return fmodel;
|
|
}
|
|
};
|
|
|
|
ceres::GradientProblemSolver::Options SolverOptions(int iterations) {
|
|
ceres::GradientProblemSolver::Options options;
|
|
options.line_search_direction_type = ceres::LBFGS;
|
|
options.max_num_iterations = iterations;
|
|
options.function_tolerance = 1e-7;
|
|
options.gradient_tolerance = 1e-10;
|
|
options.parameter_tolerance = 1e-9;
|
|
options.logging_type = ceres::SILENT;
|
|
return options;
|
|
}
|
|
|
|
// The positions or the B factors of the refinable atoms, against w times the X-ray target plus a restraint set.
|
|
// Following Liu & Nocedal (1989) Math. Program. 45, 503-528 for the minimiser (Ceres' L-BFGS)
|
|
class AtomTarget : public ceres::FirstOrderFunction {
|
|
public:
|
|
enum class Mode { Xyz, B };
|
|
AtomTarget(Context &ctx, Mode mode, const std::vector<int> ¶m_atoms, double weight, int restraint_set)
|
|
: ctx_(ctx), mode_(mode), param_atoms_(param_atoms), weight_(weight), restraint_set_(restraint_set) {}
|
|
|
|
int NumParameters() const override {
|
|
return static_cast<int>(param_atoms_.size()) * (mode_ == Mode::Xyz ? 3 : 1);
|
|
}
|
|
|
|
bool Evaluate(const double *p, double *cost, double *gradient) const override {
|
|
const RefinementStage stage(mode_ == Mode::Xyz ? "eval.xyz" : "eval.b");
|
|
for (size_t k = 0; k < param_atoms_.size(); k++) {
|
|
gemmi::Atom &a = *ctx_.atoms[param_atoms_[k]];
|
|
if (mode_ == Mode::Xyz)
|
|
a.pos = gemmi::Position(p[3 * k], p[3 * k + 1], p[3 * k + 2]);
|
|
else
|
|
a.b_iso = static_cast<float>(std::clamp(p[k], ctx_.b_floor, ctx_.b_ceiling));
|
|
}
|
|
ctx_.MoveRiders();
|
|
ctx_.evaluations++;
|
|
std::vector<std::array<double, 4>> gx, gr;
|
|
double xray = 0, restraints = 0;
|
|
{
|
|
const RefinementStage s1("eval.engine");
|
|
if (restraint_set_ >= 0)
|
|
xray = ctx_.engine->EvaluateWithRestraints(ctx_.Model(), restraint_set_, gradient != nullptr ? &gx : nullptr,
|
|
gradient != nullptr ? &gr : nullptr, restraints);
|
|
else
|
|
xray = ctx_.engine->Evaluate(ctx_.Model(), gradient != nullptr ? &gx : nullptr);
|
|
}
|
|
if (gradient != nullptr)
|
|
ctx_.FoldRiders(gx);
|
|
*cost = weight_ * xray + restraints;
|
|
if (gradient != nullptr)
|
|
for (size_t k = 0; k < param_atoms_.size(); k++) {
|
|
const int i = param_atoms_[k];
|
|
if (mode_ == Mode::Xyz)
|
|
for (int c = 0; c < 3; c++)
|
|
gradient[3 * k + c] = weight_ * gx[i][c] + (restraint_set_ >= 0 ? gr[i][c] : 0.0);
|
|
else {
|
|
// held at a bound, B does not move further out
|
|
const bool held = (p[k] <= ctx_.b_floor && gx[i][3] > 0) || (p[k] >= ctx_.b_ceiling && gx[i][3] < 0);
|
|
gradient[k] = held ? 0.0 : weight_ * gx[i][3] + (restraint_set_ >= 0 ? gr[i][3] : 0.0);
|
|
}
|
|
}
|
|
return std::isfinite(*cost);
|
|
}
|
|
|
|
private:
|
|
Context &ctx_;
|
|
Mode mode_;
|
|
const std::vector<int> ¶m_atoms_;
|
|
double weight_;
|
|
int restraint_set_;
|
|
};
|
|
|
|
void Minimise(Context &ctx, AtomTarget::Mode mode, const std::vector<int> ¶m_atoms, double weight,
|
|
int restraint_set, int iterations) {
|
|
if (param_atoms.empty())
|
|
return;
|
|
const RefinementStage stage(mode == AtomTarget::Mode::Xyz ? "minimise.xyz" : "minimise.b");
|
|
std::vector<double> p;
|
|
for (int i : param_atoms) {
|
|
const gemmi::Atom &a = *ctx.atoms[i];
|
|
if (mode == AtomTarget::Mode::Xyz) {
|
|
p.push_back(a.pos.x);
|
|
p.push_back(a.pos.y);
|
|
p.push_back(a.pos.z);
|
|
} else
|
|
p.push_back(a.b_iso);
|
|
}
|
|
ceres::GradientProblem problem(std::make_unique<AtomTarget>(ctx, mode, param_atoms, weight, restraint_set));
|
|
ceres::GradientProblemSolver::Summary summary;
|
|
ceres::Solve(SolverOptions(iterations), problem, p.data(), &summary);
|
|
// the atoms where the solver ended, not where its last line search probed
|
|
for (size_t k = 0; k < param_atoms.size(); k++) {
|
|
gemmi::Atom &a = *ctx.atoms[param_atoms[k]];
|
|
if (mode == AtomTarget::Mode::Xyz)
|
|
a.pos = gemmi::Position(p[3 * k], p[3 * k + 1], p[3 * k + 2]);
|
|
else
|
|
a.b_iso = static_cast<float>(std::clamp(p[k], ctx.b_floor, ctx.b_ceiling));
|
|
}
|
|
ctx.MoveRiders();
|
|
}
|
|
|
|
// Groups of atoms that move as rigid bodies (one per chain) against the X-ray target: per body an angle-axis
|
|
// rotation about its centroid, scaled by its rms radius so that a unit moves a typical atom by an angstrom as a
|
|
// unit of translation does, then a translation.
|
|
class RigidBodiesTarget : public ceres::FirstOrderFunction {
|
|
public:
|
|
RigidBodiesTarget(Context &ctx, const std::vector<std::vector<int>> &bodies)
|
|
: ctx_(ctx), bodies_(bodies), base_(ctx.Positions()) {
|
|
for (const std::vector<int> &b : bodies) {
|
|
gemmi::Position c;
|
|
for (int i : b)
|
|
c += base_[i];
|
|
c *= 1.0 / static_cast<double>(b.size());
|
|
double r2 = 0;
|
|
for (int i : b)
|
|
r2 += c.dist_sq(base_[i]);
|
|
centre_.push_back(c);
|
|
radius_.push_back(std::max(1.0, std::sqrt(r2 / static_cast<double>(b.size()))));
|
|
}
|
|
}
|
|
int NumParameters() const override { return static_cast<int>(6 * bodies_.size()); }
|
|
|
|
bool Evaluate(const double *q, double *cost, double *gradient) const override {
|
|
const RefinementStage stage("eval.rigid");
|
|
Place(q);
|
|
ctx_.evaluations++;
|
|
std::vector<std::array<double, 4>> g;
|
|
*cost = ctx_.engine->Evaluate(ctx_.Model(), gradient != nullptr ? &g : nullptr);
|
|
if (gradient == nullptr)
|
|
return std::isfinite(*cost);
|
|
using Jet = ceres::Jet<double, 3>;
|
|
for (size_t b = 0; b < bodies_.size(); b++) {
|
|
const Jet aa[3] = {Jet(q[6 * b] / radius_[b], 0), Jet(q[6 * b + 1] / radius_[b], 1),
|
|
Jet(q[6 * b + 2] / radius_[b], 2)};
|
|
double dq[6] = {0, 0, 0, 0, 0, 0};
|
|
for (int i : bodies_[b]) {
|
|
const gemmi::Position d = base_[i] - centre_[b];
|
|
const Jet p[3] = {Jet(d.x), Jet(d.y), Jet(d.z)};
|
|
Jet rp[3];
|
|
ceres::AngleAxisRotatePoint(aa, p, rp);
|
|
for (int c = 0; c < 3; c++) {
|
|
for (int k = 0; k < 3; k++)
|
|
dq[k] += g[i][c] * rp[c].v[k] / radius_[b];
|
|
dq[3 + c] += g[i][c];
|
|
}
|
|
}
|
|
std::copy(dq, dq + 6, gradient + 6 * b);
|
|
}
|
|
return std::isfinite(*cost);
|
|
}
|
|
|
|
void Place(const double *q) const {
|
|
for (size_t b = 0; b < bodies_.size(); b++) {
|
|
const double aa[3] = {q[6 * b] / radius_[b], q[6 * b + 1] / radius_[b], q[6 * b + 2] / radius_[b]};
|
|
for (int i : bodies_[b]) {
|
|
const gemmi::Position d = base_[i] - centre_[b];
|
|
const double p[3] = {d.x, d.y, d.z};
|
|
double rp[3];
|
|
ceres::AngleAxisRotatePoint(aa, p, rp);
|
|
ctx_.atoms[i]->pos = gemmi::Position(rp[0] + centre_[b].x + q[6 * b + 3], rp[1] + centre_[b].y + q[6 * b + 4],
|
|
rp[2] + centre_[b].z + q[6 * b + 5]);
|
|
}
|
|
}
|
|
}
|
|
|
|
private:
|
|
Context &ctx_;
|
|
const std::vector<std::vector<int>> &bodies_;
|
|
std::vector<gemmi::Position> base_;
|
|
std::vector<gemmi::Position> centre_;
|
|
std::vector<double> radius_;
|
|
};
|
|
|
|
// One B shift per group of atoms (a residue), against the X-ray target.
|
|
class GroupBTarget : public ceres::FirstOrderFunction {
|
|
public:
|
|
GroupBTarget(Context &ctx, const std::vector<std::vector<int>> &groups)
|
|
: ctx_(ctx), groups_(groups), base_(ctx.B()) {}
|
|
int NumParameters() const override { return static_cast<int>(groups_.size()); }
|
|
|
|
bool Evaluate(const double *q, double *cost, double *gradient) const override {
|
|
const RefinementStage stage("eval.group_b");
|
|
Place(q);
|
|
ctx_.evaluations++;
|
|
std::vector<std::array<double, 4>> g;
|
|
*cost = ctx_.engine->Evaluate(ctx_.Model(), gradient != nullptr ? &g : nullptr);
|
|
if (gradient != nullptr)
|
|
for (size_t k = 0; k < groups_.size(); k++) {
|
|
gradient[k] = 0;
|
|
for (int i : groups_[k]) {
|
|
const double b = base_[i] + q[k];
|
|
if (b > ctx_.b_floor && b < ctx_.b_ceiling)
|
|
gradient[k] += g[i][3];
|
|
}
|
|
}
|
|
return std::isfinite(*cost);
|
|
}
|
|
|
|
void Place(const double *q) const {
|
|
for (size_t k = 0; k < groups_.size(); k++)
|
|
for (int i : groups_[k])
|
|
ctx_.atoms[i]->b_iso = static_cast<float>(std::clamp(base_[i] + q[k], ctx_.b_floor, ctx_.b_ceiling));
|
|
}
|
|
|
|
private:
|
|
Context &ctx_;
|
|
const std::vector<std::vector<int>> &groups_;
|
|
std::vector<double> base_;
|
|
};
|
|
|
|
// The placement stage: rigid bodies, then group B.
|
|
void RefinePlacement(Context &ctx, bool group_b, RefinementResult &result) {
|
|
const RefinementStage stage("placement");
|
|
std::vector<std::vector<int>> bodies, groups;
|
|
size_t i = 0;
|
|
for (const gemmi::Chain &ch : ctx.Model().chains) {
|
|
std::vector<int> body;
|
|
for (const gemmi::Residue &r : ch.residues) {
|
|
std::vector<int> group;
|
|
for (size_t k = 0; k < r.atoms.size(); k++, i++) {
|
|
group.push_back(static_cast<int>(i));
|
|
if (!r.is_water())
|
|
body.push_back(static_cast<int>(i));
|
|
}
|
|
groups.push_back(group);
|
|
}
|
|
// A chain of a few atoms (an ion, a ligand of its own chain) stays where it is.
|
|
if (body.size() >= 50)
|
|
bodies.push_back(body);
|
|
}
|
|
// Each step is kept only where it lowers R-free beyond the tie: a deposited model was refined atom by atom
|
|
// already, and moving its chains or shifting its B factors residue by residue then only makes it worse.
|
|
RFactors before = ctx.Refit();
|
|
if (!bodies.empty()) {
|
|
const std::vector<gemmi::Position> pos0 = ctx.Positions();
|
|
std::vector<double> q(6 * bodies.size(), 0.0);
|
|
auto target = std::make_unique<RigidBodiesTarget>(ctx, bodies);
|
|
RigidBodiesTarget *rb = target.get();
|
|
ceres::GradientProblem problem(std::move(target));
|
|
ceres::GradientProblemSolver::Summary summary;
|
|
ceres::Solve(SolverOptions(30), problem, q.data(), &summary);
|
|
rb->Place(q.data());
|
|
const RFactors after = ctx.Refit();
|
|
if (after.r_free < before.r_free - R_FREE_TIE) {
|
|
result.rigid_bodies = static_cast<int>(bodies.size());
|
|
before = after;
|
|
} else {
|
|
ctx.SetPositions(pos0);
|
|
before = ctx.Refit();
|
|
}
|
|
}
|
|
if (group_b) {
|
|
const std::vector<double> b0 = ctx.B();
|
|
std::vector<double> q(groups.size(), 0.0);
|
|
auto target = std::make_unique<GroupBTarget>(ctx, groups);
|
|
GroupBTarget *gb = target.get();
|
|
ceres::GradientProblem problem(std::move(target));
|
|
ceres::GradientProblemSolver::Summary summary;
|
|
ceres::Solve(SolverOptions(30), problem, q.data(), &summary);
|
|
gb->Place(q.data());
|
|
result.group_b = ctx.Refit().r_free < before.r_free - R_FREE_TIE;
|
|
if (!result.group_b)
|
|
ctx.SetB(b0);
|
|
}
|
|
}
|
|
|
|
// The X-ray weight that balances the two gradients: the norm of the geometry's gradient (taken on the model
|
|
// shaken by 0.1 A, where a refined model's geometry would have none) over the norm of the X-ray target's.
|
|
// Following Brunger (1992)-style gradient balancing as phenix.refine starts its weight search (Afonine et al.
|
|
// (2012) Acta Cryst. D68, 352-367)
|
|
double BalancedWeight(Context &ctx, const std::vector<int> &refinable) {
|
|
std::vector<std::array<double, 4>> gx, gr;
|
|
ctx.engine->Evaluate(ctx.Model(), &gx);
|
|
std::vector<std::array<double, 4>> xyzb = ctx.XyzB();
|
|
std::mt19937 rng(1);
|
|
std::normal_distribution<double> shake(0.0, 0.1);
|
|
for (std::array<double, 4> &a : xyzb)
|
|
for (int c = 0; c < 3; c++)
|
|
a[c] += shake(rng);
|
|
ctx.engine->EvaluateRestraints(0, xyzb, &gr);
|
|
double nx = 0, ng = 0;
|
|
for (int i : refinable)
|
|
for (int c = 0; c < 3; c++) {
|
|
nx += gx[i][c] * gx[i][c];
|
|
ng += gr[i][c] * gr[i][c];
|
|
}
|
|
return nx > 0 ? std::sqrt(ng / nx) : 1.0;
|
|
}
|
|
|
|
// The model's map coefficients and its difference map, as model validation makes them (sigma_A weighted).
|
|
void MakeMaps(Context &ctx, const std::vector<MergedReflection> &merged, double wavelength_A, RefinementResult &result) {
|
|
const RefinementStage stage("maps");
|
|
ctx.Refit();
|
|
const std::vector<std::complex<float>> fmodel = ctx.FModel();
|
|
const gemmi::GroupOps gops = ctx.sg->operations();
|
|
std::vector<SigmaAReflection> sa;
|
|
for (size_t m = 0; m < ctx.refl.Size(); m++)
|
|
sa.push_back({ctx.refl.fo[m], std::abs(fmodel[m]), ctx.refl.inv_d2[m], static_cast<int>(ctx.refl.epsilon[m]),
|
|
static_cast<bool>(ctx.refl.centric[m]), static_cast<bool>(ctx.refl.free[m])});
|
|
const SigmaAResult sigma_a = EstimateSigmaA(sa, ctx.cell);
|
|
result.mean_fom = sigma_a.mean_fom;
|
|
|
|
auto maps = std::make_shared<ModelMaps>();
|
|
MtzTable &mtz = maps->coefficients;
|
|
mtz.space_group_name = ctx.sg->xhm();
|
|
mtz.unit_cell = UnitCell{static_cast<float>(ctx.cell.a), static_cast<float>(ctx.cell.b), static_cast<float>(ctx.cell.c),
|
|
static_cast<float>(ctx.cell.alpha), static_cast<float>(ctx.cell.beta),
|
|
static_cast<float>(ctx.cell.gamma)};
|
|
mtz.datasets = {{"HKL_base", "HKL_base", "HKL_base", 0.0f},
|
|
{"refinement", "refinement", "refinement", static_cast<float>(wavelength_A)}};
|
|
const std::pair<const char *, char> columns[] = {{"H", 'H'}, {"K", 'H'}, {"L", 'H'}, {"FP", 'F'}, {"FC", 'F'},
|
|
{"PHIC", 'P'}, {"FWT", 'F'}, {"PHWT", 'P'}, {"DELFWT", 'F'},
|
|
{"PHDELWT", 'P'}, {"FOM", 'W'}, {"FREE", 'I'}};
|
|
for (const auto &[label, type] : columns)
|
|
mtz.columns.push_back({label, type, type == 'H' ? 0 : 1, {}});
|
|
gemmi::AsuData<std::complex<float>> diff;
|
|
diff.unit_cell_ = ctx.cell;
|
|
diff.spacegroup_ = ctx.sg;
|
|
for (size_t m = 0; m < ctx.refl.Size(); m++) {
|
|
const double fo = ctx.refl.fo[m], fc = std::abs(fmodel[m]);
|
|
const double mw = sigma_a.weight[m].m, d = sigma_a.weight[m].d;
|
|
const double fwt = ctx.refl.centric[m] ? mw * fo : 2 * mw * fo - d * fc;
|
|
const double delfwt = mw * fo - d * fc;
|
|
const double phi = std::arg(fmodel[m]);
|
|
double phi_deg = phi * 180.0 / PI;
|
|
if (phi_deg < 0)
|
|
phi_deg += 360.0;
|
|
const float row[] = {static_cast<float>(ctx.refl.hkl[m][0]), static_cast<float>(ctx.refl.hkl[m][1]),
|
|
static_cast<float>(ctx.refl.hkl[m][2]), static_cast<float>(fo), static_cast<float>(fc),
|
|
static_cast<float>(phi_deg), static_cast<float>(fwt), static_cast<float>(phi_deg),
|
|
static_cast<float>(delfwt), static_cast<float>(phi_deg), static_cast<float>(mw),
|
|
ctx.refl.free[m] ? 0.0f : 1.0f};
|
|
for (size_t c = 0; c < mtz.columns.size(); c++)
|
|
mtz.columns[c].values.push_back(row[c]);
|
|
diff.v.push_back({ctx.refl.hkl[m], std::polar(static_cast<float>(delfwt), static_cast<float>(phi))});
|
|
}
|
|
result.maps = maps;
|
|
const std::array<int, 3> size = gemmi::get_size_for_hkl(diff, {{0, 0, 0}}, 3.0);
|
|
result.difference_map = std::make_shared<gemmi::Grid<float>>(MapFromFPhi(gemmi::get_f_phi_on_grid<float>(diff, size, true)));
|
|
result.difference_map_rms = gemmi::calculate_data_statistics(result.difference_map->data).rms;
|
|
|
|
// The anomalous difference map: F(+) - F(-) on the model's phase turned back by 90 degrees, read on the +
|
|
// rows of an anomalous merge only (as model validation reads it).
|
|
// Following ANODE, Thorn & Sheldrick (2011) J. Appl. Cryst. 44, 1285-1287
|
|
const gemmi::ReciprocalAsu asu(ctx.sg);
|
|
std::map<gemmi::Miller, float> danom;
|
|
for (const MergedReflection &r : merged) {
|
|
if (!std::isfinite(r.F_plus) || !std::isfinite(r.F_minus))
|
|
continue;
|
|
const gemmi::Miller h{{r.h, r.k, r.l}};
|
|
if (gops.is_reflection_centric(h))
|
|
continue;
|
|
const auto [hasu, plus] = asu.to_asu_sign(h, gops);
|
|
if (plus)
|
|
danom[hasu] = r.F_plus - r.F_minus;
|
|
}
|
|
if (!danom.empty()) {
|
|
gemmi::AsuData<std::complex<float>> anom;
|
|
anom.unit_cell_ = ctx.cell;
|
|
anom.spacegroup_ = ctx.sg;
|
|
for (size_t m = 0; m < ctx.refl.Size(); m++) {
|
|
const auto it = danom.find(ctx.refl.hkl[m]);
|
|
if (it != danom.end())
|
|
anom.v.push_back({ctx.refl.hkl[m], it->second * std::polar(1.0f, static_cast<float>(std::arg(fmodel[m]) - PI / 2))});
|
|
}
|
|
if (!anom.v.empty()) {
|
|
result.anomalous_map = std::make_shared<gemmi::Grid<float>>(
|
|
MapFromFPhi(gemmi::get_f_phi_on_grid<float>(anom, gemmi::get_size_for_hkl(anom, {{0, 0, 0}}, 3.0), true)));
|
|
result.anomalous_map_rms = gemmi::calculate_data_statistics(result.anomalous_map->data).rms;
|
|
}
|
|
}
|
|
}
|
|
|
|
// The engine for the model as it is now: on the GPU where it is wanted and takes the case within its budget,
|
|
// on the CPU otherwise.
|
|
void MakeEngine(Context &ctx, bool gpu, const RefinementSettings &settings, Logger &logger) {
|
|
const RefinementStage stage("engine");
|
|
ctx.engine.reset();
|
|
#ifdef JFJOCH_USE_CUDA
|
|
ctx.scale_gpu.reset();
|
|
if (gpu && get_gpu_count() > 0) {
|
|
try {
|
|
const int device = get_gpu();
|
|
const size_t card = get_gpu_total_memory(device);
|
|
const size_t usable = card / 5 * 4;
|
|
const size_t budget = usable > settings.gpu_resident_bytes ? usable - settings.gpu_resident_bytes : 0;
|
|
size_t bytes = 0;
|
|
ctx.engine = MakeRefinementEngineGPU(device, ctx.cell, *ctx.sg, ctx.d_min, ctx.refl, ctx.Model(),
|
|
ctx.b_ceiling, budget, &bytes, settings.nthreads);
|
|
ctx.device = device;
|
|
ctx.engine_bytes = bytes;
|
|
if (ctx.engine)
|
|
ctx.scale_gpu = std::make_unique<ModelScaleGPUPool>(device, ctx.refl.Size(), 1);
|
|
else
|
|
logger.Info("Refinement: on the CPU - the GPU cannot take this case or {:.1f} GB is over its budget",
|
|
bytes / 1e9);
|
|
} catch (const JFJochException &e) {
|
|
logger.Warning("Refinement: the GPU failed ({}), on the CPU", e.what());
|
|
ctx.engine.reset();
|
|
ctx.scale_gpu.reset();
|
|
}
|
|
}
|
|
#else
|
|
(void) gpu;
|
|
(void) logger;
|
|
#endif
|
|
if (!ctx.engine)
|
|
ctx.engine = MakeRefinementEngineCPU(ctx.cell, *ctx.sg, ctx.d_min, ctx.refl, settings.nthreads);
|
|
}
|
|
|
|
// The riding hydrogens placed on their parents as they are, and the offsets and B they ride with.
|
|
void PlaceRiders(Context &ctx, RefinementTopology &topo) {
|
|
const RefinementStage stage("riders");
|
|
PlaceRidingHydrogens(topo);
|
|
ctx.rider_offset.assign(ctx.atoms.size(), gemmi::Position());
|
|
for (size_t i = 0; i < ctx.atoms.size(); i++)
|
|
if (ctx.rider_parent[i] >= 0) {
|
|
ctx.rider_offset[i] = ctx.atoms[i]->pos - ctx.atoms[ctx.rider_parent[i]]->pos;
|
|
ctx.atoms[i]->b_iso = ctx.atoms[ctx.rider_parent[i]]->b_iso;
|
|
}
|
|
}
|
|
|
|
// The geometry and the repulsions of the model as it is.
|
|
RestraintSet RestraintsNow(Context &ctx, const RefinementTopology &topo) {
|
|
const RefinementStage stage("repulsions");
|
|
RestraintSet set = topo.geometry;
|
|
set.Append(BuildRepulsions(topo, ctx.Model(), ctx.cell, *ctx.sg, REPULSION_MARGIN, ctx.nthreads));
|
|
set.Finalize(ctx.atoms.size());
|
|
return set;
|
|
}
|
|
|
|
// Where a weight trial runs: a model with its topology and its engine.
|
|
struct TrialSlot {
|
|
Context *ctx;
|
|
RefinementTopology *topo;
|
|
};
|
|
|
|
// A slot of its own beside the main one.
|
|
struct OwnSlot {
|
|
Context ctx;
|
|
RefinementTopology topo;
|
|
};
|
|
|
|
// A copy of the model as it was before the topology was made (the same topology comes out of it, hydrogens
|
|
// included), with an engine of the main one's kind - a GPU's on `device`. Null where the card cannot make one.
|
|
std::unique_ptr<OwnSlot> MakeOwnSlot(const gemmi::Structure &before_topology, const Context &main, int device,
|
|
size_t nthreads, const RefinementSettings &settings) {
|
|
auto slot = std::make_unique<OwnSlot>();
|
|
Context &ctx = slot->ctx;
|
|
ctx.st = before_topology;
|
|
ctx.cell = main.cell;
|
|
ctx.sg = main.sg;
|
|
ctx.d_min = main.d_min;
|
|
ctx.refl = main.refl;
|
|
ctx.nthreads = nthreads;
|
|
ctx.b_ceiling = main.b_ceiling;
|
|
ctx.mask = std::make_unique<BulkSolventMask>(ctx.cell, *ctx.sg, ctx.d_min, nthreads);
|
|
slot->topo = BuildRefinementTopology(ctx.st, settings.riding_hydrogens);
|
|
ctx.atoms = AtomList(ctx.Model());
|
|
ctx.rider_parent = slot->topo.rider_parent;
|
|
#ifdef JFJOCH_USE_CUDA
|
|
if (main.engine->OnGPU()) {
|
|
try {
|
|
ctx.engine = MakeRefinementEngineGPU(device, ctx.cell, *ctx.sg, ctx.d_min, ctx.refl, ctx.Model(),
|
|
ctx.b_ceiling, SIZE_MAX, nullptr, nthreads);
|
|
if (!ctx.engine)
|
|
return nullptr;
|
|
ctx.scale_gpu = std::make_unique<ModelScaleGPUPool>(device, ctx.refl.Size(), 1);
|
|
} catch (const JFJochException &) {
|
|
return nullptr;
|
|
}
|
|
ctx.device = device;
|
|
}
|
|
#else
|
|
(void) device;
|
|
#endif
|
|
if (!ctx.engine)
|
|
ctx.engine = MakeRefinementEngineCPU(ctx.cell, *ctx.sg, ctx.d_min, ctx.refl, nthreads);
|
|
ctx.engine->SetRestraints(1, slot->topo.adp);
|
|
return slot;
|
|
}
|
|
|
|
// The devices of the slots beside the main one, `wanted` of them at most: the GPUs in turn from the main
|
|
// engine's, each taking engines while they fit in its budget. -1 for each one on the CPU.
|
|
std::vector<int> OwnSlotDevices(const Context &main, int wanted, const RefinementSettings &settings) {
|
|
#ifdef JFJOCH_USE_CUDA
|
|
if (main.engine->OnGPU()) {
|
|
const int n = get_gpu_count();
|
|
std::vector<size_t> used(n, 0);
|
|
used[main.device] = main.engine_bytes;
|
|
std::vector<int> devices;
|
|
for (int k = 1; k <= wanted; k++) {
|
|
const int d = (main.device + k) % n;
|
|
const size_t usable = get_gpu_total_memory(d) / 5 * 4;
|
|
const size_t budget = usable > settings.gpu_resident_bytes ? usable - settings.gpu_resident_bytes : 0;
|
|
if (used[d] + main.engine_bytes <= budget) {
|
|
used[d] += main.engine_bytes;
|
|
devices.push_back(d);
|
|
}
|
|
}
|
|
return devices;
|
|
}
|
|
#else
|
|
(void) settings;
|
|
#endif
|
|
return std::vector<int>(std::max(wanted, 0), -1);
|
|
}
|
|
|
|
// One weight trial: a full refinement from pos0 and b0, macro-cycles of the positions at X-ray weight w and,
|
|
// where adp_w is not 0, the B factors at X-ray weight adp_w.
|
|
struct TrialOutcome {
|
|
RefinementWeightTrial trial;
|
|
std::vector<gemmi::Position> pos;
|
|
std::vector<double> b;
|
|
std::string line; // what the log says of it
|
|
};
|
|
|
|
TrialOutcome RunTrial(TrialSlot slot, const std::vector<int> &refinable, const std::vector<gemmi::Position> &pos0,
|
|
const std::vector<double> &b0, double w, double adp_w, const RefinementSettings &settings) {
|
|
Context &ctx = *slot.ctx;
|
|
RefinementTopology &topo = *slot.topo;
|
|
const RefinementStage stage(adp_w > 0 ? "rung.b" : "rung.xyz");
|
|
ctx.SetPositions(pos0);
|
|
ctx.SetB(b0);
|
|
for (int cycle = 0; cycle < settings.macro_cycles; cycle++) {
|
|
PlaceRiders(ctx, topo);
|
|
ctx.Refit();
|
|
{
|
|
const RestraintSet set = RestraintsNow(ctx, topo);
|
|
const RefinementStage s1("set_restraints");
|
|
ctx.engine->SetRestraints(0, set);
|
|
}
|
|
Minimise(ctx, AtomTarget::Mode::Xyz, refinable, w, 0, settings.xyz_iterations);
|
|
if (adp_w > 0)
|
|
Minimise(ctx, AtomTarget::Mode::B, refinable, adp_w, 1, settings.b_iterations);
|
|
}
|
|
PlaceRiders(ctx, topo);
|
|
TrialOutcome out;
|
|
RefinementWeightTrial &t = out.trial;
|
|
t.weight = w;
|
|
t.adp_weight = adp_w;
|
|
const RFactors rf = ctx.Refit();
|
|
t.r_work = rf.r_work;
|
|
t.r_free = rf.r_free;
|
|
{
|
|
const RefinementStage s1("measure");
|
|
t.geometry = MeasureGeometry(topo.geometry,
|
|
BuildRepulsions(topo, ctx.Model(), ctx.cell, *ctx.sg, 0.0, ctx.nthreads), ctx.XyzB());
|
|
}
|
|
t.within_bounds = t.geometry.rmsd_bonds <= MAX_RMSD_BONDS && t.geometry.rmsd_angles <= MAX_RMSD_ANGLES;
|
|
out.line = fmt::format("Refinement, X-ray weight {:.3g}, B {}: R-work {:.4f}, R-free {:.4f}, rmsd bonds {:.4f} A, "
|
|
"angles {:.2f} deg{}",
|
|
w, adp_w > 0 ? fmt::format("refined at weight {:.3g}", adp_w) : std::string("held"),
|
|
rf.r_work, rf.r_free, t.geometry.rmsd_bonds, t.geometry.rmsd_angles,
|
|
t.within_bounds ? "" : " (geometry out of bounds)");
|
|
out.pos = ctx.Positions();
|
|
out.b = ctx.B();
|
|
return out;
|
|
}
|
|
|
|
// The trials (w, adp_w) side by side: trial i in slot i % slots, each slot's trials in turn on a thread of its
|
|
// own (on its GPU), returned in the order given. Every trial starts from pos0 and b0 and takes nothing from the
|
|
// one before it in its slot, so the outcome is the same for any number of slots.
|
|
std::vector<TrialOutcome> RunTrials(const std::vector<TrialSlot> &trial_slots,
|
|
const std::vector<std::pair<double, double>> &trials,
|
|
const std::vector<int> &refinable, const std::vector<gemmi::Position> &pos0,
|
|
const std::vector<double> &b0, const RefinementSettings &settings) {
|
|
std::vector<TrialOutcome> out(trials.size());
|
|
std::vector<std::exception_ptr> errors(trial_slots.size());
|
|
auto run_slot = [&](size_t s) {
|
|
try {
|
|
#ifdef JFJOCH_USE_CUDA
|
|
const ScopedGpuPin pin(trial_slots[s].ctx->device >= 0 ? trial_slots[s].ctx->device : get_gpu());
|
|
#endif
|
|
for (size_t i = s; i < trials.size(); i += trial_slots.size())
|
|
out[i] = RunTrial(trial_slots[s], refinable, pos0, b0, trials[i].first, trials[i].second, settings);
|
|
} catch (...) {
|
|
errors[s] = std::current_exception();
|
|
}
|
|
};
|
|
std::vector<std::thread> threads;
|
|
for (size_t s = 1; s < std::min(trial_slots.size(), trials.size()); s++)
|
|
threads.emplace_back(run_slot, s);
|
|
run_slot(0);
|
|
for (std::thread &t : threads)
|
|
t.join();
|
|
for (const std::exception_ptr &e : errors)
|
|
if (e)
|
|
std::rethrow_exception(e);
|
|
return out;
|
|
}
|
|
|
|
} // namespace
|
|
|
|
RefinementResult RefineModel(const gemmi::Structure &placed, const std::vector<MergedReflection> &merged,
|
|
const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg, double wavelength_A,
|
|
Logger &logger, const RefinementSettings &settings) {
|
|
const auto start = std::chrono::steady_clock::now();
|
|
RefinementResult result;
|
|
Context ctx;
|
|
ctx.st = PrepareModel(placed, cell, sg);
|
|
ctx.cell = cell;
|
|
ctx.sg = &sg;
|
|
ctx.nthreads = settings.nthreads;
|
|
ctx.atoms = AtomList(ctx.Model());
|
|
result.atoms = static_cast<int>(ctx.atoms.size());
|
|
double b_max = 0;
|
|
for (const gemmi::Atom *a : ctx.atoms)
|
|
b_max = std::max<double>(b_max, a->b_iso);
|
|
ctx.b_ceiling = std::clamp(2 * b_max, B_CEILING_MIN, B_CEILING_MAX);
|
|
if (ctx.atoms.empty()) {
|
|
result.failure_reason = "the model has no atoms";
|
|
return result;
|
|
}
|
|
double max_inv_d2 = 0;
|
|
for (const MergedReflection &r : merged)
|
|
if (std::isfinite(r.F) && r.d > 0)
|
|
max_inv_d2 = std::max(max_inv_d2, 1.0 / (double(r.d) * r.d));
|
|
if (max_inv_d2 <= 0) {
|
|
result.failure_reason = "no amplitudes";
|
|
return result;
|
|
}
|
|
// just past the finest reflection, so that prepare_asu_data()'s strict limit keeps it
|
|
ctx.d_min = 0.999 / std::sqrt(max_inv_d2);
|
|
result.d_min = 1.0 / std::sqrt(max_inv_d2);
|
|
ctx.refl = PrepareRefinementReflections(merged, cell, sg, ctx.d_min);
|
|
if (ctx.refl.Size() < 100) {
|
|
result.failure_reason = "too few reflections";
|
|
return result;
|
|
}
|
|
|
|
ctx.mask = std::make_unique<BulkSolventMask>(cell, sg, ctx.d_min, settings.nthreads);
|
|
|
|
// The dictionary's view of the model, and its hydrogens riding on their parents from the start, so that the
|
|
// model is scored as the restrained stage refines it.
|
|
const bool restrained = settings.restrained || settings.score_only;
|
|
// The weight trials' slots start from a copy of the model as it is before the topology is made.
|
|
const gemmi::Structure before_topology = ctx.st;
|
|
RefinementTopology topo;
|
|
if (restrained) {
|
|
std::optional<RefinementStage> topo_stage("topology");
|
|
topo = BuildRefinementTopology(ctx.st, settings.riding_hydrogens);
|
|
topo_stage.reset();
|
|
ctx.atoms = AtomList(ctx.Model());
|
|
ctx.rider_parent = topo.rider_parent;
|
|
PlaceRiders(ctx, topo);
|
|
}
|
|
|
|
MakeEngine(ctx, settings.gpu, settings, logger);
|
|
result.on_gpu = ctx.engine->OnGPU();
|
|
|
|
RFactors r = ctx.Refit();
|
|
result.n_work = r.n_work;
|
|
result.n_free = r.n_free;
|
|
result.r_work_start = r.r_work;
|
|
result.r_free_start = r.r_free;
|
|
result.score_start = ctx.Score();
|
|
logger.Info("Refinement: {} atoms ({} riding hydrogens), {} reflections to {:.2f} A ({} free), on the {}; as given "
|
|
"R-work {:.4f}, R-free {:.4f}; k_sol {:.3f}, b_sol {:.1f}, overall B {:.1f}, mask grid {:.2f} A",
|
|
ctx.atoms.size(), result.score_start.hydrogens, ctx.refl.Size(), result.d_min, r.n_free,
|
|
result.on_gpu ? "GPU" : "CPU", r.r_work, r.r_free, ctx.k_sol, ctx.b_sol, result.score_start.b_iso,
|
|
result.score_start.mask_spacing);
|
|
|
|
// The maps, the blobs and the model of where the refinement ended.
|
|
auto finish = [&]() {
|
|
MakeMaps(ctx, merged, wavelength_A, result);
|
|
{
|
|
const RefinementStage stage("blobs");
|
|
result.blobs = FindDensityBlobs(*result.difference_map, result.difference_map_rms, ctx.st, result.d_min,
|
|
result.anomalous_map.get(), result.anomalous_map_rms);
|
|
}
|
|
result.k_sol = ctx.k_sol;
|
|
result.b_sol = ctx.b_sol;
|
|
result.model = std::make_shared<gemmi::Structure>(ctx.st);
|
|
result.ok = true;
|
|
result.seconds = std::chrono::duration<double>(std::chrono::steady_clock::now() - start).count();
|
|
};
|
|
if (settings.score_only) {
|
|
result.score_start.hkl = ctx.refl.hkl;
|
|
result.score_start.fcalc = ctx.fcalc;
|
|
result.score_start.fmask = ctx.rows.fmask;
|
|
result.score_start.k_total = ctx.rows.k_total;
|
|
result.score_start.d = ctx.rows.d;
|
|
result.score_start.var = ctx.rows.var;
|
|
result.r_work_placement = result.r_work_start;
|
|
result.r_free_placement = result.r_free_start;
|
|
result.restrained_residues = topo.restrained_residues;
|
|
result.fixed_residues = topo.fixed_residues;
|
|
result.fixed_names = topo.fixed_names;
|
|
result.geometry_start = MeasureGeometry(topo.geometry, BuildRepulsions(topo, ctx.Model(), cell, sg, 0.0),
|
|
ctx.XyzB());
|
|
finish();
|
|
logger.Info("Refinement: the model scored only; geometry rmsd bonds {:.4f} A (rmsZ {:.2f}), angles {:.2f} deg "
|
|
"(rmsZ {:.2f})", result.geometry_start.rmsd_bonds, result.geometry_start.rmsz_bonds,
|
|
result.geometry_start.rmsd_angles, result.geometry_start.rmsz_angles);
|
|
return result;
|
|
}
|
|
|
|
RefinePlacement(ctx, settings.group_b, result);
|
|
if (restrained)
|
|
PlaceRiders(ctx, topo);
|
|
r = ctx.Refit();
|
|
result.r_work_placement = r.r_work;
|
|
result.r_free_placement = r.r_free;
|
|
logger.Info("Refinement, placement ({} rigid bodies kept, B per residue {}): R-work {:.4f}, R-free {:.4f}",
|
|
result.rigid_bodies, result.group_b ? "kept" : "not kept", r.r_work, r.r_free);
|
|
|
|
if (settings.restrained) {
|
|
ctx.Refit();
|
|
result.restrained_residues = topo.restrained_residues;
|
|
result.fixed_residues = topo.fixed_residues;
|
|
result.fixed_names = topo.fixed_names;
|
|
std::vector<int> refinable;
|
|
for (size_t i = 0; i < ctx.atoms.size(); i++)
|
|
if (topo.refinable[i])
|
|
refinable.push_back(static_cast<int>(i));
|
|
result.refined_atoms = static_cast<int>(refinable.size());
|
|
logger.Info("Refinement: {} residues restrained, {} held fixed{}; {} bonds, {} angles, {} torsions, {} chiral "
|
|
"volumes, {} planes",
|
|
topo.restrained_residues, topo.fixed_residues,
|
|
topo.fixed_names.empty() ? std::string() : " (" + Joined(topo.fixed_names) + ")",
|
|
topo.bonds, topo.angles, topo.torsions, topo.chirals, topo.planes);
|
|
|
|
ctx.engine->SetRestraints(1, topo.adp);
|
|
{
|
|
const RestraintSet set = RestraintsNow(ctx, topo);
|
|
ctx.engine->SetRestraints(0, set);
|
|
RestraintSet repulsions = BuildRepulsions(topo, ctx.Model(), cell, sg, 0.0, ctx.nthreads);
|
|
result.geometry_start = MeasureGeometry(topo.geometry, repulsions, ctx.XyzB());
|
|
}
|
|
double w0 = 0;
|
|
{
|
|
const RefinementStage stage("balanced_weight");
|
|
w0 = BalancedWeight(ctx, refinable);
|
|
}
|
|
logger.Info("Refinement: geometry as given rmsd bonds {:.4f} A, angles {:.2f} deg, {} clashes; balanced X-ray "
|
|
"weight {:.3g}", result.geometry_start.rmsd_bonds, result.geometry_start.rmsd_angles,
|
|
result.geometry_start.clashes, w0);
|
|
|
|
const std::vector<gemmi::Position> pos0 = ctx.Positions();
|
|
const std::vector<double> b0 = ctx.B();
|
|
|
|
// The trials run side by side, as many at a time as the longer ladder has rungs and the cards take engines:
|
|
// on the main model and engine, and on copies with engines of their own; each with a share of the threads.
|
|
std::optional<RefinementStage> slot_stage("slots");
|
|
const int wanted = static_cast<int>(std::max(settings.weight_ladder.size(), settings.adp_ladder.size()));
|
|
const std::vector<int> devices = OwnSlotDevices(ctx, wanted - 1, settings);
|
|
const size_t slot_threads = std::max<size_t>(1, settings.nthreads / (devices.size() + 1));
|
|
std::vector<std::unique_ptr<OwnSlot>> own(devices.size());
|
|
{
|
|
std::vector<std::thread> threads;
|
|
for (size_t k = 0; k < devices.size(); k++)
|
|
threads.emplace_back([&, k]() {
|
|
#ifdef JFJOCH_USE_CUDA
|
|
const ScopedGpuPin pin(devices[k] >= 0 ? devices[k] : get_gpu());
|
|
#endif
|
|
own[k] = MakeOwnSlot(before_topology, ctx, devices[k], slot_threads, settings);
|
|
});
|
|
for (std::thread &t : threads)
|
|
t.join();
|
|
}
|
|
const size_t main_threads = ctx.nthreads;
|
|
ctx.nthreads = slot_threads;
|
|
std::vector<TrialSlot> trial_slots{{&ctx, &topo}};
|
|
for (const std::unique_ptr<OwnSlot> &o : own)
|
|
if (o)
|
|
trial_slots.push_back({&o->ctx, &o->topo});
|
|
slot_stage.reset();
|
|
logger.Debug("Refinement: {} weight trials at a time", trial_slots.size());
|
|
|
|
int best = -1;
|
|
std::vector<gemmi::Position> best_pos;
|
|
std::vector<double> best_b;
|
|
// A trial is kept where it beats the best so far. The rungs come in order of trust - tighter geometry,
|
|
// then fewer parameters - and a later one displaces an earlier only by more than the tie in R-free.
|
|
auto choose = [&](std::vector<TrialOutcome> &outcomes) {
|
|
for (TrialOutcome &o : outcomes) {
|
|
logger.Info("{}", o.line);
|
|
result.trials.push_back(o.trial);
|
|
const RefinementWeightTrial &t = result.trials.back();
|
|
const int k = static_cast<int>(result.trials.size()) - 1;
|
|
const bool better = best < 0 ? true
|
|
: t.within_bounds && (!result.trials[best].within_bounds ||
|
|
t.r_free < result.trials[best].r_free - R_FREE_TIE);
|
|
if (better) {
|
|
best = k;
|
|
best_pos = std::move(o.pos);
|
|
best_b = std::move(o.b);
|
|
}
|
|
}
|
|
};
|
|
// the positions' weight, with the B factors as they are
|
|
std::vector<std::pair<double, double>> xyz_trials;
|
|
for (double ladder : settings.weight_ladder)
|
|
xyz_trials.emplace_back(w0 * ladder, 0.0);
|
|
std::vector<TrialOutcome> outcomes = RunTrials(trial_slots, xyz_trials, refinable, pos0, b0, settings);
|
|
choose(outcomes);
|
|
// then the B factors refined too, at that weight, against their own weight
|
|
const double w = result.trials[best].weight;
|
|
std::vector<std::pair<double, double>> b_trials;
|
|
for (double adp : settings.adp_ladder)
|
|
b_trials.emplace_back(w, w * adp);
|
|
outcomes = RunTrials(trial_slots, b_trials, refinable, pos0, b0, settings);
|
|
choose(outcomes);
|
|
for (size_t s = 1; s < trial_slots.size(); s++)
|
|
ctx.evaluations += trial_slots[s].ctx->evaluations;
|
|
ctx.nthreads = main_threads;
|
|
ctx.SetPositions(best_pos);
|
|
ctx.SetB(best_b);
|
|
PlaceRiders(ctx, topo);
|
|
const RefinementWeightTrial &chosen = result.trials[best];
|
|
result.restrained = true;
|
|
result.weight = chosen.weight;
|
|
result.adp_weight = chosen.adp_weight;
|
|
result.r_work = chosen.r_work;
|
|
result.r_free = chosen.r_free;
|
|
result.geometry = chosen.geometry;
|
|
}
|
|
|
|
finish();
|
|
logger.Info("Refinement: done in {:.1f} s, {} target evaluations{}", result.seconds, ctx.evaluations,
|
|
result.restrained ? fmt::format("; R-work {:.4f}, R-free {:.4f}, weight {:.3g}", result.r_work,
|
|
result.r_free, result.weight)
|
|
: std::string());
|
|
RefinementStage::Log(logger);
|
|
return result;
|
|
}
|