Jungfraujoch/tests/BraggIntegrationEngineMetalTest.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 <bit>
#include <cmath>
#include <cstdint>
#include <random>
#include <vector>

#include "../common/BraggIntegrationSettings.h"
#include "../common/DetectorSetup.h"
#include "../common/DiffractionExperiment.h"
#include "../common/PixelMask.h"
#include "../common/Reflection.h"
#include "../image_analysis/bragg_integration/BraggIntegrationEngineCPU.h"
#ifdef __APPLE__
#include "../image_analysis/bragg_integration/BraggIntegrationEngineMetal.h"
#endif

// The Metal engine's whole claim is that it IS the CPU engine, bit for bit (BraggIntegrationEngineMetal.h).
// So nothing here has a tolerance: every float of every reflection is compared as a bit pattern, and
// every count as an integer. A Mac without a Metal device fails rather than skips - the macOS CI job is
// the only place this can be checked, and a skip there would check nothing. The frames and the cases
// are built everywhere, and one test case checks on the CPU engine alone that they reach what they are
// there for.

namespace {

struct Frame {
    std::vector<int32_t> image;
    std::vector<Reflection> predicted;
};

// A crowded random frame: predictions at random sub-pixel positions (so the pixel classification meets
// every kind of rounding near the stencil boundaries), a quarter of them on whole and half pixels (so it
// meets exact ties with the radii too), with random partialities (rocking-curve tails and flux, which
// mask the rings differently), about a third of them with a spot, overlapping where they fall close. On
// a background that rises with the radius (the radial correction has a curve to fit) with noise that
// takes it below zero near the beam, as a pedestal-subtracted detector does (a negative pixel is what
// carries into the high word of the GPU's 64-bit radial sums), hot pixels in the rings (so the
// background clip rejects something), and unreadable pixels: some in the run's mask, some saturated,
// some unreadable on this frame alone (both read as an overload), inside spots as well as between them.
Frame RandomFrame(const DiffractionExperiment &x, uint32_t seed, const std::vector<size_t> &masked) {
    const int W = static_cast<int>(x.GetXPixelsNum()), H = static_cast<int>(x.GetYPixelsNum());
    const float beam_x = x.GetBeamX_pxl(), beam_y = x.GetBeamY_pxl();
    std::mt19937 rng(seed);
    std::uniform_real_distribution<float> U(0.0f, 1.0f);
    Frame f;
    f.image.resize(static_cast<size_t>(W) * H);
    for (int y = 0; y < H; ++y)
        for (int x_ = 0; x_ < W; ++x_) {
            const float r = std::hypot(x_ - beam_x, y - beam_y);
            const float mean = 0.5f + 0.01f * r;
            f.image[static_cast<size_t>(y) * W + x_] = static_cast<int32_t>(std::lround(mean + (U(rng) - 0.5f) * 8.0f));
        }
    const int nrefl = W * H / 400;
    int hkl = 1;
    for (int n = 0; n < nrefl; ++n) {
        Reflection r{};
        r.h = hkl; r.k = -hkl; r.l = hkl % 7; ++hkl;
        r.predicted_x = 2.0f + U(rng) * (W - 4);
        r.predicted_y = 2.0f + U(rng) * (H - 4);
        if (n % 4 == 1) {
            r.predicted_x = std::round(r.predicted_x * 2.0f) / 2.0f;
            r.predicted_y = std::round(r.predicted_y);
        }
        const float u = U(rng);
        r.partiality = u < 0.3f ? 0.01f * u : u;   // a third of them rocking-curve tails
        r.d = 1.2f + 3.0f * U(rng);
        r.prescaling_corr = 1.0f;
        f.predicted.push_back(r);
        if (U(rng) < 0.35f) {
            const float amp = 20.0f + 3000.0f * U(rng) * U(rng);
            const float s = 0.8f + 1.4f * U(rng);
            for (int dy = -8; dy <= 8; ++dy)
                for (int dx = -8; dx <= 8; ++dx) {
                    const int px = static_cast<int>(std::lround(r.predicted_x)) + dx;
                    const int py = static_cast<int>(std::lround(r.predicted_y)) + dy;
                    if (px < 0 || py < 0 || px >= W || py >= H) continue;
                    const float ex = px - r.predicted_x, ey = py - r.predicted_y;
                    f.image[static_cast<size_t>(py) * W + px] += static_cast<int32_t>(
                            std::lround(amp * std::exp(-(ex * ex + ey * ey) / (2.0f * s * s))));
                }
        }
    }
    std::uniform_int_distribution<size_t> P(0, f.image.size() - 1);
    for (int k = 0; k < W * H / 500; ++k)
        f.image[P(rng)] += 200 + static_cast<int32_t>(5000 * U(rng));          // hot pixels
    for (int k = 0; k < W * H / 3000; ++k)
        f.image[P(rng)] = U(rng) < 0.5f ? INT32_MIN : INT32_MAX;              // unreadable on this frame
    for (size_t i = 0; i < f.predicted.size(); i += 13) {                     // ... inside spots too
        const auto &r = f.predicted[i];
        const size_t idx = static_cast<size_t>(std::lround(r.predicted_y)) * W + std::lround(r.predicted_x) + i % 3;
        if (idx < f.image.size()) f.image[idx] = i % 2 ? INT32_MIN : INT32_MAX;
    }
    for (size_t idx : masked)
        f.image[idx] = INT32_MIN;                                             // what the run's mask holds
    return f;
}

struct Case {
    IntegratorMode mode = IntegratorMode::ProfileGaussian;
    OverlapMode overlap = OverlapMode::Exclude;
    bool radial = false;
    std::optional<float> bandwidth;
    float stencil_k = 0.0f;
    bool footprint = false;
    float clip_nsigma = 4.0f;
    bool still = false;       // a still keeps its overloads as lower bounds (saturated pixels at the limit)
    float gain = 1.0f;        // ADU per photon
    bool curve = false;       // a radial curve handed over (RadialBackgroundFromRings), which the clip follows
};

// The radial curve a case hands over: rising steeply enough with the radius that the clip it makes the
// ring follow (RadialClipOffset) moves by tens of counts across a ring.
void HandOverCurve(BraggIntegrationEngine &engine) {
    std::vector<float> rings(400);
    for (size_t i = 0; i < rings.size(); ++i)
        rings[i] = 1.0f + 0.01f * static_cast<float>(i * i);
    engine.RadialBackgroundFromRings(rings, static_cast<int32_t>(rings.size()), 0.05f, 0.01f);
}

DiffractionExperiment MakeExperiment(const DetectorSetup &det, const Case &c) {
    DiffractionExperiment x(det);
    x.DetectorDistance_mm(100.0f).IncidentEnergy_keV(WVL_1A_IN_KEV)
     .BeamX_pxl(0.4f * static_cast<float>(x.GetXPixelsNum()))
     .BeamY_pxl(0.55f * static_cast<float>(x.GetYPixelsNum()));
    // A rotation drops an overloaded observation, a still keeps it as a lower bound (Finalize).
    if (!c.still) {
        x.Goniometer(GoniometerAxis("omega", 0.0f, 0.1f, Coord(-1, 0, 0), {}));
        x.ImportIndexingSettings(IndexingSettings().RotationIndexing(true));
    }
    x.BandwidthFWHM(c.bandwidth);
    BraggIntegrationSettings s;
    s.Integrator(c.mode).Overlap(c.overlap).BackgroundRadialCorrection(c.radial).StencilKSigma(c.stencil_k)
     .BackgroundClipNSigma(c.clip_nsigma).PixelGain(c.gain);
    if (c.footprint) {
        // Wider than the r1 disk can hold from the middle out, so both stencils occur on one frame.
        SpotFootprint fp;
        fp.bin_px = 100.0f;
        for (int b = 0; b < 8; ++b) {
            fp.sigma_rad.push_back(0.9f + 0.35f * b);
            fp.sigma_tan.push_back(0.8f + 0.3f * b);
        }
        s.Footprint(fp);
    }
    x.ImportBraggIntegrationSettings(s);
    return x;
}

}   // namespace

