Jungfraujoch/tests/BraggIntegrationEngineCPUTest.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 <algorithm>
#include <cmath>
#include <cstdint>
#include <vector>

#include "../common/BraggIntegrationSettings.h"
#include "../common/DetectorSetup.h"
#include "../common/DiffractionExperiment.h"
#include "../common/JFJochMath.h"
#include "../common/Reflection.h"
#include "../image_analysis/bragg_integration/BraggIntegrationEngineCPU.h"
#include "../image_analysis/image_preprocessing/ImagePreprocessorBuffer.h"

// One spot on a flat background, inside a grid of predictions 8 px apart that put no flux on the
// frame - the tails of reflections recorded on the frames either side of a finely sliced one. Their
// r2 regions cover every background ring, so whether the spot keeps a clean ring depends only on
// whether those predictions are allowed to mask it.
TEST_CASE("BraggIntegrationEngineCPU_NeighbourMaskFollowsPartiality", "[Integration][portable]") {
    DiffractionExperiment experiment(DetJF(2));
    experiment.DetectorDistance_mm(100.0f).IncidentEnergy_keV(WVL_1A_IN_KEV).BeamX_pxl(400.0f).BeamY_pxl(400.0f);
    experiment.ImportBraggIntegrationSettings(BraggIntegrationSettings());
    const size_t width = experiment.GetXPixelsNum(), npixel = experiment.GetPixelsNum();

    const float cx = 600.3f, cy = 300.2f, amp = 800.0f, sigma = 1.3f;
    ImagePreprocessorBuffer image(npixel);
    for (size_t i = 0; i < npixel; ++i)
        image[i] = 12;
    for (int dy = -6; dy <= 6; ++dy)
        for (int dx = -6; dx <= 6; ++dx) {
            const int x = static_cast<int>(std::lround(cx)) + dx, y = static_cast<int>(std::lround(cy)) + dy;
            const float ex = x - cx, ey = y - cy;
            image[y * width + x] += static_cast<int32_t>(std::lround(amp * std::exp(-(ex * ex + ey * ey) / (2 * sigma * sigma))));
        }

    auto run = [&](float neighbour_partiality, BraggIntegrationCounts &counts) {
        std::vector<Reflection> predicted;
        for (int gy = -4; gy <= 4; ++gy)
            for (int gx = -4; gx <= 4; ++gx) {
                Reflection r{};
                r.h = gx; r.k = gy; r.l = 1;
                r.predicted_x = cx + 8.0f * gx;
                r.predicted_y = cy + 8.0f * gy;
                r.d = 2.0f;
                r.prescaling_corr = 1.0f;
                r.partiality = (gx == 0 && gy == 0) ? 0.5f : neighbour_partiality;
                predicted.push_back(r);
            }
        BraggIntegrationEngineCPU engine(experiment, PixelMask(experiment));
        const auto out = engine.Run(image, predicted, predicted.size(), 0);
        counts = engine.Counts();
        std::vector<Reflection> spot;
        for (const auto &r : out)
            if (r.h == 0 && r.k == 0)
                spot.push_back(r);
        return spot;
    };

    // Tails: no reflection is dropped, and the spot is measured against the clean flat background -
    // while the density the widened-radius guard reads still sees every prediction.
    BraggIntegrationCounts tails;
    const auto spot = run(0.01f, tails);
    CHECK(tails.bkg_starved == 0);
    CHECK(tails.bkg_starved_by_neighbour > 0);
    REQUIRE(spot.size() == 1);
    CHECK(spot[0].bkg == Catch::Approx(12.0f));
    CHECK(spot[0].I == Catch::Approx(6.2831853 * sigma * sigma * amp).epsilon(0.05));

    // The same predictions with their flux on this frame do take the rings, the spot's among them, and
    // it is the neighbours, not the detector, that starve them - the same rings the density counted.
    BraggIntegrationCounts on_frame;
    CHECK(run(1.0f, on_frame).empty());
    CHECK(on_frame.bkg_starved > 0);
    CHECK(on_frame.bkg_starved_by_neighbour == on_frame.bkg_starved);
    CHECK(on_frame.bkg_starved_by_neighbour == tails.bkg_starved_by_neighbour);
}

