From 6c8ead466a539537836ca0a3fcc6eb91e3af2ff2 Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Mon, 12 Oct 2026 16:39:03 +0200 Subject: [PATCH 1/6] rugnux --af-model: molecular replacement faster - the score's scale fit on the GPU, the expected LLG gate, the serial parts side by side The search was CPU-bound: on a 54-chain partial complex ~80% of all CPU samples were the score's FitModelScale grid (and libm's exp under it), with the GPU busy 13 of 37 s. - score_placement fits its scale on the GPU where its structure factors are made there (ModelScaleGPU's FitSolvent / Fit, as the rigid body already does); equal to rounding, the CPU build unchanged. - A model whose expected log-likelihood gain (sum sigma_A^4 / 2, McCoy et al. 2017; sigma_A^2 the part of the scattering one copy is by Matthews, times Luzzati's D^2) is below 10 is not searched for. Over the 193-set survey the lowest eLLG of a placed model is 26; every PAE domain below 24 placed nowhere, so no decision changes. - The likelihood setups made where first needed, one serving both where the reflections are the same. - After a candidate fails its null, the other placed candidates' nulls side by side. - The first copy's placements refined side by side on an engine each; the groups' translation functions side by side; CopyFactors and KnownStructureFactors over threads; the further copies' CopyFactors and gains for the peaks walked all at once; the known-structure TF setup made while the first copy is refined. All exact. Survey replay (193 sets): no decision changes; two R values differ in the 3rd decimal (GPU scale). Alternating runs, MODEL_PRIOR_SECONDS base -> new: 54-chain complex 38.3 -> 14.7, 4mtk 18.3 -> 9.7, 9v16 13.7 -> 8.3 (shared machine, indicative). Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01EizBKhTqrAYago9KJG3dA3 --- .../ModelStructureFactorsGPU.h | 2 + .../MolecularReplacement.cpp | 408 +++++++++++++----- .../MolecularReplacement.h | 3 +- 3 files changed, 316 insertions(+), 97 deletions(-) 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). From c6912e2a066008fc15a020ab6af12ee8030c008e Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Mon, 12 Oct 2026 16:49:02 +0200 Subject: [PATCH 2/6] rugnux --af-model: tNCS half-vector moves checked, the merge's own group not ruled out beside tNCS, P1's first copy judged by its orientation (WIP, validation pending) Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01EizBKhTqrAYago9KJG3dA3 --- docs/ACKNOWLEDGEMENT.md | 7 +- .../MolecularReplacement.cpp | 172 +++++++++++++++++- 2 files changed, 168 insertions(+), 11 deletions(-) diff --git a/docs/ACKNOWLEDGEMENT.md b/docs/ACKNOWLEDGEMENT.md index d6264b1b6..856cba036 100644 --- a/docs/ACKNOWLEDGEMENT.md +++ b/docs/ACKNOWLEDGEMENT.md @@ -481,7 +481,12 @@ boundaries of molecular replacement with maximum likelihood" (2001), Acta Cryst. Z-score over the same orientation at other positions, as the likelihood-enhanced translation function's is: A. J. McCoy, R. W. Grosse-Kunstleve, L. C. Storoni and R. J. Read, "Likelihood-enhanced fast translation functions" (2005), Acta Cryst. D61, 458-464 -[doi:10.1107/S0907444905001617](https://doi.org/10.1107/S0907444905001617). The logarithm of the Bessel +[doi:10.1107/S0907444905001617](https://doi.org/10.1107/S0907444905001617). A model whose expected +log-likelihood gain - the sum of sigma_A^4 / 2 over the reflections - is too small for a placement to stand +out is not searched for: A. J. McCoy, R. D. Oeffner, A. G. Wrobel, J. R. M. Ojala, K. Tryggvason, +B. Lohkamp and R. J. Read, "Ab initio solution of macromolecular crystal structures without direct +methods" (2017), Proc. Natl. Acad. Sci. USA 114, 3637-3641 +[doi:10.1073/pnas.1701640114](https://doi.org/10.1073/pnas.1701640114). The logarithm of the Bessel function I0 in it is evaluated with the polynomial approximations of M. Abramowitz and I. A. Stegun, "Handbook of Mathematical Functions" (1964), National Bureau of Standards, 9.8.1-9.8.2. A predicted model is cut into rigid domains by its predicted aligned error, as Oeffner et al. (above) do before diff --git a/image_analysis/structure_refinement/MolecularReplacement.cpp b/image_analysis/structure_refinement/MolecularReplacement.cpp index 9c11693a9..450225a19 100644 --- a/image_analysis/structure_refinement/MolecularReplacement.cpp +++ b/image_analysis/structure_refinement/MolecularReplacement.cpp @@ -102,6 +102,14 @@ 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; +// tNCS by half a lattice vector v: copies v apart make the reflections with h.v half-integral weak, and a +// placement moved by (v + o) / 2, o an origin shift of the group, fits the strong ones as well as the right one +// does; only the weak ones tell them apart. A placement is compared with every such move of itself on those +// reflections (R_tncs_weak), and one of them that fits them better by TNCS_SHIFT_MARGIN is taken instead; where the +// two best are closer than that, the placement is not decided on. A vector is taken as half-integral within +// TNCS_HALF_TOLERANCE. +constexpr double TNCS_SHIFT_MARGIN = 0.01; +constexpr double TNCS_HALF_TOLERANCE = 0.05; double seconds_since(std::chrono::steady_clock::time_point t0) { return std::chrono::duration(std::chrono::steady_clock::now() - t0).count(); @@ -161,6 +169,9 @@ struct PlacementScore { // Against the data as twinned by a law (score_placement's `law`): at the fraction that gives the lowest // R-work. double r_work_twin = NAN, r_free_twin = NAN, twin_fraction = NAN; + // Over the reflections a tNCS vector v weakens - h.v half-integral - work and free together (score_placement's + // `tncs`). + double r_tncs_weak = NAN; }; gemmi::AsuData> working_amplitudes(const std::vector &merged, @@ -205,37 +216,54 @@ void model_factors(const gemmi::Structure &st, double d_min, gemmi::AsuData> &fmodel, const gemmi::Scaling &scaling, - const gemmi::SpaceGroup &sg, const std::vector &merged, double d_min) { + const gemmi::SpaceGroup &sg, const std::vector &merged, double d_min, + const gemmi::Fractional *tncs = nullptr) { const gemmi::GroupOps gops = sg.operations(); const gemmi::ReciprocalAsu asu(&sg); - std::unordered_map> obs; + struct Obs { + double f; + bool free; + bool weak; // h.v half-integral + }; + std::unordered_map obs; for (const MergedReflection &r : merged) { if (std::isnan(r.F) || r.d < d_min) continue; gemmi::Miller h{{r.h, r.k, r.l}}; if (!asu.is_in(h)) h = asu.to_asu(h, gops).first; - obs[hkl_key(h)] = {r.F, r.rfree_flag}; + bool weak = false; + if (tncs != nullptr) { + const double hv = h[0] * tncs->x + h[1] * tncs->y + h[2] * tncs->z; + weak = std::fabs(hv - std::round(hv)) > 0.25; + } + obs[hkl_key(h)] = {r.F, r.rfree_flag, weak}; } - double num_w = 0, den_w = 0, num_f = 0, den_f = 0; + double num_w = 0, den_w = 0, num_f = 0, den_f = 0, num_t = 0, den_t = 0; for (const auto &hv : fmodel.v) { const auto it = obs.find(hkl_key(hv.hkl)); if (it == obs.end()) continue; - const double fo = it->second.first, fc = std::abs(hv.value); - if (it->second.second) { + const double fo = it->second.f, fc = std::abs(hv.value); + if (it->second.free) { num_f += std::fabs(fo - fc); den_f += fo; } else { num_w += std::fabs(fo - fc); den_w += fo; } + if (it->second.weak) { + num_t += std::fabs(fo - fc); + den_t += fo; + } } PlacementScore out; out.k_sol = scaling.k_sol; out.b_sol = scaling.b_sol; out.r_work = den_w > 0 ? num_w / den_w : 1; out.r_free = den_f > 0 ? num_f / den_f : 1; + if (tncs != nullptr) + out.r_tncs_weak = den_t > 0 ? num_t / den_t : 1; return out; } @@ -423,7 +451,7 @@ void twin_r_factors(const gemmi::AsuData> &fmodel, const gem PlacementScore score_placement(const gemmi::Structure &st, const gemmi::SpaceGroup &sg, const std::vector &merged, double d_min, size_t nthreads, ScoreEngines *engines = nullptr, const PlacementScore *solvent = nullptr, - const gemmi::Op *law = nullptr) { + const gemmi::Op *law = nullptr, const gemmi::Fractional *tncs = nullptr) { gemmi::AsuData> fmask, fmodel; #ifdef JFJOCH_USE_CUDA ScoreEngines::Engine *gpu = engines != nullptr ? engines->get(st, sg, d_min) : nullptr; @@ -454,7 +482,7 @@ PlacementScore score_placement(const gemmi::Structure &st, const gemmi::SpaceGro FitModelScale(scaling, {.b_hi = SolventBMax(st.models[0])}, nthreads); } scaling.scale_data(fmodel, &fmask); - PlacementScore out = r_factors(fmodel, scaling, sg, merged, d_min); + PlacementScore out = r_factors(fmodel, scaling, sg, merged, d_min, tncs); if (law != nullptr) twin_r_factors(fmodel, sg, merged, d_min, *law, out); return out; @@ -651,6 +679,33 @@ bool enantiomorphs(const gemmi::SpaceGroup &a, const gemmi::SpaceGroup &b) { return other != nullptr && other->number == b.number; } +// The tNCS vector where it is half a lattice vector - every component within TNCS_HALF_TOLERANCE of 0 or 1/2 - +// on those halves; empty otherwise. +std::optional half_tncs_vector(const TranslationalNcs &tncs) { + if (!tncs.present) + return std::nullopt; + double v[3] = {tncs.vector.x, tncs.vector.y, tncs.vector.z}; + for (double &x : v) { + const double h = std::round(2 * x) / 2; + if (std::fabs(x - h) > TNCS_HALF_TOLERANCE) + return std::nullopt; + x = h - std::floor(h); + } + return gemmi::Fractional(v[0], v[1], v[2]); +} + +// Whether moving the origin by `s` (a multiple of 1/24 in each component) leaves `sg`'s operators as they are: +// a placement moved by s is then the same crystal. +bool is_origin_shift(const gemmi::SpaceGroup &sg, const gemmi::Fractional &s) { + gemmi::Op shift = gemmi::Op::identity(); + const double x[3] = {s.x, s.y, s.z}; + for (int a = 0; a < 3; a++) + shift.tran[a] = static_cast(std::lround(x[a] * gemmi::Op::DEN)) % gemmi::Op::DEN; + gemmi::GroupOps ops = sg.operations(); + ops.change_basis_forward(shift); + return ops.is_same_as(sg.operations()); +} + } // namespace // --- the model --- @@ -2531,8 +2586,11 @@ MolecularReplacementResult search_model(const SearchModel &model, const std::vec cand.verdict = "CLASH"; // The likelihood's Z-score stands in for the translation function's only where the copies placed // are as many as the cell most probably holds: beside content the search has not found, a copy's - // likelihood over random positions of its own is no measure of the placement. - else if (!(cand.tf_z >= TF_Z_PLACED || (cand.llg_z >= TF_Z_PLACED && cand.copies >= cand.copies_likely))) + // likelihood over random positions of its own is no measure of the placement. In P1 the first copy sits + // at the origin, which is arbitrary there and which no function of its position can single out: one copy + // is judged by its orientation alone, against the null. + else if (!(cand.tf_z >= TF_Z_PLACED || (cand.llg_z >= TF_Z_PLACED && cand.copies >= cand.copies_likely) + || (gops.order() == 1 && cand.copies == 1))) cand.verdict = "NOT_PLACED"; else cand.verdict = "PLACED"; @@ -2630,6 +2688,79 @@ MolecularReplacementResult search_model(const SearchModel &model, const std::vec // by side, each taken where the choice comes to it. std::vector nulls(result.candidates.size(), NAN); bool first = true; + // With tNCS by half a lattice vector, candidate i against its moves by half that vector (TNCS_SHIFT_MARGIN): + // the better taken into it, and whether the two best are too close to tell apart. + // The tNCS vector is the answer's merge's, in its cell. + const std::optional half = merge == 0 ? half_tncs_vector(result.tncs) : std::nullopt; + std::vector shift_checked(result.candidates.size(), 0); + const auto check_shifts = [&](int i) -> bool { + MRCandidate &c = result.candidates[i]; + const gemmi::Fractional v = *half; + const std::vector &merged = *merges[merge].merged; + const double d_min = result.d_min_score; + struct Option { + gemmi::Fractional shift; + std::shared_ptr placed; + PlacementScore score; + double clash = 0; + }; + std::vector