Jungfraujoch/image_analysis/spot_finding/AdaptiveSpotFinderCPU.cpp

// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only

#include <algorithm>
#include <cmath>

#include "AdaptiveSpotFinderCPU.h"
#include "AdaptiveThreshold.h"

AdaptiveSpotFinderCPU::AdaptiveSpotFinderCPU(const AzimuthalIntegrationMapping &in_mapping)
    : ImageSpotFinderCPU(static_cast<int32_t>(in_mapping.GetWidth()),
                         static_cast<int32_t>(in_mapping.GetHeight())),
      mapping(in_mapping) {
    const size_t nbins = mapping.GetBinNumber();
    ring_sum.assign(nbins, 0);
    ring_sum2.assign(nbins, 0);
    ring_cnt.assign(nbins, 0);
    ring_mean.assign(nbins, 0.0f);
    ring_sigma.assign(nbins, 0.0f);
    ring_thr.assign(nbins + 1, INFINITY);   // the last entry is for pixels outside every ring
    ring_bkg.assign(nbins, NAN);
    ring_bits.assign(OutputSize(), 0);
    ring_hist.assign(nbins * HIST_VALUES, 0);
    azint_sum.assign(nbins, 0.0f);
    azint_sum2.assign(nbins, 0.0f);
    azint_count.assign(nbins, 0);
}

void AdaptiveSpotFinderCPU::GetProfile(AzimuthalIntegrationProfile &profile) const {
    profile.Clear(mapping);
    profile.Add(azint_sum, azint_sum2, azint_count);
}

// Per-ring background statistics with iterated peak exclusion following peakfinder8:
// Barty et al. (2014) J. Appl. Cryst. 47, 1118-1131
// Per-ring mean/variance from the raw (photon) image: a plain pass over every valid pixel, then
// sigma-clip passes keeping only pixels within clip_k sigma of the current ring mean, which removes
// the Bragg peaks from the background estimate.
void AdaptiveSpotFinderCPU::ResetRings() {
    std::fill(ring_sum.begin(), ring_sum.end(), 0);
    std::fill(ring_sum2.begin(), ring_sum2.end(), 0);
    std::fill(ring_cnt.begin(), ring_cnt.end(), 0);
    std::fill(ring_hist.begin(), ring_hist.end(), 0);
    ring_overflow.clear();
    if (fuse_azint) {
        std::fill(azint_sum.begin(), azint_sum.end(), 0.0f);
        std::fill(azint_sum2.begin(), azint_sum2.end(), 0.0f);
        std::fill(azint_count.begin(), azint_count.end(), 0);
    }
}

void AdaptiveSpotFinderCPU::BeginRings() {
    ResetRings();
    rings_from_blocks = true;
}

// The plain pass over pixels [first, first + n): the histogram of each ring's values, from which
// PlainRings() takes the integer sums, and the fused profile when asked for. One loop reads each pixel
// once for both.
void AdaptiveSpotFinderCPU::AccumulateRingsBlock(const ImagePreprocessorBuffer &image, size_t first, size_t n) {
    const auto &pixel_to_bin = mapping.GetPixelToBin();
    const size_t nbins = ring_sum.size();
    const float *corrections = mapping.Corrections().data();
    uint32_t *hist = ring_hist.data();

    // Consecutive pixels mostly share a ring, so the ring's profile sums are held in locals while they
    // do and written back when the ring changes: the same additions in the same order, without a store
    // and a reload of the same address on every pixel. Values outside the histogram are listed by a
    // second loop, run only when the block has any.
    bool overflow = false;
    size_t cur = nbins;   // the ring held in the locals below; nbins = none
    float az_sum = 0.0f, az_sum2 = 0.0f;
    uint32_t az_cnt = 0;
    for (size_t pxl = first; pxl < first + n; ++pxl) {
        const int32_t v = image[pxl];
        if (v == INT32_MIN || v == INT32_MAX) continue;   // bad / saturated
        const uint16_t b = pixel_to_bin[pxl];
        if (b >= nbins) continue;                          // masked / out of range (UINT16_MAX)
        if (static_cast<uint32_t>(v) < HIST_VALUES)
            hist[v * nbins + b] += 1;
        else
            overflow = true;
        if (!fuse_azint) continue;
        if (b != cur) {
            if (cur != nbins) {
                azint_sum[cur] = az_sum;
                azint_sum2[cur] = az_sum2;
                azint_count[cur] = az_cnt;
            }
            cur = b;
            az_sum = azint_sum[b];
            az_sum2 = azint_sum2[b];
            az_cnt = azint_count[b];
        }
        const float val = static_cast<float>(v) * corrections[pxl];
        const float val_sq = val * val;
        az_sum += val;
        az_sum2 += val_sq;
        ++az_cnt;
    }
    if (cur != nbins) {
        azint_sum[cur] = az_sum;
        azint_sum2[cur] = az_sum2;
        azint_count[cur] = az_cnt;
    }

    if (overflow)
        for (size_t pxl = first; pxl < first + n; ++pxl) {
            const int32_t v = image[pxl];
            if (v == INT32_MIN || v == INT32_MAX) continue;
            const uint16_t b = pixel_to_bin[pxl];
            if (b >= nbins) continue;
            if (v < 0 || v >= HIST_VALUES)
                ring_overflow.emplace_back(b, v);
        }
}