// Runs everywhere, on the CPU engine alone: the cases the Metal comparison below adds for the stills, the
// gain and the handed-over curve do reach what they are there for, so that comparison cannot pass by
// comparing nothing.
TEST_CASE("BraggIntegrationEngineMetal_CasesReachTheirModes", "[portable][Integration]") {
    Case c;
    c.still = true;
    c.gain = 4.0f;
    c.radial = true;
    c.curve = true;
    const DiffractionExperiment x = MakeExperiment(DetJF(2), c);
    const PixelMask mask(x);
    BraggIntegrationEngineCPU cpu(x, mask);
    HandOverCurve(cpu);
    cpu.BackgroundRadial(true);
    // The clip follows the curve: it moves by whole counts across a ring.
    CHECK(cpu.RadialClipOffset(500.0f, 10.0f) > 10);
    CHECK(cpu.RadialClipOffset(500.0f, -10.0f) < -10);
    ImagePreprocessorBuffer image(x.GetPixelsNum());
    const Frame f = RandomFrame(x, 7, {});
    std::copy(f.image.begin(), f.image.end(), image.data());
    const auto out = cpu.Run(image, f.predicted, f.predicted.size(), 0);
    CHECK(out.size() > 1000);
    // Saturated pixels counted at the limit: overloads kept as lower bounds.
    CHECK(cpu.Counts().overloads_kept > 0);
}

