mx/Jungfraujoch
3
2
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 81 Wiki Activity

rugnux: model validation's structure factors and maps on the GPU

Browse Source 
ModelStructureFactorsGPU computes what compute_model_factors() and
map_from_coefficients() compute on the CPU - F_calc from the model's
density (IT92, Refmac-compatible blur, unblurred as prepare_asu_data()
does) and F_mask from the Refmac bulk-solvent mask, both on the
reflections prepare_asu_data(d_min) lists, in its order; and a map from
ASU coefficients on the grid get_size_for_hkl(coef, 0, 3.0) sizes - on a
device. Made once per cell, group, resolution and model, then evaluated
as often as the coordinates change, so refinement or MR can call it in a
loop. The device is an explicit parameter; every call leaves the calling
thread's current device as it found it.

Pieces:
- ModelDensityGPU: the rigid body's deterministic brick gather, moved
  out of RigidBodyGPU.cu into a component of its own (ModelMaskGPU's
  pattern); the rigid body uses it unchanged. MAX_BRICKS_PER_AXIS 8 ->
  16, so fine grids with high-B atoms (lysozyme at 1.2 A, a 0.9 A P1
  cell) are no longer refused; existing zones are gridded identically.
- One copy of the content is gridded and the symmetry composed in
  reciprocal space (SymmetryComposition), operators applied on the fly;
  the mask is ModelMaskGPU (every image of every atom, islands, shrink).
- Maps: gemmi's get_f_phi_on_grid() in ZYX order on the host (the
  coefficients written are the same), in-place cuFFT c2r, transposed back
  to XYZ on the device. One map at a time, in the engine's buffers.

