Jungfraujoch/tests/RotationScaleWalkTest.cpp

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

#include <cmath>
#include <catch2/catch_all.hpp>

#include "../rugnux/Rugnux.h"
#include "../image_analysis/scale_merge/RotationScaleMerge.h"

namespace {
    // A synthetic sweep whose stage turned `true_scale` times the stored angles. Scored at a scale k,
    // the validation spots stay on the lattice as long as the angles track the rotation, and the share
    // that does falls off with the relative rate error; a wrong spindle angle keeps half a percent.
    struct SyntheticSweep {
        double true_scale;
        int64_t spots = 35000;
        int index_calls = 0;
        int refit_calls = 0;

        ValidationSpotEvidence IndexAt(float k) {
            ++index_calls;
            const double error = std::fabs(true_scale / k - 1.0);
            const double on = 0.9 * std::max(0.0, 1.0 - 30.0 * error);
            return ValidationSpotEvidence{spots, std::llround(on * spots), std::llround(0.005 * spots)};
        }

        // The post-refinement at k, relative to k. It reads only 70 % of the error that is left: it
        // sees only the frames the angles at k still track.
        std::optional<double> RefitAt(float k) {
            ++refit_calls;
            return 1.0 + 0.7 * (true_scale / k - 1.0);
        }

        RotationScaleWalk Walk(double first_fit, int64_t sweep_images = 3600) {
            return WalkRotationScale(first_fit, [this](float k) { return IndexAt(k); },
                                     [this](float k) { return RefitAt(k); }, 8, sweep_images);
        }
    };
}

TEST_CASE("ValidationEvidencePrefers", "[RotationScale]") {
    const ValidationSpotEvidence base{10000, 3000, 50};
    // 1 % more of the spots beyond chance is under the noise of two 30 % shares over 10000 spots
    // (sqrt(2 * 0.3 * 0.7 / 10000) = 0.65 %, times 3.29); 5 % is well over it.
    CHECK_FALSE(ValidationEvidencePrefers(base, {10000, 3100, 50}));
    CHECK(ValidationEvidencePrefers(base, {10000, 3500, 50}));
    // A candidate is judged against its own null: more spots on the lattice bought by a null that
    // rose just as much is no gain.
    CHECK_FALSE(ValidationEvidencePrefers(base, {10000, 3500, 550}));
    // Never against itself, and nothing that scored nothing wins.
    CHECK_FALSE(ValidationEvidencePrefers(base, base));
    CHECK_FALSE(ValidationEvidencePrefers(base, {}));
    CHECK(ValidationEvidencePrefers({}, base));
}

TEST_CASE("WalkRotationScale_ReachesTheFixedPoint", "[RotationScale]") {
    // A stage 3 % slow. The first fit reads 70 % of that; each refit at the adopted scale reads 70 %
    // of what is left, and the walk goes on as long as the validation frames prefer the new scale.
    SyntheticSweep sweep{0.97};
    const auto walk = sweep.Walk(sweep.RefitAt(1.0f).value());
    CHECK(walk.scale == Catch::Approx(0.97).margin(0.001));
    CHECK(walk.scale != 1.0f);
    CHECK(walk.evidence.on_lattice > sweep.IndexAt(1.0f).on_lattice);
    CHECK(sweep.refit_calls > 2);
    CHECK_FALSE(walk.trail.empty());
}

TEST_CASE("WalkRotationScale_StopsWhereTheNextStepIsUnderHalfAFrame", "[RotationScale]") {
    // Spots enough for the validation frames to prefer every step: what ends the walk is the step the
    // contraction predicts, measured against the sweep. Each refit reads 70 % of the error left, so the
    // steps contract by 0.3.
    SyntheticSweep short_sweep{0.97};
    short_sweep.spots = 1000000000;
    const auto a = short_sweep.Walk(short_sweep.RefitAt(1.0f).value(), 360);
    SyntheticSweep long_sweep{0.97};
    long_sweep.spots = 1000000000;
    const auto b = long_sweep.Walk(long_sweep.RefitAt(1.0f).value(), 36000);
    CHECK(short_sweep.refit_calls < long_sweep.refit_calls);
    // Where it stops, what is left of the error moves the end of the sweep by under a frame.
    CHECK(std::fabs(a.scale - 0.97) * 360 < 1.0);
    CHECK(std::fabs(b.scale - 0.97) * 36000 < 1.0);
}

