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 <algorithm>
#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 "BraggPanel.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.003..0.02 for the default r2=6/r3=13 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.
// In photons: on frames in ADU it is taken at the gain, as the background it bounds is.
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=13 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;

// MINPK keeps a reflection while enough of its expected profile is readable. It says nothing about
// WHERE the unreadable part is, and the two are not the same question. Renormalising over the pixels
// that remain is unbiased only while the profile MODEL is exact: lose the peak and the amplitude is
// set by the wings alone, so the estimate stops being a measurement of the reflection and becomes a
// measurement of how well the fitted shape describes it. That holds whatever made the pixel
// unreadable, so the rule below cuts on any of them. The overload is the case that motivated it and
// the one that biases hardest - a pixel over the detector's range went missing BECAUSE the reflection
// was bright, so the loss is concentrated on the strongest low-resolution reflections, which are the
// largest terms of R_meas, and there the fit reads -50% against the symmetry mates. So no unreadable
// pixel may carry more than this fraction of the profile's own peak value.
//
// A fraction of the peak rather than a radius in pixels, because the peak is as wide as the spot: for
// a Gaussian the cut sits at sqrt(-2 ln f) sigma, i.e. 0.46 sigma here, which is the peak pixel alone
// where sigma is 0.8 px and the crest of the ridge where it is 2.4 px or a bandwidth streak. It also
// needs nothing the fit does not already have, so it costs one max-reduction in the loop that
// measures the readable fraction, and it applies to the learned empirical profile unchanged.
constexpr double MINPK_MAX_MISSING_PEAK = 0.9;

// A predicted reflection sets its r2 region aside from its neighbours' background rings only where it
// puts at least this fraction of its flux on the frame. Prediction reaches +-4 sigma of the rocking
// curve, so on a finely sliced frame most predictions are the tails of reflections recorded on the
// frames either side: on a dense pattern those fill every ring while the frame shows nothing there,
// and a reflection whose ring they take is dropped. A still, whose partiality is 1, masks as before.
constexpr float NEIGHBOUR_MASK_MIN_PARTIALITY = 0.05f;
// The two marks a neighbour leaves in the reflection mask: its signal region holds a tail of a rocking
// curve only, or flux the ring has to be kept clear of. The larger wins where two regions meet.
constexpr uint8_t MASK_TAIL = 1;
constexpr uint8_t MASK_FLUX = 2;
}   // namespace bragg_engine

// How often the integrator silently dropped a reflection, or silently declined its own fit, over
// every image an engine has run. Both engines keep the same two counts, so a caller sees the same
// numbers whichever one it got. Nothing inside the engine reads them: they exist so a caller that
// CHOSE the stencil can find out what that choice cost, which no quantity available before
// integration measures. rugnux reads them out of its first pass (see Rugnux::RunAllPasses).
struct BraggIntegrationCounts {
    uint64_t predicted = 0;          // reflections offered to the engine
    uint64_t bkg_starved = 0;        // dropped whole: the r2..r3 ring kept 5 or fewer clean pixels
    // The rings the NEIGHBOURS would starve: more than five readable pixels, but five or fewer that no
    // predicted neighbour's signal region covers - counting every prediction, the tails of rocking
    // curves the background does not exclude (NEIGHBOUR_MASK_MIN_PARTIALITY) included, so it measures
    // the pattern's density whatever the ring then keeps. This is the count that answers
    // "is the aperture too wide for this pattern", because a ring the detector alone starves is module
    // gaps, the beam stop and the resolution mask - a property of the detector that a wider r1 does not
    // change.
    // Measured over the battery, that floor reaches 2.3% of all reflections while the widening that
    // costs data adds 4.1 percentage points on top of a floor of 0.02%.
    uint64_t bkg_starved_by_neighbour = 0;
    uint64_t profile_fallback = 0;   // the profile fit disagreed with the box sum; the box sum was kept
    // Signal disks a nearer neighbour's prediction takes pixels from (the overlap ownership, so counted
    // only while an overlap treatment is on), and of those the ones left with less than
    // overlap_min_peak of the disk their own - what --overlap reject drops, and XDS's MINPK. Reported
    // only (the crowding of the pattern); nothing reads them.
    uint64_t overlapped = 0;
    uint64_t overlapped_major = 0;
    // Overloaded reflections of stills, kept in the merge as a lower bound (BraggIntegrationEngine::Finalize).
    uint64_t overloads_kept = 0;
};

