ModelMaskGPU: the Refmac bulk-solvent mask on the GPU
The equivalent of PutMaskOnGrid() / gemmi SolventMasker(Refmac).put_mask_on_grid() for the GPU rigid body. Every symmetry image of every atom is masked directly (one block per image, gemmi's box and !(d2 > r2) rule), which equals gemmi's mask-then-symmetrize-with-min because the operators are isometries. Islands are removed by a 26-connected periodic union-find (atomicCAS, larger root under smaller, so each component is rooted at its lowest index), a read-only root pass, a size count capped at the island limit and a removal pass - no float atomics, bit-identical repeats. The shrink is not implemented: SetGrid() throws when gemmi's 0.8 A stencil would be non-empty, which it is not on any rigid-body grid (spacing d/3 >= 1.17 A). Atoms come in as double fractional coordinates (ModelMaskAtom): with float input, 9fhc at 3.5 A differed from gemmi at one 24-fold orbit of radius points (|d - r| ~ 2e-6 A); with double input the mask equals gemmi's on all 8 prototype sets at 3.5 and 6 A, including 6oel, 8t7r, 9hnc, 9fhc. Time per mask (RTX 5080, Compute incl. islands): 6oel 288^3 4.2 ms, 8t7r 18.4M 2.8 ms, 9fhc 8.0M 1.5 ms, 9hnc 1.9M 0.45 ms. Union-find credited to Playne & Hawick (2018) at the kernel and in ACKNOWLEDGEMENT.md. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01D1G8gJVAy6gp1K5Dz3NE5C
This commit is contained in:
Claude Opus 5.5
@@ -0,0 +1,214 @@
|
||||
// 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 <vector>
|
||||
|
||||
#include <gemmi/pdb.hpp>
|
||||
#include <gemmi/solmask.hpp>
|
||||
#include <gemmi/symmetry.hpp>
|
||||
|
||||
#include "../rugnux/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(), 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(cudaMemcpy(atoms_d, atoms.data(), atoms.size() * sizeof(ModelMaskAtom), cudaMemcpyHostToDevice) == 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();
|
||||
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(cudaMemcpy(mask_d, pre.data.data(), n * sizeof(float), cudaMemcpyHostToDevice) == 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);
|
||||
}
|
||||
|
||||
// The shrink is not implemented: a grid fine enough for it to change anything is refused.
|
||||
TEST_CASE("ModelMaskGPU_RefusesAGridTheShrinkWouldChange", "[ModelValidation][gpu]") {
|
||||
if (get_gpu_count() == 0) {
|
||||
SKIP("No GPU");
|
||||
}
|
||||
CudaStream stream;
|
||||
const gemmi::Structure st = ClusterModel(kCrysts[0]);
|
||||
const std::vector<ModelMaskOp> ops = MaskOps(*st.find_spacegroup());
|
||||
const gemmi::Grid<float> coarse = MaskGrid(st, 3.5);
|
||||
const gemmi::Grid<float> fine = MaskGrid(st, 2.0);
|
||||
ModelMaskGPU mask(stream, fine.point_count());
|
||||
CHECK_NOTHROW(mask.SetGrid(MaskGridDescription(coarse), ops));
|
||||
CHECK_THROWS(mask.SetGrid(MaskGridDescription(fine), ops));
|
||||
}
|
||||
|
||||
#endif
|
||||
Reference in New Issue
Block a user