leonarski_f
ca0a6dd343
Bitwise (the bins are fitted independently: the same bits at any thread count). The weight-trial slots Refit in lockstep while the GPU idles, each with its 1/slots share of -N for the ML bins; the bins now go to the shared worker pool with -N tasks per slot, so the threads of slots that are not fitting help the ones that are. After the Newton searches (previous commit) the ML fit is 0.1-0.5 s summed over a run, so the wall gain is within noise: 5reo 1.63-1.75 -> 1.72, 5f6m 2.60 -> 2.57, 9yl4 7.34 -> 7.28, 7d98 3.78 -> 3.78, 5f6m MD start 4.82 -> 4.76 s. Refined CIF md5 unchanged on 5reo, 5f6m, 9yl4, 7d98, 5f6m MD start (GPU) and 6o2h (CPU). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_015CeEGgpKdfkCTJb3YNMJcu (cherry picked from commit 56815f5390cc01e8da5d420fff73cdf4394c385c)
1450 lines
67 KiB
C++
1450 lines
67 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/it92.hpp"
|
|
#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 "RefinementMinimiser.h"
|
|
#include "RefinementXray.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;
|
|
// The minimiser (RefinementMinimiser): a solve stops where an iteration lowers the target by less than this,
|
|
// relative; the most a position (A) or a B (A^2) moves in the first trial of a line search.
|
|
constexpr double MINIMISER_F_TOLERANCE = 1e-7;
|
|
constexpr double MAX_STEP_XYZ = 0.3, MAX_STEP_B = 10.0;
|
|
|
|
// The figure of merit of the likelihood, <cos(dphi)>, at X = 2 Fo D Fc / var: I1(X)/I0(X) for an acentric
|
|
// reflection, tanh(X / 2) for a centric one.
|
|
// Following Murshudov, Vagin & Dodson (1997) Acta Cryst. D53, 240-255
|
|
double Fom(double x, bool centric) {
|
|
double log_i0 = 0, ratio = 0;
|
|
refine_math::BesselI0I1(x, log_i0, ratio);
|
|
return centric ? std::tanh(0.5 * x) : ratio;
|
|
}
|
|
|
|
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;
|
|
// the likelihood fit's: all of the run's, also in a trial's slot - the slots refit at the same time, the GPU
|
|
// idles meanwhile, and the bins they queue on the shared worker pool are fitted by whichever thread is free
|
|
size_t ml_threads = 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;
|
|
const std::atomic<bool> *cancel = nullptr; // RefinementSettings::cancel
|
|
|
|
bool Cancelled() const { return cancel != nullptr && cancel->load(std::memory_order_relaxed); }
|
|
|
|
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]);
|
|
}
|
|
|
|
// What Refit() fitted, to start another context from without fitting it again.
|
|
struct Fit {
|
|
double b_floor = B_FLOOR;
|
|
double k_sol = 0, b_sol = 0, k_overall = 0;
|
|
gemmi::SMat33<double> b_star{0, 0, 0, 0, 0, 0};
|
|
XrayRowModel rows;
|
|
};
|
|
Fit Fitted() const { return {b_floor, k_sol, b_sol, k_overall, b_star, rows}; }
|
|
// Refit() as it was where `fit` was taken, for the atoms as they were there.
|
|
void Restore(const Fit &fit) {
|
|
b_floor = fit.b_floor;
|
|
engine->SetBlur(RefinementBlur(cell, *sg, d_min, b_floor));
|
|
k_sol = fit.k_sol;
|
|
b_sol = fit.b_sol;
|
|
k_overall = fit.k_overall;
|
|
b_star = fit.b_star;
|
|
rows = fit.rows;
|
|
engine->SetRowModel(rows);
|
|
}
|
|
|
|
// 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");
|
|
// the GPU engine masks on the same grid as BulkSolventMask (bit for bit but at points on an atom's
|
|
// radius, ModelMaskGPU.h)
|
|
if (engine->OnGPU())
|
|
engine->Fmask(Model(), fmask);
|
|
else
|
|
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, ml_threads);
|
|
}
|
|
engine->SetRowModel(rows);
|
|
return ComputeRFactors(refl, fmodel);
|
|
}
|
|
|
|
// The overall isotropic B of the last Refit() moved into the atoms, and into the solvent's B so that F_model stays
|
|
// as it is (but where an atom meets a bound of B): left in the scale, it and the atoms' B drift against each
|
|
// other as the B factors are refined - and the B restraints, in proportion to B, depend on which holds it.
|
|
// As phenix.refine moves the overall B into the atoms (Afonine et al. (2012) Acta Cryst. D68, 352-367)
|
|
void MoveOverallB() {
|
|
const gemmi::Mat33 b = cell.orth.mat.multiply(b_star.as_mat33()).multiply(cell.orth.mat.transpose());
|
|
const double b_iso = (b[0][0] + b[1][1] + b[2][2]) / 3.0;
|
|
double b_min = b_ceiling;
|
|
for (size_t i = 0; i < atoms.size(); i++)
|
|
if (rider_parent.empty() || rider_parent[i] < 0) {
|
|
atoms[i]->b_iso = static_cast<float>(std::clamp(atoms[i]->b_iso + b_iso, B_FLOOR, b_ceiling));
|
|
b_min = std::min<double>(b_min, atoms[i]->b_iso);
|
|
}
|
|
MoveRiders();
|
|
b_floor = std::max(B_FLOOR, b_min - B_STEP_DOWN);
|
|
engine->SetBlur(RefinementBlur(cell, *sg, d_min, b_floor));
|
|
// 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());
|
|
b_star = {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]};
|
|
b_sol += b_iso;
|
|
for (size_t m = 0; m < refl.Size(); m++) {
|
|
rows.k_total[m] *= static_cast<float>(std::exp(b_iso * refl.inv_d2[m] / 4.0));
|
|
rows.fmask[m] *= static_cast<float>(std::exp(-b_iso * refl.inv_d2[m] / 4.0));
|
|
}
|
|
engine->SetRowModel(rows);
|
|
}
|
|
|
|
// 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;
|
|
int fom_n = 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, over the working reflections
|
|
const double x = 2.0 * refl.fo[m] * rows.d[m] * fc / rows.var[m];
|
|
if (!refl.free[m]) {
|
|
fom_sum += Fom(x, refl.centric[m]);
|
|
fom_n++;
|
|
}
|
|
}
|
|
sc.fom = fom_n > 0 ? fom_sum / fom_n : NAN;
|
|
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;
|
|
}
|
|
|
|
// The X-ray target's curvature in each atom's B, up to one factor (CalibrateBCurvature): the likelihood's
|
|
// curvature in |F_model| of the cycle, 2 (k_total D)^2 / var, summed in resolution shells, against the atom's own
|
|
// part of it, (s^2 / 4)^2 f(s)^2 exp(-B s^2 / 2). A hydrogen's is its parent's.
|
|
std::vector<double> XrayBCurvatureShape() const {
|
|
constexpr int SHELLS = 32;
|
|
double s2_max = 0;
|
|
for (double s2 : refl.inv_d2)
|
|
s2_max = std::max(s2_max, s2);
|
|
std::vector<double> w(SHELLS, 0), s2_sum(SHELLS, 0);
|
|
std::vector<int> n(SHELLS, 0);
|
|
for (size_t m = 0; m < refl.Size(); m++) {
|
|
if (refl.free[m])
|
|
continue;
|
|
const int k = std::min(SHELLS - 1, static_cast<int>(refl.inv_d2[m] / s2_max * SHELLS));
|
|
const double kd = rows.k_total[m] * rows.d[m];
|
|
w[k] += 2.0 * kd * kd / rows.var[m];
|
|
s2_sum[k] += refl.inv_d2[m];
|
|
n[k]++;
|
|
}
|
|
std::map<std::pair<int, signed char>, std::vector<double>> f2; // per element and charge, per shell
|
|
std::vector<double> shape(atoms.size(), 0.0);
|
|
for (size_t i = 0; i < atoms.size(); i++) {
|
|
const gemmi::Atom &a = *atoms[i];
|
|
std::vector<double> &f = f2[{static_cast<int>(a.element.elem), a.charge}];
|
|
if (f.empty()) {
|
|
const auto &coef = gemmi::IT92<double>::get(a.element, a.charge, a.serial);
|
|
f.resize(SHELLS, 0);
|
|
for (int k = 0; k < SHELLS; k++)
|
|
if (n[k] > 0)
|
|
f[k] = std::pow(coef.calculate_sf(s2_sum[k] / n[k] / 4.0), 2);
|
|
}
|
|
for (int k = 0; k < SHELLS; k++)
|
|
if (n[k] > 0) {
|
|
const double s2 = s2_sum[k] / n[k];
|
|
shape[i] += w[k] * f[k] * a.occ * a.occ * std::exp(-a.b_iso * s2 / 2.0) * s2 * s2 / 16.0;
|
|
}
|
|
}
|
|
for (size_t i = 0; i < rider_parent.size(); i++)
|
|
if (rider_parent[i] >= 0)
|
|
shape[rider_parent[i]] += shape[i];
|
|
return shape;
|
|
}
|
|
|
|
// The factor that turns XrayBCurvatureShape() into the X-ray target's curvature: its second derivative along a
|
|
// random +-1 direction of the B factors of `params`, by central differences of the gradient, over what the shape
|
|
// predicts for it.
|
|
double CalibrateBCurvature(const std::vector<int> ¶ms, const std::vector<double> &shape) {
|
|
constexpr double EPS = 0.5; // A^2
|
|
const std::vector<double> b0 = B();
|
|
std::mt19937 rng(7);
|
|
std::vector<double> v(params.size());
|
|
for (double &x : v)
|
|
x = rng() % 2 == 0 ? 1.0 : -1.0;
|
|
std::array<std::vector<std::array<double, 4>>, 2> g;
|
|
for (int side = 0; side < 2; side++) {
|
|
for (size_t k = 0; k < params.size(); k++)
|
|
atoms[params[k]]->b_iso = static_cast<float>(b0[params[k]] + (side == 0 ? EPS : -EPS) * v[k]);
|
|
MoveRiders();
|
|
engine->Evaluate(Model(), &g[side]);
|
|
FoldRiders(g[side]);
|
|
}
|
|
SetB(b0);
|
|
MoveRiders();
|
|
double num = 0, den = 0;
|
|
for (size_t k = 0; k < params.size(); k++) {
|
|
num += v[k] * (g[0][params[k]][3] - g[1][params[k]][3]) / (2 * EPS);
|
|
den += shape[params[k]];
|
|
}
|
|
return den > 0 ? std::max(num / den, 0.0) : 0.0;
|
|
}
|
|
|
|
// 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;
|
|
}
|
|
|
|
// SolverOptions(iterations) that stop at the next iteration, where the minimiser stands, once ctx is cancelled.
|
|
class CancelCheck : public ceres::IterationCallback {
|
|
public:
|
|
explicit CancelCheck(const Context &ctx) : ctx_(ctx) {}
|
|
ceres::CallbackReturnType operator()(const ceres::IterationSummary &) override {
|
|
return ctx_.Cancelled() ? ceres::SOLVER_TERMINATE_SUCCESSFULLY : ceres::SOLVER_CONTINUE;
|
|
}
|
|
ceres::GradientProblemSolver::Options Options(int iterations) const {
|
|
ceres::GradientProblemSolver::Options options = SolverOptions(iterations);
|
|
options.callbacks.push_back(const_cast<CancelCheck *>(this));
|
|
return options;
|
|
}
|
|
|
|
private:
|
|
const Context &ctx_;
|
|
};
|
|
|
|
// The positions or the B factors of the refinable atoms, against w times the X-ray target plus a restraint set
|
|
// (minimised by RefinementMinimiser).
|
|
class AtomTarget {
|
|
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) {}
|
|
|
|
// The parameters of the atoms as they are: x y z per atom, or B per atom.
|
|
std::vector<double> Get() const {
|
|
std::vector<double> p;
|
|
for (int i : param_atoms_) {
|
|
const gemmi::Atom &a = *ctx_.atoms[i];
|
|
if (mode_ == 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);
|
|
}
|
|
return p;
|
|
}
|
|
// The atoms set to the parameters, their riders with them.
|
|
void Set(const double *p) const {
|
|
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();
|
|
}
|
|
// the most a parameter moves in the first trial of a line search
|
|
double MaxStep() const { return mode_ == Mode::Xyz ? MAX_STEP_XYZ : MAX_STEP_B; }
|
|
|
|
bool Evaluate(const double *p, double *cost, double *gradient) const {
|
|
const RefinementStage stage(mode_ == Mode::Xyz ? "eval.xyz" : "eval.b");
|
|
Set(p);
|
|
ctx_.evaluations++;
|
|
std::vector<std::array<double, 4>> gx, gr;
|
|
double xray = 0, restraints = 0;
|
|
{
|
|
const RefinementStage s1("eval.engine");
|
|
xray = ctx_.engine->EvaluateWithRestraints(ctx_.Model(), restraint_set_, gradient != nullptr ? &gx : nullptr,
|
|
gradient != nullptr ? &gr : nullptr, restraints);
|
|
}
|
|
*cost = weight_ * xray + restraints;
|
|
if (gradient == nullptr)
|
|
return std::isfinite(*cost);
|
|
ctx_.FoldRiders(gx);
|
|
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] + gr[i][c];
|
|
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] + gr[i][3];
|
|
}
|
|
}
|
|
return std::isfinite(*cost);
|
|
}
|
|
|
|
private:
|
|
Context &ctx_;
|
|
Mode mode_;
|
|
const std::vector<int> ¶m_atoms_;
|
|
double weight_;
|
|
int restraint_set_;
|
|
};
|
|
|
|
// The minimiser's diagonal for the B factors of `param_atoms` (RefinementMinimiser::SetDiagonal): weight times the
|
|
// X-ray target's curvature in B plus the B restraints', each element at least a thousandth of the median. The
|
|
// positions take none: their restraints' curvature is that of stiff bonds and angles, which move atoms together, so
|
|
// as a diagonal it slows rather than scales the minimiser.
|
|
std::vector<double> BDiagonal(const std::vector<double> &xray, const std::vector<double> &restraints,
|
|
const std::vector<int> ¶m_atoms, double weight) {
|
|
std::vector<double> h;
|
|
for (int i : param_atoms)
|
|
h.push_back(weight * xray[i] + restraints[i]);
|
|
std::vector<double> sorted = h;
|
|
std::nth_element(sorted.begin(), sorted.begin() + sorted.size() / 2, sorted.end());
|
|
const double floor = std::max(1e-3 * sorted[sorted.size() / 2], 1e-12);
|
|
for (double &v : h)
|
|
v = std::max(v, floor);
|
|
return h;
|
|
}
|
|
|
|
// `minimiser` carries its memory from the call before; `diagonal` is the curvature it scales the steps by (none: 1).
|
|
void Minimise(Context &ctx, const AtomTarget &target, int iterations, RefinementMinimiser &minimiser,
|
|
std::vector<double> diagonal = {}) {
|
|
std::vector<double> p = target.Get();
|
|
if (p.empty())
|
|
return;
|
|
minimiser.SetDiagonal(diagonal.empty() ? std::vector<double>(p.size(), 1.0) : std::move(diagonal));
|
|
minimiser.Minimise([&](const std::vector<double> &x, double &f, std::vector<double> &g) {
|
|
return target.Evaluate(x.data(), &f, g.data());
|
|
},
|
|
p, iterations, MINIMISER_F_TOLERANCE, target.MaxStep(), [&]() { return ctx.Cancelled(); });
|
|
target.Set(p.data()); // where the minimiser ended, not where its last line search probed
|
|
}
|
|
|
|
// 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;
|
|
const CancelCheck cancel_check(ctx);
|
|
ceres::Solve(cancel_check.Options(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;
|
|
const CancelCheck cancel_check(ctx);
|
|
ceres::Solve(cancel_check.Options(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 B factors' X-ray weight that balances the two gradients, as BalancedWeight() does for the positions: the B
|
|
// similarity's gradient on the model with its B shaken by 20% (a flat B, as a model may come, gives none) over the
|
|
// X-ray target's.
|
|
double BalancedADPWeight(Context &ctx, const RestraintSet &b_restraints, const std::vector<int> &refinable) {
|
|
std::vector<std::array<double, 4>> gx;
|
|
ctx.engine->Evaluate(ctx.Model(), &gx);
|
|
ctx.FoldRiders(gx);
|
|
std::vector<std::array<double, 4>> xyzb = ctx.XyzB();
|
|
std::mt19937 rng(2);
|
|
std::normal_distribution<double> shake(0.0, 0.2);
|
|
for (std::array<double, 4> &a : xyzb)
|
|
a[3] *= std::exp(shake(rng));
|
|
std::vector<std::array<double, 4>> gr(xyzb.size(), {0, 0, 0, 0});
|
|
EvaluateRestraints(b_restraints, xyzb, &gr);
|
|
double nx = 0, ng = 0;
|
|
for (int i : refinable) {
|
|
nx += gx[i][3] * gx[i][3];
|
|
ng += gr[i][3] * gr[i][3];
|
|
}
|
|
return nx > 0 ? std::sqrt(ng / nx) : 1.0;
|
|
}
|
|
|
|
// The model's map coefficients and its difference map, weighted by the refinement's own likelihood: m its figure
|
|
// of merit and D its D (2mFo - DFc, mFo - DFc).
|
|
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<double> fom(ctx.refl.Size());
|
|
double fom_sum = 0;
|
|
int fom_n = 0;
|
|
for (size_t m = 0; m < ctx.refl.Size(); m++) {
|
|
const double x = 2.0 * ctx.refl.fo[m] * ctx.rows.d[m] * std::abs(fmodel[m]) / ctx.rows.var[m];
|
|
fom[m] = Fom(x, ctx.refl.centric[m]);
|
|
if (!ctx.refl.free[m]) {
|
|
fom_sum += fom[m];
|
|
fom_n++;
|
|
}
|
|
}
|
|
result.mean_fom = fom_n > 0 ? fom_sum / fom_n : 0;
|
|
|
|
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 = fom[m], d = ctx.rows.d[m];
|
|
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(topo.metal_bonds);
|
|
set.Append(BuildRepulsions(topo, ctx.Model(), ctx.cell, *ctx.sg, REPULSION_MARGIN, ctx.nthreads));
|
|
set.Finalize(ctx.atoms.size());
|
|
return set;
|
|
}
|
|
|
|
// The B similarity of the model as it is: bonded and angle-related pairs, and atoms in contact.
|
|
RestraintSet BRestraintsNow(Context &ctx, const RefinementTopology &topo) {
|
|
const RefinementStage stage("b_restraints");
|
|
RestraintSet set = topo.adp;
|
|
set.Append(BuildBContacts(topo, ctx.Model(), ctx.cell, *ctx.sg, 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.ml_threads = settings.nthreads;
|
|
ctx.cancel = main.cancel;
|
|
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: `cycles` macro-cycles from pos0 and b0, of the positions at X-ray weight w and, where adp_w is
|
|
// not 0, the B factors at X-ray weight adp_w. `fit`: what was fitted at pos0 and b0 already, null to fit it here.
|
|
struct TrialOutcome {
|
|
RefinementWeightTrial trial;
|
|
std::vector<gemmi::Position> pos;
|
|
std::vector<double> b;
|
|
Context::Fit fit; // the scale, the solvent and the likelihood's parameters fitted where it ended
|
|
std::string line; // what the log says of it
|
|
bool ran = false; // false: left out, the refinement cancelled before it
|
|
};
|
|
|
|
TrialOutcome RunTrial(TrialSlot slot, const std::vector<int> &refinable, const std::vector<gemmi::Position> &pos0,
|
|
const std::vector<double> &b0, const Context::Fit *fit, double w, double adp_w, int cycles,
|
|
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);
|
|
// the positions and the B factors in turn, each with a memory of its own kept from macro-cycle to macro-cycle
|
|
const AtomTarget xyz(ctx, AtomTarget::Mode::Xyz, refinable, w, 0), b(ctx, AtomTarget::Mode::B, refinable, adp_w, 1);
|
|
RefinementMinimiser xyz_minimiser, b_minimiser;
|
|
double xray_scale = 0;
|
|
for (int cycle = 0; cycle < cycles && !ctx.Cancelled(); cycle++) {
|
|
PlaceRiders(ctx, topo);
|
|
if (cycle == 0 && fit != nullptr)
|
|
ctx.Restore(*fit);
|
|
else {
|
|
// a trial's likelihood fit starts from nothing the slot's previous trial left
|
|
if (cycle == 0)
|
|
ctx.rows.bin_d.clear();
|
|
ctx.Refit();
|
|
}
|
|
if (adp_w > 0)
|
|
ctx.MoveOverallB();
|
|
{
|
|
const RestraintSet set = RestraintsNow(ctx, topo);
|
|
const RefinementStage s1("set_restraints");
|
|
ctx.engine->SetRestraints(0, set);
|
|
}
|
|
{
|
|
const RefinementStage s2("minimise.xyz");
|
|
Minimise(ctx, xyz, settings.xyz_iterations, xyz_minimiser);
|
|
}
|
|
if (adp_w > 0) {
|
|
const RestraintSet b_set = BRestraintsNow(ctx, topo);
|
|
ctx.engine->SetRestraints(1, b_set);
|
|
std::vector<double> xray = ctx.XrayBCurvatureShape();
|
|
if (cycle == 0)
|
|
xray_scale = ctx.CalibrateBCurvature(refinable, xray);
|
|
for (double &x : xray)
|
|
x *= xray_scale;
|
|
const RefinementStage s3("minimise.b");
|
|
Minimise(ctx, b, settings.b_iterations, b_minimiser,
|
|
BDiagonal(xray, BRestraintCurvature(b_set, ctx.XyzB()), refinable, adp_w));
|
|
}
|
|
}
|
|
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(),
|
|
&topo.metal_bonds);
|
|
}
|
|
t.within_bounds = t.geometry.rmsd_bonds <= MAX_RMSD_BONDS && t.geometry.rmsd_angles <= MAX_RMSD_ANGLES;
|
|
out.line = fmt::format("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.geometry.metal_bonds > 0 ? fmt::format(", metal bonds {:.3f} A", t.geometry.rmsd_metal)
|
|
: std::string(),
|
|
t.within_bounds ? "" : " (geometry out of bounds)");
|
|
out.pos = ctx.Positions();
|
|
out.b = ctx.B();
|
|
out.fit = ctx.Fitted();
|
|
out.ran = true;
|
|
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 Context::Fit *fit, int cycles,
|
|
const RefinementSettings &settings) {
|
|
std::vector<TrialOutcome> out(trials.size());
|
|
std::vector<std::exception_ptr> errors(trial_slots.size());
|
|
// the threads shared among the slots that run (the outcome does not depend on how many each has)
|
|
const size_t running = std::min(trial_slots.size(), trials.size());
|
|
for (size_t s = 0; s < running; s++)
|
|
trial_slots[s].ctx->nthreads = std::max<size_t>(1, settings.nthreads / running);
|
|
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())
|
|
if (!trial_slots[s].ctx->Cancelled())
|
|
out[i] = RunTrial(trial_slots[s], refinable, pos0, b0, fit, trials[i].first, trials[i].second,
|
|
cycles, 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.ml_threads = settings.nthreads;
|
|
ctx.cancel = settings.cancel;
|
|
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(), &topo.metal_bonds);
|
|
finish();
|
|
const GeometryStatistics &g = result.geometry_start;
|
|
logger.Info("Refinement: the model scored only; geometry rmsd bonds {:.4f} A (rmsZ {:.2f}, {} bonds), metal "
|
|
"bonds {:.4f} A (rmsZ {:.2f}, {}; {} links made), both {:.4f} A; angles {:.2f} deg (rmsZ {:.2f})",
|
|
g.rmsd_bonds, g.rmsz_bonds, g.bonds, g.rmsd_metal, g.rmsz_metal, g.metal_bonds, topo.metal_links,
|
|
g.RmsdAllBonds(), g.rmsd_angles, g.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, {} metal bonds ({} links made), {} "
|
|
"angles, {} torsions, {} chiral volumes, {} planes",
|
|
topo.restrained_residues, topo.fixed_residues,
|
|
topo.fixed_names.empty() ? std::string() : " (" + Joined(topo.fixed_names) + ")",
|
|
topo.bonds, topo.metal_bonds.records.size(), topo.metal_links, 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(), &topo.metal_bonds);
|
|
}
|
|
double w0 = 0, adp_w0 = 0;
|
|
{
|
|
const RefinementStage stage("balanced_weight");
|
|
w0 = BalancedWeight(ctx, refinable);
|
|
adp_w0 = BalancedADPWeight(ctx, BRestraintsNow(ctx, topo), refinable);
|
|
}
|
|
logger.Info("Refinement: geometry as given rmsd bonds {:.4f} A, angles {:.2f} deg, {} clashes; balanced X-ray "
|
|
"weights {:.3g} (positions), {:.3g} (B)", result.geometry_start.rmsd_bonds,
|
|
result.geometry_start.rmsd_angles, result.geometry_start.clashes, w0, adp_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 longest set of them has 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");
|
|
int wanted = static_cast<int>(std::max({settings.weight_ladder.size() * settings.adp_ladder.size(),
|
|
settings.final_ladder.size(), size_t{3}}));
|
|
if (settings.max_trials_at_once > 0)
|
|
wanted = std::min(wanted, static_cast<int>(settings.max_trials_at_once));
|
|
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());
|
|
|
|
// A trial displaces another only where its geometry is within bounds and its R-free lower by more than the
|
|
// tie.
|
|
auto displaces = [](const RefinementWeightTrial &t, const RefinementWeightTrial &b) {
|
|
return t.within_bounds && (!b.within_bounds || t.r_free < b.r_free - R_FREE_TIE);
|
|
};
|
|
// Of a set of trials, the one kept: they come in order of trust, the first kept unless a later one displaces
|
|
// it. -1: none ran.
|
|
auto choose = [&](std::vector<TrialOutcome> &outcomes, const char *stage) {
|
|
int best = -1;
|
|
for (size_t k = 0; k < outcomes.size(); k++) {
|
|
const TrialOutcome &o = outcomes[k];
|
|
if (!o.ran)
|
|
continue;
|
|
logger.Info("Refinement, {}: {}", stage, o.line);
|
|
result.trials.push_back(o.trial);
|
|
if (best < 0 || displaces(o.trial, outcomes[best].trial))
|
|
best = static_cast<int>(k);
|
|
}
|
|
return best;
|
|
};
|
|
const double w_max = w0 * settings.max_weight;
|
|
|
|
// 1. The weights from the placed model: the ladder's rungs, tighter geometry first, each with the B factors
|
|
// held and refined at their own balanced weight.
|
|
std::vector<std::pair<double, double>> ladder;
|
|
for (double rung : settings.weight_ladder)
|
|
for (double adp : settings.adp_ladder)
|
|
ladder.emplace_back(w0 * rung, adp_w0 * adp);
|
|
std::vector<TrialOutcome> outcomes = RunTrials(trial_slots, ladder, refinable, pos0, b0, nullptr,
|
|
settings.macro_cycles, settings);
|
|
int k = choose(outcomes, "ladder");
|
|
TrialOutcome state;
|
|
if (k >= 0)
|
|
state = std::move(outcomes[k]);
|
|
|
|
// 2. The trajectory: steps of a few macro-cycles from where the last one ended, each trying the weights as they
|
|
// are and, in turn, the positions' X-ray weight halved and doubled (halved: the geometry loosens as the
|
|
// model improves at one weight, until it leaves the bounds) and the B factors' divided and multiplied by
|
|
// four.
|
|
// A step is taken while it lowers R-free or R-work by more than the tie and its R-free stays within the tie
|
|
// of the best so far; the trajectory ends at its best step (on a deposited model, whose free set it was
|
|
// fitted to, R-work falls while R-free rises, and it ends where it began).
|
|
TrialOutcome best = state;
|
|
for (int step = 0; step < settings.max_steps && state.ran && !ctx.Cancelled(); step++) {
|
|
const double w = state.trial.weight, adp = state.trial.adp_weight;
|
|
std::vector<std::pair<double, double>> trials = {{w, adp}};
|
|
if (step % 2 == 1 && adp > 0) {
|
|
trials.emplace_back(w, adp / 4);
|
|
trials.emplace_back(w, 4 * adp);
|
|
} else {
|
|
trials.emplace_back(w / 2, adp);
|
|
if (2 * w <= w_max)
|
|
trials.emplace_back(2 * w, adp);
|
|
}
|
|
outcomes = RunTrials(trial_slots, trials, refinable, state.pos, state.b, &state.fit, settings.macro_cycles,
|
|
settings);
|
|
k = choose(outcomes, fmt::format("step {}", step + 1).c_str());
|
|
if (k < 0)
|
|
break;
|
|
const RefinementWeightTrial &t = outcomes[k].trial;
|
|
const bool gains = displaces(t, state.trial) || (t.within_bounds && t.r_work < state.trial.r_work - R_FREE_TIE);
|
|
if (!gains || !(t.within_bounds && t.r_free < best.trial.r_free + R_FREE_TIE))
|
|
break;
|
|
state = std::move(outcomes[k]);
|
|
if (displaces(state.trial, best.trial))
|
|
best = state;
|
|
}
|
|
state = std::move(best);
|
|
|
|
// 3. The final weights: a short ladder from where the trajectory ended, tighter geometry first, its choice
|
|
// kept where it displaces that.
|
|
if (state.ran && !ctx.Cancelled()) {
|
|
const double w = state.trial.weight, adp = state.trial.adp_weight;
|
|
std::vector<std::pair<double, double>> trials;
|
|
for (double f : settings.final_ladder)
|
|
if (w * f <= w_max)
|
|
trials.emplace_back(w * f, adp);
|
|
outcomes = RunTrials(trial_slots, trials, refinable, state.pos, state.b, &state.fit, settings.macro_cycles,
|
|
settings);
|
|
k = choose(outcomes, "final");
|
|
if (k >= 0 && displaces(outcomes[k].trial, state.trial))
|
|
state = std::move(outcomes[k]);
|
|
}
|
|
for (size_t s = 1; s < trial_slots.size(); s++)
|
|
ctx.evaluations += trial_slots[s].ctx->evaluations;
|
|
ctx.nthreads = main_threads;
|
|
if (state.ran) {
|
|
ctx.SetPositions(state.pos);
|
|
ctx.SetB(state.b);
|
|
PlaceRiders(ctx, topo);
|
|
result.restrained = true;
|
|
result.weight = state.trial.weight;
|
|
result.adp_weight = state.trial.adp_weight;
|
|
result.r_work = state.trial.r_work;
|
|
result.r_free = state.trial.r_free;
|
|
result.geometry = state.trial.geometry;
|
|
} else
|
|
ctx.Refit(); // no trial ran: the model as placed, its scale and solvent for the maps
|
|
}
|
|
result.cancelled = ctx.Cancelled();
|
|
if (result.cancelled)
|
|
logger.Warning("Refinement: cancelled, {} weight trials ran", result.trials.size());
|
|
|
|
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;
|
|
}
|