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/ModelValidation.cpp
+88 -8
View File
@@ -10,6 +10,7 @@
#include <array>
#include <atomic>
#include <functional>
#include <memory>
#include <future>
#include <numeric>
#include <optional>
@@ -38,8 +39,13 @@
#include "RigidBodyRefine.h"
#include "SigmaA.h"
#ifdef JFJOCH_USE_CUDA
#include "ModelStructureFactorsGPU.h"
#include "RigidBodyGPU.h"
#include "RigidBodyGPUEngine.h" // RigidBodyGPUEngine::CurrentDevice
#include "../../common/CUDAWrapper.h"
#include "../../common/JFJochException.h"
#else
class ModelStructureFactorsGPU;
#endif
namespace {
@@ -51,8 +57,19 @@ long hkl_key(const gemmi::Miller &h) {
return (h[0] + 512L) * 1048576 + (h[1] + 512L) * 1024 + (h[2] + 512L);
}
// FFT ASU map coefficients into a real-space map.
gemmi::Grid<float> map_from_coefficients(gemmi::AsuData<std::complex<float>> &coef) {
// FFT ASU map coefficients into a real-space map, on `gpu` where the validation has one.
gemmi::Grid<float> map_from_coefficients(gemmi::AsuData<std::complex<float>> &coef, ModelStructureFactorsGPU *gpu) {
#ifdef JFJOCH_USE_CUDA
if (gpu != nullptr) {
try {
return gpu->Map(coef);
} catch (const JFJochException &e) {
throw RigidBodyGPUFailure(e.what());
}
}
#else
(void) gpu;
#endif
coef.ensure_sorted();
std::array<int, 3> size = gemmi::get_size_for_hkl(coef, {{0, 0, 0}}, 3.0);
return MapFromFPhi(gemmi::get_f_phi_on_grid<float>(coef, size, true));
@@ -335,13 +352,53 @@ double frame_probe_r(const gemmi::Structure &st, const gemmi::SpaceGroup *sg,
return scaling.calculate_r_factor();
}
#ifdef JFJOCH_USE_CUDA
// The model's structure factors to the data's resolution, and the maps, on the GPU - or not, decided here,
// once, before any of them is computed, and never revisited: nothing is moved to the CPU part way through.
// Decided on what the GPU path needs against the card's TOTAL memory and not against what happens to be
// free at the moment, so the same input on the same machine always takes the same path, whatever else is
// running beside it. The structure factors may take half the card; the rigid body's engines take at most
// a quarter beside them (RigidBodyGPUPool::Create), and the last quarter is left to everything else.
std::unique_ptr<ModelStructureFactorsGPU> StructureFactorsGPU(int device, const gemmi::Model &model,
const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg,
double d_min, Logger &logger) {
auto sf = std::make_unique<ModelStructureFactorsGPU>(device, model, cell, sg, d_min);
const std::array<int, 3> n = sf->GridSize();
if (!sf->Supported()) {
logger.Info("Model validation: structure factors and maps on the CPU - the cell is too small for the "
"GPU's gridding");
return nullptr;
}
size_t total = 0;
try {
total = sf->DeviceTotalMemory();
} catch (const JFJochException &e) {
throw RigidBodyGPUFailure(e.what());
}
if (sf->DeviceBytes() > total / 2) {
logger.Info("Model validation: structure factors and maps on the CPU - on a {}x{}x{} grid they need "
"{:.2f} GB of GPU memory, over half of the card's {:.2f} GB", n[0], n[1], n[2],
sf->DeviceBytes() / 1e9, total / 1e9);
return nullptr;
}
try {
sf->Reserve();
} catch (const JFJochException &e) {
throw RigidBodyGPUFailure(e.what());
}
logger.Info("Model validation: structure factors and maps on the GPU - a {}x{}x{} grid, {:.2f} GB of the "
"card's {:.2f} GB", n[0], n[1], n[2], sf->DeviceBytes() / 1e9, total / 1e9);
return sf;
}
#endif
} // namespace
namespace {
// ValidateAgainstModel() with the rigid body on the GPU where `rigid_body_gpu` allows it and a card is
// there to take it, on the CPU otherwise - for the whole validation either way.
// ValidateAgainstModel() with the rigid body, the structure factors and the maps on the GPU where
// `rigid_body_gpu` allows it and a card is there to take them, on the CPU otherwise - for the whole
// validation either way.
ModelValidationResult Validate(const std::vector<MergedReflection> &merged,
const UnitCell &cell,
const std::string &model_path,
@@ -528,10 +585,33 @@ ModelValidationResult Validate(const std::vector<MergedReflection> &merged,
logger.Info("Model validation: {} atoms, cell a={:.2f} b={:.2f} c={:.2f}, sg {}, to {:.2f} A",
gemmi::count_atom_sites(st.models[0]), ucell.a, ucell.b, ucell.c, sg->hm, d_min);
// The structure factors to d_min and the maps on the GPU, where `rigid_body_gpu` allows it and the
// card can take them; decided here, before any is computed. Anything to a coarser limit - the null's,
// the indexing probe's - stays on the CPU, where it is cheap, so that every replicate of the null and
// the real model's side of it are computed the same way.
ModelStructureFactorsGPU *sf_gpu = nullptr;
#ifdef JFJOCH_USE_CUDA
std::unique_ptr<ModelStructureFactorsGPU> sf_engine;
// On the current device, as the rigid body's engines take it: one card for the whole validation.
if (rigid_body_gpu && get_gpu_count() > 0)
sf_engine = StructureFactorsGPU(RigidBodyGPUEngine::CurrentDevice(), st.models[0], ucell, *sg, d_min, logger);
sf_gpu = sf_engine.get();
#endif
// --- Fcalc (atomic) via electron density on a grid + FFT, plus a flat bulk-solvent mask -> Fmask.
// A lambda because the rigid-body step below moves the model and then needs both again, and it
// takes the state to work on so that a null replicate can run it on its own copy. ---
auto compute_model_factors = [&](ModelState &ms, double to_d) {
#ifdef JFJOCH_USE_CUDA
if (sf_gpu != nullptr && to_d == sf_gpu->DMin()) {
try {
sf_gpu->Compute(ms.st.models[0], ms.fcalc, ms.fmask);
} catch (const JFJochException &e) {
throw RigidBodyGPUFailure(e.what());
}
return;
}
#endif
gemmi::DensityCalculator<Table, float> dc;
dc.d_min = to_d;
dc.rate = 1.5;
@@ -1265,7 +1345,7 @@ ModelValidationResult Validate(const std::vector<MergedReflection> &merged,
// The three maps - this one, the difference map and the anomalous map - share nothing but what
// they read, so the other two are made and written on threads of their own beside this one.
std::future<void> fofc_map = std::async(std::launch::async, [&] {
write_ccp4(map_from_coefficients(mapfofc), output_prefix + "_fofc.ccp4");
write_ccp4(map_from_coefficients(mapfofc, sf_gpu), output_prefix + "_fofc.ccp4");
});
// --- anomalous difference map, where the merge kept the Bijvoet split ---
@@ -1312,7 +1392,7 @@ ModelValidationResult Validate(const std::vector<MergedReflection> &merged,
result.anomalous_pairs = static_cast<int>(mapanom.v.size());
if (!mapanom.v.empty()) {
const gemmi::Grid<float> grid = map_from_coefficients(mapanom);
const gemmi::Grid<float> grid = map_from_coefficients(mapanom, sf_gpu);
const double rms = write_ccp4(grid, output_prefix + "_anom.ccp4");
std::vector<ModelValidationResult::AnomalousSite> sites;
double scatterer_sum = 0;
@@ -1355,7 +1435,7 @@ ModelValidationResult Validate(const std::vector<MergedReflection> &merged,
}
});
const gemmi::Grid<float> grid2fofc = map_from_coefficients(map2fofc);
const gemmi::Grid<float> grid2fofc = map_from_coefficients(map2fofc, sf_gpu);
const double rms2 = write_ccp4(grid2fofc, output_prefix + "_2fofc.ccp4");
{
double s = 0; int n = 0;
@@ -1490,7 +1570,7 @@ ModelValidationResult ValidateAgainstModel(const std::vector<MergedReflection> &
double wavelength_A,
const std::vector<float> &report_shell_d_min) {
#ifdef JFJOCH_USE_CUDA
// A CUDA failure in the rigid body does not end the run: the validation is started again from the
// A CUDA failure on the GPU path does not end the run: the validation is started again from the
// model as read, on the CPU throughout. It is re-runnable, and a dead model check should not take a
// finished merge with it.
try {