// The neighbour count taken without an image (CountNeighbourCrowding) is the one the integrator
// takes while integrating, on a frame whose pixels are all readable, whatever the partialities.
TEST_CASE("BraggIntegrationEngineCPU_CrowdingCountedWithoutImage", "[Integration][portable]") {
    DiffractionExperiment experiment(DetJF(2));
    experiment.DetectorDistance_mm(100.0f).IncidentEnergy_keV(WVL_1A_IN_KEV).BeamX_pxl(400.0f).BeamY_pxl(400.0f);
    experiment.ImportBraggIntegrationSettings(BraggIntegrationSettings().R1(5.0f).R2(7.0f).R3(13.0f));
    const size_t npixel = experiment.GetPixelsNum();
    ImagePreprocessorBuffer image(npixel);
    for (size_t i = 0; i < npixel; ++i)
        image[i] = 12;

    for (float spacing : {7.0f, 11.0f, 16.0f}) {
        std::vector<Reflection> predicted;
        for (int gy = -6; gy <= 6; ++gy)
            for (int gx = -6; gx <= 6; ++gx) {
                Reflection r{};
                r.h = gx; r.k = gy; r.l = 1;
                r.predicted_x = 600.3f + spacing * gx + 0.37f * gy;
                r.predicted_y = 300.2f + spacing * gy;
                r.d = 2.0f;
                r.prescaling_corr = 1.0f;
                r.partiality = (gx + gy) % 3 == 0 ? 0.01f : 0.8f;
                predicted.push_back(r);
            }
        BraggIntegrationEngineCPU integrated(experiment, PixelMask(experiment));
        integrated.Run(image, predicted, predicted.size(), 0);
        BraggIntegrationEngineCPU counted(experiment, PixelMask(experiment));
        counted.CountNeighbourCrowding(predicted, predicted.size());
        CHECK(counted.Counts().predicted == integrated.Counts().predicted);
        CHECK(counted.Counts().bkg_starved_by_neighbour == integrated.Counts().bkg_starved_by_neighbour);
        if (spacing == 7.0f)
            CHECK(counted.Counts().bkg_starved_by_neighbour > 0);
    }
}

// A saturated pixel (INT32_MAX, as the preprocessor marks it) at the peak of a spot on rotation data: the
// reflection is not measured but kept, flagged overloaded, so that a rotation merge can drop its whole
// rocking event. The same spot without the saturated pixel is an ordinary measurement.
TEST_CASE("BraggIntegrationEngineCPU_SaturatedPeakIsFlaggedNotDropped", "[Integration][portable]") {
    DiffractionExperiment experiment(DetJF(2));
    experiment.DetectorDistance_mm(100.0f).IncidentEnergy_keV(WVL_1A_IN_KEV).BeamX_pxl(400.0f).BeamY_pxl(400.0f)
            .Goniometer(GoniometerAxis("omega", 0.0f, 0.1f, Coord(-1, 0, 0), {}));
    experiment.ImportIndexingSettings(IndexingSettings().RotationIndexing(true));
    experiment.ImportBraggIntegrationSettings(BraggIntegrationSettings());
    const size_t width = experiment.GetXPixelsNum(), npixel = experiment.GetPixelsNum();

    const int cx = 600, cy = 300;
    ImagePreprocessorBuffer image(npixel);
    for (size_t i = 0; i < npixel; ++i)
        image[i] = 12;
    for (int dy = -6; dy <= 6; ++dy)
        for (int dx = -6; dx <= 6; ++dx)
            image[(cy + dy) * width + cx + dx] += static_cast<int32_t>(std::lround(800.0 * std::exp(-(dx * dx + dy * dy) / (2 * 1.3 * 1.3))));

    Reflection r{};
    r.h = 1; r.k = 2; r.l = 3;
    r.predicted_x = cx; r.predicted_y = cy;
    r.d = 2.0f;
    r.prescaling_corr = 1.0f;
    r.partiality = 1.0f;
    const std::vector<Reflection> predicted{r};

    BraggIntegrationEngineCPU engine(experiment, PixelMask(experiment));
    const auto clean = engine.Run(image, predicted, 1, 0);
    REQUIRE(clean.size() == 1);
    CHECK_FALSE(clean[0].overloaded);

    image[cy * width + cx] = INT32_MAX;
    const auto saturated = engine.Run(image, predicted, 1, 0);
    REQUIRE(saturated.size() == 1);
    CHECK(saturated[0].overloaded);
    CHECK(saturated[0].clipped);
}

