diff --git a/image_analysis/structure_refinement/ModelFFT.cpp b/image_analysis/structure_refinement/ModelFFT.cpp index 54830fa86..aee3be8eb 100644 --- a/image_analysis/structure_refinement/ModelFFT.cpp +++ b/image_analysis/structure_refinement/ModelFFT.cpp @@ -78,6 +78,11 @@ gemmi::FPhiGrid MapToFPhi(const gemmi::Grid &map) { return hkl; } +void MapFromFPhiExecute(int nu, int nv, int nw, std::complex *in, float *out) { + static_assert(sizeof(std::complex) == sizeof(fftwf_complex)); + fftwf_execute_dft_c2r(PlanFor(nu, nv, nw, true), reinterpret_cast(in), out); +} + gemmi::Grid MapFromFPhi(const gemmi::FPhiGrid &hkl) { if (hkl.axis_order == gemmi::AxisOrder::ZYX || !hkl.half_l) gemmi::fail("MapFromFPhi(): only an XYZ half-l grid is supported"); diff --git a/image_analysis/structure_refinement/ModelFFT.h b/image_analysis/structure_refinement/ModelFFT.h index 88a83fdb5..3ebb0754f 100644 --- a/image_analysis/structure_refinement/ModelFFT.h +++ b/image_analysis/structure_refinement/ModelFFT.h @@ -23,6 +23,11 @@ gemmi::FPhiGrid MapToFPhi(const gemmi::Grid &map); // under the same planning rules. gemmi::Grid MapFromFPhi(const gemmi::FPhiGrid &hkl); +// MapFromFPhi's transform itself, for a caller that fills its input: `in`, the (nu, nv, nw/2+1) half-l coefficients +// u fastest as MapFromFPhi hands them to FFTW (conjugated, a NaN as zero), is overwritten, and `out` receives the +// (nu, nv, nw) map before MapFromFPhi's 1/V. Both from fftwf_malloc, as a plan requires. +void MapFromFPhiExecute(int nu, int nv, int nw, std::complex *in, float *out); + // The half-l grid `full` with only the indices |k|, l <= reach kept, on a grid of full.nu x (2 reach + 1) x // (reach + 1) that an index wraps into as it wraps into the whole one (the molecular transform's, which is // read no further out). diff --git a/image_analysis/structure_refinement/MolecularReplacement.cpp b/image_analysis/structure_refinement/MolecularReplacement.cpp index c04f07bbf..dee1aa224 100644 --- a/image_analysis/structure_refinement/MolecularReplacement.cpp +++ b/image_analysis/structure_refinement/MolecularReplacement.cpp @@ -20,6 +20,7 @@ #include #include +#include #include #include // read_structure_gz @@ -1433,8 +1434,48 @@ std::vector> KnownStructureFactors(const MolecularTransform return out; } -TranslationPeak TranslationFunction(const MolecularTransform &ft, const TranslationSetup &setup, - const gemmi::Mat33 &rotation, const std::vector> *known) { +namespace { + +// Where each coefficient of a translation function goes: for every point of the half grid, (l nv + v) nu + u as +// MapFromFPhi takes it, the entries landing on it, in entry order - the (term, pair) entries of setup.target_*, +// then the cross terms with a known partial structure. +struct TranslationGather { + std::vector offset, entry, cross_offset, cross_entry; +}; + +void gather_lists(size_t cells, size_t n, const auto &cell_of, std::vector &offset, std::vector &entry) { + offset.assign(cells + 1, 0); + for (size_t e = 0; e < n; e++) + if (const long c = cell_of(e); c >= 0) + offset[c + 1]++; + for (size_t c = 0; c < cells; c++) + offset[c + 1] += offset[c]; + entry.resize(offset[cells]); + std::vector fill(offset.begin(), offset.end() - 1); + for (size_t e = 0; e < n; e++) + if (const long c = cell_of(e); c >= 0) + entry[fill[c]++] = static_cast(e); +} + +TranslationGather make_gather(const TranslationSetup &setup, bool known) { + TranslationGather out; + const size_t nu = setup.nu, nv = setup.nv, cells = nu * nv * (setup.nw / 2 + 1); + gather_lists(cells, setup.target_u.size(), [&](size_t e) -> long { + return setup.target_u[e] < 0 ? -1 : setup.target_u[e] + nu * (setup.target_v[e] + nv * setup.target_l[e]); + }, out.offset, out.entry); + if (known) + gather_lists(cells, setup.cross_term.size(), [&](size_t e) -> long { + return setup.cross_u[e] + nu * (setup.cross_v[e] + nv * setup.cross_l[e]); + }, out.cross_offset, out.cross_entry); + return out; +} + +// The function at one orientation. Each point's coefficient is summed over its entries in the order the terms +// and the pairs come, the cross terms after, as the series is written out term by term, and goes straight into +// the transform's input conjugated, as MapFromFPhi takes it. +TranslationPeak translation_function(const MolecularTransform &ft, const TranslationSetup &setup, + const TranslationGather &gather, const gemmi::Mat33 &rotation, + const std::vector> *known) { TranslationPeak out; const size_t terms = setup.Terms(), nops = setup.ops, npairs = setup.pair_j.size(); if (terms == 0 || (nops < 2 && known == nullptr)) @@ -1452,51 +1493,71 @@ TranslationPeak TranslationFunction(const MolecularTransform &ft, const Translat 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]]); - } + std::vector w(terms); + for (size_t t = 0; t < terms; t++) + w[t] = setup.weight[t] / (bin_sum[setup.bin[t]] / bin_n[setup.bin[t]]); + const int nu = setup.nu, nv = setup.nv, nw = setup.nw; + const size_t cells = static_cast(nu) * nv * (nw / 2 + 1), points = static_cast(nu) * nv * nw; + auto *in = reinterpret_cast *>(fftwf_alloc_complex(cells)); + float *map = fftwf_alloc_real(points); + if (in == nullptr || map == nullptr) { + fftwf_free(in); + fftwf_free(map); + throw std::runtime_error("TranslationFunction: cannot allocate the FFT buffers"); } - if (known != nullptr) - for (size_t e = 0; e < setup.cross_term.size(); e++) { - const size_t t = setup.cross_term[e]; - const double w = setup.weight[t] / (bin_sum[setup.bin[t]] / bin_n[setup.bin[t]]); - const std::complex v = std::conj((*known)[t]) * g[t * nops + setup.cross_op[e]]; - c[setup.cross_u[e] + static_cast(nu) * (setup.cross_v[e] + static_cast(nv) * setup.cross_l[e])] - += w * (setup.cross_conj[e] ? std::conj(v) : v); + for (size_t c = 0; c < cells; c++) { + std::complex v = 0.0; + for (int k = gather.offset[c]; k < gather.offset[c + 1]; k++) { + const size_t e = gather.entry[k], t = e / npairs, p = e % npairs; + v += w[t] * 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 = ft.Grid().unit_cell; - hkl.spacegroup = &gemmi::get_spacegroup_p1(); - hkl.axis_order = gemmi::AxisOrder::XYZ; - hkl.half_l = true; - hkl.set_size_without_checking(nu, nv, nwh); - for (size_t i = 0; i < c.size(); i++) - hkl.data[i] = std::complex(std::conj(c[i])); - const gemmi::Grid map = MapFromFPhi(hkl); + if (known != nullptr) + for (int k = gather.cross_offset[c]; k < gather.cross_offset[c + 1]; k++) { + const size_t e = gather.cross_entry[k], t = setup.cross_term[e]; + const std::complex x = std::conj((*known)[t]) * g[t * nops + setup.cross_op[e]]; + v += w[t] * (setup.cross_conj[e] ? std::conj(x) : x); + } + // MapFromFPhi sums F exp(-2 pi i m.x): with F = conj(C) that is the conjugate of the series, which + // is real. + const std::complex f(std::conj(v)); + in[c] = std::isnan(f.imag()) ? std::complex(0.0f, 0.0f) : std::complex(f.real(), -f.imag()); + } + MapFromFPhiExecute(nu, nv, nw, in, map); + const float norm = static_cast(1.0 / ft.Grid().unit_cell.volume); double sum = 0, sum2 = 0; size_t best = 0; - for (size_t i = 0; i < map.data.size(); i++) { - sum += map.data[i]; - sum2 += static_cast(map.data[i]) * map.data[i]; - if (map.data[i] > map.data[best]) + for (size_t i = 0; i < points; i++) { + const float x = map[i] * norm; + sum += x; + sum2 += static_cast(x) * x; + if (x > map[best] * norm) best = i; } - const double count = static_cast(map.data.size()); + const double count = static_cast(points); const double mean = sum / count, sd = std::sqrt(std::max(0.0, sum2 / count - mean * mean)); - out.z = sd > 0 ? (map.data[best] - mean) / sd : 0.0; - const size_t u = best % map.nu, v = (best / map.nu) % map.nv, w = best / (static_cast(map.nu) * map.nv); - out.t = gemmi::Fractional(static_cast(u) / map.nu, static_cast(v) / map.nv, - static_cast(w) / map.nw); + out.z = sd > 0 ? (map[best] * norm - mean) / sd : 0.0; + fftwf_free(in); + fftwf_free(map); + const size_t u = best % nu, v = (best / nu) % nv, l = best / (static_cast(nu) * nv); + out.t = gemmi::Fractional(static_cast(u) / nu, static_cast(v) / nv, static_cast(l) / nw); + return out; +} + +} // namespace + +TranslationPeak TranslationFunction(const MolecularTransform &ft, const TranslationSetup &setup, + const gemmi::Mat33 &rotation, const std::vector> *known) { + return translation_function(ft, setup, make_gather(setup, known != nullptr), rotation, known); +} + +std::vector TranslationFunction(const MolecularTransform &ft, const TranslationSetup &setup, + const std::vector &rotations, + const std::vector> *known, size_t nthreads) { + const TranslationGather gather = make_gather(setup, known != nullptr); + std::vector out(rotations.size()); + ParallelFor(static_cast(rotations.size()), nthreads, [&](int i) { + out[i] = translation_function(ft, setup, gather, rotations[i], known); + }); return out; } @@ -2411,11 +2472,7 @@ MolecularReplacementResult search_model(const SearchModel &model, const std::vec #else (void) prepared; #endif - std::vector tf(rotations.size()); - ParallelFor(static_cast(rotations.size()), nthreads, [&](int i) { - tf[i] = TranslationFunction(ft, setup, rotations[i], known); - }); - return tf; + return TranslationFunction(ft, setup, rotations, known, nthreads); }; std::vector orientations; // of each candidate's placement, before the rigid body diff --git a/image_analysis/structure_refinement/MolecularReplacement.h b/image_analysis/structure_refinement/MolecularReplacement.h index f887e3058..ea163349a 100644 --- a/image_analysis/structure_refinement/MolecularReplacement.h +++ b/image_analysis/structure_refinement/MolecularReplacement.h @@ -181,6 +181,10 @@ struct TranslationPeak { TranslationPeak TranslationFunction(const MolecularTransform &ft, const TranslationSetup &setup, const gemmi::Mat33 &rotation, const std::vector> *known = nullptr); +// The same at each of `rotations`, `nthreads` of them at once, the grid's gather worked out once for all. +std::vector TranslationFunction(const MolecularTransform &ft, const TranslationSetup &setup, + const std::vector &rotations, + const std::vector> *known, size_t nthreads); // --- The likelihood a placement is judged by --- // The terms it is summed over: the reflections of a SearchReflections not absent in a group, each with