Jungfraujoch/image_analysis/bragg_integration/BraggIntegrationEngine.h

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

#pragma once

// =============================================================================
// BraggIntegrationEngine — box-sum + profile-fitting 2D integrator, GPU-ready
// =============================================================================
//
// A reimplementation of BraggIntegrate2D (box sum) and ProfileIntegrate2D (Kabsch profile
// fit) under one roof, following the AzIntEngine / ROIIntegration pattern: a base class that
// extracts the fixed per-experiment configuration, a plain-C++ CPU engine (the fallback and the
// numeric oracle), and a CUDA engine (BraggIntegrationEngineGPU) that reaches the same result up
// to floating-point precision.
//
// Unlike BraggIntegrate2D/ProfileIntegrate2D, which read the raw CompressedImage per pixel type
// and reject the special/saturation +/-1 band, this engine reads the already-preprocessed int32
// image held in an ImagePreprocessorBuffer (the same buffer AzIntEngineGPU/ROIIntegrationGPU
// consume): masked/bad pixels are INT32_MIN and saturated pixels INT32_MAX, so bad-pixel identity
// is owned by the preprocessor and a pixel is valid iff v != INT32_MIN && v != INT32_MAX.
//
// The integrator is selected by BraggIntegrationSettings::Integrator:
//   BoxSum          -> BraggIntegrate2D equivalent (rough disk sum minus ring-mean background)
//   ProfileGaussian -> per-reflection measured-width Gaussian profile fit (the default)
//   ProfileEmpirical-> per-shell learned empirical profile fit
// The box sum is also the seed pass (Pass A) of the two profile modes, so it always runs.
//
// This is the Bragg integrator used by the pipeline (bound in MXAnalysisWithoutFPGA: the GPU
// engine when a device is present, otherwise the CPU engine). It takes a preprocessed image +
// the predicted reflections and returns the vector<Reflection> (I, sigma, bkg, partiality, ...)
// that the downstream scaling/merge consumes unchanged.
// =============================================================================

#include <cmath>
#include <cstddef>
#include <cstdint>
#include <optional>
#include <vector>

#include "../../common/BraggIntegrationSettings.h"
#include "../../common/DiffractionExperiment.h"
#include "../../common/DiffractionGeometry.h"
#include "../../common/Reflection.h"
#include "../image_preprocessing/ImagePreprocessorBuffer.h"
#include "BraggStencil.h"

