Jungfraujoch/tests/MolecularReplacementTest.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 <atomic>
#include <cmath>
#include <cstdio>
#include <cstring>
#include <fstream>
#include <random>

#include <gemmi/dencalc.hpp>
#include <gemmi/it92.hpp>

#include "../image_analysis/structure_refinement/ModelFFT.h"
#include "../image_analysis/structure_refinement/MolecularReplacement.h"
#include "../image_analysis/structure_refinement/RigidBodyRefine.h"
#include "../common/Logger.h"
#ifdef JFJOCH_USE_CUDA
#include "../common/CUDAWrapper.h"
#include "../image_analysis/structure_refinement/MolecularReplacementGPU.h"
#include "../image_analysis/structure_refinement/RigidBodyGPU.h"
#endif

namespace {
    // A synthetic body: carbon atoms scattered in a 12 A sphere, centroid at the origin.
    SearchModel synthetic_model() {
        SearchModel m;
        std::mt19937 rng(7);
        std::uniform_real_distribution<double> u(-12.0, 12.0);
        gemmi::Model model(1);
        gemmi::Chain chain("A");
        gemmi::Residue res;
        res.name = "UNK";
        res.seqid = gemmi::SeqId(1, ' ');
        std::vector<gemmi::Vec3> pos;
        while (pos.size() < 300) {
            const gemmi::Vec3 p(u(rng), u(rng), u(rng));
            if (p.length() < 12.0)
                pos.push_back(p);
        }
        gemmi::Vec3 c;
        for (const auto &p : pos)
            c += p;
        c *= 1.0 / pos.size();
        for (const auto &p : pos) {
            gemmi::Atom a;
            a.name = "CA";   // every atom counts in the clash check
            a.element = gemmi::El::C;
            a.pos = gemmi::Position(p - c);
            a.b_iso = 15.0f;
            a.occ = 1.0f;
            res.atoms.push_back(a);
            m.radius_A = std::max(m.radius_A, (p - c).length());
        }
        chain.residues.push_back(res);
        model.chains.push_back(chain);
        m.st.models.push_back(model);
        m.residues_read = m.residues_kept = 1;
        m.mass_da = 3600;
        return m;
    }

    gemmi::Mat33 axis_angle(gemmi::Vec3 axis, double deg) {
        axis = axis.normalized();
        const double k = deg * M_PI / 180.0, s = std::sin(k), c1 = 1.0 - std::cos(k);
        const double x = axis.x, y = axis.y, z = axis.z;
        return {1 + c1 * (x * x - 1), -s * z + c1 * x * y, s * y + c1 * x * z,
                s * z + c1 * x * y, 1 + c1 * (y * y - 1), -s * x + c1 * y * z,
                -s * y + c1 * x * z, s * x + c1 * y * z, 1 + c1 * (z * z - 1)};
    }

    // |F|^2 of the placed crystal to 4 A, as merged reflections.
    std::vector<MergedReflection> observed(const gemmi::Structure &crystal) {
        gemmi::DensityCalculator<gemmi::IT92<float>, float> dc;
        dc.d_min = 3.5;
        dc.rate = 1.5;
        dc.set_grid_cell_and_spacegroup(crystal);
        dc.put_model_density_on_grid(crystal.models[0]);
        const auto f = MapToFPhi(dc.grid).prepare_asu_data(3.5, 0, false, false, false);
        std::vector<MergedReflection> out;
        for (const auto &hv : f.v) {
            MergedReflection r;
            r.h = hv.hkl[0];
            r.k = hv.hkl[1];
            r.l = hv.hkl[2];
            r.I = std::norm(hv.value);
            r.sigma = 1;
            r.F = std::abs(hv.value);
            r.d = 1.0 / std::sqrt(crystal.cell.calculate_1_d2(hv.hkl));
            out.push_back(r);
        }
        return out;
    }
}

namespace {
    constexpr size_t GB = 1000000000;

    // A search's parts, sized as a large cell measures: the transform 0.84 GB, rigid-body engines of 2.10 GB
    // (nine fits at once), score engines of up to 1.86 GB; the budget 80% of the card.
    MolecularReplacementGPUNeed LargeCellNeed(size_t card) {
        MolecularReplacementGPUNeed need;
        need.budget = card / 5 * 4;
        need.transform = 835500000;
        need.rigid_body = 2101000000;
        need.rigid_body_engines = 9;
        need.structure_factors = 1860000000;
        return need;
    }
}

TEST_CASE("MolecularReplacement_GPUPlanOnCardSizes", "[MolecularReplacement]") {
    size_t last_rb = 0;
    for (size_t card : {6 * GB, 8 * GB, 16 * GB, 24 * GB, 64 * GB}) {
        const MolecularReplacementGPUNeed need = LargeCellNeed(card);
        const MolecularReplacementGPUPlan plan = PlanMolecularReplacementGPU(need);
        CAPTURE(card, plan.rigid_body_engines, plan.bytes, plan.room);
        CHECK(plan.bytes <= need.budget);
        // Every part on the GPU on every one of these cards; only the engine count grows with it.
        CHECK(plan.search);
        CHECK(plan.structure_factors);
        CHECK(plan.rigid_body_engines >= std::max<size_t>(last_rb, 1));
        // What the transform and the engines leave is the room, and it holds the largest score engine.
        CHECK(plan.room == need.budget - need.transform - plan.rigid_body_engines * need.rigid_body);
        CHECK(plan.room >= need.structure_factors);
        last_rb = plan.rigid_body_engines;
    }
    // 16 GB, a 12.8 GB budget: one of each, 4.80 GB, and a second engine would pass half of it.
    CHECK(PlanMolecularReplacementGPU(LargeCellNeed(16 * GB)).rigid_body_engines == 1);
    CHECK(PlanMolecularReplacementGPU(LargeCellNeed(24 * GB)).rigid_body_engines == 3);
    CHECK(PlanMolecularReplacementGPU(LargeCellNeed(64 * GB)).rigid_body_engines == 9);
}

