diff --git a/docs/RUGNUX_REPORT.md b/docs/RUGNUX_REPORT.md index 516d1c699..714909a9c 100644 --- a/docs/RUGNUX_REPORT.md +++ b/docs/RUGNUX_REPORT.md @@ -972,13 +972,20 @@ data's limit). One copy in the asymmetric unit only, for now. The table has one row per group and merge. A row is `PLACED` where its translation-function Z-score (`TF_Z`) reaches 8, at most 5% of its C-alpha atoms clash with symmetry mates (`CLASH`), one copy fits the cell (`MATTHEWS`: Vm below 1.5 A^3/Da; `MULTI_COPY`: Vm of 4 or more, which is reported but not -supported yet) and its R-work stands at least 3 sd above the same model refitted from random -orientations (`NULL_SIGMA`, or `NULL` where it does not). `MODEL_PRIOR_DECISION=` is then: +supported yet). The null - the same model refitted from random orientations, whose R-work the +placement's has to beat by 3 sd (`NULL_SIGMA`) - is built for the candidate the choice falls on; one +that does not beat it is marked `NULL` and the choice falls on the next. A placement clashing with its +symmetry mates over twice the limit is not refined. The rotation and translation functions run on the +GPU where there is one (`MODEL_PRIOR_SEARCH_DEVICE=`). `MODEL_PRIOR_DECISION=` is then: - `AGREES` - the placed group with the lowest R-free is the answer's, ahead of every other group of the merge that packs by 0.03 in R-free; - `ENANTIOMORPH` - it is the answer's enantiomorph, and the answer's files take that label (a label - only, as with `--model`; no reflection moves); + only, as with `--model`; no reflection moves). It is taken only as a `--model`'s enantiomorph is: + model validation is run on the placed model, and the model has to fit against its null and the + anomalous map at its atoms must not be inverted (`MODEL_PRIOR_HAND_CHECK=`, + `MODEL_PRIOR_HAND_FIT_SIGMA=`, `MODEL_PRIOR_ANOMALOUS_INVERTED=`). Where it is not, the decision reads + `ENANTIOMORPH_VETOED` and the label stays; - `SPACE_GROUP` - another group with the answer's rotations; written as `_mr.mtz`/`.cif` from the answer's merge, the reflections that group extinguishes left out; - `TWIN_SUBGROUP` - no group of the answer's merge can hold the model at all (`MATTHEWS`/`CLASH` in diff --git a/image_analysis/structure_refinement/CMakeLists.txt b/image_analysis/structure_refinement/CMakeLists.txt index 9dcef0e26..aacd1211b 100644 --- a/image_analysis/structure_refinement/CMakeLists.txt +++ b/image_analysis/structure_refinement/CMakeLists.txt @@ -3,6 +3,7 @@ ADD_LIBRARY(JFJochStructureRefinement STATIC ModelGrid.cpp ModelGrid.h ModelScaling.cpp ModelScaling.h MolecularReplacement.cpp MolecularReplacement.h + $<$:MolecularReplacementGPU.cu> MolecularReplacementGPU.h ModelValidation.cpp ModelValidation.h RigidBodyRefine.cpp RigidBodyRefine.h $<$:RigidBodyGPU.cpp> diff --git a/image_analysis/structure_refinement/MolecularReplacement.cpp b/image_analysis/structure_refinement/MolecularReplacement.cpp index c57fe5eb9..ed7719c3c 100644 --- a/image_analysis/structure_refinement/MolecularReplacement.cpp +++ b/image_analysis/structure_refinement/MolecularReplacement.cpp @@ -27,7 +27,10 @@ #include "ModelScaling.h" #include "RigidBodyRefine.h" #ifdef JFJOCH_USE_CUDA +#include "MolecularReplacementGPU.h" #include "RigidBodyGPU.h" +#include "../../common/CUDAWrapper.h" +#include "../../common/JFJochException.h" #endif namespace { @@ -536,17 +539,24 @@ SearchReflections MakeSearchReflections(const std::vector &mer // --- the rotation function --- +std::vector CopyVectors(const SearchReflections &refl, const gemmi::UnitCell &cell, + const gemmi::SpaceGroup &sg) { + const gemmi::GroupOps gops = sg.operations(); + const size_t nops = gops.sym_ops.size(); + const gemmi::Mat33 frac_t = cell.frac.mat.transpose(); + std::vector s(refl.hkl.size() * nops); + for (size_t i = 0; i < refl.hkl.size(); i++) + for (size_t j = 0; j < nops; j++) + s[i * nops + j] = frac_t.multiply(miller_vec(rotated_index(gops.sym_ops[j], refl.hkl[i]))); + return s; +} + std::vector RotationFunction(const MolecularTransform &ft, const SearchReflections &refl, const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg, const std::vector &rotations, size_t nthreads) { - const gemmi::GroupOps gops = sg.operations(); - const size_t nops = gops.sym_ops.size(), n = refl.hkl.size(); - const gemmi::Mat33 frac_t = cell.frac.mat.transpose(); - // Copy j of the model contributes |F_model(s)|^2 at s = (h R_j) as a Cartesian reciprocal vector. - std::vector s(n * nops); - for (size_t i = 0; i < n; i++) - for (size_t j = 0; j < nops; j++) - s[i * nops + j] = frac_t.multiply(miller_vec(rotated_index(gops.sym_ops[j], refl.hkl[i]))); + const size_t n = refl.hkl.size(); + const std::vector s = CopyVectors(refl, cell, sg); + const size_t nops = s.size() / n; double w_mean = 0; for (double e2 : refl.e2) w_mean += e2; @@ -567,95 +577,112 @@ std::vector RotationFunction(const MolecularTransform &ft, const SearchR bin_sum[refl.bin[i]] += ic[i]; bin_n[refl.bin[i]]++; } - double x_mean = 0; + double x_sum = 0, x_ss = 0, xy = 0; for (size_t i = 0; i < n; i++) { - ic[i] /= bin_sum[refl.bin[i]] / bin_n[refl.bin[i]]; - x_mean += ic[i]; + const double x = ic[i] / (bin_sum[refl.bin[i]] / bin_n[refl.bin[i]]); + x_sum += x; + x_ss += x * x; + xy += x * (refl.e2[i] - w_mean); } - x_mean /= static_cast(n); - double sxy = 0, sxx = 0; - for (size_t i = 0; i < n; i++) { - sxy += (ic[i] - x_mean) * (refl.e2[i] - w_mean); - sxx += (ic[i] - x_mean) * (ic[i] - x_mean); - } - out[r] = sxx > 0 && w_ss > 0 ? sxy / std::sqrt(sxx * w_ss) : 0.0; + const double sxx = x_ss - x_sum * x_sum / static_cast(n); + out[r] = sxx > 0 && w_ss > 0 ? xy / std::sqrt(sxx * w_ss) : 0.0; }); return out; } // --- the translation function --- +// +// The full-symmetry translation function as one Fourier synthesis, following Crowther & Blow (1967) +// Acta Cryst. 23, 544-548: with the model at t, copy j adds G_j(h) exp(2 pi i (h R_j).t) exp(2 pi i h.t_j) +// to F(h), so |F(h)|^2 summed against E^2 - 1 over the sphere is, for each pair of copies j != k, a +// Fourier series in t with coefficient (E^2 - 1) g_j conj(g_k), g_j = G_j(h) exp(2 pi i h.t_j), at the +// index h R_j - h R_k. Every pair's terms go on one grid, and one FFT gives the function at every t. +// The pairs j = k are the rotation function's, constant in t. -TranslationPeak TranslationFunction(const MolecularTransform &ft, const SearchReflections &refl, - const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg, - const gemmi::Mat33 &rotation) { - // The full-symmetry translation function as one Fourier synthesis, following Crowther & Blow - // (1967) Acta Cryst. 23, 544-548: with the model at t, copy j adds G_j(h) exp(2 pi i (h R_j).t) - // exp(2 pi i h.t_j) to F(h), so |F(h)|^2 summed against E^2 - 1 over the sphere is, for each pair - // of copies j != k, a Fourier series in t with coefficient (E^2 - 1) G_j conj(G_k) exp(2 pi i - // h.(t_j - t_k)) at the index h R_j - h R_k. Every pair's terms go on one grid, and one FFT gives - // the function at every t. The pairs j = k are the rotation function's, constant in t. +TranslationSetup MakeTranslationSetup(const SearchReflections &refl, const gemmi::UnitCell &cell, + const gemmi::SpaceGroup &sg) { const gemmi::GroupOps gops = sg.operations(); - const size_t nops = gops.sym_ops.size(); - TranslationPeak out; - if (nops < 2) - return out; // P1: every position is the same + const int nops = static_cast(gops.sym_ops.size()); const gemmi::Mat33 frac_t = cell.frac.mat.transpose(); - const gemmi::Mat33 rt = rotation.transpose(); - // Per copy j: k = h R_j, and g = G_j(h) exp(2 pi i h.t_j), so a pair's coefficient is g_j conj(g_k). - struct Term { gemmi::Miller h; std::vector k; std::vector> g; double w; }; - std::vector terms; - std::vector bin_sum(refl.bins, 0.0); - std::vector bin_n(refl.bins, 0); - std::vector term_bin; - for (size_t i = 0; i < refl.hkl.size(); i++) { - if (gops.is_systematically_absent(refl.hkl[i])) - continue; - Term t{refl.hkl[i], {}, {}, refl.multiplicity[i] * (refl.e2[i] - 1.0)}; - double sum = 0; - for (const gemmi::Op &op : gops.sym_ops) { - t.k.push_back(rotated_index(op, refl.hkl[i])); - const double ht = (t.h[0] * op.tran[0] + t.h[1] * op.tran[1] + t.h[2] * op.tran[2]) - / static_cast(gemmi::Op::DEN); - t.g.push_back(ft.At(rt.multiply(frac_t.multiply(miller_vec(t.k.back())))) * std::polar(1.0, 2.0 * PI * ht)); - sum += std::norm(t.g.back()); - } - bin_sum[refl.bin[i]] += sum; - bin_n[refl.bin[i]]++; - term_bin.push_back(refl.bin[i]); - terms.push_back(std::move(t)); - } - if (terms.empty()) - return out; - int max_m[3] = {0, 0, 0}; - for (const Term &t : terms) - for (size_t j = 0; j < nops; j++) - for (size_t k = 0; k < nops; k++) - for (int a = 0; a < 3; a++) - max_m[a] = std::max(max_m[a], std::abs(t.k[j][a] - t.k[k][a])); - const int nu = good_even_size(2 * max_m[0] + 2), nv = good_even_size(2 * max_m[1] + 2), - nw = good_even_size(2 * max_m[2] + 2); - const int nwh = nw / 2 + 1; - std::vector> c(static_cast(nu) * nv * nwh, 0.0); - for (size_t i = 0; i < terms.size(); i++) { - const Term &t = terms[i]; - const double norm = bin_sum[term_bin[i]] / bin_n[term_bin[i]]; - for (size_t j = 0; j < nops; j++) - for (size_t k = 0; k < nops; k++) { - if (j == k) - continue; - const int m[3] = {t.k[j][0] - t.k[k][0], t.k[j][1] - t.k[k][1], t.k[j][2] - t.k[k][2]}; - // Only l >= 0 is stored: the pair (k, j) carries the conjugate term at -m. - if (m[2] < 0) - continue; - const std::complex v = t.w / norm * t.g[j] * std::conj(t.g[k]); - const int u = (m[0] % nu + nu) % nu, vv = (m[1] % nv + nv) % nv; - c[u + static_cast(nu) * (vv + static_cast(nv) * m[2])] += v; + TranslationSetup out; + out.ops = nops; + out.bins = refl.bins; + for (int j = 0; j < nops; j++) + for (int k = 0; k < nops; k++) + if (j != k) { + out.pair_j.push_back(j); + out.pair_k.push_back(k); } + std::vector rotated; // terms x ops + for (size_t i = 0; i < refl.hkl.size(); i++) { + const gemmi::Miller &h = refl.hkl[i]; + if (gops.is_systematically_absent(h)) + continue; + for (const gemmi::Op &op : gops.sym_ops) { + rotated.push_back(rotated_index(op, h)); + out.s.push_back(frac_t.multiply(miller_vec(rotated.back()))); + const double ht = (h[0] * op.tran[0] + h[1] * op.tran[1] + h[2] * op.tran[2]) + / static_cast(gemmi::Op::DEN); + out.phase.push_back(std::polar(1.0, 2.0 * PI * ht)); + } + out.weight.push_back(refl.multiplicity[i] * (refl.e2[i] - 1.0)); + out.bin.push_back(refl.bin[i]); + } + int max_m[3] = {0, 0, 0}; + for (size_t t = 0; t < out.Terms(); t++) + for (size_t p = 0; p < out.pair_j.size(); p++) + for (int a = 0; a < 3; a++) + max_m[a] = std::max(max_m[a], std::abs(rotated[t * nops + out.pair_j[p]][a] + - rotated[t * nops + out.pair_k[p]][a])); + out.nu = good_even_size(2 * max_m[0] + 2); + out.nv = good_even_size(2 * max_m[1] + 2); + out.nw = good_even_size(2 * max_m[2] + 2); + for (size_t t = 0; t < out.Terms(); t++) + for (size_t p = 0; p < out.pair_j.size(); p++) { + const gemmi::Miller &kj = rotated[t * nops + out.pair_j[p]], &kk = rotated[t * nops + out.pair_k[p]]; + const int l = kj[2] - kk[2]; + out.target_u.push_back(l < 0 ? -1 : ((kj[0] - kk[0]) % out.nu + out.nu) % out.nu); + out.target_v.push_back(((kj[1] - kk[1]) % out.nv + out.nv) % out.nv); + out.target_l.push_back(l); + } + return out; +} + +TranslationPeak TranslationFunction(const MolecularTransform &ft, const TranslationSetup &setup, + const gemmi::Mat33 &rotation) { + TranslationPeak out; + const size_t terms = setup.Terms(), nops = setup.ops, npairs = setup.pair_j.size(); + if (nops < 2 || terms == 0) + return out; // P1: every position is the same + const gemmi::Mat33 rt = rotation.transpose(); + std::vector> g(terms * nops); + std::vector bin_sum(setup.bins, 0.0); + std::vector bin_n(setup.bins, 0); + for (size_t t = 0; t < terms; t++) { + double sum = 0; + for (size_t j = 0; j < nops; j++) { + g[t * nops + j] = ft.At(rt.multiply(setup.s[t * nops + j])) * setup.phase[t * nops + j]; + sum += std::norm(g[t * nops + j]); + } + bin_sum[setup.bin[t]] += sum; + bin_n[setup.bin[t]]++; + } + const int nu = setup.nu, nv = setup.nv, nwh = setup.nw / 2 + 1; + std::vector> c(static_cast(nu) * nv * nwh, 0.0); + for (size_t t = 0; t < terms; t++) { + const double w = setup.weight[t] / (bin_sum[setup.bin[t]] / bin_n[setup.bin[t]]); + for (size_t p = 0; p < npairs; p++) { + const size_t e = t * npairs + p; + if (setup.target_u[e] < 0) + continue; + c[setup.target_u[e] + static_cast(nu) * (setup.target_v[e] + static_cast(nv) * setup.target_l[e])] + += w * g[t * nops + setup.pair_j[p]] * std::conj(g[t * nops + setup.pair_k[p]]); + } } // MapFromFPhi sums F exp(-2 pi i m.x): with F = conj(C) that is the conjugate of the series, which // is real. gemmi::FPhiGrid hkl; - hkl.unit_cell = cell; + hkl.unit_cell = ft.Grid().unit_cell; hkl.spacegroup = &gemmi::get_spacegroup_p1(); hkl.axis_order = gemmi::AxisOrder::XYZ; hkl.half_l = true; @@ -795,9 +822,24 @@ MolecularReplacementResult RunMolecularReplacement(const std::string &model_path t0 = std::chrono::steady_clock::now(); const MolecularTransform ft(model, result.d_min_search, nthreads); +#ifdef JFJOCH_USE_CUDA + // The rotation and translation functions on the GPU where there is one; the rest of the search is + // the same either way. + std::unique_ptr gpu; + if (get_gpu_count() > 0) { + try { + gpu = std::make_unique(ft); + } catch (const JFJochException &e) { + result.failure_reason = std::string("GPU: ") + e.what(); + return result; + } + } + result.gpu_search = gpu != nullptr; +#endif result.t_transform = seconds_since(t0); const std::vector sphere = UniformRotations(ROTATION_STEP_DEG); + std::vector orientations; // of each candidate's placement, before the rigid body // One merge: the rotation function over its point group, then for each of its groups the // translation function and the best placements through packing, the rigid body and R-free. const auto search_merge = [&](size_t m) { @@ -810,7 +852,12 @@ MolecularReplacementResult RunMolecularReplacement(const std::string &model_path result.d_min_search, result.d_max_search); const std::vector pg = CartesianPointGroup(merge.space_group, cell); const std::vector rotations = RotationAsu(sphere, pg); +#ifdef JFJOCH_USE_CUDA + const std::vector rf = gpu ? gpu->RotationFunction(refl, cell, merge.space_group, rotations) + : RotationFunction(ft, refl, cell, merge.space_group, rotations, nthreads); +#else const std::vector rf = RotationFunction(ft, refl, cell, merge.space_group, rotations, nthreads); +#endif double rf_mean = 0, rf_sd = 0; for (double x : rf) rf_mean += x; @@ -840,10 +887,23 @@ MolecularReplacementResult RunMolecularReplacement(const std::string &model_path // --- 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(groups.size() * peaks.size()); - ParallelFor(static_cast(tf_all.size()), nthreads, [&](int i) { - tf_all[i] = TranslationFunction(ft, refl, cell, groups[i / peaks.size()], rotations[peaks[i % peaks.size()]]); - }); + std::vector peak_rotations; + for (size_t p : peaks) + peak_rotations.push_back(rotations[p]); + std::vector tf_all; + for (const gemmi::SpaceGroup &sg : groups) { + 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 +#endif + ParallelFor(static_cast(peaks.size()), nthreads, [&](int i) { + tf[i] = TranslationFunction(ft, setup, peak_rotations[i]); + }); + tf_all.insert(tf_all.end(), tf.begin(), tf.end()); + } result.t_translation += seconds_since(t0); for (size_t g = 0; g < groups.size(); g++) { @@ -872,14 +932,20 @@ MolecularReplacementResult RunMolecularReplacement(const std::string &model_path const size_t i = by_tf[k]; auto placed = std::make_shared( PlaceSearchModel(model, rotations[peaks[i]], tf[i].t, cell, sg)); + // A placement clashing with its own mates over twice the limit is not refined: the rigid + // body moves a search placement by a few degrees and a fraction of an angstrom, and the + // clash fraction moves by a percent or two with it. + double clash = ClashFraction(*placed); PlacementScore score{NAN, NAN}; - if (one_copy) { - if (k == 0) - engines = make_engines(*placed, sg, *merge.merged, d_min, NULL_REPLICATES + 1, logger); + if (one_copy && 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); + clash = ClashFraction(*placed); } - const double clash = ClashFraction(*placed); - if (!have || score.r_free < cand.r_free) { // NaN (not refined) only ever on k = 0 + // 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))) { have = true; cand.rf_z = rf_sd > 0 ? (rf[peaks[i]] - rf_mean) / rf_sd : 0.0; cand.tf_z = tf[i].z; @@ -905,14 +971,7 @@ MolecularReplacementResult RunMolecularReplacement(const std::string &model_path cand.verdict = "NOT_PLACED"; else cand.verdict = "PLACED"; - if (cand.verdict == "PLACED") { - t0 = std::chrono::steady_clock::now(); - cand.null_sigma = null_sigma(*cand.placed, orientation, sg, *merge.merged, d_min, engines.get(), - nthreads); - result.t_null += seconds_since(t0); - if (!(cand.null_sigma >= NULL_SIGMA_MIN)) - cand.verdict = "NULL"; - } + orientations.push_back(orientation); result.candidates.push_back(std::move(cand)); } }; @@ -921,15 +980,31 @@ MolecularReplacementResult RunMolecularReplacement(const std::string &model_path // Within one merge the candidates are scored on the same reflections, so their R-free values can // be compared directly: the placed candidate with the lowest R-free is chosen where it leads every // other candidate of that merge that packs by R_FREE_MARGIN. + // 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; - 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); + 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(); + RigidBodyEngines engines = make_engines(*c.placed, c.space_group, *merges[merge].merged, + result.d_min_score, NULL_REPLICATES, logger); + c.null_sigma = null_sigma(*c.placed, orientations[best], c.space_group, *merges[merge].merged, + result.d_min_score, engines.get(), nthreads); + result.t_null += seconds_since(t_null); + if (c.null_sigma >= NULL_SIGMA_MIN) + break; + c.verdict = "NULL"; } - if (best < 0) - return -1; for (size_t i = 0; i < result.candidates.size(); i++) { const MRCandidate &c = result.candidates[i]; if (static_cast(i) == best || c.merge != merge || c.clash_fraction > CLASH_MAX) diff --git a/image_analysis/structure_refinement/MolecularReplacement.h b/image_analysis/structure_refinement/MolecularReplacement.h index 89323630a..3bc628534 100644 --- a/image_analysis/structure_refinement/MolecularReplacement.h +++ b/image_analysis/structure_refinement/MolecularReplacement.h @@ -69,6 +69,10 @@ public: MolecularTransform(const SearchModel &model, double d_min, size_t nthreads); // F at the Cartesian reciprocal vector s (1/A), interpolated trilinearly. std::complex At(const gemmi::Vec3 &s) const; + // The grid At() interpolates (half-l, XYZ), its box edge in A and the blur it is unblurred by. + const gemmi::FPhiGrid &Grid() const { return grid_; } + double Box() const { return box_; } + double Blur() const { return blur_; } private: gemmi::FPhiGrid grid_; // half-l, on a P1 box `box_` A wide @@ -88,21 +92,44 @@ struct SearchReflections { SearchReflections MakeSearchReflections(const std::vector &merged, const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg, double d_min, double d_max); +// Copy j of the model contributes |F_model(s)|^2 at s = h R_j, the reflection's index under the j-th +// operator of the point group as a Cartesian reciprocal vector: refl.hkl.size() x |point group|, index-major. +std::vector CopyVectors(const SearchReflections &refl, const gemmi::UnitCell &cell, + const gemmi::SpaceGroup &sg); + // The rotation function at each of `rotations`: the correlation of E^2 - 1 with the bin-normalised // sum over the point group's copies of |F_model|^2. std::vector RotationFunction(const MolecularTransform &ft, const SearchReflections &refl, const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg, const std::vector &rotations, size_t nthreads); -// The translation function of the model in orientation `rotation` in `sg`: 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). +// What the translation function of one space group needs that does not depend on the orientation: +// for every reflection not absent in the group (a "term") and every operator j, s = h R_j as a +// Cartesian vector and exp(2 pi i h.t_j); the term's weight - its multiplicity times E^2 - 1 - and bin; +// and for every term and every ordered pair of operators (j, k), j != k, where the coefficient goes: +// the index h R_j - h R_k on the function's grid, or u = -1 where its l is negative (only l >= 0 is kept). +struct TranslationSetup { + int nu = 0, nv = 0, nw = 0; // the function's grid + int ops = 0; + std::vector s; // terms x ops + std::vector> phase; // terms x ops + std::vector weight; // terms + std::vector bin; // terms + int bins = 0; + std::vector pair_j, pair_k; // the ordered pairs + std::vector target_u, target_v, target_l; // terms x pairs + size_t Terms() const { return weight.size(); } +}; +TranslationSetup MakeTranslationSetup(const SearchReflections &refl, const gemmi::UnitCell &cell, + const gemmi::SpaceGroup &sg); + +// 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). struct TranslationPeak { double z = 0; gemmi::Fractional t; }; -TranslationPeak TranslationFunction(const MolecularTransform &ft, const SearchReflections &refl, - const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg, +TranslationPeak TranslationFunction(const MolecularTransform &ft, const TranslationSetup &setup, const gemmi::Mat33 &rotation); // The model placed: rotated about its centroid by `rotation`, the centroid at `t`, in `cell` and `sg`. @@ -151,6 +178,22 @@ struct MolecularReplacementResult { // The decision: the candidate the model chooses, or -1 where it chooses none, and in words why. int chosen = -1; std::string decision; + bool gpu_search = false; // the rotation and translation functions ran on the GPU + // An ENANTIOMORPH decision is taken into the answer's label only as model validation takes a model's + // enantiomorph (ValidateAgainstModel): the placed model has to fit against its null and the anomalous + // map at its atoms must not be inverted. Filled by the caller, which runs that validation. + struct HandCheck { + bool performed = false; + bool adopted = false; + bool model_fits = false; + double fit_sigma = NAN; + bool anomalous_inverted = false; + int anomalous_pairs = 0; + double anomalous_scatterer_mean_sigma = NAN; + std::string deepest_site; + double deepest_sigma = NAN; + double seconds = 0; + } hand_check; // Seconds per step. double t_model = 0, t_transform = 0, t_rotation = 0, t_translation = 0, t_rescore = 0, t_null = 0; }; diff --git a/image_analysis/structure_refinement/MolecularReplacementGPU.cu b/image_analysis/structure_refinement/MolecularReplacementGPU.cu new file mode 100644 index 000000000..89cf87171 --- /dev/null +++ b/image_analysis/structure_refinement/MolecularReplacementGPU.cu @@ -0,0 +1,436 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#include "MolecularReplacementGPU.h" + +#include +#include +#include + +#include + +#include "../../common/JFJochException.h" +#include "../indexing/CUDAMemHelpers.h" + +namespace { + +void cuda_err(cudaError_t val) { + if (val != cudaSuccess) + throw JFJochException(JFJochExceptionCategory::GPUCUDAError, cudaGetErrorString(val)); +} + +void cufft_err(cufftResult val) { + if (val != CUFFT_SUCCESS) + throw JFJochException(JFJochExceptionCategory::GPUCUDAError, + "cuFFT error " + std::to_string(static_cast(val))); +} + +constexpr int THREADS = 256; +constexpr int MAX_BINS = 16; // SEARCH_BINS (MolecularReplacement.cpp) is 10 +// The rotation function's per-reflection scratch is made for this many orientations at a time. +constexpr size_t ROTATION_BATCH = 2048; + +struct Transform { + const float2 *data; + int nu, nv, nwh; + float box, blur; +}; + +struct Rot { + float m[9]; // R^T, row-major: the model turned by R has F_R(s) = F_model(R^T s) +}; + +__device__ __forceinline__ float2 transform_at(const Transform &g, int i, int j, int l) { + bool conj = false; + if (l < 0) { // the half-l grid holds l >= 0; F(-h) = conj F(h) + i = -i; j = -j; l = -l; + conj = true; + } + i %= g.nu; if (i < 0) i += g.nu; + j %= g.nv; if (j < 0) j += g.nv; + float2 v = g.data[i + static_cast(g.nu) * (j + static_cast(g.nv) * l)]; + if (conj) + v.y = -v.y; + return v; +} + +// MolecularTransform::At on the device. +__device__ float2 transform_interpolate(const Transform &g, const Rot &r, float3 s0) { + const float s[3] = {r.m[0] * s0.x + r.m[1] * s0.y + r.m[2] * s0.z, + r.m[3] * s0.x + r.m[4] * s0.y + r.m[5] * s0.z, + r.m[6] * s0.x + r.m[7] * s0.y + r.m[8] * s0.z}; + float f[3]; + int i0[3]; + for (int a = 0; a < 3; a++) { + const float u = s[a] * g.box; + i0[a] = static_cast(floorf(u)); + f[a] = u - i0[a]; + } + float2 out = make_float2(0, 0); + for (int dx = 0; dx < 2; dx++) + for (int dy = 0; dy < 2; dy++) + for (int dz = 0; dz < 2; dz++) { + const float w = (dx ? f[0] : 1 - f[0]) * (dy ? f[1] : 1 - f[1]) * (dz ? f[2] : 1 - f[2]); + const float2 v = transform_at(g, i0[0] + dx, i0[1] + dy, i0[2] + dz); + out.x += w * v.x; + out.y += w * v.y; + } + const float unblur = expf(g.blur * (s[0] * s[0] + s[1] * s[1] + s[2] * s[2]) / 4.0f); + return make_float2(out.x * unblur, out.y * unblur); +} + +// A block-wide sum in a fixed order: each thread's value into shared memory, then a tree. +__device__ double block_sum(double v, double *sh) { + sh[threadIdx.x] = v; + __syncthreads(); + for (int stride = THREADS / 2; stride > 0; stride /= 2) { + if (threadIdx.x < stride) + sh[threadIdx.x] += sh[threadIdx.x + stride]; + __syncthreads(); + } + const double out = sh[0]; + __syncthreads(); + return out; +} + +// One block per orientation: |F|^2 summed over the copies for every reflection, binned, normalised by +// its bin's mean, and correlated with E^2 - 1. +__global__ void rotation_function(Transform g, const Rot *rot, const float3 *s, const float *e2, const int *bin, + const int *bin_n, int n, int nops, int bins, double w_mean, double w_ss, + float *scratch, double *out) { + __shared__ double sh[THREADS]; + __shared__ double bin_mean[MAX_BINS]; + const Rot r = rot[blockIdx.x]; + float *ic = scratch + static_cast(blockIdx.x) * n; + double part[MAX_BINS]; + for (int b = 0; b < bins; b++) + part[b] = 0; + for (int i = threadIdx.x; i < n; i += THREADS) { + float sum = 0; + for (int j = 0; j < nops; j++) { + const float2 f = transform_interpolate(g, r, s[static_cast(i) * nops + j]); + sum += f.x * f.x + f.y * f.y; + } + ic[i] = sum; + part[bin[i]] += sum; + } + for (int b = 0; b < bins; b++) { + const double total = block_sum(part[b], sh); + if (threadIdx.x == 0) + bin_mean[b] = total / bin_n[b]; + } + __syncthreads(); + double x_sum = 0, x_ss = 0, xy = 0; + for (int i = threadIdx.x; i < n; i += THREADS) { + const double x = ic[i] / bin_mean[bin[i]]; + x_sum += x; + x_ss += x * x; + xy += x * (e2[i] - w_mean); + } + x_sum = block_sum(x_sum, sh); + x_ss = block_sum(x_ss, sh); + xy = block_sum(xy, sh); + if (threadIdx.x == 0) { + const double sxx = x_ss - x_sum * x_sum / n; + out[blockIdx.x] = sxx > 0 && w_ss > 0 ? xy / sqrt(sxx * w_ss) : 0.0; + } +} + +// g_j = G_j(h) exp(2 pi i h.t_j) for every term and operator, and each term's sum of |g_j|^2. +__global__ void translation_copies(Transform g, Rot r, const float3 *s, const float2 *phase, int terms, int nops, + float2 *copies, float *term_sum) { + const int t = blockIdx.x * blockDim.x + threadIdx.x; + if (t >= terms) + return; + float sum = 0; + for (int j = 0; j < nops; j++) { + const size_t e = static_cast(t) * nops + j; + const float2 f = transform_interpolate(g, r, s[e]); + const float2 p = phase[e]; + const float2 v = make_float2(f.x * p.x - f.y * p.y, f.x * p.y + f.y * p.x); + copies[e] = v; + sum += v.x * v.x + v.y * v.y; + } + term_sum[t] = sum; +} + +// Each term's weight over its bin's mean |g|^2 sum, the bins summed in a fixed order by one block. +__global__ void translation_weights(const float *term_sum, const float *weight, const int *bin, const int *bin_n, + int terms, int bins, float *scaled) { + __shared__ double sh[THREADS]; + __shared__ double bin_mean[MAX_BINS]; + double part[MAX_BINS]; + for (int b = 0; b < bins; b++) + part[b] = 0; + for (int t = threadIdx.x; t < terms; t += THREADS) + part[bin[t]] += term_sum[t]; + for (int b = 0; b < bins; b++) { + const double total = block_sum(part[b], sh); + if (threadIdx.x == 0) + bin_mean[b] = total / bin_n[b]; + } + __syncthreads(); + for (int t = threadIdx.x; t < terms; t += THREADS) + scaled[t] = static_cast(weight[t] / bin_mean[bin[t]]); +} + +// Every coefficient of the series, gathered: a grid point sums the (term, pair) entries that land on it, +// in the order they were listed. +__global__ void translation_gather(const int *offset, const int *entry, const float2 *copies, const float *scaled, + const int *pair_j, const int *pair_k, int nops, int npairs, size_t cells, + float2 *coef) { + const size_t c = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (c >= cells) + return; + float2 sum = make_float2(0, 0); + for (int e = offset[c]; e < offset[c + 1]; e++) { + const int t = entry[e] / npairs, p = entry[e] % npairs; + const float2 a = copies[static_cast(t) * nops + pair_j[p]]; + const float2 b = copies[static_cast(t) * nops + pair_k[p]]; + const float w = scaled[t]; + sum.x += w * (a.x * b.x + a.y * b.y); // a conj(b) + sum.y += w * (a.y * b.x - a.x * b.y); + } + coef[c] = sum; +} + +// Per block: the sum and the sum of squares of the function, and its highest point (the lowest index +// on a tie). +__global__ void translation_stats(const float *map, size_t n, double *sum, double *sum2, float *best, + long long *best_index) { + __shared__ double sh[THREADS]; + __shared__ float sh_best[THREADS]; + __shared__ long long sh_index[THREADS]; + double s = 0, s2 = 0; + float b = -INFINITY; + long long bi = -1; + for (size_t i = static_cast(blockIdx.x) * THREADS + threadIdx.x; i < n; + i += static_cast(gridDim.x) * THREADS) { + const float v = map[i]; + s += v; + s2 += static_cast(v) * v; + if (v > b) { + b = v; + bi = static_cast(i); + } + } + s = block_sum(s, sh); + s2 = block_sum(s2, sh); + sh_best[threadIdx.x] = b; + sh_index[threadIdx.x] = bi; + __syncthreads(); + for (int stride = THREADS / 2; stride > 0; stride /= 2) { + if (threadIdx.x < stride) { + const float o = sh_best[threadIdx.x + stride]; + const long long oi = sh_index[threadIdx.x + stride]; + if (o > sh_best[threadIdx.x] || (o == sh_best[threadIdx.x] && oi >= 0 && oi < sh_index[threadIdx.x])) { + sh_best[threadIdx.x] = o; + sh_index[threadIdx.x] = oi; + } + } + __syncthreads(); + } + if (threadIdx.x == 0) { + sum[blockIdx.x] = s; + sum2[blockIdx.x] = s2; + best[blockIdx.x] = sh_best[0]; + best_index[blockIdx.x] = sh_index[0]; + } +} + +Rot transposed(const gemmi::Mat33 &m) { + Rot r; + for (int i = 0; i < 3; i++) + for (int j = 0; j < 3; j++) + r.m[3 * i + j] = static_cast(m.a[j][i]); + return r; +} + +template +CudaDevicePtr upload(const std::vector &v, cudaStream_t stream) { + CudaDevicePtr d(std::max(v.size(), 1)); + if (!v.empty()) + cuda_err(cudaMemcpyAsync(d.get(), v.data(), v.size() * sizeof(T), cudaMemcpyHostToDevice, stream)); + return d; +} + +std::vector float3s(const std::vector &v) { + std::vector out(v.size()); + for (size_t i = 0; i < v.size(); i++) + out[i] = make_float3(static_cast(v[i].x), static_cast(v[i].y), static_cast(v[i].z)); + return out; +} + +std::vector bin_counts(const std::vector &bin, int bins) { + std::vector n(bins, 0); + for (int b : bin) + n[b]++; + return n; +} + +} // namespace + +struct MolecularReplacementGPUState { + CudaStream stream; + CudaDevicePtr grid; + Transform transform{}; +}; + +MolecularReplacementGPU::MolecularReplacementGPU(const MolecularTransform &ft) + : state_(std::make_unique()) { + const gemmi::FPhiGrid &g = ft.Grid(); + std::vector host(g.data.size()); + for (size_t i = 0; i < host.size(); i++) + host[i] = make_float2(g.data[i].real(), g.data[i].imag()); + state_->grid = upload(host, state_->stream); + state_->transform = {state_->grid.get(), g.nu, g.nv, g.nw, static_cast(ft.Box()), + static_cast(ft.Blur())}; + cuda_err(cudaStreamSynchronize(state_->stream)); +} + +MolecularReplacementGPU::~MolecularReplacementGPU() = default; + +std::vector MolecularReplacementGPU::RotationFunction(const SearchReflections &refl, + const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg, + const std::vector &rotations) { + cudaStream_t stream = state_->stream; + const int n = static_cast(refl.hkl.size()); + const std::vector s = CopyVectors(refl, cell, sg); + const int nops = static_cast(s.size() / std::max(n, 1)); + double w_mean = 0; + for (double e2 : refl.e2) + w_mean += e2; + w_mean /= n; + double w_ss = 0; + for (double e2 : refl.e2) + w_ss += (e2 - w_mean) * (e2 - w_mean); + const std::vector e2(refl.e2.begin(), refl.e2.end()); + auto d_s = upload(float3s(s), stream); + auto d_e2 = upload(e2, stream); + auto d_bin = upload(refl.bin, stream); + auto d_bin_n = upload(bin_counts(refl.bin, refl.bins), stream); + std::vector rot(rotations.size()); + for (size_t i = 0; i < rotations.size(); i++) + rot[i] = transposed(rotations[i]); + auto d_rot = upload(rot, stream); + const size_t batch = std::min(ROTATION_BATCH, std::max(rotations.size(), 1)); + CudaDevicePtr scratch(batch * std::max(n, 1)); + CudaDevicePtr d_out(std::max(rotations.size(), 1)); + for (size_t first = 0; first < rotations.size(); first += batch) { + const int count = static_cast(std::min(batch, rotations.size() - first)); + rotation_function<<>>(state_->transform, d_rot.get() + first, d_s, d_e2, d_bin, + d_bin_n, n, nops, refl.bins, w_mean, w_ss, scratch, + d_out.get() + first); + cuda_err(cudaGetLastError()); + } + std::vector out(rotations.size()); + cuda_err(cudaMemcpyAsync(out.data(), d_out.get(), out.size() * sizeof(double), cudaMemcpyDeviceToHost, stream)); + cuda_err(cudaStreamSynchronize(stream)); + return out; +} + +std::vector MolecularReplacementGPU::TranslationFunction(const TranslationSetup &setup, + const std::vector &rotations) { + std::vector out(rotations.size()); + const int terms = static_cast(setup.Terms()), nops = setup.ops; + const int npairs = static_cast(setup.pair_j.size()); + if (nops < 2 || terms == 0 || rotations.empty()) + return out; + cudaStream_t stream = state_->stream; + const int nu = setup.nu, nv = setup.nv, nw = setup.nw, nwh = nw / 2 + 1; + const size_t cells = static_cast(nu) * nv * nwh, points = static_cast(nu) * nv * nw; + + // The gather's lists: for every point of the half grid in cuFFT's layout ((u nv + v) nwh + l), the + // (term, pair) entries landing on it, in entry order. + std::vector offset(cells + 1, 0); + for (size_t e = 0; e < setup.target_u.size(); e++) + if (setup.target_u[e] >= 0) + offset[(static_cast(setup.target_u[e]) * nv + setup.target_v[e]) * nwh + setup.target_l[e] + 1]++; + for (size_t c = 0; c < cells; c++) + offset[c + 1] += offset[c]; + std::vector entry(offset[cells]); + { + std::vector fill(offset.begin(), offset.end() - 1); + for (size_t e = 0; e < setup.target_u.size(); e++) + if (setup.target_u[e] >= 0) + entry[fill[(static_cast(setup.target_u[e]) * nv + setup.target_v[e]) * nwh + + setup.target_l[e]]++] = static_cast(e); + } + std::vector phase(setup.phase.size()); + for (size_t i = 0; i < phase.size(); i++) + phase[i] = make_float2(static_cast(setup.phase[i].real()), static_cast(setup.phase[i].imag())); + const std::vector weight(setup.weight.begin(), setup.weight.end()); + auto d_s = upload(float3s(setup.s), stream); + auto d_phase = upload(phase, stream); + auto d_weight = upload(weight, stream); + auto d_bin = upload(setup.bin, stream); + auto d_bin_n = upload(bin_counts(setup.bin, setup.bins), stream); + auto d_pair_j = upload(setup.pair_j, stream); + auto d_pair_k = upload(setup.pair_k, stream); + auto d_offset = upload(offset, stream); + auto d_entry = upload(entry, stream); + CudaDevicePtr copies(static_cast(terms) * nops); + CudaDevicePtr term_sum(terms), scaled(terms); + CudaDevicePtr coef(cells); + CudaDevicePtr map(points); + const int stat_blocks = static_cast(std::min(1024, (points + THREADS - 1) / THREADS)); + const size_t nstat = static_cast(stat_blocks) * rotations.size(); + CudaDevicePtr d_sum(nstat), d_sum2(nstat); + CudaDevicePtr d_best(nstat); + CudaDevicePtr d_best_index(nstat); + + cufftHandle plan; + cufft_err(cufftPlan3d(&plan, nu, nv, nw, CUFFT_C2R)); + try { + cufft_err(cufftSetStream(plan, stream)); + for (size_t r = 0; r < rotations.size(); r++) { + const Rot rot = transposed(rotations[r]); + translation_copies<<<(terms + THREADS - 1) / THREADS, THREADS, 0, stream>>>( + state_->transform, rot, d_s, d_phase, terms, nops, copies, term_sum); + translation_weights<<<1, THREADS, 0, stream>>>(term_sum, d_weight, d_bin, d_bin_n, terms, setup.bins, + scaled); + translation_gather<<((cells + THREADS - 1) / THREADS), THREADS, 0, stream>>>( + d_offset, d_entry, copies, scaled, d_pair_j, d_pair_k, nops, npairs, cells, coef); + cufft_err(cufftExecC2R(plan, reinterpret_cast(coef.get()), map.get())); + translation_stats<<>>( + map, points, d_sum.get() + r * stat_blocks, d_sum2.get() + r * stat_blocks, + d_best.get() + r * stat_blocks, d_best_index.get() + r * stat_blocks); + cuda_err(cudaGetLastError()); + } + } catch (...) { + cufftDestroy(plan); + throw; + } + std::vector sum(nstat), sum2(nstat); + std::vector best(nstat); + std::vector best_index(nstat); + cuda_err(cudaMemcpyAsync(sum.data(), d_sum.get(), nstat * sizeof(double), cudaMemcpyDeviceToHost, stream)); + cuda_err(cudaMemcpyAsync(sum2.data(), d_sum2.get(), nstat * sizeof(double), cudaMemcpyDeviceToHost, stream)); + cuda_err(cudaMemcpyAsync(best.data(), d_best.get(), nstat * sizeof(float), cudaMemcpyDeviceToHost, stream)); + cuda_err(cudaMemcpyAsync(best_index.data(), d_best_index.get(), nstat * sizeof(long long), + cudaMemcpyDeviceToHost, stream)); + cuda_err(cudaStreamSynchronize(stream)); + cufftDestroy(plan); + + for (size_t r = 0; r < rotations.size(); r++) { + double s = 0, s2 = 0; + float b = -INFINITY; + long long bi = 0; + for (int k = 0; k < stat_blocks; k++) { + const size_t i = r * stat_blocks + k; + s += sum[i]; + s2 += sum2[i]; + if (best[i] > b) { + b = best[i]; + bi = best_index[i]; + } + } + const double count = static_cast(points); + const double mean = s / count, sd = std::sqrt(std::max(0.0, s2 / count - mean * mean)); + out[r].z = sd > 0 ? (b - mean) / sd : 0.0; + const long long w = bi % nw, v = (bi / nw) % nv, u = bi / (static_cast(nw) * nv); + out[r].t = gemmi::Fractional(static_cast(u) / nu, static_cast(v) / nv, + static_cast(w) / nw); + } + return out; +} diff --git a/image_analysis/structure_refinement/MolecularReplacementGPU.h b/image_analysis/structure_refinement/MolecularReplacementGPU.h new file mode 100644 index 000000000..38f0ebdbf --- /dev/null +++ b/image_analysis/structure_refinement/MolecularReplacementGPU.h @@ -0,0 +1,37 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#pragma once + +// The molecular-replacement search on a GPU (CUDA builds only; the header itself needs no CUDA): the +// rotation function and the translation function of MolecularReplacement.h, the same functions as the +// CPU computes - the same trilinear interpolation of the same molecular transform, the same bins, the +// same coefficients on the same grid - in float on the device, with cuFFT for FFTW. The two agree to +// rounding, not bit for bit. Every reduction runs in a fixed order and nothing is summed with an +// atomic, so the GPU is deterministic on its own. + +#include +#include + +#include "MolecularReplacement.h" + +struct MolecularReplacementGPUState; +struct MolecularReplacementGPUGroup; + +class MolecularReplacementGPU { +public: + // Uploads the molecular transform, on the calling thread's current device. Throws JFJochException on + // a CUDA failure, as every call below does. + explicit MolecularReplacementGPU(const MolecularTransform &ft); + ~MolecularReplacementGPU(); + + std::vector RotationFunction(const SearchReflections &refl, const gemmi::UnitCell &cell, + const gemmi::SpaceGroup &sg, const std::vector &rotations); + + // The translation function of one group at each of `rotations`. + std::vector TranslationFunction(const TranslationSetup &setup, + const std::vector &rotations); + +private: + std::unique_ptr state_; +}; diff --git a/rugnux/ResultReport.cpp b/rugnux/ResultReport.cpp index f1f8f32ad..b822754c6 100644 --- a/rugnux/ResultReport.cpp +++ b/rugnux/ResultReport.cpp @@ -2580,10 +2580,24 @@ ReportDocument BuildReportDocument(const std::string &output_prefix, Add(s, KeyText("MODEL_PRIOR_DECISION", mp.decision)); if (mp.chosen >= 0) Add(s, KeyText("MODEL_PRIOR_SPACE_GROUP", mp.candidates[mp.chosen].space_group.xhm())); + if (mp.hand_check.performed) { + const auto &h = mp.hand_check; + Add(s, KeyEnum("MODEL_PRIOR_HAND_CHECK", h.adopted ? "ADOPTED" : "VETOED", {"ADOPTED", "VETOED"})); + Add(s, KeyReal("MODEL_PRIOR_HAND_FIT_SIGMA", h.fit_sigma, "{:+.2f}")); + Add(s, KeyBool("MODEL_PRIOR_ANOMALOUS_INVERTED", h.anomalous_inverted)); + Add(s, KeyInt("MODEL_PRIOR_ANOMALOUS_PAIRS", h.anomalous_pairs)); + if (std::isfinite(h.anomalous_scatterer_mean_sigma)) + Add(s, KeyReal("MODEL_PRIOR_ANOMALOUS_SCATTERER_MEAN_SIGMA", h.anomalous_scatterer_mean_sigma, + "{:+.2f}")); + if (h.anomalous_inverted) + Add(s, KeyText("MODEL_PRIOR_ANOMALOUS_DEEPEST_SITE", + fmt::format("{} {:.1f}", h.deepest_site, h.deepest_sigma))); + } if (!result.model_prior_file_prefix.empty()) Add(s, KeyText("MODEL_PRIOR_FILE_PREFIX", result.model_prior_file_prefix)); + Add(s, KeyEnum("MODEL_PRIOR_SEARCH_DEVICE", mp.gpu_search ? "GPU" : "CPU", {"GPU", "CPU"})); Add(s, KeyReal("MODEL_PRIOR_SECONDS", mp.t_model + mp.t_transform + mp.t_rotation + mp.t_translation - + mp.t_rescore + mp.t_null, "{:.2f}")); + + mp.t_rescore + mp.t_null + mp.hand_check.seconds, "{:.2f}")); Add(s, Blank()); ReportEntry e; e.kind = ReportEntry::Kind::Table; @@ -2620,19 +2634,22 @@ ReportDocument BuildReportDocument(const std::string &output_prefix, " {:.1f}-{:.1f} A), and the best placements taken through the rigid body; R_WORK and R_FREE are\n" " to SCORE_D_MIN. PLACED needs TF_Z >= 8,\n" " at most 5% of the C-alpha atoms clashing with symmetry mates, one copy fitting the cell\n" - " (MATTHEWS: Vm below 1.5; MULTI_COPY: Vm of 4 or more, not supported yet) and an R-work at\n" - " least 3 sd better than the same model refitted from random orientations. The model chooses\n" + " (MATTHEWS: Vm below 1.5; MULTI_COPY: Vm of 4 or more, not supported yet) and, for the\n" + " candidate the choice falls on, an R-work at least 3 sd better than the same model refitted\n" + " from random orientations (NULL where it is not; the next candidate is then taken). The model chooses\n" " within a merge where its placed group leads every other that packs by 0.03 in R-free, and\n" " the twin hypothesis' merge - searched only then - where no group of the answer's merge can\n" " hold it at all. A group one copy cannot fill is checked for packing only, not refined.\n" - " The answer's files take the enantiomorph from it and nothing else; another group it\n" + " The answer's files take the enantiomorph from it - only where model validation of the placed\n" + " model accepts it as it accepts a --model's: a fit over its null and an anomalous map at its\n" + " atoms that is not inverted (ENANTIOMORPH_VETOED otherwise) - and nothing else; another group it\n" " chooses is written as {}.mtz/.cif with the placed model.", mp.d_max_search, mp.d_min_search, result.model_prior_file_prefix.empty() ? std::string("_mr") : result.model_prior_file_prefix))); Add(s, Prose(fmt::format( " Seconds: model {:.2f}, transform {:.2f}, rotation {:.2f}, translation {:.2f}, rigid body and\n" - " R-free {:.2f}, null {:.2f}.", mp.t_model, mp.t_transform, mp.t_rotation, mp.t_translation, - mp.t_rescore, mp.t_null), true)); + " R-free {:.2f}, null {:.2f}, hand check {:.2f}.", mp.t_model, mp.t_transform, mp.t_rotation, + mp.t_translation, mp.t_rescore, mp.t_null, mp.hand_check.seconds), true)); } doc.sections.push_back(std::move(s)); } diff --git a/rugnux/RugnuxScaleMerge.cpp b/rugnux/RugnuxScaleMerge.cpp index e0f073673..bb1f623be 100644 --- a/rugnux/RugnuxScaleMerge.cpp +++ b/rugnux/RugnuxScaleMerge.cpp @@ -20,6 +20,7 @@ #include #include #include +#include #include #include #include @@ -4000,6 +4001,46 @@ bool Rugnux::ScaleMergeAndSymmetry(PipelineLocals &p) { "The predicted model was not placed: " + mr.failure_reason}); if (mr.chosen >= 0) { const MRCandidate &c = mr.candidates[mr.chosen]; + if (mr.decision == "ENANTIOMORPH") { + // Taken into the answer as --model takes a model's enantiomorph: model validation runs on + // the placed model, in the group it chose, against the answer's merge, and its verdict + // decides - the model has to fit against its own null, and the anomalous map at its atoms + // must not come out inverted. The validation's files go to a scratch directory, removed after. + const auto t_hand = std::chrono::steady_clock::now(); + const std::filesystem::path scratch = config_.output_prefix + "_mr_hand"; + std::filesystem::create_directories(scratch); + const std::string tmp = (scratch / "m").string(); + Logger quiet = Logger::Buffered(); + WritePlacedModel(*c.placed, *result.consensus_cell, c.space_group, tmp, quiet); + const auto data_sg = experiment_.GetGemmiSpaceGroup(); + const ModelValidationResult v = ValidateAgainstModel( + sm.merged, *result.consensus_cell, tmp + "_model.cif", tmp, quiet, + data_sg ? &*data_sg : nullptr, /*probe_indexing_ambiguity=*/false, + static_cast(config_.nthreads), experiment_.GetWavelength_A()); + std::filesystem::remove_all(scratch); + auto &h = mr.hand_check; + h.performed = true; + h.adopted = v.adopted_model_enantiomorph; + h.model_fits = v.model_fits; + h.fit_sigma = v.r_work_sigma; + h.anomalous_inverted = v.anomalous_hands_disagree; + h.anomalous_pairs = v.anomalous_pairs; + h.anomalous_scatterer_mean_sigma = v.anomalous_scatterer_mean_sigma; + h.deepest_site = v.anomalous_deepest_site; + h.deepest_sigma = v.anomalous_deepest_sigma; + h.seconds = std::chrono::duration(std::chrono::steady_clock::now() - t_hand).count(); + if (!h.adopted) { + mr.decision = "ENANTIOMORPH_VETOED"; + logger.Info("Model prior: the enantiomorph {} is not taken into the answer - {} ({:.1f} s)", + c.space_group.xhm(), + h.anomalous_inverted ? fmt::format("the anomalous density at the placed model's " + "atoms is inverted ({} at {:.1f} sigma)", + h.deepest_site, h.deepest_sigma) + : fmt::format("the placed model does not fit against its " + "null ({:+.1f} sd)", h.fit_sigma), + h.seconds); + } + } if (mr.decision == "ENANTIOMORPH") { // A change of label only (see ValidateAgainstModel): the two groups index identically. experiment_.SetSpaceGroup(c.space_group); @@ -4010,8 +4051,9 @@ bool Rugnux::ScaleMergeAndSymmetry(PipelineLocals &p) { if (const gemmi::SpaceGroup *other = gemmi::find_spacegroup_by_ops(e)) twin_file->space_group = *other; } - logger.Info("Model prior: the answer is relabelled {}, the enantiomorph the placed model chose", - c.space_group.xhm()); + logger.Info("Model prior: the answer is relabelled {}, the enantiomorph the placed model chose " + "(fit {:+.1f} sd over its null, anomalous map at its atoms not inverted, {:.1f} s)", + c.space_group.xhm(), mr.hand_check.fit_sigma, mr.hand_check.seconds); } else if (mr.decision == "SPACE_GROUP") { ModelPriorFile f{sm.merged, *result.consensus_cell, c.space_group, std::nullopt, {}, c.placed}; const gemmi::GroupOps gops = c.space_group.operations(); diff --git a/tests/MolecularReplacementTest.cpp b/tests/MolecularReplacementTest.cpp index 4ca9d27c3..827df1b03 100644 --- a/tests/MolecularReplacementTest.cpp +++ b/tests/MolecularReplacementTest.cpp @@ -11,6 +11,10 @@ #include "../image_analysis/structure_refinement/ModelFFT.h" #include "../image_analysis/structure_refinement/MolecularReplacement.h" +#ifdef JFJOCH_USE_CUDA +#include "../common/CUDAWrapper.h" +#include "../image_analysis/structure_refinement/MolecularReplacementGPU.h" +#endif namespace { // A synthetic body: carbon atoms scattered in a 12 A sphere, centroid at the origin. @@ -156,7 +160,7 @@ TEST_CASE("MolecularReplacement_SyntheticPlacement", "[MolecularReplacement]") { // At the true orientation the translation function peaks at the true position, up to the // origins P 21 21 21 allows (a half in each coordinate), and stands far out of its own noise. - const TranslationPeak peak = TranslationFunction(ft, refl, cell, sg, truth); + const TranslationPeak peak = TranslationFunction(ft, MakeTranslationSetup(refl, cell, sg), truth); CHECK(peak.z > 10.0); for (double d : {peak.t.x - t.x, peak.t.y - t.y, peak.t.z - t.z}) { const double half = 2 * d - std::round(2 * d); @@ -164,3 +168,39 @@ TEST_CASE("MolecularReplacement_SyntheticPlacement", "[MolecularReplacement]") { } } + +#ifdef JFJOCH_USE_CUDA +TEST_CASE("MolecularReplacement_GPUMatchesCPU", "[MolecularReplacement]") { + if (get_gpu_count() == 0) + SKIP("no GPU"); + const SearchModel model = synthetic_model(); + const gemmi::UnitCell cell(42, 50, 58, 90, 90, 90); + const gemmi::SpaceGroup &sg = gemmi::get_spacegroup_by_name("P 21 21 21"); + const gemmi::Mat33 truth = axis_angle({1, 2, 3}, 50); + const gemmi::Fractional t(0.17, 0.31, 0.09); + const auto merged = observed(PlaceSearchModel(model, truth, t, cell, sg)); + const SearchReflections refl = MakeSearchReflections(merged, cell, sg, 4.0, 12.0); + const MolecularTransform ft(model, 4.0, 4); + MolecularReplacementGPU gpu(ft); + + const auto rotations = RotationAsu(UniformRotations(15.0), CartesianPointGroup(sg, cell)); + const auto rf_cpu = RotationFunction(ft, refl, cell, sg, rotations, 4); + const auto rf_gpu = gpu.RotationFunction(refl, cell, sg, rotations); + REQUIRE(rf_gpu.size() == rf_cpu.size()); + for (size_t i = 0; i < rf_cpu.size(); i++) + CHECK(std::fabs(rf_gpu[i] - rf_cpu[i]) < 1e-4); + + const TranslationSetup setup = MakeTranslationSetup(refl, cell, sg); + std::vector orientations{truth, rotations[3], rotations[17]}; + const auto tf_gpu = gpu.TranslationFunction(setup, orientations); + for (size_t i = 0; i < orientations.size(); i++) { + const TranslationPeak cpu = TranslationFunction(ft, setup, orientations[i]); + CHECK(std::fabs(tf_gpu[i].z - cpu.z) < 1e-3); + if (i == 0) { // the solution's peak is unique; noise peaks can tie between grid points + CHECK(tf_gpu[i].t.x == Catch::Approx(cpu.t.x)); + CHECK(tf_gpu[i].t.y == Catch::Approx(cpu.t.y)); + CHECK(tf_gpu[i].t.z == Catch::Approx(cpu.t.z)); + } + } +} +#endif