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:
leonarski_fandClaude Opus 5.5 committed
1 parent
9f00fd1f4d
commit
63fde9246a
6 files changed
+574
No files matched your search
@@ -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).
|
||||
|
||||
@@ -30,6 +30,7 @@ ADD_LIBRARY(JFJochRugnux STATIC
|
||||
HotPixels.cpp
|
||||
HotPixels.h
|
||||
$<$<BOOL:${JFJOCH_CUDA_AVAILABLE}>:HotPixelsGPU.cu>
|
||||
$<$<BOOL:${JFJOCH_CUDA_AVAILABLE}>:ModelMaskGPU.cu>
|
||||
HotPixelsGPU.h
|
||||
DiagnosticOutput.cpp
|
||||
DiagnosticOutput.h
|
||||
|
||||
@@ -0,0 +1,272 @@
|
||||
// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
|
||||
// SPDX-License-Identifier: GPL-3.0-only
|
||||
|
||||
#include "ModelMaskGPU.h"
|
||||
|
||||
#include <cmath>
|
||||
#include <cstdint>
|
||||
#include <string>
|
||||
|
||||
#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<size_t>(blockDim.x) + threadIdx.x;
|
||||
if (i < n)
|
||||
mask[i] = 1.0f;
|
||||
}
|
||||
|
||||
// One block per image (atom, operator): gemmi's use_points_around<true>(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<int>(round(nx));
|
||||
const int d = min(static_cast<int>(ceil(a.radius / g.spacing[k])), n - 1);
|
||||
lo[k] = centre - d;
|
||||
len[k] = 2 * d + 1;
|
||||
c[k] = static_cast<float>(nx - centre) + static_cast<float>(d);
|
||||
}
|
||||
__syncthreads();
|
||||
const float r2 = static_cast<float>(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<size_t>(g.nu) * (wrap(lo[1] + iv, g.nv)
|
||||
+ static_cast<size_t>(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<size_t>(blockDim.x) + threadIdx.x;
|
||||
if (i < n)
|
||||
label[i] = mask[i] == 1.0f ? static_cast<int>(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<size_t>(nu) * nv * nw;
|
||||
const size_t i = blockIdx.x * static_cast<size_t>(blockDim.x) + threadIdx.x;
|
||||
if (i >= n || mask[i] != 1.0f)
|
||||
return;
|
||||
const int u = static_cast<int>(i % nu);
|
||||
const int v = static_cast<int>((i / nu) % nv);
|
||||
const int w = static_cast<int>(i / (static_cast<size_t>(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<size_t>(nu) * (wrap(v + off[k][1], nv)
|
||||
+ static_cast<size_t>(nv) * wrap(w + off[k][2], nw));
|
||||
if (mask[j] == 1.0f)
|
||||
unite(label, static_cast<int>(i), static_cast<int>(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<size_t>(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<size_t>(blockDim.x) + threadIdx.x;
|
||||
if (i >= n)
|
||||
return;
|
||||
const int r = root[i];
|
||||
if (r >= 0 && reinterpret_cast<volatile int *>(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<size_t>(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<unsigned int>((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<ModelMaskOp> &ops) {
|
||||
npoints = static_cast<size_t>(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<float>(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<int>(std::floor(R_SHRINK / geom.spacing[0])),
|
||||
static_cast<int>(std::floor(R_SHRINK / geom.spacing[1])),
|
||||
static_cast<int>(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<double>(u) / n[0], static_cast<double>(v) / n[1],
|
||||
static_cast<double>(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<int>(static_cast<size_t>(ISLAND_MIN_VOLUME * npoints / grid.volume));
|
||||
|
||||
n_ops = static_cast<int>(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<<<Blocks(npoints), THREADS, 0, stream>>>(d_mask, npoints);
|
||||
cuda_err(cudaGetLastError());
|
||||
if (n_atoms > 0) {
|
||||
mask_images<<<static_cast<unsigned int>(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<<<Blocks(npoints), THREADS, 0, stream>>>(d_mask, label, npoints);
|
||||
cuda_err(cudaGetLastError());
|
||||
uf_merge<<<Blocks(npoints), THREADS, 0, stream>>>(d_mask, label, geom.nu, geom.nv, geom.nw);
|
||||
cuda_err(cudaGetLastError());
|
||||
uf_roots<<<Blocks(npoints), THREADS, 0, stream>>>(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<<<Blocks(npoints), THREADS, 0, stream>>>(root, size, island_limit, npoints);
|
||||
cuda_err(cudaGetLastError());
|
||||
uf_remove<<<Blocks(npoints), THREADS, 0, stream>>>(d_mask, root, size, island_limit, npoints);
|
||||
cuda_err(cudaGetLastError());
|
||||
}
|
||||
@@ -0,0 +1,81 @@
|
||||
// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
|
||||
// SPDX-License-Identifier: GPL-3.0-only
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <cstddef>
|
||||
#include <vector>
|
||||
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
#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<ModelMaskOp> &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<ModelMaskOp> ops_d;
|
||||
CudaDevicePtr<int> label; // union-find parent, then the component sizes
|
||||
CudaDevicePtr<int> root; // each solvent point's component, -1 elsewhere
|
||||
};
|
||||
@@ -135,6 +135,7 @@ ADD_EXECUTABLE(jfjoch_test
|
||||
TCPImagePusherTest.cpp
|
||||
PreflightTest.cpp
|
||||
ModelValidationTest.cpp
|
||||
ModelMaskGPUTest.cpp
|
||||
SearchSpaceGroupTest.cpp
|
||||
SearchSpaceGroupTwinTest.cpp
|
||||
SyntheticMergedReflections.h
|
||||
|
||||
@@ -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