TEST_CASE("MolecularReplacement_GPUPlanOnOtherCards", "[MolecularReplacement]") {
    // Four cards of 15.6 GB, a 12.48 GB budget each: this card holds what it holds alone (one engine), each other
    // card a copy of the transform and, within half of its budget, two engines.
    const size_t card = 15600000000;
    const MolecularReplacementGPUPlan alone = PlanMolecularReplacementGPU(LargeCellNeed(card));
    MolecularReplacementGPUNeed need = LargeCellNeed(card);
    need.other_budgets = {need.budget, need.budget, need.budget};
    const MolecularReplacementGPUPlan plan = PlanMolecularReplacementGPU(need);
    CHECK(plan.search == alone.search);
    CHECK(plan.structure_factors == alone.structure_factors);
    CHECK(plan.room == alone.room);
    CHECK(plan.bytes == alone.bytes);
    REQUIRE(plan.cards.size() == 4);
    CHECK(plan.cards[0].transform);
    CHECK(plan.cards[0].rigid_body_engines == alone.rigid_body_engines);
    CHECK(plan.cards[0].room == alone.room);
    CHECK(alone.rigid_body_engines == 1);
    for (size_t k = 1; k < 4; k++) {
        CHECK(plan.cards[k].transform);
        CHECK(plan.cards[k].rigid_body_engines == 2);
        CHECK(plan.cards[k].room == need.budget - need.transform - 2 * need.rigid_body);
    }
    CHECK(plan.rigid_body_engines == 7);
    // No more engines than the search can use: with 64 GB cards, where this card alone has room for all nine, they
    // still go to both cards in turn after this one's first.
    need = LargeCellNeed(64 * GB);
    need.other_budgets = {need.budget};
    const MolecularReplacementGPUPlan large = PlanMolecularReplacementGPU(need);
    CHECK(large.rigid_body_engines == 9);
    CHECK(large.cards[0].rigid_body_engines == 5);
    CHECK(large.cards[1].rigid_body_engines == 4);
    CHECK(large.cards[1].transform);
    CHECK(large.room == need.budget - need.transform - 5 * need.rigid_body);
    // A search on the CPU puts nothing on the other cards either.
    need = LargeCellNeed(card);
    need.transform = need.budget + 1;
    need.other_budgets = {need.budget};
    const MolecularReplacementGPUPlan cpu = PlanMolecularReplacementGPU(need);
    CHECK(!cpu.search);
    CHECK(!cpu.cards[1].transform);
}

TEST_CASE("MolecularReplacement_GPUPlanPartThatDoesNotFitGoesToTheCPU", "[MolecularReplacement]") {
    // 6 GB, a 4.8 GB budget, which the three parts fill to 4.80 GB: with a larger score engine, the scores
    // go to the CPU and the transform and the rigid body stay.
    MolecularReplacementGPUNeed need = LargeCellNeed(6 * GB);
    need.structure_factors = 2 * GB;
    MolecularReplacementGPUPlan plan = PlanMolecularReplacementGPU(need);
    CHECK(plan.search);
    CHECK(plan.rigid_body_engines == 1);
    CHECK(!plan.structure_factors);
    CHECK(plan.bytes <= need.budget);

    // A part the GPU cannot grid (0 bytes) is never placed; the room is then what the transform leaves.
    need = LargeCellNeed(16 * GB);
    need.rigid_body = 0;
    plan = PlanMolecularReplacementGPU(need);
    CHECK(plan.rigid_body_engines == 0);
    CHECK(plan.room == need.budget - need.transform);

    // No budget (no GPU, or all of it held beside): everything on the CPU.
    need = LargeCellNeed(16 * GB);
    need.budget = 0;
    plan = PlanMolecularReplacementGPU(need);
    CHECK(!plan.search);
    CHECK(plan.rigid_body_engines == 0);
    CHECK(!plan.structure_factors);
    CHECK(plan.room == 0);
}

TEST_CASE("MolecularReplacement_GPUSmallEnginesSpreadOverTheCards", "[MolecularReplacement]") {
    // A small cell's search on four 15.6 GB cards: 24 fits at once of 0.09 GB engines, which this card alone has
    // room for. They go six to a card, in turn from the next card - the fits run on all four.
    MolecularReplacementGPUNeed need;
    need.budget = 12480000000;
    need.other_budgets = {need.budget, need.budget, need.budget};
    need.transform = 10000000;
    need.rigid_body = 90000000;
    need.rigid_body_engines = 24;
    need.structure_factors = 50000000;
    const MolecularReplacementGPUPlan plan = PlanMolecularReplacementGPU(need);
    REQUIRE(plan.cards.size() == 4);
    CHECK(plan.rigid_body_engines == 24);
    for (const MolecularReplacementGPUCard &c : plan.cards) {
        CHECK(c.transform);
        CHECK(c.rigid_body_engines == 6);
    }
    CHECK(plan.room == need.budget - need.transform - 6 * need.rigid_body);
    CHECK(plan.bytes == need.transform + 6 * need.rigid_body + need.structure_factors);
    // One card: all 24 on it, as before.
    need.other_budgets.clear();
    CHECK(PlanMolecularReplacementGPU(need).rigid_body_engines == 24);
}

