v1.0.0-rc.178 (#88)

* Rugnux: Processes serial crystallography data from XFELs - SACLA (MPCCD), PAL-XFEL and European XFEL (CXI), LCLS (XTC) and SwissFEL Alvra (JUNGFRAU with a CrystFEL module geometry) - with several files or Jungfraujoch master files merged in one run (`rugnux *_master.h5`).
* Rugnux: Stills processing takes a known unit cell and space group as a starting point, finds the beam centre (and, with a cell, the detector distance) from the images, detects the saturation level in other facilities' files and refines the placement of individual detector panels.
* Rugnux: The stills report gives R_split and CC1/2 from frame half-sets, comparable with CrystFEL.
* jfjoch_viewer: Opens the same XFEL file formats.
* Validation: A stills battery checks Rugnux against deposited structures from SACLA, PAL-XFEL, LCLS, European XFEL and SwissFEL.
* Rugnux: More reliable convergence of rotation scaling, and live molecular replacement starts as soon as the model is uploaded.
* jfjoch_broker: Initialize clears the previous dataset's bearer tokens, image buffer and live Rugnux status, and a slow CUDA JIT cache is reported at start.

Reviewed-on: #88
Co-authored-by: Filip Leonarski <filip.leonarski@psi.ch>

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_t>(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<std::vector<gemmi::Fractional>(const std::vector<gemmi::Mat33> &)> translate;
+        if (c.copies == 1)
+            translate = [&](const std::vector<gemmi::Mat33> &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<TranslationPeak> tf = translation_function(
+                        MakeTranslationSetup(refl, gemmi::UnitCell(merges[merge].cell), c.space_group), turned,
+                        nullptr, nullptr);
+                std::vector<gemmi::Fractional> 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<gemmi::Structure> placed;   // the one the choice falls on
+        std::shared_ptr<gemmi::Structure> other;    // the one it is compared with; none for its null
+        std::shared_future<double> null_sigma, paired_z;
+        std::shared_future<PlacementScore> carried;
+    };
+    std::vector<Ahead> 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<MRCandidate> &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<double> *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<gemmi::Mat33> 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<int>(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<MRCandidate> 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<std::future<GroupResult>> 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<std::vector<gemmi::Fractional>(const std::vector<gemmi::Mat33> &)> translate;
-            if (c.copies == 1)
-                translate = [&](const std::vector<gemmi::Mat33> &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<TranslationPeak> tf = translation_function(
-                            MakeTranslationSetup(refl, gemmi::UnitCell(merges[merge].cell), c.space_group), turned,
-                            nullptr, nullptr);
-                    std::vector<gemmi::Fractional> 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<double> nulls(result.candidates.size(), NAN);
@@ -3317,9 +3398,20 @@ MolecularReplacementResult search_model(const SearchModel &model, const std::vec
                         if (static_cast<int>(i) != best && result.candidates[i].merge == merge
                             && result.candidates[i].verdict == "PLACED")
                             todo.push_back(static_cast<int>(i));
-                std::vector<std::future<double>> running;
-                for (int i : todo)
-                    running.push_back(std::async(std::launch::async, run_null, i));
+                std::vector<std::shared_future<double>> 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<MRCandidate> 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);
@@ -3624,8 +3713,10 @@ std::vector<TranslationPeak> SplitTranslationFunction(const std::vector<Molecula
 }
 #endif
 
-MolecularReplacementResult RunMolecularReplacement(const RugnuxModel &input, const std::vector<MRMerge> &merges,
-                                                   Logger &logger, size_t nthreads, size_t gpu_resident_bytes) {
+namespace {
+
+MolecularReplacementResult run_molecular_replacement(const RugnuxModel &input, const std::vector<MRMerge> &merges,
+                                                     Logger &logger, size_t nthreads, size_t gpu_resident_bytes) {
     const auto t0 = std::chrono::steady_clock::now();
     // What the search may hold on the card: the budget model validation keeps to. What the stream-ordered pool
     // holds with no buffer in it - freed by the steps before - goes back to the driver first, as before a
@@ -3733,6 +3824,135 @@ MolecularReplacementResult RunMolecularReplacement(const RugnuxModel &input, con
     return result;
 }
 
+} // namespace
+
+MolecularReplacementResult RunMolecularReplacement(const RugnuxModel &model, const std::vector<MRMerge> &merges,
+                                                   Logger &logger, size_t nthreads, size_t gpu_resident_bytes) {
+    WaitForModelWarmUp();
+    return run_molecular_replacement(model, merges, logger, nthreads, gpu_resident_bytes);
+}
+
+std::optional<MolecularReplacementResult> CarryMolecularReplacement(const MolecularReplacementResult &early,
+                                                                    const MRMerge &searched, const MRMerge &merge,
+                                                                    Logger &logger, size_t nthreads,
+                                                                    size_t gpu_resident_bytes) {
+    if (!early.ok || early.chosen < 0 || early.candidates[early.chosen].merge != 0) {
+        logger.Info("Molecular replacement: the search on the {} placed nothing ({}) - searched again",
+                    searched.name, early.ok ? early.decision : early.failure_reason);
+        return std::nullopt;
+    }
+    if (searched.space_group.xhm() != merge.space_group.xhm() || !merge.cell.is_close(searched.cell, 0.01, 1.0)) {
+        logger.Info("Molecular replacement: the search on the {} was made in {} ({:.2f} {:.2f} {:.2f}), this merge is in "
+                    "{} ({:.2f} {:.2f} {:.2f}) - searched again", searched.name, searched.space_group.xhm(),
+                    searched.cell.a, searched.cell.b, searched.cell.c, merge.space_group.xhm(), merge.cell.a,
+                    merge.cell.b, merge.cell.c);
+        return std::nullopt;
+    }
+    const auto t0 = std::chrono::steady_clock::now();
+    const MRCandidate &c = early.candidates[early.chosen];
+    const gemmi::UnitCell cell(merge.cell);
+    const gemmi::SpaceGroup &sg = c.space_group;   // the group it was placed in, one of the merge's (SameRotationGroups)
+    double data_d_min = 1e9;
+    for (const MergedReflection &r : *merge.merged)
+        if (r.d > 0)
+            data_d_min = std::min<double>(data_d_min, r.d);
+    const double d_min = std::max(SCORE_D_MIN, data_d_min);
+    // Every copy at its fractional position in this cell.
+    std::vector<gemmi::Structure> body = SplitCopies(*c.placed, c.copies);
+    for (gemmi::Structure &b : body) {
+        std::vector<gemmi::Position> pos = ModelPositions(b.models[0]);
+        for (gemmi::Position &x : pos)
+            x = cell.orthogonalize(b.cell.fractionalize(x));
+        SetModelPositions(b.models[0], pos);
+        b.cell = cell;
+        b.spacegroup_hm = sg.xhm();
+        b.setup_cell_images();
+    }
+    RigidBodyGPUPool *pool = nullptr;
+    ScoreEngines scorers;
+#ifdef JFJOCH_USE_CUDA
+    // One rigid-body engine and the score engine, on this card, as a search's own (size_merge_gpu).
+    std::unique_ptr<RigidBodyGPUPool> engines;
+    if (get_gpu_count() > 0) {
+        const int device = RigidBodyGPUEngine::CurrentDevice();
+        const size_t budget = ModelGPUBudget(get_gpu_total_memory(device), gpu_resident_bytes);
+        size_t observations = 0;
+        for (const MergedReflection &r : *merge.merged)
+            if (!std::isnan(r.F) && r.d >= RigidBodyLadder(d_min).back())
+                ++observations;
+        Logger quiet = Logger::Buffered();
+        const gemmi::Structure all = AssembleCopies(body);
+        engines = RigidBodyGPUPool::Plan(body[0].models[0], cell, sg, d_min, observations, quiet);
+        if (engines && engines->Cover(all.models[0], sg, quiet) && engines->EngineBytes() <= budget / 2) {
+            engines->Reserve({1});
+            pool = engines.get();
+        }
+        scorers.SetRoom({{device, budget - (pool != nullptr ? engines->EngineBytes() : 0)}}, device, true);
+    }
+#endif
+    for (size_t k = 0; k < body.size(); k++) {
+        std::vector<gemmi::Structure> rest;
+        for (size_t o = 0; o < body.size(); o++)
+            if (o != k)
+                rest.push_back(body[o]);
+        if (rest.empty()) {
+            refine_rigid_body(body[k], sg, *merge.merged, d_min, pool, nthreads);
+        } else {
+            const gemmi::Structure fixed = AssembleCopies(rest);
+            refine_rigid_body(body[k], sg, *merge.merged, d_min, pool, nthreads, nullptr, &fixed.models[0]);
+        }
+    }
+    auto placed = std::make_shared<gemmi::Structure>(AssembleCopies(body));
+    const PlacementScore score = score_placement(*placed, sg, *merge.merged, d_min, nthreads, &scorers);
+    const bool holds = score.r_free <= c.r_free + R_FREE_MARGIN;
+    logger.Info("Molecular replacement: the search on the {} ({}): its placement in {} ({} copies, R-free {:.3f} there) refined against "
+                "this merge: R-work {:.3f} R-free {:.3f} ({:.2f} s) - {}", searched.name,
+                early.decision, sg.xhm(), c.copies, c.r_free, score.r_work, score.r_free,
+                seconds_since(t0), holds ? "it holds, not searched again" : "it does not hold, searched again");
+    if (!holds)
+        return std::nullopt;
+    MolecularReplacementResult out = early;
+    MRCandidate &chosen = out.candidates[out.chosen];
+    out.r_free_searched = chosen.r_free;
+    chosen.placed = placed;
+    chosen.r_work = score.r_work;
+    chosen.r_free = score.r_free;
+    chosen.k_sol = score.k_sol;
+    chosen.b_sol = score.b_sol;
+    out.d_min_score = d_min;
+    return out;
+}
+
+MolecularReplacementWarmUp WarmUpMolecularReplacement(const RugnuxModel &model, size_t nthreads) {
+    MolecularReplacementWarmUp out;
+    const SearchModel search = PrepareSearchModel(model);
+    if (!search.failure_reason.empty()) {
+        out.failure_reason = search.failure_reason;
+        return out;
+    }
+    const double edge = 4.0 * search.radius_A;
+    out.cell = UnitCell{static_cast<float>(edge), static_cast<float>(edge + 2.0), static_cast<float>(edge + 4.0),
+                        90.0f, 100.0f, 90.0f};
+    out.space_group = *gemmi::find_spacegroup_by_name("P 1 21 1");
+    RugnuxModel crystal;
+    crystal.name = model.name;
+    crystal.structure = PlaceSearchModel(search, gemmi::Mat33(), gemmi::Fractional(0.1, 0.2, 0.3),
+                                         gemmi::UnitCell(out.cell), out.space_group);
+    Logger quiet = Logger::Buffered();
+    out.merged = ModelReferenceIntensities(crystal, std::nullopt, nullptr, 2.5, quiet);
+    for (size_t i = 0; i < out.merged.size(); i++) {
+        out.merged[i].F = std::sqrt(out.merged[i].I);
+        out.merged[i].rfree_flag = i % 20 == 0;
+    }
+    const MolecularReplacementResult mr = run_molecular_replacement(
+            model, {{"MERGE", &out.merged, out.cell, out.space_group}}, quiet, nthreads, 0);
+    if (mr.chosen >= 0)
+        out.placed = mr.candidates[mr.chosen].placed;
+    else
+        out.failure_reason = mr.ok ? mr.decision : mr.failure_reason;
+    return out;
+}
+
 void CheckEnantiomorphHand(MolecularReplacementResult &result, const std::vector<MergedReflection> &merged,
                            const gemmi::SpaceGroup &data_space_group, Logger &logger, size_t nthreads,
                            size_t gpu_resident_bytes) {