diff --git a/rugnux/ModelFFT.cpp b/rugnux/ModelFFT.cpp index 03ce9c69c..a70c4fb88 100644 --- a/rugnux/ModelFFT.cpp +++ b/rugnux/ModelFFT.cpp @@ -5,6 +5,7 @@ #include #include +#include #include #include #include @@ -17,27 +18,28 @@ 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 the process-wide planner -// lock (which also guards this cache); 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::map, fftwf_plan> plans; +// The r2c plan (or, inverse, the c2r plan) for an (nu, nv, nw) map stored u fastest, halving w - the +// layout gemmi's own transforms read and write. FFTW's planner is not thread-safe, so plans are made +// under the process-wide planner lock (which also guards this cache); executing one on new arrays +// (fftwf_execute_dft_r2c / _c2r) is, which is how every caller uses it. +fftwf_plan PlanFor(int nu, int nv, int nw, bool inverse) { + static std::map, fftwf_plan> plans; std::lock_guard lock(FFTWPlannerMutex()); - const std::array key{nu, nv, nw}; + const std::array key{nu, nv, nw, inverse}; if (const auto it = plans.find(key); it != plans.end()) return it->second; - float *in = fftwf_alloc_real(static_cast(nu) * nv * nw); - fftwf_complex *out = fftwf_alloc_complex(static_cast(nu) * nv * (nw / 2 + 1)); + float *real = fftwf_alloc_real(static_cast(nu) * nv * nw); + fftwf_complex *cplx = fftwf_alloc_complex(static_cast(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); + fftwf_plan plan = nullptr; + if (real != nullptr && cplx != nullptr) + plan = inverse ? fftwf_plan_guru_dft_c2r(3, dims, 0, nullptr, cplx, real, FFTW_ESTIMATE) + : fftwf_plan_guru_dft_r2c(3, dims, 0, nullptr, real, cplx, FFTW_ESTIMATE); + fftwf_free(real); + fftwf_free(cplx); if (plan == nullptr) - gemmi::fail("MapToFPhi(): no FFTW plan for the grid"); + gemmi::fail("ModelFFT: no FFTW plan for the grid"); plans.emplace(key, plan); return plan; } @@ -55,7 +57,7 @@ gemmi::FPhiGrid MapToFPhi(const gemmi::Grid &map) { hkl.set_size_without_checking(map.nu, map.nv, map.nw / 2 + 1); const float norm = static_cast(map.unit_cell.volume / map.point_count()); - const fftwf_plan plan = PlanFor(map.nu, map.nv, map.nw); + const fftwf_plan plan = PlanFor(map.nu, map.nv, map.nw, false); // 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()); @@ -75,3 +77,36 @@ gemmi::FPhiGrid MapToFPhi(const gemmi::Grid &map) { fftwf_free(out); return hkl; } + +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"); + gemmi::Grid map; + map.spacegroup = hkl.spacegroup; + map.unit_cell = hkl.unit_cell; + map.set_size(hkl.nu, hkl.nv, 2 * (hkl.nw - 1)); + map.axis_order = hkl.axis_order; + const float norm = static_cast(1.0 / hkl.unit_cell.volume); + + const fftwf_plan plan = PlanFor(map.nu, map.nv, map.nw, true); + fftwf_complex *in = fftwf_alloc_complex(hkl.data.size()); + float *out = fftwf_alloc_real(map.data.size()); + if (in == nullptr || out == nullptr) { + fftwf_free(in); + fftwf_free(out); + gemmi::fail("MapFromFPhi(): cannot allocate the FFT buffers"); + } + // Conjugated, as gemmi does: rho(x) = 1/V sum F exp(-2 pi i h.x), and FFTW's backward transform + // carries exp(+2 pi i h.x). A missing coefficient (NaN) counts as zero, also as gemmi does. + for (size_t i = 0; i < hkl.data.size(); ++i) { + const std::complex x = hkl.data[i]; + in[i][0] = std::isnan(x.imag()) ? 0.0f : x.real(); + in[i][1] = std::isnan(x.imag()) ? 0.0f : -x.imag(); + } + fftwf_execute_dft_c2r(plan, in, out); + for (size_t i = 0; i < map.data.size(); ++i) + map.data[i] = out[i] * norm; + fftwf_free(in); + fftwf_free(out); + return map; +} diff --git a/rugnux/ModelFFT.h b/rugnux/ModelFFT.h index fd807d1de..793d7153b 100644 --- a/rugnux/ModelFFT.h +++ b/rugnux/ModelFFT.h @@ -16,3 +16,9 @@ // one from run to run - and made once per grid size and kept, so the rigid body's hundred-odd // evaluations on one grid plan once. Safe to call from several threads at once. gemmi::FPhiGrid MapToFPhi(const gemmi::Grid &map); + +// The inverse: the real-space map of an (nu, nv, nw/2+1) half-l coefficient grid in XYZ order, as +// get_f_phi_on_grid() fills it - what gemmi::transform_f_phi_grid_to_map() returns (rho = 1/V sum F +// exp(-2 pi i h.x) on an (nu, nv, 2*(nw-1)) grid, NaN coefficients read as zero), computed with FFTW +// under the same planning rules. +gemmi::Grid MapFromFPhi(const gemmi::FPhiGrid &hkl); diff --git a/rugnux/ModelValidation.cpp b/rugnux/ModelValidation.cpp index 4d09805ab..783b19849 100644 --- a/rugnux/ModelValidation.cpp +++ b/rugnux/ModelValidation.cpp @@ -20,7 +20,7 @@ #include // MaybeGzipped #include // IT92 x-ray form factors #include // DensityCalculator -#include // get_f_phi_on_grid, transform_f_phi_grid_to_map (the maps) +#include // get_size_for_hkl, get_f_phi_on_grid (the maps) #include // SolventMasker #include // Scaling (bulk solvent + anisotropic B) #include // Ccp4 map I/O @@ -48,7 +48,7 @@ long hkl_key(const gemmi::Miller &h) { gemmi::Grid map_from_coefficients(gemmi::AsuData> &coef) { coef.ensure_sorted(); std::array size = gemmi::get_size_for_hkl(coef, {{0, 0, 0}}, 3.0); - return gemmi::transform_f_phi_grid_to_map(gemmi::get_f_phi_on_grid(coef, size, true)); + return MapFromFPhi(gemmi::get_f_phi_on_grid(coef, size, true)); } // Write a map as CCP4; return its RMS (the sigma the map is read in). diff --git a/tests/ModelValidationTest.cpp b/tests/ModelValidationTest.cpp index 3fcb7bee8..134f5ac23 100644 --- a/tests/ModelValidationTest.cpp +++ b/tests/ModelValidationTest.cpp @@ -593,10 +593,12 @@ TEST_CASE("ModelValidation_CCModelFollowsTheSignalByShell", "[ModelValidation]") // The model path's structure factors come from FFTW; they must be gemmi's own transform to float // precision, in the same layout, so prepare_asu_data() reads the same reflections from either. TEST_CASE("ModelValidation_MapToFPhiMatchesGemmi", "[ModelValidation]") { + // An even grid, and an odd one on every axis (FFTW and pocketfft split odd lengths differently). + const auto size = GENERATE(std::array{20, 24, 30}, std::array{15, 21, 27}); gemmi::Grid map; map.unit_cell.set(40.0, 50.0, 60.0, 90.0, 95.0, 90.0); map.spacegroup = gemmi::find_spacegroup_by_name("P 1"); - map.set_size(20, 24, 30); + map.set_size(size[0], size[1], size[2]); std::mt19937 rng(7); std::uniform_real_distribution dist(-1.0f, 1.0f); for (auto &x : map.data) @@ -623,6 +625,75 @@ TEST_CASE("ModelValidation_MapToFPhiMatchesGemmi", "[ModelValidation]") { CHECK(a.v[i].hkl == b.v[i].hkl); } +// The output maps come from FFTW too; each must be gemmi's own map to float precision, point for point +// in the same layout. Coefficients with arbitrary phases, negative indices, expanded by symmetry and +// Friedel, on grids that are odd along u and v (w is even by construction of a half-l grid). +TEST_CASE("ModelValidation_MapFromFPhiMatchesGemmi", "[ModelValidation]") { + const char *sg_name = GENERATE("P 1", "P 1 21 1", "P 21 21 21"); + const gemmi::SpaceGroup *sg = gemmi::find_spacegroup_by_name(sg_name); + REQUIRE(sg != nullptr); + gemmi::AsuData> coef; + coef.unit_cell_.set(31.0, 43.0, 57.0, 90.0, sg->number == 4 ? 104.0 : 90.0, 90.0); + coef.spacegroup_ = sg; + const gemmi::ReciprocalAsu asu(sg); + const gemmi::GroupOps gops = sg->operations(); + std::mt19937 rng(11); + std::uniform_real_distribution amp(0.1f, 10.0f), phase(-3.14159f, 3.14159f); + for (int h = -7; h <= 7; h++) + for (int k = -9; k <= 9; k++) + for (int l = -11; l <= 11; l++) { + const gemmi::Op::Miller hkl{{h, k, l}}; + if ((h == 0 && k == 0 && l == 0) || !asu.is_in(hkl) || gops.is_systematically_absent(hkl)) + continue; + coef.v.push_back({hkl, std::polar(amp(rng), phase(rng))}); + } + REQUIRE(coef.v.size() > 500); + // P 1 takes an odd grid on u and v; the screw axes need even factors. + const std::array size = sg->number == 1 ? std::array{17, 21, 26} + : std::array{18, 24, 26}; + gemmi::FPhiGrid grid = gemmi::get_f_phi_on_grid(coef, size, true); + grid.data[grid.index_n(2, -3, 4)] = std::complex(1.0f, NAN); // a missing coefficient + + const gemmi::Grid ours = MapFromFPhi(grid); + const gemmi::Grid ref = gemmi::transform_f_phi_grid_to_map(gemmi::FPhiGrid(grid)); + REQUIRE(ours.nu == ref.nu); + REQUIRE(ours.nv == ref.nv); + REQUIRE(ours.nw == ref.nw); + REQUIRE(ours.axis_order == ref.axis_order); + REQUIRE(ours.spacegroup == ref.spacegroup); + REQUIRE(ours.data.size() == ref.data.size()); + double largest = 0, worst = 0; + for (size_t i = 0; i < ref.data.size(); i++) { + REQUIRE(std::isfinite(ours.data[i])); + largest = std::max(largest, static_cast(std::abs(ref.data[i]))); + worst = std::max(worst, static_cast(std::abs(ours.data[i] - ref.data[i]))); + } + CHECK(largest > 0); + CHECK(worst <= 1e-5 * largest); +} + +// Map -> coefficients -> map is the identity (the V/N and 1/V scales cancel the unnormalised +// transforms), on a grid odd along u and v. +TEST_CASE("ModelValidation_ModelFFTRoundTrip", "[ModelValidation]") { + gemmi::Grid map; + map.unit_cell.set(35.0, 45.0, 55.0, 80.0, 95.0, 105.0); + map.spacegroup = gemmi::find_spacegroup_by_name("P 1"); + map.set_size(15, 21, 28); + std::mt19937 rng(3); + std::uniform_real_distribution dist(-1.0f, 1.0f); + for (auto &x : map.data) + x = dist(rng); + + const gemmi::Grid back = MapFromFPhi(MapToFPhi(map)); + REQUIRE(back.nu == map.nu); + REQUIRE(back.nv == map.nv); + REQUIRE(back.nw == map.nw); + double worst = 0; + for (size_t i = 0; i < map.data.size(); i++) + worst = std::max(worst, static_cast(std::abs(back.data[i] - map.data[i]))); + CHECK(worst <= 1e-5); +} + // The Jacobian's six columns are evaluated in parallel; the placement must be the serial one, bit for // bit. TEST_CASE("ModelValidation_RigidBodySameOnAnyNumberOfThreads", "[ModelValidation]") {