Jungfraujoch/tests/WilsonOutliersTest.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 <random>
#include <vector>

#include "../image_analysis/scale_merge/WilsonOutliers.h"

namespace {
    // Acentric Wilson intensities of mean `mean` between 3.9 and 2.0 A, measured `per_unit` times each
    // with Poisson-like noise of variance I + bkg_var.
    std::vector<WilsonObservation> WilsonPopulation(int n_units, int per_unit, double mean, double bkg_var,
                                                    uint32_t seed) {
        std::mt19937 rng(seed);
        std::exponential_distribution<double> wilson(1.0 / mean);
        std::normal_distribution<double> gauss(0.0, 1.0);
        std::vector<WilsonObservation> v;
        for (int u = 0; u < n_units; ++u) {
            const double I_true = wilson(rng);
            const float d = 3.9f - 1.9f * static_cast<float>(u) / n_units;
            for (int m = 0; m < per_unit; ++m) {
                const double sigma = std::sqrt(I_true + bkg_var);
                v.push_back({static_cast<float>(I_true + sigma * gauss(rng)), static_cast<float>(sigma), d,
                             1.0f, false, false, u});
            }
        }
        return v;
    }
}

TEST_CASE("WilsonOutliers: the artefact member of a discordant pair is dropped", "[wilson_outliers]") {
    auto obs = WilsonPopulation(10000, 2, 1000.0, 100.0, 1);
    // A pair of one ordinary observation and one two hundred times the shell mean.
    const int32_t unit = 10000;
    obs.push_back({800.0f, 30.0f, 2.5f, 1.0f, false, false, unit});
    obs.push_back({200000.0f, 450.0f, 2.5f, 1.0f, false, false, unit});

    const auto r = WilsonOutliers(obs, 0.01);
    CHECK(r.n_tested == obs.size());
    CHECK(r.tail_scale == Catch::Approx(1.0).margin(0.2));
    CHECK(r.rejected[obs.size() - 1] == 1);
    CHECK(r.rejected[obs.size() - 2] == 0);
    CHECK(r.n_rejected == 1);
}

TEST_CASE("WilsonOutliers: a clipped mate does not testify against a strong observation", "[wilson_outliers]") {
    auto obs = WilsonPopulation(10000, 2, 1000.0, 100.0, 10);
    // The strong member is real; the low one lost its saturated core to the mask.
    const int32_t unit = 10000;
    obs.push_back({800.0f, 30.0f, 2.5f, 1.0f, false, true, unit});
    obs.push_back({200000.0f, 450.0f, 2.5f, 1.0f, false, false, unit});
    const auto r = WilsonOutliers(obs, 0.01);
    CHECK(r.n_rejected == 0);
}

TEST_CASE("WilsonOutliers: two large observations of one reflection confirm each other", "[wilson_outliers]") {
    auto obs = WilsonPopulation(10000, 2, 1000.0, 100.0, 2);
    const int32_t unit = 10000;
    obs.push_back({60000.0f, 300.0f, 2.5f, 1.0f, false, false, unit});
    obs.push_back({61000.0f, 300.0f, 2.5f, 1.0f, false, false, unit});

    const auto r = WilsonOutliers(obs, 0.01);
    CHECK(r.n_rejected == 0);
}

TEST_CASE("WilsonOutliers: a Friedel pair of large observations does not confirm itself", "[wilson_outliers]") {
    auto obs = WilsonPopulation(10000, 2, 1000.0, 100.0, 2);
    // h and -h, both high - another lattice's spot and its own Friedel mate - and a symmetry mate of h
    // (other indices, measured elsewhere) that reads the crystal's ordinary intensity.
    const int32_t unit = 10000;
    obs.push_back({60000.0f, 300.0f, 2.5f, 1.0f, false, false, unit, HKLFriedelKey(5, -3, 7)});
    obs.push_back({61000.0f, 300.0f, 2.5f, 1.0f, false, false, unit, HKLFriedelKey(-5, 3, -7)});
    obs.push_back({900.0f, 40.0f, 2.5f, 1.0f, false, false, unit, HKLFriedelKey(-5, 3, 7)});
    auto r = WilsonOutliers(obs, 0.01);
    CHECK(r.rejected[obs.size() - 3] == 1);
    CHECK(r.rejected[obs.size() - 2] == 1);
    CHECK(r.rejected[obs.size() - 1] == 0);
    CHECK(r.n_rejected == 2);

    // The same two large observations under indices that are not each other's negation (h and a
    // symmetry mate) confirm each other, as before.
    obs[obs.size() - 2].hkl = HKLFriedelKey(5, -3, -7);
    r = WilsonOutliers(obs, 0.01);
    CHECK(r.n_rejected == 0);
}

TEST_CASE("WilsonOutliers: the symmetry enhancement factor keeps an axial reflection", "[wilson_outliers]") {
    auto obs = WilsonPopulation(10000, 2, 1000.0, 100.0, 3);
    // Measured once, 12x <I/epsilon> at epsilon 4: 48x the shell mean, but only E^2 = 12.
    obs.push_back({48000.0f, 250.0f, 2.5f, 4.0f, false, false, 10000});
    auto r = WilsonOutliers(obs, 0.01);
    CHECK(r.rejected.back() == 0);
    CHECK(r.e2.back() == Catch::Approx(12.0).epsilon(0.1));
    CHECK(r.n_rejected == 0);

    // The same observation on a general reflection is improbable.
    obs.back().epsilon = 1.0f;
    r = WilsonOutliers(obs, 0.01);
    CHECK(r.rejected.back() == 1);
}

