Jungfraujoch/tests/BeamCenterFFTGPUTest.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 "../common/CUDAWrapper.h"

#ifdef JFJOCH_USE_CUDA

#include <cmath>
#include <cstdint>
#include <random>
#include <vector>

#include "../image_analysis/geom_refinement/BeamCenterFFTCPU.h"
#include "../image_analysis/geom_refinement/BeamCenterFFTGPU.h"

namespace {

// An isotropic scattered background about (cx, cy) with a solvent ring, a module gap and a darker
// sextant - the same shape of image the corpus measurement runs on, at a size a test can afford.
std::vector<float> SynthesiseBackground(int64_t w, int64_t h, float cx, float cy) {
    std::vector<float> mean(static_cast<size_t>(w * h), NAN);
    std::mt19937 rng(20260913);
    std::normal_distribution<float> gauss(0.0f, 1.0f);
    for (int64_t y = 0; y < h; y++) {
        if (y >= 500 && y < 520)
            continue;   // module gap
        for (int64_t x = 0; x < w; x++) {
            const float r = std::hypot(static_cast<float>(x) - cx, static_cast<float>(y) - cy);
            const float t = (r - 330.0f) / 40.0f;
            float value = 130.0f * std::exp(-r / 260.0f) + 50.0f * std::exp(-0.5f * t * t) + 6.0f;
            const float phi = std::atan2(static_cast<float>(y) - cy, static_cast<float>(x) - cx);
            if (phi > 0.0f && phi < 1.0f)
                value *= 0.8f;
            mean[static_cast<size_t>(y * w + x)] = value + gauss(rng) * std::sqrt(value / 60.0f);
        }
    }
    return mean;
}

void CompareShortlist(const std::vector<BeamCenterFFTCandidate> &gpu,
                      const std::vector<BeamCenterFFTCandidate> &cpu, const char *what) {
    INFO(what);
    REQUIRE(gpu.size() == cpu.size());
    for (size_t i = 0; i < cpu.size(); i++) {
        INFO("candidate " << i);
        // The positions are what a caller consumes, and they are a grid index: they must be the
        // same index, not a nearby one.
        CHECK(((gpu[i].beam_x_pxl == cpu[i].beam_x_pxl)
               || (std::isnan(gpu[i].beam_x_pxl) && std::isnan(cpu[i].beam_x_pxl))));
        CHECK(((gpu[i].beam_y_pxl == cpu[i].beam_y_pxl)
               || (std::isnan(gpu[i].beam_y_pxl) && std::isnan(cpu[i].beam_y_pxl))));
        CHECK(gpu[i].score == Catch::Approx(cpu[i].score).epsilon(1e-5));
    }
}

} // namespace

// The GPU engine exists to make the capture affordable, not to answer differently: the two engines
// share the preparation, the masked Pearson, the shortlist and the margins, and differ only in how
// the four convolutions are computed. So the shortlists must be the same grid positions and the
// margins must agree - a capture that depended on which engine ran it would be a capture that could
// not be reasoned about.
TEST_CASE("BeamCenterFFTGPU_MatchesTheCPUEngine", "[BeamCenter][BeamCenterFFTGPU]") {
    if (get_gpu_count() == 0) {
        WARN("No CUDA GPU present. Skipping BeamCenterFFTGPU_MatchesTheCPUEngine");
        return;
    }

    const int64_t w = 1100, h = 900;
    const float cx = 397.25f, cy = 511.5f;
    const auto mean = SynthesiseBackground(w, h, cx, cy);

    BeamCenterFFTCPU cpu_engine;
    BeamCenterFFTGPU gpu_engine;
    const BeamCenterFFTSettings settings;
    const auto cpu = BeamCenterFFTScore(w, h, mean, settings, cpu_engine);
    const auto gpu = BeamCenterFFTScore(w, h, mean, settings, gpu_engine);

    REQUIRE(!cpu.point.empty());
    CompareShortlist(gpu.point, cpu.point, "point surface");
    CompareShortlist(gpu.line_x, cpu.line_x, "x line mirror");
    CompareShortlist(gpu.line_y, cpu.line_y, "y line mirror");

    // The margin is the surface's self-diagnostic, and it is a difference of two scores that sit
    // within a fraction of a percent of each other - the quantity most exposed to a change of
    // transform.
    CHECK(gpu.margin_point == Catch::Approx(cpu.margin_point).margin(1e-5));
    CHECK(gpu.margin_line_x == Catch::Approx(cpu.margin_line_x).margin(1e-5));
    CHECK(gpu.margin_line_y == Catch::Approx(cpu.margin_line_y).margin(1e-5));
}