#ifdef __APPLE__

namespace {

uint32_t Bits(float f) { return std::bit_cast<uint32_t>(f); }

void RequireSame(const std::vector<Reflection> &cpu, const std::vector<Reflection> &metal) {
    REQUIRE(metal.size() == cpu.size());
    for (size_t i = 0; i < cpu.size(); ++i) {
        const Reflection &a = cpu[i], &b = metal[i];
        INFO("reflection " << i << " hkl " << a.h << " " << a.k << " " << a.l
             << " I " << a.I << " / " << b.I << " bkg " << a.bkg << " / " << b.bkg);
        REQUIRE(a.h == b.h);
        REQUIRE(a.k == b.k);
        REQUIRE(a.l == b.l);
        REQUIRE(Bits(a.image_number) == Bits(b.image_number));
        REQUIRE(Bits(a.I) == Bits(b.I));
        REQUIRE(Bits(a.sigma) == Bits(b.sigma));
        REQUIRE(Bits(a.bkg) == Bits(b.bkg));
        REQUIRE(Bits(a.var_bkg) == Bits(b.var_bkg));
        REQUIRE(Bits(a.observed_x) == Bits(b.observed_x));
        REQUIRE(Bits(a.observed_y) == Bits(b.observed_y));
        REQUIRE(Bits(a.image_scale_corr) == Bits(b.image_scale_corr));
        REQUIRE(Bits(a.prescaling_corr) == Bits(b.prescaling_corr));
        REQUIRE(a.observed == b.observed);
        REQUIRE(a.overloaded == b.overloaded);
        REQUIRE(a.clipped == b.clipped);
    }
}

void RequireSame(const BraggIntegrationCounts &a, const BraggIntegrationCounts &b) {
    CHECK(a.predicted == b.predicted);
    CHECK(a.bkg_starved == b.bkg_starved);
    CHECK(a.bkg_starved_by_neighbour == b.bkg_starved_by_neighbour);
    CHECK(a.profile_fallback == b.profile_fallback);
    CHECK(a.overlapped == b.overlapped);
    CHECK(a.overlapped_major == b.overlapped_major);
    CHECK(a.overloads_kept == b.overloads_kept);
}

// Three frames through one pair of engines - so whatever the Metal engine keeps between frames (its
// masks, its buffers) has to be left as the CPU engine leaves its own - each compared exactly.
// Returns how many reflections came out, so a case can show it was not vacuous.
size_t CompareFrames(const DetectorSetup &det, const Case &c) {
    const DiffractionExperiment x = MakeExperiment(det, c);
    REQUIRE(BraggIntegrationEngineMetal::Available(x));
    const size_t W = x.GetXPixelsNum();

    // A few pixels of the run's mask, under spots and between them: unreadable there, but not overloads.
    PixelMask mask(x);
    std::vector<uint32_t> user(x.GetPixelsNum(), 0);
    std::vector<size_t> masked;
    for (size_t k = 0; k < x.GetPixelsNum(); k += 7919) {
        user[k] = 1;
        masked.push_back(k);
    }
    mask.LoadUserMask(x, user);

    BraggIntegrationEngineCPU cpu(x, mask);
    BraggIntegrationEngineMetal metal(x, mask);
    if (c.curve) {
        HandOverCurve(cpu);
        HandOverCurve(metal);
    }
    ImagePreprocessorBuffer image(x.GetPixelsNum());

    size_t n_out = 0;
    for (uint32_t frame = 0; frame < 3; ++frame) {
        INFO("frame " << frame);
        const Frame f = RandomFrame(x, 1000 * frame + static_cast<uint32_t>(W), masked);
        std::copy(f.image.begin(), f.image.end(), image.data());
        // The auto radial mode switches per image; here it is on for the radial cases, every image.
        cpu.BackgroundRadial(c.radial);
        metal.BackgroundRadial(c.radial);
        const auto out_cpu = cpu.Run(image, f.predicted, f.predicted.size(), frame);
        const auto out_metal = metal.Run(image, f.predicted, f.predicted.size(), frame);
        RequireSame(out_cpu, out_metal);
        n_out += out_cpu.size();
    }
    RequireSame(cpu.Counts(), metal.Counts());
    // Overloads kept as lower bounds count saturated pixels: the case is not vacuous there either.
    if (c.still)
        CHECK(cpu.Counts().overloads_kept > 0);
    // Not the CPU engine compared with itself: every frame went through the kernels.
    REQUIRE(metal.FramesOnGPU() == 3);
    return n_out;
}

}   // namespace