TEST_CASE("WalkRotationScale_RefitStepTheFramesCannotResolve", "[RotationScale]") {
    // Validation frames that resolve the first step off the stored angles but score every later one
    // alike: the refit is followed to its fixed point all the same.
    struct FlatPastFirstStep : SyntheticSweep {
        ValidationSpotEvidence IndexAt(float k) {
            ++index_calls;
            const double on = k == 1.0f ? 0.37 : 0.42;
            return ValidationSpotEvidence{spots, std::llround(on * spots), std::llround(0.005 * spots)};
        }
    } sweep{{0.98}};
    const auto walk = WalkRotationScale(sweep.RefitAt(1.0f).value(), [&](float k) { return sweep.IndexAt(k); },
                                        [&](float k) { return sweep.RefitAt(k); }, 8, 3600);
    CHECK(walk.scale == Catch::Approx(0.98).margin(0.0005));

    // A refit step the frames score worse is still refused.
    struct WorseAfterFirstStep : SyntheticSweep {
        ValidationSpotEvidence IndexAt(float k) {
            ++index_calls;
            const double on = k == 1.0f ? 0.37 : (k == first ? 0.42 : 0.30);
            return ValidationSpotEvidence{spots, std::llround(on * spots), std::llround(0.005 * spots)};
        }
        float first = 0.0f;
    } worse{{0.98}};
    worse.first = static_cast<float>(worse.RefitAt(1.0f).value());
    const auto refused = WalkRotationScale(worse.first, [&](float k) { return worse.IndexAt(k); },
                                           [&](float k) { return worse.RefitAt(k); }, 8, 3600);
    CHECK(refused.scale == worse.first);
}

TEST_CASE("WalkRotationScale_StoredAnglesStand", "[RotationScale]") {
    SECTION("A healthy stage: a fit off by noise scores no better than the stored angles") {
        SyntheticSweep sweep{1.0};
        const auto walk = sweep.Walk(1.0002);
        CHECK(walk.scale == 1.0f);
        CHECK(sweep.index_calls == 2);   // the stored angles and the fit, nothing more
        CHECK(sweep.refit_calls == 0);
    }
    SECTION("A real but small error the spots cannot resolve beyond their noise") {
        SyntheticSweep sweep{0.999};
        sweep.spots = 400;
        const auto walk = sweep.Walk(0.9993);
        CHECK(walk.scale == 1.0f);
    }
    SECTION("A fit that tracks something other than the rotation scores worse, and is refused") {
        SyntheticSweep sweep{1.0};
        const auto walk = sweep.Walk(0.98);
        CHECK(walk.scale == 1.0f);
        CHECK(sweep.refit_calls == 0);
    }
    SECTION("A fit of exactly one asks for no probe at all") {
        SyntheticSweep sweep{1.0};
        const auto walk = sweep.Walk(1.0);
        CHECK(walk.scale == 1.0f);
        CHECK(walk.trail.empty());
        CHECK(sweep.index_calls == 0);
    }
}

TEST_CASE("SmoothLogScale_FollowsInformationBridgesGaps", "[RotationScale]") {
    const int n = 60;
    // A ramp is no curvature, so any amount of smoothing keeps it exactly.
    std::vector<double> ramp(n), J(n, 1.0);
    for (int f = 0; f < n; ++f) ramp[f] = -0.1 * f;
    auto x = RotationScaleMerge::SmoothLogScale(ramp, J, 1e6);
    for (int f = 0; f < n; ++f) CHECK(x[f] == Catch::Approx(ramp[f]).margin(1e-6));
    // A stretch with no information is bridged by the straight line through its neighbours.
    std::vector<double> Jgap(J);
    for (int f = 20; f < 40; ++f) Jgap[f] = 0.0;
    x = RotationScaleMerge::SmoothLogScale(ramp, Jgap, 1.0);
    CHECK(x[30] == Catch::Approx(-3.0).margin(1e-6));
    // A frame with far more information than its neighbours keeps its own value.
    std::vector<double> y(n, 0.0), Jone(n, 1.0);
    y[30] = 1.0; Jone[30] = 1e6;
    x = RotationScaleMerge::SmoothLogScale(y, Jone, 10.0);
    CHECK(x[30] == Catch::Approx(1.0).margin(1e-3));
    // Fewer than two frames with information: nothing to smooth against.
    std::vector<double> Jsingle(n, 0.0);
    Jsingle[5] = 1.0;
    CHECK(RotationScaleMerge::SmoothLogScale(y, Jsingle, 1.0) == y);
}

