Jungfraujoch/tests/BraggIntegrationEngineGPUTest.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 <chrono>
#include <cmath>
#include <vector>

#include "../common/BraggIntegrationSettings.h"
#include "../common/DetectorSetup.h"
#include "../common/DiffractionExperiment.h"
#include "../common/Reflection.h"
#include "../image_analysis/bragg_integration/BraggIntegrationEngineCPU.h"
#include "../image_analysis/bragg_integration/BraggIntegrationEngineGPU.h"
#include "../image_analysis/image_preprocessing/ImagePreprocessorBufferGPU.h"

namespace {

// A grid of clean Gaussian spots on a flat background, each seeding one predicted reflection.
struct Scene {
    std::vector<int32_t> image;
    std::vector<Reflection> predicted;
    size_t width = 0, height = 0;
};

Reflection MakeReflection(float x, float y, float d, int hkl) {
    Reflection r{};
    r.h = hkl; r.k = hkl; r.l = hkl;
    r.predicted_x = x;
    r.predicted_y = y;
    r.d = d;
    r.prescaling_corr = 1.0f;
    r.partiality = 1.0f;
    return r;
}

// companion_dx > 0 puts a second spot that many pixels beside every grid spot, so their r1 signal
// disks share pixels while the background rings still see clean sky - which is what a dense pattern
// actually looks like (crowded along one reciprocal axis, sparse across it).
// clip_spots punches unreadable pixels into the spots themselves rather than into empty sky: the
// centre of every 5th, a mid-profile pixel of every 7th and a disk-edge pixel of every 11th. That is
// the MINPK rescue's own case - a reflection kept and fitted over the pixels it has - and with it the
// peak-loss rule, which has to fire on the same reflections in both engines.
Scene BuildScene(size_t width, size_t height, int spacing = 60, float companion_dx = 0.0f,
                 bool clip_spots = false) {
    Scene s;
    s.width = width;
    s.height = height;
    s.image.assign(width * height, 12);   // flat background

    // A grid of spots, well separated so background rings do not overlap the neighbours' disks.
    // A spread of intensities (some weak, some very strong) and a spread of d (so several resolution
    // shells are populated) exercises the strong-spot selection, shell learning and the fit.
    const int margin = 45;
    int hkl = 1;
    for (int gy = 0; margin + gy * spacing < static_cast<int>(height) - margin; ++gy) {
        for (int gx = 0; margin + gx * spacing < static_cast<int>(width) - margin; ++gx) {
            const float cx = static_cast<float>(margin + gx * spacing) + 0.3f;   // sub-pixel offset
            const float cy = static_cast<float>(margin + gy * spacing) - 0.2f;
            const double amp = 150.0 + 60.0 * ((gx * 7 + gy * 13) % 30);   // 150..1890
            const double sigma = 1.3;
            for (int dy = -6; dy <= 6; ++dy)
                for (int dx = -6; dx <= 6; ++dx) {
                    const int x = static_cast<int>(std::lround(cx)) + dx;
                    const int y = static_cast<int>(std::lround(cy)) + dy;
                    if (x < 0 || y < 0 || x >= static_cast<int>(width) || y >= static_cast<int>(height)) continue;
                    const double ex = x - cx, ey = y - cy;
                    const double g = amp * std::exp(-(ex * ex + ey * ey) / (2.0 * sigma * sigma));
                    s.image[y * width + x] += static_cast<int32_t>(std::lround(g));
                }
            const float d = 1.4f + 0.12f * static_cast<float>((gx + gy) % 12);   // 1.4..2.72 A
            s.predicted.push_back(MakeReflection(cx, cy, d, hkl++));
            if (companion_dx > 0.0f) {
                const float ccx = cx + companion_dx;
                for (int dy = -6; dy <= 6; ++dy)
                    for (int dx = -6; dx <= 6; ++dx) {
                        const int x = static_cast<int>(std::lround(ccx)) + dx;
                        const int y = static_cast<int>(std::lround(cy)) + dy;
                        if (x < 0 || y < 0 || x >= static_cast<int>(width) || y >= static_cast<int>(height)) continue;
                        const double ex = x - ccx, ey = y - cy;
                        const double g = 0.6 * amp * std::exp(-(ex * ex + ey * ey) / (2.0 * sigma * sigma));
                        s.image[y * width + x] += static_cast<int32_t>(std::lround(g));
                    }
                s.predicted.push_back(MakeReflection(ccx, cy, d, hkl++));
            }
        }
    }

    // A few masked (INT32_MIN) and saturated (INT32_MAX) pixels in background gaps to exercise the
    // validity rejection in both engines identically.
    for (int k = 0; k < 20; ++k) {
        const size_t idx = (static_cast<size_t>(k) * 2654435761u) % s.image.size();
        s.image[idx] = (k % 2) ? INT32_MIN : INT32_MAX;
    }

    if (clip_spots)
        for (size_t n = 0; n < s.predicted.size(); ++n) {
            int dx = 0, dy = 0;
            if (n % 5 == 0)       { dx = 0; dy = 0; }      // the peak itself: the rule must reject
            else if (n % 7 == 0)  { dx = 1; dy = 1; }      // ~1.1 sigma out: near the rule's boundary
            else if (n % 11 == 0) { dx = 3; dy = -2; }     // disk edge: MINPK keeps it, the rule does not fire
            else continue;
            const int x = static_cast<int>(std::lround(s.predicted[n].predicted_x)) + dx;
            const int y = static_cast<int>(std::lround(s.predicted[n].predicted_y)) + dy;
            if (x < 0 || y < 0 || x >= static_cast<int>(width) || y >= static_cast<int>(height)) continue;
            s.image[y * width + x] = (n % 2) ? INT32_MAX : INT32_MIN;
        }
    return s;
}

// clip_nsigma 0 selects the OTHER background-ring estimator, the symmetric trim, so the two branches
// the CPU and GPU each implement separately are both covered.
DiffractionExperiment MakeExperiment(IntegratorMode mode, std::optional<float> bandwidth_fwhm,
                                     float clip_nsigma = 4.0f,
                                     bool radial = false,
                                     const DetectorSetup &det = DetJF(2),
                                     float stencil_k = 0.0f,
                                     float r1 = 0.0f, float r2 = 0.0f, float r3 = 0.0f,
                                     OverlapMode overlap = OverlapMode::Off) {
    DiffractionExperiment experiment(det);   // DetJF(2) (small) keeps the correctness test fast
    experiment.DetectorDistance_mm(100.0f).IncidentEnergy_keV(WVL_1A_IN_KEV)
              .BeamX_pxl(400.0f).BeamY_pxl(400.0f);
    experiment.BandwidthFWHM(bandwidth_fwhm);
    BraggIntegrationSettings settings;
    settings.Integrator(mode);
    if (r1 > 0.0f)
        settings.R1(r1).R2(r2).R3(r3);
    if (clip_nsigma > 0.0f)
        settings.BackgroundClipNSigma(clip_nsigma);
    else
        settings.BackgroundTrimFraction(0.10f);
    settings.BackgroundRadialCorrection(radial);
    settings.StencilKSigma(stencil_k);
    settings.Overlap(overlap);
    experiment.ImportBraggIntegrationSettings(settings);
    return experiment;
}

void CompareCpuVsGpu(IntegratorMode mode, std::optional<float> bandwidth_fwhm,
                     float clip_nsigma = 4.0f, bool radial = false, int spacing = 60,
                     float stencil_k = 0.0f,
                     float r1 = 0.0f, float r2 = 0.0f, float r3 = 0.0f,
                     OverlapMode overlap = OverlapMode::Off, float companion_dx = 0.0f,
                     bool clip_spots = false) {
    const DiffractionExperiment experiment =
            MakeExperiment(mode, bandwidth_fwhm, clip_nsigma, radial, DetJF(2), stencil_k, r1, r2, r3,
                           overlap);
    const size_t width = experiment.GetXPixelsNum();
    const size_t height = experiment.GetYPixelsNum();
    const size_t npixel = experiment.GetPixelsNum();
    REQUIRE(npixel == width * height);

    const Scene scene = BuildScene(width, height, spacing, companion_dx, clip_spots);
    REQUIRE(scene.image.size() == npixel);
    REQUIRE(scene.predicted.size() > 60);

    // CPU reference
    ImagePreprocessorBuffer cpu_image(npixel);
    for (size_t i = 0; i < npixel; ++i)
        cpu_image[i] = scene.image[i];
    BraggIntegrationEngineCPU cpu(experiment);
    const auto out_cpu = cpu.Run(cpu_image, scene.predicted, scene.predicted.size(), 5);

    // GPU under test, identical input uploaded to the device
    auto stream = std::make_shared<CudaStream>();
    ImagePreprocessorBufferGPU gpu_image(npixel);
    for (size_t i = 0; i < npixel; ++i)
        gpu_image[i] = scene.image[i];
    REQUIRE(cudaMemcpyAsync(gpu_image.getGPUBuffer(), gpu_image.getBuffer().data(),
                            npixel * sizeof(int32_t), cudaMemcpyHostToDevice, *stream) == cudaSuccess);
    BraggIntegrationEngineGPU gpu(experiment, stream);
    const auto out_gpu = gpu.Run(gpu_image, scene.predicted, scene.predicted.size(), 5);

    // The ok/observed decisions are deterministic geometry, so both engines return the same set in
    // the same (predicted-index) order. Intensities differ only by float rounding and the unordered
    // atomic summation of the learned profile, so compare up to a small tolerance.
    REQUIRE(out_gpu.size() == out_cpu.size());
    REQUIRE(out_cpu.size() > 40);
    if (clip_spots) {
        // Guard against the coverage going vacuous: the punched pixels have to actually cost some
        // reflections, or the two engines are being compared on a case neither of them meets.
        const Scene clean_scene = BuildScene(width, height, spacing, companion_dx, false);
        ImagePreprocessorBuffer clean_image(npixel);
        for (size_t i = 0; i < npixel; ++i)
            clean_image[i] = clean_scene.image[i];
        BraggIntegrationEngineCPU clean_cpu(experiment);
        const auto out_clean = clean_cpu.Run(clean_image, clean_scene.predicted,
                                             clean_scene.predicted.size(), 5);
        CHECK(out_cpu.size() < out_clean.size());
    }
    for (size_t i = 0; i < out_cpu.size(); ++i) {
        INFO("mode " << static_cast<int>(mode) << " reflection " << i << " hkl " << out_cpu[i].h);
        CHECK(out_gpu[i].h == out_cpu[i].h);
        CHECK(out_gpu[i].image_number == out_cpu[i].image_number);
        CHECK(out_gpu[i].bkg == Catch::Approx(out_cpu[i].bkg).epsilon(0.02).margin(0.5));
        CHECK(out_gpu[i].I == Catch::Approx(out_cpu[i].I).epsilon(0.03).margin(2.0));
        CHECK(out_gpu[i].sigma == Catch::Approx(out_cpu[i].sigma).epsilon(0.03).margin(0.5));
    }
}

} // namespace

