Jungfraujoch/tests/MergeScaleTest.cpp

// SPDX-FileCopyrightText: 2024 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only

#include <catch2/catch_all.hpp>

#include <random>
#include "../image_analysis/scale_merge/HKLKey.h"
#include "../image_analysis/scale_merge/Merge.h"
#include "../image_analysis/scale_merge/ResolutionCutoff.h"
#include "gemmi/reciproc.hpp"

TEST_CASE("HKLKey_NoSG_noMergeFriedel") {
    HKLKeyGenerator hkl_key_gen(false, *gemmi::find_spacegroup_by_number(1));
    CHECK(hkl_key_gen(-1, -2, -3) != hkl_key_gen(1,2,3));
    CHECK(hkl_key_gen(-1,-2,-3) == hkl_key_gen(-1,-2,-3));
    CHECK(hkl_key_gen(-1,-2,-3) != hkl_key_gen(1,-2,-3));
}

TEST_CASE("HKLKey_NoSG_MergeFriedel") {
    HKLKeyGenerator hkl_key_gen(true, *gemmi::find_spacegroup_by_number(1));
    CHECK(hkl_key_gen(-1, -2, -3) == hkl_key_gen(1,2,3));
    CHECK(hkl_key_gen(-1,-2,-3) == hkl_key_gen(-1,-2,-3));
    CHECK(hkl_key_gen(-1,-2,-3) != hkl_key_gen(1,-2,-3));
}

TEST_CASE("HKLKey_SG1_MergeFriedel") {
    HKLKeyGenerator hkl_key_gen(true, *gemmi::find_spacegroup_by_number(1));
    CHECK(hkl_key_gen(-1, -2, -3) == hkl_key_gen(1,2,3));
    CHECK(hkl_key_gen(-1,-2,-3) == hkl_key_gen(-1,-2,-3));
    CHECK(hkl_key_gen(-1,-2,-3) != hkl_key_gen(1,-2,-3));
}

TEST_CASE("HKLKey_SG1_NoMergeFriedel") {
    HKLKeyGenerator hkl_key_gen(false, *gemmi::find_spacegroup_by_number(1));
    CHECK(hkl_key_gen(-1, -2, -3) != hkl_key_gen(1,2,3));
    CHECK(hkl_key_gen(-1,-2,-3) == hkl_key_gen(-1,-2,-3));
    CHECK(hkl_key_gen(-1,-2,-3) != hkl_key_gen(1,-2,-3));
}

TEST_CASE("HKLKey_SG96_MergeFriedel") {
    HKLKeyGenerator hkl_key_gen(true, *gemmi::find_spacegroup_by_number(96));
    CHECK(hkl_key_gen(-1, -2, -3) == hkl_key_gen(1,2,3));
    CHECK(hkl_key_gen(-1,-2,-3) == hkl_key_gen(-1,-2,-3));
    CHECK(hkl_key_gen(1,2,3) == hkl_key_gen(-2,1,3));
    CHECK(hkl_key_gen(1,2,3) == hkl_key_gen(-1,-2,3));
    CHECK(hkl_key_gen(1,2,3) == hkl_key_gen(2,-1,3));
    CHECK(hkl_key_gen(1,2,3) == hkl_key_gen(1,-2,-3));
    CHECK(hkl_key_gen(1,2,3) == hkl_key_gen(-1,2,-3));
    CHECK(hkl_key_gen(1,2,3) == hkl_key_gen(2,1,-3));
    CHECK(hkl_key_gen(1,2,3) == hkl_key_gen(-2, -1, -3));
    
    CHECK(hkl_key_gen(1,2,3) == hkl_key_gen(-2,-1,3));
    CHECK(hkl_key_gen(1,2,3) == hkl_key_gen(2, 1, 3));
}

