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

Merge branch 'gpu-modelval' into rc175: --model validation's scale fits on the GPU, change-of-basis scoring gridded once in parallel

Browse Source 
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
This commit is contained in:
leonarski_f
2026-10-13 02:47:06 +02:00
co-authored by Claude Opus 5.5
parent 85e714f072 1fe73bab41
commit 9cb19753c7
4 changed files with 282 additions and 35 deletions
Download Patch File Download Diff File Expand all files Collapse all files
image_analysis/structure_refinement/ModelValidation.cpp
+105 -35
View File
@@ -344,49 +344,63 @@ constexpr double FRAME_PROBE_D_MAX = 8.0;
constexpr size_t FRAME_PROBE_MIN_REFLECTIONS = 200;
// R of the model against the observed amplitudes after an overall scale and an isotropic B, over
// that coarse shell. Only the ranking is ever used, never the value.
// that coarse shell - the lowest over `labellings`, the labellings of the data it is scored against,
// each on the same F_calc. Only the ranking is ever used, never the value.
double frame_probe_r(const gemmi::Structure &st, const gemmi::SpaceGroup *sg,
const std::vector<MergedReflection> &obs, double d_min) {
const std::vector<const std::vector<MergedReflection> *> &labellings, double d_min,
size_t nthreads) {
const double probe_d_min = std::max(d_min, FRAME_PROBE_D_MIN);
gemmi::DensityCalculator<Table, float> dc;
dc.d_min = probe_d_min;
dc.rate = 1.5;
dc.set_grid_cell_and_spacegroup(st);
dc.set_refmac_compatible_blur(st.models[0]);
dc.put_model_density_on_grid(st.models[0]);
// put_model_density_on_grid() spread over the threads, the same grid bit for bit (ModelGrid.h); the
// orbits are those of the grid initialize_grid() sizes.
gemmi::Grid<float> sized;
sized.unit_cell = dc.grid.unit_cell;
sized.spacegroup = dc.grid.spacegroup;
sized.set_size_from_spacing(dc.requested_grid_spacing(), gemmi::GridSizeRounding::Up);
PutModelDensityOnGrid(dc, st.models[0], OrbitLeaders(sized, nthreads), nthreads);
gemmi::AsuData<std::complex<float>> fcalc =
MapToFPhi(dc.grid)
.prepare_asu_data(probe_d_min, dc.blur, false, false, false);
gemmi::GroupOps gops = sg->operations();
gemmi::ReciprocalAsu asu(sg);
gemmi::AsuData<gemmi::ValueSigma<float>> fobs;
fobs.unit_cell_ = st.cell;
fobs.spacegroup_ = sg;
// The low-resolution cap is dropped, rather than the probe abandoned, when a small cell does not
// put enough reflections in the shell.
for (double d_max : {FRAME_PROBE_D_MAX, 1e9}) {
fobs.v.clear();
for (const MergedReflection &r : obs) {
if (std::isnan(r.F) || r.d < probe_d_min || r.d > d_max)
continue;
gemmi::Miller h{{r.h, r.k, r.l}};
if (!asu.is_in(h))
h = asu.to_asu(h, gops).first;
fobs.v.push_back({h, {r.F, 1.0f}});
double best = 0;
for (size_t i = 0; i < labellings.size(); i++) {
gemmi::AsuData<gemmi::ValueSigma<float>> fobs;
fobs.unit_cell_ = st.cell;
fobs.spacegroup_ = sg;
// The low-resolution cap is dropped, rather than the probe abandoned, when a small cell does not
// put enough reflections in the shell.
for (double d_max : {FRAME_PROBE_D_MAX, 1e9}) {
fobs.v.clear();
for (const MergedReflection &r : *labellings[i]) {
if (std::isnan(r.F) || r.d < probe_d_min || r.d > d_max)
continue;
gemmi::Miller h{{r.h, r.k, r.l}};
if (!asu.is_in(h))
h = asu.to_asu(h, gops).first;
fobs.v.push_back({h, {r.F, 1.0f}});
}
if (fobs.v.size() >= FRAME_PROBE_MIN_REFLECTIONS)
break;
}
if (fobs.v.size() >= FRAME_PROBE_MIN_REFLECTIONS)
break;
double r = 1.0;
if (!fobs.v.empty()) {
fobs.ensure_asu();
fobs.ensure_sorted();
gemmi::Scaling<float> scaling(st.cell, sg);
scaling.use_solvent = false;
scaling.prepare_points(fcalc, fobs, nullptr);
scaling.fit_isotropic_b_approximately();
r = scaling.calculate_r_factor();
}
best = i == 0 ? r : std::min(best, r);
}
if (fobs.v.empty())
return 1.0;
fobs.ensure_asu();
fobs.ensure_sorted();
gemmi::Scaling<float> scaling(st.cell, sg);
scaling.use_solvent = false;
scaling.prepare_points(fcalc, fobs, nullptr);
scaling.fit_isotropic_b_approximately();
return scaling.calculate_r_factor();
return best;
}
// The anomalous difference map's coefficients on the reflections of `fmodel`: F(+) - F(-) carried on the
@@ -516,6 +530,33 @@ std::unique_ptr<ModelStructureFactorsGPU> StructureFactorsGPU(int device, const
"card's {:.2f} GB", d_min, n[0], n[1], n[2], sf->DeviceBytes() / 1e9, total / 1e9);
return sf;
}
// The scale fits on the GPU (ModelScaleGPUPool) - or not, decided here, once, on the card's total memory as
// the structure factors are: SCALE_GPU_ENGINES engines, fewer where they would take over an eighth of the
// card, none where not even one fits.
constexpr size_t SCALE_GPU_ENGINES = 4;
std::unique_ptr<ModelScaleGPUPool> ScaleFitsGPU(int device, size_t max_points, Logger &logger) {
const size_t total = get_gpu_total_memory(device);
size_t engines = SCALE_GPU_ENGINES;
while (engines > 0 && ModelScaleGPUPool::DeviceBytes(max_points, engines) > total / 8)
engines--;
if (engines == 0) {
logger.Info("Model validation: scale fits on the CPU - one fit of {} reflections needs {:.2f} GB of GPU "
"memory, over an eighth of the card's {:.2f} GB", max_points,
ModelScaleGPUPool::DeviceBytes(max_points, 1) / 1e9, total / 1e9);
return nullptr;
}
std::unique_ptr<ModelScaleGPUPool> pool;
try {
pool = std::make_unique<ModelScaleGPUPool>(device, max_points, engines);
} catch (const JFJochException &e) {
throw RigidBodyGPUFailure(e.what());
}
logger.Info("Model validation: scale fits on the GPU, {} engine(s), {:.2f} GB of the card's {:.2f} GB",
engines, pool->Bytes() / 1e9, total / 1e9);
return pool;
}
#endif
} // namespace
@@ -671,9 +712,20 @@ ModelValidationResult Validate(const std::vector<MergedReflection> &merged,
sg->xhm(), frames.size() - 1);
// The candidates are scored at the same time, each on its own copy of the model, and read
// in their own order below, so the ranking and its tie-break are those of one after another.
// On threads of their own, not on ParallelFor's pool, so that each one's gridding spreads over
// the pool (a parallel pass reached from a pool worker runs inline, ParallelFor.h).
std::vector<double> frame_r(frames.size(), 0.0);
std::vector<const gemmi::SpaceGroup *> frame_sg(frames.size(), nullptr);
ParallelFor(static_cast<int>(frames.size()), nthreads, [&](int i) {
// The data as merged and, where the alternative indexings are left to the probe below, in each
// of them: the lowest R over these is the frame's.
std::vector<std::vector<MergedReflection>> relabelled;
if (same_point_group)
for (const gemmi::Op &law : reindex_ops)
relabelled.push_back(ReindexReflections(obs, law));
std::vector<const std::vector<MergedReflection> *> labellings{&obs};
for (const auto &v : relabelled)
labellings.push_back(&v);
auto score_frame = [&](size_t i) {
gemmi::Structure trial = st;
const gemmi::SpaceGroup *trial_sg = sg;
if (i > 0 && !change_model_basis(trial, trial_sg, frames[i]))
@@ -685,13 +737,14 @@ ModelValidationResult Validate(const std::vector<MergedReflection> &merged,
// as every other and the ranking is noise. Measured on a pseudo-orthorhombic P2
// crystal: every frame read R 0.77-0.81 in the merged indexing, and the frame chosen
// on that noise swapped two axes the model shares with the data.
double r = frame_probe_r(trial, trial_sg, obs, d_min);
if (same_point_group)
for (const gemmi::Op &law : reindex_ops)
r = std::min(r, frame_probe_r(trial, trial_sg, ReindexReflections(obs, law), d_min));
frame_r[i] = r;
frame_r[i] = frame_probe_r(trial, trial_sg, labellings, d_min, nthreads);
frame_sg[i] = trial_sg;
});
};
std::vector<std::future<void>> scored;
for (size_t i = 0; i < frames.size(); i++)
scored.push_back(std::async(std::launch::async, score_frame, i));
for (std::future<void> &f : scored)
f.get();
size_t best = 0;
double best_r = 0;
for (size_t i = 0; i < frames.size(); i++) {
@@ -739,6 +792,10 @@ ModelValidationResult Validate(const std::vector<MergedReflection> &merged,
sf_engine = StructureFactorsGPU(RigidBodyGPUEngine::CurrentDevice(), st.models[0], ucell, *sg, d_min, 0,
logger);
sf_gpu = sf_engine.get();
// The scale fits beside them, on the same card, for as many reflections as the data have.
std::unique_ptr<ModelScaleGPUPool> scale_gpu;
if (rigid_body_gpu && get_gpu_count() > 0)
scale_gpu = ScaleFitsGPU(RigidBodyGPUEngine::CurrentDevice(), obs.size(), logger);
#endif
// --- Fcalc (atomic) via electron density on a grid + FFT, plus a flat bulk-solvent mask -> Fmask.
@@ -876,6 +933,15 @@ ModelValidationResult Validate(const std::vector<MergedReflection> &merged,
gemmi::Scaling<float> scaling(ucell, sg);
scaling.use_solvent = true;
scaling.prepare_points(out.fmodel, fobs_work, &ms.fmask);
#ifdef JFJOCH_USE_CUDA
if (scale_gpu) {
try {
scale_gpu->Fit(scaling, SolventBMax(ms.st.models[0]), nthreads);
} catch (const JFJochException &e) {
throw RigidBodyGPUFailure(e.what());
}
} else
#endif
FitModelScale(scaling, {.b_hi = SolventBMax(ms.st.models[0])}, nthreads);
scaling.scale_data(out.fmodel, &ms.fmask); // out.fmodel now holds the scaled, solvent-corrected Fmodel
out.k_sol = scaling.k_sol;
@@ -1042,6 +1108,10 @@ ModelValidationResult Validate(const std::vector<MergedReflection> &merged,
planned += (decision_pending ? 2 : 1) * sf_gpu->DeviceBytes();
card = sf_gpu->DeviceTotalMemory();
}
if (scale_gpu) {
planned += scale_gpu->Bytes();
card = get_gpu_total_memory(RigidBodyGPUEngine::CurrentDevice());
}
if (planned > 0)
may_run_beside = 2 * planned <= static_cast<size_t>(get_gpu_count()) * card / 2;
#endif