diff --git a/image_analysis/structure_refinement/ModelStructureFactorsGPU.h b/image_analysis/structure_refinement/ModelStructureFactorsGPU.h index 499d4e8aa..12349b56a 100644 --- a/image_analysis/structure_refinement/ModelStructureFactorsGPU.h +++ b/image_analysis/structure_refinement/ModelStructureFactorsGPU.h @@ -58,6 +58,8 @@ public: // work area and the reflections. The largest map Map() can be given shares the same buffers. size_t DeviceBytes() const { return device_bytes_; } double DMin() const { return d_min_; } + // The reflections Compute() gives F_calc and F_mask at. + size_t Reflections() const { return rows_.size(); } int Device() const { return device_; } // The device's total memory, which DeviceBytes() is to be judged against. size_t DeviceTotalMemory() const { return device_total_; } diff --git a/image_analysis/structure_refinement/MolecularReplacement.cpp b/image_analysis/structure_refinement/MolecularReplacement.cpp index 23e1b4600..9c11693a9 100644 --- a/image_analysis/structure_refinement/MolecularReplacement.cpp +++ b/image_analysis/structure_refinement/MolecularReplacement.cpp @@ -40,6 +40,7 @@ #include "RigidBodyGPU.h" #include "RigidBodyGPUEngine.h" #include "ModelStructureFactorsGPU.h" +#include "ModelScaleGPU.h" #include "../../common/CUDAWrapper.h" #include "../../common/JFJochException.h" #endif @@ -98,6 +99,9 @@ constexpr double NULL_D_MIN = 3.5; // a new draw lies t(n-1) sqrt(1 + 1/n) sample sds above their mean at that probability, 4.28 x 1.054 for // n = 9 - the gate model validation's fit is read at (ModelValidation.cpp, MODEL_FIT_SIGMA). constexpr double NULL_SIGMA_MIN = 4.51; +// A model whose expected log-likelihood gain (ExpectedLogLikelihoodGain) is below this is not searched for. The +// lowest of a model the search did place, over 193 survey sets: 26; every PAE domain below 24 placed nowhere. +constexpr double ELLG_MIN = 10.0; double seconds_since(std::chrono::steady_clock::time_point t0) { return std::chrono::duration(std::chrono::steady_clock::now() - t0).count(); @@ -236,13 +240,23 @@ PlacementScore r_factors(const gemmi::AsuData> &fmodel, cons } // Model validation's structure-factor engines on the GPU, one per group, copy count and resolution, made -// where the first placement of that kind is scored - or none, where there is no GPU, the cell is too small -// for its gridding or one would take over a quarter of the card. Decided on the card's total memory, so -// the same input always takes the same path. +// where the first placement of that kind is scored, each with the scale fit on the GPU beside it - or none, +// where there is no GPU, the cell is too small for its gridding or the two would take over a quarter of the +// card. Decided on the card's total memory, so the same input always takes the same path. #ifdef JFJOCH_USE_CUDA class ScoreEngines { public: - ModelStructureFactorsGPU *get(const gemmi::Structure &st, const gemmi::SpaceGroup &sg, double d_min) { + struct Engine { + std::unique_ptr sf; + // The scale fit of the reflections the engine's F_calc are matched to, on its card: FitModelScale and + // gemmi's fit at a held solvent as ModelScaleGPU makes them - the same fits, equal to rounding. + std::mutex scale_m; + std::unique_ptr stream; + std::unique_ptr scale; + CudaDevicePtr> fcalc, fmask; + }; + + Engine *get(const gemmi::Structure &st, const gemmi::SpaceGroup &sg, double d_min) { if (get_gpu_count() == 0) return nullptr; size_t atoms = 0; @@ -256,25 +270,98 @@ public: try { auto sf = std::make_unique(RigidBodyGPUEngine::CurrentDevice(), st.models[0], st.cell, sg, d_min); - if (sf->Supported() && sf->DeviceBytes() <= sf->DeviceTotalMemory() / 4) { + const size_t points = sf->Reflections(); + const size_t bytes = sf->DeviceBytes() + ModelScaleGPU::DeviceBytes(points) + + 2 * points * sizeof(std::complex); + if (sf->Supported() && bytes <= sf->DeviceTotalMemory() / 4) { sf->Reserve(); - e.sf = std::move(sf); + const int device = RigidBodyGPUEngine::CurrentDevice(); + set_gpu(sf->Device()); + e.engine = std::make_unique(); + e.engine->stream = std::make_unique(); + e.engine->scale = std::make_unique(*e.engine->stream, points); + e.engine->fcalc = CudaDevicePtr>(points); + e.engine->fmask = CudaDevicePtr>(points); + e.engine->sf = std::move(sf); + set_gpu(device); } } catch (const JFJochException &) { + e.engine.reset(); cuda_clear_error(); } } - return e.sf.get(); + return e.engine.get(); } private: struct Entry { bool tried = false; - std::unique_ptr sf; + std::unique_ptr engine; }; std::map, Entry> engines_; std::mutex m_; }; + +// The scale of `scaling`'s points fitted on the engine's card: FitModelScale's grid over the solvent box, or +// with `solvent`, the overall scale and anisotropic B at its k_sol and b_sol held. +void fit_scale_gpu(ScoreEngines::Engine &e, gemmi::Scaling &scaling, const PlacementScore *solvent, + double b_hi, size_t nthreads) { + const size_t n = scaling.points.size(); + std::vector> hkl(n); + std::vector stol2(n); + std::vector fobs(n), sigma(n); + std::vector> fc(n), fm(n); + for (size_t i = 0; i < n; i++) { + const auto &p = scaling.points[i]; + hkl[i] = {p.hkl[0], p.hkl[1], p.hkl[2]}; + stol2[i] = p.stol2; + fobs[i] = p.fobs; + sigma[i] = p.sigma; + fc[i] = p.fcmol; + fm[i] = p.fmask; + } + std::vector> constraints; + for (const gemmi::Vec6 &c : scaling.constraint_matrix) + constraints.push_back(c); + double frac[9]; + for (int i = 0; i < 3; i++) + for (int j = 0; j < 3; j++) + frac[3 * i + j] = scaling.cell.frac.mat[i][j]; + + std::unique_lock lock(e.scale_m); + const int device = RigidBodyGPUEngine::CurrentDevice(); + set_gpu(e.sf->Device()); + e.scale->SetPoints(hkl, stol2, fobs, sigma, constraints, frac); + if (cudaMemcpyAsync(e.fcalc.get(), fc.data(), n * sizeof(std::complex), cudaMemcpyHostToDevice, + *e.stream) != cudaSuccess + || cudaMemcpyAsync(e.fmask.get(), fm.data(), n * sizeof(std::complex), cudaMemcpyHostToDevice, + *e.stream) != cudaSuccess) + throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "Model prior: the scale fit's upload failed"); + const auto *d_fc = reinterpret_cast(e.fcalc.get()); + const auto *d_fm = reinterpret_cast(e.fmask.get()); + ModelScaleParams p; + if (solvent != nullptr) { + scaling.k_sol = solvent->k_sol; + scaling.b_sol = solvent->b_sol; + p = e.scale->Fit(d_fc, d_fm, solvent->k_sol, solvent->b_sol); + } else { + try { + const ModelSolventFit fit = e.scale->FitSolvent(d_fc, d_fm, b_hi); + scaling.k_sol = fit.k_sol; + scaling.b_sol = fit.b_sol; + p = fit.scale; + } catch (const ModelScaleGPUTooFewReflections &) { + // FitModelScale walks such a grid in order, each point from where the last ended. + set_gpu(device); + lock.unlock(); + FitModelScale(scaling, {.b_hi = b_hi}, nthreads); + return; + } + } + set_gpu(device); + scaling.k_overall = p.k_overall; + scaling.b_star = {p.b_star[0], p.b_star[1], p.b_star[2], p.b_star[3], p.b_star[4], p.b_star[5]}; +} #else class ScoreEngines {}; #endif @@ -339,8 +426,9 @@ PlacementScore score_placement(const gemmi::Structure &st, const gemmi::SpaceGro const gemmi::Op *law = nullptr) { gemmi::AsuData> fmask, fmodel; #ifdef JFJOCH_USE_CUDA - if (ModelStructureFactorsGPU *gpu = engines != nullptr ? engines->get(st, sg, d_min) : nullptr) - gpu->Compute(st.models[0], fmodel, fmask); + ScoreEngines::Engine *gpu = engines != nullptr ? engines->get(st, sg, d_min) : nullptr; + if (gpu) + gpu->sf->Compute(st.models[0], fmodel, fmask); else #endif model_factors(st, d_min, fmodel, fmask, nthreads); @@ -349,6 +437,11 @@ PlacementScore score_placement(const gemmi::Structure &st, const gemmi::SpaceGro gemmi::Scaling scaling(st.cell, &sg); scaling.use_solvent = true; scaling.prepare_points(fmodel, fobs, &fmask); +#ifdef JFJOCH_USE_CUDA + if (gpu) + fit_scale_gpu(*gpu, scaling, solvent, SolventBMax(st.models[0]), nthreads); + else +#endif if (solvent != nullptr) { // The bulk solvent of an earlier fit of the same crystal, held, as the rigid body holds it. scaling.k_sol = solvent->k_sol; @@ -1067,19 +1160,27 @@ TranslationSetup MakeTranslationSetup(const SearchReflections &refl, const gemmi std::vector> KnownStructureFactors(const MolecularTransform &ft, const TranslationSetup &setup, const gemmi::UnitCell &cell, const std::vector &rotations, - const std::vector &positions) { + const std::vector &positions, + size_t nthreads) { const size_t terms = setup.Terms(), nops = setup.ops; - std::vector> out(terms, 0.0); + std::vector rt; + std::vector x; for (size_t c = 0; c < rotations.size(); c++) { - const gemmi::Mat33 rt = rotations[c].transpose(); - const gemmi::Position x = cell.orthogonalize(positions[c]); - for (size_t t = 0; t < terms; t++) - for (size_t j = 0; j < nops; j++) { - const gemmi::Vec3 &sj = setup.s[t * nops + j]; - // (h R_j).t = s_j.x for the Cartesian s_j and x - out[t] += ft.At(rt.multiply(sj)) * setup.phase[t * nops + j] * std::polar(1.0, 2.0 * PI * sj.dot(x)); - } + rt.push_back(rotations[c].transpose()); + x.push_back(cell.orthogonalize(positions[c])); } + std::vector> out(terms, 0.0); + ParallelChunks(static_cast(terms), ThreadsForWork(terms * nops * rotations.size(), nthreads), + [&](int lo, int hi) { + for (int t = lo; t < hi; t++) + for (size_t c = 0; c < rotations.size(); c++) + for (size_t j = 0; j < nops; j++) { + const gemmi::Vec3 &sj = setup.s[t * nops + j]; + // (h R_j).t = s_j.x for the Cartesian s_j and x + out[t] += ft.At(rt[c].multiply(sj)) * setup.phase[t * nops + j] + * std::polar(1.0, 2.0 * PI * sj.dot(x[c])); + } + }); return out; } @@ -1218,8 +1319,10 @@ std::vector> CopyFactors(const SearchModel &model, const Li std::vector> g(setup.s.size()); if (ft != nullptr && d_min >= ft->DMin()) { const gemmi::Mat33 rt = rotation.transpose(); - for (size_t e = 0; e < g.size(); e++) - g[e] = ft->At(rt.multiply(setup.s[e])) * setup.phase[e]; + ParallelChunks(static_cast(g.size()), ThreadsForWork(g.size(), nthreads), [&](int lo, int hi) { + for (int e = lo; e < hi; e++) + g[e] = ft->At(rt.multiply(setup.s[e])) * setup.phase[e]; + }); return g; } gemmi::Structure st = model.st; @@ -1832,6 +1935,28 @@ std::vector SameRotationGroups(const gemmi::SpaceGroup &sg) { namespace { +// The expected log-likelihood gain of one copy of the model on the reflections of `merge` the search is made +// on: sum sigma_A^4 / 2, with sigma_A^2 the part of the scattering the copy is - one over the copy number the +// cell most probably holds, by Matthews - times Luzzati's D^2 for the coordinate error the pLDDTs predict. +// Following McCoy et al. (2017) Proc. Natl. Acad. Sci. USA 114, 3637-3641 +double ExpectedLogLikelihoodGain(const SearchModel &model, const MRMerge &merge, double d_min, double d_max) { + const gemmi::UnitCell cell(merge.cell); + const SearchReflections refl = MakeSearchReflections(*merge.merged, cell, merge.space_group, d_min, d_max); + const double asu_volume = cell.volume / static_cast(merge.space_group.operations().order()); + const std::vector pm = MatthewsProbabilities(asu_volume, model.mass_da); + size_t likely = 0; + for (size_t n = 0; n < pm.size(); n++) + if (pm[n] > pm[likely]) + likely = n; + const double fraction = 1.0 / static_cast(likely + 1); + double ellg = 0; + for (const gemmi::Miller &h : refl.hkl) { + const double d = std::exp(-2.0 * PI * PI * model.error_A * model.error_A * cell.calculate_1_d2(h) / 3.0); + ellg += fraction * fraction * d * d * d * d / 2.0; + } + return ellg; +} + // The search and the decision for one search model. MolecularReplacementResult search_model(const SearchModel &model, const std::vector &merges, Logger &logger, size_t nthreads) { @@ -1854,6 +1979,15 @@ MolecularReplacementResult search_model(const SearchModel &model, const std::vec result.d_min_score = std::max(SCORE_D_MIN, data_d_min); logger.Info("Model prior: {}: {} of {} residues kept at pLDDT >= 70 ({:.1f} kDa, radius {:.1f} A)", model.description, model.residues_kept, model.residues_read, model.mass_da / 1000.0, model.radius_A); + // What a placement of the model can gain at best: a model too small a part of the cell, or too far from + // it, for its placement to stand out is not searched for. + const double ellg = ExpectedLogLikelihoodGain(model, merges[0], result.d_min_search, result.d_max_search); + logger.Info("Model prior: {}: expected LLG {:.1f}", model.description, ellg); + if (ellg < ELLG_MIN) { + result.ok = true; + result.decision = fmt::format("NOT_PLACED: an expected LLG of {:.1f}, below {:.0f}", ellg, ELLG_MIN); + return result; + } t0 = std::chrono::steady_clock::now(); std::optional ft_made; @@ -1990,18 +2124,25 @@ MolecularReplacementResult search_model(const SearchModel &model, const std::vec // --- translation, for every group and every peak orientation --- t0 = std::chrono::steady_clock::now(); const std::vector groups = SameRotationGroups(merge.space_group); - std::vector tf_all; - for (const gemmi::SpaceGroup &sg : groups) { - const TranslationSetup setup = MakeTranslationSetup(refl, cell, sg); - std::vector tf(peaks.size()); + // Each group's on a thread of its own. + std::vector>> tf_running; + for (const gemmi::SpaceGroup &sg : groups) + tf_running.push_back(std::async(std::launch::async, [&, sg] { + const TranslationSetup setup = MakeTranslationSetup(refl, cell, sg); + std::vector tf(peaks.size()); #ifdef JFJOCH_USE_CUDA - if (gpu) - tf = gpu->TranslationFunction(setup, peak_rotations); - else + if (gpu) + tf = gpu->TranslationFunction(setup, peak_rotations); + else #endif - ParallelFor(static_cast(peaks.size()), nthreads, [&](int i) { - tf[i] = TranslationFunction(ft, setup, peak_rotations[i]); - }); + ParallelFor(static_cast(peaks.size()), nthreads, [&](int i) { + tf[i] = TranslationFunction(ft, setup, peak_rotations[i]); + }); + return tf; + })); + std::vector tf_all; + for (auto &r : tf_running) { + const std::vector tf = r.get(); tf_all.insert(tf_all.end(), tf.begin(), tf.end()); } result.t_translation += seconds_since(t0); @@ -2062,11 +2203,22 @@ MolecularReplacementResult search_model(const SearchModel &model, const std::vec by_tf.insert(by_tf.begin(), first.begin(), first.end()); } const TranslationPeak *tf = tfs.data(); - const LikelihoodSetup lsetup = MakeLikelihoodSetup(lrefl, cell, sg); - const std::vector> no_known(lsetup.Terms(), 0.0); // The further copies' peaks are ranked by the likelihood on the search's own reflections, read off - // the molecular transform. - const LikelihoodSetup rsetup = MakeLikelihoodSetup(refl, cell, sg); + // the molecular transform. Each made where it is first needed; where the likelihood's reflections are + // the search's own, one serves both. + std::optional lsetup_made, rsetup_made; + const auto lsetup = [&]() -> const LikelihoodSetup & { + if (!lsetup_made) + lsetup_made.emplace(MakeLikelihoodSetup(lrefl, cell, sg)); + return *lsetup_made; + }; + const auto rsetup = [&]() -> const LikelihoodSetup & { + if (d_llg == result.d_min_search) + return lsetup(); + if (!rsetup_made) + rsetup_made.emplace(MakeLikelihoodSetup(refl, cell, sg)); + return *rsetup_made; + }; // --- the first copy: the best few through packing, the rigid body and R-free --- RigidBodyEngines engines; @@ -2088,27 +2240,53 @@ MolecularReplacementResult search_model(const SearchModel &model, const std::vec // search placement by a few degrees and a fraction of an angstrom, and the clash fraction moves // by a percent or two with it - and does not count, so that a group is CLASH only where no // orientation the rotation function offers packs in it. - int refined = 0; - for (size_t k = 0; k < by_tf.size() && refined < (fits ? RESCORED_PER_GROUP : 1); k++) { - const size_t i = by_tf[k]; - auto placed = std::make_shared( - PlaceSearchModel(model, rots[i], tf[i].t, cell, sg)); - double clash = ClashFraction(*placed); + // The placements that go through it are known before any does, so they run through it side by side, + // on an engine each, and are then taken in order. + struct First { + size_t i = 0; + std::shared_ptr placed; + double clash = 1; + bool refine = false; PlacementScore score{NAN, NAN, NAN, NAN}; gemmi::Mat33 moved; - if (fits && clash <= 2 * CLASH_MAX) { - if (engines.get() == nullptr) - engines = make_engines(*placed, sg, *merge.merged, d_min, 1, logger); - score = refine_placement(*placed, sg, *merge.merged, d_min, engines.get(), nthreads, &scorers, - &moved); - clash = ClashFraction(*placed); + }; + std::vector first; + int refined = 0; + for (size_t k = 0; k < by_tf.size() && refined < (fits ? RESCORED_PER_GROUP : 1); k++) { + First f; + f.i = by_tf[k]; + f.placed = std::make_shared(PlaceSearchModel(model, rots[f.i], tf[f.i].t, cell, sg)); + f.clash = ClashFraction(*f.placed); + f.refine = fits && f.clash <= 2 * CLASH_MAX; + if (f.refine || !fits) refined++; - } else if (fits && have) { + first.push_back(std::move(f)); + } + // The further copies' translation function, made while the first copy is refined. + std::future setup2_made; + if (fits && cand.copies_max > 1) + setup2_made = std::async(std::launch::async, [&] { return MakeTranslationSetup(refl, cell, sg, true); }); + std::vector to_refine; + for (First &f : first) + if (f.refine) + to_refine.push_back(&f); + if (!to_refine.empty()) { + engines = make_engines(*to_refine[0]->placed, sg, *merge.merged, d_min, to_refine.size(), logger); + std::vector> running; + for (First *f : to_refine) + running.push_back(std::async(std::launch::async, [&, f] { + f->score = refine_placement(*f->placed, sg, *merge.merged, d_min, engines.get(), nthreads, + &scorers, &f->moved); + f->clash = ClashFraction(*f->placed); + })); + for (auto &r : running) + r.get(); + } + for (const First &f : first) { + if (!f.refine && fits && have) continue; - } else if (!fits) { - refined++; - } - // The refined placement's likelihood, judged against the same orientation elsewhere. + const size_t i = f.i; + const PlacementScore &score = f.score; // The refined placement with the lowest R-free; the best translation-function one where // none was refined. if (!have || (std::isfinite(score.r_free) && !(cand.r_free <= score.r_free))) { @@ -2116,29 +2294,30 @@ MolecularReplacementResult search_model(const SearchModel &model, const std::vec cand.rf_z = rf_sd > 0 ? (rfs[i] - rf_mean) / rf_sd : 0.0; cand.tf_z = tf[i].z; cand.copy_tf_z = {tf[i].z}; - cand.clash_fraction = clash; + cand.clash_fraction = f.clash; cand.r_work = score.r_work; cand.r_free = score.r_free; cand.k_sol = score.k_sol; cand.b_sol = score.b_sol; - cand.placed = placed; + cand.placed = f.placed; cand.copies = 1; orientation = rots[i]; - copies = {*placed}; - copy_rotation = {moved.multiply(rots[i])}; - copy_position = {centroid(*placed)}; + copies = {*f.placed}; + copy_rotation = {f.moved.multiply(rots[i])}; + copy_position = {centroid(*f.placed)}; copy_g = {{}}; copy_rg = {{}}; } glog.Info("Model prior: {:<12} RF Z {:5.1f} TF Z {:5.1f} clash {:4.1f}% R-work {:.3f} R-free {:.3f}", - sg.xhm(), rf_sd > 0 ? (rfs[i] - rf_mean) / rf_sd : 0.0, tf[i].z, 100 * clash, + sg.xhm(), rf_sd > 0 ? (rfs[i] - rf_mean) / rf_sd : 0.0, tf[i].z, 100 * f.clash, score.r_work, score.r_free); } // The likelihood of the placement kept, judged against the same orientation elsewhere - only where // its translation function's Z-score does not already stand out. if (have && std::isfinite(cand.r_free) && cand.tf_z < TF_Z_PLACED) { - copy_g[0] = CopyFactors(model, lsetup, cell, copy_rotation[0], d_llg, nthreads, &ft); - const LikelihoodPeak lk = LikelihoodTranslationZ(lsetup, no_known, 0.0, copy_g[0], copy_position[0], + copy_g[0] = CopyFactors(model, lsetup(), cell, copy_rotation[0], d_llg, nthreads, &ft); + const std::vector> no_known(lsetup().Terms(), 0.0); + const LikelihoodPeak lk = LikelihoodTranslationZ(lsetup(), no_known, 0.0, copy_g[0], copy_position[0], model.error_A, nthreads); cand.llg = lk.llg; cand.llg_z = lk.z; @@ -2186,9 +2365,9 @@ MolecularReplacementResult search_model(const SearchModel &model, const std::vec std::vector rots2 = rots; rots2.insert(rots2.end(), copy_rotation.begin(), copy_rotation.end()); if (!setup2) - setup2 = MakeTranslationSetup(refl, cell, sg, true); + setup2 = setup2_made.get(); const std::vector> known = KnownStructureFactors(ft, *setup2, cell, copy_rotation, - copy_position); + copy_position, nthreads); std::vector tf2(rots2.size()); #ifdef JFJOCH_USE_CUDA if (gpu) @@ -2218,28 +2397,42 @@ MolecularReplacementResult search_model(const SearchModel &model, const std::vec // likelihood gain of the copy beside the known ones, and the best few of those through the // rigid body. A copy that only repeats a known one - moved by one of the group's origin shifts, // it doubles F on half the reflections and cancels it on the other half - gains nothing. - std::vector> rknown(rsetup.Terms(), 0.0); + std::vector> rknown(rsetup().Terms(), 0.0); for (size_t c = 0; c < copies.size(); c++) { if (copy_rg[c].empty()) - copy_rg[c] = CopyFactors(model, rsetup, cell, copy_rotation[c], result.d_min_search, nthreads, &ft); - AddCopy(rsetup, copy_rg[c], copy_position[c], rknown); + copy_rg[c] = CopyFactors(model, rsetup(), cell, copy_rotation[c], result.d_min_search, nthreads, &ft); + AddCopy(rsetup(), copy_rg[c], copy_position[c], rknown); } double known_scale = 0; - const double rknown_llg = RefinedLogLikelihoodGain(rsetup, rknown, model.error_A, &known_scale); + const double rknown_llg = RefinedLogLikelihoodGain(rsetup(), rknown, model.error_A, &known_scale); // Down the highest peaks until COPY_TRIALS that gain likelihood - a copy that lowers it is not // one, and a peak on a copy already placed (or a mate of one) does - and pack. The gain that // ranks them is taken at the known copies' sigma_A. std::vector trials; + // The CopyFactors the walk down the peaks may need, all made at once. + const auto factors_of = [&](size_t i) -> std::vector> & { + return i < rots.size() ? rot_g[i] : copy_rg[i - rots.size()]; + }; + std::vector missing; + for (size_t k = 0; k < by_tf2.size() && k < 2 * COPY_TRIALS; k++) + if (factors_of(by_tf2[k]).empty()) + missing.push_back(by_tf2[k]); + rsetup(); + ParallelFor(static_cast(missing.size()), nthreads, [&](int m) { + factors_of(missing[m]) = CopyFactors(model, rsetup(), cell, rots2[missing[m]], + result.d_min_search, nthreads, &ft); + }); + // Their gains, all taken at once. + std::vector gains(std::min(by_tf2.size(), 2 * COPY_TRIALS)); + ParallelFor(static_cast(gains.size()), nthreads, [&](int k) { + std::vector> f = rknown; + AddCopy(rsetup(), factors_of(by_tf2[k]), tf2[by_tf2[k]].t, f); + gains[k] = LogLikelihoodGain(rsetup(), f, model.error_A, known_scale) - rknown_llg; + }); for (size_t k = 0; k < by_tf2.size() && k < 2 * COPY_TRIALS && trials.size() < COPY_TRIALS; k++) { Trial trial; trial.i = by_tf2[k]; - std::vector> &g = trial.i < rots.size() ? rot_g[trial.i] - : copy_rg[trial.i - rots.size()]; - if (g.empty()) - g = CopyFactors(model, rsetup, cell, rots2[trial.i], result.d_min_search, nthreads, &ft); - std::vector> f = rknown; - AddCopy(rsetup, g, tf2[trial.i].t, f); - trial.gain = LogLikelihoodGain(rsetup, f, model.error_A, known_scale) - rknown_llg; + trial.gain = gains[k]; if (!(trial.gain > 0)) continue; trial.copy = PlaceSearchModel(model, rots2[trial.i], tf2[trial.i].t, cell, sg); @@ -2272,16 +2465,16 @@ MolecularReplacementResult search_model(const SearchModel &model, const std::vec trial.refined.gain = trial.refined.llg = trial.refined.z = NAN; if (tf2[trial.i].z < TF_Z_PLACED) { if (!lknown) { - lknown.emplace(lsetup.Terms(), 0.0); + lknown.emplace(lsetup().Terms(), 0.0); for (size_t c = 0; c < copies.size(); c++) { if (copy_g[c].empty()) - copy_g[c] = CopyFactors(model, lsetup, cell, copy_rotation[c], d_llg, nthreads, &ft); - AddCopy(lsetup, copy_g[c], copy_position[c], *lknown); + copy_g[c] = CopyFactors(model, lsetup(), cell, copy_rotation[c], d_llg, nthreads, &ft); + AddCopy(lsetup(), copy_g[c], copy_position[c], *lknown); } - known_llg = RefinedLogLikelihoodGain(lsetup, *lknown, model.error_A); + known_llg = RefinedLogLikelihoodGain(lsetup(), *lknown, model.error_A); } - trial.g = CopyFactors(model, lsetup, cell, trial.rotation, d_llg, nthreads, &ft); - trial.refined = LikelihoodTranslationZ(lsetup, *lknown, known_llg, trial.g, centroid(trial.copy), + trial.g = CopyFactors(model, lsetup(), cell, trial.rotation, d_llg, nthreads, &ft); + trial.refined = LikelihoodTranslationZ(lsetup(), *lknown, known_llg, trial.g, centroid(trial.copy), model.error_A, nthreads); } glog.Info("Model prior: {:<12} copy {} TF Z {:5.1f} LLG gain {:7.1f}, refined {:7.1f} Z {:5.1f} clash {:4.1f}% R-work {:.3f} R-free {:.3f}", @@ -2401,19 +2594,9 @@ 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 { - int best = -1; - while (true) { - best = -1; - for (size_t i = 0; i < result.candidates.size(); i++) { - const MRCandidate &c = result.candidates[i]; - if (c.merge == merge && c.verdict == "PLACED" - && (best < 0 || c.r_free < result.candidates[best].r_free)) - best = static_cast(i); - } - if (best < 0) - return -1; - MRCandidate &c = result.candidates[best]; - const auto t_null = std::chrono::steady_clock::now(); + // The null of candidate i. + const auto run_null = [&](int i) { + const MRCandidate &c = result.candidates[i]; RigidBodyEngines engines = make_engines(*c.placed, c.space_group, *merges[merge].merged, result.d_min_score, NULL_REPLICATES, logger); // One copy: its replicates placed by the translation function, as it was. @@ -2440,9 +2623,42 @@ MolecularReplacementResult search_model(const SearchModel &model, const std::vec out.push_back(p.t); return out; }; - c.null_sigma = null_sigma(*c.placed, orientations[best], c.space_group, *merges[merge].merged, - result.d_min_score, engines.get(), nthreads, &scorers, translate); - result.t_null += seconds_since(t_null); + return null_sigma(*c.placed, orientations[i], c.space_group, *merges[merge].merged, result.d_min_score, + engines.get(), 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); + bool first = true; + int best = -1; + while (true) { + best = -1; + for (size_t i = 0; i < result.candidates.size(); i++) { + const MRCandidate &c = result.candidates[i]; + if (c.merge == merge && c.verdict == "PLACED" + && (best < 0 || c.r_free < result.candidates[best].r_free)) + best = static_cast(i); + } + if (best < 0) + return -1; + MRCandidate &c = result.candidates[best]; + if (std::isnan(nulls[best])) { + const auto t_null = std::chrono::steady_clock::now(); + std::vector todo{best}; + if (!first) + for (size_t i = 0; i < result.candidates.size(); i++) + 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)); + for (size_t k = 0; k < todo.size(); k++) + nulls[todo[k]] = running[k].get(); + result.t_null += seconds_since(t_null); + first = false; + } + c.null_sigma = nulls[best]; if (c.null_sigma >= NULL_SIGMA_MIN) break; c.verdict = "NULL"; diff --git a/image_analysis/structure_refinement/MolecularReplacement.h b/image_analysis/structure_refinement/MolecularReplacement.h index d775e8b09..ddb9576d6 100644 --- a/image_analysis/structure_refinement/MolecularReplacement.h +++ b/image_analysis/structure_refinement/MolecularReplacement.h @@ -159,7 +159,8 @@ TranslationSetup MakeTranslationSetup(const SearchReflections &refl, const gemmi std::vector> KnownStructureFactors(const MolecularTransform &ft, const TranslationSetup &setup, const gemmi::UnitCell &cell, const std::vector &rotations, - const std::vector &positions); + const std::vector &positions, + size_t nthreads = 1); // The translation function of the model in orientation `rotation`: its highest point as a Z-score over // every point of the function, and where it is (the fractional position of the model's centroid).