mx/Jungfraujoch
3
2
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 81 Wiki Activity

ModelMaskGPU: gemmi's shrink step, so fine rigid-body zones stay on the GPU

Browse Source 
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
This commit is contained in:
leonarski_f
2026-10-09 12:38:37 +02:00
co-authored by Claude Opus 5.5
parent 91e63b6f52
commit 6025d51a48
7 changed files with 184 additions and 69 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/ModelMaskGPUTest.cpp
+57 -12
View File
@@ -10,6 +10,7 @@
#include <cstring>
#include <random>
#include <string>
#include <utility>
#include <vector>
#include <gemmi/pdb.hpp>
@@ -202,21 +203,65 @@ TEST_CASE("ModelMaskGPU_MatchesGemmi", "[ModelValidation][gpu]") {
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]") {
// 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;
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));
// The same answer ahead of time, which is what sends such a grid to the CPU before any engine is made.
CHECK(ModelMaskGPU::ShrinkIsNoOp(MaskGridDescription(coarse)));
CHECK_FALSE(ModelMaskGPU::ShrinkIsNoOp(MaskGridDescription(fine)));
// 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