TEST_CASE("HKLKey_SG96_NoMergeFriedel") {
    HKLKeyGenerator hkl_key_gen(false, *gemmi::find_spacegroup_by_number(96));
    CHECK(hkl_key_gen(-1, -2, -3) != hkl_key_gen(1,2,3));
    CHECK(hkl_key_gen(-1,-2,-3) == hkl_key_gen(-1,-2,-3));
    CHECK(hkl_key_gen(1,2,3) == hkl_key_gen(-2,1,3));
    CHECK(hkl_key_gen(1,2,3) == hkl_key_gen(-1,-2,3));
    CHECK(hkl_key_gen(1,2,3) == hkl_key_gen(2,-1,3));
    CHECK(hkl_key_gen(1,2,3) == hkl_key_gen(1,-2,-3));
    CHECK(hkl_key_gen(1,2,3) == hkl_key_gen(-1,2,-3));
    CHECK(hkl_key_gen(1,2,3) == hkl_key_gen(2,1,-3));
    CHECK(hkl_key_gen(1,2,3) == hkl_key_gen(-2, -1, -3));
    
    CHECK(hkl_key_gen(1,2,3) != hkl_key_gen(-2,-1,3));
    CHECK(hkl_key_gen(1,2,3) != hkl_key_gen(2, 1, 3));
}

TEST_CASE("HKLKey_pack_friedel") {
    HKLKeyGenerator hkl_key_gen(false, *gemmi::find_spacegroup_by_number(1));
    CHECK(hkl_key_gen(-1, -2, -3).pack() != hkl_key_gen(1,2,3).pack());
    CHECK(hkl_key_gen(-1,-2,-3).pack() == hkl_key_gen(-1,-2,-3).pack());
    CHECK(hkl_key_gen(-1,-2,-3).pack() != hkl_key_gen(1,-2,-3).pack());
}

TEST_CASE("HKLKey_pack_no_friedel") {
    HKLKeyGenerator hkl_key_gen(true, *gemmi::find_spacegroup_by_number(1));
    CHECK(hkl_key_gen(-1, -2, -3).pack() == hkl_key_gen(1,2,3).pack());
    CHECK(hkl_key_gen(-1,-2,-3).pack() == hkl_key_gen(-1,-2,-3).pack());
    CHECK(hkl_key_gen(-1,-2,-3).pack() != hkl_key_gen(1,-2,-3).pack());
}


TEST_CASE("HKLKey_sys_absence_P212121") {
    HKLKeyGenerator hkl_key_gen(false, *gemmi::find_spacegroup_by_number(19));
    CHECK(hkl_key_gen.IsSystematicallyAbsent(5,0,0));
    CHECK(!hkl_key_gen.IsSystematicallyAbsent(6,0,0));
    CHECK(hkl_key_gen.IsSystematicallyAbsent(0,5,0));
    CHECK(hkl_key_gen.IsSystematicallyAbsent(0,0,5));
    CHECK(!hkl_key_gen.IsSystematicallyAbsent(0,4,0));
    CHECK(!hkl_key_gen.IsSystematicallyAbsent(5,5,5));
}
TEST_CASE("AcceptReflection_ResolutionLimits") {
    Reflection r{};
    r.I = 100.0f;
    r.sigma = 5.0f;
    r.prescaling_corr = 1.0f;
    r.d = 20.0f;

    // No limits: only the finiteness checks apply.
    CHECK(AcceptReflection(r, std::nullopt, std::nullopt));

    // Low-resolution limit rejects anything coarser than the limit, and is exclusive at it.
    CHECK_FALSE(AcceptReflection(r, std::nullopt, std::optional<double>(15.0)));
    CHECK(AcceptReflection(r, std::nullopt, std::optional<double>(20.0)));
    CHECK(AcceptReflection(r, std::nullopt, std::optional<double>(50.0)));

    // High-resolution limit still rejects anything finer, in the same direction as before.
    CHECK_FALSE(AcceptReflection(r, std::optional<double>(25.0), std::nullopt));
    CHECK(AcceptReflection(r, std::optional<double>(2.0), std::optional<double>(50.0)));

    // The plain-double overload treats 0 as "no limit" at both ends.
    CHECK(AcceptReflection(r, 0.0, 0.0));
    CHECK_FALSE(AcceptReflection(r, 0.0, 15.0));
    CHECK(AcceptReflection(r, 2.0, 50.0));
}

// --- Completeness denominator --------------------------------------------------------------------

