Jungfraujoch/image_analysis/geom_refinement/BeamCenterFFT.cpp

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

#include "BeamCenterFFT.h"

#include <algorithm>
#include <bit>
#include <cmath>
#include <cstdint>
#include <limits>
#include <stdexcept>

#include "BeamCenterFFTCPU.h"
#include "BeamCenterFFTEngine.h"
#include "../../common/CUDAWrapper.h"
#include <spdlog/spdlog.h>
#ifdef JFJOCH_USE_CUDA
#include "BeamCenterFFTGPU.h"
#include "../../common/JFJochException.h"
#endif

namespace {

// The image prepared for scoring: hot pixels clipped at the 99.9th percentile, the median
// background level removed, negatives clamped, masked pixels zero. The final subtraction of the
// valid-pixel mean changes nothing mathematically - the masked Pearson is exactly invariant under
// a global shift of the image (both D*C - S^2 and D*Q - S^2 cancel the shift) - but it removes
// most of the large-term cancellation that single-precision FFTs would otherwise have to survive.
struct Prepared {
    std::vector<float> a;   // image, 0 on masked pixels
    std::vector<float> m;   // 1 on valid pixels, 0 on masked
    double variance = 0.0;  // of a over the valid pixels, after the whole preparation
};

// Order-preserving: unsigned comparison of the keys orders the floats as < does (-0 below +0).
uint32_t SortKey(float v) {
    const uint32_t u = std::bit_cast<uint32_t>(v);
    return (u & 0x80000000u) ? ~u : (u | 0x80000000u);
}

// For each k, the value nth_element would leave at position k of the finite values of x - without
// copying or reordering the image. Counting the finite values by the upper 16 bits of their key
// finds the bucket holding the k-th; counting that bucket by the lower 16 bits finds its key,
// which is the value. Two passes over x for all k together.
std::vector<float> KthFinite(const std::vector<float> &x, std::vector<size_t> k) {
    std::vector<size_t> upper(1 << 16, 0);
    for (const float v : x)
        if (std::isfinite(v))
            upper[SortKey(v) >> 16]++;
    std::vector<uint32_t> bucket(k.size(), 0);
    for (size_t i = 0; i < k.size(); i++)
        while (k[i] >= upper[bucket[i]])
            k[i] -= upper[bucket[i]++];

    std::vector<std::vector<size_t>> lower(k.size(), std::vector<size_t>(1 << 16, 0));
    for (const float v : x)
        if (std::isfinite(v))
            for (size_t i = 0; i < k.size(); i++)
                if (SortKey(v) >> 16 == bucket[i])
                    lower[i][SortKey(v) & 0xFFFF]++;
    std::vector<float> out;
    for (size_t i = 0; i < k.size(); i++) {
        uint32_t b = 0;
        while (k[i] >= lower[i][b])
            k[i] -= lower[i][b++];
        const uint32_t key = (bucket[i] << 16) | b;
        out.push_back(std::bit_cast<float>((key & 0x80000000u) ? (key & 0x7FFFFFFFu) : ~key));
    }
    return out;
}

Prepared PrepareImage(const std::vector<float> &mean) {
    Prepared out;
    const size_t n = mean.size();
    out.a.assign(n, 0.0f);
    out.m.assign(n, 0.0f);

    const size_t n_valid = std::count_if(mean.begin(), mean.end(), [](float v) { return std::isfinite(v); });
    if (n_valid < 2)
        return out;

    // Where these fall among equal values a -0 can come out for a +0, which the clipping below
    // turns into the same image.
    const std::vector<float> q = KthFinite(mean, {static_cast<size_t>(0.999 * (n_valid - 1)), n_valid / 2});
    const float clip_hi = q[0];
    const float median = q[1];

    double sum = 0.0;
    size_t count = 0;
    for (size_t i = 0; i < n; i++) {
        if (!std::isfinite(mean[i]))
            continue;
        const float v = std::max(0.0f, std::min(std::max(mean[i], 0.0f), clip_hi) - median);
        out.a[i] = v;
        out.m[i] = 1.0f;
        sum += v;
        count++;
    }
    const float shift = static_cast<float>(sum / static_cast<double>(count));
    double var = 0.0;
    for (size_t i = 0; i < n; i++)
        if (out.m[i] != 0.0f) {
            out.a[i] -= shift;
            var += static_cast<double>(out.a[i]) * static_cast<double>(out.a[i]);
        }
    out.variance = var / static_cast<double>(count);
    return out;
}

float Margin(const std::vector<BeamCenterFFTCandidate> &peaks) {
    if (peaks.size() < 2 || peaks[0].score <= 0.0f)
        return 0.0f;
    return (peaks[0].score - peaks[1].score) / peaks[0].score;
}

} // namespace

std::vector<float> BeamCenterPointScore(const BeamCenterConvSurfaces2D &conv,
                                        float min_pair_fraction, double global_variance) {
    const size_t n = conv.C.size();
    float d_max = 0.0f;
    for (size_t i = 0; i < n; i++)
        d_max = std::max(d_max, conv.D[i]);
    const double lim = min_pair_fraction * d_max;
    std::vector<float> r(n, -std::numeric_limits<float>::infinity());
    for (size_t i = 0; i < n; i++) {
        const double d = conv.D[i];
        if (d <= lim)
            continue;
        const double s = conv.S[i];
        const double num = d * conv.C[i] - s * s;
        const double den = d * conv.Q[i] - s * s;
        if (den > d * d * BEAM_CENTER_VARIANCE_FLOOR * global_variance)
            r[i] = static_cast<float>(num / den);
    }
    return r;
}