TEST_CASE("MolecularReplacement_GPUPackedTransformWhereItIsExpectedFaster", "[MolecularReplacement]") {
    // The large cell's transform packed, 2.28 GB, and its rotation function: 17255 orientations x 244183
    // reflections x 4 operators, the search the constants were measured on.
    const size_t packed = 2276000000;
    const double large_search = 17255.0 * 244183.0 * 4.0;
    auto need_on = [&](size_t card, size_t others, double interpolations) {
        MolecularReplacementGPUNeed need = LargeCellNeed(card);
        need.other_budgets.assign(others, need.budget);
        need.rotation_interpolations = interpolations;
        return need;
    };
    // One card, where packing keeps the plan's engines: 75.8 s of rotation function down to 15.1 s, packed.
    MolecularReplacementGPUNeed one = need_on(15600000000, 0, large_search);
    const MolecularReplacementGPUPlan one_plan = PlanMolecularReplacementGPU(one);
    CHECK(MolecularReplacementGPUSeconds(one, one_plan, false) == Catch::Approx(75.8 + 19.7).epsilon(1e-6));
    CHECK(PackMolecularTransformGPU(one, packed));
    CHECK(PackMolecularTransformGPU(need_on(8 * GB, 0, large_search), packed));
    // On a 6 GB card the score engine would no longer fit beside it: plain, however fast packing would be.
    CHECK(!PackMolecularTransformGPU(need_on(6 * GB, 0, large_search), packed));

    // Four 15.6 GB cards: packed, each copy leaves room for one engine in place of two (4 engines, not 7), but the
    // rotation function over four cards goes from about 19 s to about 4 s, more than the engines cost.
    MolecularReplacementGPUNeed four = need_on(15600000000, 3, large_search);
    MolecularReplacementGPUNeed four_packed = four;
    four_packed.transform = packed;
    const MolecularReplacementGPUPlan plain_plan = PlanMolecularReplacementGPU(four);
    const MolecularReplacementGPUPlan packed_plan = PlanMolecularReplacementGPU(four_packed);
    CHECK(plain_plan.rigid_body_engines == 7);
    CHECK(packed_plan.rigid_body_engines == 4);
    const double plain_s = MolecularReplacementGPUSeconds(four, plain_plan, false);
    const double packed_s = MolecularReplacementGPUSeconds(four_packed, packed_plan, true);
    CAPTURE(plain_s, packed_s);
    CHECK(plain_s == Catch::Approx(75.8 / 4 + 19.7 / 7).epsilon(1e-6));
    CHECK(packed_s == Catch::Approx(15.1 / 4 + 19.7 / 4).epsilon(1e-6));
    CHECK(PackMolecularTransformGPU(four, packed));
    // A search a hundred times smaller with the same engines: the rotation function no longer pays for them.
    CHECK(!PackMolecularTransformGPU(need_on(15600000000, 3, large_search / 100), packed));
    // On one 24 GB card, the same: packed holds two engines in place of three.
    CHECK(PackMolecularTransformGPU(need_on(24 * GB, 0, large_search), packed));
    CHECK(!PackMolecularTransformGPU(need_on(24 * GB, 0, large_search / 100), packed));

    // The rigid body on the CPU in both layouts: decided on the rotation function alone.
    MolecularReplacementGPUNeed no_rigid_body = need_on(15600000000, 3, large_search / 100);
    no_rigid_body.rigid_body = 0;
    CHECK(PackMolecularTransformGPU(no_rigid_body, packed));
    // Where packing would put the rigid body on the CPU, plain: a budget just short of the packed transform and an
    // engine (the scores on the CPU in both).
    MolecularReplacementGPUNeed tight = need_on(15600000000, 0, large_search);
    tight.structure_factors = 0;
    tight.budget = packed + tight.rigid_body - 1;
    CHECK(PlanMolecularReplacementGPU(tight).rigid_body_engines == 1);
    CHECK(!PackMolecularTransformGPU(tight, packed));
    // No budget: nothing on the GPU, so nothing packed.
    MolecularReplacementGPUNeed none = need_on(16 * GB, 0, large_search);
    none.budget = 0;
    CHECK(!PackMolecularTransformGPU(none, packed));
}

TEST_CASE("MolecularReplacement_UniformRotationsCoverSO3", "[MolecularReplacement]") {
    const auto rotations = UniformRotations(6.0);
    const double step = 6.0 * M_PI / 180.0;
    // 8 pi^2 is the volume of SO(3) in the measure a 6 deg cube of Euler-like angles samples.
    const double expected = 8 * M_PI * M_PI / (step * step * step);
    CHECK(rotations.size() > 0.8 * expected);
    CHECK(rotations.size() < 1.2 * expected);
    // Every rotation has a sampled one nearby.
    std::mt19937 rng(11);
    std::normal_distribution<double> g;
    const std::vector<gemmi::Mat33> identity{gemmi::Mat33()};
    double worst = 0;
    for (int i = 0; i < 100; i++) {
        const gemmi::Mat33 r = axis_angle({g(rng), g(rng), g(rng)}, std::uniform_real_distribution<double>(0, 180)(rng));
        double best = 180;
        for (const auto &s : rotations)
            best = std::min(best, RotationDistanceDeg(r, s, identity));
        worst = std::max(worst, best);
    }
    CHECK(worst < 6.0);
}

TEST_CASE("MolecularReplacement_RotationAsu", "[MolecularReplacement]") {
    const gemmi::UnitCell cell(40, 40, 60, 90, 90, 90);
    const auto pg = CartesianPointGroup(gemmi::get_spacegroup_by_name("P 41 2 2"), cell);
    REQUIRE(pg.size() == 8);
    const auto all = UniformRotations(10.0);
    const auto asu = RotationAsu(all, pg);
    // One class in eight, give or take the boundaries.
    CHECK(asu.size() > all.size() / 8 * 0.9);
    CHECK(asu.size() < all.size() / 8 * 1.3);
    // Every sampled rotation has a class member in the ASU close by.
    for (size_t i = 0; i < all.size(); i += 37) {
        double best = 180;
        for (const auto &s : asu)
            best = std::min(best, RotationDistanceDeg(s, all[i], pg));
        CHECK(best < 10.0);
    }
}