namespace {
    // A merged set built straight out of the reflections the cell and space group can give, so the
    // test says exactly which of them were measured: every unique reflection between d_min and
    // d_max_measured and none outside. Without Friedel merging an acentric contributes both hands.
    std::vector<MergedReflection> MeasuredBetween(const gemmi::SpaceGroup &sg, const UnitCell &cell,
                                                  double d_min, double d_max_measured,
                                                  bool merge_friedel) {
        const gemmi::UnitCell gemmi_cell = cell;
        const gemmi::GroupOps gops = sg.operations();
        std::vector<MergedReflection> out;
        for (const auto &hkl: gemmi::make_miller_vector(gemmi_cell, &sg, d_min, d_max_measured, true)) {
            MergedReflection r;
            r.h = hkl[0];
            r.k = hkl[1];
            r.l = hkl[2];
            r.d = static_cast<float>(gemmi_cell.calculate_d(hkl));
            r.I = 100.0f;
            r.sigma = 10.0f;
            r.I_half[0] = 100.0f;
            r.I_half[1] = 100.0f;
            out.push_back(r);
            if (!merge_friedel && !gops.is_reflection_centric(hkl)) {
                r.h = -hkl[0];
                r.k = -hkl[1];
                r.l = -hkl[2];
                out.push_back(r);
            }
        }
        return out;
    }

    MergeStatistics StatsWithLowLimit(const gemmi::SpaceGroup &sg, const std::optional<UnitCell> &cell,
                                      const std::vector<MergedReflection> &merged,
                                      std::optional<double> low_limit, bool merge_friedel) {
        DiffractionExperiment x;
        x.SetSpaceGroup(sg);
        ScalingSettings s = x.GetScalingSettings();
        s.LowResolutionLimit_A(low_limit);
        s.MergeFriedel(merge_friedel);
        x.ImportScalingSettings(s);

        MergeOnTheFly merge(x);
        merge.ReferenceCell(cell);
        return merge.MergeStats(merged, {});
    }

    double Completeness(const MergeStatisticsShell &s) {
        return s.possible_unique_reflections > 0
            ? 100.0 * s.unique_reflections / s.possible_unique_reflections : 0.0;
    }

    const gemmi::SpaceGroup &TestSpaceGroup() { return gemmi::get_spacegroup_by_name("P 1 2 1"); }
    constexpr UnitCell TEST_CELL{40, 45, 50, 90, 100, 90};   // synthetic; coarsest reflection ~49 A
}

// The low-resolution terms a beam stop ate must count as missing: the denominator is the declared
// range, so widening the declared range lowers completeness rather than leaving it alone.
TEST_CASE("MergeStats_CompletenessFallsWhenTheLowBoundCrossesAMaskedRegion") {
    const auto &sg = TestSpaceGroup();
    // Nothing coarser than 20 A was measured - it is all behind the stop.
    const auto merged = MeasuredBetween(sg, TEST_CELL, 2.0, 20.0, true);
    REQUIRE(!merged.empty());

    const auto at_20 = StatsWithLowLimit(sg, TEST_CELL, merged, 20.0, true);
    const auto at_50 = StatsWithLowLimit(sg, TEST_CELL, merged, 50.0, true);

    // Declared exactly where the data stop: everything possible was measured.
    CHECK(Completeness(at_20.overall) > 99.0);

    // Declared out to 50 A: the 20-50 A shell is in the denominator and in nothing else.
    CHECK(at_50.overall.possible_unique_reflections > at_20.overall.possible_unique_reflections);
    CHECK(at_50.overall.unique_reflections == at_20.overall.unique_reflections);
    CHECK(Completeness(at_50.overall) < Completeness(at_20.overall));
    // The innermost shell is where it bites.
    CHECK(Completeness(at_50.shells.front()) < Completeness(at_20.shells.front()));
}

