diff --git a/docs/ACKNOWLEDGEMENT.md b/docs/ACKNOWLEDGEMENT.md index 4f4dc67ff..7a5c5a5e7 100644 --- a/docs/ACKNOWLEDGEMENT.md +++ b/docs/ACKNOWLEDGEMENT.md @@ -357,6 +357,11 @@ the result, follows Jon Wright (ESRF): "Experiences with GPU decompression for b data", HDF5 User Group meeting (2021), and [bslz4decoders](https://github.com/jonwright/bslz4decoders). The CUDA kernels in Jungfraujoch are its own, but the approach is his. +Removing the small islands of solvent from the bulk-solvent mask on the GPU labels the connected +components with a parallel union-find, following D. P. Playne and K. Hawick, "A New Algorithm for +Parallel Connected-Component Labelling on GPUs" (2018), IEEE Trans. Parallel Distrib. Syst. 29, +1217-1230 [doi:10.1109/TPDS.2018.2799216](https://doi.org/10.1109/TPDS.2018.2799216). + This software uses the Viridis, Magma and Inferno colormaps from Matplotlib under its BSD-compatible license. J. D. Hunter, "Matplotlib: A 2D graphics environment" (2007), Comput. Sci. Eng. 9, 90-95 [doi:10.1109/MCSE.2007.55](https://doi.org/10.1109/MCSE.2007.55). diff --git a/rugnux/CMakeLists.txt b/rugnux/CMakeLists.txt index ad785e24a..ae877c757 100644 --- a/rugnux/CMakeLists.txt +++ b/rugnux/CMakeLists.txt @@ -30,6 +30,7 @@ ADD_LIBRARY(JFJochRugnux STATIC HotPixels.cpp HotPixels.h $<$:HotPixelsGPU.cu> + $<$:ModelMaskGPU.cu> HotPixelsGPU.h DiagnosticOutput.cpp DiagnosticOutput.h diff --git a/rugnux/ModelMaskGPU.cu b/rugnux/ModelMaskGPU.cu new file mode 100644 index 000000000..e1cf91466 --- /dev/null +++ b/rugnux/ModelMaskGPU.cu @@ -0,0 +1,272 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#include "ModelMaskGPU.h" + +#include +#include +#include + +#include "../common/JFJochException.h" + +namespace { + +void cuda_err(cudaError_t val) { + if (val != cudaSuccess) + throw JFJochException(JFJochExceptionCategory::GPUCUDAError, cudaGetErrorString(val)); +} + +constexpr int THREADS = 256; +constexpr int IMAGE_THREADS = 64; + +// gemmi SolventMasker(AtomicRadiiSet::Refmac): rshrink and island_min_volume. +constexpr double R_SHRINK = 0.8; +constexpr double ISLAND_MIN_VOLUME = 50.0; + +__device__ int wrap(int a, int n) { + const int r = a % n; + return r < 0 ? r + n : r; +} + +__global__ void fill_solvent(float *mask, size_t n) { + const size_t i = blockIdx.x * static_cast(blockDim.x) + threadIdx.x; + if (i < n) + mask[i] = 1.0f; +} + +// One block per image (atom, operator): gemmi's use_points_around(fpos, radius) around the image, +// setting to 0 every point of the box with !(dist_sq > radius^2). The box is sized from the interplanar +// spacing, clamped to n - 1 and centred on the rounded grid position, as gemmi does it. +__global__ void mask_images(const ModelMaskAtom *atoms, const ModelMaskOp *ops, int n_ops, + ModelMaskGPU::Geometry g, float *mask) { + const ModelMaskAtom a = atoms[blockIdx.x / n_ops]; + // The image's box, one axis per thread: double arithmetic is slow on the GPU, so it is done once + // per image rather than once per point. + __shared__ int lo[3], len[3]; + __shared__ float c[3]; // the image's grid position less the box's first point + if (threadIdx.x < 3) { + const int k = threadIdx.x; + const ModelMaskOp &op = ops[blockIdx.x % n_ops]; + const int n = k == 0 ? g.nu : k == 1 ? g.nv : g.nw; + const double x = op.rot[3 * k] * a.x + op.rot[3 * k + 1] * a.y + op.rot[3 * k + 2] * a.z + op.tran[k]; + const double nx = x * n; + const int centre = static_cast(round(nx)); + const int d = min(static_cast(ceil(a.radius / g.spacing[k])), n - 1); + lo[k] = centre - d; + len[k] = 2 * d + 1; + c[k] = static_cast(nx - centre) + static_cast(d); + } + __syncthreads(); + const float r2 = static_cast(a.radius * a.radius); + const int total = len[0] * len[1] * len[2]; + for (int t = threadIdx.x; t < total; t += blockDim.x) { + const int iu = t % len[0], iv = (t / len[0]) % len[1], iw = t / (len[0] * len[1]); + const float fu = c[0] - iu, fv = c[1] - iv, fw = c[2] - iw; + const float x = g.orth_n[0] * fu + g.orth_n[1] * fv + g.orth_n[2] * fw; + const float y = g.orth_n[3] * fu + g.orth_n[4] * fv + g.orth_n[5] * fw; + const float z = g.orth_n[6] * fu + g.orth_n[7] * fv + g.orth_n[8] * fw; + if (x * x + y * y + z * z > r2) + continue; + const size_t idx = wrap(lo[0] + iu, g.nu) + + static_cast(g.nu) * (wrap(lo[1] + iv, g.nv) + + static_cast(g.nv) * wrap(lo[2] + iw, g.nw)); + mask[idx] = 0.0f; + } +} + +// Union-find labelling of the solvent after Playne & Hawick (2018) IEEE TPDS 29, 1217-1230: every +// solvent point starts as its own set, neighbouring points are united with atomicCAS, always linking the +// larger root under the smaller, so each component ends up rooted at its lowest index whatever the +// order of the unions. + +// Path halving: a non-root entry is pointed at its grandparent. Entries only ever point to a lower +// index of the same set, so the racing writes of other threads are harmless and a root is changed +// only by the atomicCAS in unite(). +__device__ int find_root(int *label, int i) { + int p = label[i]; + while (p != i) { + const int gp = label[p]; + if (gp != p) + label[i] = gp; + i = p; + p = gp; + } + return i; +} + +__device__ void unite(int *label, int a, int b) { + while (true) { + a = find_root(label, a); + b = find_root(label, b); + if (a == b) + return; + if (a < b) { + const int t = a; + a = b; + b = t; + } + const int old = atomicCAS(&label[a], a, b); + if (old == a) + return; + a = old; + } +} + +__global__ void uf_init(const float *mask, int *label, size_t n) { + const size_t i = blockIdx.x * static_cast(blockDim.x) + threadIdx.x; + if (i < n) + label[i] = mask[i] == 1.0f ? static_cast(i) : -1; +} + +// gemmi's FloodFill joins a run of points to the runs of the rows v +- 1, w +- 1 over u - 1 .. u + len, +// with wrap-around: 26-connectivity on the periodic grid. Each point unites with the 13 neighbours +// ahead of it; the other 13 unite with it from their side. +__global__ void uf_merge(const float *mask, int *label, int nu, int nv, int nw) { + const size_t n = static_cast(nu) * nv * nw; + const size_t i = blockIdx.x * static_cast(blockDim.x) + threadIdx.x; + if (i >= n || mask[i] != 1.0f) + return; + const int u = static_cast(i % nu); + const int v = static_cast((i / nu) % nv); + const int w = static_cast(i / (static_cast(nu) * nv)); + const int off[13][3] = {{1, 0, 0}, {-1, 1, 0}, {0, 1, 0}, {1, 1, 0}, + {-1, -1, 1}, {0, -1, 1}, {1, -1, 1}, {-1, 0, 1}, {0, 0, 1}, + {1, 0, 1}, {-1, 1, 1}, {0, 1, 1}, {1, 1, 1}}; + for (int k = 0; k < 13; k++) { + const size_t j = wrap(u + off[k][0], nu) + + static_cast(nu) * (wrap(v + off[k][1], nv) + + static_cast(nv) * wrap(w + off[k][2], nw)); + if (mask[j] == 1.0f) + unite(label, static_cast(i), static_cast(j)); + } +} + +// The roots, read-only and into a separate array: writing them over the labels while other threads +// still walk the labels would let a path-halving write replace a written root with a mere ancestor. +__global__ void uf_roots(const int *label, int *root, size_t n) { + const size_t i = blockIdx.x * static_cast(blockDim.x) + threadIdx.x; + if (i >= n) + return; + int r = label[i]; + if (r >= 0) + while (label[r] != r) + r = label[r]; + root[i] = r; +} + +// Component sizes, counted only while they are at most limit: every island is counted in full, and the +// one large component stops being counted almost at once instead of queueing millions of atomics on one +// address. Only size <= limit is ever read, which the capping does not change. +__global__ void uf_count(const int *root, int *size, int limit, size_t n) { + const size_t i = blockIdx.x * static_cast(blockDim.x) + threadIdx.x; + if (i >= n) + return; + const int r = root[i]; + if (r >= 0 && reinterpret_cast(size)[r] <= limit) + atomicAdd(&size[r], 1); +} + +__global__ void uf_remove(float *mask, const int *root, const int *size, int limit, size_t n) { + const size_t i = blockIdx.x * static_cast(blockDim.x) + threadIdx.x; + if (i < n && root[i] >= 0 && size[root[i]] <= limit) + mask[i] = 0.0f; +} + +unsigned int Blocks(size_t n) { + return static_cast((n + THREADS - 1) / THREADS); +} + +} // namespace + +size_t ModelMaskGPU::DeviceBytes(size_t max_points) { + return 2 * max_points * sizeof(int) + MAX_OPS * sizeof(ModelMaskOp); +} + +ModelMaskGPU::ModelMaskGPU(cudaStream_t stream, size_t max_points) + : stream(stream), max_points(max_points), ops_d(MAX_OPS), label(max_points), root(max_points) { + if (max_points > INT32_MAX) + throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, + "Solvent mask grid too large for int32 labels"); +} + +void ModelMaskGPU::SetGrid(const ModelMaskGrid &grid, const std::vector &ops) { + npoints = static_cast(grid.nu) * grid.nv * grid.nw; + if (npoints > max_points) + throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, + "Solvent mask grid larger than the buffers"); + if (ops.empty() || ops.size() > MAX_OPS) + throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, + "Solvent mask needs 1 to " + std::to_string(MAX_OPS) + " operators"); + + const int n[3] = {grid.nu, grid.nv, grid.nw}; + geom.nu = grid.nu; + geom.nv = grid.nv; + geom.nw = grid.nw; + for (int i = 0; i < 3; i++) + for (int j = 0; j < 3; j++) + geom.orth_n[3 * i + j] = static_cast(grid.orth[3 * i + j] / n[j]); + + // gemmi Grid::calculate_spacing(): 1 / (n |a*|), where a* is row k of orth's inverse, which is the + // cross product of the other two columns over the determinant (= volume). + const double *o = grid.orth; + const double col[3][3] = {{o[0], o[3], o[6]}, {o[1], o[4], o[7]}, {o[2], o[5], o[8]}}; + for (int k = 0; k < 3; k++) { + const double *p = col[(k + 1) % 3], *q = col[(k + 2) % 3]; + const double cx = p[1] * q[2] - p[2] * q[1], cy = p[2] * q[0] - p[0] * q[2], cz = p[0] * q[1] - p[1] * q[0]; + geom.spacing[k] = grid.volume / (n[k] * std::sqrt(cx * cx + cy * cy + cz * cz)); + } + + // gemmi's shrink, set_margin_around(r_shrink), only ever looks at the grid offsets other than 0 + // within r_shrink, in the box of floor(r_shrink / spacing). No such offset, no change. + const int d[3] = {static_cast(std::floor(R_SHRINK / geom.spacing[0])), + static_cast(std::floor(R_SHRINK / geom.spacing[1])), + static_cast(std::floor(R_SHRINK / geom.spacing[2]))}; + for (int w = -d[2]; w <= d[2]; w++) + for (int v = -d[1]; v <= d[1]; v++) + for (int u = -d[0]; u <= d[0]; u++) { + const double f[3] = {static_cast(u) / n[0], static_cast(v) / n[1], + static_cast(w) / n[2]}; + double r2 = 0; + for (int i = 0; i < 3; i++) { + const double x = o[3 * i] * f[0] + o[3 * i + 1] * f[1] + o[3 * i + 2] * f[2]; + r2 += x * x; + } + if (r2 <= R_SHRINK * R_SHRINK && r2 != 0.0) + throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, + "Solvent mask grid too fine: the GPU mask has no shrink step"); + } + + // gemmi SolventMasker::remove_islands(), rounded the same way. + island_limit = static_cast(static_cast(ISLAND_MIN_VOLUME * npoints / grid.volume)); + + n_ops = static_cast(ops.size()); + cuda_err(cudaMemcpyAsync(ops_d, ops.data(), ops.size() * sizeof(ModelMaskOp), cudaMemcpyHostToDevice, stream)); +} + +void ModelMaskGPU::Compute(const ModelMaskAtom *d_atoms, int n_atoms, float *d_mask) { + fill_solvent<<>>(d_mask, npoints); + cuda_err(cudaGetLastError()); + if (n_atoms > 0) { + mask_images<<(n_atoms) * n_ops, IMAGE_THREADS, 0, stream>>>( + d_atoms, ops_d, n_ops, geom, d_mask); + cuda_err(cudaGetLastError()); + } + RemoveIslands(d_mask); +} + +// gemmi SolventMasker::remove_islands(): every component of solvent with at most island_limit points +// becomes macromolecule. +void ModelMaskGPU::RemoveIslands(float *d_mask) { + uf_init<<>>(d_mask, label, npoints); + cuda_err(cudaGetLastError()); + uf_merge<<>>(d_mask, label, geom.nu, geom.nv, geom.nw); + cuda_err(cudaGetLastError()); + uf_roots<<>>(label, root, npoints); + cuda_err(cudaGetLastError()); + int *size = label; // the parents are no longer needed + cuda_err(cudaMemsetAsync(size, 0, npoints * sizeof(int), stream)); + uf_count<<>>(root, size, island_limit, npoints); + cuda_err(cudaGetLastError()); + uf_remove<<>>(d_mask, root, size, island_limit, npoints); + cuda_err(cudaGetLastError()); +} diff --git a/rugnux/ModelMaskGPU.h b/rugnux/ModelMaskGPU.h new file mode 100644 index 000000000..c7210d9a8 --- /dev/null +++ b/rugnux/ModelMaskGPU.h @@ -0,0 +1,81 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#pragma once + +#include +#include + +#include + +#include "../image_analysis/indexing/CUDAMemHelpers.h" + +// The bulk-solvent mask of PutMaskOnGrid() (ModelGrid.h), i.e. gemmi's +// SolventMasker(AtomicRadiiSet::Refmac).put_mask_on_grid(), on the GPU. gemmi masks the atoms and then +// symmetrizes the grid with the minimum; the operators are isometries, so that is the same as masking +// every symmetry image of every atom, which is what is done here - no orbits, no symmetrize. The island +// removal is gemmi's (26-connected, periodic, the same size limit) as a union-find. The shrink is not +// implemented: it is a no-op on every rigid-body grid, and SetGrid() refuses a grid where it would not be. +// +// Every write of the masking is the same idempotent store and the union-find partition is unique, so +// the mask is bit-identical from run to run. It can differ from gemmi's at points lying exactly at an +// atom's radius, where float and double distances round differently. + +struct ModelMaskGrid { + int nu = 0, nv = 0, nw = 0; // index = u + nu * (v + nv * w) + double orth[9]; // gemmi UnitCell::orth.mat, row-major (Cartesian = orth * fractional) + double volume = 0; // UnitCell::volume, A^3 +}; + +// One atom to mask: fractional x, y, z wrapped into [0,1) and the mask radius in A. Double, because +// float fractional coordinates are off by up to 1e-5 A in a 200 A cell, enough to move points that lie at +// the radius; the distances themselves are float. (CUDA's double4 is deprecated from CUDA 13 and its +// replacement does not exist before, hence a struct of our own.) +struct ModelMaskAtom { + double x, y, z, radius; +}; + +// A fractional operator x' = rot * x + tran: every symmetry operator combined with every centring +// vector, identity included. +struct ModelMaskOp { + double rot[9]; // row-major + double tran[3]; +}; + +class ModelMaskGPU { +public: + // Enough for Fm-3m, 48 operators times 4 centring vectors. + static constexpr int MAX_OPS = 192; + + static size_t DeviceBytes(size_t max_points); + ModelMaskGPU(cudaStream_t stream, size_t max_points); + + // Per zone. Throws if the grid has more than max_points points, if there are more than MAX_OPS + // operators, or if gemmi's shrink (r_shrink = 0.8 A) would change anything on this grid. + void SetGrid(const ModelMaskGrid &grid, const std::vector &ops); + + // d_atoms: hydrogens and unoccupied atoms already left out. d_mask: the whole grid, 1 = solvent, 0 = macromolecule. Queued on the stream. + void Compute(const ModelMaskAtom *d_atoms, int n_atoms, float *d_mask); + + // The island removal alone, on a mask of 0 and 1 already on the grid. Compute() ends with it. + void RemoveIslands(float *d_mask); + + // What the masking kernel needs of the grid. + struct Geometry { + int nu, nv, nw; + float orth_n[9]; // orth * diag(1/nu, 1/nv, 1/nw): Cartesian of a grid-step offset + double spacing[3]; // gemmi Grid::spacing, the interplanar distance of the grid planes + }; + +private: + cudaStream_t stream; + size_t max_points; + Geometry geom{}; + size_t npoints = 0; + int n_ops = 0; + int island_limit = 0; + + CudaDevicePtr ops_d; + CudaDevicePtr label; // union-find parent, then the component sizes + CudaDevicePtr root; // each solvent point's component, -1 elsewhere +}; diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 7ce420525..83fbd99b8 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -135,6 +135,7 @@ ADD_EXECUTABLE(jfjoch_test TCPImagePusherTest.cpp PreflightTest.cpp ModelValidationTest.cpp + ModelMaskGPUTest.cpp SearchSpaceGroupTest.cpp SearchSpaceGroupTwinTest.cpp SyntheticMergedReflections.h diff --git a/tests/ModelMaskGPUTest.cpp b/tests/ModelMaskGPUTest.cpp new file mode 100644 index 000000000..793187cd6 --- /dev/null +++ b/tests/ModelMaskGPUTest.cpp @@ -0,0 +1,214 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#include +#include "../common/CUDAWrapper.h" + +#ifdef JFJOCH_USE_CUDA + +#include +#include +#include +#include +#include + +#include +#include +#include + +#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 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 MaskGrid(const gemmi::Structure &st, double d_min) { + gemmi::Grid 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 MaskAtoms(const gemmi::Structure &st) { + const gemmi::SolventMasker masker(gemmi::AtomicRadiiSet::Refmac); + std::vector 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 MaskOps(const gemmi::SpaceGroup &sg) { + const gemmi::GroupOps gops = sg.operations(); + std::vector 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(op.rot[i][j]) / gemmi::Op::DEN; + m.tran[i] = static_cast(op.tran[i] + cen[i]) / gemmi::Op::DEN; + } + ops.push_back(m); + } + return ops; +} + +ModelMaskGrid MaskGridDescription(const gemmi::Grid &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 &grid, const std::vector &atoms, + const std::vector &ops, int u, int v, int w) { + const double p[3] = {static_cast(u) / grid.nu, static_cast(v) / grid.nv, + static_cast(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 atoms = MaskAtoms(st); + const std::vector ops = MaskOps(*sg); + CudaDevicePtr 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 pre = MaskGrid(st, d_min); + masker.clear(pre); + masker.mask_points(pre, st.models[0]); + masker.symmetrize(pre); + gemmi::Grid 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 mask_d(n); + std::vector 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(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(i % ref.nu), v = static_cast(i / ref.nu % ref.nv), + w = static_cast(i / (static_cast(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 ops = MaskOps(*st.find_spacegroup()); + const gemmi::Grid coarse = MaskGrid(st, 3.5); + const gemmi::Grid 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