Jungfraujoch/tests/ModelValidationTest.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 <cmath>
#include <cstdio>
#include <filesystem>
#include <fstream>
#include <random>

#include <gemmi/mmread_gz.hpp>

#include "../common/Logger.h"
#include "../rugnux/ModelValidation.h"
#include "../rugnux/RigidBodyRefine.h"
#include "../rugnux/SigmaA.h"
#include "../rugnux/WriteModel.h"

namespace {
    // A synthetic P1 cell with two carbon atoms - enough for a reader to produce a Structure with
    // atoms, a cell and a space group, which is all these tests are about. Neutral by construction:
    // no real specimen's parameters are involved.
    const char *kMmcif = R"(data_test
_cell.length_a      40.000
_cell.length_b      50.000
_cell.length_c      60.000
_cell.angle_alpha   90.000
_cell.angle_beta    90.000
_cell.angle_gamma   90.000
_symmetry.space_group_name_H-M   'P 1'
loop_
_atom_site.group_PDB
_atom_site.id
_atom_site.type_symbol
_atom_site.label_alt_id
_atom_site.label_atom_id
_atom_site.label_comp_id
_atom_site.label_asym_id
_atom_site.label_seq_id
_atom_site.Cartn_x
_atom_site.Cartn_y
_atom_site.Cartn_z
_atom_site.occupancy
_atom_site.B_iso_or_equiv
ATOM 1 C . CA GLY A 1 10.000 12.000 14.000 1.00 20.00
ATOM 2 C . CB GLY A 1 12.000 14.000 16.000 1.00 20.00
)";

    const char *kPdb =
        "CRYST1   40.000   50.000   60.000  90.00  90.00  90.00 P 1           1\n"
        "ATOM      1  CA  GLY A   1      10.000  12.000  14.000  1.00 20.00           C\n"
        "ATOM      2  CB  GLY A   1      12.000  14.000  16.000  1.00 20.00           C\n"
        "END\n";

    // A synthetic "protein": carbons filling one asymmetric unit of a small cell. The space group
    // matters - in P1 the origin is free in all three directions, so |F| does not change when the
    // whole content is translated and there is no translation to recover. No specimen is involved;
    // the positions come out of a fixed seed.
    const char *kCryst      = "CRYST1   30.000   34.000   38.000  90.00  90.00  90.00 P 21 21 21    4\n";
    // The same cell in a polar group: b is the unique axis, so the origin is free along y alone.
    const char *kPolarCryst = "CRYST1   30.000   34.000   38.000  90.00  90.00  90.00 P 1 2 1       2\n";

    std::string ClusterPdb(const char *cryst = kCryst) {
        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++) {
            std::snprintf(line, sizeof line,
                          "ATOM  %5d  C   UNK A   1    %8.3f%8.3f%8.3f  1.00 20.00           C\n",
                          i, x(rng), y(rng), z(rng));
            pdb += line;
        }
        return pdb + "END\n";
    }

    // The same two atoms as kPdb, plus an anisotropic U, a partial occupancy and a water: a written
    // model has to carry what the input carried, not only its coordinates.
    const char *kPdbRich =
        "CRYST1   40.000   50.000   60.000  90.00  90.00  90.00 P 1           1\n"
        "ATOM      1  CA  GLY A   1      10.000  12.000  14.000  1.00 20.00           C\n"
        "ANISOU    1  CA  GLY A   1     1000   1200   1400    100    200    300       C\n"
        "ATOM      2  CB  GLY A   1      12.000  14.000  16.000  0.60 25.00           C\n"
        "HETATM    3  O   HOH A 101      20.000  22.000  24.000  1.00 30.00           O\n"
        "END\n";

    std::string WriteTemp(const std::string &name, const char *content) {
        std::ofstream f(name);
        f << content;
        f.close();
        return name;
    }
}