TEST_CASE("BraggIntegrationEngineMetal_BitIdenticalToCPU", "[portable][Integration][Metal]") {
    // A JUNGFRAU 1M, and a small odd-sized detector, where most boxes meet an edge.
    const DetectorSetup large = DetJF(2);
    const DetectorSetup small = DetDECTRIS(211, 197, "small", {});

    std::vector<std::pair<std::string, Case>> cases;
    for (IntegratorMode mode : {IntegratorMode::BoxSum, IntegratorMode::ProfileGaussian, IntegratorMode::ProfileEmpirical})
        for (OverlapMode overlap : {OverlapMode::Off, OverlapMode::Reject, OverlapMode::Exclude})
            for (bool radial : {false, true}) {
                Case c;
                c.mode = mode;
                c.overlap = overlap;
                c.radial = radial;
                cases.emplace_back("mode " + std::to_string(static_cast<int>(mode)) + " overlap "
                                   + std::to_string(static_cast<int>(overlap)) + (radial ? " radial" : ""), c);
            }
    // Rings elongated by the bandwidth, and the measured footprint: per-reflection stencils.
    Case broadband;
    broadband.bandwidth = 0.01f;
    broadband.stencil_k = 3.0f;
    broadband.radial = true;
    cases.emplace_back("broadband stencil radial", broadband);
    Case footprint;
    footprint.footprint = true;
    footprint.radial = true;
    cases.emplace_back("footprint radial", footprint);
    Case footprint_box = footprint;
    footprint_box.mode = IntegratorMode::BoxSum;
    cases.emplace_back("footprint box sum", footprint_box);
    // No clip: the plain ring mean, and no clip kernel at all.
    Case no_clip;
    no_clip.clip_nsigma = 0.0f;
    cases.emplace_back("no clip", no_clip);
    // A radial curve handed over, which the ring's clip follows (RadialClipOffset), under every overlap
    // mode and every integrator.
    for (IntegratorMode mode : {IntegratorMode::BoxSum, IntegratorMode::ProfileGaussian, IntegratorMode::ProfileEmpirical})
        for (OverlapMode overlap : {OverlapMode::Off, OverlapMode::Exclude}) {
            Case curve;
            curve.mode = mode;
            curve.overlap = overlap;
            curve.radial = true;
            curve.curve = true;
            cases.emplace_back("curve mode " + std::to_string(static_cast<int>(mode)) + " overlap "
                               + std::to_string(static_cast<int>(overlap)), curve);
        }
    // Stills: overloads kept as lower bounds, every saturated pixel counted. At a gain of one, and in
    // detector units at two gains (the clip, the floors and the output's units all move with it), with
    // and without the handed-over curve.
    for (IntegratorMode mode : {IntegratorMode::BoxSum, IntegratorMode::ProfileGaussian, IntegratorMode::ProfileEmpirical})
        for (float gain : {1.0f, 4.0f, 1.7f}) {
            Case still;
            still.mode = mode;
            still.still = true;
            still.gain = gain;
            still.radial = gain != 1.0f;
            still.curve = gain == 4.0f;
            cases.emplace_back("still mode " + std::to_string(static_cast<int>(mode)) + " gain "
                               + std::to_string(gain) + (still.radial ? " radial" : "")
                               + (still.curve ? " curve" : ""), still);
        }
    // Footprint stencils on a still with a curve: the clip follows it on a per-reflection stencil.
    Case footprint_still = footprint;
    footprint_still.still = true;
    footprint_still.curve = true;
    footprint_still.gain = 4.0f;
    cases.emplace_back("footprint still gain 4 curve", footprint_still);

    for (const auto &[name, c] : cases) {
        DYNAMIC_SECTION(name) {
            CHECK(CompareFrames(large, c) > 1000);
            CHECK(CompareFrames(small, c) > 50);
        }
    }
}

// The trimmed-mean background sorts each ring, which only the CPU engine does: such a run keeps the CPU
// engine (MXAnalysisWithoutFPGA asks Available).
TEST_CASE("BraggIntegrationEngineMetal_TrimmedMeanStaysOnCPU", "[portable][Integration][Metal]") {
    Case trim;
    trim.clip_nsigma = 0.0f;
    DiffractionExperiment x = MakeExperiment(DetJF(2), trim);
    BraggIntegrationSettings s = x.GetBraggIntegrationSettings();
    s.BackgroundTrimFraction(0.1f);
    x.ImportBraggIntegrationSettings(s);
    CHECK(!BraggIntegrationEngineMetal::Available(x));
}

#endif