std::vector<BeamCenterFFTCandidate> BeamCenterShortlist2D(std::vector<float> &surface, int64_t h,
                                                          int64_t w, float nms_pxl, int budget) {
    std::vector<BeamCenterFFTCandidate> out;
    const int64_t d = std::lround(2.0f * nms_pxl);
    // Each round takes the first maximum of the whole surface in index order. Kept per row - its
    // maximum and the first index holding it - so a round reads the rows' maxima, and only the rows
    // the suppression touched are scanned again. Rows in order and the first index within a row is
    // the first index overall, the same pick as one scan of the whole surface.
    std::vector<float> row_v(h);
    std::vector<int64_t> row_i(h);
    const auto scan_row = [&](int64_t y) {
        float v = -std::numeric_limits<float>::infinity();
        int64_t at = -1;
        for (int64_t x = 0; x < w; x++)
            if (surface[static_cast<size_t>(y * w + x)] > v) {
                v = surface[static_cast<size_t>(y * w + x)];
                at = y * w + x;
            }
        row_v[y] = v;
        row_i[y] = at;
    };
    for (int64_t y = 0; y < h; y++)
        scan_row(y);
    for (int k = 0; k < budget; k++) {
        int64_t best = -1;
        float best_v = -std::numeric_limits<float>::infinity();
        for (int64_t y = 0; y < h; y++)
            if (row_i[y] >= 0 && row_v[y] > best_v) {
                best_v = row_v[y];
                best = row_i[y];
            }
        if (best < 0 || !std::isfinite(best_v))
            break;
        const int64_t iy = best / w, ix = best % w;
        out.push_back({static_cast<float>(ix) / 2.0f, static_cast<float>(iy) / 2.0f, best_v});
        for (int64_t y = std::max<int64_t>(0, iy - d); y <= std::min(h - 1, iy + d); y++) {
            for (int64_t x = std::max<int64_t>(0, ix - d); x <= std::min(w - 1, ix + d); x++)
                surface[static_cast<size_t>(y * w + x)] =
                        -std::numeric_limits<float>::infinity();
            scan_row(y);
        }
    }
    return out;
}

// The default route, and the whole of the CPU path: the four convolutions come back, and the
// combination and the search run here.
std::vector<BeamCenterFFTCandidate>
BeamCenterFFTEngine::PointShortlist(const std::vector<float> &a, const std::vector<float> &m,
                                    int64_t h, int64_t w, const BeamCenterFFTSettings &settings,
                                    double global_variance) {
    const auto conv = PointSurfaces(a, m, h, w);
    auto surface = BeamCenterPointScore(conv, settings.min_pair_fraction, global_variance);
    return BeamCenterShortlist2D(surface, 2 * h, 2 * w, settings.nms_radius_pxl,
                                 settings.candidates_point);
}

int64_t BeamCenterFFTPadSize(int64_t n) {
    for (int64_t k = n;; ++k) {
        int64_t v = k;
        for (int p : {2, 3, 5, 7})
            while (v % p == 0)
                v /= p;
        if (v == 1)
            return k;
    }
}

BeamCenterFFTResult BeamCenterFFTScore(int64_t width, int64_t height,
                                       const std::vector<float> &mean,
                                       const BeamCenterFFTSettings &settings,
                                       BeamCenterFFTEngine &engine) {
    if (width <= 0 || height <= 0
        || mean.size() != static_cast<size_t>(width) * static_cast<size_t>(height))
        throw std::runtime_error("BeamCenterFFT: image dimensions do not match the projection");

    BeamCenterFFTResult result;
    const Prepared prep = PrepareImage(mean);

    result.point = engine.PointShortlist(prep.a, prep.m, height, width, settings, prep.variance);
    result.margin_point = Margin(result.point);
    return result;
}

BeamCenterFFTResult BeamCenterFFTScore(int64_t width, int64_t height,
                                       const std::vector<float> &mean,
                                       const BeamCenterFFTSettings &settings) {
#ifdef JFJOCH_USE_CUDA
    if (get_gpu_count() > 0) {
        if (BeamCenterFFTGPU::FitsInDeviceMemory(width, height)) {
            BeamCenterFFTGPU gpu;
            return BeamCenterFFTScore(width, height, mean, settings, gpu);
        }
        spdlog::info("Beam centre FFT: scoring on the CPU - on the GPU it needs {:.2f} GB, over half of the card",
                     BeamCenterFFTGPU::DeviceMemoryNeeded(width, height) / 1e9);
    }
#endif
    BeamCenterFFTCPU cpu;
    return BeamCenterFFTScore(width, height, mean, settings, cpu);
}

BeamCenterFFTResult FindBeamCenterFFT(const DiffractionExperiment &experiment,
                                      const std::vector<float> &mean,
                                      const BeamCenterFFTSettings &settings) {
    // The pre-scan projection is in converted geometry (ShadowFinder's frame).
    return BeamCenterFFTScore(experiment.GetXPixelsNumConv(), experiment.GetYPixelsNumConv(),
                              mean, settings);
}