// --model used to call gemmi::read_pdb unconditionally, so a deposited model handed over as mmCIF -
// which is how the PDB serves coordinates by default - was refused outright. Both formats now go
// through the same reader, which decides on the file's content.
TEST_CASE("ModelValidation_ReadsPdbAndMmcif", "[ModelValidation]") {
    Logger logger("ModelValidation_ReadsPdbAndMmcif");

    const auto pdb = WriteTemp("model_validation_test.pdb", kPdb);
    const auto cif = WriteTemp("model_validation_test.cif", kMmcif);

    const auto from_pdb = ModelReferenceIntensities(pdb, {}, {}, 4.0, logger);
    const auto from_cif = ModelReferenceIntensities(cif, {}, {}, 4.0, logger);

    REQUIRE_FALSE(from_pdb.empty());
    REQUIRE_FALSE(from_cif.empty());
    // The same structure either way, so the same reflections come out of it.
    CHECK(from_cif.size() == from_pdb.size());

    // The extension is not what decides: the same mmCIF under a .pdb name still reads.
    const auto misnamed = WriteTemp("model_validation_test_misnamed.pdb", kMmcif);
    CHECK_FALSE(ModelReferenceIntensities(misnamed, {}, {}, 4.0, logger).empty());

    std::filesystem::remove(pdb);
    std::filesystem::remove(cif);
    std::filesystem::remove(misnamed);
}

// A model that cannot be used must say why. Returning an empty result and logging was enough to lose
// the fact entirely: the run finished successfully with no R-free and no maps, which is exactly what
// a run without --model looks like.
TEST_CASE("ModelValidation_UnusableModelGivesAReason", "[ModelValidation]") {
    Logger logger("ModelValidation_UnusableModelGivesAReason");

    const auto empty_file = WriteTemp("model_validation_test_bogus.pdb", "not a coordinate file\n");

    const auto result = ValidateAgainstModel({}, UnitCell{.a = 40, .b = 50, .c = 60,
                                                          .alpha = 90, .beta = 90, .gamma = 90},
                                             empty_file, "", logger);
    CHECK_FALSE(result.ok);
    CHECK_FALSE(result.failure_reason.empty());
    CHECK_THAT(result.failure_reason, Catch::Matchers::ContainsSubstring(empty_file));

    const auto missing = ValidateAgainstModel({}, UnitCell{.a = 40, .b = 50, .c = 60,
                                                           .alpha = 90, .beta = 90, .gamma = 90},
                                              "model_validation_test_does_not_exist.pdb", "", logger);
    CHECK_FALSE(missing.ok);
    CHECK_FALSE(missing.failure_reason.empty());

    std::filesystem::remove(empty_file);
}

// A rigid-body step is only worth taking if it can find a shift it was not told about. The check is
// closed: the "observed" amplitudes are the model's own, so the answer is known to be zero shift.
TEST_CASE("ModelValidation_RigidBodyRecoversASmallShift", "[ModelValidation]") {
    Logger logger("ModelValidation_RigidBodyRecoversASmallShift");

    const auto path = WriteTemp("rigid_body_test.pdb", ClusterPdb().c_str());
    gemmi::Structure st = gemmi::read_structure_gz(path, gemmi::CoorFormat::Detect);
    const gemmi::SpaceGroup *sg = st.find_spacegroup();
    REQUIRE(sg != nullptr);
    st.setup_cell_images();

    // "Observed" amplitudes: the model's own structure factors, so the target's minimum is exactly
    // where the model started.
    const auto ref = ModelReferenceIntensities(path, {}, {}, 3.0, logger);
    REQUIRE_FALSE(ref.empty());
    gemmi::AsuData<gemmi::ValueSigma<float>> fobs;
    fobs.unit_cell_ = st.cell;
    fobs.spacegroup_ = sg;
    for (const auto &r : ref)
        fobs.v.push_back({{{r.h, r.k, r.l}}, {std::sqrt(r.I), 1.0f}});
    fobs.ensure_sorted();

    const std::vector<gemmi::Position> original = ModelPositions(st.models[0]);
    std::vector<gemmi::Position> displaced;
    for (const gemmi::Position &p : original)
        displaced.emplace_back(p.x + 0.40, p.y - 0.30, p.z + 0.20);   // 0.54 A off
    SetModelPositions(st.models[0], displaced);

    const RigidBodyRefineResult result =
        RefineRigidBody(st.models[0], st.cell, *sg, fobs, 3.0, logger);
    CHECK(result.converged);

    const std::vector<gemmi::Position> refined = ModelPositions(st.models[0]);
    double before = 0, after = 0;
    for (size_t i = 0; i < original.size(); i++) {
        before += original[i].dist_sq(displaced[i]);
        after += original[i].dist_sq(refined[i]);
    }
    before = std::sqrt(before / original.size());
    after = std::sqrt(after / original.size());
    logger.Info("Rigid-body test: rmsd from the truth {:.3f} A -> {:.3f} A", before, after);
    CHECK(after < 0.2 * before);

    std::filesystem::remove(path);
}

