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,