// The same identity test the CPU engine carries, on the device: an exactly centrosymmetric image
// scores exactly 1 at its centre and the capture lands on that grid point. It is the cheapest
// check there is for an indexing or sign error that only the device path has.
TEST_CASE("BeamCenterFFTGPU_RecoversAnExactlySymmetricCentre", "[BeamCenter][BeamCenterFFTGPU]") {
    if (get_gpu_count() == 0) {
        WARN("No CUDA GPU present. Skipping BeamCenterFFTGPU_RecoversAnExactlySymmetricCentre");
        return;
    }

    // Keyed on the canonical member of the +- pair, so the image is exactly its own mirror about
    // (cx, cy) - the same construction the CPU identity test uses.
    const int64_t w = 220, h = 180;
    const float cx = 83.5f, cy = 61.5f;
    std::vector<float> mean(static_cast<size_t>(w * h));
    for (int64_t y = 0; y < h; y++)
        for (int64_t x = 0; x < w; x++) {
            int64_t dx = std::lround(2.0f * (static_cast<float>(x) - cx));
            int64_t dy = std::lround(2.0f * (static_cast<float>(y) - cy));
            if (dy < 0 || (dy == 0 && dx < 0)) {
                dx = -dx;
                dy = -dy;
            }
            uint64_t k = static_cast<uint64_t>(dx) * 0x9e3779b97f4a7c15ULL
                         + static_cast<uint64_t>(dy) * 0xc2b2ae3d27d4eb4fULL;
            k ^= k >> 31;
            k *= 0xbf58476d1ce4e5b9ULL;
            k ^= k >> 29;
            mean[static_cast<size_t>(y * w + x)] = 100.0f + static_cast<float>(k % 4096u) * 0.25f;
        }

    BeamCenterFFTGPU engine;
    const auto r = BeamCenterFFTScore(w, h, mean, BeamCenterFFTSettings{}, engine);
    REQUIRE(!r.point.empty());
    CHECK(r.point[0].beam_x_pxl == cx);
    CHECK(r.point[0].beam_y_pxl == cy);
    CHECK(r.point[0].score == Catch::Approx(1.0f).margin(1e-4));
}

// The card is shared, so the engine has to be able to say that an image will not fit before it
// starts allocating - that is what sends the capture to the CPU instead of taking the run down.
TEST_CASE("BeamCenterFFTGPU_SizesItsOwnFootprint", "[BeamCenter][BeamCenterFFTGPU]") {
    if (get_gpu_count() == 0) {
        WARN("No CUDA GPU present. Skipping BeamCenterFFTGPU_SizesItsOwnFootprint");
        return;
    }

    // A 16 Mpixel detector is the largest this runs on, and it has to be a bounded cost.
    CHECK(BeamCenterFFTGPU::DeviceMemoryNeeded(4148, 4362) < (3ull << 30));
    // Bigger images need more.
    CHECK(BeamCenterFFTGPU::DeviceMemoryNeeded(4148, 4362)
          > BeamCenterFFTGPU::DeviceMemoryNeeded(1030, 1064));
    // Nothing on this card can hold a hundred-gigapixel image.
    CHECK(!BeamCenterFFTGPU::FitsInDeviceMemory(300000, 300000));
}

#endif