TEST_CASE("BraggIntegrationEngineGPU_MatchesCPU") {
    if (get_gpu_count() == 0) {
        WARN("No CUDA GPU present. Skipping BraggIntegrationEngineGPU_MatchesCPU");
        return;
    }

    SECTION("BoxSum")                    { CompareCpuVsGpu(IntegratorMode::BoxSum, std::nullopt); }
    SECTION("ProfileGaussian mono")      { CompareCpuVsGpu(IntegratorMode::ProfileGaussian, std::nullopt); }
    SECTION("ProfileGaussian broadband") { CompareCpuVsGpu(IntegratorMode::ProfileGaussian, 0.03f); }
    // An elongated background ring: the classification, the bounding box, the neighbour mask and the
    // shared-memory radial window all become reflection-dependent, and the two engines have to agree
    // on every one of them. Spots spaced wider so the grown rings stay clear of the neighbours -
    // what is under test is the stencil, not the crowding.
    SECTION("ProfileGaussian stencil broadband") {
        CompareCpuVsGpu(IntegratorMode::ProfileGaussian, 0.005f, 4.0f, false, 120, 3.0f);
    }
    // A monochromatic beam has no streak, so k_sigma changes nothing - the point of the section is
    // that both engines agree that it changes nothing.
    SECTION("ProfileGaussian stencil mono") {
        CompareCpuVsGpu(IntegratorMode::ProfileGaussian, std::nullopt, 4.0f, false, 120, 3.0f);
    }
    // Crowded: at the default spacing the grown rings DO overlap their neighbours, so the elongated
    // neighbour mask, the shrinking background-pixel count and the n_bkg acceptance gate are all in
    // play. That is the case the feature meets at high resolution, and the wide-spacing sections
    // above deliberately avoid it.
    SECTION("ProfileGaussian stencil crowded") {
        CompareCpuVsGpu(IntegratorMode::ProfileGaussian, 0.02f, 4.0f, false, 60, 4.0f);
    }
    SECTION("BoxSum stencil") {
        CompareCpuVsGpu(IntegratorMode::BoxSum, 0.005f, 4.0f, false, 120, 3.0f);
    }
    // The trimmed-mean ring is sorted in a fixed-size shared buffer on the GPU; an elongated ring
    // holds more pixels, so both engines have to fall back to the plain mean at the same place.
    SECTION("ProfileGaussian stencil trim") {
        CompareCpuVsGpu(IntegratorMode::ProfileGaussian, 0.005f, 0.0f, false, 120, 3.0f);
    }
    // A ring wide enough to overflow the GPU's fixed trimmed-mean buffer, so the fallback to the
    // plain ring mean is exercised - and has to happen in both engines at the same reflection. The
    // growth cap keeps the default 6/10 ring under the buffer at any bandwidth, so this needs the
    // wider stills radii to be reachable at all.
    SECTION("ProfileGaussian stencil trim overflow") {
        CompareCpuVsGpu(IntegratorMode::ProfileGaussian, 0.04f, 0.0f, false, 120, 4.0f, 6.0f, 8.0f, 12.0f);
    }
    SECTION("ProfileEmpirical")          { CompareCpuVsGpu(IntegratorMode::ProfileEmpirical, std::nullopt); }
    // Overlap treatment: companions 4 px apart put each reflection's centre inside its neighbour's
    // signal disk, so the owner map, the excluded pixels and the profile fraction the two modes act on
    // all have to come out the same in both engines - the ownership atomic in particular is settled by
    // an atomicMin on the GPU and a serial minimum on the CPU.
    SECTION("ProfileGaussian overlap exclude") {
        CompareCpuVsGpu(IntegratorMode::ProfileGaussian, std::nullopt, 4.0f, false, 60, 0.0f,
                        0.0f, 0.0f, 0.0f, OverlapMode::Exclude, 4.0f);
    }
    SECTION("ProfileGaussian overlap reject") {
        CompareCpuVsGpu(IntegratorMode::ProfileGaussian, std::nullopt, 4.0f, false, 60, 0.0f,
                        0.0f, 0.0f, 0.0f, OverlapMode::Reject, 4.0f);
    }
    SECTION("BoxSum overlap reject") {
        CompareCpuVsGpu(IntegratorMode::BoxSum, std::nullopt, 4.0f, false, 60, 0.0f,
                        0.0f, 0.0f, 0.0f, OverlapMode::Reject, 4.0f);
    }
    // Nothing shares a pixel at this spacing, so an overlap treatment has to leave the result alone.
    SECTION("ProfileGaussian overlap inert") {
        CompareCpuVsGpu(IntegratorMode::ProfileGaussian, std::nullopt, 4.0f, false, 60, 0.0f,
                        0.0f, 0.0f, 0.0f, OverlapMode::Exclude);
    }
    SECTION("ProfileGaussian mono trim") { CompareCpuVsGpu(IntegratorMode::ProfileGaussian, std::nullopt, 0.0f); }
    // Unreadable pixels inside the signal disks themselves: the MINPK rescue keeps the reflection and
    // fits it over what is left, and the peak-loss rule throws back the ones that lost the profile's
    // maximum. Both decisions are per-reflection cuts on a reduction over the profile grid, computed
    // independently in the two engines (serial max vs an atomicMax on the float bit pattern), so they
    // have to reject exactly the same reflections - a mismatch shows up as a size mismatch here.
    SECTION("ProfileGaussian clipped disks") {
        CompareCpuVsGpu(IntegratorMode::ProfileGaussian, std::nullopt, 4.0f, false, 60, 0.0f,
                        0.0f, 0.0f, 0.0f, OverlapMode::Off, 0.0f, true);
    }
    SECTION("ProfileEmpirical clipped disks") {
        CompareCpuVsGpu(IntegratorMode::ProfileEmpirical, std::nullopt, 4.0f, false, 60, 0.0f,
                        0.0f, 0.0f, 0.0f, OverlapMode::Off, 0.0f, true);
    }
    // The same, with an elongated profile: the peak is then a ridge, so the fraction-of-peak test has
    // to protect a crest rather than one pixel, and the grid it reduces over is reflection-dependent.
    SECTION("ProfileGaussian clipped disks stencil") {
        CompareCpuVsGpu(IntegratorMode::ProfileGaussian, 0.005f, 4.0f, false, 120, 3.0f,
                        0.0f, 0.0f, 0.0f, OverlapMode::Off, 0.0f, true);
    }
    SECTION("BoxSum clipped disks") {
        CompareCpuVsGpu(IntegratorMode::BoxSum, std::nullopt, 4.0f, false, 60, 0.0f,
                        0.0f, 0.0f, 0.0f, OverlapMode::Off, 0.0f, true);
    }
    // The radial background curvature correction is computed independently in the two engines
    // (host loop vs radial_correct kernel), so it needs its own parity coverage.
    SECTION("BoxSum radial")             { CompareCpuVsGpu(IntegratorMode::BoxSum, std::nullopt, 4.0f, true); }
    SECTION("ProfileGaussian radial")    { CompareCpuVsGpu(IntegratorMode::ProfileGaussian, std::nullopt, 4.0f, true); }
    // With an elongated ring the radial-curvature kernel is a table indexed per reflection, and the
    // shared window boxsum accumulates the curve in is sized from the widest aperture on the
    // detector. Both are computed independently in the two engines.
    SECTION("ProfileGaussian radial stencil") {
        CompareCpuVsGpu(IntegratorMode::ProfileGaussian, 0.005f, 4.0f, true, 120, 3.0f);
    }
}

