Files
Jungfraujoch/image_analysis/structure_refinement/DensityBlobs.cpp
T
leonarski_f b4b68d368f 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)
2026-10-25 07:20:32 +01:00

111 lines
4.9 KiB
C++

// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only
#include "DensityBlobs.h"
#include <algorithm>
#include <cmath>
#include "gemmi/blob.hpp"
#include "gemmi/neighbor.hpp"
namespace {
constexpr double CUTOFF_SIGMA = 3.0;
constexpr double PEAK_SIGMA = 4.0;
// Closer than this to a model atom, the density is the model's own atom out of place.
constexpr double MODEL_ERROR_DISTANCE = 1.6;
// Waters and ions sit at least this far from the model's atoms.
constexpr double CONTACT_MIN = 2.2;
// An ion's peak, in sigma, where the anomalous map gives no help.
constexpr double ION_PEAK_SIGMA = 8.0;
constexpr double ION_ANOMALOUS_SIGMA = 4.0;
// The smallest blob called ligand-like, A^3: at high resolution a strong single atom's blob already spans four of
// the resolution's spheres.
constexpr double LIGAND_MIN_VOLUME = 8.0;
bool Polar(const gemmi::Element &el) {
return el == gemmi::El::N || el == gemmi::El::O;
}
} // namespace
const char *DensityBlobKindName(DensityBlob::Kind kind) {
switch (kind) {
case DensityBlob::Kind::Water: return "WATER";
case DensityBlob::Kind::Ion: return "ION";
case DensityBlob::Kind::Ligand: return "LIGAND";
case DensityBlob::Kind::ModelError: return "MODEL_ERROR";
case DensityBlob::Kind::Unassigned: break;
}
return "UNASSIGNED";
}
std::vector<DensityBlob> FindDensityBlobs(const gemmi::Grid<float> &difference, double rms, const gemmi::Structure &model,
double d_min, const gemmi::Grid<float> *anomalous, double anomalous_rms) {
std::vector<DensityBlob> out;
if (!(rms > 0) || model.models.empty())
return out;
gemmi::Grid<float> map = difference;
for (float &v : map.data)
v = static_cast<float>(v / rms);
gemmi::BlobCriteria criteria;
criteria.cutoff = CUTOFF_SIGMA;
criteria.min_volume = 0.0;
criteria.min_score = 0.0;
criteria.min_peak = PEAK_SIGMA;
const std::vector<gemmi::Blob> blobs = gemmi::find_blobs_by_flood_fill(map, criteria);
gemmi::Model copy = model.models[0];
gemmi::UnitCell cell = map.unit_cell;
cell.set_cell_images_from_spacegroup(map.spacegroup);
gemmi::NeighborSearch ns(copy, cell, 5.0);
ns.populate(false);
// A single atom's blob at this resolution: a sphere of half the resolution's radius. Four of them is ligand size.
const double atom_volume = 4.0 / 3.0 * 3.14159265358979 * std::pow(0.5 * d_min, 3);
for (const gemmi::Blob &b : blobs) {
DensityBlob d;
d.peak_sigma = b.peak_value;
d.volume = b.volume;
d.score = b.score;
d.centroid = b.centroid;
// the copy of the blob beside the atom nearest to it
if (gemmi::NeighborSearch::Mark *m = ns.find_nearest_atom(b.centroid, 12.0)) {
const gemmi::CRA cra = m->to_cra(copy);
const gemmi::FTransform t = ns.get_image_transformation(m->image_idx);
const gemmi::Fractional fc = cell.fractionalize(b.centroid);
const gemmi::Fractional fi = t.apply(cell.fractionalize(cra.atom->pos));
const gemmi::Fractional shifted(fc.x - std::round(fc.x - fi.x), fc.y - std::round(fc.y - fi.y),
fc.z - std::round(fc.z - fi.z));
d.centroid = cell.orthogonalize(gemmi::FTransform(t.inverse()).apply(shifted));
d.nearest_atom_distance = d.centroid.dist(cra.atom->pos);
d.nearest_atom = cra.atom->name + " " + cra.residue->name + " " + cra.chain->name +
cra.residue->seqid.str();
} else
d.nearest_atom_distance = INFINITY;
ns.for_each(d.centroid, '\0', 3.4, [&](gemmi::NeighborSearch::Mark &m, double dist_sq) {
const double r = std::sqrt(dist_sq);
if (Polar(m.element) && r >= 2.4)
d.polar_contacts++;
if ((Polar(m.element) || m.element == gemmi::El::S) && r >= 1.8 && r <= 2.6)
d.metal_contacts++;
});
if (anomalous != nullptr && anomalous_rms > 0)
d.anomalous_sigma = anomalous->interpolate_value(b.peak_pos) / anomalous_rms;
if (d.nearest_atom_distance < MODEL_ERROR_DISTANCE)
d.kind = DensityBlob::Kind::ModelError;
else if (d.nearest_atom_distance >= 1.8 && d.metal_contacts >= 2 &&
(d.anomalous_sigma >= ION_ANOMALOUS_SIGMA || d.peak_sigma >= ION_PEAK_SIGMA))
d.kind = DensityBlob::Kind::Ion;
else if (d.volume > std::max(4 * atom_volume, LIGAND_MIN_VOLUME) && d.nearest_atom_distance >= CONTACT_MIN)
d.kind = DensityBlob::Kind::Ligand;
else if (d.nearest_atom_distance >= CONTACT_MIN && d.polar_contacts >= 1)
d.kind = DensityBlob::Kind::Water;
out.push_back(d);
}
std::sort(out.begin(), out.end(), [](const DensityBlob &a, const DensityBlob &b) { return a.score > b.score; });
return out;
}