// An empty aperture on a background that is a curved function of the radius alone: the flat ring mean
// reads it wrong by the curvature, and the radial correction with the curve taken from a per-ring
// profile of the whole frame (as the spot finder hands it over on stills) takes that back out.
TEST_CASE("BraggIntegrationEngineCPU_RadialCurveFromRings", "[Integration][portable]") {
    DiffractionExperiment experiment(DetJF(2));
    experiment.DetectorDistance_mm(100.0f).IncidentEnergy_keV(WVL_1A_IN_KEV).BeamX_pxl(400.0f).BeamY_pxl(400.0f);
    const float bx = 400.0f, by = 400.0f;
    const auto background = [](double r) { return 50.0 + 0.05 * (r - 300.0) * (r - 300.0); };
    const size_t width = experiment.GetXPixelsNum(), npixel = experiment.GetPixelsNum();
    ImagePreprocessorBuffer image(npixel);
    for (size_t i = 0; i < npixel; ++i)
        image[i] = static_cast<int32_t>(std::lround(
                background(std::hypot(static_cast<double>(i % width) - bx, static_cast<double>(i / width) - by))));

    // The profile per q bin, as the spot finder's rings would hold it.
    const float low_q = 0.05f, dq = 0.01f;
    const double F_px = 100.0 / 0.075, lambda = 1.0;
    std::vector<float> rings(500);
    for (size_t i = 0; i < rings.size(); ++i) {
        const double q = low_q + (i + 0.5) * dq;
        rings[i] = static_cast<float>(background(F_px * std::tan(2.0 * std::asin(q * lambda / (4.0 * PI)))));
    }
    rings[100] = NAN;   // a ring with no value is bridged

    Reflection r{};
    r.h = 1; r.k = 0; r.l = 0;
    r.predicted_x = bx + 300.0f;
    r.predicted_y = by;
    r.d = 2.0f;
    r.prescaling_corr = 1.0f;
    r.partiality = 1.0f;
    const std::vector<Reflection> predicted{r};

    auto integrate = [&](bool correct) {
        experiment.ImportBraggIntegrationSettings(
                BraggIntegrationSettings().Integrator(IntegratorMode::BoxSum).BackgroundRadialCorrection(correct));
        BraggIntegrationEngineCPU engine(experiment, PixelMask(experiment));
        engine.RadialBackgroundFromRings(rings, static_cast<int32_t>(rings.size()), low_q, dq);
        REQUIRE(!engine.GetRadialBackground().empty());
        const auto out = engine.Run(image, predicted, predicted.size(), 0);
        REQUIRE(out.size() == 1);
        return out[0].I;
    };
    // The curvature is along the radius only, so the error is n_inner x 0.05 x (<u^2>_ring - <u^2>_disk)
    // with u the radial offset: 49 px x 0.05 x ((6^2 + 13^2) / 4 - 4^2 / 4) = 116 counts read too low.
    CHECK(integrate(false) < -80.0f);
    CHECK(std::fabs(integrate(true)) < 5.0f);

    // An empty profile hands the curve back to the reflections' own rings.
    BraggIntegrationEngineCPU engine(experiment, PixelMask(experiment));
    engine.RadialBackgroundFromRings({}, 500, low_q, dq);
    CHECK(engine.GetRadialBackground().empty());
}