TEST_CASE("MolecularReplacement_SameRotationGroups", "[MolecularReplacement]") {
    const auto groups = SameRotationGroups(gemmi::get_spacegroup_by_name("P 41"));
    REQUIRE(groups.size() == 4);
    CHECK(groups[0].xhm() == "P 41");
    const auto i = SameRotationGroups(gemmi::get_spacegroup_by_name("I 41 2 2"));
    CHECK(i.size() == 2);   // I 4 2 2, I 41 2 2
    // Each group once, not again on another origin: P 2 2 2, P 2 2 21 in three settings, P 21 21 2 in
    // three, P 21 21 21.
    CHECK(SameRotationGroups(gemmi::get_spacegroup_by_name("P 21 21 2")).size() == 8);
    CHECK(SameRotationGroups(gemmi::get_spacegroup_by_name("I 2 2 2")).size() == 2);
    CHECK(SameRotationGroups(gemmi::get_spacegroup_by_name("P 42 21 2")).size() == 8);
}

TEST_CASE("MolecularReplacement_SyntheticPlacement", "[MolecularReplacement]") {
    const SearchModel model = synthetic_model();
    const gemmi::UnitCell cell(42, 50, 58, 90, 90, 90);
    const gemmi::SpaceGroup &sg = gemmi::get_spacegroup_by_name("P 21 21 21");
    const gemmi::Mat33 truth = axis_angle({1, 2, 3}, 50);
    const gemmi::Fractional t(0.17, 0.31, 0.09);
    const gemmi::Structure crystal = PlaceSearchModel(model, truth, t, cell, sg);
    const auto merged = observed(crystal);
    const SearchReflections refl = MakeSearchReflections(merged, cell, sg, 4.0, 12.0);
    REQUIRE(refl.hkl.size() > 300);
    const MolecularTransform ft(model, 4.0, 4);

    // The rotation function's best orientation is the true one, to the sampling.
    const auto pg = CartesianPointGroup(sg, cell);
    const auto rotations = RotationAsu(UniformRotations(6.0), pg);
    const auto rf = RotationFunction(ft, refl, cell, sg, rotations, 4);
    size_t best = 0;
    for (size_t i = 1; i < rf.size(); i++)
        if (rf[i] > rf[best])
            best = i;
    CHECK(RotationDistanceDeg(rotations[best], truth, pg) < 6.0);

    // At the true orientation the translation function peaks at the true position, up to the
    // origins P 21 21 21 allows (a half in each coordinate), and stands far out of its own noise.
    const TranslationPeak peak = TranslationFunction(ft, MakeTranslationSetup(refl, cell, sg), truth);
    CHECK(peak.z > 10.0);
    for (double d : {peak.t.x - t.x, peak.t.y - t.y, peak.t.z - t.z}) {
        const double half = 2 * d - std::round(2 * d);
        CHECK(std::fabs(half) < 0.06);
    }
}

// Two copies in the asymmetric unit: with the first one known, the translation function of the second
// at its true orientation peaks at its true position - not merely at one of the group's origins, which
// only the cross term with the known copy tells apart.
TEST_CASE("MolecularReplacement_SecondCopyBesideKnownFirst", "[MolecularReplacement]") {
    const SearchModel model = synthetic_model();
    const gemmi::UnitCell cell(52, 60, 70, 90, 90, 90);
    const gemmi::SpaceGroup &sg = gemmi::get_spacegroup_by_name("P 21 21 21");
    const gemmi::Mat33 r1 = axis_angle({1, 2, 3}, 50), r2 = axis_angle({-2, 1, 1}, 110);
    const gemmi::Fractional t1(0.17, 0.31, 0.09), t2(0.41, 0.12, 0.33);
    const gemmi::Structure crystal = AssembleCopies({PlaceSearchModel(model, r1, t1, cell, sg),
                                                     PlaceSearchModel(model, r2, t2, cell, sg)});
    const auto merged = observed(crystal);
    const SearchReflections refl = MakeSearchReflections(merged, cell, sg, 4.0, 12.0);
    const MolecularTransform ft(model, 4.0, 4);
    const TranslationSetup setup = MakeTranslationSetup(refl, cell, sg, true);
    const auto known = KnownStructureFactors(ft, setup, cell, {r1}, {t1});
    const TranslationPeak peak = TranslationFunction(ft, setup, r2, &known);
    CHECK(peak.z > 10.0);
    for (double d : {peak.t.x - t2.x, peak.t.y - t2.y, peak.t.z - t2.z})
        CHECK(std::fabs(d - std::round(d)) < 0.04);
    // Two copies of the same body in one cell do not clash where they are placed apart.
    CHECK(ClashFraction(crystal, 2) < 0.05);
    CHECK(ClashFraction(AssembleCopies({PlaceSearchModel(model, r1, t1, cell, sg),
                                        PlaceSearchModel(model, r2, t1, cell, sg)}), 2) > 0.5);
}

// The likelihood of the true placement stands out over the same orientation elsewhere, and the gain is zero
// where sigma_A is.
TEST_CASE("MolecularReplacement_LikelihoodPlacement", "[MolecularReplacement]") {
    const SearchModel model = synthetic_model();
    const gemmi::UnitCell cell(42, 50, 58, 90, 90, 90);
    const gemmi::SpaceGroup &sg = gemmi::get_spacegroup_by_name("P 21 21 21");
    const gemmi::Mat33 truth = axis_angle({1, 2, 3}, 50);
    const gemmi::Fractional t(0.17, 0.31, 0.09);
    const auto merged = observed(PlaceSearchModel(model, truth, t, cell, sg));
    const LikelihoodSetup setup = MakeLikelihoodSetup(MakeSearchReflections(merged, cell, sg, 3.5, 12.0), cell, sg);
    REQUIRE(setup.Terms() > 500);
    const std::vector<std::complex<double>> none(setup.Terms(), 0.0);
    const auto g = CopyFactors(model, setup, cell, truth, 3.5, 4);

    const LikelihoodPeak peak = LikelihoodTranslationZ(setup, none, 0.0, g, t, 0.3, 4);
    CHECK(peak.llg > 100.0);
    CHECK(peak.z > 8.0);
    CHECK(peak.scale > 0.8);
    const double wrong = LikelihoodGain(setup, none, 0.0, g, gemmi::Fractional(0.4, 0.05, 0.3), 0.3);
    CHECK(wrong < peak.llg / 4);
    std::vector<std::complex<double>> f = none;
    AddCopy(setup, g, t, f);
    CHECK(LogLikelihoodGain(setup, f, 0.3, 0.0) == Catch::Approx(0.0).margin(1e-9));
}

