From 04655bd8bf12be616ea759986fd4e1bdbc63b489 Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Sun, 18 Oct 2026 12:51:49 +0200 Subject: [PATCH] MolecularReplacement: the rotation function on the Apple GPU (Metal), equal to rounding On macOS, where a Metal device may be used (not under --no-gpu), each merge's rotation stage - the sphere and the peaks' refinement - runs on the GPU: the same trilinear interpolation of the same transform in float, one threadgroup per orientation, every bin's sums of |F|^2, |F|^4 and |F|^2 (E^2 - 1) in a fixed order, and the correlation made of them in double on the host. The GPU reads its own copy of the transform in 2 x 2 x 2 bricks (1.3x faster than the grid's rows, as the eight corners then lie on fewer cache lines), made for the rotation stage and given back before the translation function; the seeds' orientations later on stay on the CPU. Where the GPU fails, the stage goes on on the CPU with a warning. The kernel is MolecularReplacementMetal.metal, carried in the binary as the Bragg integration's are. Elsewhere nothing changes. Measured on an Apple M6 (16 GB), alternating with rc176 + the CPU prefetch, two runs each: - 4nwv (the timing benchmark's moved model): rotation 78.9 -> 30.7 s, wall 344.5 -> 298.8 s; peak footprint 16.1-16.8 -> 18.5 GB. - 9s02: rotation 2.99 -> 0.28 s, wall 79.2 -> 76.4 s. p.hkl, the placed model and the maps are identical to the CPU's on both, and so is every line of the molecular-replacement log (RF Z, TF Z, R values). The RF itself agrees with the CPU's to 1e-4, with the same best orientation (MolecularReplacementMetal_MatchesCPU, [portable]). Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_018My7tqt8a4LyKKjP3FhScT --- .../structure_refinement/CMakeLists.txt | 20 ++ .../MolecularReplacement.cpp | 47 +++- .../MolecularReplacementMetal.h | 38 ++++ .../MolecularReplacementMetal.metal | 92 ++++++++ .../MolecularReplacementMetal.mm | 215 ++++++++++++++++++ tests/CMakeLists.txt | 1 + tests/MolecularReplacementMetalTest.cpp | 121 ++++++++++ 7 files changed, 530 insertions(+), 4 deletions(-) create mode 100644 image_analysis/structure_refinement/MolecularReplacementMetal.h create mode 100644 image_analysis/structure_refinement/MolecularReplacementMetal.metal create mode 100644 image_analysis/structure_refinement/MolecularReplacementMetal.mm create mode 100644 tests/MolecularReplacementMetalTest.cpp diff --git a/image_analysis/structure_refinement/CMakeLists.txt b/image_analysis/structure_refinement/CMakeLists.txt index 0644f809d..d4583917d 100644 --- a/image_analysis/structure_refinement/CMakeLists.txt +++ b/image_analysis/structure_refinement/CMakeLists.txt @@ -21,3 +21,23 @@ ADD_LIBRARY(JFJochStructureRefinement STATIC # JFJochIndexing for CUDAMemHelpers.h and, under CUDA, to put cuFFT (RigidBodyGPU.cu) after these # objects on the link line. TARGET_LINK_LIBRARIES(JFJochStructureRefinement JFJochScaleMerge JFJochIndexing JFJochCommon gemmi fftw3f) + +IF (APPLE) + # The rotation function on the Apple GPU (MolecularReplacementMetal.h). Its kernel is compiled when it is first + # used, from source carried in the binary, as image_analysis/bragg_integration does: the .metal file copied into + # a generated .cpp here; editing it re-runs the configure. + FILE(READ ${CMAKE_CURRENT_SOURCE_DIR}/MolecularReplacementMetal.metal text) + SET_PROPERTY(DIRECTORY APPEND PROPERTY CMAKE_CONFIGURE_DEPENDS ${CMAKE_CURRENT_SOURCE_DIR}/MolecularReplacementMetal.metal) + # Written only when it changed, so a configure does not rebuild the library. + FILE(WRITE ${CMAKE_CURRENT_BINARY_DIR}/MolecularReplacementMetalSource.cpp.tmp + "// Generated from image_analysis/structure_refinement by its CMakeLists.txt - do not edit.\n" + "extern const char *const mr_metal_source;\n" + "const char *const mr_metal_source = R\"MR_METAL(\n${text})MR_METAL\";\n") + CONFIGURE_FILE(${CMAKE_CURRENT_BINARY_DIR}/MolecularReplacementMetalSource.cpp.tmp + ${CMAKE_CURRENT_BINARY_DIR}/MolecularReplacementMetalSource.cpp COPYONLY) + TARGET_SOURCES(JFJochStructureRefinement PRIVATE MolecularReplacementMetal.mm MolecularReplacementMetal.h + MolecularReplacementMetal.metal ${CMAKE_CURRENT_BINARY_DIR}/MolecularReplacementMetalSource.cpp) + SET_SOURCE_FILES_PROPERTIES(MolecularReplacementMetal.mm PROPERTIES COMPILE_OPTIONS -fobjc-arc) + SET_SOURCE_FILES_PROPERTIES(MolecularReplacementMetal.metal PROPERTIES HEADER_FILE_ONLY ON) + TARGET_LINK_LIBRARIES(JFJochStructureRefinement "-framework Metal" "-framework Foundation") +ENDIF() diff --git a/image_analysis/structure_refinement/MolecularReplacement.cpp b/image_analysis/structure_refinement/MolecularReplacement.cpp index e5fc047e9..672c29a9c 100644 --- a/image_analysis/structure_refinement/MolecularReplacement.cpp +++ b/image_analysis/structure_refinement/MolecularReplacement.cpp @@ -38,6 +38,9 @@ #include "ModelValidation.h" #include "RigidBodyRefine.h" #include "MolecularReplacementGPU.h" +#ifdef __APPLE__ +#include "MolecularReplacementMetal.h" +#endif #ifdef JFJOCH_USE_CUDA #include "RigidBodyGPU.h" #include "RigidBodyGPUEngine.h" @@ -2616,7 +2619,36 @@ MolecularReplacementResult search_model(const SearchModel &model, const std::vec const SearchReflections lrefl = MakeSearchReflections(*merge.merged, cell, merge.space_group, d_llg, result.d_max_search); plan_gpu(m); - const std::vector rf = rotation_function(refl, cell, merge.space_group, rotations); +#ifdef __APPLE__ + // On the Apple GPU where there is one, from its copy of the transform, made for this stage and given back + // before the translation function: the seeds' few orientations later on are scored on the CPU. Where the + // GPU fails, on the CPU from there on - a Mac has no other GPU to try. + std::unique_ptr metal; + const auto metal_failed = [&](const JFJochException &e) { + logger.Warning("Molecular replacement: the rotation function on the Apple GPU failed, on the CPU instead: {}", + e.what()); + metal.reset(); + }; + try { + if (MolecularReplacementMetal::Available(ft)) + metal = std::make_unique(ft, nthreads); + } catch (const JFJochException &e) { + metal_failed(e); + } +#endif + const auto rotation_stage = [&](const std::vector &r) { +#ifdef __APPLE__ + if (metal) { + try { + return metal->RotationFunction(refl, cell, merge.space_group, r); + } catch (const JFJochException &e) { + metal_failed(e); + } + } +#endif + return rotation_function(refl, cell, merge.space_group, r); + }; + const std::vector rf = rotation_stage(rotations); double rf_mean = 0, rf_sd = 0; for (double x : rf) rf_mean += x; @@ -2658,7 +2690,7 @@ MolecularReplacementResult search_model(const SearchModel &model, const std::vec for (int axis = 0; axis < 3; axis++) for (double sign : {1.0, -1.0}) trial.push_back(axis_rotation(axis, sign * step).multiply(peak_rotations[k])); - const std::vector v = rotation_function(refl, cell, merge.space_group, trial); + const std::vector v = rotation_stage(trial); size_t next = 0; for (size_t k = 0; k < peaks.size(); k++) { if (!active[k]) @@ -2677,10 +2709,17 @@ MolecularReplacementResult search_model(const SearchModel &model, const std::vec } } } +#ifdef __APPLE__ + const bool on_metal = metal != nullptr; + metal.reset(); +#else + const bool on_metal = false; +#endif result.t_rotation += seconds_since(t0); - logger.Info("Molecular replacement: {} in {}: {} reflections {:.1f}-{:.1f} A, {} orientations, best RF Z {:.1f}", + logger.Info("Molecular replacement: {} in {}: {} reflections {:.1f}-{:.1f} A, {} orientations{}, best RF Z {:.1f}", merge.name, merge.space_group.xhm(), refl.hkl.size(), result.d_max_search, result.d_min_search, - rotations.size(), rf_sd > 0 ? (peak_rf.front() - rf_mean) / rf_sd : 0.0); + rotations.size(), on_metal ? " (rotation function on the Apple GPU)" : "", + rf_sd > 0 ? (peak_rf.front() - rf_mean) / rf_sd : 0.0); // --- translation, for every group and every peak orientation --- t0 = std::chrono::steady_clock::now(); diff --git a/image_analysis/structure_refinement/MolecularReplacementMetal.h b/image_analysis/structure_refinement/MolecularReplacementMetal.h new file mode 100644 index 000000000..2fc2c76a0 --- /dev/null +++ b/image_analysis/structure_refinement/MolecularReplacementMetal.h @@ -0,0 +1,38 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#pragma once + +// The rotation function of MolecularReplacement.h on an Apple GPU (macOS only): the same trilinear interpolation +// of the same molecular transform, the same bins and the same correlation, in float on the device. It agrees with +// the CPU to rounding, not bit for bit. The device sums every bin's |F|^2, |F|^4 and |F|^2 (E^2 - 1) for each +// orientation in a fixed order, and the host makes the correlation of them in double. +// +// The device reads its own copy of the transform, in 2 x 2 x 2 bricks: the eight corners of an interpolation then +// lie on fewer cache lines than in the grid's rows, which the GPU, waiting on memory as the CPU does, reads 1.3x +// faster. The copy takes as much memory as the transform and is held as long as this object. + +#include +#include + +#include "MolecularReplacement.h" + +class MolecularReplacementMetal { +public: + // Whether the rotation function of `ft` can run on the GPU: a Metal device the process may use (rugnux + // --no-gpu hides it, see hide_gpus) whose buffers hold the transform. + static bool Available(const MolecularTransform &ft); + + // Copies the transform into the device's layout. Throws JFJochException on a Metal failure, as every call + // below does. + MolecularReplacementMetal(const MolecularTransform &ft, size_t nthreads); + ~MolecularReplacementMetal(); + + // RotationFunction() of MolecularReplacement.h; may be called from several threads at once. + std::vector RotationFunction(const SearchReflections &refl, const gemmi::UnitCell &cell, + const gemmi::SpaceGroup &sg, const std::vector &rotations) const; + +private: + struct State; + std::unique_ptr state_; +}; diff --git a/image_analysis/structure_refinement/MolecularReplacementMetal.metal b/image_analysis/structure_refinement/MolecularReplacementMetal.metal new file mode 100644 index 000000000..c0280c03c --- /dev/null +++ b/image_analysis/structure_refinement/MolecularReplacementMetal.metal @@ -0,0 +1,92 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +// The rotation function's kernel (MolecularReplacementMetal.h), compiled when it is first used from this source, +// which CMakeLists.txt copies into the binary. + +#include +using namespace metal; + +struct Params { + int nu, nv, bu, bv; // the grid, and its bricks along u and v + uint nops, bins; + float box, blur; +}; + +// MolecularTransform::At()'s corner, from the bricked grid. +static float2 transform_at(device const float2 *g, constant Params &p, 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; + } + if (i < 0 || i >= p.nu) { i %= p.nu; if (i < 0) i += p.nu; } + if (j < 0 || j >= p.nv) { j %= p.nv; if (j < 0) j += p.nv; } + const uint brick = uint(i >> 1) + uint(p.bu) * (uint(j >> 1) + uint(p.bv) * uint(l >> 1)); + float2 v = g[brick * 8 + uint(i & 1) + 2 * uint(j & 1) + 4 * uint(l & 1)]; + if (conj) + v.y = -v.y; + return v; +} + +// One threadgroup per orientation, as MolecularReplacementGPU.cu's rotation_function; the reflections in the order +// of their bins, every bin's sums taken in a fixed order: each thread's own, then the SIMD group's, then the groups'. +// For orientation r, bin b: out[(r bins + b) 3 + k] = the sum over the bin's reflections of ic, ic^2 and +// ic (E^2 - 1 - mean), ic being |F|^2 summed over the copies. +kernel void rotation_function(device const float2 *g [[buffer(0)]], constant Params &p [[buffer(1)]], + device const float3x3 *rot [[buffer(2)]], device const packed_float3 *s [[buffer(3)]], + device const float *w [[buffer(4)]], device const uint *bin_first [[buffer(5)]], + device float *out [[buffer(6)]], + uint r [[threadgroup_position_in_grid]], uint t [[thread_index_in_threadgroup]], + uint nt [[threads_per_threadgroup]], uint lane [[thread_index_in_simdgroup]], + uint sg [[simdgroup_index_in_threadgroup]], uint nsg [[simdgroups_per_threadgroup]]) { + threadgroup float sh[3][32]; + const float3x3 m = rot[r]; // R^T: the model turned by R has F_R(s) = F_model(R^T s) + for (uint b = 0; b < p.bins; b++) { + float s1 = 0, s2 = 0, sw = 0; + for (uint i = bin_first[b] + t; i < bin_first[b + 1]; i += nt) { + float ic = 0; + for (uint j = 0; j < p.nops; j++) { + const float3 v = m * float3(s[i * p.nops + j]); + const float3 u = v * p.box; + const float3 fl = floor(u); + const float3 f = u - fl; + const int3 i0 = int3(fl); + float2 acc = 0; + for (int dx = 0; dx < 2; dx++) + for (int dy = 0; dy < 2; dy++) + for (int dz = 0; dz < 2; dz++) { + const float wt = (dx ? f.x : 1 - f.x) * (dy ? f.y : 1 - f.y) * (dz ? f.z : 1 - f.z); + acc += wt * transform_at(g, p, i0.x + dx, i0.y + dy, i0.z + dz); + } + // |F exp(blur s^2 / 4)|^2 + ic += dot(acc, acc) * exp(p.blur * dot(v, v) / 2.0f); + } + s1 += ic; + s2 += ic * ic; + sw += ic * w[i]; + } + s1 = simd_sum(s1); + s2 = simd_sum(s2); + sw = simd_sum(sw); + if (lane == 0) { + sh[0][sg] = s1; + sh[1][sg] = s2; + sh[2][sg] = sw; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (t == 0) { + float a = 0, c = 0, d = 0; + for (uint k = 0; k < nsg; k++) { + a += sh[0][k]; + c += sh[1][k]; + d += sh[2][k]; + } + device float *o = out + (r * p.bins + b) * 3; + o[0] = a; + o[1] = c; + o[2] = d; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } +} diff --git a/image_analysis/structure_refinement/MolecularReplacementMetal.mm b/image_analysis/structure_refinement/MolecularReplacementMetal.mm new file mode 100644 index 000000000..768fdd143 --- /dev/null +++ b/image_analysis/structure_refinement/MolecularReplacementMetal.mm @@ -0,0 +1,215 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#import + +#include "MolecularReplacementMetal.h" + +#include +#include +#include +#include +#include +#include + +#include "../../common/CUDAWrapper.h" // gpus_are_hidden +#include "../../common/JFJochException.h" +#include "../../common/ParallelFor.h" + +// The kernel's source, MolecularReplacementMetal.metal, copied in at configure time (CMakeLists.txt) so the +// binary needs no file beside it. +extern const char *const mr_metal_source; + +namespace { + +// The grid in bricks of 2 x 2 x 2 points, brick (i/2, j/2, l/2) holding point (i%2, j%2, l%2) at +// (i%2) + 2 (j%2) + 4 (l%2); the bricks in the grid's order. The kernel's index is 32-bit. +size_t BrickedPoints(const gemmi::FPhiGrid &g) { + return static_cast((g.nu + 1) / 2) * ((g.nv + 1) / 2) * ((g.nw + 1) / 2) * 8; +} + +constexpr size_t THREADS = 256; +// Orientations per command buffer: half a second of the GPU at most on the largest transforms. +constexpr size_t BATCH = 256; + +struct Params { + int32_t nu, nv, bu, bv; + uint32_t nops, bins; + float box, blur; +}; + +struct Kernel { + id device = nil; // nil: no Metal device + id pipeline = nil; +}; + +// Compiled once per process, the first time it is asked for. +const Kernel &GetKernel() { + static const Kernel kernel = [] { + Kernel k; + k.device = MTLCreateSystemDefaultDevice(); + if (k.device == nil) + return k; + MTLCompileOptions *options = [MTLCompileOptions new]; + options.languageVersion = MTLLanguageVersion3_0; + NSError *error = nil; + id library = [k.device newLibraryWithSource:@(mr_metal_source) options:options error:&error]; + if (library == nil) + throw JFJochException(JFJochExceptionCategory::GPUMetalError, + std::string("rotation function kernel: ") + error.localizedDescription.UTF8String); + id function = [library newFunctionWithName:@"rotation_function"]; + k.pipeline = [k.device newComputePipelineStateWithFunction:function error:&error]; + if (k.pipeline == nil) + throw JFJochException(JFJochExceptionCategory::GPUMetalError, + std::string("rotation_function: ") + error.localizedDescription.UTF8String); + return k; + }(); + return kernel; +} + +id Upload(id device, const void *data, size_t bytes) { + id b = [device newBufferWithBytes:data length:std::max(bytes, 4) + options:MTLResourceStorageModeShared]; + if (b == nil) + throw JFJochException(JFJochExceptionCategory::GPUMetalError, "rotation function: out of GPU memory"); + return b; +} + +} // namespace + +struct MolecularReplacementMetal::State { + id queue = nil; + id grid = nil; + Params params{}; +}; + +bool MolecularReplacementMetal::Available(const MolecularTransform &ft) { + if (gpus_are_hidden()) + return false; + const Kernel &k = GetKernel(); + const size_t points = BrickedPoints(ft.Grid()); + return k.device != nil && points <= std::numeric_limits::max() + && points * sizeof(std::complex) <= k.device.maxBufferLength; +} + +MolecularReplacementMetal::MolecularReplacementMetal(const MolecularTransform &ft, size_t nthreads) + : state_(std::make_unique()) { + const Kernel &k = GetKernel(); + const gemmi::FPhiGrid &g = ft.Grid(); + const int bu = (g.nu + 1) / 2, bv = (g.nv + 1) / 2, bw = (g.nw + 1) / 2; + state_->params = {g.nu, g.nv, bu, bv, 0, 0, static_cast(ft.Box()), static_cast(ft.Blur())}; + state_->queue = [k.device newCommandQueue]; + // Points past the grid's edge, in its last bricks, are never read. + state_->grid = [k.device newBufferWithLength:BrickedPoints(g) * sizeof(std::complex) + options:MTLResourceStorageModeShared]; + if (state_->queue == nil || state_->grid == nil) + throw JFJochException(JFJochExceptionCategory::GPUMetalError, "rotation function: out of GPU memory"); + auto *out = static_cast *>(state_->grid.contents); + ParallelFor(bw, nthreads, [&](int bl) { + for (int l = 2 * bl; l < std::min(2 * bl + 2, g.nw); l++) + for (int j = 0; j < g.nv; j++) + for (int i = 0; i < g.nu; i++) { + const size_t brick = (i >> 1) + static_cast(bu) * ((j >> 1) + static_cast(bv) * (l >> 1)); + out[brick * 8 + (i & 1) + 2 * (j & 1) + 4 * (l & 1)] = + g.data[i + static_cast(g.nu) * (j + static_cast(g.nv) * l)]; + } + }); +} + +MolecularReplacementMetal::~MolecularReplacementMetal() = default; + +std::vector MolecularReplacementMetal::RotationFunction(const SearchReflections &refl, + const gemmi::UnitCell &cell, + const gemmi::SpaceGroup &sg, + const std::vector &rotations) const { + const Kernel &k = GetKernel(); + 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; + w_mean /= static_cast(n); + double w_ss = 0; + for (double e2 : refl.e2) + w_ss += (e2 - w_mean) * (e2 - w_mean); + + // The reflections in the order of their bins, each bin's from bin_first[b] to bin_first[b + 1]. + std::vector order(n); + std::iota(order.begin(), order.end(), 0); + std::stable_sort(order.begin(), order.end(), [&](size_t a, size_t b) { return refl.bin[a] < refl.bin[b]; }); + std::vector bin_first(refl.bins + 1, 0); + for (size_t i = 0; i < n; i++) + bin_first[refl.bin[i] + 1]++; + for (int b = 0; b < refl.bins; b++) + bin_first[b + 1] += bin_first[b]; + std::vector sf(3 * n * nops), w(n); + for (size_t i = 0; i < n; i++) { + for (size_t j = 0; j < nops; j++) { + const gemmi::Vec3 &v = s[order[i] * nops + j]; + sf[3 * (i * nops + j)] = static_cast(v.x); + sf[3 * (i * nops + j) + 1] = static_cast(v.y); + sf[3 * (i * nops + j) + 2] = static_cast(v.z); + } + w[i] = static_cast(refl.e2[order[i]] - w_mean); + } + // R^T as Metal's float3x3: three columns of four floats. + std::vector rot(12 * rotations.size(), 0.0f); + for (size_t r = 0; r < rotations.size(); r++) + for (int col = 0; col < 3; col++) + for (int row = 0; row < 3; row++) + rot[12 * r + 4 * col + row] = static_cast(rotations[r].a[col][row]); + + Params params = state_->params; + params.nops = static_cast(nops); + params.bins = static_cast(refl.bins); + id d_s = Upload(k.device, sf.data(), sf.size() * sizeof(float)); + id d_w = Upload(k.device, w.data(), w.size() * sizeof(float)); + id d_bin_first = Upload(k.device, bin_first.data(), bin_first.size() * sizeof(uint32_t)); + id d_rot = Upload(k.device, rot.data(), rot.size() * sizeof(float)); + id d_out = [k.device newBufferWithLength:std::max(rotations.size() * refl.bins * 3 * sizeof(float), 4) + options:MTLResourceStorageModeShared]; + if (d_out == nil) + throw JFJochException(JFJochExceptionCategory::GPUMetalError, "rotation function: out of GPU memory"); + std::vector> running; + for (size_t first = 0; first < rotations.size(); first += BATCH) { + const size_t count = std::min(BATCH, rotations.size() - first); + id cb = [state_->queue commandBuffer]; + id e = [cb computeCommandEncoder]; + [e setComputePipelineState:k.pipeline]; + [e setBuffer:state_->grid offset:0 atIndex:0]; + [e setBytes:¶ms length:sizeof(params) atIndex:1]; + [e setBuffer:d_rot offset:first * 12 * sizeof(float) atIndex:2]; + [e setBuffer:d_s offset:0 atIndex:3]; + [e setBuffer:d_w offset:0 atIndex:4]; + [e setBuffer:d_bin_first offset:0 atIndex:5]; + [e setBuffer:d_out offset:first * refl.bins * 3 * sizeof(float) atIndex:6]; + [e dispatchThreadgroups:MTLSizeMake(count, 1, 1) threadsPerThreadgroup:MTLSizeMake(THREADS, 1, 1)]; + [e endEncoding]; + [cb commit]; + running.push_back(cb); + } + for (id cb : running) { + [cb waitUntilCompleted]; + if (cb.status != MTLCommandBufferStatusCompleted) + throw JFJochException(JFJochExceptionCategory::GPUMetalError, + cb.error ? cb.error.localizedDescription.UTF8String : "command buffer failed"); + } + + // The correlation of E^2 - 1 with ic over its bin's mean, from the bins' sums: x = ic / mean_b. + const auto *sums = static_cast(d_out.contents); + std::vector out(rotations.size()); + for (size_t r = 0; r < rotations.size(); r++) { + double x_sum = 0, x_ss = 0, xy = 0; + for (int b = 0; b < refl.bins; b++) { + const float *o = sums + (r * refl.bins + b) * 3; + const double mean = o[0] / static_cast(bin_first[b + 1] - bin_first[b]); + x_sum += o[0] / mean; + x_ss += o[1] / (mean * mean); + xy += o[2] / mean; + } + 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; +} diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 5fd306d3f..7dd31e085 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -21,6 +21,7 @@ ADD_EXECUTABLE(jfjoch_portable_test EXCLUDE_FROM_ALL BraggIntegrationEngineCompressedImageTest.cpp BraggIntegrationEngineCPUTest.cpp BraggIntegrationEngineMetalTest.cpp + MolecularReplacementMetalTest.cpp WriteReflectionsTest.cpp LoadReferenceMtzFreeFlagsTest.cpp LoadReferenceSfCifTest.cpp diff --git a/tests/MolecularReplacementMetalTest.cpp b/tests/MolecularReplacementMetalTest.cpp new file mode 100644 index 000000000..bfe30a50c --- /dev/null +++ b/tests/MolecularReplacementMetalTest.cpp @@ -0,0 +1,121 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#include + +#ifdef __APPLE__ + +#include +#include +#include +#include + +#include +#include + +#include "../image_analysis/structure_refinement/ModelFFT.h" +#include "../image_analysis/structure_refinement/MolecularReplacement.h" +#include "../image_analysis/structure_refinement/MolecularReplacementMetal.h" + +// The rotation function on the Apple GPU is the CPU's in float: equal to rounding, with the same best orientation. +// A Mac without a Metal device fails rather than skips - the macOS CI job is the only place this can be checked, +// as in BraggIntegrationEngineMetalTest.cpp. The search as in MolecularReplacementTest.cpp. + +namespace { + // A synthetic body: carbon atoms scattered in a 12 A sphere, centroid at the origin. + SearchModel synthetic_model() { + SearchModel m; + std::mt19937 rng(7); + std::uniform_real_distribution u(-12.0, 12.0); + gemmi::Model model(1); + gemmi::Chain chain("A"); + gemmi::Residue res; + res.name = "UNK"; + res.seqid = gemmi::SeqId(1, ' '); + std::vector pos; + while (pos.size() < 300) { + const gemmi::Vec3 p(u(rng), u(rng), u(rng)); + if (p.length() < 12.0) + pos.push_back(p); + } + gemmi::Vec3 c; + for (const auto &p : pos) + c += p; + c *= 1.0 / pos.size(); + for (const auto &p : pos) { + gemmi::Atom a; + a.name = "CA"; // every atom counts in the clash check + a.element = gemmi::El::C; + a.pos = gemmi::Position(p - c); + a.b_iso = 15.0f; + a.occ = 1.0f; + res.atoms.push_back(a); + m.radius_A = std::max(m.radius_A, (p - c).length()); + } + chain.residues.push_back(res); + model.chains.push_back(chain); + m.st.models.push_back(model); + m.residues_read = m.residues_kept = 1; + m.mass_da = 3600; + return m; + } + + gemmi::Mat33 axis_angle(gemmi::Vec3 axis, double deg) { + axis = axis.normalized(); + const double k = deg * M_PI / 180.0, s = std::sin(k), c1 = 1.0 - std::cos(k); + const double x = axis.x, y = axis.y, z = axis.z; + return {1 + c1 * (x * x - 1), -s * z + c1 * x * y, s * y + c1 * x * z, + s * z + c1 * x * y, 1 + c1 * (y * y - 1), -s * x + c1 * y * z, + -s * y + c1 * x * z, s * x + c1 * y * z, 1 + c1 * (z * z - 1)}; + } + + // |F|^2 of the placed crystal to 4 A, as merged reflections. + std::vector observed(const gemmi::Structure &crystal) { + gemmi::DensityCalculator, float> dc; + dc.d_min = 3.5; + dc.rate = 1.5; + dc.set_grid_cell_and_spacegroup(crystal); + dc.put_model_density_on_grid(crystal.models[0]); + const auto f = MapToFPhi(dc.grid).prepare_asu_data(3.5, 0, false, false, false); + std::vector out; + for (const auto &hv : f.v) { + MergedReflection r; + r.h = hv.hkl[0]; + r.k = hv.hkl[1]; + r.l = hv.hkl[2]; + r.I = std::norm(hv.value); + r.sigma = 1; + r.F = std::abs(hv.value); + r.d = 1.0 / std::sqrt(crystal.cell.calculate_1_d2(hv.hkl)); + out.push_back(r); + } + return out; + } +} + +TEST_CASE("MolecularReplacementMetal_MatchesCPU", "[portable][MolecularReplacement][Metal]") { + 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 auto merged = observed(PlaceSearchModel(model, truth, gemmi::Fractional(0.17, 0.31, 0.09), cell, sg)); + const SearchReflections refl = MakeSearchReflections(merged, cell, sg, 4.0, 12.0); + const MolecularTransform ft(model, 4.0, 4); + REQUIRE(MolecularReplacementMetal::Available(ft)); + const MolecularReplacementMetal metal(ft, 4); + + const auto pg = CartesianPointGroup(sg, cell); + const auto rotations = RotationAsu(UniformRotations(6.0), pg); + const auto rf_cpu = RotationFunction(ft, refl, cell, sg, rotations, 4); + const auto rf_metal = metal.RotationFunction(refl, cell, sg, rotations); + REQUIRE(rf_metal.size() == rf_cpu.size()); + double worst = 0; + for (size_t i = 0; i < rf_cpu.size(); i++) + worst = std::max(worst, std::fabs(rf_metal[i] - rf_cpu[i])); + CHECK(worst < 1e-4); + const size_t best = std::max_element(rf_metal.begin(), rf_metal.end()) - rf_metal.begin(); + CHECK(best == static_cast(std::max_element(rf_cpu.begin(), rf_cpu.end()) - rf_cpu.begin())); + CHECK(RotationDistanceDeg(rotations[best], truth, pg) < 6.0); +} + +#endif