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