From 4b9ef226ef1174c332f93386e8d4e5d0d13e93d2 Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Wed, 28 Oct 2026 05:48:33 +0100 Subject: [PATCH] Refinement: trials placed over the cards - a queue of slots, longest first, dependent trials run ahead where the slots take them On 4 cards (4YBB, 4x T4, 12 slots) two phases ran as two dependent waves, the second leaving card 3 idle: the ladder's NCS rungs (run at the B weight the rungs without NCS chose) and the displacement trials' intermediate TLS cuts (run where the finest cut beats one body per chain). Where the slots take the whole set at once, these now run with the others, and those not needed are left out with their evaluations (not counted, so the logged count is the same whatever the slots). In RunTrials the first wave is one trial per slot in slot order (breadth-first over the cards), then each slot takes the next trial as it is done (instead of a fixed j += slots), ordered TLS, U, B refined, B held. Every trial starts from its given model and fit and takes nothing from its slot, so the refined model, the trials kept and the report are unchanged; new test Refinement_trials_run_ahead (CPU, 8 slots vs 1) checks it bit for bit on two copies of a tRNA. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01HVeDQqN9qRD57vnShzTC37 --- .../structure_refinement/Refinement.cpp | 73 ++++++++++++++---- tests/RefinementTest.cpp | 74 +++++++++++++++++++ 2 files changed, 131 insertions(+), 16 deletions(-) diff --git a/image_analysis/structure_refinement/Refinement.cpp b/image_analysis/structure_refinement/Refinement.cpp index e8cb56cf8..1584b78dc 100644 --- a/image_analysis/structure_refinement/Refinement.cpp +++ b/image_analysis/structure_refinement/Refinement.cpp @@ -4,6 +4,7 @@ #include "Refinement.h" #include +#include #include #include #include @@ -1528,6 +1529,7 @@ struct TrialOutcome { Context::Fit fit; // the scale, the solvent and the likelihood's parameters fitted where it ended std::string line; // what the log says of it bool ran = false; // false: left out, the refinement cancelled before it + int evaluations = 0; }; TrialOutcome RunTrial(TrialSlot slot, const std::vector &refinable, const std::vector &pos0, @@ -1537,6 +1539,7 @@ TrialOutcome RunTrial(TrialSlot slot, const std::vector &refinable, const s Context &ctx = *slot.ctx; RefinementTopology &topo = *slot.topo; const RefinementStage stage(hyp != nullptr ? "rung.adp" : adp_w > 0 ? "rung.b" : "rung.xyz"); + const int evaluations_before = ctx.evaluations; ctx.SetPositions(pos0); ctx.SetB(b0); ctx.SetAniso(aniso0); @@ -1678,12 +1681,14 @@ TrialOutcome RunTrial(TrialSlot slot, const std::vector &refinable, const s out.tls_b = ctx.tls_b; out.fit = ctx.Fitted(); out.ran = true; + out.evaluations = ctx.evaluations - evaluations_before; return out; } -// The trials (w, adp_w) side by side, the longest to the slots first and 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. +// The trials (w, adp_w) side by side, each slot on a thread of its own (on its GPU), returned in the order given: the +// longest first, one to each slot, then each slot takes the next one as it is done. 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 and +// whichever slot runs it. std::vector RunTrials(const std::vector &trial_slots, const std::vector &trials, const std::vector &refinable, const std::vector &pos0, @@ -1696,18 +1701,23 @@ std::vector RunTrials(const std::vector &trial_slots, const size_t running = std::min(trial_slots.size(), trials.size()); for (size_t s = 0; s < running; s++) trial_slots[s].ctx->nthreads = std::max(1, settings.nthreads / running); - // The trials that refine B are the longest: they go to the slots first, so that two of them share a slot (or a - // card) only where the slots are fewer than they are; the first slots are on cards of their own (OwnSlotPlaces). + // The longest go to the slots first - TLS bodies, then a U per atom, then the B factors refined, then held - so + // that two of them share a card only where the cards are fewer than they are: the slots are on the cards in turn + // (OwnSlotPlaces), a card's second slot only after every card's first. + auto length = [&](size_t i) { + const AdpHypothesis::Kind kind = hyps.empty() ? AdpHypothesis::Kind::Isotropic : hyps[i]->kind; + return kind == AdpHypothesis::Kind::Tls ? 3 : kind == AdpHypothesis::Kind::Anisotropic ? 2 : trials[i].adp > 0; + }; std::vector order(trials.size()); std::iota(order.begin(), order.end(), 0); - std::stable_sort(order.begin(), order.end(), - [&](size_t a, size_t b) { return (trials[a].adp > 0) > (trials[b].adp > 0); }); + std::stable_sort(order.begin(), order.end(), [&](size_t a, size_t b) { return length(a) > length(b); }); + std::atomic next{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 j = s; j < trials.size(); j += trial_slots.size()) + for (size_t j = s; j < trials.size(); j = next++) if (const size_t i = order[j]; !trial_slots[s].ctx->Cancelled()) out[i] = RunTrial(trial_slots[s], refinable, pos0, b0, aniso0, fit, trials[i].w, trials[i].adp, cycles, settings, ncs, trials[i].ncs, hyps.empty() ? nullptr : hyps[i]); @@ -2371,6 +2381,14 @@ RefinementResult RefineOnce(const gemmi::Structure &placed, const std::vectorctx, &o->topo}); slot_stage.reset(); logger.Debug("Refinement: {} weight trials at a time", trial_slots.size()); + // Trials whose need is decided by others are run with them where the slots take them all at once, and left + // out where it is decided they were not needed - with their evaluations, which do not count: the outcome and + // the count are the same whatever the slots. + int ahead_evaluations = 0; + auto drop_ahead = [&](TrialOutcome &o) { + ahead_evaluations += o.evaluations; + o = TrialOutcome(); + }; // A trial displaces another only where its geometry is within bounds and its R-free lower by more than the // tie. @@ -2399,7 +2417,8 @@ RefinementResult RefineOnce(const gemmi::Structure &placed, const std::vector ladder; for (double rung : settings.weight_ladder) for (double adp : settings.adp_ladder) { @@ -2420,11 +2439,23 @@ RefinementResult RefineOnce(const gemmi::Structure &placed, const std::vector= ladder.size(); + run_ladder([&](const TrialWeights &l) { return l.ncs == 0 || ncs_ahead; }); + // the best of those without NCS restraints, those with them set aside + std::vector with_ncs(ladder.size()); + for (size_t i = 0; i < ladder.size(); i++) + if (ladder[i].ncs > 0) + std::swap(with_ncs[i], outcomes[i]); if (const int b = best_of(outcomes); ncs_p != nullptr && b >= 0) { const double adp = outcomes[b].trial.adp_weight; - run_ladder([&](const TrialWeights &l) { return l.ncs > 0 && l.adp == adp; }); + for (size_t i = 0; i < ladder.size(); i++) + if (ladder[i].ncs > 0 && ladder[i].adp == adp) + std::swap(with_ncs[i], outcomes[i]); + if (!ncs_ahead) + run_ladder([&](const TrialWeights &l) { return l.ncs > 0 && l.adp == adp; }); } + for (TrialOutcome &o : with_ncs) + drop_ahead(o); int k = choose(outcomes, "ladder"); TrialOutcome state; if (k >= 0) @@ -2780,8 +2811,9 @@ RefinementResult RefineOnce(const gemmi::Structure &placed, const std::vector= trials.size(); std::vector now(trials.size()); for (size_t k = 0; k < trials.size(); k++) - now[k] = !later[k]; + now[k] = !later[k] || later_ahead; run(now); auto fits_better = [&](int a, int b) { return outcomes[a].ran && outcomes[b].ran && outcomes[a].trial.nll_free < outcomes[b].trial.nll_free; }; - if (finest > coarsest + 1 && fits_better(finest, coarsest)) - run(later); + if (finest > coarsest + 1 && fits_better(finest, coarsest)) { + if (!later_ahead) + run(later); + } else + for (size_t k = 0; k < trials.size(); k++) + if (later[k]) + drop_ahead(outcomes[k]); int best = -1; // where the weights ended auto trial = [&](int k) -> const RefinementWeightTrial & { return k < 0 ? state.trial : outcomes[k].trial; }; for (size_t k = 0; k < outcomes.size(); k++) { @@ -2879,6 +2917,9 @@ RefinementResult RefineOnce(const gemmi::Structure &placed, const std::vectorevaluations; + ctx.evaluations -= ahead_evaluations; + if (ahead_evaluations > 0) + logger.Debug("Refinement: {} target evaluations in trials run ahead and not needed", ahead_evaluations); ctx.nthreads = main_threads; if (state.ran) { ctx.SetPositions(state.pos); diff --git a/tests/RefinementTest.cpp b/tests/RefinementTest.cpp index e86a7242c..ca8ef9887 100644 --- a/tests/RefinementTest.cpp +++ b/tests/RefinementTest.cpp @@ -841,6 +841,10 @@ TEST_CASE("Refinement_restrained_recovers_model", "[Refinement]") { const RefinementResult s = RefineModel(start, merged, cell, sg, 1.0, logger, one); REQUIRE(s.ok); CHECK(s.r_free == r.r_free); + // the same trials kept, those run ahead where the slots take them left out + CHECK(s.trials.size() == r.trials.size()); + CHECK(s.adp_trials.size() == r.adp_trials.size()); + CHECK(s.adp_model == r.adp_model); const std::vector a_side = atoms_of(r.model->models[0]), a_one = atoms_of(s.model->models[0]); REQUIRE(a_one.size() == a_side.size()); size_t same = 0; @@ -871,6 +875,76 @@ TEST_CASE("Refinement_restrained_recovers_model", "[Refinement]") { } } +// Two copies of a tRNA: with a slot for every ladder trial, the NCS rungs run ahead at both B weights and the +// intermediate TLS cuts with the others, those not needed left out; one trial at a time runs only what is needed. +// The same trials kept and the same model, bit for bit. +TEST_CASE("Refinement_trials_run_ahead", "[Refinement]") { + gemmi::Structure truth = gemmi::read_structure_gz("../../tests/test_data/ligands/1ehz_psu.pdb.gz"); + truth.models.resize(1); + gemmi::remove_hydrogens(truth.models[0]); + truth.cell = gemmi::UnitCell(120, 120, 120, 90, 90, 90); + truth.spacegroup_hm = "P 1"; + truth.setup_cell_images(); + gemmi::Chain copy = truth.models[0].chains.at(0); + copy.name = "B"; + const double a = 0.7; + const gemmi::Mat33 rot(std::cos(a), -std::sin(a), 0, std::sin(a), std::cos(a), 0, 0, 0, 1); + for (gemmi::Residue &r : copy.residues) + for (gemmi::Atom &at : r.atoms) + at.pos = gemmi::Position(rot.multiply(at.pos)) + gemmi::Position(40, 10, -5); + truth.models[0].chains.push_back(copy); + truth.connections.clear(); // the LINK records name chain A only + gemmi::Structure start = truth; + start.models[0] = shaken(truth.models[0], 0.3, 5.0, 21); + + const gemmi::SpaceGroup &sg = *gemmi::find_spacegroup_by_name("P 1"); + const double d_min = 5.0; + std::vector merged = asu_reflections(truth.cell, sg, d_min); + { + RefinementReflections refl = PrepareRefinementReflections(merged, truth.cell, sg, d_min * 0.999); + auto cpu = MakeRefinementEngineCPU(truth.cell, sg, d_min * 0.999, refl, 8); + cpu->SetBlur(RefinementBlur(truth.cell, sg, d_min, 5.0)); + std::vector> f; + cpu->Fcalc(truth.models[0], f); + std::map amplitude; + for (size_t m = 0; m < refl.Size(); m++) + amplitude[refl.hkl[m]] = std::abs(f[m]); + for (MergedReflection &r : merged) { + const auto it = amplitude.find(gemmi::Miller{{r.h, r.k, r.l}}); + r.F = it == amplitude.end() ? NAN : static_cast(it->second); + r.sigmaF = 0.03f * r.F + 0.01f; + } + } + Logger logger("RefinementTest"); + RefinementSettings settings; + settings.gpu = false; + settings.nthreads = 8; + settings.weight_ladder = {1.0 / 128, 1.0 / 16}; + settings.macro_cycles = 1; + settings.max_steps = 1; + settings.final_ladder = {1.0}; + settings.riding_hydrogens = false; + settings.solvent = false; + const RefinementResult ahead = RefineModel(start, merged, truth.cell, sg, 1.0, logger, settings); + RefinementSettings one = settings; + one.max_trials_at_once = 1; + const RefinementResult needed = RefineModel(start, merged, truth.cell, sg, 1.0, logger, one); + REQUIRE(ahead.ok); + REQUIRE(needed.ok); + CHECK(ahead.ncs_copies == 2); + CHECK(ahead.trials.size() == needed.trials.size()); + CHECK(ahead.adp_trials.size() == needed.adp_trials.size()); + CHECK(ahead.adp_model == needed.adp_model); + CHECK(ahead.r_free == needed.r_free); + std::vector x = atoms_of(ahead.model->models[0]), y = atoms_of(needed.model->models[0]); + REQUIRE(x.size() == y.size()); + size_t same = 0; + for (size_t i = 0; i < x.size(); i++) + same += x[i]->pos.x == y[i]->pos.x && x[i]->pos.y == y[i]->pos.y && x[i]->pos.z == y[i]->pos.z && + x[i]->b_iso == y[i]->b_iso; + CHECK(same == x.size()); +} + // A K+ on the 4-fold axis of P4 at occupancy 0.25 (its four coinciding copies one ion), 0.2 A off the axis in the // model as given, with a water 2.8 A from it: the topology puts the ion on the axis, its contact with the water is // listed once (not once per copy) and none with itself; refinement keeps it on the axis at the occupancy given,