// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute // 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 (I, sigma, bkg, partiality, ...) // that the downstream scaling/merge consumes unchanged. // ============================================================================= #include #include #include #include #include #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.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. 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; } // 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; // --- 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; // --- 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 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 hist_disk; double sum_disk = 0.0; DiffractionGeometry geom; // kept for the per-reflection polarization correction std::optional polarization; // Assemble output reflections from the per-reflection fit results (polarization + scale corr). std::vector Finalize(const std::vector &predicted, size_t npredicted, const std::vector &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 Run(const ImagePreprocessorBuffer &image, const std::vector &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; } };