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
122 lines
4.9 KiB
C++
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
|