// A sigma-clip pass over the plain pass's values: each distinct value of a ring meets the same test the
// pixels holding it would, and its pixels are added as a count. Integer sums, so with nothing clipped
// (clip_k = INFINITY) they are the plain sums over the pixels themselves.
void AdaptiveSpotFinderCPU::ClipRings(float clip_k) {
    const size_t nbins = ring_sum.size();

    std::fill(ring_sum.begin(), ring_sum.end(), 0);
    std::fill(ring_sum2.begin(), ring_sum2.end(), 0);
    std::fill(ring_cnt.begin(), ring_cnt.end(), 0);

    const auto keep = [&](uint16_t b, int32_t v) {
        if (std::isinf(clip_k)) return true;
        const float lo = ring_mean[b] - clip_k * ring_sigma[b];
        const float hi = ring_mean[b] + clip_k * ring_sigma[b];
        return !(v < lo || v > hi);                        // exclude peaks / outliers
    };
    for (int32_t v = 0; v < HIST_VALUES; ++v)
        for (size_t b = 0; b < nbins; ++b) {
            const uint32_t n = ring_hist[v * nbins + b];
            if (n == 0 || !keep(static_cast<uint16_t>(b), v)) continue;
            ring_sum[b] += static_cast<int64_t>(n) * v;
            ring_sum2[b] += static_cast<uint64_t>(n) * static_cast<uint64_t>(static_cast<int64_t>(v) * v);
            ring_cnt[b] += n;
        }
    for (const auto &[b, v] : ring_overflow) {
        if (!keep(b, v)) continue;
        ring_sum[b] += v;
        ring_sum2[b] += static_cast<uint64_t>(static_cast<int64_t>(v) * v);
        ring_cnt[b] += 1;
    }
}

void AdaptiveSpotFinderCPU::UpdateRingStatistics() {
    const size_t nbins = ring_sum.size();
    for (size_t b = 0; b < nbins; ++b) {
        if (ring_cnt[b] > 0) {
            const double m = static_cast<double>(ring_sum[b]) / ring_cnt[b];
            const double var = std::max(0.0, static_cast<double>(ring_sum2[b]) / ring_cnt[b] - m * m);
            ring_mean[b] = static_cast<float>(m);
            ring_sigma[b] = static_cast<float>(std::sqrt(var));
        }
    }
}

