leonarski_fandClaude Fable 5.1
00f8c8ba4b
The trajectory ends where a step gains nothing, and that step ran from the best state at the weights the final ladder then tries from the same state: on a positional step every rung, on a B step the middle one. A trial of the same weights from the same model is the same computation (RunTrials: the outcome is the same on any slot), so the final takes those outcomes from the step and runs only the rungs it has not got. The maps' Refit() then starts from the fit the main slot would hold had it run its rung (Context::Restore of that rung's fit) - the solvent grid and the likelihood's D it starts from otherwise depend on which trial the main slot ran last, and that showed as last-bit differences in the map coefficients. RunTrials gives the trials that refine B, the longest, to the slots first: on a machine whose cards hold a slot each the ladder's two B rungs no longer share one card while a positional rung has one of its own. Byte-identical (p.mtz, p.hkl, p.cif, p_model, p_refined_model, p_refined_maps.mtz, p_blobs.json, report minus timing) on three rotation sets run live-path, and on two of them in --mode refine. Refinement wall on one RTX 5080: 1.86 -> 1.71 s, 2.42 -> 2.17 s, 50.0 -> 40.8 s (one stage of the trajectory each). Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_015CeEGgpKdfkCTJb3YNMJcu
1520 lines
71 KiB
C++
1520 lines
71 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 <future>
|
|
#include <map>
|
|
#include <optional>
|
|
#include <cmath>
|
|
#include <numeric>
|
|
#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;
|
|
// where the last Refit()'s solvent grid ended (before MoveOverallB()): the next one's fine pass is around it
|
|
// (negative: the next one searches the whole box)
|
|
double grid_k_sol = -1, grid_b_sol = -1;
|
|
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;
|
|
double grid_k_sol = -1, grid_b_sol = -1;
|
|
};
|
|
Fit Fitted() const { return {b_floor, k_sol, b_sol, k_overall, b_star, rows, grid_k_sol, grid_b_sol}; }
|
|
// 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;
|
|
grid_k_sol = fit.grid_k_sol;
|
|
grid_b_sol = fit.grid_b_sol;
|
|
engine->SetRowModel(rows);
|
|
}
|
|
// The next Refit() from nothing an earlier one found: the whole solvent box, the likelihood from the moments.
|
|
void ForgetStarts() {
|
|
grid_k_sol = grid_b_sol = -1;
|
|
rows.bin_d.clear();
|
|
}
|
|
|
|
// 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, grid_k_sol, grid_b_sol);
|
|
else
|
|
#endif
|
|
FitModelScale(scaling, {.b_hi = SolventBMax(Model()), .robust = true, .k_start = grid_k_sol,
|
|
.b_start = grid_b_sol}, nthreads);
|
|
s3.reset();
|
|
k_sol = grid_k_sol = scaling.k_sol;
|
|
b_sol = grid_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: by the whole target, restraints included
|
|
const double g = weight_ * gx[i][3] + gr[i][3];
|
|
const bool held = (p[k] <= ctx_.b_floor && g > 0) || (p[k] >= ctx_.b_ceiling && g < 0);
|
|
gradient[k] = held ? 0.0 : g;
|
|
}
|
|
}
|
|
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, Karplus & Petsko (1989) Acta Cryst. A45, 50-61, 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).
|
|
// Following Read (1986) Acta Cryst. A42, 140-149
|
|
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;
|
|
// The difference map beside the anomalous one below: neither reads the other.
|
|
std::future<void> difference_map = std::async(std::launch::async, [&diff, &result] {
|
|
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;
|
|
}
|
|
}
|
|
difference_map.get();
|
|
}
|
|
|
|
// 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 std::exception &) {
|
|
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, where the main one and
|
|
// those beside it are no more than the threads.
|
|
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();
|
|
// a slot's engine and its scale pool, the robust fit's buffers included
|
|
const size_t slot_bytes = main.engine_bytes + ModelScaleGPUPool::DeviceBytes(main.refl.Size(), 1, true);
|
|
std::vector<size_t> used(n, 0);
|
|
used[main.device] = slot_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] + slot_bytes <= budget) {
|
|
used[d] += slot_bytes;
|
|
devices.push_back(d);
|
|
}
|
|
}
|
|
return devices;
|
|
}
|
|
#endif
|
|
// on the CPU, no more trials at once than threads: each slot is a model, a topology and an engine of its own
|
|
const int beside = std::min(wanted, static_cast<int>(settings.nthreads) - 1);
|
|
return std::vector<int>(std::max(beside, 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 starts from nothing the slot's previous trial left
|
|
if (cycle == 0)
|
|
ctx.ForgetStarts();
|
|
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, the longest to the slots first and 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);
|
|
// The trials that refine B are the longest: they go to the slots first, so that two of them share a slot (or a
|
|
// card) only where the slots are fewer than they are.
|
|
std::vector<size_t> order(trials.size());
|
|
std::iota(order.begin(), order.end(), 0);
|
|
std::stable_sort(order.begin(), order.end(),
|
|
[&](size_t a, size_t b) { return (trials[a].second > 0) > (trials[b].second > 0); });
|
|
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 j = s; j < trials.size(); j += trial_slots.size())
|
|
if (const size_t i = order[j]; !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;
|
|
}
|
|
|
|
// 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. The
|
|
// copies are made while the placement runs: they start from the model as it was before the topology was made
|
|
// and take nothing from the main context that the placement changes.
|
|
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 = settings.restrained ? OwnSlotDevices(ctx, wanted - 1, settings) : std::vector<int>();
|
|
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> slot_builders;
|
|
// joined before the slots are taken, or here where the placement throws
|
|
struct JoinAll {
|
|
std::vector<std::thread> &threads;
|
|
~JoinAll() {
|
|
for (std::thread &t : threads)
|
|
if (t.joinable())
|
|
t.join();
|
|
}
|
|
} join_slot_builders{slot_builders};
|
|
for (size_t k = 0; k < devices.size(); k++)
|
|
slot_builders.emplace_back([&, k]() {
|
|
#ifdef JFJOCH_USE_CUDA
|
|
const ScopedGpuPin pin(devices[k] >= 0 ? devices[k] : get_gpu());
|
|
#endif
|
|
// a slot that cannot be made is left out, as where the card cannot take its engine
|
|
try {
|
|
own[k] = MakeOwnSlot(before_topology, ctx, devices[k], slot_threads, settings);
|
|
} catch (const std::exception &e) {
|
|
logger.Warning("Refinement: a trial slot could not be made ({}), the trials share fewer", e.what());
|
|
}
|
|
});
|
|
|
|
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();
|
|
|
|
std::optional<RefinementStage> slot_stage("slots");
|
|
for (std::thread &t : slot_builders)
|
|
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;
|
|
// The trials of the step that went from the best so far: the final ladder goes from there too, and a trial
|
|
// of the same weights from the same model is the same computation.
|
|
std::vector<TrialOutcome> from_best;
|
|
bool state_is_best = true;
|
|
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);
|
|
if (state_is_best)
|
|
from_best = outcomes;
|
|
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]);
|
|
state_is_best = displaces(state.trial, best.trial);
|
|
if (state_is_best) {
|
|
best = state;
|
|
from_best.clear();
|
|
}
|
|
}
|
|
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. A rung the step from there tried already is taken from that step.
|
|
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.assign(trials.size(), TrialOutcome());
|
|
std::vector<std::pair<double, double>> missing;
|
|
for (size_t i = 0; i < trials.size(); i++) {
|
|
for (const TrialOutcome &o : from_best)
|
|
if (o.ran && o.trial.weight == trials[i].first && o.trial.adp_weight == trials[i].second)
|
|
outcomes[i] = o;
|
|
if (!outcomes[i].ran)
|
|
missing.push_back(trials[i]);
|
|
}
|
|
std::vector<TrialOutcome> ran = RunTrials(trial_slots, missing, refinable, state.pos, state.b, &state.fit,
|
|
settings.macro_cycles, settings);
|
|
for (size_t i = 0, m = 0; i < trials.size(); i++)
|
|
if (!outcomes[i].ran)
|
|
outcomes[i] = std::move(ran[m++]);
|
|
k = choose(outcomes, "final");
|
|
// The main slot's fit as it stands where its last rung ran on it, for the maps' Refit() to start from
|
|
// (the solvent grid and the likelihood's D) whether the rung ran or was taken from the step.
|
|
const size_t main_last = (trials.size() - 1) / trial_slots.size() * trial_slots.size();
|
|
if (outcomes[main_last].ran)
|
|
ctx.Restore(outcomes[main_last].fit);
|
|
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;
|
|
}
|