Jungfraujoch/tests/ModelScalingTest.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 <algorithm>
#include <complex>
#include <random>
#include <utility>

#include "../image_analysis/structure_refinement/ModelScaling.h"
#include "gemmi/asumask.hpp"
#include "gemmi/model.hpp"
#include "gemmi/scaling.hpp"
#include "gemmi/symmetry.hpp"
#include "gemmi/unitcell.hpp"

namespace {

// A synthetic set of scaling points: |Fobs| is generated from a known scale, isotropic B, k_sol
// and b_sol, so the fit has a right answer to find. The "structure factors" are a hash of the
// index, so the set is reproduced bit for bit.
void MakePoints(gemmi::Scaling<float> &scaling, const gemmi::UnitCell &cell,
                double k_overall, double b_iso, double k_sol, double b_sol,
                double d_min, double d_max) {
    scaling.points.clear();
    const int hmax = static_cast<int>(cell.a / d_min) + 1;
    const int kmax = static_cast<int>(cell.b / d_min) + 1;
    const int lmax = static_cast<int>(cell.c / d_min) + 1;
    uint32_t seed = 12345;
    auto rnd = [&seed]() {
        seed = seed * 1664525u + 1013904223u;
        return static_cast<double>((seed >> 8) & 0xffff) / 65535.0;
    };
    for (int h = 0; h <= hmax; ++h)
        for (int k = 0; k <= kmax; ++k)
            for (int l = 0; l <= lmax; ++l) {
                if (h == 0 && k == 0 && l == 0) continue;
                const gemmi::Miller hkl{{h, k, l}};
                const double d = cell.calculate_d(hkl);
                if (d < d_min || d > d_max) continue;
                const double stol2 = cell.calculate_stol_sq(hkl);
                const std::complex<double> fc(20.0 + 80.0 * rnd(), 0.0);
                const std::complex<double> fm(5.0 + 5.0 * rnd(), 0.0);
                const std::complex<double> total = fc + k_sol * std::exp(-b_sol * stol2) * fm;
                const double fobs = k_overall * std::exp(-b_iso * stol2) * std::abs(total);
                gemmi::Scaling<float>::Point p{};
                p.hkl = hkl;
                p.stol2 = stol2;
                p.fcmol = std::complex<float>(static_cast<float>(fc.real()), 0.f);
                p.fmask = std::complex<float>(static_cast<float>(fm.real()), 0.f);
                p.fobs = static_cast<float>(fobs);
                p.sigma = 1.f;
                scaling.points.push_back(p);
            }
}

gemmi::UnitCell Cell() {
    gemmi::UnitCell c;
    c.set(60.0, 70.0, 80.0, 90.0, 90.0, 90.0);
    return c;
}

}  // namespace

// The fit has to find a solvent pair it was given, not merely land somewhere legal.
TEST_CASE("ModelScaling recovers a known bulk solvent") {
    const gemmi::UnitCell cell = Cell();
    const gemmi::SpaceGroup *sg = gemmi::find_spacegroup_by_name("P 21 21 21");
    REQUIRE(sg != nullptr);
    gemmi::Scaling<float> scaling(cell, sg);
    scaling.use_solvent = true;
    MakePoints(scaling, cell, /*k_overall=*/0.5, /*b_iso=*/20.0,
               /*k_sol=*/0.35, /*b_sol=*/45.0, /*d_min=*/2.0, /*d_max=*/50.0);
    REQUIRE(scaling.points.size() > 500);

    const ModelScaleReport report = FitModelScale(scaling);
    CHECK(report.n_grid > 0);
    // The grid steps are 0.025 in k_sol and 5 in b_sol, so this is one step of tolerance.
    CHECK(scaling.k_sol == Catch::Approx(0.35).margin(0.03));
    CHECK(scaling.b_sol == Catch::Approx(45.0).margin(6.0));
    CHECK(report.r_work_fit < 0.05);
}

// The regression this exists for: gemmi's own fit_parameters() is an unbounded Levenberg-Marquardt
// and, where the data cannot determine the solvent, walks k_sol/b_sol out of the range a flat
// solvent model means anything in (b_sol of 1707 A^2 was measured on real data). Here the solvent
// contribution is made unidentifiable - only high-resolution data, where exp(-b_sol * s^2) is
// indistinguishable from zero for any large b_sol - and the fit must still return a physical pair.
TEST_CASE("ModelScaling stays physical when the solvent is unidentifiable") {
    const gemmi::UnitCell cell = Cell();
    const gemmi::SpaceGroup *sg = gemmi::find_spacegroup_by_name("P 21 21 21");
    REQUIRE(sg != nullptr);
    gemmi::Scaling<float> scaling(cell, sg);
    scaling.use_solvent = true;
    // Nothing below 3 A: the bulk solvent is a low-resolution feature, so it has almost
    // no leverage here.
    MakePoints(scaling, cell, /*k_overall=*/0.5, /*b_iso=*/20.0,
               /*k_sol=*/0.35, /*b_sol=*/45.0, /*d_min=*/1.2, /*d_max=*/3.0);
    REQUIRE(scaling.points.size() > 500);

    FitModelScale(scaling);
    CHECK(scaling.k_sol >= 0.15);
    CHECK(scaling.k_sol <= 0.50);
    CHECK(scaling.b_sol >= 10.0);
    CHECK(scaling.b_sol <= 80.0);

    // ... where gemmi's own unbounded fit is free to leave that range.
    gemmi::Scaling<float> unbounded(cell, sg);
    unbounded.use_solvent = true;
    unbounded.points = scaling.points;
    unbounded.fit_isotropic_b_approximately();
    unbounded.fit_parameters();
    CHECK(std::isfinite(unbounded.b_sol));   // it converges; it is simply not bounded
}

