diff --git a/image_analysis/structure_refinement/MolecularReplacement.cpp b/image_analysis/structure_refinement/MolecularReplacement.cpp index 5fc50422f..63b7689a0 100644 --- a/image_analysis/structure_refinement/MolecularReplacement.cpp +++ b/image_analysis/structure_refinement/MolecularReplacement.cpp @@ -2500,7 +2500,10 @@ MolecularReplacementResult search_model(const SearchModel &model, const std::vec need.other_budgets = other_budgets; need.transform = transform_bytes; need.rigid_body = size.rigid_body; - need.rigid_body_engines = size.rigid_body_engines; + // Engines for the fits the choice runs at once besides: a null's replicates, beside them a carry and a group + // searched again (search_merge's work ahead of the choice). The layout of the transform was decided on the + // search's own fits. + need.rigid_body_engines = std::max(size.rigid_body_engines, NULL_REPLICATES + 1 + RESCORED_PER_GROUP); need.structure_factors = size.structure_factors; const MolecularReplacementGPUPlan plan = PlanMolecularReplacementGPU(need); const int cards = get_gpu_count(); @@ -2599,10 +2602,101 @@ MolecularReplacementResult search_model(const SearchModel &model, const std::vec return TranslationFunction(ft, setup, rotations, known, nthreads); }; + // The null of a candidate (null_sigma) placed in the orientation `orientation` before the rigid body. + const auto null_of = [&](const MRCandidate &c, const gemmi::Mat33 &orientation) { + const int merge = c.merge; + // One copy: its replicates placed by the translation function, as it was. + std::function(const std::vector &)> translate; + if (c.copies == 1) + translate = [&](const std::vector &turned) { + const SearchReflections refl = MakeSearchReflections(*merges[merge].merged, + gemmi::UnitCell(merges[merge].cell), + merges[merge].space_group, + result.d_min_search, result.d_max_search); + const std::vector tf = translation_function( + MakeTranslationSetup(refl, gemmi::UnitCell(merges[merge].cell), c.space_group), turned, + nullptr, nullptr); + std::vector out; + for (const TranslationPeak &p : tf) + out.push_back(p.t); + return out; + }; + return null_sigma(*c.placed, c.copies, orientation, c.space_group, *merges[merge].merged, + result.d_min_score, rigid_body_pool(merge), nthreads, &scorers, translate); + }; + // The comparisons of a candidate `b` the choice falls on with another of its merge, `c`, as choose_in() makes + // them: b's copies carried into c's group, where c placed fewer beside tNCS, and the paired R, where the two are + // within the R-free margin with as many copies. + const auto needs_carry = [&](const MRCandidate &c, const MRCandidate &b) { + return result.tncs.present && c.copies < b.copies + && (TncsRelatesGroups(c.space_group, b.space_group, result.tncs.vector) + || enantiomorphs(c.space_group, b.space_group)); + }; + const auto needs_paired = [&](const MRCandidate &c, const MRCandidate &b) { + return !(c.r_free - b.r_free >= R_FREE_MARGIN) && c.copies == b.copies; + }; + const auto carry_of = [&](const MRCandidate &c, const MRCandidate &b) { + return carry_copies(*b.placed, b.copies, c.space_group, *merges[b.merge].merged, result.d_min_score, + rigid_body_pool(b.merge), nthreads, &scorers); + }; + const auto paired_of = [&](const MRCandidate &c, const MRCandidate &b) { + return paired_r_z(*c.placed, c.space_group, *b.placed, b.space_group, *merges[b.merge].merged, + result.d_min_score, nthreads, &scorers); + }; + // The work of the choice started before the choice comes to it, beside what decides where it falls: the null of + // the placement it is expected to fall on, and that placement's comparisons with the others. Each is taken where + // the choice comes to the same placements; a null or a comparison is the same number whenever it is computed, so + // this changes when they run, never what they give. All are waited for before a merge's rigid-body engines are + // given back. + struct Ahead { + std::shared_ptr placed; // the one the choice falls on + std::shared_ptr other; // the one it is compared with; none for its null + std::shared_future null_sigma, paired_z; + std::shared_future carried; + }; + std::vector ahead; + const auto find_ahead = [&](const MRCandidate &b, const MRCandidate *c) -> const Ahead * { + for (const Ahead &a : ahead) + if (a.placed == b.placed && a.other == (c != nullptr ? c->placed : nullptr)) + return &a; + return nullptr; + }; + // `b`'s null (with its orientation before the rigid body), where it is not already started, and its comparisons + // with `others` - those of its merge that pack. + const auto start_ahead = [&](const MRCandidate &b, const gemmi::Mat33 *orientation, + const std::vector &others) { + if (orientation != nullptr && find_ahead(b, nullptr) == nullptr) + ahead.push_back({b.placed, nullptr, + std::async(std::launch::async, [&, b, o = *orientation] { return null_of(b, o); }).share(), + {}, {}}); + for (const MRCandidate &c : others) { + if (c.placed == b.placed || c.clash_fraction > CLASH_MAX || find_ahead(b, &c) != nullptr) + continue; + Ahead a{b.placed, c.placed, {}, {}, {}}; + if (needs_carry(c, b)) + a.carried = std::async(std::launch::async, [&, b, c] { return carry_of(c, b); }).share(); + if (needs_paired(c, b)) + a.paired_z = std::async(std::launch::async, [&, b, c] { return paired_of(c, b); }).share(); + if (a.carried.valid() || a.paired_z.valid()) + ahead.push_back(std::move(a)); + } + }; + const auto wait_ahead = [&] { + for (Ahead &a : ahead) + for (const std::shared_future *f : {&a.null_sigma, &a.paired_z}) + if (f->valid()) + f->wait(); + for (Ahead &a : ahead) + if (a.carried.valid()) + a.carried.wait(); + ahead.clear(); + }; + std::vector orientations; // of each candidate's placement, before the rigid body // One merge: the rotation function over its point group, then for each of its groups the // translation function and the best placements through packing, the rigid body and R-free. const auto search_merge = [&](size_t m) { + wait_ahead(); // before plan_gpu() gives the last merge's engines back const MRMerge &merge = merges[m]; const gemmi::UnitCell cell(merge.cell); const double d_min = result.d_min_score; @@ -3158,6 +3252,15 @@ MolecularReplacementResult search_model(const SearchModel &model, const std::vec for (size_t g = 0; g < found.size(); g++) if (found[g].cand.verdict == "PLACED" && (lead < 0 || found[g].cand.r_free < found[lead].cand.r_free)) lead = static_cast(g); + // The choice falls first on the placed candidate of lowest R-free (choose_in), which is this one unless a + // search again below finds a better: its null and its comparisons with the others are started now, beside + // those searches. + if (lead >= 0) { + std::vector others; + for (const GroupResult &gr : found) + others.push_back(gr.cand); + start_ahead(found[lead].cand, &found[lead].orientation, others); + } if (lead >= 0 && result.tncs.present) { std::vector> again(found.size()); for (size_t g = 0; g < found.size(); g++) @@ -3192,28 +3295,6 @@ MolecularReplacementResult search_model(const SearchModel &model, const std::vec // The null is built for the candidate the choice would fall on, and only for it: a placed candidate // that does not beat its own null is set aside, and the choice falls on the next. const auto choose_in = [&](int merge) -> int { - // The null of candidate i. - const auto run_null = [&](int i) { - const MRCandidate &c = result.candidates[i]; - // One copy: its replicates placed by the translation function, as it was. - std::function(const std::vector &)> translate; - if (c.copies == 1) - translate = [&](const std::vector &turned) { - const SearchReflections refl = MakeSearchReflections(*merges[merge].merged, - gemmi::UnitCell(merges[merge].cell), - merges[merge].space_group, - result.d_min_search, result.d_max_search); - const std::vector tf = translation_function( - MakeTranslationSetup(refl, gemmi::UnitCell(merges[merge].cell), c.space_group), turned, - nullptr, nullptr); - std::vector out; - for (const TranslationPeak &p : tf) - out.push_back(p.t); - return out; - }; - return null_sigma(*c.placed, c.copies, orientations[i], c.space_group, *merges[merge].merged, - result.d_min_score, rigid_body_pool(merge), nthreads, &scorers, translate); - }; // The first candidate's null alone; where it fails, the nulls of all the other placed candidates side // by side, each taken where the choice comes to it. std::vector nulls(result.candidates.size(), NAN); @@ -3317,9 +3398,20 @@ MolecularReplacementResult search_model(const SearchModel &model, const std::vec if (static_cast(i) != best && result.candidates[i].merge == merge && result.candidates[i].verdict == "PLACED") todo.push_back(static_cast(i)); - std::vector> running; - for (int i : todo) - running.push_back(std::async(std::launch::async, run_null, i)); + std::vector> running; + for (int i : todo) { + const MRCandidate &ci = result.candidates[i]; + const Ahead *a = find_ahead(ci, nullptr); + running.push_back(a != nullptr ? a->null_sigma + : std::async(std::launch::async, null_of, std::cref(ci), + std::cref(orientations[i])).share()); + } + // Its comparisons, beside its null, where they were not started before. + std::vector others; + for (const MRCandidate &o : result.candidates) + if (o.merge == merge) + others.push_back(o); + start_ahead(c, nullptr, others); for (size_t k = 0; k < todo.size(); k++) nulls[todo[k]] = running[k].get(); result.t_null += seconds_since(t_null); @@ -3337,16 +3429,14 @@ MolecularReplacementResult search_model(const SearchModel &model, const std::vec // With tNCS one copy alone fits no better than at random, so a group whose search stopped at // fewer copies than the chosen one's has not been shown unable to hold them - where the tNCS // vector makes up the groups' screw translations, or where they are each other's enantiomorph. - if (result.tncs.present && c.copies < result.candidates[best].copies - && (TncsRelatesGroups(c.space_group, result.candidates[best].space_group, result.tncs.vector) - || enantiomorphs(c.space_group, result.candidates[best].space_group))) { + const Ahead *a = find_ahead(result.candidates[best], &c); + if (needs_carry(c, result.candidates[best])) { // Unless the chosen copies, carried into this group as they are and refined, build a crystal // that fits clearly worse: where the tNCS vector does make up the difference between the groups, // they build the same one. const MRCandidate &b = result.candidates[best]; - const PlacementScore carried = carry_copies(*b.placed, b.copies, c.space_group, *merges[merge].merged, - result.d_min_score, rigid_body_pool(merge), nthreads, - &scorers); + const PlacementScore carried = a != nullptr && a->carried.valid() ? a->carried.get() + : carry_of(c, b); logger.Info("Molecular replacement: {}'s {} copies carried into {}: R-free {:.3f} against {:.3f}", b.space_group.xhm(), b.copies, c.space_group.xhm(), carried.r_free, b.r_free); if (!(carried.r_free - b.r_free >= R_FREE_MARGIN) || !(b.r_free < CARRY_R_FREE_MAX)) { @@ -3361,8 +3451,7 @@ MolecularReplacementResult search_model(const SearchModel &model, const std::vec // the chosen placement fits every reflection better by PAIRED_Z_MIN standard errors. const MRCandidate &b = result.candidates[best]; if (c.copies == b.copies) { - const double z = paired_r_z(*c.placed, c.space_group, *b.placed, b.space_group, - *merges[merge].merged, result.d_min_score, nthreads, &scorers); + const double z = a != nullptr && a->paired_z.valid() ? a->paired_z.get() : paired_of(c, b); logger.Info("Molecular replacement: {} and {} within {:.2f} in R-free, {} copies each: {} fits every " "reflection better by {:.1f} standard errors", b.space_group.xhm(), c.space_group.xhm(), R_FREE_MARGIN, b.copies, b.space_group.xhm(), z);