Molecular replacement: the CPU translation function gathers its grid once per setup, straight into FFTW's input, output unchanged

On the CPU every orientation's translation function scattered its terms into a complex<double> grid,
converted that to an FPhiGrid, which MapFromFPhi copied into FFTW's input, transformed, and copied out
again into a map the statistics then read: 24 bytes a grid point, 1.6 GB an orientation on 4nwv's cell,
with the 2 x 32 orientations of the groups' functions in flight at once.

Now, once per setup and for all its orientations (a new TranslationFunction over a list of rotations, which
the search calls), each half-grid point's entries are listed in the order the terms and pairs come, cross
terms after; an orientation sums each point over its own list - the same additions in the same order as
the scatter - writes it conjugated into FFTW's input as MapFromFPhi would (MapFromFPhiExecute, the same
plan) and reads the statistics off FFTW's output scaled as MapFromFPhi scales it: 8 bytes a point.

CPU-only (--no-gpu), timing benchmark on 4nwv: the translation stage 61.5 -> 17.7 s, provisional (load
average 24-31 during both runs, and every other stage moved by as much again with no change). Reports,
MR candidate lines, reflection files, maps and placed models byte-identical on myob, cytc and 4nwv,
CPU-only and on the GPU. [MolecularReplacement] passes.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

image_analysis/structure_refinement/ModelFFT.cpp

@@ -78,6 +78,11 @@ gemmi::FPhiGrid<float> MapToFPhi(const gemmi::Grid<float> &map) {
     return hkl;
 }
 
+void MapFromFPhiExecute(int nu, int nv, int nw, std::complex<float> *in, float *out) {
+    static_assert(sizeof(std::complex<float>) == sizeof(fftwf_complex));
+    fftwf_execute_dft_c2r(PlanFor(nu, nv, nw, true), reinterpret_cast<fftwf_complex *>(in), out);
+}
+
 gemmi::Grid<float> MapFromFPhi(const gemmi::FPhiGrid<float> &hkl) {
     if (hkl.axis_order == gemmi::AxisOrder::ZYX || !hkl.half_l)
         gemmi::fail("MapFromFPhi(): only an XYZ half-l grid is supported");

image_analysis/structure_refinement/ModelFFT.h

@@ -23,6 +23,11 @@ gemmi::FPhiGrid<float> MapToFPhi(const gemmi::Grid<float> &map);
 // under the same planning rules.
 gemmi::Grid<float> MapFromFPhi(const gemmi::FPhiGrid<float> &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<float> *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).

image_analysis/structure_refinement/MolecularReplacement.cpp

@@ -20,6 +20,7 @@
 #include <unordered_map>
 #include <fstream>
 
+#include <fftw3.h>
 #include <nlohmann/json.hpp>
 
 #include <gemmi/mmread_gz.hpp>   // read_structure_gz
@@ -1433,8 +1434,48 @@ std::vector<std::complex<double>> KnownStructureFactors(const MolecularTransform
     return out;
 }
 
-TranslationPeak TranslationFunction(const MolecularTransform &ft, const TranslationSetup &setup,
-                                    const gemmi::Mat33 &rotation, const std::vector<std::complex<double>> *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<int> offset, entry, cross_offset, cross_entry;
+};
+
+void gather_lists(size_t cells, size_t n, const auto &cell_of, std::vector<int> &offset, std::vector<int> &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<int> 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<int>(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<std::complex<double>> *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<std::complex<double>> c(static_cast<size_t>(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<size_t>(nu) * (setup.target_v[e] + static_cast<size_t>(nv) * setup.target_l[e])]
-                    += w * g[t * nops + setup.pair_j[p]] * std::conj(g[t * nops + setup.pair_k[p]]);
-        }
+    std::vector<double> 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<size_t>(nu) * nv * (nw / 2 + 1), points = static_cast<size_t>(nu) * nv * nw;
+    auto *in = reinterpret_cast<std::complex<float> *>(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<double> v = std::conj((*known)[t]) * g[t * nops + setup.cross_op[e]];
-            c[setup.cross_u[e] + static_cast<size_t>(nu) * (setup.cross_v[e] + static_cast<size_t>(nv) * setup.cross_l[e])]
-                    += w * (setup.cross_conj[e] ? std::conj(v) : v);
+    for (size_t c = 0; c < cells; c++) {
+        std::complex<double> 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<float> 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<float>(std::conj(c[i]));
-    const gemmi::Grid<float> 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<double> 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<float> f(std::conj(v));
+        in[c] = std::isnan(f.imag()) ? std::complex<float>(0.0f, 0.0f) : std::complex<float>(f.real(), -f.imag());
+    }
+    MapFromFPhiExecute(nu, nv, nw, in, map);
+    const float norm = static_cast<float>(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<double>(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<double>(x) * x;
+        if (x > map[best] * norm)
             best = i;
     }
-    const double count = static_cast<double>(map.data.size());
+    const double count = static_cast<double>(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<size_t>(map.nu) * map.nv);
-    out.t = gemmi::Fractional(static_cast<double>(u) / map.nu, static_cast<double>(v) / map.nv,
-                              static_cast<double>(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<size_t>(nu) * nv);
+    out.t = gemmi::Fractional(static_cast<double>(u) / nu, static_cast<double>(v) / nv, static_cast<double>(l) / nw);
+    return out;
+}
+
+}  // namespace
+
+TranslationPeak TranslationFunction(const MolecularTransform &ft, const TranslationSetup &setup,
+                                    const gemmi::Mat33 &rotation, const std::vector<std::complex<double>> *known) {
+    return translation_function(ft, setup, make_gather(setup, known != nullptr), rotation, known);
+}
+
+std::vector<TranslationPeak> TranslationFunction(const MolecularTransform &ft, const TranslationSetup &setup,
+                                                 const std::vector<gemmi::Mat33> &rotations,
+                                                 const std::vector<std::complex<double>> *known, size_t nthreads) {
+    const TranslationGather gather = make_gather(setup, known != nullptr);
+    std::vector<TranslationPeak> out(rotations.size());
+    ParallelFor(static_cast<int>(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<TranslationPeak> tf(rotations.size());
-        ParallelFor(static_cast<int>(rotations.size()), nthreads, [&](int i) {
-            tf[i] = TranslationFunction(ft, setup, rotations[i], known);
-        });
-        return tf;
+        return TranslationFunction(ft, setup, rotations, known, nthreads);
     };
 
     std::vector<gemmi::Mat33> orientations;   // of each candidate's placement, before the rigid body

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<std::complex<double>> *known = nullptr);
+// The same at each of `rotations`, `nthreads` of them at once, the grid's gather worked out once for all.
+std::vector<TranslationPeak> TranslationFunction(const MolecularTransform &ft, const TranslationSetup &setup,
+                                                 const std::vector<gemmi::Mat33> &rotations,
+                                                 const std::vector<std::complex<double>> *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