// A polar space group leaves the origin free along one direction: moving the whole cell content
// along it multiplies every structure factor by a phase and changes no amplitude, so the data
// cannot say where the body sits along it and the refinement must not pretend otherwise. The check
// is closed - the "observed" amplitudes are the model's own - and the model is displaced in all
// three directions at once, so the same run says both what is recovered and what is left alone.
TEST_CASE("ModelValidation_RigidBodyLeavesThePolarDirectionAlone", "[ModelValidation]") {
    Logger logger("ModelValidation_RigidBodyLeavesThePolarDirectionAlone");

    const auto path = WriteTemp("rigid_body_polar_test.pdb", ClusterPdb(kPolarCryst).c_str());
    gemmi::Structure st = gemmi::read_structure_gz(path, gemmi::CoorFormat::Detect);
    const gemmi::SpaceGroup *sg = st.find_spacegroup();
    REQUIRE(sg != nullptr);
    st.setup_cell_images();

    const auto ref = ModelReferenceIntensities(path, {}, {}, 3.0, logger);
    REQUIRE_FALSE(ref.empty());
    gemmi::AsuData<gemmi::ValueSigma<float>> fobs;
    fobs.unit_cell_ = st.cell;
    fobs.spacegroup_ = sg;
    for (const auto &r : ref)
        fobs.v.push_back({{{r.h, r.k, r.l}}, {std::sqrt(r.I), 1.0f}});
    fobs.ensure_sorted();

    // b is the unique axis of P 1 2 1, so y is the free direction and x and z are determined.
    const std::vector<gemmi::Position> original = ModelPositions(st.models[0]);
    std::vector<gemmi::Position> displaced;
    for (const gemmi::Position &p : original)
        displaced.emplace_back(p.x + 0.35, p.y + 0.50, p.z - 0.30);
    SetModelPositions(st.models[0], displaced);

    const RigidBodyRefineResult result =
        RefineRigidBody(st.models[0], st.cell, *sg, fobs, 3.0, logger);
    CHECK(result.converged);

    const std::vector<gemmi::Position> refined = ModelPositions(st.models[0]);
    gemmi::Vec3 left;
    for (size_t i = 0; i < original.size(); i++)
        left += gemmi::Vec3(refined[i]) - gemmi::Vec3(original[i]);
    left *= 1.0 / static_cast<double>(original.size());
    logger.Info("Rigid-body polar test: left over ({:.3f}, {:.3f}, {:.3f}) A", left.x, left.y, left.z);
    CHECK(std::fabs(left.x) < 0.10);
    CHECK(std::fabs(left.z) < 0.10);
    // Along b nothing was refined away, because there is nothing there to refine.
    CHECK(left.y == Catch::Approx(0.50).margin(0.02));

    std::filesystem::remove(path);
}