// No low-resolution limit means the whole sphere. The cell has no reflection coarser than 50 A, so
// freeing the 50 A limit must count the same set - the fix does not presuppose either default.
TEST_CASE("MergeStats_CompletenessWithNoDeclaredLowLimit") {
    const auto &sg = TestSpaceGroup();
    const auto merged = MeasuredBetween(sg, TEST_CELL, 2.0, 20.0, true);

    const auto at_50 = StatsWithLowLimit(sg, TEST_CELL, merged, 50.0, true);
    const auto unlimited = StatsWithLowLimit(sg, TEST_CELL, merged, std::nullopt, true);

    CHECK(unlimited.overall.possible_unique_reflections == at_50.overall.possible_unique_reflections);
    CHECK(Completeness(unlimited.overall) == Catch::Approx(Completeness(at_50.overall)));
    // The shell table stays finite even though the bound is not.
    CHECK(std::isfinite(unlimited.shells.front().d_max));
}

// Counting the two Bijvoet mates of an acentric separately doubles the denominator too, so a fully
// measured anomalous set is 100% complete and not 200%.
TEST_CASE("MergeStats_CompletenessNeverExceeds100") {
    const auto &sg = TestSpaceGroup();
    for (const bool merge_friedel: {true, false}) {
        const auto merged = MeasuredBetween(sg, TEST_CELL, 2.0, 50.0, merge_friedel);
        const auto stats = StatsWithLowLimit(sg, TEST_CELL, merged, 50.0, merge_friedel);
        INFO("merge_friedel = " << merge_friedel);
        CHECK(Completeness(stats.overall) <= 100.0);
        CHECK(Completeness(stats.overall) > 99.0);
        for (const auto &sh: stats.shells)
            CHECK(Completeness(sh) <= 100.0);
    }
}

// The shell grid and the denominator share their bounds, so every possible reflection lands in a
// shell: the shells sum to the overall, and the overall is the sphere the run declared.
TEST_CASE("MergeStats_PossibleSumsOverTheShellsToTheDeclaredSphere") {
    const auto &sg = TestSpaceGroup();
    const auto merged = MeasuredBetween(sg, TEST_CELL, 2.0, 20.0, true);
    const auto stats = StatsWithLowLimit(sg, TEST_CELL, merged, 50.0, true);

    int sum = 0;
    for (const auto &sh: stats.shells)
        sum += sh.possible_unique_reflections;
    CHECK(sum == stats.overall.possible_unique_reflections);

    // Counted independently over the same declared range - nothing is lost between the two.
    const gemmi::UnitCell gemmi_cell = TEST_CELL;
    const int expected = gemmi::count_reflections(gemmi_cell, &sg, stats.overall.d_min * 0.999, 50.0, true);
    CHECK(stats.overall.possible_unique_reflections == expected);
}

// Without a reference cell there is no set to count against; completeness stays unmeasured rather
// than becoming a number, with or without a declared low limit.
TEST_CASE("MergeStats_NoReferenceCellLeavesCompletenessUnmeasured") {
    const auto &sg = TestSpaceGroup();
    const auto merged = MeasuredBetween(sg, TEST_CELL, 2.0, 20.0, true);

    for (const std::optional<double> low_limit: {std::optional<double>(50.0), std::optional<double>()}) {
        const auto stats = StatsWithLowLimit(sg, std::nullopt, merged, low_limit, true);
        CHECK(stats.overall.possible_unique_reflections == 0);
        CHECK(stats.overall.unique_reflections > 0);
        CHECK(Completeness(stats.overall) == 0.0);
    }
}

