Restrained refinement on the GPU: ML amplitude target, FFT gradient, dictionary restraints, R1 + R2 driver, blobs
RefineModel (Refinement.h) refines a placed model against the merge, reading and writing nothing: - placement stage (R1): every chain of 50+ atoms as a rigid body at once (L-BFGS on six parameters per chain, their gradient from the per-atom one through Ceres jets), then one B shift per residue; - restrained stage (R2): positions and isotropic B of every atom of a residue the embedded dictionary describes, the rest held fixed but kept in Fcalc and the mask; riding hydrogens added by gemmi's topology, moved with their parents, re-placed every macro-cycle. Macro-cycles refit the mask, the scale and bulk solvent (FitModelScale / ModelScaleGPUPool) and the likelihood's D and beta (ML on the free set in resolution bins), rebuild the repulsion pairs (symmetry mates included), then L-BFGS on positions and, where chosen, on B. - weights: a ladder of full refinements at multiples of the gradient-balanced X-ray weight, chosen by R-free within bond/angle bounds (ties to the tighter geometry), then the B refinement tried at its own weights and kept only where it beats B held by more than the tie. The target and its gradient (RefinementXray, RefinementGPU): the Rice/centric amplitude likelihood (RefinementMath.h, shared host/device code), the gradient by Agarwal's FFT method - the adjoint of the reciprocal-space symmetry composition into a half-l (CPU) or half-u (cuFFT) spectrum, one inverse real transform, then each atom's density derivatives (xyz and B) summed over its box. On the GPU everything per evaluation stays on the device: ModelDensityGPU's gather, r2c, a fused compose + likelihood kernel, a CSR gather for the adjoint (no atomics), c2r, one warp per atom for the gradient; the restraints are one kernel per restraint plus a per-atom gather of gradient slots; cub reductions. Deterministic. Restraints (RefinementRestraints): bonds, angles, torsions (with esd, no phi/psi, no alternative sugar puckers), chiral volumes, least-squares planes from gemmi's Topo; repulsions beyond 1-4 (vdW radii from the dictionary's energy library, 2.6 A for H-bond pairs, 1.9 A metal-N/O/S); B similarity of bonded and angle-related atoms. DensityBlobs (D1): gemmi's flood-fill blobs of mFo-DFc above 3 sigma peaking above 4 sigma, placed beside the model and labelled WATER / ION / LIGAND / MODEL_ERROR by distance, contacts, volume and the anomalous map. tools/rugnux_refine_test refines a model against a finished run's p.mtz (developer utility, not installed). Tests: likelihood derivative, the embedded library, restraint values on ideal residues and every restraint kind's gradient against finite differences (CPU, and the GPU against the CPU), the X-ray gradient against finite differences in P1, P212121, C2 and P6122 (CPU and GPU), and a synthetic refinement that recovers a shaken model. Credits in docs/ACKNOWLEDGEMENT.md and at the algorithms. Measured on the open arm's battery outputs against REFMAC 5.8, 10 cycles, same p.mtz, free set and start model (R-free): 5reo 0.2228 vs 0.2236, 8dyz 0.1733 vs 0.1651, 9chw 0.1995 vs 0.1969 (CPU engine). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_015CeEGgpKdfkCTJb3YNMJcu (cherry picked from commit 8efe6f2322b3bfdae580bb34f256e67730dfc3c2)
This commit is contained in:
leonarski_f committed
1 parent
e58453a693
commit
b4b68d368f
19 files changed
+4220
No files matched your search
@@ -0,0 +1,283 @@
|
||||
// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
|
||||
// SPDX-License-Identifier: GPL-3.0-only
|
||||
|
||||
#include "RefinementGPU.h"
|
||||
|
||||
#include <algorithm>
|
||||
#include <cmath>
|
||||
|
||||
#include "gemmi/dencalc.hpp"
|
||||
#include "gemmi/it92.hpp"
|
||||
#include "gemmi/solmask.hpp" // SolventMasker, refmac_radius_for_bulk_solvent
|
||||
|
||||
#include "RefinementGPUEngine.h"
|
||||
|
||||
namespace {
|
||||
|
||||
using Table = gemmi::IT92<float>;
|
||||
|
||||
// One atom's density as the gather takes it, isotropic, with `blur`.
|
||||
ModelDensityAtom DensityAtom(const gemmi::Atom &atom, double b_iso, const gemmi::DensityCalculator<Table, float> &dc) {
|
||||
using CReal = Table::Coef::coef_type;
|
||||
const auto &coef = Table::get(atom.element, atom.charge, atom.serial);
|
||||
const CReal b = static_cast<CReal>(b_iso + dc.blur);
|
||||
const auto precal = coef.precalculate_density_iso(b, 0.0f);
|
||||
const double radius = dc.estimate_radius(precal, b);
|
||||
ModelDensityAtom a{};
|
||||
a.occ = atom.occ;
|
||||
a.aniso = 0;
|
||||
a.radius = static_cast<float>(radius);
|
||||
a.max_r2 = static_cast<float>(radius * radius);
|
||||
for (int k = 0; k < 5; k++) {
|
||||
a.a[k] = precal.a[k];
|
||||
a.b[k][0] = precal.b[k];
|
||||
}
|
||||
return a;
|
||||
}
|
||||
|
||||
std::vector<const gemmi::Atom *> Atoms(const gemmi::Model &model) {
|
||||
std::vector<const gemmi::Atom *> atoms;
|
||||
for (const gemmi::Chain &ch : model.chains)
|
||||
for (const gemmi::Residue &r : ch.residues)
|
||||
for (const gemmi::Atom &a : r.atoms)
|
||||
atoms.push_back(&a);
|
||||
return atoms;
|
||||
}
|
||||
|
||||
class RefinementEngineGPU : public RefinementEngine {
|
||||
public:
|
||||
RefinementEngineGPU(int device, const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg, double d_min,
|
||||
const RefinementReflections &refl, const gemmi::Model &model, double b_max)
|
||||
: device_(device), cell_(cell), sg_(sg), d_min_(d_min) {
|
||||
const gemmi::Grid<float> grid = RigidBodyZoneGrid(cell, sg, d_min);
|
||||
setup_.grid.nu = grid.nu;
|
||||
setup_.grid.nv = grid.nv;
|
||||
setup_.grid.nw = grid.nw;
|
||||
for (int i = 0; i < 3; i++)
|
||||
for (int j = 0; j < 3; j++) {
|
||||
setup_.grid.orth[3 * i + j] = cell.orth.mat[i][j];
|
||||
setup_.grid.frac[3 * i + j] = cell.frac.mat[i][j];
|
||||
}
|
||||
setup_.volume = cell.volume;
|
||||
const gemmi::GroupOps gops = sg.operations();
|
||||
setup_.den = gemmi::Op::DEN;
|
||||
setup_.centring = static_cast<double>(gops.cen_ops.size());
|
||||
for (const gemmi::Op &op : gops.sym_ops) {
|
||||
std::array<int, 12> o{};
|
||||
for (int i = 0; i < 3; i++) {
|
||||
for (int j = 0; j < 3; j++)
|
||||
o[3 * i + j] = op.rot[i][j];
|
||||
o[9 + i] = op.tran[i];
|
||||
}
|
||||
setup_.sym_ops.push_back(o);
|
||||
}
|
||||
for (const gemmi::Op::Tran &cen : gops.cen_ops)
|
||||
for (const gemmi::Op &op : gops.sym_ops) {
|
||||
ModelMaskOp m{};
|
||||
for (int i = 0; i < 3; i++) {
|
||||
for (int j = 0; j < 3; j++)
|
||||
m.rot[3 * i + j] = static_cast<double>(op.rot[i][j]) / gemmi::Op::DEN;
|
||||
m.tran[i] = static_cast<double>(op.tran[i] + cen[i]) / gemmi::Op::DEN;
|
||||
}
|
||||
setup_.mask_ops.push_back(m);
|
||||
}
|
||||
for (size_t m = 0; m < refl.Size(); m++) {
|
||||
setup_.rows.push_back({refl.hkl[m][0], refl.hkl[m][1], refl.hkl[m][2]});
|
||||
setup_.inv_d2.push_back(refl.inv_d2[m]);
|
||||
setup_.fo.push_back(refl.fo[m]);
|
||||
setup_.work.push_back(!refl.free[m]);
|
||||
setup_.centric.push_back(refl.centric[m]);
|
||||
}
|
||||
|
||||
// The adjoint of the composition on cuFFT's half-u spectrum: the term of row m and operator o, at
|
||||
// k = hR, goes to k where k_u > 0, conjugated to -k where k_u < 0, and to both on the planes k_u = 0
|
||||
// and k_u = nu/2, which hold their own Friedel mates.
|
||||
const int nu = grid.nu, nv = grid.nv, nw = grid.nw;
|
||||
auto point = [&](int ku, int kv, int kw) {
|
||||
return ku + (nu / 2 + 1) * (gemmi::modulo(kv, nv) + nv * gemmi::modulo(kw, nw));
|
||||
};
|
||||
std::vector<std::pair<int, int>> entries;
|
||||
const int n_ops = static_cast<int>(gops.sym_ops.size());
|
||||
for (size_t m = 0; m < refl.Size(); m++)
|
||||
for (int o = 0; o < n_ops; o++) {
|
||||
const gemmi::Miller k = gops.sym_ops[o].apply_to_hkl(refl.hkl[m]);
|
||||
const int term = (static_cast<int>(m) * n_ops + o) * 2;
|
||||
if (k[0] > 0 && 2 * k[0] != nu)
|
||||
entries.emplace_back(point(k[0], k[1], k[2]), term);
|
||||
else if (k[0] < 0 && -2 * k[0] != nu)
|
||||
entries.emplace_back(point(-k[0], -k[1], -k[2]), term + 1);
|
||||
else {
|
||||
const int ku = std::abs(k[0]);
|
||||
entries.emplace_back(point(ku, k[1], k[2]), term);
|
||||
entries.emplace_back(point(ku, -k[1], -k[2]), term + 1);
|
||||
}
|
||||
}
|
||||
std::stable_sort(entries.begin(), entries.end(),
|
||||
[](const auto &a, const auto &b) { return a.first < b.first; });
|
||||
for (size_t e = 0; e < entries.size(); e++) {
|
||||
if (e == 0 || entries[e].first != entries[e - 1].first) {
|
||||
setup_.scatter_point.push_back(entries[e].first);
|
||||
setup_.scatter_start.push_back(static_cast<int>(e));
|
||||
}
|
||||
setup_.scatter_term.push_back(entries[e].second);
|
||||
}
|
||||
setup_.scatter_start.push_back(static_cast<int>(entries.size()));
|
||||
|
||||
// The gather's reserve for every atom at b_max with the largest blur a refinement sets (B down to 0).
|
||||
dc_ = Calculator(RefinementBlur(cell, sg, d_min, 0.0));
|
||||
std::vector<ModelDensityAtom> widest;
|
||||
for (const gemmi::Atom *a : Atoms(model))
|
||||
widest.push_back(DensityAtom(*a, b_max, dc_));
|
||||
setup_.max_atoms = widest.size();
|
||||
setup_.max_pairs = ModelDensityGPU::PairBound(setup_.grid, widest);
|
||||
supported_ = !setup_.rows.empty() && ModelDensityGPU::Supports(setup_.grid, widest) &&
|
||||
setup_.mask_ops.size() <= ModelMaskGPU::MAX_OPS;
|
||||
if (supported_) {
|
||||
fft_work_ = RefinementGPUEngine::FFTWorkBytes(device, setup_.grid);
|
||||
// the restraints: about thirty records and slots per atom
|
||||
bytes_ = RefinementGPUEngine::DeviceBytes(setup_, fft_work_, 30 * widest.size(), 60 * widest.size());
|
||||
}
|
||||
}
|
||||
|
||||
bool Supported() const { return supported_; }
|
||||
size_t Bytes() const { return bytes_; }
|
||||
void Reserve() { engine_ = std::make_unique<RefinementGPUEngine>(device_, setup_, fft_work_); }
|
||||
|
||||
void SetBlur(double blur) override {
|
||||
dc_ = Calculator(blur);
|
||||
last_b_.clear();
|
||||
engine_->SetBlur(blur);
|
||||
}
|
||||
|
||||
void SetRowModel(const XrayRowModel &rows) override {
|
||||
std::vector<std::array<float, 2>> fmask(rows.fmask.size());
|
||||
for (size_t m = 0; m < fmask.size(); m++)
|
||||
fmask[m] = {rows.fmask[m].real(), rows.fmask[m].imag()};
|
||||
engine_->SetRowModel(rows.k_total, fmask, rows.d, rows.var);
|
||||
}
|
||||
|
||||
bool OnGPU() const override { return true; }
|
||||
|
||||
void Fcalc(const gemmi::Model &model, std::vector<std::complex<float>> &f) override {
|
||||
SetAtoms(model);
|
||||
std::vector<std::array<float, 2>> out;
|
||||
engine_->Fcalc(out);
|
||||
f.resize(out.size());
|
||||
for (size_t m = 0; m < out.size(); m++)
|
||||
f[m] = {out[m][0], out[m][1]};
|
||||
}
|
||||
|
||||
void Fmask(const gemmi::Model &model, std::vector<std::complex<float>> &f) override {
|
||||
const gemmi::SolventMasker masker(gemmi::AtomicRadiiSet::Refmac);
|
||||
std::vector<ModelMaskAtom> mask_atoms;
|
||||
for (const gemmi::Atom *a : Atoms(model)) {
|
||||
if ((masker.ignore_hydrogen && a->is_hydrogen()) || (masker.ignore_zero_occupancy_atoms && a->occ <= 0))
|
||||
continue;
|
||||
const gemmi::Fractional f0 = cell_.fractionalize(a->pos);
|
||||
double frac[3] = {f0.x, f0.y, f0.z};
|
||||
for (double &x : frac) {
|
||||
x -= std::floor(x);
|
||||
if (x >= 1.0)
|
||||
x = 0.0;
|
||||
}
|
||||
mask_atoms.push_back({frac[0], frac[1], frac[2],
|
||||
masker.constant_r + masker.rprobe +
|
||||
gemmi::refmac_radius_for_bulk_solvent(a->element.elem)});
|
||||
}
|
||||
std::vector<std::array<float, 2>> out;
|
||||
engine_->Fmask(mask_atoms, out);
|
||||
f.resize(out.size());
|
||||
for (size_t m = 0; m < out.size(); m++)
|
||||
f[m] = {out[m][0], out[m][1]};
|
||||
}
|
||||
|
||||
double Evaluate(const gemmi::Model &model, std::vector<std::array<double, 4>> *grad) override {
|
||||
SetAtoms(model);
|
||||
std::vector<std::array<float, 4>> g;
|
||||
const double value = engine_->Evaluate(grad != nullptr ? &g : nullptr);
|
||||
if (grad != nullptr) {
|
||||
grad->resize(g.size());
|
||||
for (size_t i = 0; i < g.size(); i++)
|
||||
(*grad)[i] = {g[i][0], g[i][1], g[i][2], g[i][3]};
|
||||
}
|
||||
return value;
|
||||
}
|
||||
|
||||
void SetRestraints(int which, const RestraintSet &set) override {
|
||||
engine_->SetRestraints(which, set.records, set.plane_atoms, set.slot_start, set.slot_index, set.slots);
|
||||
}
|
||||
|
||||
double EvaluateRestraints(int which, const std::vector<std::array<double, 4>> &xyzb,
|
||||
std::vector<std::array<double, 4>> *grad) override {
|
||||
std::vector<std::array<float, 4>> x(xyzb.size()), g;
|
||||
for (size_t i = 0; i < xyzb.size(); i++)
|
||||
x[i] = {static_cast<float>(xyzb[i][0]), static_cast<float>(xyzb[i][1]), static_cast<float>(xyzb[i][2]),
|
||||
static_cast<float>(xyzb[i][3])};
|
||||
const double value = engine_->EvaluateRestraints(which, x, grad != nullptr ? &g : nullptr);
|
||||
if (grad != nullptr) {
|
||||
grad->resize(g.size());
|
||||
for (size_t i = 0; i < g.size(); i++)
|
||||
(*grad)[i] = {g[i][0], g[i][1], g[i][2], g[i][3]};
|
||||
}
|
||||
return value;
|
||||
}
|
||||
|
||||
private:
|
||||
gemmi::DensityCalculator<Table, float> Calculator(double blur) const {
|
||||
gemmi::DensityCalculator<Table, float> dc;
|
||||
dc.d_min = d_min_;
|
||||
dc.rate = 1.5;
|
||||
dc.grid.unit_cell = cell_;
|
||||
dc.grid.spacegroup = &sg_;
|
||||
dc.blur = blur;
|
||||
return dc;
|
||||
}
|
||||
|
||||
// The atoms' densities (made again only for the atoms whose B moved) and grid positions, to the device.
|
||||
void SetAtoms(const gemmi::Model &model) {
|
||||
const std::vector<const gemmi::Atom *> atoms = Atoms(model);
|
||||
if (last_b_.size() != atoms.size()) {
|
||||
last_b_.assign(atoms.size(), NAN);
|
||||
atoms_.resize(atoms.size());
|
||||
}
|
||||
std::vector<std::array<float, 4>> pos(atoms.size());
|
||||
const int n3[3] = {setup_.grid.nu, setup_.grid.nv, setup_.grid.nw};
|
||||
for (size_t i = 0; i < atoms.size(); i++) {
|
||||
if (!(atoms[i]->b_iso == last_b_[i])) {
|
||||
atoms_[i] = DensityAtom(*atoms[i], atoms[i]->b_iso, dc_);
|
||||
last_b_[i] = atoms[i]->b_iso;
|
||||
}
|
||||
const gemmi::Fractional f = cell_.fractionalize(atoms[i]->pos);
|
||||
pos[i] = {static_cast<float>(f.x * n3[0]), static_cast<float>(f.y * n3[1]), static_cast<float>(f.z * n3[2]), 0};
|
||||
}
|
||||
engine_->SetAtoms(atoms_, pos);
|
||||
}
|
||||
|
||||
int device_;
|
||||
gemmi::UnitCell cell_;
|
||||
const gemmi::SpaceGroup &sg_;
|
||||
double d_min_;
|
||||
RefinementGPUSetup setup_;
|
||||
gemmi::DensityCalculator<Table, float> dc_;
|
||||
bool supported_ = false;
|
||||
size_t fft_work_ = 0, bytes_ = 0;
|
||||
std::unique_ptr<RefinementGPUEngine> engine_;
|
||||
std::vector<float> last_b_;
|
||||
std::vector<ModelDensityAtom> atoms_;
|
||||
};
|
||||
|
||||
} // namespace
|
||||
|
||||
std::unique_ptr<RefinementEngine> MakeRefinementEngineGPU(int device, const gemmi::UnitCell &cell,
|
||||
const gemmi::SpaceGroup &sg, double d_min,
|
||||
const RefinementReflections &refl,
|
||||
const gemmi::Model &model, double b_max, size_t budget,
|
||||
size_t *device_bytes) {
|
||||
auto engine = std::make_unique<RefinementEngineGPU>(device, cell, sg, d_min, refl, model, b_max);
|
||||
if (device_bytes != nullptr)
|
||||
*device_bytes = engine->Bytes();
|
||||
if (!engine->Supported() || engine->Bytes() > budget)
|
||||
return nullptr;
|
||||
engine->Reserve();
|
||||
return engine;
|
||||
}
|
||||
Reference in new issue
Block a user