// The bulk solvent the placement is scored through has to stay inside the range a flat solvent model
// means anything in. gemmi's own scaler is an unbounded Levenberg-Marquardt and reaches b_sol of
// hundreds or thousands of A^2, which does not corrupt a reported number here but distorts the target
// that decides where the model goes, at every one of the hundreds of evaluations.
TEST_CASE("ModelValidation_RigidBodySolventStaysPhysical", "[ModelValidation]") {
    Logger logger("ModelValidation_RigidBodySolventStaysPhysical");

    const auto path = WriteTemp("rigid_body_solvent_test.pdb", ClusterPdb().c_str());
    gemmi::Structure st = gemmi::read_structure_gz(path, gemmi::CoorFormat::Detect);
    const gemmi::SpaceGroup *sg = st.find_spacegroup();
    REQUIRE(sg != nullptr);
    st.setup_cell_images();

    const auto ref = ModelReferenceIntensities(path, {}, {}, 3.0, logger);
    REQUIRE_FALSE(ref.empty());
    gemmi::AsuData<gemmi::ValueSigma<float>> fobs;
    fobs.unit_cell_ = st.cell;
    fobs.spacegroup_ = sg;
    for (const auto &r : ref)
        fobs.v.push_back({{{r.h, r.k, r.l}}, {std::sqrt(r.I), 1.0f}});
    fobs.ensure_sorted();

    // Turned right around, so the model explains nothing: that is where an unbounded solvent fit
    // has nothing to hold it and runs away.
    std::vector<gemmi::Position> turned = ModelPositions(st.models[0]);
    gemmi::Vec3 centre;
    for (const gemmi::Position &p : turned)
        centre += p;
    centre *= 1.0 / static_cast<double>(turned.size());
    for (gemmi::Position &p : turned)
        p = gemmi::Position(centre - (gemmi::Vec3(p) - centre));
    SetModelPositions(st.models[0], turned);

    const RigidBodyRefineResult result =
        RefineRigidBody(st.models[0], st.cell, *sg, fobs, 3.0, logger);
    logger.Info("Rigid-body solvent test: k_sol {:.3f}, b_sol {:.1f} A^2", result.k_sol, result.b_sol);
    CHECK(result.k_sol >= 0.10);
    CHECK(result.k_sol <= 0.60);
    CHECK(result.b_sol >= 10.0);
    CHECK(result.b_sol <= 80.0);

    std::filesystem::remove(path);
}

// sigma_A is what says how much of the model to believe, so the two ends of its range are what the
// weighting has to get right: a model that explains the data completely, and one that explains none
// of it.
TEST_CASE("ModelValidation_SigmaAWeightsFollowTheModelsAgreement", "[ModelValidation]") {
    gemmi::UnitCell cell(40, 50, 60, 90, 90, 90);
    std::mt19937 rng(12345);
    std::normal_distribution<double> normal(0.0, 1.0);

    auto weights = [&](bool agreeing) {
        std::vector<SigmaAReflection> refl;
        for (int i = 0; i < 2000; i++) {
            SigmaAReflection r;
            r.f_calc = std::fabs(normal(rng)) * 100;
            r.f_obs = agreeing ? r.f_calc : std::fabs(normal(rng)) * 100;
            r.inv_d2 = 0.01 + 0.2 * (i / 2000.0);
            r.free = (i % 20) == 0;                 // the usual 5 %
            refl.push_back(r);
        }
        return EstimateSigmaA(refl, cell);
    };

    const SigmaAResult perfect = weights(true);
    const SigmaAResult useless = weights(false);
    CHECK(perfect.mean_fom > 0.85);
    CHECK(useless.mean_fom < 0.2);
    CHECK(perfect.shells == 2);   // 100 free reflections, 50 to a shell

    // No free reflections to estimate on: the coefficients are left alone rather than weighted by a
    // number that was never measured.
    std::vector<SigmaAReflection> no_free;
    for (int i = 0; i < 100; i++)
        no_free.push_back({100.0, 100.0, 0.05, 1, false, false});
    const SigmaAResult unweighted = EstimateSigmaA(no_free, cell);
    CHECK(unweighted.weight.size() == no_free.size());
    CHECK(unweighted.weight[0].m == 1.0);
    CHECK(unweighted.weight[0].d == 1.0);
}