// The reference-range table (--report-resolution) is binned over the range it is GIVEN, whether or
// not the data reach it: its last shell ends at the declared d_min, and past the run's own limit
// nothing is read from anywhere - the reflections the run did not keep count as missing.
TEST_CASE("MergeStats_ReferenceRangeIsBinnedOverTheDeclaredRange") {
    const auto &sg = TestSpaceGroup();
    const auto merged = MeasuredBetween(sg, TEST_CELL, 2.0, 20.0, true);   // what the run kept
    DiffractionExperiment x;
    x.SetSpaceGroup(sg);
    MergeOnTheFly merge(x);
    merge.ReferenceCell(TEST_CELL);
    const auto own = merge.MergeStats(merged, {});

    // A range coarser than the run's own: a complete subset, on a grid ending at the declared bound.
    const auto within = merge.MergeStats(merged, {}, {}, std::nullopt, ReportResolutionRange{3.0, 50.0});
    CHECK(within.shells.back().d_min == Catch::Approx(3.0));
    CHECK(within.shells.front().d_max == Catch::Approx(50.0));
    int coarser_than_3 = 0;
    for (const auto &m: merged)
        if (m.d > 3.0f) ++coarser_than_3;
    CHECK(within.overall.unique_reflections == coarser_than_3);
    CHECK(within.overall.unique_reflections < own.overall.unique_reflections);
    CHECK(Completeness(within.overall) > 99.0);

    // A range finer than the run kept: the grid still ends at 1.5 A, the outer shell is empty and its
    // possible reflections are in the denominator, and the measured range says where the data stop.
    const auto beyond = merge.MergeStats(merged, {}, {}, std::nullopt, ReportResolutionRange{1.5, 50.0});
    CHECK(beyond.shells.back().d_min == Catch::Approx(1.5));
    CHECK(beyond.overall.unique_reflections == own.overall.unique_reflections);
    CHECK(beyond.overall.possible_unique_reflections > own.overall.possible_unique_reflections);
    CHECK(Completeness(beyond.overall) < Completeness(own.overall));
    CHECK(beyond.shells.back().unique_reflections == 0);
    CHECK(beyond.shells.back().possible_unique_reflections > 0);
    CHECK(beyond.overall.d_min == Catch::Approx(own.overall.d_min));

    // A range the data never reach at all is still a grid, not an error.
    const auto empty = merge.MergeStats(merged, {}, {}, std::nullopt, ReportResolutionRange{1.0, 1.5});
    CHECK(empty.overall.unique_reflections == 0);
    CHECK(empty.overall.possible_unique_reflections > 0);
    CHECK(empty.overall.d_min == 0.0f);
}

// ---------------------------------------------------------------- the automatic resolution cutoff
namespace {
    // Half-set pairs spread uniformly in s = 1/d^2 over [s_from, s_to), either correlated with each
    // other (signal) or drawn independently (noise, CC1/2 ~ 0), so a whole CC1/2 curve can be built
    // band by band.
    void AddBand(std::vector<MergedReflection> &v, std::mt19937 &rng,
                 double s_from, double s_to, int n, bool correlated, double sigma = 1.0) {
        std::normal_distribution<double> g(0.0, sigma);
        for (int j = 0; j < n; ++j) {
            MergedReflection m;
            m.d = static_cast<float>(1.0 / std::sqrt(s_from + (j + 0.5) * (s_to - s_from) / n));
            const double a = g(rng), b = g(rng);
            m.I_half[0] = static_cast<float>(a);
            m.I_half[1] = static_cast<float>(correlated ? a : b);
            // The pair's weight in a CC1/2: the precision it would have had at the typical frame scale
            // over the precision it has - 1/sigma^2 for a band scaled up by sigma from dead frames.
            m.cc_weight = static_cast<float>(1.0 / (sigma * sigma));
            v.push_back(m);
        }
    }
}

// A clean fall-off: CC1/2 crosses the target where the signal stops, and the cut is written one
// shell past it.
TEST_CASE("ResolutionCutoff_CleanFallOff") {
    Logger logger("test");
    std::mt19937 rng(12345);
    std::vector<MergedReflection> merged;
    AddBand(merged, rng, 0.01, 0.25, 480, true);    // signal to 1/sqrt(0.25) = 2.00 A
    AddBand(merged, rng, 0.25, 0.51, 520, false);   // noise beyond it

    const auto rc = ComputeCCHalfLogisticCutoff(merged, 0.30, logger);
    REQUIRE(rc.d_fit);
    CHECK(*rc.d_fit == Catch::Approx(2.0).margin(0.15));
    REQUIRE(rc.d_cut);
    CHECK(*rc.d_cut < *rc.d_fit);                   // the deliberate one-shell extension
    CHECK(*rc.d_cut == Catch::Approx(1.92).margin(0.15));
}

