mx/Jungfraujoch
3
2
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases 81 Wiki Activity
Files
6025d51a4890934684cbff971061013261629d89
Jungfraujoch/tests/ModelMaskGPUTest.cpp
T
leonarski_fandClaude Opus 5.5 6025d51a48 ModelMaskGPU: gemmi's shrink step, so fine rigid-body zones stay on the GPU 
The GPU bulk-solvent mask had no shrink step (SolventMasker::shrink(),
set_margin_around() with Refmac's r_shrink = 0.8 A), so every rigid-body
zone whose grid has an offset within 0.8 A was sent to the CPU. That is
not only fine grids: in an oblique setting (P2_1, beta ~141 deg) the
3.5 A zone's lattice-plane spacing is 0.69 A, and on such a set the real
fit and all nine null replicates ran their rigid bodies on the CPU.

The shrink is now two passes over the grid: mark the solvent points with
a macromolecule point among gemmi's near offsets, then turn into solvent
every macromolecule point with such an edge point at any stencil offset
(near or far, gemmi's split at the coarsest grid step). That is gemmi's
rule in both of its branches, and the stencil is built on the host in
double exactly as gemmi builds it. ModelMaskGPU::ShrinkIsNoOp() and
RigidBodyGPUEngine::MaskSupports() are gone; the rigid body no longer
refuses such zones.

Verified (ModelMaskGPU_ShrinkMatchesGemmi): the shrink alone on gemmi's
post-island mask is bit-identical, and the whole GPU mask equals gemmi's
put_mask_on_grid() bit for bit (0 differing points) on the five test
groups at 1.5 A and an oblique P2_1 cell at 3.5 A; repeats are identical.

Measured on the oblique-setting set of the open arm (P2_1, 14k atoms,
1.66 A): model validation 17.2 s -> 5.2 s with the GPU structure factors
of the next commit in place (the first validation's ten rigid bodies,
~12 s on the CPU, now take well under a second). p.mtz md5-identical;
the null moves from +73.1 to +86.1 sigma (the replicates are refined on
the GPU, which agrees with the CPU to rounding), verdict unchanged.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01SVmAWnzCmRKAXVUCdc4iNi
2026-10-09 12:38:37 +02:00

268 lines
13 KiB
C++
Raw Blame History

// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only
#include <catch2/catch_all.hpp>
#include "../common/CUDAWrapper.h"
#ifdef JFJOCH_USE_CUDA
#include <cmath>
#include <cstring>
#include <random>
#include <string>
#include <utility>
#include <vector>
#include <gemmi/pdb.hpp>
#include <gemmi/solmask.hpp>
#include <gemmi/symmetry.hpp>
#include "../image_analysis/structure_refinement/ModelGrid.h"
#include "../image_analysis/structure_refinement/ModelMaskGPU.h"
namespace {
// The five groups of ModelValidation_ParallelGriddingMatchesGemmi: none, a centring, screws, a cubic
// body centring and a cubic face centring with a diamond glide.
const char *kCrysts[] = {
"CRYST1 40.000 50.000 60.000 90.00 90.00 90.00 P 1 1\n",
"CRYST1 40.000 50.000 60.000 90.00 100.00 90.00 C 1 2 1 4\n",
"CRYST1 40.000 50.000 60.000 90.00 90.00 90.00 P 21 21 21 4\n",
"CRYST1 60.000 60.000 60.000 90.00 90.00 90.00 I 2 3 24\n",
"CRYST1 80.000 80.000 80.000 90.00 90.00 90.00 F 41 3 2 96\n",
};
// ModelValidationTest.cpp's ClusterPdb(): random atoms in one corner of the cell, from a fixed seed, with
// every fifth atom a nitrogen and every seventh an oxygen so that more than one radius is in play. A
// cluster this loose leaves many small pockets of solvent between its atoms: the islands.
gemmi::Structure ClusterModel(const char *cryst) {
std::string pdb = cryst;
std::mt19937 rng(20260902);
std::uniform_real_distribution<double> x(2, 14), y(2, 16), z(2, 18);
char line[96];
for (int i = 1; i <= 150; i++) {
const char *el = i % 5 == 0 ? "N" : i % 7 == 0 ? "O" : "C";
std::snprintf(line, sizeof line,
"ATOM %5d %s UNK A 1 %8.3f%8.3f%8.3f 1.00 20.00 %s\n",
i, el, x(rng), y(rng), z(rng), el);
pdb += line;
}
pdb += "END\n";
gemmi::Structure st = gemmi::read_pdb_string(pdb, "cluster");
return st;
}
gemmi::Grid<float> MaskGrid(const gemmi::Structure &st, double d_min) {
gemmi::Grid<float> grid;
grid.unit_cell = st.cell;
grid.spacegroup = st.find_spacegroup();
grid.set_size_from_spacing(d_min / (2 * 1.5), gemmi::GridSizeRounding::Up);
return grid;
}
// What the rigid body hands ModelMaskGPU: the atoms gemmi masks, fractional and wrapped, with gemmi's
// Refmac radius; every symmetry operator with every centring vector.
std::vector<ModelMaskAtom> MaskAtoms(const gemmi::Structure &st) {
const gemmi::SolventMasker masker(gemmi::AtomicRadiiSet::Refmac);
std::vector<ModelMaskAtom> atoms;
for (const gemmi::Chain &ch : st.models[0].chains)
for (const gemmi::Residue &res : ch.residues)
for (const gemmi::Atom &a : res.atoms) {
if (a.is_hydrogen() || a.occ <= 0)
continue;
const gemmi::Fractional f = st.cell.fractionalize(a.pos).wrap_to_unit();
const double r = masker.constant_r + masker.rprobe
+ gemmi::refmac_radius_for_bulk_solvent(a.element.elem);
atoms.push_back(ModelMaskAtom{f.x, f.y, f.z, r});
}
return atoms;
}
std::vector<ModelMaskOp> MaskOps(const gemmi::SpaceGroup &sg) {
const gemmi::GroupOps gops = sg.operations();
std::vector<ModelMaskOp> ops;
for (const auto &cen : gops.cen_ops)
for (const gemmi::Op &op : gops.sym_ops) {
ModelMaskOp m{};
for (int i = 0; i < 3; i++) {
for (int j = 0; j < 3; j++)
m.rot[3 * i + j] = static_cast<double>(op.rot[i][j]) / gemmi::Op::DEN;
m.tran[i] = static_cast<double>(op.tran[i] + cen[i]) / gemmi::Op::DEN;
}
ops.push_back(m);
}
return ops;
}
ModelMaskGrid MaskGridDescription(const gemmi::Grid<float> &grid) {
ModelMaskGrid g;
g.nu = grid.nu;
g.nv = grid.nv;
g.nw = grid.nw;
for (int i = 0; i < 3; i++)
for (int j = 0; j < 3; j++)
g.orth[3 * i + j] = grid.unit_cell.orth.mat[i][j];
g.volume = grid.unit_cell.volume;
return g;
}
// |distance - radius| to the nearest image of any atom that reaches point (u, v, w), in A: how close
// the point lies to the edge of the mask.
double DistanceToMaskEdge(const gemmi::Grid<float> &grid, const std::vector<ModelMaskAtom> &atoms,
const std::vector<ModelMaskOp> &ops, int u, int v, int w) {
const double p[3] = {static_cast<double>(u) / grid.nu, static_cast<double>(v) / grid.nv,
static_cast<double>(w) / grid.nw};
double best = INFINITY;
for (const ModelMaskAtom &a : atoms)
for (const ModelMaskOp &op : ops) {
const double f[3] = {a.x, a.y, a.z};
gemmi::Fractional d;
for (int k = 0; k < 3; k++) {
const double x = op.rot[3 * k] * f[0] + op.rot[3 * k + 1] * f[1] + op.rot[3 * k + 2] * f[2]
+ op.tran[k];
d.at(k) = p[k] - x - std::round(p[k] - x);
}
best = std::min(best, std::fabs(grid.unit_cell.orthogonalize_difference(d).length() - a.radius));
}
return best;
}
} // namespace
// The whole mask against gemmi's put_mask_on_grid() - and so against the CPU path's PutMaskOnGrid() -
// the island step alone against gemmi's on gemmi's own pre-island mask (exact), and three repeats bit
// for bit. The masking itself may differ from gemmi only at points that sit at an atom's radius, where
// the float distance rounds the other way.
TEST_CASE("ModelMaskGPU_MatchesGemmi", "[ModelValidation][gpu]") {
if (get_gpu_count() == 0)
SKIP("No GPU");
CudaStream stream;
bool any_islands = false, any_odd_size = false;
for (const char *cryst : kCrysts) {
const gemmi::Structure st = ClusterModel(cryst);
const gemmi::SpaceGroup *sg = st.find_spacegroup();
REQUIRE(sg != nullptr);
const std::vector<ModelMaskAtom> atoms = MaskAtoms(st);
const std::vector<ModelMaskOp> ops = MaskOps(*sg);
CudaDevicePtr<ModelMaskAtom> atoms_d(atoms.size());
REQUIRE(cudaMemcpyAsync(atoms_d, atoms.data(), atoms.size() * sizeof(ModelMaskAtom), cudaMemcpyHostToDevice, stream) == cudaSuccess);
for (double d_min : {6.0, 3.5}) {
const gemmi::SolventMasker masker(gemmi::AtomicRadiiSet::Refmac);
gemmi::Grid<float> pre = MaskGrid(st, d_min);
masker.clear(pre);
masker.mask_points(pre, st.models[0]);
masker.symmetrize(pre);
gemmi::Grid<float> ref = pre;
const int islands = masker.remove_islands(ref);
masker.shrink(ref);
const size_t n = ref.data.size();
// The rigid body's CPU path masks with PutMaskOnGrid, which is gemmi's put_mask_on_grid bit
// for bit - checked here on this grid, so what the GPU is held to below is the CPU path too.
gemmi::Grid<float> cpu = MaskGrid(st, d_min);
PutMaskOnGrid(cpu, st.models[0], OrbitLeaders(cpu, 4), 4);
REQUIRE(std::memcmp(cpu.data.data(), ref.data.data(), n * sizeof(float)) == 0);
any_islands |= islands > 0;
any_odd_size |= ref.nu % 8 != 0 || ref.nv % 8 != 0 || ref.nw % 8 != 0;
ModelMaskGPU mask(stream, n);
mask.SetGrid(MaskGridDescription(ref), ops);
CudaDevicePtr<float> mask_d(n);
std::vector<float> out(n), first;
// The island step alone, on gemmi's pre-island mask: exact.
REQUIRE(cudaMemcpyAsync(mask_d, pre.data.data(), n * sizeof(float), cudaMemcpyHostToDevice, stream) == cudaSuccess);
mask.RemoveIslands(mask_d);
REQUIRE(cudaMemcpyAsync(out.data(), mask_d, n * sizeof(float), cudaMemcpyDeviceToHost, stream) == cudaSuccess);
REQUIRE(cudaStreamSynchronize(stream) == cudaSuccess);
CHECK(std::memcmp(out.data(), ref.data.data(), n * sizeof(float)) == 0);
for (int repeat = 0; repeat < 3; repeat++) {
mask.Compute(atoms_d, static_cast<int>(atoms.size()), mask_d);
REQUIRE(cudaMemcpyAsync(out.data(), mask_d, n * sizeof(float), cudaMemcpyDeviceToHost, stream) == cudaSuccess);
REQUIRE(cudaStreamSynchronize(stream) == cudaSuccess);
if (repeat == 0) {
first = out;
size_t differ = 0;
for (size_t i = 0; i < n; i++)
if (out[i] != ref.data[i]) {
differ++;
const int u = static_cast<int>(i % ref.nu), v = static_cast<int>(i / ref.nu % ref.nv),
w = static_cast<int>(i / (static_cast<size_t>(ref.nu) * ref.nv));
CHECK(DistanceToMaskEdge(ref, atoms, ops, u, v, w) < 1e-4);
}
INFO(sg->hm << " d " << d_min << " grid " << ref.nu << "x" << ref.nv << "x" << ref.nw
<< " islands " << islands << ": " << differ << " points differ");
CHECK(differ <= 15 * n / 10000000 + 1);
} else
CHECK(std::memcmp(out.data(), first.data(), n * sizeof(float)) == 0);
}
}
}
CHECK(any_islands);
CHECK(any_odd_size);
}
// Grids fine enough for gemmi's shrink to act. The shrink alone, on gemmi's own mask after its island
// step, is exact; the whole mask may differ from gemmi's only where the masking does (points at an atom's
// radius) and at the neighbours the shrink then treats differently - measured, it does not differ at
// all here - and is the same bits twice.
TEST_CASE("ModelMaskGPU_ShrinkMatchesGemmi", "[ModelValidation][gpu]") {
if (get_gpu_count() == 0)
SKIP("No GPU");
CudaStream stream;
// The five groups at 1.5 A, and a strongly oblique monoclinic cell at 3.5 A, where the planes of the
// grid along a and c are closer than r_shrink although its steps along the axes are not.
std::vector<std::pair<const char *, double>> cases;
for (const char *cryst : kCrysts)
cases.emplace_back(cryst, 1.5);
cases.emplace_back("CRYST1 40.000 50.000 60.000 90.00 141.00 90.00 P 1 21 1 2\n", 3.5);
for (const auto &[cryst, d_min] : cases) {
const gemmi::Structure st = ClusterModel(cryst);
const gemmi::SpaceGroup *sg = st.find_spacegroup();
REQUIRE(sg != nullptr);
const std::vector<ModelMaskAtom> atoms = MaskAtoms(st);
const std::vector<ModelMaskOp> ops = MaskOps(*sg);
CudaDevicePtr<ModelMaskAtom> atoms_d(atoms.size());
REQUIRE(cudaMemcpyAsync(atoms_d, atoms.data(), atoms.size() * sizeof(ModelMaskAtom), cudaMemcpyHostToDevice, stream) == cudaSuccess);
const gemmi::SolventMasker masker(gemmi::AtomicRadiiSet::Refmac);
gemmi::Grid<float> islands = MaskGrid(st, d_min);
masker.clear(islands);
masker.mask_points(islands, st.models[0]);
masker.symmetrize(islands);
masker.remove_islands(islands);
gemmi::Grid<float> ref = islands;
masker.shrink(ref);
const size_t n = ref.data.size();
REQUIRE(std::memcmp(ref.data.data(), islands.data.data(), n * sizeof(float)) != 0); // the shrink acts
ModelMaskGPU mask(stream, n);
mask.SetGrid(MaskGridDescription(ref), ops);
CudaDevicePtr<float> mask_d(n);
std::vector<float> out(n), first;
REQUIRE(cudaMemcpyAsync(mask_d, islands.data.data(), n * sizeof(float), cudaMemcpyHostToDevice, stream) == cudaSuccess);
mask.Shrink(mask_d);
REQUIRE(cudaMemcpyAsync(out.data(), mask_d, n * sizeof(float), cudaMemcpyDeviceToHost, stream) == cudaSuccess);
REQUIRE(cudaStreamSynchronize(stream) == cudaSuccess);
CHECK(std::memcmp(out.data(), ref.data.data(), n * sizeof(float)) == 0);
for (int repeat = 0; repeat < 2; repeat++) {
mask.Compute(atoms_d, static_cast<int>(atoms.size()), mask_d);
REQUIRE(cudaMemcpyAsync(out.data(), mask_d, n * sizeof(float), cudaMemcpyDeviceToHost, stream) == cudaSuccess);
REQUIRE(cudaStreamSynchronize(stream) == cudaSuccess);
if (repeat == 0) {
first = out;
size_t differ = 0;
for (size_t i = 0; i < n; i++)
differ += out[i] != ref.data[i];
INFO(sg->hm << " grid " << ref.nu << "x" << ref.nv << "x" << ref.nw << ": " << differ << " points differ");
CHECK(differ <= 27 * (15 * n / 10000000 + 1));
} else
CHECK(std::memcmp(out.data(), first.data(), n * sizeof(float)) == 0);
}
}
}
#endif
Reference in New Issue View Git Blame Copy Permalink