// The mask and the owner map are cleared by the run that marked them rather than at the start of the
// next one, so a reused engine has to give the same answer as a fresh one. A first frame whose spots
// are somewhere else entirely is what would show a leftover mark: a stale mask pixel is read as a
// neighbour's signal and dropped from the background ring, a stale owner steals a pixel outright.
TEST_CASE("BraggIntegrationEngineGPU_ReusedEngineMatchesFresh") {
    if (get_gpu_count() == 0) {
        WARN("No CUDA GPU present. Skipping BraggIntegrationEngineGPU_ReusedEngineMatchesFresh");
        return;
    }

    for (OverlapMode ovl : {OverlapMode::Off, OverlapMode::Exclude}) {
        const DiffractionExperiment experiment =
                MakeExperiment(IntegratorMode::ProfileGaussian, std::nullopt, 4.0f, false, DetJF(2),
                               0.0f, 0.0f, 0.0f, 0.0f, ovl);
        const size_t width = experiment.GetXPixelsNum();
        const size_t height = experiment.GetYPixelsNum();
        const size_t npixel = experiment.GetPixelsNum();

        // Two frames whose spot grids do not line up, so the first frame's marks fall on the second
        // frame's background rings rather than back onto its own disks.
        const Scene first = BuildScene(width, height, 47);
        const Scene second = BuildScene(width, height, 60);
        REQUIRE(first.predicted.size() > 60);
        REQUIRE(second.predicted.size() > 60);

        auto integrate = [&](BraggIntegrationEngineGPU &engine, const Scene &scene,
                             const std::shared_ptr<CudaStream> &stream) {
            ImagePreprocessorBufferGPU img(npixel);
            for (size_t i = 0; i < npixel; ++i) img[i] = scene.image[i];
            REQUIRE(cudaMemcpyAsync(img.getGPUBuffer(), img.getBuffer().data(),
                                    npixel * sizeof(int32_t), cudaMemcpyHostToDevice, *stream) == cudaSuccess);
            return engine.Run(img, scene.predicted, scene.predicted.size(), 7);
        };

        auto stream_fresh = std::make_shared<CudaStream>();
        BraggIntegrationEngineGPU fresh(experiment, stream_fresh);
        const auto out_fresh = integrate(fresh, second, stream_fresh);

        auto stream_reused = std::make_shared<CudaStream>();
        BraggIntegrationEngineGPU reused(experiment, stream_reused);
        integrate(reused, first, stream_reused);
        const auto out_reused = integrate(reused, second, stream_reused);

        INFO("overlap mode " << static_cast<int>(ovl));
        REQUIRE(out_reused.size() == out_fresh.size());
        for (size_t i = 0; i < out_fresh.size(); ++i) {
            INFO("reflection " << i);
            CHECK(out_reused[i].h == out_fresh[i].h);
            CHECK(out_reused[i].I == out_fresh[i].I);
            CHECK(out_reused[i].sigma == out_fresh[i].sigma);
            CHECK(out_reused[i].bkg == out_fresh[i].bkg);
        }
    }
}