// A ring on a background falling steeply away from the beam: the flat high-side clip cuts the beam-side
// half of the ring and the background reads low. With the correction on, the clip follows the frame's
// radial curve (RadialClipOffset) and the empty aperture reads zero.
TEST_CASE("BraggIntegrationEngineCPU_ClipFollowsRadialCurve", "[Integration][portable]") {
    DiffractionExperiment experiment(DetJF(2));
    experiment.DetectorDistance_mm(100.0f).IncidentEnergy_keV(WVL_1A_IN_KEV).BeamX_pxl(400.0f).BeamY_pxl(400.0f);
    const float bx = 400.0f, by = 400.0f;
    const auto background = [](double r) { return 40000.0 - 100.0 * r; };
    const size_t width = experiment.GetXPixelsNum(), npixel = experiment.GetPixelsNum();
    ImagePreprocessorBuffer image(npixel);
    for (size_t i = 0; i < npixel; ++i)
        image[i] = static_cast<int32_t>(std::lround(std::max(0.0,
                background(std::hypot(static_cast<double>(i % width) - bx, static_cast<double>(i / width) - by)))));
    const float low_q = 0.05f, dq = 0.01f;
    const double F_px = 100.0 / 0.075, lambda = 1.0;
    std::vector<float> rings(500);
    for (size_t i = 0; i < rings.size(); ++i) {
        const double q = low_q + (i + 0.5) * dq;
        rings[i] = static_cast<float>(background(F_px * std::tan(2.0 * std::asin(q * lambda / (4.0 * PI)))));
    }
    Reflection r{};
    r.h = 1; r.k = 0; r.l = 0;
    r.predicted_x = bx + 200.0f;
    r.predicted_y = by;
    r.d = 2.0f;
    r.prescaling_corr = 1.0f;
    r.partiality = 1.0f;
    const std::vector<Reflection> predicted{r};
    auto background_read = [&](bool correct) {
        experiment.ImportBraggIntegrationSettings(BraggIntegrationSettings().BackgroundRadialCorrection(correct));
        BraggIntegrationEngineCPU engine(experiment, PixelMask(experiment));
        engine.RadialBackgroundFromRings(rings, static_cast<int32_t>(rings.size()), low_q, dq);
        const auto out = engine.Run(image, predicted, predicted.size(), 0);
        REQUIRE(out.size() == 1);
        return out[0].bkg;
    };
    // The kernel models the radial offset only; the ring's tangential reach adds the circle's own
    // curvature, slope * <v^2> / (2 r) = 100 * 51 / 400 = 13 counts, which it leaves.
    const double truth = background(200.0);
    CHECK(background_read(false) < truth - 100.0);
    CHECK(std::fabs(background_read(true) - truth) < 20.0);
}