TEST_CASE("MolecularReplacement_ModelKind", "[MolecularReplacement]") {
    // Three alanines along x, written as PDB with a CRYST1 line and the B column given per atom.
    const auto write = [](const std::string &path, const char *cryst1, const float b[3], bool water, bool altloc) {
        std::ofstream out(path);
        out << cryst1 << "\n";
        int serial = 1;
        for (int r = 0; r < 3; r++)
            for (const char *name : {" N  ", " CA ", " C  ", " O  "}) {
                char line[128];
                const char alt = altloc && r == 1 ? 'A' : ' ';
                std::snprintf(line, sizeof line, "ATOM  %5d %4s%cALA A%4d    %8.3f%8.3f%8.3f%6.2f%6.2f           %c",
                              serial++, name, alt, r + 1, 3.8 * r + (name[1] == 'C' && name[2] == ' ' ? 1.5 : 0.0),
                              name[1] == 'O' ? 1.2 : 0.0, 0.0, 1.0, b[r], name[1]);
                out << line << "\n";
                if (alt == 'A') {
                    std::snprintf(line, sizeof line, "ATOM  %5d %4sBALA A%4d    %8.3f%8.3f%8.3f%6.2f%6.2f           %c",
                                  serial++, name, r + 1, 3.8 * r + 0.3, 0.3, 0.0, 0.5, b[r], name[1]);
                    out << line << "\n";
                }
            }
        if (water)
            out << "HETATM  100  O   HOH A 101      10.000  10.000  10.000  1.00 30.00           O\n";
        out << "END\n";
    };
    const float plddt[3] = {90, 85, 95};
    write("MolecularReplacement_ModelKind_af.pdb", "CRYST1    1.000    1.000    1.000  90.00  90.00  90.00 P 1           1",
          plddt, false, false);
    const SearchModel af = PrepareSearchModel("MolecularReplacement_ModelKind_af.pdb");
    CHECK(af.failure_reason.empty());
    CHECK(af.predicted);
    CHECK(af.residues_kept == 3);

    // In a cell, with refined B-factors (one above 100), a water and an alternative conformer.
    const float refined[3] = {25, 112, 40};
    write("MolecularReplacement_ModelKind_xtal.pdb", "CRYST1   50.000   60.000   70.000  90.00  90.00  90.00 P 21 21 21    4",
          refined, true, true);
    const SearchModel x = PrepareSearchModel("MolecularReplacement_ModelKind_xtal.pdb");
    CHECK(x.failure_reason.empty());
    CHECK(!x.predicted);
    CHECK(x.residues_kept == 3);
    size_t atoms = 0;
    for (const gemmi::Chain &ch : x.st.models[0].chains)
        for (const gemmi::Residue &r : ch.residues)
            atoms += r.atoms.size();
    CHECK(atoms == 12);   // no water, one conformer
    CHECK(x.error_A == Catch::Approx(0.67).margin(0.01));
    std::remove("MolecularReplacement_ModelKind_af.pdb");
    std::remove("MolecularReplacement_ModelKind_xtal.pdb");
}

TEST_CASE("MolecularReplacement_PaeDomains", "[MolecularReplacement]") {
    CHECK(PaePathFor("/x/AF-P12345-F1-model_v6.cif") == "/x/AF-P12345-F1-predicted_aligned_error_v6.json");
    CHECK(PaePathFor("AF-P12345-F1-model_v4.pdb.gz") == "AF-P12345-F1-predicted_aligned_error_v4.json");
    CHECK(PaePathFor("/x/model.cif").empty());

    // A prediction of 120 residues: two confident blocks, 1-50 and 61-120, placed relative to each other
    // with no confidence, and a loop 51-60 the trim removed.
    SearchModel m;
    gemmi::Model model(1);
    gemmi::Chain chain("A");
    for (int i = 0; i < 120; i++) {
        if (i >= 50 && i < 60)
            continue;
        gemmi::Residue r;
        r.name = "ALA";
        r.seqid = gemmi::SeqId(i + 1, ' ');
        gemmi::Atom a;
        a.name = "CA";
        a.element = gemmi::El::C;
        a.pos = gemmi::Position(3.8 * i, i < 50 ? 0.0 : 20.0, 0.0);
        a.b_iso = 10.0f;
        a.occ = 1.0f;
        r.atoms.push_back(a);
        chain.residues.push_back(r);
        m.residue_index.push_back(i);
    }
    model.chains.push_back(chain);
    m.st.models.push_back(model);
    m.residues_read = 120;
    m.residues_kept = 110;
    const std::string path = "MolecularReplacement_PaeDomains.json";
    {
        std::ofstream out(path);
        out << "[{\"predicted_aligned_error\":[";
        for (int i = 0; i < 120; i++) {
            out << (i ? "," : "") << "[";
            for (int j = 0; j < 120; j++)
                out << (j ? "," : "") << ((i < 55) == (j < 55) ? 2 : 25);
            out << "]";
        }
        out << "],\"max_predicted_aligned_error\":31.75}]";
    }
    const auto domains = PaeDomains(path, m);
    REQUIRE(domains.size() == 2);
    CHECK(domains[0].size() == 60);
    CHECK(domains[0].front() == 60);
    CHECK(domains[1].size() == 50);
    const SearchModel d = DomainModel(m, domains[1]);
    CHECK(d.residues_kept == 50);
    CHECK(d.description == "PAE domain of 50 residues (1-50)");
    std::remove(path.c_str());
}

