mx/Jungfraujoch
3
2 Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases 88 Wiki Activity
Files
rc181
Jungfraujoch/image_analysis/structure_refinement/Refinement.cpp
T
leonarski_fandClaude Opus 5.5 b822a775ac rugnux: solvent build optional, off by default (--refine-solvent / --no-refine-solvent) 
The refinement's solvent step (waters built and removed, ions relabelled, and the
solvent audit) now runs only with --refine-solvent; --no-refine-solvent is the
default. When off, no audit is made, no <prefix>_solvent.json is written, and the
report carries REFINE_SOLVENT=NOT_PERFORMED in place of the REFINE_SOLVENT_ keys.
RefinementSettings::solvent and RugnuxConfig::refine_solvent default to false, so
the broker's live rugnux follows; RefinementResult::solvent becomes optional.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01HVeDQqrcrEZy6XBUCMZvyT
2026-11-01 22:07:04 +01:00

3239 lines
167 KiB
C++
Raw Permalink Blame History

// 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 <atomic>
#include <chrono>
#include <set>
#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 "../../common/ParallelFor.h"
#include "ModelFFT.h"
#include "ModelGrid.h" // IonFormFactor, AnisoRadiusB
#include "ModelScaling.h"
#include "MonomerLibrary.h"
#include "RefinementADP.h"
#include "RefinementMath.h"
#include "RefinementMinimiser.h"
#include "RefinementNcs.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;
// The minimiser's iterations per macro-cycle where the solvent build refines the solvent alone.
constexpr int SOLVENT_ITERATIONS = 10;
// 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 loosest NCS restraint a trajectory step tries (A); looser than that, the copies are refined independently.
constexpr double NCS_SIGMA_MAX = 0.4;
// 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 charge an atom scatters with (IonFormFactor): a single-atom ion's, a metal or a halide, as its dictionary gives
// it (MG 2+, NA 1+, CL 1-); 0 for every other atom, whatever the file says (a carboxylate's O1-, a lysine's N1+).
signed char ScatteringCharge(const gemmi::Residue &r, const gemmi::Atom &a) {
if (!(a.element.is_metal() || a.element == gemmi::El::F || a.element == gemmi::El::Cl ||
a.element == gemmi::El::Br || a.element == gemmi::El::I))
return 0;
const auto cc = EmbeddedMonomerLibrary().monomers.find(r.name);
if (cc == EmbeddedMonomerLibrary().monomers.end() || cc->second.atoms.size() != 1)
return 0;
return static_cast<signed char>(std::lround(cc->second.atoms[0].charge));
}
// The model as refinement takes it: the first model only, without hydrogens and atoms of unknown element, every B
// inside the allowed range - an anisotropic atom's equivalent isotropic B, its U kept (Context::TakeAniso() takes it) -
// and the single-atom ions charged and nothing else (ScatteringCharge).
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) {
a.charge = ScatteringCharge(r, a);
if (a.aniso.nonzero())
a.b_iso = static_cast<float>(gemmi::u_to_b() * a.aniso.trace() / 3.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::shared_ptr<ModelScaleGPUPool> scale_gpu; // shared by the slots on the engine's card
#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, beside what it shares with the engines like it
size_t card_bytes = 0; // on the same card, and a lane of working buffers (RefinementGPUBytes)
size_t lane_bytes = 0;
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<int> rider_start, rider_index; // SetRiderParents()
std::vector<gemmi::Position> rider_offset;
int evaluations = 0;
const std::atomic<bool> *cancel = nullptr; // RefinementSettings::cancel
// The anisotropic atoms (empty: none), which the engine reads (RefinementEngine::SetAniso), and per atom how far
// its B stays inside the bounds so that its U does: its smallest eigenvalue above b_floor, its largest below
// b_ceiling (in B, A^2).
std::vector<AtomAniso> aniso;
std::vector<double> b_lo, b_hi;
double BLow(size_t i) const { return b_floor + (b_lo.empty() ? 0.0 : b_lo[i]); }
double BHigh(size_t i) const { return std::max(BLow(i), b_ceiling - (b_hi.empty() ? 0.0 : b_hi[i])); }
// The anisotropy of `a` (empty: none) for the atoms as they are, and the bounds of their B with it.
void SetAniso(const std::vector<AtomAniso> &a) {
aniso = a;
b_lo.clear();
b_hi.clear();
if (aniso.empty())
return;
b_lo.assign(atoms.size(), 0.0);
b_hi.assign(atoms.size(), 0.0);
for (size_t i = 0; i < aniso.size(); i++)
if (aniso[i].on) {
const std::array<float, 6> &d = aniso[i].dev;
const gemmi::SMat33<double> dev{d[0], d[1], d[2], d[3], d[4], d[5]};
const std::array<double, 3> ev = dev.calculate_eigenvalues();
b_lo[i] = std::max(0.0, -gemmi::u_to_b() * std::min({ev[0], ev[1], ev[2]}));
b_hi[i] = std::max(0.0, gemmi::u_to_b() * std::max({ev[0], ev[1], ev[2]}));
}
}
// TLS bodies (empty: none) and per atom the part of its B they hold as they were last applied (8 pi^2 tr(U_TLS) / 3,
// and what bringing U inside the bounds added): the rest of its B is its residual.
std::vector<TlsBody> tls;
std::vector<double> tls_b;
// The bodies' U at the atoms as they are now, each atom's residual B kept: its anisotropy the traceless part of
// U_TLS, its B the residual plus the trace.
void ApplyTls() {
if (tls.empty())
return;
if (aniso.size() != atoms.size())
aniso.assign(atoms.size(), AtomAniso());
tls_b.resize(atoms.size(), 0.0);
for (const TlsBody &body : tls)
for (int i : body.atoms) {
const gemmi::SMat33<double> u = body.U(atoms[i]->pos);
double t = u.trace() / 3.0;
const double residual = atoms[i]->b_iso - tls_b[i];
// the total U inside the bounds: its largest eigenvalue (in B) at most the ceiling, its smallest at least
// the floor - the whole U scaled down, or raised by a multiple of I
gemmi::SMat33<double> b = u.scaled(gemmi::u_to_b()).added_kI(residual);
std::array<double, 3> ev = b.calculate_eigenvalues();
const double top = std::max({ev[0], ev[1], ev[2]});
if (top > b_ceiling) {
b = b.scaled(b_ceiling / top);
ev = b.calculate_eigenvalues();
}
const double lowest = std::min({ev[0], ev[1], ev[2]});
if (!(lowest >= b_floor))
b = b.added_kI(b_floor - (std::isfinite(lowest) ? lowest : 0.0));
t = b.trace() / 3.0;
atoms[i]->b_iso = static_cast<float>(t);
// what the bodies hold of B, the bounds' share included: the residual stays as it was
tls_b[i] = atoms[i]->b_iso - residual;
aniso[i].on = true;
aniso[i].dev = {static_cast<float>((b.u11 - t) / gemmi::u_to_b()), static_cast<float>((b.u22 - t) / gemmi::u_to_b()),
static_cast<float>((b.u33 - t) / gemmi::u_to_b()), static_cast<float>(b.u12 / gemmi::u_to_b()),
static_cast<float>(b.u13 / gemmi::u_to_b()), static_cast<float>(b.u23 / gemmi::u_to_b())};
}
MoveRiders();
SetAniso(aniso);
}
// The U of every atom (AnisoU for the anisotropic, B / (8 pi^2) I for the others), in A^2.
std::vector<std::array<double, 6>> UOfAtoms() const {
std::vector<std::array<double, 6>> u(atoms.size());
for (size_t i = 0; i < atoms.size(); i++) {
const double b = atoms[i]->b_iso / gemmi::u_to_b();
if (!aniso.empty() && aniso[i].on) {
const gemmi::SMat33<float> m = AnisoU(aniso[i], atoms[i]->b_iso);
u[i] = {m.u11, m.u22, m.u33, m.u12, m.u13, m.u23};
} else
u[i] = {b, b, b, 0, 0, 0};
}
return u;
}
// The anisotropic U the atoms of the model bring (PrepareModel() keeps it) as their AtomAniso, taken off them.
void TakeAniso() {
std::vector<AtomAniso> a(atoms.size());
bool any = false;
for (size_t i = 0; i < atoms.size(); i++) {
gemmi::Atom &at = *atoms[i];
if (!at.aniso.nonzero())
continue;
// its widest axis no wider than the refinement allows
const double top = AnisoRadiusB(at.aniso.scaled(gemmi::u_to_b()));
if (top > B_CEILING_MAX) {
at.aniso = at.aniso.scaled(static_cast<float>(B_CEILING_MAX / top));
at.b_iso = static_cast<float>(gemmi::u_to_b() * at.aniso.trace() / 3.0);
}
const double t = at.aniso.trace() / 3.0;
a[i].on = true;
a[i].dev = {static_cast<float>(at.aniso.u11 - t), static_cast<float>(at.aniso.u22 - t),
static_cast<float>(at.aniso.u33 - t), at.aniso.u12, at.aniso.u13, at.aniso.u23};
at.aniso = {0, 0, 0, 0, 0, 0};
any = true;
}
SetAniso(any ? a : std::vector<AtomAniso>());
}
bool Cancelled() const { return cancel != nullptr && cancel->load(std::memory_order_relaxed); }
gemmi::Model &Model() { return st.models[0]; }
// The riders' parents, and per parent its riders in index order (rider_start, rider_index: CSR over the atoms).
void SetRiderParents(const std::vector<int> &parents) {
rider_parent = parents;
rider_start.assign(atoms.size() + 1, 0);
for (int p : rider_parent)
if (p >= 0)
rider_start[p + 1]++;
for (size_t i = 0; i < atoms.size(); i++)
rider_start[i + 1] += rider_start[i];
rider_index.resize(rider_start.back());
std::vector<int> fill(rider_start.begin(), rider_start.end() - 1);
for (size_t i = 0; i < rider_parent.size(); i++)
if (rider_parent[i] >= 0)
rider_index[fill[rider_parent[i]]++] = static_cast<int>(i);
}
// The hydrogens moved with their parents, at the offsets they were last placed at, and given their B.
void MoveRiders() {
const int n = static_cast<int>(rider_parent.size());
ParallelChunks(n, ThreadsForWork(n, nthreads, 16384), [&](int lo, int hi) {
for (int i = lo; i < hi; 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: each parent takes its riders in index order, as
// the serial walk over the atoms added them (a rider has no riders of its own).
void FoldRiders(std::vector<std::array<double, 4>> &g) const {
const int n = static_cast<int>(rider_parent.size());
ParallelChunks(n, ThreadsForWork(n, nthreads, 16384), [&](int lo, int hi) {
for (int i = lo; i < hi; i++)
for (int r = rider_start[i]; r < rider_start[i + 1]; r++)
for (int c = 0; c < 4; c++)
g[i][c] += g[rider_index[r]][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 (size_t i = 0; i < atoms.size(); i++)
b_min = std::min<double>(b_min, atoms[i]->b_iso - (b_lo.empty() ? 0.0 : b_lo[i]));
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.RowHkl(), fmask);
}
std::optional<RefinementStage> s3("refit.scale");
// the per-reflection loops on the threads, each element on its own
const int n_refl = static_cast<int>(refl.Size());
const size_t refl_threads = ThreadsForWork(refl.Size(), nthreads);
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;
fc.v.resize(refl.Size());
fm.v.resize(refl.Size());
ParallelChunks(n_refl, refl_threads, [&](int lo, int hi) {
for (int m = lo; m < hi; m++) {
fc.v[m] = {refl.hkl[m], fcalc[m]};
fm.v[m] = {refl.hkl[m], fmask[m]};
}
});
fo.v.reserve(refl.Size());
for (size_t m = 0; m < refl.Size(); 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;
// the scale of the score: by robust least squares (ModelScaleBox::robust)
auto fit_scale = [&]() {
scaling.points.clear();
scaling.prepare_points(fc, fo, &fm);
#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);
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.Rows());
ParallelChunks(n_refl, refl_threads, [&](int lo, int hi) {
for (int m = lo; m < hi; 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];
}
});
for (size_t m = refl.Size(); m < refl.Rows(); m++) {
const gemmi::Miller &h = refl.mate_hkl[m - refl.Size()];
rows.fmask[m] = static_cast<float>(scaling.get_solvent_scale(cell.calculate_stol_sq(h))) * fmask[m];
}
};
fit_scale();
if (refl.Twinned()) {
// The scale and the solvent inside the twin sum: the fraction for the untwinned fit, then the scale
// fitted again against Fo detwinned by the model as it stands, Fo |F_model(h)| / |F_twin(h)| - which
// the twinned model at that scale meets as the untwinned one meets Fo - and the fraction again.
rows.twin_fraction = static_cast<float>(FitTwinFraction(refl, rows, fcalc));
const std::vector<std::complex<float>> twinned = TwinnedModel(refl, rows, fcalc);
fo.v.clear();
for (size_t m = 0; m < refl.Size(); m++) {
if (refl.free[m])
continue;
const double own = std::abs(rows.k_total[m] * (fcalc[m] + rows.fmask[m]));
const double tw = std::abs(twinned[m]);
fo.v.push_back({refl.hkl[m], {static_cast<float>(tw > 0 ? refl.fo[m] * own / tw : refl.fo[m]), 1.0f}});
}
fit_scale();
rows.twin_fraction = static_cast<float>(FitTwinFraction(refl, rows, fcalc));
}
s3.reset();
const std::vector<std::complex<float>> fmodel = TwinnedModel(refl, rows, fcalc);
{
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) {
const double lo = b_lo.empty() ? 0.0 : b_lo[i], hi = b_hi.empty() ? 0.0 : b_hi[i];
atoms[i]->b_iso = static_cast<float>(std::clamp(atoms[i]->b_iso + b_iso, B_FLOOR + lo,
std::max(B_FLOOR + lo, b_ceiling - hi)));
b_min = std::min<double>(b_min, atoms[i]->b_iso - lo);
}
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));
}
for (size_t m = refl.Size(); m < refl.Rows(); m++)
rows.fmask[m] *= static_cast<float>(std::exp(-b_iso * cell.calculate_1_d2(refl.mate_hkl[m - refl.Size()]) / 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;
sc.twin_fraction = rows.twin_fraction;
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 = IonFormFactor<double>(a.element.elem, a.charge);
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> &params, 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);
return TwinnedModel(refl, rows, fcalc);
}
};
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 };
// `topo`: its atoms on special positions move only within their site's subspace (SpecialPosition).
AtomTarget(Context &ctx, Mode mode, const std::vector<int> &param_atoms, double weight, int restraint_set,
const RefinementTopology &topo)
: ctx_(ctx), mode_(mode), param_atoms_(param_atoms), weight_(weight), restraint_set_(restraint_set), topo_(topo) {}
// 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 {
const int n = static_cast<int>(param_atoms_.size());
ParallelChunks(n, ThreadsForWork(n, ctx_.nthreads, 16384), [&](int lo, int hi) {
for (int k = lo; k < hi; 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]);
if (const SpecialPosition *sp = Special(param_atoms_[k]))
a.pos = sp->OnSite(a.pos); // where the shifts, all within its subspace, left it by rounding
} else
a.b_iso = static_cast<float>(std::clamp(p[k], ctx_.BLow(param_atoms_[k]), ctx_.BHigh(param_atoms_[k])));
}
});
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 = gx_, &gr = 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);
const int n = static_cast<int>(param_atoms_.size());
ParallelChunks(n, ThreadsForWork(n, ctx_.nthreads, 16384), [&](int lo, int hi) {
for (int k = lo; k < hi; k++) {
const int i = param_atoms_[k];
if (mode_ == Mode::Xyz) {
gemmi::Vec3 g(weight_ * gx[i][0] + gr[i][0], weight_ * gx[i][1] + gr[i][1],
weight_ * gx[i][2] + gr[i][2]);
if (const SpecialPosition *sp = Special(i))
g = sp->projector.multiply(g);
gradient[3 * k] = g.x;
gradient[3 * k + 1] = g.y;
gradient[3 * k + 2] = g.z;
} 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_.BLow(i) && g > 0) || (p[k] >= ctx_.BHigh(i) && g < 0);
gradient[k] = held ? 0.0 : g;
}
}
});
return std::isfinite(*cost);
}
private:
const SpecialPosition *Special(int atom) const {
const int k = topo_.special_index.empty() ? -1 : topo_.special_index[atom];
return k >= 0 ? &topo_.special[k] : nullptr;
}
Context &ctx_;
Mode mode_;
const std::vector<int> &param_atoms_;
double weight_;
int restraint_set_;
const RefinementTopology &topo_;
// the gradients per atom, kept from evaluation to evaluation (7 MB each on a large model: made anew, their
// pages fault in every time)
mutable std::vector<std::array<double, 4>> gx_, gr_;
};
// The X-ray target's derivatives in U of every atom, with what the target says of a riding hydrogen's B (its parent's
// equivalent isotropic B, 8 pi^2 tr(U) / 3) said of its parent's U.
std::vector<std::array<double, 6>> FoldedUGradient(const Context &ctx, const std::vector<std::array<double, 4>> &gx) {
std::vector<std::array<double, 6>> g = ctx.engine->UGradient();
g.resize(ctx.atoms.size(), {0, 0, 0, 0, 0, 0});
for (size_t r = 0; r < ctx.rider_parent.size(); r++)
if (const int p = ctx.rider_parent[r]; p >= 0)
for (int c = 0; c < 3; c++)
g[p][c] += gx[r][3] * gemmi::u_to_b() / 3.0;
return g;
}
// The displacements of the anisotropic atoms against w times the X-ray target. Mode::U: a U per atom of
// `param_atoms`, as 8 pi^2 U (in the units of B), plus the U restraints; its smallest eigenvalue held at the cycle's
// floor and its diagonal under the ceiling. Mode::Tls: the parameters of each of ctx.tls, in the units of B at the
// body's radius (T, L r^2, S r, each times 8 pi^2), the atoms' residual B held; X-ray only.
class AdpTarget {
public:
enum class Mode { U, Tls };
AdpTarget(Context &ctx, Mode mode, const std::vector<int> &param_atoms, double weight, const URestraints *restraints)
: ctx_(ctx), mode_(mode), param_atoms_(param_atoms), weight_(weight), restraints_(restraints) {
for (const TlsBody &body : ctx.tls) {
double r2 = 0;
for (int i : body.atoms)
r2 += body.origin.dist_sq(ctx.atoms[i]->pos);
radius_.push_back(std::max(1.0, std::sqrt(r2 / std::max<size_t>(body.atoms.size(), 1))));
}
}
std::vector<double> Get() const {
std::vector<double> p;
if (mode_ == Mode::U) {
const std::vector<std::array<double, 6>> u = ctx_.UOfAtoms();
for (int i : param_atoms_)
for (int c = 0; c < 6; c++)
p.push_back(gemmi::u_to_b() * u[i][c]);
} else
for (size_t b = 0; b < ctx_.tls.size(); b++) {
const std::array<double, TlsBody::PARAMETERS> q = ctx_.tls[b].Get();
for (int k = 0; k < TlsBody::PARAMETERS; k++)
p.push_back(q[k] * Scale(b, k));
}
return p;
}
void Set(const double *p) const {
if (mode_ == Mode::U) {
const int n = static_cast<int>(param_atoms_.size());
ParallelChunks(n, ThreadsForWork(n, ctx_.nthreads, 16384), [&](int lo, int hi) {
for (int k = lo; k < hi; k++)
SetU(param_atoms_[k], p + 6 * k);
});
ctx_.MoveRiders();
ctx_.SetAniso(ctx_.aniso);
} else {
for (size_t b = 0; b < ctx_.tls.size(); b++) {
std::array<double, TlsBody::PARAMETERS> q{};
for (int k = 0; k < TlsBody::PARAMETERS; k++)
q[k] = p[b * TlsBody::PARAMETERS + k] / Scale(b, k);
ctx_.tls[b].Set(q);
}
ctx_.ApplyTls();
}
}
double MaxStep() const { return MAX_STEP_B; }
bool Evaluate(const double *p, double *cost, double *gradient) const {
const RefinementStage stage(mode_ == Mode::U ? "eval.u" : "eval.tls");
Set(p);
ctx_.evaluations++;
std::vector<std::array<double, 4>> gx;
const double xray = ctx_.engine->Evaluate(ctx_.Model(), gradient != nullptr ? &gx : nullptr);
double restraints = 0;
std::vector<std::array<double, 6>> gr;
if (mode_ == Mode::U && restraints_ != nullptr) {
if (gradient != nullptr)
gr.assign(ctx_.atoms.size(), {0, 0, 0, 0, 0, 0});
restraints = EvaluateURestraints(*restraints_, ctx_.Positions(), ctx_.UOfAtoms(), gradient != nullptr ? &gr : nullptr);
}
*cost = weight_ * xray + restraints;
if (gradient == nullptr)
return std::isfinite(*cost);
const std::vector<std::array<double, 6>> gu = FoldedUGradient(ctx_, gx);
if (mode_ == Mode::U) {
for (size_t k = 0; k < param_atoms_.size(); k++)
for (int c = 0; c < 6; c++)
gradient[6 * k + c] = (weight_ * gu[param_atoms_[k]][c] + (gr.empty() ? 0.0 : gr[param_atoms_[k]][c])) /
gemmi::u_to_b();
} else {
for (size_t b = 0; b < ctx_.tls.size(); b++) {
const std::array<double, TlsBody::PARAMETERS> d = ctx_.tls[b].Gradient(ctx_.atoms, gu);
for (int k = 0; k < TlsBody::PARAMETERS; k++)
gradient[b * TlsBody::PARAMETERS + k] = weight_ * d[k] / Scale(b, k);
}
}
return std::isfinite(*cost);
}
private:
double Scale(size_t body, int k) const {
const double r = radius_[body];
return gemmi::u_to_b() * (k < 6 ? 1.0 : k < 12 ? r * r : r);
}
// Atom i's U from 8 pi^2 U = p[0..5]: its smallest eigenvalue raised to the cycle's floor and its largest brought
// down to the ceiling where they are outside, then its B and its anisotropy.
void SetU(int i, const double *p) const {
gemmi::SMat33<double> b{p[0], p[1], p[2], p[3], p[4], p[5]};
const std::array<double, 3> ev = b.calculate_eigenvalues();
const double lowest = std::min({ev[0], ev[1], ev[2]});
if (!(lowest >= ctx_.b_floor))
b = b.added_kI(ctx_.b_floor - (std::isfinite(lowest) ? lowest : 0.0));
const std::array<double, 3> ev2 = b.calculate_eigenvalues();
const double top = std::max({ev2[0], ev2[1], ev2[2]});
if (top > ctx_.b_ceiling)
b = b.scaled(ctx_.b_ceiling / top);
const double t = b.trace() / 3.0;
ctx_.atoms[i]->b_iso = static_cast<float>(t);
AtomAniso &a = ctx_.aniso[i];
a.on = true;
a.dev = {static_cast<float>((b.u11 - t) / gemmi::u_to_b()), static_cast<float>((b.u22 - t) / gemmi::u_to_b()),
static_cast<float>((b.u33 - t) / gemmi::u_to_b()), static_cast<float>(b.u12 / gemmi::u_to_b()),
static_cast<float>(b.u13 / gemmi::u_to_b()), static_cast<float>(b.u23 / gemmi::u_to_b())};
}
Context &ctx_;
Mode mode_;
const std::vector<int> &param_atoms_;
double weight_;
const URestraints *restraints_;
std::vector<double> radius_;
};
// 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> &param_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).
template <typename Target>
void Minimise(Context &ctx, const Target &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_.BLow(i) && b < ctx_.BHigh(i))
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_.BLow(i), ctx_.BHigh(i)));
}
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));
// an atom on a special position stays on it
if (!r.is_water() && ctx.st.cell.is_special_position(r.atoms[k].pos, SPECIAL_POSITION_TOLERANCE) == 0)
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;
}
// A phase in degrees as an MTZ 'P' column holds it: in [0, 360), also after the rounding to float.
float MtzPhase(double deg) {
const auto f = static_cast<float>(deg - 360.0 * std::floor(deg / 360.0));
return f < 360.0f ? f : 0.0f;
}
// 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). With a twin law, Fo is detwinned by the model in proportion to its
// own share of the twinned intensity, Fo |F_model(h)| / |F_twin(h)|, and Fc is the model's own |F_model(h)|.
// 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, fwt_data;
diff.unit_cell_ = fwt_data.unit_cell_ = ctx.cell;
diff.spacegroup_ = fwt_data.spacegroup_ = ctx.sg;
for (size_t m = 0; m < ctx.refl.Size(); m++) {
double fo = ctx.refl.fo[m], fc = std::abs(fmodel[m]);
if (ctx.refl.Twinned()) {
const double own = std::abs(ctx.rows.k_total[m] * (ctx.fcalc[m] + ctx.rows.fmask[m]));
if (fc > 0)
fo *= own / fc;
fc = own;
}
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]);
const double phi_deg = phi * 180.0 / PI;
// A negative coefficient is written as its magnitude on the phase turned by 180 degrees: an MTZ amplitude
// column is non-negative by convention, and not every reader accepts otherwise.
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),
MtzPhase(phi_deg), static_cast<float>(std::abs(fwt)), MtzPhase(fwt < 0 ? phi_deg + 180.0 : phi_deg),
static_cast<float>(std::abs(delfwt)), MtzPhase(delfwt < 0 ? phi_deg + 180.0 : 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))});
fwt_data.v.push_back({ctx.refl.hkl[m], std::polar(static_cast<float>(fwt), static_cast<float>(phi))});
}
result.maps = maps;
// The difference map and the 2mFo-DFc map beside the anomalous one below: none reads another.
const std::array<int, 3> size = gemmi::get_size_for_hkl(diff, {{0, 0, 0}}, 3.0);
std::future<void> difference_map = std::async(std::launch::async, [&diff, &size, &result] {
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;
});
std::future<void> map = std::async(std::launch::async, [&fwt_data, &size, &result] {
result.map = std::make_shared<gemmi::Grid<float>>(MapFromFPhi(gemmi::get_f_phi_on_grid<float>(fwt_data, size, true)));
result.map_rms = gemmi::calculate_data_statistics(result.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();
map.get();
}
// The null of a twinned refinement, for the model as the last Refit() left it: its R-free at the same fraction with
// every reflection's mate replaced by an unrelated observed one of the same resolution and the same set (free with
// free, work with work, as a law's twin pairs are). The reflections of each set are put in order of 1/d^2 with the two
// members of a twin pair next to each other, and paired TWIN_NULL_SHIFT apart - past the reflection's own pair, and no
// further, as the free set is sparse and a wider step takes the mate to another resolution, which a null pays for.
// What the averaging of two intensities buys on its own, whatever they are.
constexpr size_t TWIN_NULL_SHIFT = 2;
// The twinned refinement is taken only where the law beats its null on the free set by this many standard errors -
// or by TWIN_NULL_Z_DATA where the data's own intensity statistics read twinned, a second, independent piece of
// evidence (AdoptTwinLaw).
constexpr double TWIN_NULL_Z = 3.0;
constexpr double TWIN_NULL_Z_DATA = 2.0;
void TwinNull(const Context &ctx, RefinementResult &result) {
RefinementReflections null = ctx.refl;
const size_t n = null.Size();
null.mate_hkl.clear();
null.mate.assign(n, 0);
for (const bool free : {false, true}) {
std::vector<int> order;
for (size_t m = 0; m < n; m++)
if ((null.free[m] != 0) == free)
order.push_back(static_cast<int>(m));
// a twin pair is ordered at the 1/d^2 of its first member
auto pair_of = [&](int m) { return ctx.refl.mate[m] < static_cast<int>(n) ? std::min(m, ctx.refl.mate[m]) : m; };
std::stable_sort(order.begin(), order.end(), [&](int a, int b) {
const int pa = pair_of(a), pb = pair_of(b);
return std::make_pair(null.inv_d2[pa], pa) < std::make_pair(null.inv_d2[pb], pb);
});
for (size_t r = 0; r < order.size(); r++) {
// in blocks of twice the shift, the first half with the second
const size_t p = r % (2 * TWIN_NULL_SHIFT) < TWIN_NULL_SHIFT ? r + TWIN_NULL_SHIFT : r - TWIN_NULL_SHIFT;
null.mate[order[r]] = order[p < order.size() ? p : r];
}
}
// at the law's fraction: averaging any two intensities lifts the weak reflections' model amplitudes off zero
// by itself (the fraction the unrelated mates ask for, reported beside it, is about 0)
XrayRowModel rows = ctx.rows;
result.twin_fraction_null = FitTwinFraction(null, rows, ctx.fcalc);
const std::vector<std::complex<float>> f_null = TwinnedModel(null, rows, ctx.fcalc);
result.twin_r_free_null = ComputeRFactors(null, f_null).r_free;
// How far the law beats its null on the free reflections, paired: the mean of |Fo - |F_null|| - |Fo - |F_twin||
// over its standard error. The two members of a free twin pair share their two F_calc, so their terms are summed
// into one: the standard error counts twin pairs, not reflections.
const std::vector<std::complex<float>> f_law = TwinnedModel(ctx.refl, ctx.rows, ctx.fcalc);
std::vector<double> d(n, 0.0);
for (size_t m = 0; m < n; m++)
if (ctx.refl.free[m])
d[m] = std::fabs(ctx.refl.fo[m] - std::abs(f_null[m])) - std::fabs(ctx.refl.fo[m] - std::abs(f_law[m]));
double sum = 0, sum2 = 0;
int n_pairs = 0;
for (size_t m = 0; m < n; m++) {
const size_t t = ctx.refl.mate[m];
if (!ctx.refl.free[m] || (t < n && t < m))
continue; // the pair was counted at its first member
const double pair = t < n && t != m ? d[m] + d[t] : d[m];
sum += pair;
sum2 += pair * pair;
n_pairs++;
}
if (n_pairs > 1) {
const double mean = sum / n_pairs, var = (sum2 - n_pairs * mean * mean) / (n_pairs - 1);
result.twin_null_z = var > 0 ? mean / std::sqrt(var / n_pairs) : 0.0;
}
}
// 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. `restraint_records`, `restraint_slots`: what its restraint sets hold at most (0: the GPU
// engine's estimate per atom).
void MakeEngine(Context &ctx, bool gpu, const RefinementSettings &settings, Logger &logger, size_t restraint_records = 0,
size_t restraint_slots = 0) {
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;
RefinementGPUBytes bytes;
ctx.engine = MakeRefinementEngineGPU(device, ctx.cell, *ctx.sg, ctx.d_min, ctx.refl, ctx.Model(),
ctx.b_ceiling, budget, &bytes, settings.nthreads, restraint_records,
restraint_slots);
ctx.device = device;
ctx.engine_bytes = bytes.engine;
ctx.card_bytes = bytes.card;
ctx.lane_bytes = bytes.lane;
if (ctx.engine) {
ctx.scale_gpu = std::make_shared<ModelScaleGPUPool>(device, ctx.refl.Size(), 1);
logger.Debug("Refinement: GPU engine {:.2f} GB, {:.2f} GB shared with the slots on its card and {:.2f} GB "
"per lane, of a {:.2f} GB budget (scale pool {:.2f} GB)", bytes.engine / 1e9,
bytes.card / 1e9, bytes.lane / 1e9, budget / 1e9,
ModelScaleGPUPool::DeviceBytes(ctx.refl.Size(), 1, true) / 1e9);
} else
logger.Info("Refinement: on the CPU - the GPU cannot take this case or {:.1f} GB is over its budget",
(bytes.engine + bytes.card + bytes.lane) / 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);
ctx.engine->SetAniso(&ctx.aniso);
}
// 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, and where `ncs_sigma` is not 0 the NCS restraints at it.
RestraintSet RestraintsNow(Context &ctx, const RefinementTopology &topo, const NcsCopies *ncs = nullptr,
double ncs_sigma = 0) {
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));
if (ncs != nullptr && ncs_sigma > 0)
set.Append(NcsRestraintsAt(*ncs, ncs_sigma));
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;
};
// The scale pools of the slots' cards: one per card, shared by the slots there (a fit waits for the one before it).
struct ScalePools {
#ifdef JFJOCH_USE_CUDA
std::mutex m;
std::map<int, std::shared_ptr<ModelScaleGPUPool>> pool;
std::shared_ptr<ModelScaleGPUPool> Get(int device, size_t points) {
const std::lock_guard lock(m);
std::shared_ptr<ModelScaleGPUPool> &p = pool[device];
if (!p)
p = std::make_shared<ModelScaleGPUPool>(device, points, 1);
return p;
}
#endif
};
// 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`, like the main one's and sharing with the
// engines on that card what they share (RefinementGPUBytes), and the card's scale pool. Null where the card cannot
// make one.
std::unique_ptr<OwnSlot> MakeOwnSlot(const gemmi::Structure &before_topology, const Context &main, int device,
bool own_lane, size_t nthreads, const RefinementSettings &settings,
ScalePools &pools) {
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.SetRiderParents(slot->topo.rider_parent);
ctx.TakeAniso(); // each trial sets the anisotropy it starts from
#ifdef JFJOCH_USE_CUDA
if (main.engine->OnGPU()) {
try {
ctx.engine = MakeRefinementEngineGPU(*main.engine, device, nthreads, own_lane);
ctx.scale_gpu = pools.Get(device, ctx.refl.Size());
} catch (const std::exception &) {
return nullptr;
}
ctx.device = device;
}
#else
(void) device;
(void) own_lane;
(void) pools;
#endif
if (!ctx.engine)
ctx.engine = MakeRefinementEngineCPU(ctx.cell, *ctx.sg, ctx.d_min, ctx.refl, nthreads);
ctx.engine->SetAniso(&ctx.aniso);
ctx.engine->SetRestraints(1, slot->topo.adp);
return slot;
}
// Where a slot beside the main one goes: its GPU (-1: the CPU), and whether its engine has a lane of working buffers
// of its own there.
struct SlotPlace {
int device = -1;
bool own_lane = false;
};
// The places 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 - so the first slots are on cards of their own, where RunTrials
// puts the long trials. A slot takes its engine's own part and a lane of its own where both fit, else its engine
// alone on a lane it shares (one evaluation at a time: lanes side by side keep the card busier, more slots overlap
// more host work); the first one on another card also the part the engines there share and a scale pool. On the
// CPU, the main one and those beside it are no more than the threads.
std::vector<SlotPlace> OwnSlotPlaces(const Context &main, int wanted, const RefinementSettings &settings) {
#ifdef JFJOCH_USE_CUDA
if (main.engine->OnGPU()) {
const int n = get_gpu_count();
// the robust fit's buffers included
const size_t pool = ModelScaleGPUPool::DeviceBytes(main.refl.Size(), 1, true);
std::vector<size_t> used(n, 0);
used[main.device] = main.engine_bytes + main.lane_bytes + main.card_bytes + pool;
std::vector<SlotPlace> places;
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;
const size_t first = used[d] == 0 ? main.card_bytes + pool : 0;
if (used[d] + first + main.engine_bytes + main.lane_bytes <= budget) {
used[d] += first + main.engine_bytes + main.lane_bytes;
places.push_back({d, true});
} else if (first == 0 && used[d] + main.engine_bytes <= budget) {
used[d] += main.engine_bytes;
places.push_back({d, false});
}
}
return places;
}
#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<SlotPlace>(std::max(beside, 0));
}
// The weights of one trial: the positions' X-ray weight, the B factors' (0: held) and the NCS restraints' sigma (0:
// none).
struct TrialWeights {
double w = 0, adp = 0, ncs = 0;
};
// A displacement model a trial refines (RefinementSettings::adp_hypotheses): the B factors as they are (isotropic, or
// the shapes of the anisotropic atoms held), TLS bodies with a residual B per atom, or a U per atom.
struct AdpHypothesis {
enum class Kind { Isotropic, Tls, Anisotropic };
Kind kind = Kind::Isotropic;
std::string name;
std::vector<TlsBody> bodies; // Tls, with the part of each atom's B they hold at the start (Context::tls_b)
std::vector<double> tls_b;
std::vector<std::pair<int, int>> pairs; // Anisotropic: the pairs of the U restraints
int segments = 0; // Tls: the TLSMD cut (segments per chain at most), 0 the model's own
};
// 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;
std::vector<AtomAniso> aniso;
std::vector<TlsBody> tls;
std::vector<double> tls_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
int evaluations = 0;
};
TrialOutcome RunTrial(TrialSlot slot, const std::vector<int> &refinable, const std::vector<gemmi::Position> &pos0,
const std::vector<double> &b0, const std::vector<AtomAniso> &aniso0, const Context::Fit *fit,
double w, double adp_w, int cycles, const RefinementSettings &settings, const NcsCopies *ncs,
double ncs_sigma, const AdpHypothesis *hyp = nullptr, bool hold_positions = false) {
Context &ctx = *slot.ctx;
RefinementTopology &topo = *slot.topo;
const RefinementStage stage(hyp != nullptr ? "rung.adp" : adp_w > 0 ? "rung.b" : "rung.xyz");
const int evaluations_before = ctx.evaluations;
ctx.SetPositions(pos0);
ctx.SetB(b0);
ctx.SetAniso(aniso0);
ctx.tls.clear();
ctx.tls_b.clear();
const AdpHypothesis::Kind kind = hyp != nullptr ? hyp->kind : AdpHypothesis::Kind::Isotropic;
// The displacement model of the hypothesis: its TLS bodies applied to their atoms, or a U of its own for every
// refinable atom (the isotropic ones from their B) with the U restraints.
std::vector<int> u_atoms;
URestraints u_restraints;
if (kind == AdpHypothesis::Kind::Tls) {
PlaceRiders(ctx, topo);
ctx.tls = hyp->bodies;
ctx.tls_b = hyp->tls_b;
ctx.ApplyTls();
} else if (kind == AdpHypothesis::Kind::Anisotropic) {
if (ctx.aniso.size() != ctx.atoms.size())
ctx.aniso.assign(ctx.atoms.size(), AtomAniso());
std::vector<char> refines_u(ctx.atoms.size(), 0);
for (int i : refinable)
if (!ctx.atoms[i]->is_hydrogen()) {
ctx.aniso[i].on = true;
refines_u[i] = 1;
u_atoms.push_back(i);
}
ctx.SetAniso(ctx.aniso);
u_restraints = MakeURestraints(hyp->pairs, refines_u);
}
// 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, topo),
b(ctx, AtomTarget::Mode::B, refinable, adp_w, 1, topo);
const AdpTarget tls_target(ctx, AdpTarget::Mode::Tls, refinable, 1.0, nullptr);
const AdpTarget u_target(ctx, AdpTarget::Mode::U, u_atoms, adp_w, &u_restraints);
RefinementMinimiser xyz_minimiser(20, ctx.nthreads), b_minimiser(20, ctx.nthreads), tls_minimiser(20, ctx.nthreads);
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();
if (!hold_positions) {
{
const RestraintSet set = RestraintsNow(ctx, topo, ncs, ncs_sigma);
const RefinementStage s1("set_restraints");
ctx.engine->SetRestraints(0, set);
}
const RefinementStage s2("minimise.xyz");
Minimise(ctx, xyz, settings.xyz_iterations, xyz_minimiser);
}
if (kind == AdpHypothesis::Kind::Tls) {
ctx.ApplyTls(); // the bodies' U at the atoms as they are
const RefinementStage s4("minimise.tls");
Minimise(ctx, tls_target, settings.b_iterations, tls_minimiser);
}
if (adp_w > 0 && kind == AdpHypothesis::Kind::Anisotropic) {
// each U's six elements take a third of the curvature of its B
std::vector<double> xray = ctx.XrayBCurvatureShape();
if (cycle == 0)
xray_scale = ctx.CalibrateBCurvature(refinable, xray);
// and the restraints' curvature in 8 pi^2 U: per pair 2 / sigma^2 of the similarity and about a third of
// the rigid bond's, per isotropic atom 2 / sigma^2
std::vector<double> restraint(xray.size(), 0.0);
const double u2 = gemmi::u_to_b() * gemmi::u_to_b();
const double per_pair = (2.0 / (URestraints::SIGMA_SIMILAR * URestraints::SIGMA_SIMILAR) +
2.0 / (3.0 * URestraints::SIGMA_RIGID * URestraints::SIGMA_RIGID)) / u2;
for (const auto &[i, j] : u_restraints.pairs) {
restraint[i] += per_pair;
restraint[j] += per_pair;
}
for (int i : u_restraints.isotropic)
restraint[i] += 2.0 / (URestraints::SIGMA_ISOTROPIC * URestraints::SIGMA_ISOTROPIC) / u2;
for (double &x : xray)
x *= xray_scale / 3.0;
std::vector<double> diagonal;
for (double h : BDiagonal(xray, restraint, u_atoms, adp_w))
diagonal.insert(diagonal.end(), 6, h);
const RefinementStage s3("minimise.u");
Minimise(ctx, u_target, settings.b_iterations, b_minimiser, diagonal);
} else 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;
t.ncs_sigma = ncs_sigma;
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;
if (ncs != nullptr && !ncs->Empty())
t.ncs_rmsd = NcsRmsd(*ncs, ctx.XyzB());
if (hyp != nullptr) {
t.adp = hyp->name;
t.nll_free = ctx.Score().nll_free;
}
out.line = fmt::format("{}X-ray weight {:.3g}, {} {}{}: R-work {:.4f}, R-free {:.4f}{}, rmsd bonds {:.4f} A, "
"angles {:.2f} deg{}{}",
hyp != nullptr ? hyp->name + ", " : std::string(), w,
kind == AdpHypothesis::Kind::Anisotropic ? "U" : kind == AdpHypothesis::Kind::Tls ? "residual B" : "B",
adp_w > 0 ? fmt::format("refined at weight {:.3g}", adp_w) : std::string("held"),
ncs == nullptr || ncs->Empty() ? std::string()
: ncs_sigma > 0 ? fmt::format(", NCS at {:.3g} A (copies {:.3f} A apart)", ncs_sigma, t.ncs_rmsd)
: fmt::format(", no NCS (copies {:.3f} A apart)", t.ncs_rmsd),
rf.r_work, rf.r_free,
hyp != nullptr ? fmt::format(", -LL free {:.1f}", t.nll_free) : std::string(),
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.aniso = ctx.aniso;
out.tls = ctx.tls;
out.tls_b = ctx.tls_b;
out.fit = ctx.Fitted();
out.ran = true;
out.evaluations = ctx.evaluations - evaluations_before;
return out;
}
// The trials (w, adp_w) side by side, each slot on a thread of its own (on its GPU), returned in the order given: the
// longest first, one to each slot, then each slot takes the next one as it is done. 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 and
// whichever slot runs it.
std::vector<TrialOutcome> RunTrials(const std::vector<TrialSlot> &trial_slots,
const std::vector<TrialWeights> &trials,
const std::vector<int> &refinable, const std::vector<gemmi::Position> &pos0,
const std::vector<double> &b0, const std::vector<AtomAniso> &aniso0,
const Context::Fit *fit, int cycles, const RefinementSettings &settings,
const NcsCopies *ncs, const std::vector<const AdpHypothesis *> &hyps = {},
bool hold_positions = false) {
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 longest go to the slots first - TLS bodies, then a U per atom, then the B factors refined, then held - so
// that two of them share a card only where the cards are fewer than they are: the slots are on the cards in turn
// (OwnSlotPlaces), a card's second slot only after every card's first.
auto length = [&](size_t i) {
const AdpHypothesis::Kind kind = hyps.empty() ? AdpHypothesis::Kind::Isotropic : hyps[i]->kind;
return kind == AdpHypothesis::Kind::Tls ? 3 : kind == AdpHypothesis::Kind::Anisotropic ? 2 : trials[i].adp > 0;
};
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 length(a) > length(b); });
std::atomic<size_t> next{trial_slots.size()};
auto run_slot = [&](size_t s) {
try {
#ifdef JFJOCH_USE_CUDA
const ScopedGpuPin pin(trial_slots[s].ctx->device >= 0 ? trial_slots[s].ctx->device : get_gpu());
#endif
for (size_t j = s; j < trials.size(); j = next++)
if (const size_t i = order[j]; !trial_slots[s].ctx->Cancelled())
out[i] = RunTrial(trial_slots[s], refinable, pos0, b0, aniso0, fit, trials[i].w, trials[i].adp,
cycles, settings, ncs, trials[i].ncs, hyps.empty() ? nullptr : hyps[i],
hold_positions);
} 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
namespace {
// The solvent build inside the refinement (SolventBuild.h): the atom list of a refinement is fixed - its engines are
// made for it - so the waters a round may add are there from the start, as placeholders: one water each in a chain
// of their own at the end of the model, at occupancy 0, which puts no density, no mask and no contact anywhere and is
// not refined, until a round puts one at a peak. A water taken out goes to occupancy 0 the same way, and a site
// relabelled keeps its place in the list.
const char *const PLACEHOLDER_CHAIN = "~";
// As many placeholders as a fifteenth of the model's heavy atoms, at least 64: the depositions' builds add up to a
// tenth (2R8S, 1HR2), most of it in the first round.
int PlaceholderCount(const gemmi::Model &model) {
int heavy = 0;
for (const gemmi::Chain &ch : model.chains)
for (const gemmi::Residue &r : ch.residues)
for (const gemmi::Atom &a : r.atoms)
heavy += !a.is_hydrogen();
return std::max(64, heavy / 15);
}
void AddPlaceholders(gemmi::Structure &st, int n) {
gemmi::Chain ch(PLACEHOLDER_CHAIN);
for (int k = 0; k < n; k++) {
gemmi::Residue r;
r.name = "HOH";
r.seqid = gemmi::SeqId(k + 1, ' ');
r.het_flag = 'H';
r.entity_type = gemmi::EntityType::Water;
gemmi::Atom a;
a.name = "O";
a.element = gemmi::El::O;
a.occ = 0;
a.b_iso = 20;
r.atoms.push_back(a);
ch.residues.push_back(r);
}
st.models[0].chains.push_back(ch);
}
// The placeholders of a model (the last `n` atoms) and where those put in go: the chain of their nearest atom and the
// number they take there.
struct SolventPool {
int first = 0, n = 0;
std::vector<std::string> chain; // empty: not put in
std::vector<gemmi::SeqId> seqid;
std::map<std::string, int> last; // per chain name, the last residue number taken
int Free() const {
for (int k = 0; k < n; k++)
if (chain[k].empty())
return k;
return -1;
}
};
// The model as the solvent build and the output see it: the placeholders put in moved into their chains (the last of
// the name) under their numbers, those not put in and the waters taken out (occupancy 0) left out. Each atom's serial
// is its index in `st`.
gemmi::Structure SolventView(const gemmi::Structure &st, const SolventPool &pool) {
gemmi::Structure view = st;
gemmi::Model &model = view.models[0];
int index = 0;
for (gemmi::Chain &ch : model.chains)
for (gemmi::Residue &r : ch.residues)
for (gemmi::Atom &a : r.atoms)
a.serial = index++;
std::vector<gemmi::Residue> moved;
std::vector<std::string> into;
for (gemmi::Chain &ch : model.chains)
if (ch.name == PLACEHOLDER_CHAIN)
for (size_t k = 0; k < ch.residues.size(); k++)
if (!pool.chain[k].empty() && ch.residues[k].atoms[0].occ > 0) {
moved.push_back(ch.residues[k]);
moved.back().seqid = pool.seqid[k];
into.push_back(pool.chain[k]);
}
std::erase_if(model.chains, [](const gemmi::Chain &ch) { return ch.name == PLACEHOLDER_CHAIN; });
for (size_t k = 0; k < moved.size(); k++) {
gemmi::Chain *target = nullptr;
for (gemmi::Chain &ch : model.chains)
if (ch.name == into[k])
target = &ch;
if (target == nullptr) {
model.chains.emplace_back(into[k]);
target = &model.chains.back();
}
for (const gemmi::Residue &w : target->residues)
if (w.is_water()) {
moved[k].subchain = w.subchain;
moved[k].entity_id = w.entity_id;
}
target->residues.push_back(moved[k]);
}
for (gemmi::Chain &ch : model.chains) {
for (gemmi::Residue &r : ch.residues)
if (r.is_water())
std::erase_if(r.atoms, [](const gemmi::Atom &a) { return !(a.occ > 0); });
std::erase_if(ch.residues, [](const gemmi::Residue &r) { return r.atoms.empty(); });
}
return view;
}
// TLS bodies with their atoms (the refinable ones but hydrogens), about their centroid, T = t I with t half of what
// the lowest B among them holds, L = S = 0: so they leave every B as it is (tls_b), the residual taking the rest.
AdpHypothesis TlsHypothesis(const Context &ctx, std::string name, std::vector<std::vector<int>> groups) {
AdpHypothesis h;
h.kind = AdpHypothesis::Kind::Tls;
h.name = std::move(name);
h.tls_b.assign(ctx.atoms.size(), 0.0);
for (std::vector<int> &atoms : groups) {
if (atoms.empty())
continue;
TlsBody body;
body.id = std::to_string(h.bodies.size() + 1);
double b_min = ctx.b_ceiling;
for (int i : atoms) {
body.origin += ctx.atoms[i]->pos;
b_min = std::min<double>(b_min, ctx.atoms[i]->b_iso);
}
body.origin /= static_cast<double>(atoms.size());
const double t = 0.5 * b_min / gemmi::u_to_b();
body.T = {t, t, t, 0, 0, 0};
for (int i : atoms)
h.tls_b[i] = gemmi::u_to_b() * t;
body.atoms = std::move(atoms);
h.bodies.push_back(std::move(body));
}
return h;
}
// The displacement models the refinement chooses among once its weights are chosen, from the model as it stands: the
// TLS bodies of the model's own groups (`input`, where it has them); one body per chain of at least MIN_TLS_ATOMS
// refinable atoms, and the chains cut into 4, 8 and 16 TLSMD segments of at least MIN_TLS_RESIDUES residues
// (TlsSegments) - each cut a hypothesis where it differs from the one before; and a U per atom. (Cut into 2, the
// chains changed the choice on 6 of 364 refinements of the battery, none by more than 0.001 in R-free.)
std::vector<AdpHypothesis> AdpHypotheses(const Context &ctx, const RefinementTopology &topo,
const std::vector<TlsBody> &input) {
constexpr size_t MIN_TLS_ATOMS = 50;
constexpr int MIN_TLS_RESIDUES = 10;
std::vector<AdpHypothesis> out;
if (!input.empty()) {
AdpHypothesis h;
h.kind = AdpHypothesis::Kind::Tls;
h.name = fmt::format("TLS, the model's {} groups", input.size());
h.bodies = input;
h.tls_b.assign(ctx.atoms.size(), 0.0);
for (const TlsBody &body : h.bodies)
for (int i : body.atoms)
h.tls_b[i] = gemmi::u_to_b() * body.U(ctx.atoms[i]->pos).trace() / 3.0;
out.push_back(std::move(h));
}
// per chain its residues' refinable atoms but hydrogens
std::vector<std::vector<std::vector<int>>> chains;
int i = 0;
for (const gemmi::Chain &ch : ctx.st.models[0].chains) {
std::vector<std::vector<int>> residues;
size_t atoms = 0;
for (const gemmi::Residue &r : ch.residues) {
std::vector<int> ra;
for (size_t k = 0; k < r.atoms.size(); k++, i++)
if (!r.is_water() && topo.refinable[i] && !r.atoms[k].is_hydrogen())
ra.push_back(i);
if (!ra.empty()) {
atoms += ra.size();
residues.push_back(std::move(ra));
}
}
if (atoms >= MIN_TLS_ATOMS)
chains.push_back(std::move(residues));
}
if (!chains.empty()) {
const std::vector<gemmi::Position> pos = ctx.Positions();
const std::vector<double> b = ctx.B();
size_t last = 0;
for (int k : {1, 4, 8, 16}) {
std::vector<std::vector<int>> groups;
for (const std::vector<std::vector<int>> &residues : chains) {
std::vector<int> starts = TlsSegments(residues, pos, b, k, MIN_TLS_RESIDUES);
starts.push_back(static_cast<int>(residues.size()));
for (size_t s = 0; s + 1 < starts.size(); s++) {
std::vector<int> g;
for (int r = starts[s]; r < starts[s + 1]; r++)
g.insert(g.end(), residues[r].begin(), residues[r].end());
groups.push_back(std::move(g));
}
}
if (groups.size() == last)
continue;
last = groups.size();
std::string name = k == 1 ? fmt::format("TLS, {} groups (one per chain)", groups.size())
: fmt::format("TLS, {} groups (chains in up to {} segments)", groups.size(), k);
out.push_back(TlsHypothesis(ctx, std::move(name), std::move(groups)));
out.back().segments = k;
}
}
AdpHypothesis aniso;
aniso.kind = AdpHypothesis::Kind::Anisotropic;
aniso.name = "anisotropic";
for (const RestraintRecord &rec : topo.adp.records)
if (rec.kind == static_cast<int>(RestraintKind::BSimilarity))
aniso.pairs.emplace_back(rec.atom[0], rec.atom[1]);
out.push_back(std::move(aniso));
return out;
}
// The U X-ray weight that balances the gradients, as BalancedADPWeight() does for B: the U restraints' gradient with
// every U given an anisotropy of 20% of its B at random, over the X-ray target's, for the atoms that refine a U.
double BalancedUWeight(Context &ctx, const URestraints &restraints, const std::vector<int> &u_atoms) {
std::vector<std::array<double, 4>> gx;
ctx.engine->Evaluate(ctx.Model(), &gx);
const std::vector<std::array<double, 6>> gu = FoldedUGradient(ctx, gx);
std::vector<std::array<double, 6>> u = ctx.UOfAtoms();
std::mt19937 rng(3);
std::normal_distribution<double> shake(0.0, 0.2);
for (int i : u_atoms) {
const double t = (u[i][0] + u[i][1] + u[i][2]) / 3.0;
for (int c = 0; c < 6; c++)
u[i][c] += t * shake(rng);
}
std::vector<std::array<double, 6>> gr(u.size(), {0, 0, 0, 0, 0, 0});
EvaluateURestraints(restraints, ctx.Positions(), u, &gr);
double nx = 0, ng = 0;
for (int i : u_atoms)
for (int c = 0; c < 6; c++) {
nx += gu[i][c] * gu[i][c];
ng += gr[i][c] * gr[i][c];
}
return nx > 0 ? std::sqrt(ng / nx) : 1.0;
}
// The displacements of the model as given, read where it has TLS groups whose atoms bring no U of their own: their U
// from the groups, its B column taken as the total B (the U's equivalent isotropic B, as phenix.refine writes it) or as
// the residual to the groups (as REFMAC writes it), and the isotropic B as it stands - the one of the three whose
// R-free is lowest, beyond the tie in that order. `tls`: the groups, their atoms those of ctx.atoms but the hydrogens
// (which ride on their parents' B). Returns the R-factors of the reading kept.
// Following Winn, Isupov & Murshudov (2001) Acta Cryst. D57, 122-133 (the B column total or residual to the TLS)
RFactors ReadInputDisplacements(Context &ctx, const std::vector<TlsBody> &tls, Logger &logger, RefinementResult &result) {
int own = 0;
for (const AtomAniso &a : ctx.aniso)
own += a.on;
result.adp_input = own > 0 ? "ANISOU" : "isotropic";
result.adp_input_aniso_atoms = own;
std::vector<std::pair<int, gemmi::SMat33<double>>> tls_u; // atom, U_TLS
for (const TlsBody &body : tls)
for (int i : body.atoms)
if (ctx.aniso.empty() || !ctx.aniso[i].on)
tls_u.emplace_back(i, body.U(ctx.atoms[i]->pos));
RFactors r = ctx.Refit();
if (tls_u.empty())
return r;
const std::vector<double> b0 = ctx.B();
const std::vector<AtomAniso> aniso0 = ctx.aniso;
struct Reading {
const char *name;
std::vector<double> b;
std::vector<AtomAniso> aniso;
};
std::vector<Reading> readings;
for (int residual = 0; residual < 2; residual++) {
Reading read{residual ? "TLS, B residual" : "TLS, B total", b0,
aniso0.empty() ? std::vector<AtomAniso>(ctx.atoms.size()) : aniso0};
for (const auto &[i, u] : tls_u) {
const double t = u.trace() / 3.0;
read.aniso[i].on = true;
read.aniso[i].dev = {static_cast<float>(u.u11 - t), static_cast<float>(u.u22 - t),
static_cast<float>(u.u33 - t), static_cast<float>(u.u12), static_cast<float>(u.u13),
static_cast<float>(u.u23)};
if (residual)
read.b[i] = std::min(b0[i] + gemmi::u_to_b() * t, ctx.b_ceiling);
}
readings.push_back(std::move(read));
}
readings.push_back({result.adp_input == "ANISOU" ? "ANISOU" : "isotropic", b0, aniso0});
int best = -1;
std::vector<RFactors> rf;
for (size_t k = 0; k < readings.size(); k++) {
ctx.SetB(readings[k].b);
ctx.SetAniso(readings[k].aniso);
for (size_t i = 0; i < ctx.atoms.size(); i++)
if (!ctx.rider_parent.empty() && ctx.rider_parent[i] >= 0)
ctx.atoms[i]->b_iso = ctx.atoms[ctx.rider_parent[i]]->b_iso;
ctx.ForgetStarts();
rf.push_back(ctx.Refit());
logger.Info("Refinement: the model's {} TLS groups ({} atoms) read as {}: R-work {:.4f}, R-free {:.4f}", tls.size(),
tls_u.size(), readings[k].name, rf[k].r_work, rf[k].r_free);
if (best < 0 || rf[k].r_free < rf[best].r_free - R_FREE_TIE)
best = static_cast<int>(k);
}
ctx.SetB(readings[best].b);
ctx.SetAniso(readings[best].aniso);
for (size_t i = 0; i < ctx.atoms.size(); i++)
if (!ctx.rider_parent.empty() && ctx.rider_parent[i] >= 0)
ctx.atoms[i]->b_iso = ctx.atoms[ctx.rider_parent[i]]->b_iso;
ctx.ForgetStarts();
r = ctx.Refit();
if (best < 2) {
result.adp_input = readings[best].name;
result.adp_input_aniso_atoms = own + static_cast<int>(tls_u.size());
}
return r;
}
// 1/d^2 of the finest amplitude of `merged`, 0 without any.
double MaxInvD2(const std::vector<MergedReflection> &merged) {
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));
return max_inv_d2;
}
// The d_min a refinement composes to: just past the finest reflection, so that prepare_asu_data()'s strict limit
// keeps it.
double RefinementDMin(double max_inv_d2) {
return 0.999 / std::sqrt(max_inv_d2);
}
// One refinement of a model, with the solvent build in it where settings.solvent asks for it.
// `twin_law`: against the data as twinned by it, null for untwinned.
RefinementResult RefineOnce(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 gemmi::Op *twin_law = nullptr) {
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->aniso.nonzero() ? AnisoRadiusB(a->aniso.scaled(gemmi::u_to_b())) : a->b_iso);
// The model's TLS groups, where it has them; read as the residual to them, its B may go up by their trace.
const std::vector<TlsBody> tls_bodies =
ctx.st.meta.get_tls_groups() != nullptr ? TlsBodiesFromModel(ctx.Model(), *ctx.st.meta.get_tls_groups())
: std::vector<TlsBody>();
for (const TlsBody &body : tls_bodies)
for (int i : body.atoms)
b_max = std::max(b_max, ctx.atoms[i]->b_iso + gemmi::u_to_b() * body.U(ctx.atoms[i]->pos).trace() / 3.0);
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;
}
const double max_inv_d2 = MaxInvD2(merged);
if (max_inv_d2 <= 0) {
result.failure_reason = "no amplitudes";
return result;
}
ctx.d_min = RefinementDMin(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;
}
if (twin_law != nullptr) {
result.twin_split_pairs = SetTwinLaw(ctx.refl, cell, sg, ctx.d_min, *twin_law);
result.twinned = true;
result.twin_law = twin_law->as_hkl().triplet('h');
}
ctx.mask = std::make_unique<BulkSolventMask>(cell, sg, ctx.d_min, settings.nthreads);
// The solvent build's placeholders, before the slots copy the model.
const bool build = settings.solvent && settings.restrained && !settings.score_only &&
result.d_min <= SOLVENT_BUILD_D_MIN;
SolventPool pool;
if (build) {
pool.n = PlaceholderCount(ctx.Model());
for (const gemmi::Chain &ch : ctx.Model().chains)
for (const gemmi::Residue &r : ch.residues)
if (r.seqid.num != gemmi::SeqId::OptionalNum::None)
pool.last[ch.name] = std::max(pool.last[ch.name], static_cast<int>(r.seqid.num));
AddPlaceholders(ctx.st, pool.n);
pool.chain.assign(pool.n, "");
pool.seqid.assign(pool.n, gemmi::SeqId());
}
// 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.SetRiderParents(topo.rider_parent);
PlaceRiders(ctx, topo);
}
pool.first = static_cast<int>(ctx.atoms.size()) - pool.n; // the placeholders are the last atoms (no hydrogens)
// The copies of a molecule, restrained to each other in trials beside those that refine them independently.
NcsCopies ncs;
if (settings.restrained && settings.ncs)
ncs = FindNcsCopies(ctx.Model(), topo);
const NcsCopies *ncs_p = ncs.Empty() ? nullptr : &ncs;
ctx.TakeAniso();
// The engine's restraint sets as the model has them now - the geometry with its repulsions and the NCS restraints,
// and the B restraints - with a quarter more for the contacts a trial or the solvent build adds.
size_t restraint_records = 0, restraint_slots = 0;
if (restrained) {
for (const RestraintSet &set : {RestraintsNow(ctx, topo, ncs_p, 1.0), BRestraintsNow(ctx, topo)}) {
restraint_records += set.records.size() + set.plane_atoms.size();
restraint_slots += set.n_slots;
}
restraint_records += restraint_records / 4;
restraint_slots += restraint_slots / 4;
}
MakeEngine(ctx, settings.gpu, settings, logger, restraint_records, restraint_slots);
result.on_gpu = ctx.engine->OnGPU();
std::vector<TlsBody> input_tls;
if (ctx.st.meta.get_tls_groups() != nullptr) {
input_tls = TlsBodiesFromModel(ctx.Model(), *ctx.st.meta.get_tls_groups());
for (TlsBody &body : input_tls)
std::erase_if(body.atoms, [&](int i) { return ctx.atoms[i]->is_hydrogen(); });
}
RFactors r = ReadInputDisplacements(ctx, input_tls, logger, result);
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.refl.Twinned() ? fmt::format(" against the data twinned by {} (fraction {:.3f}, {} mates not "
"observed)", result.twin_law, ctx.rows.twin_fraction,
ctx.refl.mate_hkl.size())
: std::string(),
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 model as given, scaled: a candidate of the final choice beside the trials (the restrained stage, 5.), with
// its atoms and residues as they are, which the solvent build edits in place.
TrialOutcome given;
given.trial.r_work = r.r_work;
given.trial.r_free = r.r_free;
given.trial.nll_free = result.score_start.nll_free;
given.pos = ctx.Positions();
given.b = ctx.B();
given.aniso = ctx.aniso;
given.fit = ctx.Fitted();
given.ran = true;
std::vector<gemmi::Atom> given_atoms;
std::vector<std::pair<std::string, gemmi::EntityType>> given_residues;
const SolventPool given_pool = pool;
if (settings.restrained && !settings.score_only)
for (const gemmi::Chain &ch : ctx.Model().chains)
for (const gemmi::Residue &res : ch.residues) {
given_residues.emplace_back(res.name, res.entity_type);
given_atoms.insert(given_atoms.end(), res.atoms.begin(), res.atoms.end());
}
// The waters and ions the build relabelled and kept (SolventEdit::site their old name), for the metal records the
// refined model is written with.
std::vector<SolventEdit> relabels;
// The maps, the blobs and the model of where the refinement ended.
auto finish = [&]() {
MakeMaps(ctx, merged, wavelength_A, result);
if (ctx.refl.Twinned()) {
result.twin_fraction = ctx.rows.twin_fraction;
TwinNull(ctx, result);
}
// the anisotropic atoms' U into the model written out (ANISOU), and the TLS groups refined (the model's own
// stay where the U of their atoms was held, and go where every atom refined a U of its own)
ApplyAniso(ctx.Model(), ctx.aniso);
for (const AtomAniso &a : ctx.aniso)
result.aniso_atoms += a.on;
if (!ctx.tls.empty()) {
if (ctx.st.meta.refinement.empty())
ctx.st.meta.refinement.emplace_back();
ctx.st.meta.refinement[0].tls_groups = TlsBodiesToGroups(ctx.Model(), ctx.tls);
for (size_t k = 1; k < ctx.st.meta.refinement.size(); k++)
ctx.st.meta.refinement[k].tls_groups.clear();
} else if (result.adp_model == "anisotropic")
for (gemmi::RefinementInfo &ref : ctx.st.meta.refinement)
ref.tls_groups.clear();
if (const std::vector<gemmi::TlsGroup> *groups = ctx.st.meta.get_tls_groups())
result.tls_groups = static_cast<int>(groups->size());
// the model as the build left it: the waters put in in their chains, those taken out out, and each metal it
// relabelled with the records of its inner sphere in place of its old ones
gemmi::Structure model = build ? SolventView(ctx.st, pool) : ctx.st;
int record = 0;
for (const SolventEdit &e : relabels) {
const std::string from = e.site.substr(0, e.site.find(' '));
std::erase_if(model.connections, [&](const gemmi::Connection &c) {
auto is_site = [&](const gemmi::AtomAddress &a) {
return a.chain_name == e.chain && a.res_id.seqid == e.seqid && a.res_id.name == from;
};
return is_site(c.partner1) || is_site(c.partner2);
});
for (const gemmi::Chain &ch : model.models[0].chains)
for (const gemmi::Residue &r : ch.residues)
if (ch.name == e.chain && r.seqid == e.seqid && r.name == e.species)
AddMetalRecords(model, ch, r, record);
}
{
const RefinementStage stage("blobs");
result.blobs = FindDensityBlobs(*result.difference_map, result.difference_map_rms, model, result.d_min,
result.anomalous_map.get(), result.anomalous_map_rms);
}
if (settings.solvent) {
const RefinementStage stage("solvent");
const SolventMaps maps{result.map.get(), result.map_rms, result.difference_map.get(),
result.difference_map_rms, result.anomalous_map.get(), result.anomalous_map_rms,
wavelength_A};
result.solvent = AnalyseSolventSites(model, result.d_min, maps, ctx.nthreads);
}
result.b_ceiling = ctx.b_ceiling;
result.k_sol = ctx.k_sol;
result.b_sol = ctx.b_sol;
if (build)
for (gemmi::Chain &ch : model.models[0].chains)
for (gemmi::Residue &r : ch.residues)
for (gemmi::Atom &a : r.atoms)
a.serial = 0; // SolventView's indices
result.model = std::make_shared<gemmi::Structure>(std::move(model));
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() *
(ncs_p != nullptr ? 2 : 1),
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<SlotPlace> devices = settings.restrained ? OwnSlotPlaces(ctx, wanted - 1, settings)
: std::vector<SlotPlace>();
const size_t slot_threads = std::max<size_t>(1, settings.nthreads / (devices.size() + 1));
ScalePools pools;
#ifdef JFJOCH_USE_CUDA
if (ctx.scale_gpu)
pools.pool[ctx.device] = ctx.scale_gpu;
#endif
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].device >= 0 ? devices[k].device : 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].device, devices[k].own_lane, slot_threads, settings,
pools);
} 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 - pool.n; // not the placeholders
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, {} records "
"beyond the inner shell left as contacts), {} "
"angles, {} torsions, {} chiral volumes, {} planes",
result.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.metal_links_refused, topo.angles, topo.torsions, topo.chirals, topo.planes);
if (ncs_p != nullptr) {
result.ncs_groups = ncs.groups;
result.ncs_copies = ncs.copies;
result.ncs_restraints = static_cast<int>(ncs.distances.records.size());
result.ncs_chains = ncs.description;
result.ncs_rmsd_start = NcsRmsd(ncs, ctx.XyzB());
logger.Info("Refinement: NCS, {} molecule{} in {} copies ({}), {} local distance restraints; copies {:.3f} A "
"apart (C-alpha/P rmsd after superposition)", ncs.groups, ncs.groups == 1 ? "" : "s",
ncs.copies, ncs.description, result.ncs_restraints, result.ncs_rmsd_start);
}
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);
}
// (the placement moves rigid bodies, which leaves the bonds and angles as given)
given.trial.geometry = result.geometry_start;
given.trial.within_bounds = given.trial.geometry.rmsd_bonds <= MAX_RMSD_BONDS &&
given.trial.geometry.rmsd_angles <= MAX_RMSD_ANGLES;
given.trial.adp = result.adp_input;
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());
// Trials whose need is decided by others are run with them where the slots take them all at once - but not
// where the slots share one card, whose time they would take from the trials that are needed (a model of 43k
// atoms: the ladder 10 s ahead, 5 s not) - and left out where it is decided they were not needed, with their
// evaluations, which do not count: the outcome and the count are the same whatever the slots.
std::set<int> cards;
for (const TrialSlot &s : trial_slots)
cards.insert(s.ctx->device);
const bool one_card = cards.size() == 1 && *cards.begin() >= 0;
auto ahead = [&](size_t trials) { return !one_card && trial_slots.size() >= trials; };
int ahead_evaluations = 0;
auto drop_ahead = [&](TrialOutcome &o) {
ahead_evaluations += o.evaluations;
o = TrialOutcome();
};
// 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 best_of = [&](const std::vector<TrialOutcome> &outcomes) {
int best = -1;
for (size_t k = 0; k < outcomes.size(); k++)
if (outcomes[k].ran && (best < 0 || displaces(outcomes[k].trial, outcomes[best].trial)))
best = static_cast<int>(k);
return best;
};
auto choose = [&](std::vector<TrialOutcome> &outcomes, const char *stage) {
for (const TrialOutcome &o : outcomes)
if (o.ran) {
logger.Info("Refinement, {}: {}", stage, o.line);
result.trials.push_back(o.trial);
}
return best_of(outcomes);
};
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; where there are copies, each rung without NCS restraints
// and then with them - those with them at the B factors' weight the rungs without chose, which NCS does
// not change (on 39 refinements of 21 structures with copies it never did; the X-ray weight it does), and
// at the rung they chose and the next looser one (on 198 refinements of the battery with copies, NCS was
// never kept at any other).
// The B factors are refined first at the middle rung only, and at the others where that comes within two
// ties of holding them there: on 364 refinements of the battery B held was kept by 281, and the rule
// changed the rung kept on 4 (all within 0.006 in R-free).
std::vector<TrialWeights> ladder;
for (double rung : settings.weight_ladder)
for (double adp : settings.adp_ladder) {
ladder.push_back({w0 * rung, adp_w0 * adp, 0});
if (ncs_p != nullptr)
ladder.push_back({w0 * rung, adp_w0 * adp, settings.ncs_sigma});
}
const std::vector<AtomAniso> aniso0 = ctx.aniso;
std::vector<TrialOutcome> outcomes(ladder.size());
auto run_ladder = [&](const std::function<bool(const TrialWeights &)> &these) {
std::vector<TrialWeights> t;
for (const TrialWeights &l : ladder)
if (these(l))
t.push_back(l);
std::vector<TrialOutcome> ran = RunTrials(trial_slots, t, refinable, pos0, b0, aniso0, nullptr,
settings.macro_cycles, settings, ncs_p);
for (size_t i = 0, m = 0; i < ladder.size(); i++)
if (these(ladder[i]))
outcomes[i] = std::move(ran[m++]);
};
const double mid = w0 * settings.weight_ladder[settings.weight_ladder.size() / 2];
auto first = [&](const TrialWeights &l) { return l.ncs == 0 && (l.adp == 0 || l.w == mid); };
const bool all_ahead = ahead(ladder.size());
run_ladder([&](const TrialWeights &l) { return first(l) || all_ahead; });
// B refined against B held at the middle rung (where both ran)
const TrialOutcome *held = nullptr, *refined = nullptr;
for (size_t i = 0; i < ladder.size(); i++)
if (ladder[i].w == mid && ladder[i].ncs == 0)
(ladder[i].adp == 0 ? held : refined) = &outcomes[i];
const bool refine_b = held == nullptr || refined == nullptr ||
(refined->ran && refined->trial.within_bounds &&
(!held->trial.within_bounds || refined->trial.r_free < held->trial.r_free + 2 * R_FREE_TIE));
auto rest_b = [&](const TrialWeights &l) { return l.ncs == 0 && !first(l); };
if (refine_b && !all_ahead)
run_ladder(rest_b);
// the best of those without NCS restraints, those with them set aside
std::vector<TrialOutcome> with_ncs(ladder.size());
for (size_t i = 0; i < ladder.size(); i++)
if (ladder[i].ncs > 0 || (!refine_b && rest_b(ladder[i])))
std::swap(with_ncs[i], outcomes[i]);
if (const int b = best_of(outcomes); ncs_p != nullptr && b >= 0) {
const double adp = outcomes[b].trial.adp_weight, w = outcomes[b].trial.weight;
double looser = w;
for (double rung : settings.weight_ladder)
if (w0 * rung > w && (looser == w || w0 * rung < looser))
looser = w0 * rung;
auto these = [&](const TrialWeights &l) { return l.ncs > 0 && l.adp == adp && (l.w == w || l.w == looser); };
for (size_t i = 0; i < ladder.size(); i++)
if (these(ladder[i]))
std::swap(with_ncs[i], outcomes[i]);
if (!all_ahead)
run_ladder(these);
}
for (TrialOutcome &o : with_ncs)
drop_ahead(o);
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 only in the first step, from the
// ladder's coarse rung: on 364 refinements of the battery no later step kept it, of 91 that tried it) and
// the B factors' divided and multiplied by four; where there are copies, also the NCS restraints changed:
// switched on where they are off, and in turn tightened and loosened (to off beyond NCS_SIGMA_MAX) where
// they are on.
// 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).
// 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;
// `fresh`: the first step fits the scale, the solvent and the likelihood anew (the model changed since).
auto trajectory = [&](TrialOutcome state, bool fresh, std::vector<TrialOutcome> &from_best) {
TrialOutcome best = state;
from_best.clear();
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, ncs_s = state.trial.ncs_sigma;
std::vector<TrialWeights> trials = {{w, adp, ncs_s}};
if (step % 2 == 1 && adp > 0) {
trials.push_back({w, adp / 4, ncs_s});
trials.push_back({w, 4 * adp, ncs_s});
} else {
if (step == 0)
trials.push_back({w / 2, adp, ncs_s});
if (2 * w <= w_max)
trials.push_back({2 * w, adp, ncs_s});
}
if (ncs_p != nullptr) {
const double changed = ncs_s == 0 ? settings.ncs_sigma : step % 2 == 0 ? ncs_s / 2 : 2 * ncs_s;
trials.push_back({w, adp, changed <= NCS_SIGMA_MAX ? changed : 0});
}
std::vector<TrialOutcome> outcomes = RunTrials(trial_slots, trials, refinable, state.pos, state.b, state.aniso,
fresh && step == 0 ? nullptr : &state.fit,
settings.macro_cycles, settings, ncs_p);
if (state_is_best)
from_best = outcomes;
const int k = choose(outcomes, fmt::format("step {}", step + 1).c_str());
if (k < 0)
break;
const RefinementWeightTrial &t = outcomes[k].trial;
// after a change of the model the first step is taken whatever it gains: it is where the change is
// judged from
const bool gains = (fresh && step == 0) || displaces(t, state.trial) ||
(t.within_bounds && t.r_work < state.trial.r_work - R_FREE_TIE);
if (!gains || !((fresh && step == 0) || (t.within_bounds && t.r_free < best.trial.r_free + R_FREE_TIE)))
break;
state = std::move(outcomes[k]);
state_is_best = (fresh && step == 0) || displaces(state.trial, best.trial);
if (state_is_best) {
best = state;
from_best.clear();
}
}
return best;
};
state = trajectory(std::move(state), false, from_best);
// 2a. The solvent build (SolventBuild.h), from where the trajectory ended: waters put in at the difference
// map's peaks, taken out where the maps no longer hold them, and waters and ions relabelled, on the maps
// and the solvent audit of that model; then the solvent alone refined, the rest held, at the weights the
// trajectory ended at. A round is kept where that lowers R-free, and undone otherwise (the build ends
// there); at most SOLVENT_BUILD_ROUNDS of them. What they kept is then refined with the rest by the final
// ladder (the trajectory run again from there gained nothing on twelve sets). Each round thus costs a map,
// an audit and a refinement of the solvent, where a refinement of its own cost the whole weight search
// (the build was 2-4 times the refinement).
if (build && state.ran && !ctx.Cancelled()) {
const auto build_start = std::chrono::steady_clock::now();
const RefinementStage build_stage("solvent_build");
SolventBuildResult &b = result.solvent_build;
b.ran = true;
std::set<std::string> added; // the waters put in and still there, by name
std::set<std::string> relabelled; // the sites relabelled, by their new names
std::vector<SolventEdit> not_solvent;
// per slot, the residue of each atom (the atoms stay where they are in memory)
std::vector<std::vector<gemmi::Residue *>> residue_of(trial_slots.size());
for (size_t sl = 0; sl < trial_slots.size(); sl++)
for (gemmi::Chain &ch : trial_slots[sl].ctx->Model().chains)
for (gemmi::Residue &r : ch.residues)
for (size_t a = 0; a < r.atoms.size(); a++)
residue_of[sl].push_back(&r);
for (int round = 1; !ctx.Cancelled(); round++) {
if (round > SOLVENT_BUILD_ROUNDS) {
b.stop = fmt::format("{} rounds", SOLVENT_BUILD_ROUNDS);
break;
}
// the maps and the audit of where the trajectory ended, on all of the threads (the trials left the
// main slot its share of them)
ctx.nthreads = main_threads;
ctx.SetPositions(state.pos);
ctx.SetB(state.b);
PlaceRiders(ctx, topo);
// the maps' Refit() from the trajectory's own fit, whichever slot ran it last
ctx.Restore(state.fit);
RefinementResult scratch;
MakeMaps(ctx, merged, wavelength_A, scratch);
std::optional<RefinementStage> pick_stage("build.view");
const gemmi::Structure view = SolventView(ctx.st, pool);
const SolventMaps maps{scratch.map.get(), scratch.map_rms, scratch.difference_map.get(),
scratch.difference_map_rms, scratch.anomalous_map.get(),
scratch.anomalous_map_rms, wavelength_A};
pick_stage.emplace("build.audit");
const std::vector<SolventSite> sites = AnalyseSolventSites(view, result.d_min, maps, settings.nthreads);
pick_stage.emplace("build.pick");
SolventRound picked = PickSolvent(view, maps, sites, added, relabelled, ctx.b_ceiling, round);
pick_stage.reset();
not_solvent.clear();
for (const SolventEdit &e : picked.edits)
if (e.kind == SolventEdit::Kind::NotSolvent)
not_solvent.push_back(e);
b.not_solvent = picked.not_solvent;
std::string refused;
for (const auto &[why, n] : picked.refused)
refused += fmt::format("{}{} {}", refused.empty() ? "" : ", ", n, why);
logger.Info("Solvent build, round {}: {} waters to add, {} to remove, {} sites to relabel, {} peaks "
"not solvent{}", round, picked.added, picked.removed, picked.relabelled,
picked.not_solvent, refused.empty() ? std::string() : "; peaks refused: " + refused);
if (picked.added + picked.removed + picked.relabelled == 0) {
b.stop = "nothing to change";
break;
}
// the atoms the round edits (an index in the model of every slot), each with what it was
std::map<std::string, int> by_name; // "HOH A401" (and "|altloc" where it has one) -> atom index
for (const gemmi::Chain &ch : view.models[0].chains)
for (const gemmi::Residue &r : ch.residues)
for (const gemmi::Atom &a : r.atoms)
if (!a.is_hydrogen())
by_name[r.name + " " + ch.name + r.seqid.str() +
(a.altloc ? std::string("|") + a.altloc : std::string())] = a.serial;
struct Undo {
int index;
gemmi::Atom atom;
std::string resname;
gemmi::EntityType entity;
};
std::vector<std::vector<Undo>> undo(trial_slots.size());
std::vector<RefinementTopology> topo_before(trial_slots.size());
const SolventPool pool_before = pool;
// the edits, in order, made below in every slot - each slot on a thread of its own
using Edit = std::function<void(gemmi::Atom &, gemmi::Residue &, size_t)>;
std::vector<std::pair<int, Edit>> edits;
std::vector<char> retyped(ctx.atoms.size(), 0); // the atoms whose B similarities go
auto edit_atom = [&](int index, Edit f) { edits.emplace_back(index, std::move(f)); };
std::vector<SolventEdit> kept;
for (SolventEdit &e : picked.edits) {
if (e.kind == SolventEdit::Kind::Added) {
const int k = pool.Free();
if (k < 0)
continue; // no placeholder left: the weaker peaks wait
pool.chain[k] = e.chain;
pool.seqid[k] = gemmi::SeqId(++pool.last[e.chain], ' ');
e.site = "HOH " + e.chain + pool.seqid[k].str();
// on a special position (PickSolvent put it exactly there), held on it
SpecialPosition site = SiteSymmetry(view.cell, e.pos);
site.atom = pool.first + k;
const int index = pool.first + k;
edit_atom(index, [&trial_slots, e, index, site](gemmi::Atom &a, gemmi::Residue &, size_t sl) {
a.pos = e.pos;
a.occ = static_cast<float>(e.occupancy);
a.b_iso = static_cast<float>(e.b_iso);
RefinementTopology &t = *trial_slots[sl].topo;
t.refinable[index] = 1;
if (!site.ops.empty()) {
t.special_index[index] = static_cast<int>(t.special.size());
t.special.push_back(site);
}
});
} else if (e.kind == SolventEdit::Kind::Removed) {
const auto it = by_name.find(e.site + (e.altloc ? std::string("|") + e.altloc : std::string()));
if (it == by_name.end())
continue;
edit_atom(it->second, [&trial_slots, index = it->second](gemmi::Atom &a, gemmi::Residue &, size_t sl) {
a.occ = 0;
trial_slots[sl].topo->refinable[index] = 0;
});
} else if (e.kind == SolventEdit::Kind::Relabelled) {
const auto it = by_name.find(e.site);
if (it == by_name.end())
continue;
const int i = it->second;
// its inner sphere, for a metal's bonds, read in the view and numbered as the slots number it
std::vector<int> donors;
if (gemmi::Element(SolventSpeciesElement(e.species)).is_metal())
for (const gemmi::Chain &ch : view.models[0].chains)
for (const gemmi::Residue &r : ch.residues)
if (ch.name == e.chain && r.seqid == e.seqid && r.atoms.size() == 1) {
gemmi::Residue as_metal = r;
as_metal.atoms[0].element = gemmi::Element(SolventSpeciesElement(e.species));
for (const InnerSphereDonor &d : InnerSphereDonors(view.models[0], r, as_metal.atoms[0]))
donors.push_back(view.models[0].chains[d.chain].residues[d.residue].atoms[d.atom].serial);
}
retyped[i] = 1;
edit_atom(i, [&trial_slots, e, i, donors](gemmi::Atom &a, gemmi::Residue &r, size_t sl) {
const bool water = e.species == "HOH";
r.name = e.species;
r.entity_type = water ? gemmi::EntityType::Water : gemmi::EntityType::NonPolymer;
a.element = gemmi::Element(SolventSpeciesElement(e.species));
a.name = water ? "O" : e.species;
a.charge = ScatteringCharge(r, a);
RetypeAtom(*trial_slots[sl].topo, r, a, i, donors);
});
} else
continue;
kept.push_back(e);
}
std::optional<RefinementStage> edit_stage("build.edit");
std::vector<std::thread> editors;
for (size_t sl = 0; sl < trial_slots.size(); sl++)
editors.emplace_back([&, sl]() {
topo_before[sl] = *trial_slots[sl].topo;
for (const auto &[index, f] : edits) {
gemmi::Residue &r = *residue_of[sl][index];
gemmi::Atom &a = *trial_slots[sl].ctx->atoms[index];
undo[sl].push_back({index, a, r.name, r.entity_type});
f(a, r, sl);
}
if (std::find(retyped.begin(), retyped.end(), 1) != retyped.end())
DropBRestraints(*trial_slots[sl].topo, retyped);
});
for (std::thread &t : editors)
t.join();
edit_stage.reset();
refinable.clear();
for (size_t i = 0; i < ctx.atoms.size(); i++)
if (topo.refinable[i])
refinable.push_back(static_cast<int>(i));
// the solvent alone refined, the rest held, at the weights the trajectory ended at: what the round's
// changes are judged by
std::vector<int> solvent;
for (int i : refinable) {
const gemmi::Residue &r = *residue_of[0][i];
if (r.is_water() || (r.atoms.size() == 1 && (r.atoms[0].element.is_metal() || r.atoms[0].element == gemmi::El::Cl)))
solvent.push_back(i);
}
const std::vector<gemmi::Position> changed_pos = ctx.Positions();
const std::vector<double> changed_b = ctx.B();
// the round as built, nothing refined: where refining the solvent alone loses what the sites gained,
// it is judged as built
ctx.Restore(state.fit);
const RFactors built = ctx.Refit();
const Context::Fit built_fit = ctx.Fitted();
// (B refined, at its balanced weight where the trajectory held B: a new water's starting B is a
// guess; fewer iterations, the waters' few parameters converging fast)
RefinementSettings solvent_settings = settings;
solvent_settings.xyz_iterations = SOLVENT_ITERATIONS;
solvent_settings.b_iterations = SOLVENT_ITERATIONS;
std::vector<TrialOutcome> local =
RunTrials(trial_slots, {{state.trial.weight, state.trial.adp_weight > 0 ? state.trial.adp_weight : adp_w0,
state.trial.ncs_sigma}},
solvent, changed_pos, changed_b, state.aniso, nullptr, settings.macro_cycles, solvent_settings,
ncs_p);
logger.Info("Solvent build, round {}: the solvent refined alone: {}", round, local[0].line);
TrialOutcome after = std::move(local[0]);
if (after.ran && !(after.trial.r_free < state.trial.r_free) && built.r_free < state.trial.r_free) {
logger.Info("Solvent build, round {}: kept as built, R-work {:.4f}, R-free {:.4f}", round,
built.r_work, built.r_free);
after.pos = changed_pos;
after.b = changed_b;
after.aniso = state.aniso;
after.fit = built_fit;
after.trial.r_work = built.r_work;
after.trial.r_free = built.r_free;
after.trial.geometry = state.trial.geometry;
after.trial.within_bounds = state.trial.within_bounds;
}
if (!after.ran || !(after.trial.r_free < state.trial.r_free)) {
logger.Info("Solvent build, round {}: not kept (R-free {:.4f}, before {:.4f})", round,
after.trial.r_free, state.trial.r_free);
b.stop = fmt::format("round {} raised R-free", round);
// undone, in every slot
for (size_t sl = 0; sl < trial_slots.size(); sl++) {
for (auto u = undo[sl].rbegin(); u != undo[sl].rend(); ++u) {
*trial_slots[sl].ctx->atoms[u->index] = u->atom;
residue_of[sl][u->index]->name = u->resname;
residue_of[sl][u->index]->entity_type = u->entity;
}
*trial_slots[sl].topo = topo_before[sl];
}
pool = pool_before;
refinable.clear();
for (size_t i = 0; i < ctx.atoms.size(); i++)
if (topo.refinable[i])
refinable.push_back(static_cast<int>(i));
break;
}
for (const SolventEdit &e : kept) {
if (e.kind == SolventEdit::Kind::Added)
added.insert(e.site);
else if (e.kind == SolventEdit::Kind::Removed) {
if (added.erase(e.site) == 0)
b.removed++;
} else {
const std::string now = e.species + e.site.substr(e.site.find(' '));
if (added.erase(e.site) > 0)
added.insert(now);
relabelled.insert(now);
relabels.push_back(e);
b.relabelled++;
}
b.edits.push_back(e);
}
b.rounds = round;
// the solvent's B refined for the judge only: the rest goes on with the B factors as the trajectory
// had them
after.trial.adp_weight = state.trial.adp_weight;
state = std::move(after);
from_best.clear();
}
b.added = static_cast<int>(added.size());
b.edits.insert(b.edits.end(), not_solvent.begin(), not_solvent.end());
b.seconds = std::chrono::duration<double>(std::chrono::steady_clock::now() - build_start).count();
logger.Info("Solvent build: {} waters added, {} removed, {} sites relabelled in {} round{} kept ({}), {} "
"peaks not solvent; R-free {:.4f}; {:.1f} s", b.added, b.removed, b.relabelled, b.rounds,
b.rounds == 1 ? "" : "s", b.stop, b.not_solvent, state.trial.r_free, b.seconds);
}
// 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, ncs_s = state.trial.ncs_sigma;
std::vector<TrialWeights> trials;
for (double f : settings.final_ladder)
if (w * f <= w_max)
trials.push_back({w * f, adp, ncs_s});
outcomes.assign(trials.size(), TrialOutcome());
std::vector<TrialWeights> missing;
for (size_t i = 0; i < trials.size(); i++) {
for (const TrialOutcome &o : from_best)
if (o.ran && o.trial.weight == trials[i].w && o.trial.adp_weight == trials[i].adp &&
o.trial.ncs_sigma == trials[i].ncs)
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.aniso,
&state.fit, settings.macro_cycles, settings, ncs_p);
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");
if (k >= 0 && displaces(outcomes[k].trial, state.trial))
state = std::move(outcomes[k]);
}
// 4. The displacement model, decided late, on the free set: from where the weights ended, TLS bodies (the
// model's own groups, one per chain, the chains in TLSMD segments) with a residual B, and a U per atom at
// its balanced weight and a quarter of it - each trial as long as a step, on the positions where the weights
// ended, held, so that the trials differ in the displacement model alone; the model chosen then takes a
// step with the positions refined, kept where it displaces its trial (against trials a macro-cycle longer
// that refined the positions too, on 50 refinements of the battery R-free 0.001 lower on average and none
// higher by more than 0.0015, at 78% of the evaluations). Where the geometry is out of bounds there, which
// held positions cannot mend, each trial refines them too, a macro-cycle longer than a step, and no step
// follows. One displaces the one before it (the first
// where the weights ended, the B factors as they are) only where its geometry is within bounds and both its
// R-free (beyond the tie) and the -ln L of the free reflections are lower: the free set is the test of the
// parameters it adds.
// Following Merritt (2012) Acta Cryst. D68, 468-477 (the displacement model the data support, decided by test)
// The trials that would not change the choice are left out: the model's own TLS groups are tried only where
// one body per chain or the finest cut lowers R-free, the cuts between those only where also the finest
// fits the free set better, the U at its balanced weight only where the U at a quarter of it lowers R-free
// or the -ln L of the free reflections (run ahead where the slots take them), and the U at four times its
// balanced weight not at all (on 49 refinements of 33 structures, no choice taken changed).
if (settings.adp_hypotheses && state.ran && !ctx.Cancelled()) {
const RefinementStage adp_stage("adp_hypotheses");
ctx.SetPositions(state.pos);
ctx.SetB(state.b);
ctx.SetAniso(state.aniso);
ctx.Restore(state.fit);
PlaceRiders(ctx, topo);
state.trial.nll_free = ctx.Score().nll_free;
logger.Info("Refinement, displacements: from where the weights ended, R-free {:.4f}, -LL free {:.1f}",
state.trial.r_free, state.trial.nll_free);
std::vector<AdpHypothesis> hyps = AdpHypotheses(ctx, topo, input_tls);
const double w = state.trial.weight, adp = state.trial.adp_weight > 0 ? state.trial.adp_weight : adp_w0;
// the U weight balanced on the model with every refinable atom given a U
double u_w0 = adp;
{
const std::vector<AtomAniso> keep = ctx.aniso;
std::vector<AtomAniso> all = keep.empty() ? std::vector<AtomAniso>(ctx.atoms.size()) : keep;
std::vector<char> refines_u(ctx.atoms.size(), 0);
std::vector<int> u_atoms;
for (int i : refinable)
if (!ctx.atoms[i]->is_hydrogen()) {
all[i].on = true;
refines_u[i] = 1;
u_atoms.push_back(i);
}
ctx.SetAniso(all);
u_w0 = BalancedUWeight(ctx, MakeURestraints(hyps.back().pairs, refines_u), u_atoms);
ctx.SetAniso(keep);
}
const double ncs_s = state.trial.ncs_sigma;
// every trial in order of trust; the model's own TLS groups and the cuts between one TLS body per chain
// and the finest, and the U at its balanced weight, run after the others where those ask for them
std::vector<TrialWeights> trials;
std::vector<const AdpHypothesis *> which;
int own = -1, coarsest = -1, finest = -1, u_quarter = -1;
for (const AdpHypothesis &h : hyps) {
if (h.kind == AdpHypothesis::Kind::Anisotropic) {
u_quarter = static_cast<int>(trials.size());
for (double f : {0.25, 1.0}) {
trials.push_back({w, u_w0 * f, ncs_s});
which.push_back(&h);
}
} else {
if (h.segments == 0)
own = static_cast<int>(trials.size());
else {
if (coarsest < 0)
coarsest = static_cast<int>(trials.size());
finest = static_cast<int>(trials.size());
}
trials.push_back({w, adp, ncs_s});
which.push_back(&h);
}
}
std::vector<char> later(trials.size(), 0);
for (int k = coarsest + 1; k < finest; k++)
later[k] = 1;
if (own >= 0 && coarsest >= 0)
later[own] = 1;
if (u_quarter >= 0)
later[u_quarter + 1] = 1;
outcomes.assign(trials.size(), TrialOutcome());
// the positions held only where the geometry is within its bounds where the weights ended: a trial that
// holds them cannot bring it there, one that refines them can (and is then the whole trial, no step after)
const bool hold = state.trial.within_bounds;
auto run = [&](const std::vector<char> &these) {
std::vector<TrialWeights> t;
std::vector<const AdpHypothesis *> h;
for (size_t k = 0; k < trials.size(); k++)
if (these[k]) {
t.push_back(trials[k]);
h.push_back(which[k]);
}
if (t.empty())
return;
std::vector<TrialOutcome> ran = RunTrials(trial_slots, t, refinable, state.pos, state.b, state.aniso,
&state.fit, settings.macro_cycles + (hold ? 0 : 1), settings,
ncs_p, h, hold);
for (size_t k = 0, m = 0; k < trials.size(); k++)
if (these[k])
outcomes[k] = std::move(ran[m++]);
};
const bool later_ahead = ahead(trials.size());
std::vector<char> now(trials.size());
for (size_t k = 0; k < trials.size(); k++)
now[k] = !later[k] || later_ahead;
run(now);
// the model's own groups where one body per chain or the finest cut lowers R-free from where the weights
// ended, the cuts between where also the finest fits the free set better than one body per chain, the U
// at its balanced weight where the U at a quarter of it lowers R-free or -ln L free (R-free alone left out
// the trial that would have been kept where the quarter lowered -ln L free but not R-free)
auto lowers = [&](int k) {
return k >= 0 && outcomes[k].ran && outcomes[k].trial.within_bounds &&
outcomes[k].trial.r_free < state.trial.r_free;
};
// the U at a quarter fitting the free reflections better than where the weights ended asks for the
// balanced weight as well: the choice below is made on -ln L free as much as on R-free
auto fits_better = [&](int k) {
return k >= 0 && outcomes[k].ran && outcomes[k].trial.within_bounds &&
outcomes[k].trial.nll_free < state.trial.nll_free;
};
const bool tls = lowers(coarsest) || lowers(finest), u_full = lowers(u_quarter) || fits_better(u_quarter);
const bool cuts = tls && finest > coarsest && outcomes[finest].ran && outcomes[coarsest].ran &&
outcomes[finest].trial.nll_free < outcomes[coarsest].trial.nll_free;
std::vector<char> needed(trials.size(), 0);
for (int k = 0; k < static_cast<int>(trials.size()); k++)
needed[k] = later[k] && (k == u_quarter + 1 ? u_full : k == own ? tls : cuts);
if (!later_ahead)
run(needed);
for (size_t k = 0; k < trials.size(); k++)
if (later[k] && !needed[k])
drop_ahead(outcomes[k]);
int best = -1; // where the weights ended
auto trial = [&](int k) -> const RefinementWeightTrial & { return k < 0 ? state.trial : outcomes[k].trial; };
for (size_t k = 0; k < outcomes.size(); k++) {
const RefinementWeightTrial &t = outcomes[k].trial;
if (!outcomes[k].ran)
continue;
logger.Info("Refinement, displacements: {}", outcomes[k].line);
result.adp_trials.push_back(t);
if (t.within_bounds && (!trial(best).within_bounds ||
(t.r_free < trial(best).r_free - R_FREE_TIE && t.nll_free < trial(best).nll_free)))
best = static_cast<int>(k);
}
result.adp_model = result.adp_input == "isotropic" ? "isotropic" : "the model's U held";
if (best >= 0) {
result.adp_model = outcomes[best].trial.adp;
state = std::move(outcomes[best]);
}
if (best >= 0 && hold) {
// the model chosen refined on with the positions, a step from where its trial ended, kept where it
// displaces that
AdpHypothesis chosen = *which[best];
chosen.bodies = state.tls;
chosen.tls_b = state.tls_b;
std::vector<TrialOutcome> on = RunTrials(trial_slots, {trials[best]}, refinable, state.pos, state.b,
state.aniso, &state.fit, settings.macro_cycles, settings,
ncs_p, {&chosen});
if (on[0].ran) {
logger.Info("Refinement, displacements, positions refined: {}", on[0].line);
if (displaces(on[0].trial, state.trial))
state = std::move(on[0]);
}
}
logger.Info("Refinement: displacement model {}", result.adp_model);
}
// 5. The model as given, first in order of trust: the refinement is kept where it displaces it by the rule every
// trial is held to, else the model is given back as it came (the solvent build's edits undone), so that what
// comes out is never worse on the free set than what went in - where the free set is free to the model as
// given (RefinementSettings::as_given_competes).
if (settings.as_given_competes && state.ran && !ctx.Cancelled() && !displaces(state.trial, given.trial)) {
logger.Info("Refinement: the model as given kept (R-work {:.4f}, R-free {:.4f}), the refinement did not "
"displace it (R-work {:.4f}, R-free {:.4f})", given.trial.r_work, given.trial.r_free,
state.trial.r_work, state.trial.r_free);
size_t a = 0, k = 0;
for (gemmi::Chain &ch : ctx.Model().chains)
for (gemmi::Residue &res : ch.residues) {
std::tie(res.name, res.entity_type) = given_residues[k++];
for (gemmi::Atom &at : res.atoms)
at = given_atoms[a++];
}
pool = given_pool;
relabels.clear();
result.solvent_build = SolventBuildResult();
result.adp_model = result.adp_input == "isotropic" ? "isotropic" : "the model's U held";
result.as_given = true;
state = given;
}
for (size_t s = 1; s < trial_slots.size(); s++)
ctx.evaluations += trial_slots[s].ctx->evaluations;
ctx.evaluations -= ahead_evaluations;
if (ahead_evaluations > 0)
logger.Debug("Refinement: {} target evaluations in trials run ahead and not needed", ahead_evaluations);
ctx.nthreads = main_threads;
if (state.ran) {
ctx.SetPositions(state.pos);
ctx.SetB(state.b);
ctx.SetAniso(state.aniso);
ctx.tls = state.tls;
ctx.tls_b = state.tls_b;
PlaceRiders(ctx, topo);
// the maps' Refit() from the kept trial's own fit (the solvent grid and the likelihood's D it starts
// from), not from whichever trial the main slot ran last: the same whatever the slots
ctx.Restore(state.fit);
result.restrained = true;
result.weight = state.trial.weight;
result.adp_weight = state.trial.adp_weight;
result.ncs_sigma = state.trial.ncs_sigma;
if (ncs_p != nullptr) {
result.ncs_rmsd = NcsRmsd(ncs, ctx.XyzB());
result.ncs_released = NcsReleasedResidues(ncs, ctx.XyzB());
logger.Info("Refinement: NCS {}; copies {:.3f} A apart (as placed {:.3f} A), {} residues released",
result.ncs_sigma > 0 ? fmt::format("restrained at {:.3g} A", result.ncs_sigma)
: std::string("not kept, the copies refined independently"),
result.ncs_rmsd, result.ncs_rmsd_start, result.ncs_released);
}
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;
}
// The R-free a refinement ended at: the restrained stage's, or the placement's where that did not run.
double FinalRFree(const RefinementResult &r) {
return r.restrained ? r.r_free : r.r_free_placement;
}
double FinalRWork(const RefinementResult &r) {
return r.restrained ? r.r_work : r.r_work_placement;
}
// The twin hypothesis beside the untwinned refinement `result`: the placed model refined again, solvent build and
// displacement model included, against the data as twinned by a law of the lattice, and the twinned refinement taken
// (into `result`) where the law beats its own null (TwinNull) on the free reflections by TWIN_NULL_Z standard errors,
// or by TWIN_NULL_Z_DATA where the caller says the data read twinned (AdoptTwinLaw) - not on its R-free against the
// untwinned one's, which compares two refinement paths and so carries their noise - so that neither the drop in R
// any twin target buys by averaging, nor a pseudo-symmetry the model already carries (it predicts I(Th) ~ I(h), and
// the law then changes nothing) is read as a twin. A law is tested only where the free set keeps its twin pairs
// together; of several, the one the placed model asks the largest fraction of.
void DecideTwin(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, RefinementResult &result) {
const RefinementStage stage("twin");
result.untwinned_r_work = FinalRWork(result);
result.untwinned_r_free = FinalRFree(result);
const std::vector<gemmi::Op> laws = RefinementTwinLaws(cell, sg);
result.twin_laws = static_cast<int>(laws.size());
if (laws.empty()) {
result.twin_decision = "NO_TWIN_LAW";
return;
}
// the laws whose twin pairs the free set keeps together
RefinementReflections refl = PrepareRefinementReflections(merged, cell, sg, RefinementDMin(MaxInvD2(merged)));
std::vector<gemmi::Op> usable;
for (const gemmi::Op &law : laws) {
const int split = SetTwinLaw(refl, cell, sg, RefinementDMin(MaxInvD2(merged)), law);
if (split == 0)
usable.push_back(law);
else {
logger.Info("Refinement, twin law {}: not tested - {} of its twin pairs are split between the working and "
"the free set", law.as_hkl().triplet('h'), split);
result.twin_law = law.as_hkl().triplet('h');
result.twin_split_pairs = split;
}
}
if (usable.empty()) {
result.twin_decision = "FREE_SET_SPLIT";
return;
}
result.twin_split_pairs = 0;
gemmi::Op law = usable[0];
if (usable.size() > 1) {
RefinementSettings score = settings;
score.score_only = true;
double best = -1;
for (const gemmi::Op &candidate : usable) {
const RefinementResult s = RefineOnce(placed, merged, cell, sg, wavelength_A, logger, score, &candidate);
logger.Info("Refinement, twin law {}: the placed model asks a fraction of {:.3f}",
candidate.as_hkl().triplet('h'), s.score_start.twin_fraction);
if (s.ok && s.score_start.twin_fraction > best) {
best = s.score_start.twin_fraction;
law = candidate;
}
}
}
result.twin_law = law.as_hkl().triplet('h');
RefinementResult twinned = RefineOnce(placed, merged, cell, sg, wavelength_A, logger, settings, &law);
if (!twinned.ok || twinned.cancelled) {
result.twin_decision = "UNTWINNED";
return;
}
result.twin_fraction = twinned.twin_fraction;
result.twin_r_work = FinalRWork(twinned);
result.twin_r_free = FinalRFree(twinned);
result.twin_r_free_null = twinned.twin_r_free_null;
result.twin_fraction_null = twinned.twin_fraction_null;
result.twin_null_z = twinned.twin_null_z;
const bool adopt = AdoptTwinLaw(twinned.twin_fraction, result.twin_null_z, settings.twin_data_indicated);
result.twin_decision = twinned.twin_fraction == 0 ? "ZERO_FRACTION" : adopt ? "TWINNED" : "UNTWINNED";
if (twinned.twin_fraction > 0)
result.twin_rule = adopt && result.twin_null_z < TWIN_NULL_Z ? "NULL+LTEST" : "NULL";
logger.Info("Refinement, twin law {}: fraction {:.3f}, R-work {:.4f}, R-free {:.4f}; untwinned R-free {:.4f}; null "
"(mates unrelated, fraction {:.3f}) R-free {:.4f}, the law ahead of it by {:.1f} standard errors{} -> {}",
result.twin_law, result.twin_fraction, result.twin_r_work, result.twin_r_free, result.untwinned_r_free,
result.twin_fraction_null, result.twin_r_free_null, result.twin_null_z,
settings.twin_data_indicated ? ", the data read twinned" : "", result.twin_decision);
if (!adopt)
return;
// the twinned refinement is the result, with the evidence that chose it
twinned.twin_decision = result.twin_decision;
twinned.twin_rule = result.twin_rule;
twinned.twin_laws = result.twin_laws;
twinned.twin_r_work = result.twin_r_work;
twinned.twin_r_free = result.twin_r_free;
twinned.untwinned_r_work = result.untwinned_r_work;
twinned.untwinned_r_free = result.untwinned_r_free;
result = std::move(twinned);
}
} // namespace
bool AdoptTwinLaw(double twin_fraction, double null_z, bool data_indicated) {
return twin_fraction > 0 && (null_z >= TWIN_NULL_Z || (data_indicated && null_z >= TWIN_NULL_Z_DATA));
}
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) {
RefinementResult result = RefineOnce(placed, merged, cell, sg, wavelength_A, logger, settings);
// Each arm, untwinned and twinned, is a whole refinement - the solvent built and the displacement model chosen
// against its own target - and the decision is taken on where they ended. A model only scored is scored both
// ways and decided by the same rule, so the R-free it reports is not an untwinned one by default.
if (settings.twin && result.ok && !result.cancelled)
DecideTwin(placed, merged, cell, sg, wavelength_A, logger, settings, result);
result.map.reset();
return result;
}
Reference in new issue View Git Blame Copy Permalink