Files
Jungfraujoch/image_analysis/structure_refinement/RefinementXray.cpp
T
leonarski_f 5569f3d1f2 Refinement GPU: the X-ray target and the restraints in one submission with one wait
RefinementEngine::EvaluateWithRestraints evaluates both targets at the model's atoms. On the GPU the X-ray
chain is queued first, the restraints' coordinates are written into a pinned buffer while the device works
on it, the restraint kernels follow on the same stream, and the host waits once; the sums and gradients
come back through pinned buffers. Same kernels on the same inputs, same bits. The CPU engine keeps calling
the two in turn.

9yl4: evaluation 5.1 -> 4.7 ms, refinement 14.2 -> 13.4 s; 5reo 4.0 -> 3.9 s (RTX 5080, -N 8).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_015CeEGgpKdfkCTJb3YNMJcu
(cherry picked from commit 51434b98b62a9e6322ffe811bf272a45ff6ce003)
(cherry picked from commit 0bcfa3f8041cec4bb4bc376b8559f30c36d3ac80)
2026-10-26 01:38:23 +01:00

433 lines
19 KiB
C++

// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only
#include "RefinementXray.h"
#include <algorithm>
#include <cmath>
#include <map>
#include <mutex>
#include <fftw3.h>
#include "gemmi/dencalc.hpp"
#include "gemmi/it92.hpp"
#include "gemmi/math.hpp"
#include "../../common/Logger.h"
#include "../../common/ParallelFor.h"
#include "ModelFFT.h"
#include "ModelGrid.h"
#include "RefinementMath.h"
namespace {
using Table = gemmi::IT92<float>;
// The free reflections a likelihood bin is fitted on, at least.
constexpr size_t FREE_PER_BIN = 100;
constexpr size_t MAX_BINS = 30;
// The density calculator the refinement grids atoms with: DensityCalculator's settings for the grid of
// RigidBodyZoneGrid(), and the given blur.
gemmi::DensityCalculator<Table, float> Calculator(const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg,
double d_min, double blur) {
gemmi::DensityCalculator<Table, float> dc;
dc.d_min = d_min;
dc.rate = 1.5; // RigidBodyRefine.cpp's GRID_RATE, which RigidBodyZoneGrid sizes the grid with
dc.grid.unit_cell = cell;
dc.grid.spacegroup = &sg;
dc.blur = blur;
return dc;
}
// -ln L summed over rows `idx` at D and beta.
double BinNegLogLikelihood(const RefinementReflections &refl, const std::vector<double> &fc,
const std::vector<size_t> &idx, double d, double beta) {
double sum = 0, dfc = 0;
for (size_t i : idx) {
const double var = refl.epsilon[i] * beta + double(refl.sigma[i]) * refl.sigma[i];
sum += refine_math::AmplitudeNegLogLikelihood<double>(refl.fo[i], fc[i], d, var, refl.centric[i], dfc);
}
return sum;
}
// The minimum of f over [lo, hi] by golden-section search.
template <typename F>
double GoldenMinimum(F f, double lo, double hi, int iterations) {
const double g = 0.5 * (std::sqrt(5.0) - 1.0);
double a = lo, b = hi;
double c = b - g * (b - a), d = a + g * (b - a);
double fc = f(c), fd = f(d);
for (int i = 0; i < iterations; i++) {
if (fc < fd) {
b = d;
d = c;
fd = fc;
c = b - g * (b - a);
fc = f(c);
} else {
a = c;
c = d;
fc = fd;
d = a + g * (b - a);
fd = f(d);
}
}
return 0.5 * (a + b);
}
class RefinementEngineCPU : public RefinementEngine {
public:
RefinementEngineCPU(const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg, double d_min,
const RefinementReflections &refl, size_t nthreads)
: cell_(cell), sg_(sg), d_min_(d_min), refl_(refl), nthreads_(nthreads),
grid_(RigidBodyZoneGrid(cell, sg, d_min)), composition_(grid_, d_min, refl.hkl),
orbit_leaders_(OrbitLeaders(grid_, nthreads)) {}
void SetBlur(double blur) override { blur_ = blur; }
void SetRowModel(const XrayRowModel &rows) override { rows_ = rows; }
bool OnGPU() const override { return false; }
void Fcalc(const gemmi::Model &model, std::vector<std::complex<float>> &f) override {
std::vector<std::complex<double>> fd;
Compose(model, fd);
f.assign(fd.begin(), fd.end());
}
void Fmask(const gemmi::Model &model, std::vector<std::complex<float>> &f) override {
gemmi::Grid<float> mask = grid_;
PutMaskOnGrid(mask, model, orbit_leaders_, nthreads_);
const gemmi::FPhiGrid<float> fm = MapToFPhi(mask);
f.clear();
for (const gemmi::Miller &h : refl_.hkl)
f.push_back(fm.get_value_by_hkl(h));
}
double Evaluate(const gemmi::Model &model, std::vector<std::array<double, 4>> *grad) override {
std::vector<std::complex<double>> f;
Compose(model, f);
const size_t n = f.size();
// The target and, per reflection, dT/dF of the atoms' structure factor (dT/dRe F + i dT/dIm F).
std::vector<std::complex<double>> dt_df(n, 0.0);
double sum = 0;
for (size_t m = 0; m < n; m++) {
if (refl_.free[m])
continue;
const double k = rows_.k_total[m];
const std::complex<double> fm = k * (f[m] + std::complex<double>(rows_.fmask[m]));
const double fc = std::abs(fm);
double dfc = 0;
sum += refine_math::AmplitudeNegLogLikelihood<double>(refl_.fo[m], fc, rows_.d[m], rows_.var[m],
refl_.centric[m], dfc);
if (fc > 0)
dt_df[m] = k * dfc * fm / fc;
}
if (grad == nullptr)
return sum;
// The FFT method's map, dT/drho(x) = (V/N) Re sum_m sum_ops conj(dT/dF_m) scale_m exp(2 pi i h.t)
// exp(2 pi i (hR).x), as the inverse real transform of a half-l coefficient grid: a term A at index k
// goes in as A/2 at k, or as conj(A)/2 at -k where k is in the half not stored; on the planes l = 0 and
// l = nw/2, which hold both, as both.
const int nu = grid_.nu, nv = grid_.nv, nw = grid_.nw, nwh = nw / 2 + 1;
const size_t ncomplex = static_cast<size_t>(nu) * nv * nwh;
std::complex<float> *in = reinterpret_cast<std::complex<float> *>(fftwf_alloc_complex(ncomplex));
float *out = fftwf_alloc_real(static_cast<size_t>(nu) * nv * nw);
std::fill(in, in + ncomplex, std::complex<float>(0, 0));
const std::vector<SymmetryComposition::Term> &terms = composition_.Terms();
const size_t ops = composition_.Ops();
auto add = [&](int ku, int kv, int kw, std::complex<double> value) {
in[gemmi::modulo(ku, nu) + static_cast<size_t>(nu) * (gemmi::modulo(kv, nv) + static_cast<size_t>(nv) * kw)] +=
std::complex<float>(value);
};
for (size_t m = 0; m < n; m++) {
if (dt_df[m] == 0.0)
continue;
const double scale = composition_.Centring() * std::exp(blur_ * 0.25 * composition_.InvD2()[m]);
const std::complex<double> c = std::conj(dt_df[m]) * scale;
for (size_t o = 0; o < ops; o++) {
const SymmetryComposition::Term &t = terms[m * ops + o];
const std::complex<double> a = 0.5 * c * t.phase;
const gemmi::Miller &k = t.k;
if (k[2] > 0 && 2 * k[2] != nw)
add(k[0], k[1], k[2], a);
else if (k[2] < 0 && -2 * k[2] != nw)
add(-k[0], -k[1], -k[2], std::conj(a));
else {
const int l = k[2] < 0 ? -k[2] : k[2];
add(k[0], k[1], l, a);
add(-k[0], -k[1], l, std::conj(a));
}
}
}
MapFromFPhiExecute(nu, nv, nw, in, out);
gemmi::Grid<float> map = grid_;
const float norm = static_cast<float>(cell_.volume / static_cast<double>(map.point_count()));
for (size_t i = 0; i < map.data.size(); i++)
map.data[i] = out[i] * norm;
fftwf_free(in);
fftwf_free(out);
AtomGradientsFromMap(map, model, blur_, nthreads_, *grad);
return sum;
}
private:
void Compose(const gemmi::Model &model, std::vector<std::complex<double>> &f) const {
gemmi::DensityCalculator<Table, float> dc = Calculator(cell_, sg_, d_min_, blur_);
PutModelDensityOnGrid(dc, model, {}, nthreads_); // one copy: the symmetry is composed below
composition_.Compose(MapToFPhi(dc.grid), blur_, f, nullptr, nthreads_);
}
gemmi::UnitCell cell_;
const gemmi::SpaceGroup &sg_;
double d_min_;
const RefinementReflections &refl_;
size_t nthreads_;
gemmi::Grid<float> grid_;
SymmetryComposition composition_;
std::vector<size_t> orbit_leaders_;
double blur_ = 0;
XrayRowModel rows_;
};
} // namespace
RefinementReflections PrepareRefinementReflections(const std::vector<MergedReflection> &merged,
const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg,
double d_min) {
const gemmi::GroupOps gops = sg.operations();
const gemmi::ReciprocalAsu asu(&sg);
// ASU, last one wins where a reflection is listed twice (both hands of an anomalous merge)
std::map<gemmi::Miller, const MergedReflection *> by_hkl;
for (const MergedReflection &r : merged) {
if (!std::isfinite(r.F) || r.F < 0)
continue;
gemmi::Miller h{{r.h, r.k, r.l}};
if (!asu.is_in(h))
h = asu.to_asu(h, gops).first;
by_hkl[h] = &r;
}
std::vector<gemmi::Miller> hkl;
for (const auto &[h, r] : by_hkl)
hkl.push_back(h);
const SymmetryComposition composition(RigidBodyZoneGrid(cell, sg, d_min), d_min, hkl);
RefinementReflections out;
for (size_t i = 0; i < hkl.size(); i++) {
if (composition.Row()[i] < 0)
continue;
const MergedReflection &r = *by_hkl[hkl[i]];
out.hkl.push_back(hkl[i]);
out.fo.push_back(r.F);
out.sigma.push_back(std::isfinite(r.sigmaF) ? r.sigmaF : 0.0f);
out.free.push_back(r.rfree_flag);
out.centric.push_back(gops.is_reflection_centric(hkl[i]));
out.epsilon.push_back(static_cast<float>(gops.epsilon_factor(hkl[i])));
out.inv_d2.push_back(cell.calculate_1_d2(hkl[i]));
}
return out;
}
void EstimateLikelihoodParameters(const RefinementReflections &refl, const std::vector<std::complex<float>> &fmodel,
XrayRowModel &rows, size_t nthreads) {
const size_t n = refl.Size();
std::vector<double> fc(n);
for (size_t i = 0; i < n; i++)
fc[i] = std::abs(fmodel[i]);
std::vector<size_t> free;
for (size_t i = 0; i < n; i++)
if (refl.free[i] && fc[i] > 0)
free.push_back(i);
// Too few free reflections to bin: the working set, which says the model is better than it is, but
// still gives the likelihood a scale.
if (free.size() < FREE_PER_BIN)
for (size_t i = 0; i < n; i++)
if (!refl.free[i] && fc[i] > 0)
free.push_back(i);
std::sort(free.begin(), free.end(), [&](size_t a, size_t b) { return refl.inv_d2[a] < refl.inv_d2[b]; });
const size_t bins = std::clamp<size_t>(free.size() / FREE_PER_BIN, 1, MAX_BINS);
std::vector<double> centre(bins), bin_d(bins), bin_beta(bins);
// each bin is fitted on its own
ParallelFor(static_cast<int>(bins), nthreads, [&](int b) {
const std::vector<size_t> idx(free.begin() + b * free.size() / bins, free.begin() + (b + 1) * free.size() / bins);
double s2 = 0, sfo = 0, sfc2 = 0, sfofc = 0;
for (size_t i : idx) {
s2 += refl.inv_d2[i];
sfo += double(refl.fo[i]) * refl.fo[i] / refl.epsilon[i];
sfc2 += fc[i] * fc[i] / refl.epsilon[i];
sfofc += refl.fo[i] * fc[i] / refl.epsilon[i];
}
centre[b] = s2 / static_cast<double>(idx.size());
double d = sfc2 > 0 ? std::max(sfofc / sfc2, 0.05) : 1.0;
double beta = std::max((sfo - d * d * sfc2) / static_cast<double>(idx.size()), 1e-3 * sfo / idx.size());
// Alternating one-dimensional minimisations of -ln L over ln D and ln beta.
for (int round = 0; round < 4; round++) {
const double ld = GoldenMinimum([&](double x) { return BinNegLogLikelihood(refl, fc, idx, std::exp(x), beta); },
std::log(d) - 2.0, std::log(std::min(d * 7.4, 3.0)), 40);
d = std::exp(ld);
const double lb = GoldenMinimum([&](double x) { return BinNegLogLikelihood(refl, fc, idx, d, std::exp(x)); },
std::log(beta) - 3.0, std::log(beta) + 3.0, 40);
beta = std::exp(lb);
}
bin_d[b] = d;
bin_beta[b] = beta;
});
rows.d.resize(n);
rows.var.resize(n);
for (size_t i = 0; i < n; i++) {
// linear in 1/d^2 between the bins' centres, flat beyond the outermost
const double s = refl.inv_d2[i];
double d = bin_d[0], beta = bin_beta[0];
if (s >= centre[bins - 1]) {
d = bin_d[bins - 1];
beta = bin_beta[bins - 1];
} else if (s > centre[0]) {
const size_t b = std::upper_bound(centre.begin(), centre.end(), s) - centre.begin();
const double t = (s - centre[b - 1]) / (centre[b] - centre[b - 1]);
d = bin_d[b - 1] + t * (bin_d[b] - bin_d[b - 1]);
beta = bin_beta[b - 1] + t * (bin_beta[b] - bin_beta[b - 1]);
}
rows.d[i] = static_cast<float>(d);
rows.var[i] = static_cast<float>(refl.epsilon[i] * beta + double(refl.sigma[i]) * refl.sigma[i]);
}
}
RFactors ComputeRFactors(const RefinementReflections &refl, const std::vector<std::complex<float>> &fmodel) {
double nw = 0, dw = 0, nf = 0, df = 0;
RFactors r;
for (size_t i = 0; i < refl.Size(); i++) {
const double diff = std::fabs(refl.fo[i] - std::abs(fmodel[i]));
if (refl.free[i]) {
nf += diff;
df += refl.fo[i];
r.n_free++;
} else {
nw += diff;
dw += refl.fo[i];
r.n_work++;
}
}
r.r_work = dw > 0 ? nw / dw : 1;
r.r_free = df > 0 ? nf / df : 1;
return r;
}
std::vector<std::array<double, 4>> ModelXyzB(const gemmi::Model &model) {
std::vector<std::array<double, 4>> out;
for (const gemmi::Chain &ch : model.chains)
for (const gemmi::Residue &r : ch.residues)
for (const gemmi::Atom &a : r.atoms)
out.push_back({a.pos.x, a.pos.y, a.pos.z, a.b_iso});
return out;
}
namespace {
std::mutex stage_mutex;
std::map<std::string, std::pair<double, int>> stage_times;
} // namespace
RefinementStage::RefinementStage(const char *name)
: name_(name), start_(std::chrono::steady_clock::now()), nvtx_(NvtxRangeStart(name)) {}
RefinementStage::~RefinementStage() {
NvtxRangeEnd(nvtx_);
const double seconds = std::chrono::duration<double>(std::chrono::steady_clock::now() - start_).count();
std::lock_guard lock(stage_mutex);
std::pair<double, int> &t = stage_times[name_];
t.first += seconds;
t.second++;
}
void RefinementStage::Log(Logger &logger) {
std::lock_guard lock(stage_mutex);
for (const auto &[name, t] : stage_times)
logger.Debug("TIMING refine.{} {:.3f} {}", name, t.first, t.second);
stage_times.clear();
}
double RefinementBlur(const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg, double d_min, double b_floor) {
// set_refmac_compatible_blur(): u_to_b / 1.1 * spacing^2 - b_min, not below 0, the spacing the one asked for
const gemmi::DensityCalculator<Table, float> dc = Calculator(cell, sg, d_min, 0);
const double spacing = dc.requested_grid_spacing();
return std::max(0.0, gemmi::u_to_b() / 1.1 * spacing * spacing - b_floor);
}
std::unique_ptr<RefinementEngine> MakeRefinementEngineCPU(const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg,
double d_min, const RefinementReflections &refl,
size_t nthreads) {
return std::make_unique<RefinementEngineCPU>(cell, sg, d_min, refl, nthreads);
}
void AtomGradientsFromMap(const gemmi::Grid<float> &map, const gemmi::Model &model, double blur, size_t nthreads,
std::vector<std::array<double, 4>> &grad) {
const gemmi::DensityCalculator<Table, float> dc = Calculator(map.unit_cell, *map.spacegroup, 1.0, blur);
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);
grad.assign(atoms.size(), {0, 0, 0, 0});
const int nu = map.nu, nv = map.nv, nw = map.nw;
// As ModelGrid.cpp walks a density box: the orthogonalization over the grid size in single precision.
const gemmi::UpperTriangularMat33 &om = map.orth_n;
const float m00 = static_cast<float>(om.a11), m01 = static_cast<float>(om.a12), m02 = static_cast<float>(om.a13),
m11 = static_cast<float>(om.a22), m12 = static_cast<float>(om.a23), m22 = static_cast<float>(om.a33);
constexpr double FOUR_PI2 = 4 * refine_math::PI * refine_math::PI;
ParallelChunks(static_cast<int>(atoms.size()), nthreads, [&](int lo, int hi) {
for (int i = lo; i < hi; i++) {
const gemmi::Atom &atom = *atoms[i];
const float b = static_cast<float>(atom.b_iso + blur);
const auto precal = Table::get(atom.element, atom.charge, atom.serial).precalculate_density_iso(b, 0.0f);
const double radius = dc.estimate_radius(precal, b);
const float max_r2 = static_cast<float>(radius * radius);
const gemmi::Fractional fpos = map.unit_cell.fractionalize(atom.pos);
const double n3[3] = {fpos.x * nu, fpos.y * nv, fpos.z * nw};
float nctr[3];
int c[3], d[3];
for (int k = 0; k < 3; k++) {
nctr[k] = static_cast<float>(n3[k]);
c[k] = gemmi::iround(n3[k]);
}
d[0] = std::min(static_cast<int>(std::ceil(radius / map.spacing[0])), nu - 1);
d[1] = std::min(static_cast<int>(std::ceil(radius / map.spacing[1])), nv - 1);
d[2] = std::min(static_cast<int>(std::ceil(radius / map.spacing[2])), nw - 1);
double gx = 0, gy = 0, gz = 0, gb = 0;
for (int w = c[2] - d[2]; w <= c[2] + d[2]; w++) {
const float fz = nctr[2] - static_cast<float>(w);
const float z = m22 * fz;
const int w_ = gemmi::modulo(w, nw);
for (int v = c[1] - d[1]; v <= c[1] + d[1]; v++) {
const float fy = nctr[1] - static_cast<float>(v);
const float y = m11 * fy + m12 * fz;
const float yz = y * y + z * z;
if (yz > max_r2)
continue;
const float x_yz = m01 * fy + m02 * fz;
const size_t row = static_cast<size_t>(nu) * (gemmi::modulo(v, nv) + static_cast<size_t>(nv) * w_);
for (int u = c[0] - d[0]; u <= c[0] + d[0]; u++) {
const float x = m00 * (nctr[0] - static_cast<float>(u)) + x_yz;
const float r2 = yz + x * x;
if (r2 > max_r2)
continue;
const double value = map.data[row + gemmi::modulo(u, nu)];
double dr2 = 0, db = 0;
for (size_t q = 0; q < std::size(precal.a); q++) {
const double e = precal.a[q] * gemmi::unsafe_expapprox(std::max(precal.b[q] * r2, -88.f));
dr2 += e * precal.b[q];
db += e * precal.b[q] * (1.5 + precal.b[q] * r2);
}
gx += value * dr2 * x;
gy += value * dr2 * y;
gz += value * dr2 * z;
gb += value * db;
}
}
}
// rho = occ sum_q a_q exp(b_q r^2): d rho / d(centre) = occ 2 (sum a_q b_q e_q) (centre - point), and
// with b_q = -4 pi^2 / (B_q + B + blur), d rho / dB = occ sum a_q e_q b_q (1.5 + b_q r^2) / (4 pi^2).
grad[i] = {2.0 * atom.occ * gx, 2.0 * atom.occ * gy, 2.0 * atom.occ * gz, atom.occ * gb / FOUR_PI2};
}
});
}