// The model rugnux scored has to reach disk, or a user overlaying their input model on rugnux's maps
// is wrong by the whole rigid-body shift. Two things have to hold: the file carries what the input
// carried, and it carries the cell and space group the reflection files beside it are written in -
// which, once --model has adopted the model's enantiomorph, is neither the data's original label nor
// necessarily the input model's.
TEST_CASE("WriteModel_KeepsTheContentAndTakesTheGivenFrame", "[ModelValidation]") {
    Logger logger("WriteModel_KeepsTheContentAndTakesTheGivenFrame");

    const auto input = WriteTemp("write_model_test_input.pdb", kPdbRich);
    gemmi::Structure st = gemmi::read_structure_gz(input, gemmi::CoorFormat::Detect);

    // A frame that is neither the model's (P 1, 40/50/60) nor anything derived from it: the tetragonal
    // lysozyme cell and one enantiomorph of its group, standing in for what AdoptModelFrame settled.
    const UnitCell data_cell{.a = 79, .b = 79, .c = 38, .alpha = 90, .beta = 90, .gamma = 90};
    const gemmi::SpaceGroup *sg = gemmi::find_spacegroup_by_name("P 43 21 2");
    REQUIRE(sg != nullptr);

    WritePlacedModel(st, data_cell, *sg, "write_model_test", logger);
    const std::string written = "write_model_test_model.cif";
    REQUIRE(std::filesystem::exists(written));

    const gemmi::Structure back = gemmi::read_structure_gz(written, gemmi::CoorFormat::Detect);
    CHECK(back.cell.a == Catch::Approx(79.0));
    CHECK(back.cell.c == Catch::Approx(38.0));
    REQUIRE(back.find_spacegroup() != nullptr);
    CHECK(back.find_spacegroup()->number == 96);

    REQUIRE(back.models.size() == 1);
    REQUIRE(back.models[0].chains.size() == 1);
    const gemmi::Chain &chain = back.models[0].chains[0];
    REQUIRE(chain.residues.size() == 2);           // GLY and the water
    REQUIRE(chain.residues[0].atoms.size() == 2);
    REQUIRE(chain.residues[1].atoms.size() == 1);

    // The coordinates are written as they stand - WritePlacedModel changes the label, not the atoms.
    const gemmi::Atom &ca = chain.residues[0].atoms[0];
    const gemmi::Atom &cb = chain.residues[0].atoms[1];
    CHECK(ca.pos.x == Catch::Approx(10.0));
    CHECK(ca.pos.z == Catch::Approx(14.0));
    CHECK(cb.pos.y == Catch::Approx(14.0));
    // B-factors, occupancies and the anisotropic U survive the round trip.
    CHECK(ca.b_iso == Catch::Approx(20.0));
    CHECK(cb.b_iso == Catch::Approx(25.0));
    CHECK(cb.occ == Catch::Approx(0.60));
    CHECK(ca.aniso.nonzero());
    CHECK(ca.aniso.u11 == Catch::Approx(0.1000).margin(1e-4));
    CHECK(ca.aniso.u12 == Catch::Approx(0.0100).margin(1e-4));
    CHECK(chain.residues[1].name == "HOH");

    std::filesystem::remove(input);
    std::filesystem::remove(written);
}