// A placement of a group carried into a subgroup's merge: in P 21 21 21 read as P 1 21 1 on the same cell each
// copy becomes two, the second its image under a two-fold the subgroup does not have; in C 2 2 21 read as
// P 1 21 1 on the primitive monoclinic cell half the volume holds the same two.
TEST_CASE("MolecularReplacement_ExpandIntoSubgroup", "[MolecularReplacement]") {
    const SearchModel model = synthetic_model();
    {
        const gemmi::UnitCell cell(42, 50, 58, 90, 90, 90);
        const gemmi::SpaceGroup &g = gemmi::get_spacegroup_by_name("P 21 21 21");
        const gemmi::SpaceGroup &h = gemmi::get_spacegroup_by_name("P 1 21 1");
        const auto setting = SubgroupBasis(cell, g, cell, h);
        REQUIRE(setting);
        const gemmi::Structure placed = PlaceSearchModel(model, axis_angle({1, 2, 3}, 50), gemmi::Fractional(0.17, 0.31, 0.09), cell, g);
        int n = 0;
        const gemmi::Structure expanded = ExpandIntoSubgroup(placed, 1, g, *setting, cell, h, &n);
        CHECK(n == 2);
        // The same crystal: the copies and their P 1 21 1 mates are the P 21 21 21 crystal, which packs.
        CHECK(ClashFraction(expanded, 2) < 0.05);
    }
    {
        const gemmi::UnitCell g_cell(56, 158, 144, 90, 90, 90);
        const gemmi::SpaceGroup &g = gemmi::get_spacegroup_by_name("C 2 2 21");
        const double half = std::sqrt(56.0 * 56.0 + 158.0 * 158.0) / 2;
        const double beta = 180.0 - std::acos(28.0 / half) * 180.0 / M_PI;
        const gemmi::UnitCell h_cell(half, 144, 56, 90, beta, 90);
        const gemmi::SpaceGroup &h = gemmi::get_spacegroup_by_name("P 1 21 1");
        const auto setting = SubgroupBasis(g_cell, g, h_cell, h);
        REQUIRE(setting);
        CHECK(std::fabs(setting->basis.determinant() - 0.5) < 1e-9);
        const gemmi::Structure placed = PlaceSearchModel(model, axis_angle({1, 2, 3}, 50), gemmi::Fractional(0.17, 0.31, 0.09), g_cell, g);
        int n = 0;
        const gemmi::Structure expanded = ExpandIntoSubgroup(placed, 1, g, *setting, h_cell, h, &n);
        CHECK(n == 2);
        CHECK(ClashFraction(expanded, 2) < 0.05);
    }
}

TEST_CASE("MolecularReplacement_MatthewsProbabilities", "[MolecularReplacement]") {
    // Vm 2.5 A^3/Da for one copy: one copy is by far the most probable, four impossible (Vm 0.6).
    const auto p = MatthewsProbabilities(2.5 * 20000, 20000);
    REQUIRE(p.size() == 1);
    CHECK(p[0] == Catch::Approx(1.0));
    // Vm 9 for one copy: three or four copies (Vm 3, 2.25) are the likely ones, one copy is not.
    const auto q = MatthewsProbabilities(9.0 * 20000, 20000);
    REQUIRE(q.size() == 6);
    CHECK(q[0] < 0.01);
    CHECK(std::max_element(q.begin(), q.end()) - q.begin() == 3);
}

TEST_CASE("MolecularReplacement_TranslationalNcs", "[MolecularReplacement]") {
    const SearchModel model = synthetic_model();
    const gemmi::UnitCell cell(52, 60, 90, 90, 90, 90);
    const gemmi::SpaceGroup &sg = gemmi::get_spacegroup_by_name("P 21 21 21");
    const gemmi::Mat33 r = axis_angle({1, 2, 3}, 50);
    const gemmi::Fractional t(0.17, 0.31, 0.09), v(0.0, 0.0, 0.5);
    const auto one = observed(PlaceSearchModel(model, r, t, cell, sg));
    CHECK_FALSE(DetectTranslationalNcs(one, cell, sg).present);
    const auto two = observed(AssembleCopies({PlaceSearchModel(model, r, t, cell, sg),
                                              PlaceSearchModel(model, r, gemmi::Fractional(t.x, t.y, t.z + v.z), cell, sg)}));
    const TranslationalNcs tncs = DetectTranslationalNcs(two, cell, sg);
    REQUIRE(tncs.present);
    CHECK(std::fabs(tncs.vector.x - std::round(tncs.vector.x)) < 0.03);
    CHECK(std::fabs(tncs.vector.y - std::round(tncs.vector.y)) < 0.03);
    CHECK(std::fabs(tncs.vector.z - 0.5) < 0.03);
}

TEST_CASE("MolecularReplacement_TncsRelatesGroups", "[MolecularReplacement]") {
    const auto &p3121 = gemmi::get_spacegroup_by_name("P 31 2 1"), &p3221 = gemmi::get_spacegroup_by_name("P 32 2 1");
    // A translation of c/3 turns a 3(1) axis into a 3(2) one, wherever it lies in the plane.
    CHECK(TncsRelatesGroups(p3121, p3221, gemmi::Fractional(0.3, 0.1, 1.0 / 3)));
    CHECK_FALSE(TncsRelatesGroups(p3121, p3221, gemmi::Fractional(0.3, 0.1, 0.1)));
    // A translation of a/2 turns 2(1) along a into 2, and nothing else.
    const auto &p212121 = gemmi::get_spacegroup_by_name("P 21 21 21");
    CHECK(TncsRelatesGroups(gemmi::get_spacegroup_by_name("P 2 21 21"), p212121, gemmi::Fractional(0.5, 0, 0)));
    CHECK_FALSE(TncsRelatesGroups(gemmi::get_spacegroup_by_name("P 21 2 2"), p212121, gemmi::Fractional(0.5, 0, 0)));
}