// Hidden ([.]) benchmark: the raison d'etre of the GPU port is < 2 ms/frame (vs ~142 ms on the CPU
// for ProfileIntegrate2D). Run explicitly with:  ./jfjoch_test "[bragg_bench]"
TEST_CASE("BraggIntegrationEngineGPU_Benchmark", "[.][bragg_bench]") {
    if (get_gpu_count() == 0) {
        WARN("No CUDA GPU present. Skipping benchmark");
        return;
    }
    // The overlap treatment is priced here too: it adds an owner map over the whole frame plus one
    // atomic per claimed pixel, so what it costs is a property of the frame more than of the crowding.
    for (OverlapMode ovl : {OverlapMode::Off, OverlapMode::Reject, OverlapMode::Exclude}) {
    const DiffractionExperiment experiment = MakeExperiment(IntegratorMode::ProfileGaussian, std::nullopt,
                                                            4.0f, false, DetJF4M(), 0.0f, 0.0f, 0.0f, 0.0f,
                                                            ovl);
    const size_t width = experiment.GetXPixelsNum();
    const size_t height = experiment.GetYPixelsNum();
    const size_t npixel = experiment.GetPixelsNum();
    REQUIRE(npixel == width * height);

    auto stream = std::make_shared<CudaStream>();
    BraggIntegrationEngineGPU gpu(experiment, stream);
    for (int spacing : {28, 60}) {
        const Scene scene = BuildScene(width, height, spacing);
        const size_t nrefl = scene.predicted.size();

        ImagePreprocessorBufferGPU gpu_image(npixel);
        for (size_t i = 0; i < npixel; ++i) gpu_image[i] = scene.image[i];
        REQUIRE(cudaMemcpyAsync(gpu_image.getGPUBuffer(), gpu_image.getBuffer().data(),
                                npixel * sizeof(int32_t), cudaMemcpyHostToDevice, *stream) == cudaSuccess);
        cudaStreamSynchronize(*stream);

        auto run = [&] { return gpu.Run(gpu_image, scene.predicted, nrefl, 0); };
        for (int i = 0; i < 5; ++i) run();   // warm-up (allocations, JIT)

        const int iters = 100;
        const auto t0 = std::chrono::steady_clock::now();
        size_t observed = 0;
        for (int i = 0; i < iters; ++i) observed += run().size();
        const auto t1 = std::chrono::steady_clock::now();
        const double ms = std::chrono::duration<double, std::milli>(t1 - t0).count() / iters;

        BraggIntegrationEngineCPU cpu(experiment);
        ImagePreprocessorBuffer cpu_image(npixel);
        for (size_t i = 0; i < npixel; ++i) cpu_image[i] = scene.image[i];
        const auto c0 = std::chrono::steady_clock::now();
        const size_t cpu_observed = cpu.Run(cpu_image, scene.predicted, nrefl, 0).size();
        const auto c1 = std::chrono::steady_clock::now();
        const double cpu_ms = std::chrono::duration<double, std::milli>(c1 - c0).count();

        WARN((int) ovl << " | " << width << "x" << height << " | " << nrefl << " refl ("
             << observed / iters << " obs) | GPU " << ms << " ms | CPU " << cpu_ms << " ms ("
             << cpu_observed << " obs) | speedup " << cpu_ms / ms << "x");
    }
    }
}

#endif