TEST_CASE("ChooseLogScaleSmoothing_SmoothsNoiseFollowsSignal", "[RotationScale]") {
    const int n = 400;
    std::vector<double> J(n, 1.0), noisy(n), step(n);
    // Deterministic noise about a flat scale: the chosen curve is close to flat.
    for (int f = 0; f < n; ++f) noisy[f] = 0.3 * std::sin(12.9898 * f) * std::cos(78.233 * f);
    const double l_noise = RotationScaleMerge::ChooseLogScaleSmoothing(noisy, J, 1);
    const auto flat = RotationScaleMerge::SmoothLogScale(noisy, J, l_noise);
    double rms = 0.0;
    for (double v : flat) rms += v * v;
    CHECK(std::sqrt(rms / n) < 0.05);
    // A precise slow wave is followed.
    for (int f = 0; f < n; ++f) step[f] = 2.0 * std::sin(f / 30.0);
    const double l_wave = RotationScaleMerge::ChooseLogScaleSmoothing(step, J, 1);
    const auto wave = RotationScaleMerge::SmoothLogScale(step, J, l_wave);
    CHECK(wave[47] == Catch::Approx(step[47]).margin(0.01));
    CHECK(l_wave < l_noise);
}

TEST_CASE("ValidationEvidenceBeatsChance_ExchangeableArrangements", "[RotationScale]") {
    // Four reflections taken at the recorded angles and none at any wrong one: (1/6)^4 < 0.001.
    CHECK(ValidationEvidenceBeatsChance({.reflections = 4, .reflections_by_chance = 0}));
    CHECK_FALSE(ValidationEvidenceBeatsChance({.reflections = 3, .reflections_by_chance = 0}));
    // The same four with two more spread over the wrong angles is what chance does once in ~100.
    CHECK_FALSE(ValidationEvidenceBeatsChance({.reflections = 4, .reflections_by_chance = 2}));
    CHECK_FALSE(ValidationEvidenceBeatsChance({}));
    // A real lattice: thousands against a handful.
    CHECK(ValidationEvidenceBeatsChance({.reflections = 3000, .reflections_by_chance = 250}));
    // As many at the recorded angles as at an average wrong one is no lattice.
    CHECK_FALSE(ValidationEvidenceBeatsChance({.reflections = 100, .reflections_by_chance = 500}));
}

TEST_CASE("ValidationEvidenceBeatsChance_NoCrystalSignificance", "[RotationScale]") {
    // A handful of stray reflections passes 0.001 but not the no-crystal test's own level.
    CHECK(ValidationEvidenceBeatsChance({.reflections = 5, .reflections_by_chance = 1}));
    CHECK_FALSE(ValidationEvidenceBeatsChance({.reflections = 5, .reflections_by_chance = 1}, NO_CRYSTAL_SIGNIFICANCE));
    // Against two by chance it takes 15 (tail 10^-9.7), and 14 is not enough (10^-8.96).
    CHECK(ValidationEvidenceBeatsChance({.reflections = 15, .reflections_by_chance = 2}, NO_CRYSTAL_SIGNIFICANCE));
    CHECK_FALSE(ValidationEvidenceBeatsChance({.reflections = 14, .reflections_by_chance = 2}, NO_CRYSTAL_SIGNIFICANCE));
    // A weak real lattice on a few frames is still tens of decades past it.
    CHECK(ValidationEvidenceBeatsChance({.reflections = 115, .reflections_by_chance = 4}, NO_CRYSTAL_SIGNIFICANCE));
}