// A model at very high atomic B, where the solvent edge is smeared as much as the atoms: SolventBMax
// raises the b_sol ceiling to twice the model's mean B, and the fit then finds a b_sol far above 80
// that the default box cannot represent - which leaves it on the box edge with a worse R.
TEST_CASE("ModelScaling finds a high b_sol for a high-B model") {
    gemmi::Model model(1);
    model.chains.emplace_back("A");
    model.chains[0].residues.emplace_back();
    for (double b : {280.0, 320.0}) {
        gemmi::Atom atom;
        atom.b_iso = static_cast<float>(b);
        model.chains[0].residues[0].atoms.push_back(atom);
    }
    CHECK(SolventBMax(model) == Catch::Approx(600.0));
    for (auto &atom : model.chains[0].residues[0].atoms)
        atom.b_iso = 20.0f;
    CHECK(SolventBMax(model) == 80.0);

    const gemmi::UnitCell cell = Cell();
    const gemmi::SpaceGroup *sg = gemmi::find_spacegroup_by_name("P 21 21 21");
    REQUIRE(sg != nullptr);
    gemmi::Scaling<float> scaling(cell, sg);
    scaling.use_solvent = true;
    MakePoints(scaling, cell, /*k_overall=*/0.5, /*b_iso=*/300.0,
               /*k_sol=*/0.35, /*b_sol=*/300.0, /*d_min=*/3.5, /*d_max=*/50.0);
    REQUIRE(scaling.points.size() > 500);

    gemmi::Scaling<float> capped = scaling;
    const ModelScaleReport capped_report = FitModelScale(capped);
    CHECK(capped.b_sol == 80.0);

    const ModelScaleReport report = FitModelScale(scaling, {.b_hi = 600.0});
    // The b_sol steps are 590/14 A^2 in the refinement pass, so this is one step of tolerance.
    CHECK(scaling.b_sol == Catch::Approx(300.0).margin(43.0));
    CHECK(scaling.k_sol == Catch::Approx(0.35).margin(0.03));
    CHECK(report.r_work_fit < capped_report.r_work_fit);
}

// A densely packed crystal with no disordered solvent channels - a small molecule, say - has an
// empty solvent mask, and then Fmask is zero for every reflection. The bulk-solvent contribution
// k_sol * exp(-b_sol s^2) * Fmask is then identically zero WHATEVER k_sol and b_sol come out as,
// so a solvent term fitted where there is no solvent cannot add anything to the model. That is why
// this needs no solvent-content threshold and no switch: the mask already carries the answer.
TEST_CASE("ModelScaling adds nothing when there is no solvent to model") {
    const gemmi::UnitCell cell = Cell();
    const gemmi::SpaceGroup *sg = gemmi::find_spacegroup_by_name("P 21 21 21");
    REQUIRE(sg != nullptr);
    gemmi::Scaling<float> scaling(cell, sg);
    scaling.use_solvent = true;
    MakePoints(scaling, cell, /*k_overall=*/0.5, /*b_iso=*/20.0,
               /*k_sol=*/0.35, /*b_sol=*/45.0, /*d_min=*/2.0, /*d_max=*/50.0);
    REQUIRE(scaling.points.size() > 500);
    for (auto &p : scaling.points)     // an empty mask: no solvent-accessible volume
        p.fmask = {0.f, 0.f};

    FitModelScale(scaling);
    // Whatever the search settled on, the model it produces is the solvent-free one.
    double worst = 0;
    for (const auto &p : scaling.points)
        worst = std::max(worst, static_cast<double>(
                            std::fabs(std::abs(scaling.get_fcalc(p)) - std::abs(p.fcmol))));
    CHECK(worst == Catch::Approx(0.0).margin(1e-6));
    // and the parameters are still reported inside the physical box, not at some arbitrary value
    CHECK(scaling.k_sol >= 0.10);
    CHECK(scaling.k_sol <= 0.60);
    CHECK(scaling.b_sol >= 10.0);
    CHECK(scaling.b_sol <= 80.0);
}

