Jungfraujoch/image_analysis/geom_refinement/BeamCenterFFTEngine.h

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

#pragma once

#include <cstdint>
#include <vector>

#include "BeamCenterFFT.h"

// The transform half of BeamCenterFFTScore, with a CPU (fftw3f) and a GPU (cuFFT) implementation
// chosen the way the FFT indexer's are. Everything that decides anything - the preparation of the
// image, the masked Pearson combination, the shortlist and the margins - is shared, so the two
// engines can differ only in how the four convolutions are computed. That is what makes them
// comparable: the parity test scores one image with each and compares the shortlists.
//
// The masked Pearson needs four convolutions of the prepared image `a` and its valid-pixel mask
// `m`: C = a (*) a, S = a (*) m, Q = a^2 (*) m and D = m (*) m.

// The 2D point-inversion score: each surface cropped to 2h x 2w, row-major.
struct BeamCenterConvSurfaces2D {
    std::vector<float> C, S, Q, D;
};

// The 1D line-mirror score: each sequence correlated with its own mirror and the four accumulators
// summed over the other coordinate, so each surface is 2 x (the mirrored extent) long.
struct BeamCenterConvSurfaces1D {
    std::vector<double> C, S, Q, D;
};

// Which coordinate the 1D score mirrors: Rows mirrors y (each column is one sequence) and gives
// the beam's y; Columns mirrors x and gives its x.
enum class BeamCenterMirror { Rows, Columns };

class BeamCenterFFTEngine {
public:
    virtual ~BeamCenterFFTEngine() = default;

    // `a` and `m` are h x w row-major, as PrepareImage leaves them.
    virtual BeamCenterConvSurfaces2D PointSurfaces(const std::vector<float> &a,
                                                   const std::vector<float> &m, int64_t h,
                                                   int64_t w) = 0;
    virtual BeamCenterConvSurfaces1D LineSurfaces(const std::vector<float> &a,
                                                  const std::vector<float> &m, int64_t h, int64_t w,
                                                  BeamCenterMirror mirror) = 0;

    // The 2D point score all the way to its shortlist. The default - and the whole of the CPU
    // path - is PointSurfaces followed by BeamCenterPointScore and BeamCenterShortlist2D below.
    // An engine whose surfaces are already somewhere else overrides it rather than moving them:
    // on a 16 Mpixel detector the two host steps are 1.8 s and bringing the four surfaces back is
    // another 0.9 s, against 26 ms of transform.
    virtual std::vector<BeamCenterFFTCandidate>
    PointShortlist(const std::vector<float> &a, const std::vector<float> &m, int64_t h, int64_t w,
                   const BeamCenterFFTSettings &settings, double global_variance);
};

// The masked Pearson r(t) = (D*C - S^2) / (D*Q - S^2) from the four convolution surfaces, combined
// elementwise in double: the numerator and denominator are cancellations of large near-equal terms,
// and the shortlist margins ride on differences of ~0.3 %. Centres whose mirror overlap holds fewer
// than `min_pair_fraction` of the largest pair count, and those whose overlap carries no variance
// (see BEAM_CENTER_VARIANCE_FLOOR), score -infinity.
//
// This is the reference implementation. An engine that computes the same thing where its surfaces
// already are is held to it element by element by the parity test.
[[nodiscard]] std::vector<float> BeamCenterPointScore(const BeamCenterConvSurfaces2D &conv,
                                                      float min_pair_fraction,
                                                      double global_variance);

// Greedy shortlist of a 2D score surface: strongest peak first, a square of half-width `nms_pxl`
// around each taken peak suppressed. Ties go to the lowest index, and a non-finite entry is never a
// peak. The surface is indexed by t = 2c, so the centre grid is half-pixel. Reference, as above.
// `surface` is consumed - the suppression is written into it.
[[nodiscard]] std::vector<BeamCenterFFTCandidate> BeamCenterShortlist2D(std::vector<float> &surface,
                                                                        int64_t h, int64_t w,
                                                                        float nms_pxl, int budget);

// The denominator D*Q - S^2 equals D^2 times the image variance over the mirror overlap, so an
// overlap with (nearly) no variance - a mirrored empty corner, everything clamped to zero by the
// preparation - carries no evidence and its r is pure roundoff (measured: such a region minted
// r = 4..274 on a synthetic image, in float64 as much as float32). This is the fraction of the
// whole image's variance an overlap must reach to be scored at all; on real detector images every
// overlap that passes the pair-count gate is orders of magnitude above it.
constexpr double BEAM_CENTER_VARIANCE_FLOOR = 1e-4;

// Smallest 2-3-5-7-smooth transform length >= n. Both engines pad to it, so both score the same
// grid. A smooth length is not a micro-optimisation here: the naive next-power-of-two pad of a
// 16 Mpixel detector costs 3x the time of the smooth one on both FFTW and cuFFT, and 2.5 GB of
// cuFFT plan memory against a few hundred MB (measured).
[[nodiscard]] int64_t BeamCenterFFTPadSize(int64_t n);

// BeamCenterFFTScore with the engine named rather than chosen, for the GPU/CPU parity test.
[[nodiscard]] BeamCenterFFTResult BeamCenterFFTScore(int64_t width, int64_t height,
                                                     const std::vector<float> &mean,
                                                     const BeamCenterFFTSettings &settings,
                                                     BeamCenterFFTEngine &engine);