Files
Jungfraujoch/tests/MolecularReplacementMetalTest.cpp
T
leonarski_fandClaude Opus 5.5 04655bd8bf 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 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_018My7tqt8a4LyKKjP3FhScT
2026-10-18 12:51:49 +02:00

122 lines
4.9 KiB
C++

// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only
#include <catch2/catch_all.hpp>
#ifdef __APPLE__
#include <algorithm>
#include <cmath>
#include <random>
#include <vector>
#include <gemmi/dencalc.hpp>
#include <gemmi/it92.hpp>
#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<double> 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<gemmi::Vec3> 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<MergedReflection> observed(const gemmi::Structure &crystal) {
gemmi::DensityCalculator<gemmi::IT92<float>, 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<MergedReflection> 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<size_t>(std::max_element(rf_cpu.begin(), rf_cpu.end()) - rf_cpu.begin()));
CHECK(RotationDistanceDeg(rotations[best], truth, pg) < 6.0);
}
#endif