TEST_CASE("WilsonOutliers: a weak shell does not misfire", "[wilson_outliers]") {
    // <I> = 20 under a background of sd 100: noise alone reaches twenty times <I>, while the shell mean
    // is still established.
    const auto obs = WilsonPopulation(20000, 1, 20.0, 10000.0, 4);
    const auto r = WilsonOutliers(obs, 0.01);
    CHECK(r.n_tested == obs.size());
    CHECK(r.n_rejected == 0);
}

TEST_CASE("WilsonOutliers: an artefact among several mates is dropped", "[wilson_outliers]") {
    auto obs = WilsonPopulation(5000, 4, 1000.0, 100.0, 5);
    // One precise-looking artefact beside three ordinary mates: the mates out-vote it.
    obs[0].I = 500000.0f;
    obs[0].sigma = 700.0f;
    const auto r = WilsonOutliers(obs, 0.01);
    CHECK(r.n_tested == obs.size());
    CHECK(r.rejected[0] == 1);
    CHECK(r.n_rejected == 1);
}

TEST_CASE("WilsonOutliers: a reflection whose observations are mostly large is kept", "[wilson_outliers]") {
    auto obs = WilsonPopulation(5000, 3, 1000.0, 100.0, 8);
    // Two of three observations large, the third low (a partial that caught little, say).
    obs[0].I = 60000.0f; obs[0].sigma = 300.0f;
    obs[1].I = 62000.0f; obs[1].sigma = 300.0f;
    obs[2].I = 500.0f; obs[2].sigma = 30.0f;
    const auto r = WilsonOutliers(obs, 0.01);
    CHECK(r.n_rejected == 0);
}

TEST_CASE("WilsonOutliers: a shell without a measured mean is not judged", "[wilson_outliers]") {
    // Pure noise: <I> = 0 within its error, so nothing can be improbable against it.
    auto obs = WilsonPopulation(10000, 1, 1e-6, 10000.0, 9);
    obs[0].I = 5000.0f; obs[0].sigma = 100.0f;
    const auto r = WilsonOutliers(obs, 0.01);
    CHECK(r.n_tested == 0);
    CHECK(r.n_rejected == 0);
}

TEST_CASE("WilsonOutliers: a heavier-tailed population widens the bound", "[wilson_outliers]") {
    // Half the reflections at twice the mean, half at a tenth - the intensity classes of a strong
    // pseudo-translation. Wilson's single exponential would call the top of the strong class improbable.
    auto strong = WilsonPopulation(10000, 1, 2000.0, 100.0, 6);
    const auto weak = WilsonPopulation(10000, 1, 100.0, 100.0, 7);
    for (auto o : weak) {
        o.unit += 10000;
        strong.push_back(o);
    }
    const auto r = WilsonOutliers(strong, 0.01);
    CHECK(r.tail_scale > 1.5);
    CHECK(r.n_rejected == 0);
}

#ifdef JFJOCH_USE_CUDA
#include <cstring>
#include "../common/CUDAWrapper.h"
#include "../image_analysis/scale_merge/WilsonOutliersGPU.h"

// The device test against the host one, on a population that takes every branch: reflections measured
// one to six times, resolutions on a coarse grid (so many observations share a d and the tie order
// counts) reaching past the 4 A limit, symmetry factors 1-4, centrics, clipped mates, observations
// outside the test (no reflection, NaN, zero sigma), a heavy-tailed mixture and planted artefacts. The
// two results must be the same to the bit.
TEST_CASE("WilsonOutliers: the device test is the host test", "[wilson_outliers]") {
    if (get_gpu_count() == 0) {
        WARN("No CUDA GPU present. Skipping the device Wilson test");
        return;
    }
    for (uint32_t seed : {11u, 12u}) {
        std::mt19937 rng(seed);
        std::uniform_real_distribution<double> uni(0.0, 1.0);
        std::normal_distribution<double> gauss(0.0, 1.0);
        std::vector<WilsonObservation> obs;
        for (int u = 0; u < 120000; ++u) {
            const float d = 1.5f + 0.01f * static_cast<float>(static_cast<int>(uni(rng) * 300));
            const double mean = 1000.0 * std::exp(-2.0 / (d * d)) * (u % 2 ? 2.0 : 0.1);
            const float epsilon = static_cast<float>(1 << static_cast<int>(uni(rng) * 3));
            const bool centric = uni(rng) < 0.1;
            const double I_true = -mean * std::log(1.0 - uni(rng)) * epsilon;
            const int mult = 1 + static_cast<int>(uni(rng) * 6);
            for (int m = 0; m < mult; ++m) {
                const double sigma = std::sqrt(I_true + 100.0);
                double I = I_true + sigma * gauss(rng);
                if (uni(rng) < 0.002) I *= 200.0;
                obs.push_back({static_cast<float>(I), static_cast<float>(sigma), d, epsilon, centric,
                               uni(rng) < 0.05, u, HKLFriedelKey(m % 2 ? -1 : 1, m % 2 ? -2 : 2, m % 3 ? 3 : -3)});
            }
        }
        obs[3].unit = -1;
        obs[5].I = NAN;
        obs[7].sigma = 0.0f;

        const auto host = WilsonOutliers(obs, 0.01, 8);
        const auto dev = WilsonOutliersGPU(obs, 0.01);
        REQUIRE(host.n_rejected > 0);
        CHECK(host.tail_scale > 1.0);
        CHECK(dev.n_tested == host.n_tested);
        CHECK(dev.n_rejected == host.n_rejected);
        CHECK(dev.tail_scale == host.tail_scale);
        CHECK(dev.bound == host.bound);
        CHECK(dev.rejected == host.rejected);
        REQUIRE(dev.e2.size() == host.e2.size());
        CHECK(std::memcmp(dev.e2.data(), host.e2.data(), host.e2.size() * sizeof(float)) == 0);
    }
}
#endif