Jungfraujoch/tests/ModelMaskGPUTest.cpp

// 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