// CC(model, data) has to follow where the signal actually is, or it cannot support the one-sided
// claim it exists for. The check is closed: the "observed" intensities are the model's own with
// Gaussian noise added, and how much noise is chosen per shell - almost none in the first, some in
// the second, enough to bury the signal in the third - so the answer is known before the run.
TEST_CASE("ModelValidation_CCModelFollowsTheSignalByShell", "[ModelValidation]") {
    Logger logger("ModelValidation_CCModelFollowsTheSignalByShell");

    const auto path = WriteTemp("cc_model_test.pdb", ClusterPdb().c_str());
    auto obs = ModelReferenceIntensities(path, {}, {}, 2.5, logger);
    REQUIRE(obs.size() > 1000);

    // The shells the correlation is reported on, coarse to fine, and the noise each one gets as a
    // multiple of the r.m.s. intensity of that shell. Nothing coarser than the first shell is kept:
    // the reference intensities carry the bulk solvent at fixed constants while the validation fits
    // its own, and below about 6 A that difference is a large part of |F| and would decorrelate a
    // shell this test needs to be clean.
    const std::vector<float> shells{5.0f, 3.2f, 2.5f};
    const double noise[3] = {0.02, 1.0, 30.0};

    auto shell_of = [&](float d) {
        for (size_t s = 0; s < shells.size(); s++)
            if (d > shells[s]) return s;
        return shells.size();
    };
    std::erase_if(obs, [&](const MergedReflection &r) { return r.d > 6.0f || shell_of(r.d) >= shells.size(); });
    REQUIRE(obs.size() > 500);

    std::vector<double> sum_i2(shells.size(), 0.0);
    std::vector<int> count(shells.size(), 0);
    for (const auto &r : obs) {
        sum_i2[shell_of(r.d)] += static_cast<double>(r.I) * r.I;
        ++count[shell_of(r.d)];
    }

    std::mt19937 rng(20260907);
    std::normal_distribution<double> normal(0.0, 1.0);
    for (size_t i = 0; i < obs.size(); i++) {
        const size_t bin = shell_of(obs[i].d);
        const double sd = noise[bin] * std::sqrt(sum_i2[bin] / count[bin]);
        obs[i].I = static_cast<float>(obs[i].I + sd * normal(rng));
        obs[i].sigma = static_cast<float>(std::max(1.0, sd));
        obs[i].F = std::sqrt(std::max(0.0f, obs[i].I));
        obs[i].rfree_flag = (i % 20) == 0;
    }

    const std::string prefix = (std::filesystem::temp_directory_path() / "cc_model_test").string();
    const auto result =
        ValidateAgainstModel(obs, UnitCell{.a = 30, .b = 34, .c = 38,
                                           .alpha = 90, .beta = 90, .gamma = 90},
                             path, prefix, logger, gemmi::find_spacegroup_by_name("P 21 21 21"),
                             /*probe_indexing_ambiguity=*/false, 1, 1.0, shells);
    REQUIRE(result.ok);
    REQUIRE(result.cc_model_shells.size() == shells.size());
    int n_total = 0;
    for (size_t s = 0; s < shells.size(); s++) {
        const auto &sh = result.cc_model_shells[s];
        logger.Info("CC(model,data) {:.2f} A: {:.3f} on {} refl, {:+.1f} sigma",
                    sh.d_min, sh.cc, sh.n, sh.sigma);
        CHECK(sh.d_min == shells[s]);
        CHECK(sh.n > 20);
        n_total += sh.n;
    }
    CHECK(n_total == result.cc_model_n);
    // Essentially noiseless: the model is the data, so the correlation is high and hugely significant.
    CHECK(result.cc_model_shells[0].cc > 0.9);
    CHECK(result.cc_model_shells[0].sigma > 10.0);
    // Noise at the shell's own r.m.s. still leaves plenty to see.
    CHECK(result.cc_model_shells[1].cc > 0.25);
    CHECK(result.cc_model_shells[1].sigma > 5.0);
    // Buried: the shell must NOT come out significant, or the one-sided test would fire on noise.
    CHECK(std::fabs(result.cc_model_shells[2].cc) < 0.15);
    CHECK(std::fabs(result.cc_model_shells[2].sigma) < 4.0);

    // No shells asked for, none reported: a run that did not measure it writes no key.
    const auto no_shells =
        ValidateAgainstModel(obs, UnitCell{.a = 30, .b = 34, .c = 38,
                                           .alpha = 90, .beta = 90, .gamma = 90},
                             path, prefix, logger, gemmi::find_spacegroup_by_name("P 21 21 21"),
                             false, 1, 1.0);
    CHECK(no_shells.ok);
    CHECK(no_shells.cc_model_shells.empty());

    for (const char *suffix : {"_2fofc.ccp4", "_fofc.ccp4", "_anom.ccp4", "_maps.mtz"})
        std::filesystem::remove(prefix + suffix);
    std::filesystem::remove(path);
}