Decided once, up front, per card, from its TOTAL memory: the engine's
bytes (16 N + cuFFT work + reflections, N the larger of the structure-
factor and map grids) must be at most half the card - the rigid body's
engines take at most a quarter beside it. Otherwise, or where the gather
cannot grid the cell, the CPU path runs, logged with needed vs total.
Anything to a resolution other than d_min (the null's 3.5 A fits) stays
on the CPU, so all replicates and the real model's side of the null are
computed the same way. A CUDA failure takes the existing path: the
validation restarts on the CPU.

Measured, model validation total per run (CPU path -> GPU), 16 GB card:
  F432 215 A cubic, 1.30 A, 500^3 grid: 47.7 -> 15.6 s (two validations;
     14.3 -> 2.4 and 33.4 -> 13.2, the rest of the second is writing the
     three 0.5 GB maps); F_calc + F_mask 5.7 s -> 0.05 s
  C2 1.11 A: 23.9 -> 13.6 s; P3_2 1.55 A: 18.2 -> 10.2 s;
  P2_1 1.25 A: 13.4 -> 6.5 s; F4_132 328 A: 13.0 -> 5.5 s;
  P6_5: 8.8 -> 4.2 s; P4_3 0.97 A: 4.6 -> 2.5 s; P1 0.92 A: 4.2 -> 2.2 s;
  small P1: 3.2 -> 1.3 s; P6_1: 8.8 -> 6.2 s; lysozyme: 1.8 -> 1.4 s.
p.mtz md5-identical on all 13 sets. Against the CPU path: FC within
1e-4 of mean |F|, phases of the strong half within 0.003 deg, maps within
1e-4 (2mFo-DFc) and 7e-4 (mFo-DFc) of their rms; every logged R, CC,
FOM, k_sol and anomalous site list identical at the printed precision,
except where a rigid-body commit sat on an exact R-free tie (0.2155 ->
0.2155) and fell the other way (R-work 0.2127 vs 0.2129). The GPU result
is bit-identical run to run and with -N 8 (maps, map MTZ, placed model).
Peak device memory of the engine: 2.5 GB at 500^3 (process total peaked
at 14.4 GB with what the merge still holds).

Tests: ModelStructureFactorsGPU_MatchesCPU (five groups, 3.5 and 1.5 A:
same reflections, F_calc <= 1e-5 of mean |F|, F_mask 2e-7 rms, repeat
bit-identical), ModelStructureFactorsGPU_MapMatchesCPU (<= 5e-6 of rms).

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01SVmAWnzCmRKAXVUCdc4iNi
This commit is contained in:
leonarski_f
2026-10-09 12:38:37 +02:00
co-authored by Claude Opus 5.5
parent 6025d51a48
commit d5fcf2f05b
13 changed files with 1427 additions and 359 deletions
Download Patch File Download Diff File Expand all files Collapse all files
image_analysis/structure_refinement/ModelStructureFactorsGPU.cpp
+254
View File
@@ -0,0 +1,254 @@
// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only
// The reflection list follows ReciprocalGrid::prepare_asu_data() of GEMMI's recgrid.hpp
// (https://github.com/project-gemmi/gemmi/blob/master/include/gemmi/recgrid.hpp)
// (c) Global Phasing Ltd., Mozilla Public License Version 2.0
#include "ModelStructureFactorsGPU.h"
#include <algorithm>
#include <cmath>
#include "gemmi/fourier.hpp" // get_size_for_hkl, get_f_phi_on_grid
#include "gemmi/solmask.hpp" // SolventMasker, refmac_radius_for_bulk_solvent
#include "RigidBodyRefine.h" // RigidBodyZoneGrid
namespace {
using Table = gemmi::IT92<float>;
// The density calculator model validation grids F_calc with: d_min, rate 1.5 and the blur gemmi's
// set_refmac_compatible_blur() gives this model.
gemmi::DensityCalculator<Table, float> Calculator(const gemmi::Model &model, const gemmi::UnitCell &cell,
const gemmi::SpaceGroup &sg, double d_min) {
gemmi::DensityCalculator<Table, float> dc;
dc.d_min = d_min;
dc.rate = 1.5;
dc.grid.unit_cell = cell;
dc.grid.spacegroup = &sg;
dc.set_refmac_compatible_blur(model);
return dc;
}
// The reflections prepare_asu_data(d_min) lists for the half-l transform of an (nu, nv, nw) grid, in its
// order (h, then k, then l ascending): in the group's reciprocal ASU, strictly inside d_min, inside the
// grid's Nyquist box, not systematically absent and not 000.
std::vector<gemmi::Miller> AsuReflections(const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg, double d_min,
int nu, int nv, int nw) {
const gemmi::Miller lim = cell.get_hkl_limits(d_min);
const int max_h = std::min((nu - 1) / 2, lim[0]);
const int max_k = std::min((nv - 1) / 2, lim[1]);
const int max_l = std::min(nw / 2, lim[2]);
const double max_1_d2 = 1.0 / (d_min * d_min);
const gemmi::ReciprocalAsu asu(&sg);
const gemmi::GroupOps gops = sg.operations();
std::vector<gemmi::Miller> rows;
gemmi::Miller hkl;
for (hkl[0] = -max_h; hkl[0] <= max_h; ++hkl[0])
for (hkl[1] = -max_k; hkl[1] <= max_k; ++hkl[1])
for (hkl[2] = -max_l; hkl[2] <= max_l; ++hkl[2])
if (asu.is_in(hkl) && cell.calculate_1_d2(hkl) < max_1_d2 && !gops.is_systematically_absent(hkl) &&
!(hkl[0] == 0 && hkl[1] == 0 && hkl[2] == 0))
rows.push_back(hkl);
return rows;
}
} // namespace
void ModelDensityAtoms(const gemmi::Model &model, const gemmi::DensityCalculator<Table, float> &dc,
std::vector<ModelDensityAtom> &atoms, std::vector<int> &mask_atom,
std::vector<float> &mask_radius) {
atoms.clear();
mask_atom.clear();
mask_radius.clear();
const gemmi::SolventMasker masker(gemmi::AtomicRadiiSet::Refmac);
int index = 0;
for (const gemmi::Chain &ch : model.chains)
for (const gemmi::Residue &r : ch.residues)
for (const gemmi::Atom &atom : r.atoms) {
using CReal = Table::Coef::coef_type;
const auto &coef = Table::get(atom.element, atom.charge, atom.serial);
const float addend = dc.addends.get(atom.element);
ModelDensityAtom a{};
a.occ = atom.occ;
a.aniso = atom.aniso.nonzero();
if (!a.aniso) {
const CReal b = static_cast<CReal>(atom.b_iso + dc.blur);
const auto precal = coef.precalculate_density_iso(b, addend);
a.radius = dc.estimate_radius(precal, b);
for (int k = 0; k < 5; k++) {
a.a[k] = precal.a[k];
a.b[k][0] = precal.b[k];
}
} else {
const auto aniso_b = atom.aniso.scaled(CReal(gemmi::u_to_b())).added_kI(CReal(dc.blur));
const CReal b_max = std::max(std::max(aniso_b.u11, aniso_b.u22), aniso_b.u33);
a.radius = static_cast<float>(dc.estimate_radius(coef.precalculate_density_iso(b_max, addend), b_max));
const auto precal = coef.precalculate_density_aniso_b(aniso_b, addend);
for (int k = 0; k < 5; k++) {
a.a[k] = precal.a[k];
const gemmi::SMat33<float> &m = precal.b[k];
const float e[6] = {m.u11, m.u22, m.u33, m.u12, m.u13, m.u23};
std::copy(e, e + 6, a.b[k]);
}
}
atoms.push_back(a);
if (!((masker.ignore_hydrogen && atom.is_hydrogen()) ||
(masker.ignore_zero_occupancy_atoms && atom.occ <= 0))) {
mask_atom.push_back(index);
mask_radius.push_back(static_cast<float>(
masker.constant_r + masker.rprobe + gemmi::refmac_radius_for_bulk_solvent(atom.element.elem)));
}
++index;
}
}
ModelStructureFactorsGPU::ModelStructureFactorsGPU(int device, const gemmi::Model &model, const gemmi::UnitCell &cell,
const gemmi::SpaceGroup &sg, double d_min)
: device_(device), cell_(cell), sg_(&sg), d_min_(d_min), dc_(Calculator(model, cell, sg, d_min)) {
const gemmi::Grid<float> grid = RigidBodyZoneGrid(cell, sg, d_min); // DensityCalculator's for 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;
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);
}
rows_ = AsuReflections(cell, sg, d_min, grid.nu, grid.nv, grid.nw);
const double n_cen = static_cast<double>(gops.cen_ops.size());
for (const gemmi::Miller &h : rows_) {
setup_.rows.push_back({h[0], h[1], h[2]});
// prepare_asu_data()'s unblur, exp(B_blur |s|^2 / 4)
setup_.row_scale.push_back(n_cen * std::exp(dc_.blur * 0.25 * cell.calculate_1_d2(h)));
}
std::vector<ModelDensityAtom> atoms;
std::vector<int> mask_atom;
std::vector<float> mask_radius;
ModelDensityAtoms(model, dc_, atoms, mask_atom, mask_radius);
setup_.max_atoms = atoms.size();
setup_.max_pairs = ModelDensityGPU::PairBound(setup_.grid, atoms);
supported_ = !rows_.empty() && ModelDensityGPU::Supports(setup_.grid, atoms) &&
setup_.mask_ops.size() <= ModelMaskGPU::MAX_OPS;
if (!supported_)
return;
// The largest map: Map()'s coefficients are on a subset of these reflections, and get_size_for_hkl()
// only grows with the indices and the resolution it is given.
gemmi::AsuData<std::complex<float>> all;
all.unit_cell_ = cell;
all.spacegroup_ = &sg;
for (const gemmi::Miller &h : rows_)
all.v.push_back({h, {1.0f, 0.0f}});
map_size_ = gemmi::get_size_for_hkl(all, {{0, 0, 0}}, 3.0);
// MapFromFPhi() transforms a half-l grid of nw / 2 + 1 planes back to 2 (nw / 2) of them.
const int map_nw = 2 * (map_size_[2] / 2);
setup_.map_points = static_cast<size_t>(map_size_[0]) * map_size_[1] * map_nw;
setup_.map_complex_points = static_cast<size_t>(map_size_[0]) * map_size_[1] * (map_nw / 2 + 1);
fft_work_bytes_ = ModelStructureFactorsGPUEngine::FFTWorkBytes(device_, setup_.grid, map_size_[0], map_size_[1], map_nw);
device_bytes_ = ModelStructureFactorsGPUEngine::DeviceBytes(setup_, fft_work_bytes_);
}
ModelStructureFactorsGPU::~ModelStructureFactorsGPU() = default;
void ModelStructureFactorsGPU::Reserve() {
std::lock_guard lock(m_);
if (!engine_)
engine_ = std::make_unique<ModelStructureFactorsGPUEngine>(device_, setup_, fft_work_bytes_);
}
void ModelStructureFactorsGPU::Compute(const gemmi::Model &model, gemmi::AsuData<std::complex<float>> &fcalc,
gemmi::AsuData<std::complex<float>> &fmask) {
std::vector<ModelDensityAtom> atoms;
std::vector<int> mask_atom;
std::vector<float> mask_radius;
ModelDensityAtoms(model, dc_, atoms, mask_atom, mask_radius);
// Each atom in grid units and, for the mask, fractional - both wrapped into the cell.
const int n3[3] = {setup_.grid.nu, setup_.grid.nv, setup_.grid.nw};
std::vector<std::array<float, 4>> pos;
std::vector<std::array<double, 3>> frac;
for (const gemmi::Chain &ch : model.chains)
for (const gemmi::Residue &r : ch.residues)
for (const gemmi::Atom &a : r.atoms) {
const gemmi::Fractional f0 = cell_.fractionalize(a.pos);
std::array<double, 3> f{f0.x, f0.y, f0.z};
std::array<float, 4> g{};
for (int k = 0; k < 3; k++) {
f[k] -= std::floor(f[k]);
if (f[k] >= 1.0) // a tiny negative coordinate wraps to 1 - epsilon, which rounds to 1
f[k] = 0.0;
g[k] = static_cast<float>(f[k] * n3[k]);
if (g[k] >= n3[k])
g[k] -= n3[k];
}
frac.push_back(f);
pos.push_back(g);
}
std::vector<ModelMaskAtom> mask_atoms;
for (size_t j = 0; j < mask_atom.size(); j++) {
const std::array<double, 3> &f = frac[mask_atom[j]];
mask_atoms.push_back({f[0], f[1], f[2], mask_radius[j]});
}
std::vector<std::array<float, 2>> fc, fm;
{
std::lock_guard lock(m_);
engine_->Compute(atoms, pos, mask_atoms, fc, fm);
}
fcalc.v.clear();
fmask.v.clear();
fcalc.v.reserve(rows_.size());
fmask.v.reserve(rows_.size());
for (size_t i = 0; i < rows_.size(); i++) {
fcalc.v.push_back({rows_[i], {fc[i][0], fc[i][1]}});
fmask.v.push_back({rows_[i], {fm[i][0], fm[i][1]}});
}
fcalc.unit_cell_ = fmask.unit_cell_ = cell_;
fcalc.spacegroup_ = fmask.spacegroup_ = sg_;
}
gemmi::Grid<float> ModelStructureFactorsGPU::Map(gemmi::AsuData<std::complex<float>> &coef) {
coef.ensure_sorted();
const std::array<int, 3> size = gemmi::get_size_for_hkl(coef, {{0, 0, 0}}, 3.0);
// ZYX: l fastest and halved, the layout a c2r transform takes, where gemmi's XYZ grid halves the slowest
// axis. The coefficients written are the same.
const gemmi::FPhiGrid<float> hkl = gemmi::get_f_phi_on_grid<float>(coef, size, true, gemmi::AxisOrder::ZYX);
gemmi::Grid<float> map;
map.spacegroup = coef.spacegroup_;
map.unit_cell = coef.unit_cell_;
// As MapFromFPhi(): 2 (nw - 1) planes back from nw stored ones.
map.set_size(size[0], size[1], 2 * (hkl.nu - 1));
map.axis_order = gemmi::AxisOrder::XYZ;
std::lock_guard lock(m_);
engine_->Map(map.nu, map.nv, map.nw, map.unit_cell.volume, reinterpret_cast<const float *>(hkl.data.data()),
map.data.data());
return map;
}