From 8933a72a895a05cb966fe49c57de5ce16b28c9a0 Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Sun, 25 Oct 2026 13:48:50 +0100 Subject: [PATCH] Refinement: the weight trials side by side, each on a model, topology and engine of its own Every rung of the weight ladder starts from the same model and depends on no other, so the six position trials, and then the three B trials at the chosen weight, run at the same time: on the main model and engine and on up to five slots beside it, each a copy of the model before the topology was made (the topology is built again on it, hydrogens included) with its own engine. The slots go on the GPUs in turn from the main engine's, as many as each card's budget takes, so a machine with several cards spreads them; a slot whose engine cannot be made is left out, and with none the trials run one after another on the main engine as before. The threads are shared out among the slots. The log lines and the choice are made afterwards in ladder order, so the log, the trial table and the refined model are the same, bit for bit (GPU and CPU engines checked). 5reo 3.9 -> 2.5 s, 9yl4 13.4 -> 7.8 s, 6o2h 3.1 -> 1.9 s, 8dyz 4.6 -> 2.7 s, 9rcs 2.9 -> 2.2 s (one RTX 5080, -N 8). Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_015CeEGgpKdfkCTJb3YNMJcu (cherry picked from commit 0ea6a6db8aba640afc5249fe748922ddfaa2537a) (cherry picked from commit 0c0176c8bf571ed26f8a68a9de40222487c6200e) --- .../structure_refinement/Refinement.cpp | 312 ++++++++++++++---- 1 file changed, 245 insertions(+), 67 deletions(-) diff --git a/image_analysis/structure_refinement/Refinement.cpp b/image_analysis/structure_refinement/Refinement.cpp index b6b0b6d78..a04a80182 100644 --- a/image_analysis/structure_refinement/Refinement.cpp +++ b/image_analysis/structure_refinement/Refinement.cpp @@ -9,6 +9,7 @@ #include #include #include +#include #include #include @@ -101,6 +102,8 @@ struct Context { std::unique_ptr scale_gpu; #endif size_t nthreads = 1; + int device = -1; // the engine's GPU + size_t engine_bytes = 0; // and what it holds there XrayRowModel rows; std::vector> fcalc, fmask; double k_sol = 0, b_sol = 0; @@ -611,6 +614,8 @@ void MakeEngine(Context &ctx, bool gpu, const RefinementSettings &settings, Logg size_t bytes = 0; ctx.engine = MakeRefinementEngineGPU(device, ctx.cell, *ctx.sg, ctx.d_min, ctx.refl, ctx.Model(), ctx.b_ceiling, budget, &bytes, settings.nthreads); + ctx.device = device; + ctx.engine_bytes = bytes; if (ctx.engine) ctx.scale_gpu = std::make_unique(device, ctx.refl.Size(), 1); else @@ -630,6 +635,187 @@ void MakeEngine(Context &ctx, bool gpu, const RefinementSettings &settings, Logg ctx.engine = MakeRefinementEngineCPU(ctx.cell, *ctx.sg, ctx.d_min, ctx.refl, settings.nthreads); } +// The riding hydrogens placed on their parents as they are, and the offsets and B they ride with. +void PlaceRiders(Context &ctx, RefinementTopology &topo) { + const RefinementStage stage("riders"); + PlaceRidingHydrogens(topo); + ctx.rider_offset.assign(ctx.atoms.size(), gemmi::Position()); + for (size_t i = 0; i < ctx.atoms.size(); i++) + if (ctx.rider_parent[i] >= 0) { + ctx.rider_offset[i] = ctx.atoms[i]->pos - ctx.atoms[ctx.rider_parent[i]]->pos; + ctx.atoms[i]->b_iso = ctx.atoms[ctx.rider_parent[i]]->b_iso; + } +} + +// The geometry and the repulsions of the model as it is. +RestraintSet RestraintsNow(Context &ctx, const RefinementTopology &topo) { + const RefinementStage stage("repulsions"); + RestraintSet set = topo.geometry; + set.Append(BuildRepulsions(topo, ctx.Model(), ctx.cell, *ctx.sg, REPULSION_MARGIN, ctx.nthreads)); + set.Finalize(ctx.atoms.size()); + return set; +} + +// Where a weight trial runs: a model with its topology and its engine. +struct TrialSlot { + Context *ctx; + RefinementTopology *topo; +}; + +// A slot of its own beside the main one. +struct OwnSlot { + Context ctx; + RefinementTopology topo; +}; + +// A copy of the model as it was before the topology was made (the same topology comes out of it, hydrogens +// included), with an engine of the main one's kind - a GPU's on `device`. Null where the card cannot make one. +std::unique_ptr MakeOwnSlot(const gemmi::Structure &before_topology, const Context &main, int device, + size_t nthreads, const RefinementSettings &settings) { + auto slot = std::make_unique(); + Context &ctx = slot->ctx; + ctx.st = before_topology; + ctx.cell = main.cell; + ctx.sg = main.sg; + ctx.d_min = main.d_min; + ctx.refl = main.refl; + ctx.nthreads = nthreads; + ctx.b_ceiling = main.b_ceiling; + slot->topo = BuildRefinementTopology(ctx.st, settings.riding_hydrogens); + ctx.atoms = AtomList(ctx.Model()); + ctx.rider_parent = slot->topo.rider_parent; +#ifdef JFJOCH_USE_CUDA + if (main.engine->OnGPU()) { + try { + ctx.engine = MakeRefinementEngineGPU(device, ctx.cell, *ctx.sg, ctx.d_min, ctx.refl, ctx.Model(), + ctx.b_ceiling, SIZE_MAX, nullptr, nthreads); + if (!ctx.engine) + return nullptr; + ctx.scale_gpu = std::make_unique(device, ctx.refl.Size(), 1); + } catch (const JFJochException &) { + return nullptr; + } + ctx.device = device; + } +#else + (void) device; +#endif + if (!ctx.engine) + ctx.engine = MakeRefinementEngineCPU(ctx.cell, *ctx.sg, ctx.d_min, ctx.refl, nthreads); + ctx.engine->SetRestraints(1, slot->topo.adp); + return slot; +} + +// The devices of the slots beside the main one, `wanted` of them at most: the GPUs in turn from the main +// engine's, each taking engines while they fit in its budget. -1 for each one on the CPU. +std::vector OwnSlotDevices(const Context &main, int wanted, const RefinementSettings &settings) { +#ifdef JFJOCH_USE_CUDA + if (main.engine->OnGPU()) { + const int n = get_gpu_count(); + std::vector used(n, 0); + used[main.device] = main.engine_bytes; + std::vector devices; + for (int k = 1; k <= wanted; k++) { + const int d = (main.device + k) % n; + const size_t usable = get_gpu_total_memory(d) / 5 * 4; + const size_t budget = usable > settings.gpu_resident_bytes ? usable - settings.gpu_resident_bytes : 0; + if (used[d] + main.engine_bytes <= budget) { + used[d] += main.engine_bytes; + devices.push_back(d); + } + } + return devices; + } +#else + (void) settings; +#endif + return std::vector(std::max(wanted, 0), -1); +} + +// One weight trial: a full refinement from pos0 and b0, macro-cycles of the positions at X-ray weight w and, +// where adp_w is not 0, the B factors at X-ray weight adp_w. +struct TrialOutcome { + RefinementWeightTrial trial; + std::vector pos; + std::vector b; + std::string line; // what the log says of it +}; + +TrialOutcome RunTrial(TrialSlot slot, const std::vector &refinable, const std::vector &pos0, + const std::vector &b0, double w, double adp_w, const RefinementSettings &settings) { + Context &ctx = *slot.ctx; + RefinementTopology &topo = *slot.topo; + const RefinementStage stage(adp_w > 0 ? "rung.b" : "rung.xyz"); + ctx.SetPositions(pos0); + ctx.SetB(b0); + for (int cycle = 0; cycle < settings.macro_cycles; cycle++) { + PlaceRiders(ctx, topo); + ctx.Refit(); + { + const RestraintSet set = RestraintsNow(ctx, topo); + const RefinementStage s1("set_restraints"); + ctx.engine->SetRestraints(0, set); + } + Minimise(ctx, AtomTarget::Mode::Xyz, refinable, w, 0, settings.xyz_iterations); + if (adp_w > 0) + Minimise(ctx, AtomTarget::Mode::B, refinable, adp_w, 1, settings.b_iterations); + } + PlaceRiders(ctx, topo); + TrialOutcome out; + RefinementWeightTrial &t = out.trial; + t.weight = w; + t.adp_weight = adp_w; + const RFactors rf = ctx.Refit(); + t.r_work = rf.r_work; + t.r_free = rf.r_free; + { + const RefinementStage s1("measure"); + t.geometry = MeasureGeometry(topo.geometry, + BuildRepulsions(topo, ctx.Model(), ctx.cell, *ctx.sg, 0.0, ctx.nthreads), ctx.XyzB()); + } + t.within_bounds = t.geometry.rmsd_bonds <= MAX_RMSD_BONDS && t.geometry.rmsd_angles <= MAX_RMSD_ANGLES; + out.line = fmt::format("Refinement, X-ray weight {:.3g}, B {}: R-work {:.4f}, R-free {:.4f}, rmsd bonds {:.4f} A, " + "angles {:.2f} deg{}", + w, adp_w > 0 ? fmt::format("refined at weight {:.3g}", adp_w) : std::string("held"), + rf.r_work, rf.r_free, t.geometry.rmsd_bonds, t.geometry.rmsd_angles, + t.within_bounds ? "" : " (geometry out of bounds)"); + out.pos = ctx.Positions(); + out.b = ctx.B(); + return out; +} + +// The trials (w, adp_w) side by side: trial i in slot i % slots, each slot's trials in turn on a thread of its +// own (on its GPU), returned in the order given. Every trial starts from pos0 and b0 and takes nothing from the +// one before it in its slot, so the outcome is the same for any number of slots. +std::vector RunTrials(const std::vector &trial_slots, + const std::vector> &trials, + const std::vector &refinable, const std::vector &pos0, + const std::vector &b0, const RefinementSettings &settings) { + std::vector out(trials.size()); + std::vector errors(trial_slots.size()); + auto run_slot = [&](size_t s) { + try { +#ifdef JFJOCH_USE_CUDA + const ScopedGpuPin pin(trial_slots[s].ctx->device >= 0 ? trial_slots[s].ctx->device : get_gpu()); +#endif + for (size_t i = s; i < trials.size(); i += trial_slots.size()) + out[i] = RunTrial(trial_slots[s], refinable, pos0, b0, trials[i].first, trials[i].second, settings); + } catch (...) { + errors[s] = std::current_exception(); + } + }; + std::vector threads; + for (size_t s = 1; s < std::min(trial_slots.size(), trials.size()); s++) + threads.emplace_back(run_slot, s); + run_slot(0); + for (std::thread &t : threads) + t.join(); + for (const std::exception_ptr &e : errors) + if (e) + std::rethrow_exception(e); + return out; +} + } // namespace RefinementResult RefineModel(const gemmi::Structure &placed, const std::vector &merged, @@ -689,23 +875,14 @@ RefinementResult RefineModel(const gemmi::Structure &placed, const std::vector topo_stage("topology"); RefinementTopology topo = BuildRefinementTopology(ctx.st, settings.riding_hydrogens); topo_stage.reset(); // the hydrogens it added: the atoms and the engine made again for them ctx.atoms = AtomList(ctx.Model()); ctx.rider_parent = topo.rider_parent; - auto place_riders = [&]() { - const RefinementStage stage("riders"); - PlaceRidingHydrogens(topo); - ctx.rider_offset.assign(ctx.atoms.size(), gemmi::Position()); - for (size_t i = 0; i < ctx.atoms.size(); i++) - if (ctx.rider_parent[i] >= 0) { - ctx.rider_offset[i] = ctx.atoms[i]->pos - ctx.atoms[ctx.rider_parent[i]]->pos; - ctx.atoms[i]->b_iso = ctx.atoms[ctx.rider_parent[i]]->b_iso; - } - }; - place_riders(); + PlaceRiders(ctx, topo); if (static_cast(ctx.atoms.size()) != result.atoms) { MakeEngine(ctx, result.on_gpu, settings, logger); result.on_gpu = ctx.engine->OnGPU(); @@ -725,16 +902,9 @@ RefinementResult RefineModel(const gemmi::Structure &placed, const std::vectorSetRestraints(1, topo.adp); { - const RestraintSet set = restraints_now(); + const RestraintSet set = RestraintsNow(ctx, topo); ctx.engine->SetRestraints(0, set); RestraintSet repulsions = BuildRepulsions(topo, ctx.Model(), cell, sg, 0.0, ctx.nthreads); result.geometry_start = MeasureGeometry(topo.geometry, repulsions, ctx.XyzB()); @@ -750,67 +920,75 @@ RefinementResult RefineModel(const gemmi::Structure &placed, const std::vector pos0 = ctx.Positions(); const std::vector b0 = ctx.B(); + + // The trials run side by side, as many at a time as the longer ladder has rungs and the cards take engines: + // on the main model and engine, and on copies with engines of their own; each with a share of the threads. + std::optional slot_stage("slots"); + const int wanted = static_cast(std::max(settings.weight_ladder.size(), settings.adp_ladder.size())); + const std::vector devices = OwnSlotDevices(ctx, wanted - 1, settings); + const size_t slot_threads = std::max(1, settings.nthreads / (devices.size() + 1)); + std::vector> own(devices.size()); + { + std::vector threads; + for (size_t k = 0; k < devices.size(); k++) + threads.emplace_back([&, k]() { +#ifdef JFJOCH_USE_CUDA + const ScopedGpuPin pin(devices[k] >= 0 ? devices[k] : get_gpu()); +#endif + own[k] = MakeOwnSlot(before_topology, ctx, devices[k], slot_threads, settings); + }); + for (std::thread &t : threads) + t.join(); + } + const size_t main_threads = ctx.nthreads; + ctx.nthreads = slot_threads; + std::vector trial_slots{{&ctx, &topo}}; + for (const std::unique_ptr &o : own) + if (o) + trial_slots.push_back({&o->ctx, &o->topo}); + slot_stage.reset(); + logger.Debug("Refinement: {} weight trials at a time", trial_slots.size()); + int best = -1; std::vector best_pos; std::vector best_b; - // One full refinement from the start: macro-cycles of the positions at X-ray weight w and, where adp_w is - // not 0, the B factors at X-ray weight adp_w; kept where it beats the best so far. The rungs come in order - // of trust - tighter geometry, then fewer parameters - and a later one displaces an earlier only by more - // than the tie in R-free. - auto trial = [&](double w, double adp_w) { - const RefinementStage stage(adp_w > 0 ? "rung.b" : "rung.xyz"); - ctx.SetPositions(pos0); - ctx.SetB(b0); - for (int cycle = 0; cycle < settings.macro_cycles; cycle++) { - place_riders(); - ctx.Refit(); - { - const RestraintSet set = restraints_now(); - const RefinementStage s1("set_restraints"); - ctx.engine->SetRestraints(0, set); + // A trial is kept where it beats the best so far. The rungs come in order of trust - tighter geometry, + // then fewer parameters - and a later one displaces an earlier only by more than the tie in R-free. + auto choose = [&](std::vector &outcomes) { + for (TrialOutcome &o : outcomes) { + logger.Info("{}", o.line); + result.trials.push_back(o.trial); + const RefinementWeightTrial &t = result.trials.back(); + const int k = static_cast(result.trials.size()) - 1; + const bool better = best < 0 ? true + : t.within_bounds && (!result.trials[best].within_bounds || + t.r_free < result.trials[best].r_free - R_FREE_TIE); + if (better) { + best = k; + best_pos = std::move(o.pos); + best_b = std::move(o.b); } - Minimise(ctx, AtomTarget::Mode::Xyz, refinable, w, 0, settings.xyz_iterations); - if (adp_w > 0) - Minimise(ctx, AtomTarget::Mode::B, refinable, adp_w, 1, settings.b_iterations); - } - place_riders(); - RefinementWeightTrial t; - t.weight = w; - t.adp_weight = adp_w; - const RFactors rf = ctx.Refit(); - t.r_work = rf.r_work; - t.r_free = rf.r_free; - { - const RefinementStage s1("measure"); - t.geometry = MeasureGeometry(topo.geometry, BuildRepulsions(topo, ctx.Model(), cell, sg, 0.0, ctx.nthreads), ctx.XyzB()); - } - t.within_bounds = t.geometry.rmsd_bonds <= MAX_RMSD_BONDS && t.geometry.rmsd_angles <= MAX_RMSD_ANGLES; - logger.Info("Refinement, X-ray weight {:.3g}, B {}: R-work {:.4f}, R-free {:.4f}, rmsd bonds {:.4f} A, " - "angles {:.2f} deg{}", - w, adp_w > 0 ? fmt::format("refined at weight {:.3g}", adp_w) : std::string("held"), rf.r_work, - rf.r_free, t.geometry.rmsd_bonds, t.geometry.rmsd_angles, - t.within_bounds ? "" : " (geometry out of bounds)"); - result.trials.push_back(t); - const int k = static_cast(result.trials.size()) - 1; - const bool better = best < 0 ? true - : t.within_bounds && (!result.trials[best].within_bounds || - t.r_free < result.trials[best].r_free - R_FREE_TIE); - if (better) { - best = k; - best_pos = ctx.Positions(); - best_b = ctx.B(); } }; // the positions' weight, with the B factors as they are + std::vector> xyz_trials; for (double ladder : settings.weight_ladder) - trial(w0 * ladder, 0.0); + xyz_trials.emplace_back(w0 * ladder, 0.0); + std::vector outcomes = RunTrials(trial_slots, xyz_trials, refinable, pos0, b0, settings); + choose(outcomes); // then the B factors refined too, at that weight, against their own weight const double w = result.trials[best].weight; + std::vector> b_trials; for (double adp : settings.adp_ladder) - trial(w, w * adp); + b_trials.emplace_back(w, w * adp); + outcomes = RunTrials(trial_slots, b_trials, refinable, pos0, b0, settings); + choose(outcomes); + for (size_t s = 1; s < trial_slots.size(); s++) + ctx.evaluations += trial_slots[s].ctx->evaluations; + ctx.nthreads = main_threads; ctx.SetPositions(best_pos); ctx.SetB(best_b); - place_riders(); + PlaceRiders(ctx, topo); const RefinementWeightTrial &chosen = result.trials[best]; result.restrained = true; result.weight = chosen.weight;