leonarski_fandClaude Opus 5.5
08594fe6c9
FWT/PHWT in <prefix>_maps.mtz and <prefix>_refined_maps.mtz now carries, beside the observed reflections, every reflection of the data's resolution range that the data do not have, at D Fc - its expected value given the model - with D that of the observed reflection nearest in resolution. This is the convention of phenix.refine (fill_missing_f_obs) and BUSTER: a missing reflection left at zero biases the map towards it, which at atomic resolution shows as broader, lower atoms where completeness is short. The filled rows have FP, FOM, FREE and the difference coefficients missing. FWT_no_fill/PHWT_no_fill is the same map on the observed reflections only (phenix's 2FOFCWT_no_fill), and DELFWT/PHDELWT is never filled. For the refined model, F_model at the filled reflections is computed on the CPU with the bulk solvent and scale of the last fit. The map-coefficient tests check the filled FWT and the missing FWT_no_fill/DELFWT at unmeasured reflections. The R-factors, the blobs and the maps the refinement reads are unchanged. On an 86%-complete 0.99 A set the median FWHM of well-ordered carbons drops from 1.061 to 1.034 A and their peak rises by 4%; on a 99%-complete set nothing changes. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01UqXYdL4qBGYRiSXSGXfCLu
3524 lines
183 KiB
C++
3524 lines
183 KiB
C++
// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
|
|
// SPDX-License-Identifier: GPL-3.0-only
|
|
|
|
#include "Refinement.h"
|
|
|
|
#include <algorithm>
|
|
#include <atomic>
|
|
#include <chrono>
|
|
#include <set>
|
|
#include <future>
|
|
#include <limits>
|
|
#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_fractions = FitTwinFractions(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_fractions = FitTwinFractions(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 = 0;
|
|
for (float a : rows.twin_fractions)
|
|
sc.twin_fraction += a;
|
|
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> ¶ms, const std::vector<double> &shape) {
|
|
constexpr double EPS = 0.5; // A^2
|
|
const std::vector<double> b0 = B();
|
|
std::mt19937 rng(7);
|
|
std::vector<double> v(params.size());
|
|
for (double &x : v)
|
|
x = rng() % 2 == 0 ? 1.0 : -1.0;
|
|
std::array<std::vector<std::array<double, 4>>, 2> g;
|
|
for (int side = 0; side < 2; side++) {
|
|
for (size_t k = 0; k < params.size(); k++)
|
|
atoms[params[k]]->b_iso = static_cast<float>(b0[params[k]] + (side == 0 ? EPS : -EPS) * v[k]);
|
|
MoveRiders();
|
|
engine->Evaluate(Model(), &g[side]);
|
|
FoldRiders(g[side]);
|
|
}
|
|
SetB(b0);
|
|
MoveRiders();
|
|
double num = 0, den = 0;
|
|
for (size_t k = 0; k < params.size(); k++) {
|
|
num += v[k] * (g[0][params[k]][3] - g[1][params[k]][3]) / (2 * EPS);
|
|
den += shape[params[k]];
|
|
}
|
|
return den > 0 ? std::max(num / den, 0.0) : 0.0;
|
|
}
|
|
|
|
// F_model with the cycle's scale and solvent, for the atoms as they are now.
|
|
std::vector<std::complex<float>> FModel() {
|
|
engine->Fcalc(Model(), fcalc);
|
|
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> ¶m_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> ¶m_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> ¶m_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> ¶m_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> ¶m_atoms, double weight) {
|
|
std::vector<double> h;
|
|
for (int i : param_atoms)
|
|
h.push_back(weight * xray[i] + restraints[i]);
|
|
std::vector<double> sorted = h;
|
|
std::nth_element(sorted.begin(), sorted.begin() + sorted.size() / 2, sorted.end());
|
|
const double floor = std::max(1e-3 * sorted[sorted.size() / 2], 1e-12);
|
|
for (double &v : h)
|
|
v = std::max(v, floor);
|
|
return h;
|
|
}
|
|
|
|
// `minimiser` carries its memory from the call before; `diagonal` is the curvature it scales the steps by (none: 1).
|
|
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;
|
|
}
|
|
|
|
// A reflection of the data's resolution range that the data do not have: F_model on the scale of the map
|
|
// coefficients (the overall anisotropy divided out), and the D of the observed reflection nearest in resolution.
|
|
struct MissingReflection {
|
|
gemmi::Miller hkl;
|
|
std::complex<double> f_model;
|
|
double d = 0;
|
|
};
|
|
|
|
// The reflections 2mFo - DFc is filled at, with D F_model - the reflection's expected value given the model - as
|
|
// phenix.refine (fill_missing_f_obs) and BUSTER fill their maps. F_model is computed on the CPU, its bulk solvent
|
|
// and overall scale those of the last Refit(), `b_aniso` the anisotropic part of the scale.
|
|
std::vector<MissingReflection> MissingReflections(Context &ctx, const gemmi::SMat33<double> &b_aniso) {
|
|
if (ctx.refl.Size() == 0)
|
|
return {};
|
|
const auto [s_lo, s_hi] = std::minmax_element(ctx.refl.inv_d2.begin(), ctx.refl.inv_d2.end());
|
|
const std::set<gemmi::Miller> observed(ctx.refl.hkl.begin(), ctx.refl.hkl.end());
|
|
const gemmi::ReciprocalAsu asu(ctx.sg);
|
|
const gemmi::GroupOps gops = ctx.sg->operations();
|
|
RefinementReflections missing;
|
|
const double q_max = std::sqrt(*s_hi);
|
|
const int h_max = static_cast<int>(ctx.cell.a * q_max) + 1, k_max = static_cast<int>(ctx.cell.b * q_max) + 1,
|
|
l_max = static_cast<int>(ctx.cell.c * q_max) + 1;
|
|
for (int h = -h_max; h <= h_max; h++)
|
|
for (int k = -k_max; k <= k_max; k++)
|
|
for (int l = -l_max; l <= l_max; l++) {
|
|
const gemmi::Miller hkl{{h, k, l}};
|
|
const double s = ctx.cell.calculate_1_d2(hkl);
|
|
if (s < *s_lo || s > *s_hi || !asu.is_in(hkl) || gops.is_systematically_absent(hkl)
|
|
|| observed.count(hkl) > 0)
|
|
continue;
|
|
missing.hkl.push_back(hkl);
|
|
missing.inv_d2.push_back(s);
|
|
}
|
|
if (missing.hkl.empty())
|
|
return {};
|
|
|
|
std::unique_ptr<RefinementEngine> engine = MakeRefinementEngineCPU(ctx.cell, *ctx.sg, ctx.d_min, missing, ctx.nthreads);
|
|
engine->SetBlur(RefinementBlur(ctx.cell, *ctx.sg, ctx.d_min, ctx.b_floor));
|
|
engine->SetAniso(&ctx.aniso);
|
|
std::vector<std::complex<float>> fcalc, fmask;
|
|
engine->Fcalc(ctx.Model(), fcalc);
|
|
ctx.mask->Fmask(ctx.Model(), missing.hkl, fmask);
|
|
|
|
std::vector<std::pair<double, double>> d_by_s; // (1/d^2, D) of the observed reflections, by resolution
|
|
for (size_t m = 0; m < ctx.refl.Size(); m++)
|
|
d_by_s.emplace_back(ctx.refl.inv_d2[m], ctx.rows.d[m]);
|
|
std::sort(d_by_s.begin(), d_by_s.end());
|
|
std::vector<MissingReflection> out;
|
|
for (size_t i = 0; i < missing.hkl.size(); i++) {
|
|
const double s = missing.inv_d2[i];
|
|
auto it = std::lower_bound(d_by_s.begin(), d_by_s.end(), std::make_pair(s, -1.0));
|
|
if (it == d_by_s.end() || (it != d_by_s.begin() && s - std::prev(it)->first < it->first - s))
|
|
--it;
|
|
const double k = ctx.k_overall * std::exp(-0.25 * (ctx.b_star.r_u_r(missing.hkl[i]) - b_aniso.r_u_r(missing.hkl[i])));
|
|
const std::complex<double> f_model = k * (std::complex<double>(fcalc[i])
|
|
+ ctx.k_sol * std::exp(-ctx.b_sol * s / 4.0) * std::complex<double>(fmask[i]));
|
|
out.push_back({missing.hkl[i], f_model, it->second});
|
|
}
|
|
return out;
|
|
}
|
|
|
|
// 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
|
|
// The coefficients written to the file have the overall anisotropy of the fit (the anisotropic part of k_total)
|
|
// divided out, as phenix.refine's are: left in, it stretches every atom along the data's weak direction. The maps
|
|
// the refinement reads itself (result.map, result.difference_map) keep the observed scale. FWT/PHWT is filled at
|
|
// the reflections the data do not have (MissingReflections), FWT_no_fill/PHWT_no_fill is the same map without
|
|
// them, and the difference map is never filled.
|
|
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'}, {"FWT_no_fill", 'F'},
|
|
{"PHWT_no_fill", '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;
|
|
const gemmi::SMat33<double> b_aniso = AnisotropicPart(ctx.b_star, ctx.cell);
|
|
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;
|
|
const double k_aniso = std::exp(-0.25 * b_aniso.r_u_r(ctx.refl.hkl[m]));
|
|
const double fwt_file = fwt / k_aniso, delfwt_file = delfwt / k_aniso;
|
|
// 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_file)),
|
|
MtzPhase(fwt_file < 0 ? phi_deg + 180.0 : phi_deg), static_cast<float>(std::abs(fwt_file)),
|
|
MtzPhase(fwt_file < 0 ? phi_deg + 180.0 : phi_deg), static_cast<float>(std::abs(delfwt_file)),
|
|
MtzPhase(delfwt_file < 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;
|
|
// Its coefficients in the maps file too (ANOM/PHANOM, as model validation writes them), missing (NaN) where
|
|
// the reflection has no Bijvoet pair.
|
|
const float nan = std::numeric_limits<float>::quiet_NaN();
|
|
MtzTable::Column anom_col{"ANOM", 'F', 1, {}}, phanom_col{"PHANOM", 'P', 1, {}};
|
|
for (size_t m = 0; m < ctx.refl.Size(); m++) {
|
|
const auto it = danom.find(ctx.refl.hkl[m]);
|
|
if (it == danom.end()) {
|
|
anom_col.values.push_back(nan);
|
|
phanom_col.values.push_back(nan);
|
|
continue;
|
|
}
|
|
const std::complex<float> c = it->second * std::polar(1.0f, static_cast<float>(std::arg(fmodel[m]) - PI / 2));
|
|
anom.v.push_back({ctx.refl.hkl[m], c});
|
|
anom_col.values.push_back(std::abs(c));
|
|
phanom_col.values.push_back(MtzPhase(std::arg(c) * 180.0 / PI));
|
|
}
|
|
if (!anom.v.empty()) {
|
|
mtz.columns.push_back(std::move(anom_col));
|
|
mtz.columns.push_back(std::move(phanom_col));
|
|
result.anomalous_map = std::make_shared<gemmi::Grid<float>>(
|
|
MapFromFPhi(gemmi::get_f_phi_on_grid<float>(anom, gemmi::get_size_for_hkl(anom, {{0, 0, 0}}, 3.0), true)));
|
|
result.anomalous_map_rms = gemmi::calculate_data_statistics(result.anomalous_map->data).rms;
|
|
}
|
|
}
|
|
// The filled reflections after the observed ones: FWT/PHWT and the model's FC/PHIC, every other column missing.
|
|
const float nan = std::numeric_limits<float>::quiet_NaN();
|
|
for (const MissingReflection &r : MissingReflections(ctx, b_aniso)) {
|
|
const double phi_deg = std::arg(r.f_model) * 180.0 / PI;
|
|
const double k_aniso = std::exp(-0.25 * b_aniso.r_u_r(r.hkl));
|
|
const float row[] = {static_cast<float>(r.hkl[0]), static_cast<float>(r.hkl[1]), static_cast<float>(r.hkl[2]), nan,
|
|
static_cast<float>(std::abs(r.f_model) * k_aniso), MtzPhase(phi_deg),
|
|
static_cast<float>(r.d * std::abs(r.f_model)), MtzPhase(phi_deg)};
|
|
for (size_t c = 0; c < mtz.columns.size(); c++)
|
|
mtz.columns[c].values.push_back(c < std::size(row) ? row[c] : nan);
|
|
}
|
|
difference_map.get();
|
|
map.get();
|
|
}
|
|
|
|
// The fractions of the twin laws, together.
|
|
double FractionSum(const XrayRowModel &rows) {
|
|
double sum = 0;
|
|
for (float a : rows.twin_fractions)
|
|
sum += a;
|
|
return sum;
|
|
}
|
|
|
|
// law=fraction per twin law, ';'-separated.
|
|
std::string DomainText(const RefinementReflections &refl, const XrayRowModel &rows) {
|
|
std::string out;
|
|
for (size_t d = 0; d < refl.twin_laws.size(); d++)
|
|
out += fmt::format("{}{}={:.3f}", out.empty() ? "" : ";", refl.twin_laws[d].as_hkl().triplet('h'),
|
|
d < rows.twin_fractions.size() ? rows.twin_fractions[d] : 0.0f);
|
|
return out;
|
|
}
|
|
|
|
// A domain added to the adopted law is kept only at this fraction or more, on top of its null.
|
|
// As REFMAC5 drops twin domains of small fraction, Murshudov et al. (2011) Acta Cryst. D67, 355-367
|
|
constexpr double TWIN_DOMAIN_MIN = 0.05;
|
|
|
|
// 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;
|
|
// `first`: the first law whose mates are replaced - the laws before it keep theirs (the null of the domains added to
|
|
// them). With several laws a twin group holds up to Laws() + 1 reflections, kept next to each other, and law first + k
|
|
// is paired (Laws() + 1) (k + 1) apart: past the reflection's own group, and no two laws drawing the same partner.
|
|
void TwinNull(const Context &ctx, RefinementResult &result, size_t first = 0) {
|
|
RefinementReflections null = ctx.refl;
|
|
const size_t n = null.Size();
|
|
// each observed reflection's twin group, named by its lowest observed row
|
|
std::vector<int> group(n);
|
|
for (size_t m = 0; m < n; m++) {
|
|
group[m] = static_cast<int>(m);
|
|
for (size_t d = 0; d < ctx.refl.Laws(); d++)
|
|
if (ctx.refl.mates[d][m] < static_cast<int>(n))
|
|
group[m] = std::min(group[m], ctx.refl.mates[d][m]);
|
|
}
|
|
for (size_t d = first; d < null.Laws(); d++)
|
|
std::iota(null.mates[d].begin(), null.mates[d].end(), 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 group is ordered at the 1/d^2 of its first member
|
|
std::stable_sort(order.begin(), order.end(), [&](int a, int b) {
|
|
const int pa = group[a], pb = group[b];
|
|
return std::make_pair(null.inv_d2[pa], pa) < std::make_pair(null.inv_d2[pb], pb);
|
|
});
|
|
for (size_t d = first; d < null.Laws(); d++) {
|
|
const size_t shift = (null.Laws() + 1) * (d - first + 1);
|
|
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 * shift) < shift ? r + shift : r - shift;
|
|
null.mates[d][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;
|
|
{
|
|
const std::vector<float> a = FitTwinFractions(null, rows, ctx.fcalc);
|
|
result.twin_fraction_null = 0;
|
|
for (size_t d = first; d < a.size(); d++)
|
|
result.twin_fraction_null += a[d];
|
|
}
|
|
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 laws beat their null on the free reflections, paired: the mean of |Fo - |F_null|| - |Fo - |F_twin||
|
|
// over its standard error. The members of a free twin group share their F_calc, so their terms are summed into
|
|
// one: the standard error counts twin groups, not reflections.
|
|
const std::vector<std::complex<float>> f_law = TwinnedModel(ctx.refl, ctx.rows, ctx.fcalc);
|
|
std::vector<double> sum_group(n, 0.0);
|
|
for (size_t m = 0; m < n; m++)
|
|
if (ctx.refl.free[m])
|
|
sum_group[group[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_groups = 0;
|
|
for (size_t m = 0; m < n; m++)
|
|
if (ctx.refl.free[m] && group[m] == static_cast<int>(m)) {
|
|
sum += sum_group[m];
|
|
sum2 += sum_group[m] * sum_group[m];
|
|
n_groups++;
|
|
}
|
|
if (n_groups > 1) {
|
|
const double mean = sum / n_groups, var = (sum2 - n_groups * mean * mean) / (n_groups - 1);
|
|
result.twin_null_z = var > 0 ? mean / std::sqrt(var / n_groups) : 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 water chain (PickSolvent's) 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();
|
|
}
|
|
// a water the subchain and entity of the waters already there, where there are any; else none, for the
|
|
// writer to give the chain's residues subchains of their own (setup_entities), not the placeholders'
|
|
moved[k].subchain.clear();
|
|
moved[k].entity_id.clear();
|
|
for (const gemmi::Residue &w : target->residues)
|
|
if (w.is_water() && moved[k].is_water()) {
|
|
moved[k].subchain = w.subchain;
|
|
moved[k].entity_id = w.entity_id;
|
|
}
|
|
target->residues.push_back(moved[k]);
|
|
}
|
|
// In a chain that took any, the waters after its other residues (a site relabelled from a water put in), so that
|
|
// the writer's subchains are each one run of residues.
|
|
for (gemmi::Chain &ch : model.chains)
|
|
if (std::find(into.begin(), into.end(), ch.name) != into.end())
|
|
std::stable_partition(ch.residues.begin(), ch.residues.end(),
|
|
[](const gemmi::Residue &r) { return !r.is_water(); });
|
|
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_laws`: against the data as twinned by them (one domain each beside the model's own), empty for untwinned;
|
|
// `null_first`: the first of them the null (TwinNull) replaces the mates of.
|
|
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 std::vector<gemmi::Op> &twin_laws = {}, size_t null_first = 0) {
|
|
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_laws.empty()) {
|
|
result.twin_split_pairs = SetTwinLaws(ctx.refl, cell, sg, ctx.d_min, twin_laws);
|
|
result.twinned = true;
|
|
for (const gemmi::Op &law : twin_laws)
|
|
result.twin_law += (result.twin_law.empty() ? "" : ";") + 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, FractionSum(ctx.rows),
|
|
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);
|
|
result.score = ctx.Score(); // at the scale and solvent MakeMaps refitted
|
|
if (ctx.refl.Twinned()) {
|
|
result.twin_fraction = FractionSum(ctx.rows);
|
|
result.twin_multi_domains = DomainText(ctx.refl, ctx.rows);
|
|
result.twin_domain_fractions.assign(ctx.rows.twin_fractions.begin(), ctx.rows.twin_fractions.end());
|
|
TwinNull(ctx, result, null_first);
|
|
}
|
|
// 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 hypotheses beside the untwinned refinement `result`, as a ladder. First one law: 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. Where that law is
|
|
// taken and the lattice has more domains (RefinementTwinDomains), the next rung: every domain at once, each with its
|
|
// fraction, taken where the domains added beat their own null - their mates unrelated, the law's kept - by TWIN_NULL_Z
|
|
// standard errors, each added domain at a fraction of TWIN_DOMAIN_MIN at least.
|
|
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);
|
|
const std::vector<gemmi::Op> domains = RefinementTwinDomains(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_free_closed = result.twin_free_closed;
|
|
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);
|
|
|
|
// the next rung: every domain of the lattice, the adopted law's first
|
|
std::vector<gemmi::Op> all{law};
|
|
for (const gemmi::Op &op : domains)
|
|
if (op.rot != law.rot)
|
|
all.push_back(op);
|
|
if (all.size() < 2)
|
|
return;
|
|
// Refined, the added domains under TWIN_DOMAIN_MIN dropped and the rest refined again, until every one left
|
|
// has it; none left: the single law stands.
|
|
RefinementResult multi;
|
|
while (true) {
|
|
multi = RefineOnce(placed, merged, cell, sg, wavelength_A, logger, settings, all, 1);
|
|
if (!multi.ok || multi.cancelled)
|
|
return;
|
|
std::vector<gemmi::Op> kept{law};
|
|
for (size_t d = 1; d < all.size(); d++)
|
|
if (d < multi.twin_domain_fractions.size() && multi.twin_domain_fractions[d] >= TWIN_DOMAIN_MIN)
|
|
kept.push_back(all[d]);
|
|
if (kept.size() == all.size())
|
|
break;
|
|
logger.Info("Refinement, {} twin domains ({}): {} below a fraction of {:.2f}{}", all.size() + 1,
|
|
multi.twin_multi_domains, all.size() - kept.size(), TWIN_DOMAIN_MIN,
|
|
kept.size() > 1 ? ", refined again without" : " - the single law stands");
|
|
if (kept.size() == 1) {
|
|
result.twin_multi_decision = "MULTI_UNTWINNED";
|
|
result.twin_multi_domains = multi.twin_multi_domains;
|
|
return;
|
|
}
|
|
all = kept;
|
|
}
|
|
result.twin_multi_domains = multi.twin_multi_domains;
|
|
result.twin_multi_r_work = FinalRWork(multi);
|
|
result.twin_multi_r_free = FinalRFree(multi);
|
|
result.twin_multi_r_free_null = multi.twin_r_free_null;
|
|
result.twin_multi_null_z = multi.twin_null_z;
|
|
const bool adopt_multi = multi.twin_null_z >= TWIN_NULL_Z;
|
|
result.twin_multi_decision = adopt_multi ? "MULTI_TWINNED" : "MULTI_UNTWINNED";
|
|
logger.Info("Refinement, {} twin domains ({}): R-work {:.4f}, R-free {:.4f}; one law R-free {:.4f}; null (the added "
|
|
"domains' mates unrelated) R-free {:.4f}, the domains ahead of it by {:.1f} standard errors -> {}",
|
|
all.size() + 1, multi.twin_multi_domains, result.twin_multi_r_work, result.twin_multi_r_free,
|
|
result.twin_r_free, result.twin_multi_r_free_null, result.twin_multi_null_z, result.twin_multi_decision);
|
|
if (!adopt_multi)
|
|
return;
|
|
multi.twin_decision = result.twin_decision;
|
|
multi.twin_rule = result.twin_rule;
|
|
multi.twin_laws = result.twin_laws;
|
|
multi.twin_free_closed = result.twin_free_closed;
|
|
multi.twin_fraction_null = result.twin_fraction_null;
|
|
multi.twin_r_free_null = result.twin_r_free_null;
|
|
multi.twin_null_z = result.twin_null_z;
|
|
multi.twin_r_work = result.twin_r_work;
|
|
multi.twin_r_free = result.twin_r_free;
|
|
multi.untwinned_r_work = result.untwinned_r_work;
|
|
multi.untwinned_r_free = result.untwinned_r_free;
|
|
multi.twin_multi_decision = result.twin_multi_decision;
|
|
multi.twin_multi_r_work = result.twin_multi_r_work;
|
|
multi.twin_multi_r_free = result.twin_multi_r_free;
|
|
multi.twin_multi_r_free_null = result.twin_multi_r_free_null;
|
|
multi.twin_multi_null_z = result.twin_multi_null_z;
|
|
result = std::move(multi);
|
|
}
|
|
|
|
} // namespace
|
|
|
|
int CloseFreeSet(std::vector<MergedReflection> &merged, const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg) {
|
|
const std::vector<gemmi::Op> laws = RefinementTwinDomains(cell, sg);
|
|
if (laws.empty())
|
|
return 0;
|
|
const gemmi::GroupOps gops = sg.operations();
|
|
const gemmi::ReciprocalAsu asu(&sg);
|
|
auto key = [&](const gemmi::Miller &h) { return asu.to_asu(h, gops).first; };
|
|
// each reflection's orbit, by the ASU index of every reflection the laws reach from it; named by its lowest index
|
|
std::map<gemmi::Miller, int> orbit_of;
|
|
std::vector<gemmi::Miller> name;
|
|
std::vector<int> orbit(merged.size());
|
|
for (size_t i = 0; i < merged.size(); i++) {
|
|
const gemmi::Miller h = key(gemmi::Miller{{merged[i].h, merged[i].k, merged[i].l}});
|
|
if (const auto it = orbit_of.find(h); it != orbit_of.end()) {
|
|
orbit[i] = it->second;
|
|
continue;
|
|
}
|
|
const int o = static_cast<int>(name.size());
|
|
name.push_back(h);
|
|
std::vector<gemmi::Miller> todo{h};
|
|
orbit_of[h] = o;
|
|
while (!todo.empty()) {
|
|
const gemmi::Miller x = todo.back();
|
|
todo.pop_back();
|
|
name[o] = std::min(name[o], x);
|
|
for (const gemmi::Op &law : laws)
|
|
if (const gemmi::Miller y = key(law.apply_to_hkl(x)); orbit_of.emplace(y, o).second)
|
|
todo.push_back(y);
|
|
}
|
|
orbit[i] = o;
|
|
}
|
|
std::vector<int> free(name.size(), 0), all(name.size(), 0);
|
|
for (size_t i = 0; i < merged.size(); i++) {
|
|
free[orbit[i]] += merged[i].rfree_flag;
|
|
all[orbit[i]]++;
|
|
}
|
|
// splitmix64 of the orbit's name -> uniform in [0, 1), salted: unsalted it is AssignRfreeFlags' own draw of a
|
|
// reflection the name often is, and an orbit would go free whenever that member was
|
|
auto uniform = [](const gemmi::Miller &h) {
|
|
uint64_t z = (static_cast<uint64_t>(h[0] + (1 << 20)) << 42 | static_cast<uint64_t>(h[1] + (1 << 20)) << 21
|
|
| static_cast<uint64_t>(h[2] + (1 << 20))) + 0x9e3779b97f4a7c15ULL * 2;
|
|
z = (z ^ (z >> 30)) * 0xbf58476d1ce4e5b9ULL;
|
|
z = (z ^ (z >> 27)) * 0x94d049bb133111ebULL;
|
|
return static_cast<double>((z ^ (z >> 31)) >> 11) * (1.0 / 9007199254740992.0);
|
|
};
|
|
int moved = 0;
|
|
for (size_t i = 0; i < merged.size(); i++) {
|
|
const int o = orbit[i];
|
|
const bool flag = free[o] == all[o] || (free[o] > 0 && uniform(name[o]) * all[o] < free[o]);
|
|
moved += merged[i].rfree_flag != flag;
|
|
merged[i].rfree_flag = flag;
|
|
}
|
|
return moved;
|
|
}
|
|
|
|
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) {
|
|
// The free set closed under the lattice's twin domains once, before either arm, so that both share it.
|
|
std::vector<MergedReflection> closed;
|
|
int moved = 0;
|
|
if (settings.twin) {
|
|
closed = merged;
|
|
moved = CloseFreeSet(closed, cell, sg);
|
|
if (moved > 0)
|
|
logger.Info("Refinement: the free set closed under the lattice's twin domains, each group of twin-related "
|
|
"reflections in one set - {} reflections changed set", moved);
|
|
}
|
|
const std::vector<MergedReflection> &data = moved > 0 ? closed : merged;
|
|
RefinementResult result = RefineOnce(placed, data, 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, data, cell, sg, wavelength_A, logger, settings, result);
|
|
result.twin_free_closed = moved;
|
|
result.map.reset();
|
|
return result;
|
|
}
|