// 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;
    bool  overloaded = false;   // a signal-disk pixel was saturated (Reflection::overloaded)
    int   saturated = 0;        // signal-disk pixels at the saturation limit, left out of I
};

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)
    // ADU per photon (BraggIntegrationSettings::pixel_gain): every variance the engine forms is taken
    // at it, the background clip and the variance floors are a photon wide in it, and Finalize converts
    // the reflections to photons with it. 1 changes nothing, bit for bit.
    float  pixel_gain = 1.0f;
    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;

    // --- overlap treatment (BraggIntegrationSettings, rugnux --overlap) ---
    // Off leaves the signal region exactly as it always was. Reject and Exclude both need the owner
    // map (BraggStencil.h): which of two touching reflections a shared pixel belongs to. `claim` is
    // how far a reflection claims pixels - the fit grid's half size, so ownership is decided
    // everywhere the fit looks.
    OverlapMode overlap = OverlapMode::Off;
    float overlap_min_peak = 0.0f;   // Reject: least clean profile fraction that is kept
    float claim = 0.0f, inv_claim = 0.0f;

    // Overloaded reflections of stills are kept as a lower bound (Finalize), with every saturated pixel
    // counted at the saturation limit, in the units of the preprocessed image.
    bool keep_overloads = false;
    double saturation_limit = 0.0;

    // --- 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
    bool  bkg_radial_built = false;  // the kernel table was sized for the rings this engine uses
    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
    // The radial background B the kernels are applied to, one value per pixel of distance from the
    // beam, where the caller handed one over (RadialBackgroundFromRings). Empty: the curve is
    // accumulated from the clipped ring pixels of the predicted reflections themselves.
    std::vector<float> radial_curve;

    // 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;

    // The panels of the image (DetectorModules), empty for a flat detector, and the box on the plane
    // they span - where the far corner is, for anything sized by the largest radius on the detector.
    DetectorModules modules;
    PlaneBox plane_bounds;
    // The panel of each prediction, found from its position in the image (BraggPanel).
    std::vector<BraggPanel> Panels(const std::vector<Reflection> &predicted, size_t npredicted) const;

    // Accumulated over every image this engine has run. An engine belongs to one worker thread, so
    // these are plain counters and the caller sums over the engines it made.
    BraggIntegrationCounts counts;

    // 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<BraggPanel> &panels,
                                     const std::vector<BraggFitResult> &results,
                                     int64_t image_number);

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. Positions in and out are
    // the image's: a reflection is on the panel whose box holds its predicted position.
    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.
    //
    // It cannot be raised on an engine the constructor did not build the kernel table for: that
    // table would hold a single CIRCULAR kernel for rings that may be elongated, and on the GPU the
    // radial buffers were never allocated, so the CPU would correct and the GPU would not.
    void BackgroundRadial(bool on) { bkg_radial = on && bkg_radial_built; }
    [[nodiscard]] bool IsBackgroundRadialAuto() const { return bkg_radial_auto; }

    // The radial background curve for the images that follow, from a peak-excluded background per
    // resolution ring of the whole detector (the adaptive spot finder's GetRingBackground): q bin i
    // holds the ring at low_q + (i + 0.5) * dq, entries past the first q_bins are further azimuthal
    // bins and are averaged in, and a ring with no value (NAN) is bridged by its neighbours. On a
    // serial frame the predicted reflections' own ring pixels are a sparse sample of the same curve -
    // a few dozen reflections, against every pixel of the frame here. An empty profile goes back to
    // accumulating the curve from the reflections.
    void RadialBackgroundFromRings(const std::vector<float> &ring_bkg, int32_t q_bins, float low_q, float dq);
    // Where that curve is in use (the correction on), the ring's high-side clip follows it as well: a ring
    // pixel at radial offset `rad` from a reflection at radius r0 is clipped against the threshold plus
    // the curve's rise from r0 to it, rounded to whole counts so that every engine keeps the same pixels.
    // On a background falling steeply away from the beam the flat threshold otherwise cuts the beam-side
    // half of the ring as outliers, and the background reads low. 0 where no curve is in use.
    [[nodiscard]] int32_t RadialClipOffset(float r0, float rad) const {
        if (!bkg_radial || radial_curve.empty())
            return 0;
        const int n = static_cast<int>(radial_curve.size());
        const int b = std::clamp(static_cast<int>(std::lround(r0 + rad)), 0, n - 1);
        const int b0 = std::clamp(static_cast<int>(std::lround(r0)), 0, n - 1);
        return static_cast<int32_t>(std::lround(radial_curve[b] - radial_curve[b0]));
    }
    [[nodiscard]] const std::vector<float> &GetRadialBackground() const { return radial_curve; }

    // What this engine has counted since it was built (BraggIntegrationCounts). The GPU engine keeps
    // its counts on the device and brings them back here, so this is a synchronising call - ask for it
    // once a pass, not once an image.
    [[nodiscard]] virtual BraggIntegrationCounts Counts() const { return counts; }
};