// A still has no rocking event to drop, so an overloaded observation is kept as a lower bound: the box sum
// with every saturated pixel at the saturation limit. A spot with its core clipped at the limit, the
// clipped pixels marked saturated as the preprocessor marks them: the still gives exactly the clipped
// sum, unflagged and clipped (it may be low), and below the unclipped spot; rotation data keep the flag.
TEST_CASE("BraggIntegrationEngineCPU_StillKeepsOverloadAsLowerBound", "[Integration][portable]") {
    DiffractionExperiment experiment(DetJF(2));
    experiment.DetectorDistance_mm(100.0f).IncidentEnergy_keV(WVL_1A_IN_KEV).BeamX_pxl(400.0f).BeamY_pxl(400.0f);
    experiment.ImportBraggIntegrationSettings(BraggIntegrationSettings());
    const int32_t limit = 500, bkg = 12;
    experiment.Detector().SaturationLimit(limit);
    const size_t width = experiment.GetXPixelsNum(), npixel = experiment.GetPixelsNum();

    const int cx = 600, cy = 300;
    const float r1 = experiment.GetBraggIntegrationSettings().GetR1();
    ImagePreprocessorBuffer image(npixel);
    for (size_t i = 0; i < npixel; ++i)
        image[i] = bkg;
    double clipped_sum = 0.0, true_sum = 0.0;
    int n_saturated = 0;
    for (int dy = -6; dy <= 6; ++dy)
        for (int dx = -6; dx <= 6; ++dx) {
            const int32_t v = bkg + static_cast<int32_t>(std::lround(2000.0 * std::exp(-(dx * dx + dy * dy) / (2 * 1.3 * 1.3))));
            image[(cy + dy) * width + cx + dx] = v >= limit ? INT32_MAX : v;
            if (dx * dx + dy * dy < r1 * r1) {
                clipped_sum += std::min(v, limit) - bkg;
                true_sum += v - bkg;
                n_saturated += v >= limit;
            }
        }
    REQUIRE(n_saturated > 1);

    Reflection r{};
    r.h = 1; r.k = 2; r.l = 3;
    r.predicted_x = cx; r.predicted_y = cy;
    r.d = 2.0f;
    r.prescaling_corr = 1.0f;
    r.partiality = 1.0f;
    const std::vector<Reflection> predicted{r};

    BraggIntegrationEngineCPU still(experiment, PixelMask(experiment));
    const auto kept = still.Run(image, predicted, 1, 0);
    REQUIRE(kept.size() == 1);
    CHECK_FALSE(kept[0].overloaded);
    CHECK(kept[0].clipped);
    CHECK(kept[0].I == Catch::Approx(clipped_sum).epsilon(0.01));
    CHECK(kept[0].I < 0.8 * true_sum);
    CHECK(kept[0].sigma >= std::sqrt(kept[0].I));
    CHECK(still.Counts().overloads_kept == 1);

    experiment.Goniometer(GoniometerAxis("omega", 0.0f, 0.1f, Coord(-1, 0, 0), {}))
            .ImportIndexingSettings(IndexingSettings().RotationIndexing(true));
    BraggIntegrationEngineCPU rotation(experiment, PixelMask(experiment));
    const auto flagged = rotation.Run(image, predicted, 1, 0);
    REQUIRE(flagged.size() == 1);
    CHECK(flagged[0].overloaded);
    CHECK(rotation.Counts().overloads_kept == 0);
}

// A still keeps an overloaded observation as a lower bound only where it has pixels at the saturation
// limit. One overloaded by a pixel unreadable on this frame alone has nothing to count at the limit:
// it stays flagged and is dropped.
TEST_CASE("BraggIntegrationEngineCPU_StillKeepsOnlySaturatedOverloads", "[Integration][portable]") {
    DiffractionExperiment experiment(DetJF(2));
    experiment.DetectorDistance_mm(100.0f).IncidentEnergy_keV(WVL_1A_IN_KEV).BeamX_pxl(400.0f).BeamY_pxl(400.0f);
    experiment.ImportBraggIntegrationSettings(BraggIntegrationSettings());
    const int32_t limit = 500, bkg = 12;
    experiment.Detector().SaturationLimit(limit);
    const size_t width = experiment.GetXPixelsNum(), npixel = experiment.GetPixelsNum();

    ImagePreprocessorBuffer image(npixel);
    for (size_t i = 0; i < npixel; ++i)
        image[i] = bkg;
    auto spot = [&](int cx, int cy, double amplitude) {
        for (int dy = -6; dy <= 6; ++dy)
            for (int dx = -6; dx <= 6; ++dx) {
                const int32_t v = bkg + static_cast<int32_t>(std::lround(amplitude * std::exp(-(dx * dx + dy * dy) / (2 * 1.3 * 1.3))));
                image[(cy + dy) * width + cx + dx] = v >= limit ? INT32_MAX : v;
            }
    };
    spot(600, 300, 2000.0);                  // clipped: a plateau of saturated pixels in a bright spot
    spot(300, 300, 40.0);
    image[300 * width + 300] = INT32_MIN;    // its peak unreadable on this frame, not in the mask

    std::vector<Reflection> predicted;
    for (const auto &[x, y] : {std::pair{600, 300}, std::pair{300, 300}}) {
        Reflection r{};
        r.h = x; r.k = y; r.l = 1;
        r.predicted_x = x; r.predicted_y = y;
        r.d = 2.0f;
        r.prescaling_corr = 1.0f;
        r.partiality = 1.0f;
        predicted.push_back(r);
    }

    BraggIntegrationEngineCPU still(experiment, PixelMask(experiment));
    const auto out = still.Run(image, predicted, predicted.size(), 0);
    REQUIRE(out.size() == 2);
    CHECK_FALSE(out[0].overloaded);
    CHECK(out[1].overloaded);
    CHECK(still.Counts().overloads_kept == 1);
}