namespace bragg_engine {
// Shared with both engines so the CPU and GPU paths stay numerically aligned.
constexpr int    N_SHELL = 6;                 // resolution shells for per-shell profile learning
constexpr double STRONG_I_OVER_SIGMA = 5.0;   // strong-spot threshold that seeds the profile
constexpr int    MIN_STRONG_PER_SHELL = 30;   // below this a shell falls back to the global profile
constexpr double C_CAPTURE = 2.5;             // weak-spot radial capture term (coefficient of tan^2(2theta))
// Lower bound on the background term of the Kabsch fit weights (v = max(bkg, floor) + signal). It
// guards the background ESTIMATE, not the detector: the r2..r3 ring mean of a high-angle reflection
// can come out exactly zero, and v = 0 makes the weights P^2/v diverge. A ring of n pixels cannot
// resolve a background below ~1/n (0.005..0.02 for the default r2=6/r3=10 stencil), so that is the
// scale the floor has to work at. Anything larger over-regularizes: the floor multiplies the reported
// variance by floor/bkg for every pixel below it, so the previous 1/12 inflated sigma by 1.3x at
// 0.05 ct/px and 1.7x at 0.03 - exactly where the weakest high-resolution data live. Digitisation
// noise, where a detector has it, is additive on top of the background and does not belong here.
constexpr double PIXEL_VARIANCE_FLOOR = 0.01;
// The plug-in signal term of the fit weights may lower the per-pixel variance as well as raise it,
// but not below this fraction of the background. Half-wave rectifying it (max(0, I)) instead makes
// the weights - and so the reported 1/den - respond only to upward fluctuations of a noisy intensity
// estimate, which adds ~0.4*sigma*sum(P^3)/sum(P^2)^2 to every sigma whatever the count rate.
constexpr double WEIGHT_VARIANCE_MIN_FRACTION = 0.5;

// Most background-ring pixels a block can hold for the GPU trimmed-mean sort. Shared with the CPU so
// that a ring which overflows it falls back to the plain mean in BOTH engines: the CPU sorts an
// unbounded vector and would otherwise keep trimming where the GPU had silently stopped. Note that
// an elongated ring makes this a function of resolution rather than a property of the dataset - the
// ring area grows with the elongation, so on a wide enough stencil the estimator changes at a fixed
// detector radius. The growth cap below keeps the default r2=6/r3=10 ring under the bound at any
// bandwidth; the wider stills radii can cross it, and only ever with --background-trim, which is
// off by default and kept for back compatibility.
constexpr int BKG_TRIM_MAX = 512;
// Ceiling on how far the background ring may be pushed out radially, as a multiple of r3 - so the
// outer ellipse never exceeds (1 + this) * r3. Bounds what a mis-declared bandwidth can do to the
// per-reflection bounding box, and with it the shared memory the GPU sizes from the widest ring.
constexpr float MAX_STENCIL_GROW_OVER_R3 = 2.0f;

// Guard against profile-fit runaways: on a weak / near-zero reflection the reweighted Kabsch iteration
// has no real peak to lock onto and can manufacture intensity the box sum never sees. Fall back to the
// summation (box-sum) intensity when the profile result disagrees with the summation seed by more than
// this many box-sum sigmas (a real fit agrees within counting noise, so the margin is generous).
constexpr double PROFILE_SUMMATION_MAX_NSIGMA = 10.0;
}   // namespace bragg_engine

// One reflection's extracted intensity, produced by the derived engine and turned into a
// Reflection by Finalize() (which owns the polarization correction and scale bookkeeping).
struct BraggFitResult {
    float I = 0.0f;
    float sigma = NAN;
    float bkg = 0.0f;
    float observed_x = 0.0f;   // intensity-weighted centroid (BoxSum mode only)
    float observed_y = 0.0f;
    // Variance of everything in `sigma` that is NOT the reflection's own Poisson signal (the
    // background and the error of its estimate). The merge rebuilds each partial's variance at the
    // pooled intensity and needs this term; back-deriving it as sigma^2 - I only works while that
    // identity holds exactly, which it does not for a profile fit.
    float var_bkg = 0.0f;
    bool  ok = false;
    bool  has_observed = false;
};

class BraggIntegrationEngine {
protected:
    // --- fixed configuration extracted from the experiment (see ProfileIntegrate2D) ---
    IntegratorMode mode;
    bool empirical;              // ProfileEmpirical (vs ProfileGaussian)

    size_t xpixel, ypixel, npixel;

    float r1_sq;
    float r2, r2_sq;
    float r3, r3_sq;
    int   R, G, GG;              // profile-grid half-size, edge (2R+1) and area (G*G)

    double bw_sigma;             // bandwidth sigma [dimensionless, * Rpx -> px]
    float  bkg_clip_nsigma;      // high-outlier background sigma-clip multiplier (0 = no clip)
    bool   use_ellipse;          // radially elongate the per-reflection Gaussian