// The grid points are fitted in parallel; the answer must be the serial one, bit for bit.
TEST_CASE("ModelScaling is the same on any number of threads") {
    const gemmi::UnitCell cell = Cell();
    const gemmi::SpaceGroup *sg = gemmi::find_spacegroup_by_name("P 21 21 21");
    REQUIRE(sg != nullptr);
    gemmi::Scaling<float> serial(cell, sg);
    serial.use_solvent = true;
    MakePoints(serial, cell, /*k_overall=*/0.5, /*b_iso=*/20.0,
               /*k_sol=*/0.35, /*b_sol=*/45.0, /*d_min=*/2.0, /*d_max=*/50.0);
    gemmi::Scaling<float> parallel = serial;

    const ModelScaleReport r1 = FitModelScale(serial, {}, 1);
    const ModelScaleReport r8 = FitModelScale(parallel, {}, 8);
    CHECK(r1.n_grid == r8.n_grid);
    CHECK(r1.r_work_fit == r8.r_work_fit);
    CHECK(serial.k_sol == parallel.k_sol);
    CHECK(serial.b_sol == parallel.b_sol);
    CHECK(serial.k_overall == parallel.k_overall);
    CHECK(serial.b_star.u11 == parallel.b_star.u11);
    CHECK(serial.b_star.u22 == parallel.b_star.u22);
    CHECK(serial.b_star.u33 == parallel.b_star.u33);
}

// FitModelScale fits each solvent pair of its grid on a copy of gemmi::Scaling's target that takes the
// solvent term once per pair (ModelScaling.cpp). On a box of a single pair it has to give exactly what
// gemmi's own Scaling gives for that pair - to the bit, since the validation's numbers are made of it.
TEST_CASE("ModelScaling fits a grid point exactly as gemmi's Scaling does") {
    const gemmi::UnitCell cell(61.0, 72.0, 83.0, 90.0, 101.0, 90.0);
    const gemmi::SpaceGroup *sg = gemmi::find_spacegroup_by_name("P 1 21 1");
    REQUIRE(sg != nullptr);
    const gemmi::GroupOps gops = sg->operations();
    const gemmi::ReciprocalAsu asu(sg);
    gemmi::Scaling<float> scaling(cell, sg);
    scaling.use_solvent = true;
    std::mt19937 rng(20261008);
    std::uniform_real_distribution<double> uniform(0.0, 1.0);
    const gemmi::SMat33<double> b_star = gemmi::SMat33<double>{22, 31, 27, 0, 3, 0}.transformed_by(cell.frac.mat);
    const gemmi::Miller lim = cell.get_hkl_limits(2.5);
    for (int h = -lim[0]; h <= lim[0]; h++)
        for (int k = -lim[1]; k <= lim[1]; k++)
            for (int l = -lim[2]; l <= lim[2]; l++) {
                const gemmi::Miller m{{h, k, l}};
                if ((h == 0 && k == 0 && l == 0) || cell.calculate_d(m) < 2.5 || !asu.is_in(m)
                    || gops.is_systematically_absent(m))
                    continue;
                const double stol2 = cell.calculate_stol_sq(m);
                const std::complex<float> fc(std::polar(100.0 * std::exp(-8.0 * stol2) * (0.2 + uniform(rng)),
                                                        6.283 * uniform(rng)));
                const std::complex<float> fm(std::polar(900.0 * std::exp(-20.0 * stol2), 6.283 * uniform(rng)));
                const std::complex<double> total = std::complex<double>(fc)
                                                   + 0.4 * std::exp(-50.0 * stol2) * std::complex<double>(fm);
                const double fobs = 2.5 * std::exp(-0.25 * b_star.r_u_r(m)) * std::abs(total)
                                    * (0.9 + 0.2 * uniform(rng));
                gemmi::Scaling<float>::Point p{};
                p.hkl = m;
                p.stol2 = stol2;
                p.fcmol = fc;
                p.fmask = fm;
                p.fobs = static_cast<float>(fobs);
                p.sigma = static_cast<float>(0.05 * fobs + 0.5);
                scaling.points.push_back(p);
            }
    REQUIRE(scaling.points.size() > 1000);

    for (const auto &[k_sol, b_sol] : {std::pair{0.35, 40.0}, std::pair{0.20, 70.0}}) {
        gemmi::Scaling<float> ours = scaling;
        FitModelScale(ours, ModelScaleBox{k_sol, k_sol, b_sol, b_sol}, 4);

        gemmi::Scaling<float> gemmis = scaling;
        gemmis.fix_k_sol = true;
        gemmis.fix_b_sol = true;
        gemmis.k_sol = k_sol;
        gemmis.b_sol = b_sol;
        gemmis.fit_isotropic_b_approximately();
        gemmis.fit_parameters();

        CHECK(ours.k_sol == k_sol);
        CHECK(ours.b_sol == b_sol);
        CHECK(ours.k_overall == gemmis.k_overall);
        CHECK(ours.b_star.u11 == gemmis.b_star.u11);
        CHECK(ours.b_star.u22 == gemmis.b_star.u22);
        CHECK(ours.b_star.u33 == gemmis.b_star.u33);
        CHECK(ours.b_star.u12 == gemmis.b_star.u12);
        CHECK(ours.b_star.u13 == gemmis.b_star.u13);
        CHECK(ours.b_star.u23 == gemmis.b_star.u23);
        CHECK(ours.calculate_r_factor() == gemmis.calculate_r_factor());
    }
}