// A detector that reports gain ADU per photon: the same spot read at gain 4 leaves the integrator in
// photons - a quarter of the ADU intensity, at the Poisson sigma of that many photons, i.e. half the
// sigma the ADU counted as photons gave.
TEST_CASE("BraggIntegrationEngineCPU_PixelGainGivesPhotons", "[Integration][portable]") {
    DiffractionExperiment experiment(DetJF(2));
    experiment.DetectorDistance_mm(100.0f).IncidentEnergy_keV(WVL_1A_IN_KEV).BeamX_pxl(400.0f).BeamY_pxl(400.0f);
    const size_t width = experiment.GetXPixelsNum(), npixel = experiment.GetPixelsNum();

    const int cx = 600, cy = 300;
    ImagePreprocessorBuffer image(npixel);
    for (size_t i = 0; i < npixel; ++i)
        image[i] = 40;
    for (int dy = -6; dy <= 6; ++dy)
        for (int dx = -6; dx <= 6; ++dx)
            image[(cy + dy) * width + cx + dx] += static_cast<int32_t>(std::lround(800.0 * std::exp(-(dx * dx + dy * dy) / (2 * 1.3 * 1.3))));

    Reflection r{};
    r.h = 1; r.k = 2; r.l = 3;
    r.predicted_x = cx; r.predicted_y = cy;
    r.d = 2.0f;
    r.prescaling_corr = 1.0f;
    r.partiality = 1.0f;
    const std::vector<Reflection> predicted{r};

    for (const auto mode : {IntegratorMode::BoxSum, IntegratorMode::ProfileGaussian}) {
        BraggIntegrationSettings settings;
        settings.Integrator(mode);
        experiment.ImportBraggIntegrationSettings(settings);
        BraggIntegrationEngineCPU adu(experiment, PixelMask(experiment));
        const auto out_adu = adu.Run(image, predicted, 1, 0);
        settings.PixelGain(4.0f);
        experiment.ImportBraggIntegrationSettings(settings);
        BraggIntegrationEngineCPU photons(experiment, PixelMask(experiment));
        const auto out_ph = photons.Run(image, predicted, 1, 0);
        REQUIRE(out_adu.size() == 1);
        REQUIRE(out_ph.size() == 1);
        CHECK(out_ph[0].I == Catch::Approx(out_adu[0].I / 4.0f));
        CHECK(out_ph[0].sigma == Catch::Approx(out_adu[0].sigma / 2.0f));
        CHECK(out_ph[0].bkg == Catch::Approx(out_adu[0].bkg / 4.0f));
        CHECK(out_ph[0].var_bkg == Catch::Approx(out_adu[0].var_bkg / 4.0f));
    }
}