Model validation: scale fits on the GPU, change-of-basis probe gridded once per frame on all threads

The scale fit of every model-validation fit (fit_model: the real model, the indexing probe's relabellings and every null replicate) was FitModelScale on the host - a long-double Levenberg-Marquardt per solvent grid point, ~80% of the validation's CPU on a 400 A cell. It now runs on the GPU through ModelScaleGPU::FitSolvent, the engine the rigid body and molecular replacement already use, from a small pool (ModelScaleGPUPool: up to four engines with their own streams, sized on the card's total memory, at most an eighth of it). The grid, the fits and the winner rule are FitModelScale's; the sums are taken in double on the device, so R-factors agree with the host's to rounding. The CPU build, and a GPU build without a card, keep FitModelScale. The change-of-basis probe gridded each candidate frame with gemmi's serial put_model_density_on_grid, once for the data as merged and once more for each alternative indexing. It now grids each frame once, with ModelGrid's parallel PutModelDensityOnGrid (the same grid, bit for bit), from threads outside the pool so the gridding spreads over it, and scores every labelling on that F_calc. Exact. 4nwv (122k atoms, I222, 400 A cell): model validation 33.2 s -> 14.0 s; every logged R, sigma and decision the same. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01EizBKhTqrAYago9KJG3dA3
2026-10-12 21:40:54 +02:00
parent f567771d56
commit 1fe73bab41
4 changed files with 282 additions and 35 deletions
@@ -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
@@ -669,9 +710,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]))
@@ -683,13 +735,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++) {
@@ -737,6 +790,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.
@@ -874,6 +931,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;
@@ -1046,6 +1112,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