void AdaptiveSpotFinderCPU::Detect(const ImagePreprocessorBuffer &image,
                                   const SpotFindingSettings &settings) {
    const size_t nbins = ring_sum.size();

    // --- Stage A: robust per-ring background (one plain pass + two sigma-clip passes) ---
    if (!rings_from_blocks) {
        ResetRings();
        AccumulateRingsBlock(image, 0, static_cast<size_t>(width) * height);
    }
    rings_from_blocks = false;
    ClipRings(INFINITY);
    UpdateRingStatistics();
    ClipRings(3.0f);
    UpdateRingStatistics();
    ClipRings(3.0f);
    UpdateRingStatistics();

    // --- Stage B: per-ring threshold from the single portable knob E (false pixels / frame) ---
    int64_t n_total = 0;
    double g_sum = 0.0, g_sum2 = 0.0;
    for (size_t b = 0; b < nbins; ++b) {
        n_total += ring_cnt[b];
        g_sum += static_cast<double>(ring_sum[b]);
        g_sum2 += static_cast<double>(ring_sum2[b]);
    }
    if (n_total == 0) {
        // Nothing valid to threshold against: leave no strong pixels for ExtractSpots to build on.
        std::fill(output_buffer.begin(), output_buffer.end(), 0);
        std::fill(ring_bkg.begin(), ring_bkg.end(), NAN);
        return;
    }

    for (size_t b = 0; b < nbins; ++b)
        ring_bkg[b] = (ring_cnt[b] < adaptive_threshold::MIN_RING_PIXELS) ? NAN : ring_mean[b];

    const double E = std::max(1.0f, settings.false_pixels_per_frame);
    double p = E / static_cast<double>(n_total);
    p = std::min(std::max(p, 1e-9), 0.1);
    const float z = static_cast<float>(adaptive_threshold::NormalQuantile(1.0 - p));

    // whole-frame fallback background for rings too sparse to trust on their own
    const double g_mean = g_sum / n_total;
    const double g_sigma = std::sqrt(std::max(0.0, g_sum2 / n_total - g_mean * g_mean));
    const float g_thr = adaptive_threshold::RingThreshold(static_cast<float>(g_mean),
                                                          static_cast<float>(g_sigma), p, z);

    for (size_t b = 0; b < nbins; ++b)
        ring_thr[b] = (ring_cnt[b] < adaptive_threshold::MIN_RING_PIXELS)
                          ? g_thr
                          : adaptive_threshold::RingThreshold(ring_mean[b], ring_sigma[b], p, z);

    // --- Stage C: the ring threshold, intersected with the classic local-box SNR test ---
    std::fill(ring_bits.begin(), ring_bits.end(), 0);

    if (settings.signal_to_noise_threshold <= 0.0f) {
        // No local test asked for: the ring threshold alone decides, as the fixed photon floor
        // alone would in the classic finder.
        for (int32_t row = 0; row < height; row++)
            FlagRow(image, row);
        output_buffer = ring_bits;
        return;
    }

    // The ring threshold IS the photon floor here, so the local pass must not apply another one.
    SpotFindingSettings local = settings;
    local.photon_count_threshold = 0;
    // Only the ring pixels survive the intersection, so the local test is asked of those alone. They
    // are flagged row by row as the local pass reaches each row, which saves reading the image for it.
    DetectAt(image, local, ring_bits, [&](int32_t row) { FlagRow(image, row); });
}

void AdaptiveSpotFinderCPU::FlagRow(const ImagePreprocessorBuffer &image, int32_t row) {
    const auto &pixel_to_bin = mapping.GetPixelToBin();
    const auto nbins = static_cast<uint32_t>(ring_thr.size() - 1);   // ring_thr[nbins] is +inf
    const float *thr = ring_thr.data();
    const int32_t *img = image.data();
    const uint16_t *bin = pixel_to_bin.data();
    const size_t first = static_cast<size_t>(row) * width;
    const size_t end = first + width;

    // Saturated is strong, bad is not, and a pixel outside every ring (bin >= nbins) meets the +inf
    // threshold. Written without branches and a word of 32 pixels at a time, so that the loop vectorises.
    const auto strong = [&](size_t pxl) -> uint32_t {
        const int32_t v = img[pxl];
        const uint32_t b = std::min<uint32_t>(bin[pxl], nbins);
        return (v == INT32_MAX) | ((v != INT32_MIN) & (v >= thr[b]));
    };
    size_t pxl = first;
    while (pxl < end && pxl % 32 != 0) {
        ring_bits[pxl / 32] |= strong(pxl) << (pxl % 32);
        ++pxl;
    }
    for (; pxl + 32 <= end; pxl += 32) {
        uint32_t word = 0;
        for (uint32_t j = 0; j < 32; ++j)
            word |= strong(pxl + j) << j;
        ring_bits[pxl / 32] |= word;
    }
    for (; pxl < end; ++pxl)
        ring_bits[pxl / 32] |= strong(pxl) << (pxl % 32);
}