// A copy refined beside another one held fixed goes back where it belongs, on the CPU and on the GPU alike.
TEST_CASE("MolecularReplacement_RigidBodyBesideFixedCopy", "[MolecularReplacement]") {
    const SearchModel model = synthetic_model();
    const gemmi::UnitCell cell(52, 60, 70, 90, 90, 90);
    const gemmi::SpaceGroup &sg = gemmi::get_spacegroup_by_name("P 21 21 21");
    const gemmi::Mat33 r1 = axis_angle({1, 2, 3}, 50), r2 = axis_angle({-2, 1, 1}, 110);
    const gemmi::Fractional t1(0.17, 0.31, 0.09), t2(0.41, 0.12, 0.33);
    const gemmi::Structure first = PlaceSearchModel(model, r1, t1, cell, sg);
    const gemmi::Structure truth = PlaceSearchModel(model, r2, t2, cell, sg);
    const auto merged = observed(AssembleCopies({first, truth}));
    gemmi::AsuData<gemmi::ValueSigma<float>> fobs;
    fobs.unit_cell_ = cell;
    fobs.spacegroup_ = &sg;
    for (const MergedReflection &m : merged)
        fobs.v.push_back({{{m.h, m.k, m.l}}, {m.F, 1.0f}});
    fobs.ensure_asu();
    fobs.ensure_sorted();
    // 3 deg and 0.5 A out.
    const gemmi::Structure start = PlaceSearchModel(model, axis_angle({0, 0, 1}, 3).multiply(r2),
                                                    gemmi::Fractional(t2.x + 0.5 / 52, t2.y, t2.z), cell, sg);
    const auto worst = [&](const gemmi::Structure &st) {
        const auto a = ModelPositions(st.models[0]), b = ModelPositions(truth.models[0]);
        double w = 0;
        for (size_t i = 0; i < a.size(); i++)
            w = std::max(w, a[i].dist(b[i]));
        return w;
    };
    Logger logger("MolecularReplacement_RigidBodyBesideFixedCopy");
    gemmi::Structure cpu = start;
    RefineRigidBody(cpu.models[0], cell, sg, fobs, 3.0, logger, 4, nullptr, &first.models[0]);
    CHECK(worst(start) > 1.0);
    CHECK(worst(cpu) < 0.1);
#ifdef JFJOCH_USE_CUDA
    if (get_gpu_count() > 0) {
        gemmi::Structure gpu = start;
        auto pool = RigidBodyGPUPool::Plan(gpu.models[0], cell, sg, 3.0, fobs.v.size(), logger);
        REQUIRE(pool);
        pool->Reserve({1});
        RefineRigidBody(gpu.models[0], cell, sg, fobs, 3.0, logger, 4, pool.get(), &first.models[0]);
        CHECK(worst(gpu) < 0.1);
    }
#endif
}

#ifdef JFJOCH_USE_CUDA
TEST_CASE("MolecularReplacement_GPUMatchesCPU", "[MolecularReplacement]") {
    if (get_gpu_count() == 0)
        SKIP("no GPU");
    const SearchModel model = synthetic_model();
    const gemmi::UnitCell cell(42, 50, 58, 90, 90, 90);
    const gemmi::SpaceGroup &sg = gemmi::get_spacegroup_by_name("P 21 21 21");
    const gemmi::Mat33 truth = axis_angle({1, 2, 3}, 50);
    const gemmi::Fractional t(0.17, 0.31, 0.09);
    const auto merged = observed(PlaceSearchModel(model, truth, t, cell, sg));
    const SearchReflections refl = MakeSearchReflections(merged, cell, sg, 4.0, 12.0);
    const MolecularTransform ft(model, 4.0, 4);
    MolecularReplacementGPU gpu(ft);

    const auto rotations = RotationAsu(UniformRotations(15.0), CartesianPointGroup(sg, cell));
    const auto rf_cpu = RotationFunction(ft, refl, cell, sg, rotations, 4);
    const auto rf_gpu = gpu.RotationFunction(refl, cell, sg, rotations);
    REQUIRE(rf_gpu.size() == rf_cpu.size());
    for (size_t i = 0; i < rf_cpu.size(); i++)
        CHECK(std::fabs(rf_gpu[i] - rf_cpu[i]) < 1e-4);

    const TranslationSetup setup = MakeTranslationSetup(refl, cell, sg);
    std::vector<gemmi::Mat33> orientations{truth, rotations[3], rotations[17]};
    const auto tf_gpu = gpu.TranslationFunction(setup, orientations);
    for (size_t i = 0; i < orientations.size(); i++) {
        const TranslationPeak cpu = TranslationFunction(ft, setup, orientations[i]);
        CHECK(std::fabs(tf_gpu[i].z - cpu.z) < 1e-3);
        // The solution's peak is the same up to the origins P 21 21 21 allows, which tie exactly; noise peaks
        // can tie between grid points.
        if (i == 0)
            for (double d : {tf_gpu[i].t.x - cpu.t.x, tf_gpu[i].t.y - cpu.t.y, tf_gpu[i].t.z - cpu.t.z})
                CHECK(std::fabs(2 * d - std::round(2 * d)) < 1e-6);
    }
    // The packed transform gives the same bits as the plain one.
    MolecularReplacementGPU packed(ft, true);
    CHECK(packed.RotationFunction(refl, cell, sg, rotations) == rf_gpu);
    const auto tf_packed = packed.TranslationFunction(setup, orientations);
    for (size_t i = 0; i < orientations.size(); i++) {
        CHECK(tf_packed[i].z == tf_gpu[i].z);
        CHECK(tf_packed[i].t.x == tf_gpu[i].t.x);
        CHECK(tf_packed[i].t.y == tf_gpu[i].t.y);
        CHECK(tf_packed[i].t.z == tf_gpu[i].t.z);
    }
    CHECK(MolecularReplacementGPU::TransformBytes(ft, true) > MolecularReplacementGPU::TransformBytes(ft, false));
    // The molecular transform's FFT.
    gemmi::DensityCalculator<gemmi::IT92<float>, float> dc;
    dc.d_min = 4.0;
    dc.rate = 1.5;
    gemmi::Structure st = PlaceSearchModel(model, truth, t, cell, gemmi::get_spacegroup_p1());
    dc.set_grid_cell_and_spacegroup(st);
    dc.put_model_density_on_grid(st.models[0]);
    const int reach = dc.grid.nv / 3;   // inside the half grid, as the transform's own reach is
    const auto f_cpu = CropHalfL(MapToFPhi(dc.grid), reach), f_gpu = MapToFPhiGPU(dc.grid, reach);
    REQUIRE(f_gpu.data.size() == f_cpu.data.size());
    double worst = 0, largest = 0;
    for (size_t i = 0; i < f_cpu.data.size(); i++) {
        worst = std::max(worst, static_cast<double>(std::abs(f_gpu.data[i] - f_cpu.data[i])));
        largest = std::max(largest, static_cast<double>(std::abs(f_cpu.data[i])));
    }
    CHECK(worst < 1e-4 * largest);
    // Beside a known copy.
    const TranslationSetup setup2 = MakeTranslationSetup(refl, cell, sg, true);
    const auto known = KnownStructureFactors(ft, setup2, cell, {rotations[5]}, {gemmi::Fractional(0.3, 0.1, 0.4)});
    const auto tf2_gpu = gpu.TranslationFunction(setup2, orientations, &known);
    for (size_t i = 0; i < orientations.size(); i++) {
        const TranslationPeak cpu = TranslationFunction(ft, setup2, orientations[i], &known);
        CHECK(std::fabs(tf2_gpu[i].z - cpu.z) < 1e-3);
    }
}
#endif