    double c_radial;             // radial variance coefficient of tan^2(2theta): parallax + capture
    double F_px;                 // detector distance expressed in pixels
    float  beam_x, beam_y;
    double r_max;                // distance from the beam centre to the far corner [px]
    // Per-reflection signal/background geometry (BraggStencil.h). With k_sigma = 0 this is the fixed
    // r1 disk + r2..r3 ring the integrator has always used, bit for bit; above 0 the RING is
    // elongated radially, per reflection, by the analytic radial smear. Both engines build every
    // stencil through MakeBraggStencil, so the geometry has one definition.
    BraggStencilParams stencil;
    // Effective symmetric trimmed-mean background fraction (BraggIntegrationSettings), used only when
    // the high-side clip is switched off - the two are alternatives. 0 = plain ring mean. Read by both
    // the CPU and GPU engines.
    float  bkg_trim = 0.0f;

    // --- radial background curvature correction (BraggIntegrationSettings) ---
    // The disk and the annulus are concentric, so any background LINEAR in position cancels between
    // them; what survives is the curvature of the radial background. mean_annulus(B) - mean_disk(B)
    // of a radial B is a kernel over radial offset:
    //     bkg_error = sum_k k_diff[k] * B(r0 + k - k_off)
    // That is one short dot product per reflection and reads no pixels. Built in the constructor,
    // so bkg_radial can be flipped between images at no cost - which is what the auto mode does,
    // applying the correction only to the images whose background really is a smooth function of
    // radius (see BackgroundRadial below). It is built only when that mode, or an explicit setting,
    // could ever raise the correction; an engine that can never apply it keeps the single circular
    // kernel and never reads it.
    //
    // With a circular stencil ONE kernel serves every reflection. An elongated ring does not: its
    // radial-offset histogram depends on how far that reflection's ring was grown. So k_diff holds
    // n_kern kernels of k_len each, indexed by the growth quantized to whole pixels
    // (BraggStencilKernelIndex); n_kern is 1 when nothing is elongated, which is the old
    // single-kernel layout unchanged. The azimuthal average is kept, and still means what it did:
    // the stencil is rebuilt in the reflection's own frame at each azimuth, so it stays radially
    // aligned and what is averaged over is the sub-pixel phase of the detector grid against the
    // radius. What an elongated ring rules out is one kernel for ALL of them, not the average.
    bool  bkg_radial = false;
    bool  bkg_radial_auto = false;   // settings left it unset: decide per image from the ice score
    int   k_off = 0;                 // index of offset 0 within one kernel
    int   k_len = 0;                 // entries per kernel
    int   n_kern = 1;                // kernels in the table (1 = circular stencil)
    std::vector<float> k_diff;       // n_kern * k_len, annulus-minus-disk weight per radial offset

    // One radial-offset kernel for a ring grown by `grow` px, appended to k_diff. Kept out of line
    // so the circular and elongated cases cannot drift apart. The signal disk does not change with
    // the growth, so its histogram is built on the first call and reused.
    void BuildRadialKernel(float grow);
    std::vector<double> hist_disk;
    double sum_disk = 0.0;

    DiffractionGeometry geom;    // kept for the per-reflection polarization correction
    std::optional<float> polarization;

    // Assemble output reflections from the per-reflection fit results (polarization + scale corr).
    std::vector<Reflection> Finalize(const std::vector<Reflection> &predicted, size_t npredicted,
                                     const std::vector<BraggFitResult> &results,
                                     int64_t image_number) const;

public:
    explicit BraggIntegrationEngine(const DiffractionExperiment &experiment);
    virtual ~BraggIntegrationEngine() = default;

    // predicted[0..npredicted) are the reflections to extract; image is the preprocessed int32
    // frame (image.size() == npixel). Returns only the observed reflections.
    virtual std::vector<Reflection> Run(const ImagePreprocessorBuffer &image,
                                        const std::vector<Reflection> &predicted, size_t npredicted,
                                        int64_t image_number) = 0;

    // Turn the radial background correction on or off for the images that follow. The caller owns
    // the decision; in the auto mode the analysis sets it per image from that image's ice score.
    void BackgroundRadial(bool on) { bkg_radial = on; }
    [[nodiscard]] bool IsBackgroundRadialAuto() const { return bkg_radial_auto; }
};