TEST_CASE("ValidationEvidenceLog10Tail_Exact", "[RotationScale]") {
    CHECK(ValidationEvidenceLog10Tail({.reflections = 4, .reflections_by_chance = 2}) == Catch::Approx(-2.0604).margin(1e-4));
    CHECK(ValidationEvidenceLog10Tail({.reflections = 5, .reflections_by_chance = 2}) == Catch::Approx(-2.6981).margin(1e-4));
    // Far below the smallest double, still finite.
    CHECK(ValidationEvidenceLog10Tail({.reflections = 1354, .reflections_by_chance = 777}) == Catch::Approx(-509.691).margin(1e-3));
    CHECK(ValidationEvidenceLog10Tail({}) == 0.0);
}

namespace {
    // A frame of `strong` reflections at I/sigma 10, `up` and `down` at +4 and -4 (noise in both
    // tails), and enough at +-0.5 to make 200 in all, every one on a background of `bkg`.
    IntegrationOutcome FrameOf(int strong, float bkg, int up = 0, int down = 0) {
        IntegrationOutcome io;
        auto add = [&](float z) {
            Reflection r{};
            r.I = z;
            r.sigma = 1.0f;
            r.bkg = bkg;
            io.reflections.push_back(r);
        };
        for (int i = 0; i < strong; ++i) add(10.0f);
        for (int i = 0; i < up; ++i) add(4.0f);
        for (int i = 0; i < down; ++i) add(-4.0f);
        for (int i = strong + up + down; i < 200; ++i) add(i % 2 ? 0.5f : -0.5f);
        return io;
    }
}

TEST_CASE("TrimBlankSweepEnds_ShutterEndsOnly", "[RotationScale]") {
    std::vector<IntegrationOutcome> sweep;
    sweep.push_back(FrameOf(0, 0.0f));      // 0: the shutter not yet open
    sweep.push_back(FrameOf(20, 3.0f));     // 1: partly exposed - it recorded something, so it stays
    for (int f = 2; f < 37; ++f)            // 20: no beam inside the sweep - the ledger's, not ours
        sweep.push_back(f == 20 ? FrameOf(0, 0.0f) : FrameOf(40, 10.0f));
    sweep.push_back(FrameOf(0, 2.0f));      // 37: closed; a dark level, as a CCD leaves
    sweep.push_back(FrameOf(0, 0.0f, 4, 4));   // 38: closed, noise in both tails
    sweep.emplace_back();                   // 39: nothing integrated - passed over, not counted

    const auto [at_start, at_end] = TrimBlankSweepEnds(sweep, 2);
    CHECK(at_start == 1);
    CHECK(at_end == 2);
    CHECK(sweep[0].reflections.empty());
    CHECK(sweep[1].reflections.size() == 200);
    CHECK(sweep[20].reflections.size() == 200);
    CHECK(sweep[36].reflections.size() == 200);
    CHECK(sweep[37].reflections.empty());
    CHECK(sweep[38].reflections.empty());
}

TEST_CASE("TrimBlankSweepEnds_CrystalOutOfTheBeamStays", "[RotationScale]") {
    // No diffraction at the end, but the beam was on - the background is still there. That is the crystal
    // leaving the beam, which the sweep ledger reports; it is not a shutter.
    std::vector<IntegrationOutcome> sweep;
    for (int f = 0; f < 36; ++f)
        sweep.push_back(FrameOf(40, 10.0f));
    for (int f = 36; f < 40; ++f)
        sweep.push_back(FrameOf(0, 8.0f));
    const auto [at_start, at_end] = TrimBlankSweepEnds(sweep, 2);
    CHECK(at_start == 0);
    CHECK(at_end == 0);
    CHECK(sweep.back().reflections.size() == 200);
}

TEST_CASE("TrimBlankSweepEnds_WeakCrystalKeepsItsEnds", "[RotationScale]") {
    // Three strong reflections a frame: one frame cannot tell this crystal from no crystal, so even a
    // frame without beam at an end is not called.
    std::vector<IntegrationOutcome> sweep;
    sweep.push_back(FrameOf(0, 0.0f));
    for (int f = 1; f < 39; ++f)
        sweep.push_back(FrameOf(3, 10.0f));
    sweep.push_back(FrameOf(0, 0.0f));
    const auto [at_start, at_end] = TrimBlankSweepEnds(sweep, 2);
    CHECK(at_start == 0);
    CHECK(at_end == 0);
    CHECK(sweep.front().reflections.size() == 200);
    CHECK(sweep.back().reflections.size() == 200);
}