// A sweep with a long stretch where the crystal barely diffracted: the reflections measured only there
// are noise scaled up by 1/G, sigmas with them, and there are many of them at every resolution. Counted
// by the information they carry they must not hide where the well-measured reflections stop.
TEST_CASE("ResolutionCutoff_ScaledUpNoiseDoesNotHideTheFallOff") {
    Logger logger("test");
    std::mt19937 rng(12345);
    std::vector<MergedReflection> merged;
    AddBand(merged, rng, 0.01, 0.25, 480, true);
    AddBand(merged, rng, 0.25, 0.51, 520, false);
    AddBand(merged, rng, 0.01, 0.51, 300, false, 50.0);

    const auto rc = ComputeCCHalfLogisticCutoff(merged, 0.30, logger);
    REQUIRE(rc.d_fit);
    CHECK(*rc.d_fit == Catch::Approx(2.0).margin(0.15));

    // Counted as equals, the same reflections hide it: the curve reads noise from the first bin on.
    for (auto &m : merged) m.cc_weight = 1.0f;
    const auto unweighted = ComputeCCHalfLogisticCutoff(merged, 0.30, logger);
    CHECK((!unweighted.d_fit || *unweighted.d_fit > 3.0));
}

// The frame factor of the CC1/2 weight is 1 on a sweep whose frames all sit at one scale, however that
// scale is spread over the frames, and 1 for any frame brighter than typical.
TEST_CASE("CCHalfFrameFactors_UniformScaleIsOne") {
    const auto f = CCHalfFrameFactors({1.7, 1.7, 1.7, 1.7}, {100, 3, 250, 0});
    for (double x : f) CHECK(x == 1.0);
    const auto g = CCHalfFrameFactors({1.0, 2.0}, {100, 100});
    CHECK(g[1] == 1.0);
    CHECK(g[0] == Catch::Approx(9.0 / 5.0 * 9.0 / 5.0));   // G_ref = (1 + 8) / (1 + 4)
}

// Most of the sweep at 2% of the good frames' scale: the typical scale is still the good frames', not
// the run median, and an observation from a dead frame carries 2500x the variance of a good one.
TEST_CASE("CCHalfFrameFactors_DeadStretchDoesNotSetTheTypicalScale") {
    std::vector<double> scale;
    std::vector<int64_t> n;
    for (int f = 0; f < 100; ++f) {
        scale.push_back(f < 60 ? 0.02 : 1.0);
        n.push_back(50);
    }
    scale.push_back(NAN);   // a frame without a scale counts as G = 1
    n.push_back(50);
    const auto factor = CCHalfFrameFactors(scale, n);
    CHECK(factor[0] == Catch::Approx(2500.0).epsilon(0.01));
    CHECK(factor[99] == Catch::Approx(1.0).epsilon(1e-3));
    CHECK(factor[100] == Catch::Approx(1.0).epsilon(1e-3));
}

// A fall-off region a logistic cannot follow: CC1/2 drops through the target and comes straight back
// up. The fitted crossing is then an extrapolation far past the bins it was made over, and reading
// the cut off it writes the data deep into the noise; the crossing the bins themselves show is where
// the signal stopped, and that is what must be used.
TEST_CASE("ResolutionCutoff_RaggedFallOffIsReadOffTheBins") {
    Logger logger("test");
    std::mt19937 rng(12345);
    std::vector<MergedReflection> merged;
    AddBand(merged, rng, 0.01, 0.13, 240, true);    // signal to 1/sqrt(0.13) = 2.77 A
    AddBand(merged, rng, 0.13, 0.17, 80, false);    // a hole below the target
    AddBand(merged, rng, 0.17, 0.25, 160, true);    // correlated again - not a fall-off
    AddBand(merged, rng, 0.25, 0.51, 520, false);

    const auto rc = ComputeCCHalfLogisticCutoff(merged, 0.30, logger);
    REQUIRE(rc.d_fit);
    CHECK(*rc.d_fit == Catch::Approx(2.77).margin(0.20));
    REQUIRE(rc.d_cut);
    CHECK(*rc.d_cut > 2.30);                        // coarser than the band that correlates again
}