Files
Jungfraujoch/rugnux/ModelFFT.cpp
T
leonarski_fandClaude Opus 5 6c8a32124c rugnux --model: model structure factors with FFTW instead of pocketfft
Every map -> structure-factor transform on the model path - Fcalc from the
model density and Fmask from the bulk-solvent mask, in the fit, the rigid-body
evaluations (two per evaluation, ~120 evaluations per placement), the frame
probe and the stills model reference - went through gemmi's vendored pocketfft,
single-threaded. They now go through MapToFPhi (rugnux/ModelFFT.{h,cpp}): an
FFTW r2c (fftw3f, already fetched and linked) planned with the guru interface on
the same layout gemmi uses (u fastest, w halved), scaled by V/N, into the same
FPhiGrid, so gemmi's prepare_asu_data() extracts the reflections exactly as
before. gemmi's vendored code is not modified. The F -> map transforms of the
output maps are left on gemmi.

Plans are FFTW_ESTIMATE only (MEASURE times candidates and could choose
differently between runs), made under a mutex - FFTW's planner is not
thread-safe, and the null replicates and the parallel Jacobian call this
concurrently - cached per grid size for the life of the process, and executed
with fftwf_execute_dft_r2c on fftwf_malloc buffers (new-array execution needs
the planning arrays' alignment, which a std::vector does not promise).

NOT bit-identical with pocketfft: the two libraries order their arithmetic
differently, so structure factors move in the last float bits and everything
downstream (scale fit, rigid body, R-factors, maps) can move in the last
printed digit. Verify by tolerance, separately from the parallelisation commits:
MODEL_* keys equal to printed precision and the same MODEL_FIT / hand /
indexing decisions on the audit set.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_013nW6FNRP1bBJJ8pfHiByAT
2026-09-20 18:45:03 +02:00

74 lines
2.8 KiB
C++

// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only
#include "ModelFFT.h"
#include <algorithm>
#include <array>
#include <complex>
#include <map>
#include <mutex>
#include <fftw3.h>
#include "gemmi/fail.hpp"
namespace {
// The r2c plan for an (nu, nv, nw) map stored u fastest, halving w - the layout gemmi's own transform
// reads and writes. FFTW's planner is not thread-safe, so plans are made under a lock; executing one
// on new arrays (fftwf_execute_dft_r2c) is, which is how every caller uses it.
fftwf_plan PlanFor(int nu, int nv, int nw) {
static std::mutex m;
static std::map<std::array<int, 3>, fftwf_plan> plans;
std::lock_guard lock(m);
const std::array<int, 3> key{nu, nv, nw};
if (const auto it = plans.find(key); it != plans.end())
return it->second;
float *in = fftwf_alloc_real(static_cast<size_t>(nu) * nv * nw);
fftwf_complex *out = fftwf_alloc_complex(static_cast<size_t>(nu) * nv * (nw / 2 + 1));
// The last dimension is the one r2c halves; the strides put u fastest on both sides.
fftwf_iodim dims[3] = {{nu, 1, 1}, {nv, nu, nu}, {nw, nu * nv, nu * nv}};
fftwf_plan plan = (in != nullptr && out != nullptr)
? fftwf_plan_guru_dft_r2c(3, dims, 0, nullptr, in, out, FFTW_ESTIMATE)
: nullptr;
fftwf_free(in);
fftwf_free(out);
if (plan == nullptr)
gemmi::fail("MapToFPhi(): no FFTW plan for the grid");
plans.emplace(key, plan);
return plan;
}
} // namespace
gemmi::FPhiGrid<float> MapToFPhi(const gemmi::Grid<float> &map) {
if (map.axis_order == gemmi::AxisOrder::ZYX)
gemmi::fail("MapToFPhi(): ZYX order is not supported");
gemmi::FPhiGrid<float> hkl;
hkl.unit_cell = map.unit_cell;
hkl.spacegroup = map.spacegroup;
hkl.axis_order = map.axis_order;
hkl.half_l = true;
hkl.set_size_without_checking(map.nu, map.nv, map.nw / 2 + 1);
const float norm = static_cast<float>(map.unit_cell.volume / map.point_count());
const fftwf_plan plan = PlanFor(map.nu, map.nv, map.nw);
// Buffers from fftwf_malloc: a plan may only be executed on arrays aligned as the ones it was made
// on were, which the vectors' own allocations do not promise.
float *in = fftwf_alloc_real(map.data.size());
fftwf_complex *out = fftwf_alloc_complex(hkl.data.size());
if (in == nullptr || out == nullptr) {
fftwf_free(in);
fftwf_free(out);
gemmi::fail("MapToFPhi(): cannot allocate the FFT buffers");
}
std::copy(map.data.begin(), map.data.end(), in);
fftwf_execute_dft_r2c(plan, in, out);
for (size_t i = 0; i < hkl.data.size(); ++i)
hkl.data[i] = std::complex<float>(out[i][0], out[i][1]) * norm;
fftwf_free(in);
fftwf_free(out);
return hkl;
}