#ifdef JFJOCH_USE_CUDA
TEST_CASE("MolecularReplacement_SplitOverCardsMatchesOneCard", "[MolecularReplacement]") {
    if (get_gpu_count() == 0)
        SKIP("no GPU");
    const SearchModel model = synthetic_model();
    const gemmi::UnitCell cell(42, 50, 58, 90, 90, 90);
    const gemmi::SpaceGroup &sg = gemmi::get_spacegroup_by_name("P 21 21 21");
    const gemmi::Mat33 truth = axis_angle({1, 2, 3}, 50);
    const auto merged = observed(PlaceSearchModel(model, truth, gemmi::Fractional(0.17, 0.31, 0.09), cell, sg));
    const SearchReflections refl = MakeSearchReflections(merged, cell, sg, 4.0, 12.0);
    const MolecularTransform ft(model, 4.0, 4);
    // Four "cards", each a search of its own on this one: the split's slices, gather and order are those of four
    // cards, so this is the multi-card path, run where there is one card.
    MolecularReplacementGPU one(ft);
    std::vector<std::unique_ptr<MolecularReplacementGPU>> copies;
    std::vector<MolecularReplacementGPUOnCard> cards;
    for (int k = 0; k < 4; k++) {
        copies.push_back(std::make_unique<MolecularReplacementGPU>(ft));
        cards.push_back({0, copies.back().get()});   // the test thread's device
    }
    std::atomic<size_t> holds{0};
    const MolecularReplacementGPUHold hold = [&](int, size_t bytes) -> std::shared_ptr<void> {
        CHECK(bytes > 0);
        holds++;
        return nullptr;
    };
    const auto all = RotationAsu(UniformRotations(15.0), CartesianPointGroup(sg, cell));
    REQUIRE(all.size() % 4 != 0);   // the last slice shorter than the others
    for (size_t n : {all.size(), size_t(3), size_t(1)}) {
        const std::vector<gemmi::Mat33> rotations(all.begin(), all.begin() + n);
        const std::vector<double> want = one.RotationFunction(refl, cell, sg, rotations);
        for (size_t parts : {1, 2, 3, 4}) {
            const std::vector<MolecularReplacementGPUOnCard> some(cards.begin(), cards.begin() + parts);
            const std::vector<double> got = SplitRotationFunction(some, refl, cell, sg, rotations, hold);
            CAPTURE(n, parts);
            REQUIRE(got.size() == want.size());
            CHECK(std::memcmp(got.data(), want.data(), want.size() * sizeof(double)) == 0);
        }
    }
    CHECK(holds > 0);

    const TranslationSetup setup = MakeTranslationSetup(refl, cell, sg);
    const std::vector<gemmi::Mat33> orientations{truth, all[3], all[17], all[20], all[31]};
    const TranslationSetup setup2 = MakeTranslationSetup(refl, cell, sg, true);
    const auto known = KnownStructureFactors(ft, setup2, cell, {all[5]}, {gemmi::Fractional(0.3, 0.1, 0.4)});
    for (const std::vector<std::complex<double>> *k : {(const std::vector<std::complex<double>> *) nullptr, &known}) {
        const TranslationSetup &s = k == nullptr ? setup : setup2;
        const std::vector<TranslationPeak> want = one.TranslationFunction(s, orientations, k);
        for (size_t parts : {2, 4}) {
            const std::vector<MolecularReplacementGPUOnCard> some(cards.begin(), cards.begin() + parts);
            const std::vector<TranslationPeak> got = SplitTranslationFunction(some, s, orientations, k, nullptr, hold);
            CAPTURE(parts, k != nullptr);
            REQUIRE(got.size() == want.size());
            for (size_t i = 0; i < want.size(); i++) {
                CHECK(std::memcmp(&got[i].z, &want[i].z, sizeof(double)) == 0);
                CHECK(got[i].t.x == want[i].t.x);
                CHECK(got[i].t.y == want[i].t.y);
                CHECK(got[i].t.z == want[i].t.z);
            }
        }
    }
}
#endif