Jungfraujoch/image_analysis/bragg_integration/BraggIntegrationEngine.cpp

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

#include "BraggIntegrationEngine.h"

#include <algorithm>
#include <cmath>
#include <map>
#include <mutex>
#include <numeric>
#include <string>
#include <tuple>

#include "../../common/JFJochMath.h"   // PI (M_PI is not standard, and MSVC does not define it)
#include "../SensorAbsorption.h"
#include "RingBackground.h"

namespace {

// Radial parallax broadening as the coefficient of tan^2(2theta), i.e. Var(z)/pixel^2 [px^2].
// Copied verbatim from ProfileIntegrate2D: a photon converts at a random depth z (exponential,
// attenuation length L, truncated at the sensor thickness), shifting the recorded spot radially by
// z*tan(2theta). L comes from the tabulated NIST attenuation coefficients (SensorAbsorption.h); the
// lambda^3 approximation this used before is within 0.2% for silicon above 10 keV but overstates
// the attenuation length of CdTe by up to a factor of two, and by six above the Cd K edge, which
// made this variance 1.9x too large on 750 um CdTe data.
double parallax_var_px2(const std::string &material, double thickness_um, double lambda_A, double pixel_um) {
    if (!(thickness_um > 0.0) || !(pixel_um > 0.0) || !(lambda_A > 0.0))
        return 0.0;
    const double L = sensor_absorption::AttenuationLength_um(material, lambda_A);
    return sensor_absorption::ConversionDepthVariance_um2(L, thickness_um) / (pixel_um * pixel_um);
}

// The radial-offset kernels below are a pure function of these six numbers, and one engine is built
// per worker per pass - 96 of them on a two-pass run - so the table was built 96 times over from the
// same inputs. Build it once and let the rest copy it; it is a few hundred floats. Two workers can
// still race to build the same table, which costs nothing but the second build: the values are
// identical, and emplace keeps whichever arrived first.
struct RadialKernelKey {
    float r1_sq, r2, r3;
    int n_kern, k_off, k_len;

    bool operator<(const RadialKernelKey &o) const {
        return std::tie(r1_sq, r2, r3, n_kern, k_off, k_len)
             < std::tie(o.r1_sq, o.r2, o.r3, o.n_kern, o.k_off, o.k_len);
    }
};

std::mutex radial_kernel_mutex;
std::map<RadialKernelKey, std::vector<float>> radial_kernel_cache;

}   // namespace

BraggIntegrationEngine::BraggIntegrationEngine(const DiffractionExperiment &experiment)
    : geom(experiment.GetDiffractionGeometry()), modules(experiment.GetModules()),
      plane_bounds(modules.PlaneBounds(experiment.GetXPixelsNum(), experiment.GetYPixelsNum())) {
    const auto settings = experiment.GetBraggIntegrationSettings();
    const auto &det = experiment.GetDetectorSetup();

    mode = settings.GetIntegrator();
    empirical = mode == IntegratorMode::ProfileEmpirical;

    // Same frame as the reflections' predicted_x/predicted_y and the ImagePreprocessorBuffer that
    // feeds this engine (MXAnalysisWithoutFPGA sizes that buffer to GetPixelsNum()).
    xpixel = experiment.GetXPixelsNum();
    ypixel = experiment.GetYPixelsNum();
    npixel = experiment.GetPixelsNum();

    r1_sq = settings.GetR1() * settings.GetR1();
    r2 = settings.GetR2();
    r2_sq = r2 * r2;
    r3 = settings.GetR3();
    r3_sq = r3 * r3;
    R = static_cast<int>(std::ceil(r2));
    G = 2 * R + 1;
    GG = G * G;

    // The X-ray bandwidth enters ONE place: it smears a reflection radially by bw_sigma * Rpx, which
    // the per-reflection Gaussian carries as part of its radial variance. It is not a mode switch -
    // the background estimator and the parallax/capture term below are the same whatever the beam is.
    bw_sigma = experiment.GetBandwidthFWHM().value_or(0.0f) / 2.3548f;

    const double c_par = parallax_var_px2(det.GetSensorMaterial(), det.GetSensorThickness_um(),
                                          geom.GetWavelength_A(), geom.GetPixelSize_mm() * 1000.0);
    c_radial = c_par + bragg_engine::C_CAPTURE;
    F_px = geom.GetDetectorDistance_mm() / std::max(1e-6f, geom.GetPixelSize_mm());
    beam_x = geom.GetBeamX_pxl();
    beam_y = geom.GetBeamY_pxl();
    use_ellipse = !empirical;

    // Per-reflection signal/background geometry: the ring elongated radially by k_sigma times the
    // beam's own radial streak, capped. k_sigma = 0 is the fixed circular stencil, bit for bit, and
    // so is any monochromatic beam, where the streak is zero.
    stencil.beam_x = beam_x;
    stencil.beam_y = beam_y;
    stencil.r2 = r2;
    stencil.r3 = r3;
    stencil.bw_sigma = static_cast<float>(bw_sigma);
    stencil.k_sigma = settings.GetStencilKSigma();
    stencil.max_grow = bragg_engine::MAX_STENCIL_GROW_OVER_R3 * r3;
    // The measured spot footprint (SpotFootprint.h). It moves the ring only where a spot outgrows the
    // r1 disk, so a pattern of compact spots integrates exactly as without it.
    static_assert(SpotFootprint::MAX_BINS <= BRAGG_FOOTPRINT_MAX_BINS);
    stencil.r1 = settings.GetR1();
    const auto &fp = settings.GetFootprint();
    if (!empirical && !fp.empty()) {
        stencil.fp_n = static_cast<int>(fp.sigma_rad.size());
        stencil.fp_bin_px = fp.bin_px;
        for (int i = 0; i < stencil.fp_n; ++i) {
            stencil.fp_sigma_rad[i] = fp.sigma_rad[i];
            stencil.fp_sigma_tan[i] = fp.sigma_tan[i];
        }
    }

    // Robust background ring, one estimator or the other (see BraggIntegrationSettings): a high-side
    // sigma-clip (rugnux --background-clip, the default) or, when the clip is switched off, a
    // symmetric trimmed mean (rugnux --background-trim). The caller owns the choice - the engine no
    // longer overrides it for broadband data.
    bkg_clip_nsigma = settings.GetBackgroundClipNSigma();
    bkg_pixel_gain = settings.GetBackgroundPixelGain();
    pixel_gain = settings.GetPixelGain();
    bkg_trim = bkg_clip_nsigma > 0.0f ? 0.0f : settings.GetBackgroundTrimFraction();

    // Overlap treatment. Ownership is decided out to the fit grid's half size, which is where the
    // profile fit reads pixels; beyond it a pixel that nobody claims is this reflection's own.
    // Excluding the shared pixels needs a profile to renormalise, so it cannot act on a box sum -
    // drop it to Off there rather than build an owner map nothing will read.
    overlap = settings.GetOverlap();
    if (overlap == OverlapMode::Exclude && mode == IntegratorMode::BoxSum)
        overlap = OverlapMode::Off;
    overlap_min_peak = settings.GetOverlapMinPeak();
    keep_overloads = !experiment.IsRotationIndexing();
    saturation_limit = static_cast<double>(experiment.GetSaturationLimit());
    claim = static_cast<float>(R);
    inv_claim = 1.0f / claim;

    // Radial-offset kernels for the background curvature correction. A stencil pixel at (dx, dy)
    // sits at radial offset dx*cos(phi) + dy*sin(phi) from the reflection, where phi is the
    // reflection's azimuth; averaging over phi makes the kernels position-independent, which is
    // exact to the extent the stencil is small against the reflection's radius (r3 = 10 px vs
    // hundreds). k_diff is the annulus histogram minus the disk histogram, each normalised, so
    // dot(k_diff, B) is directly mean_annulus(B) - mean_disk(B).
    // Unset = auto: start off, and let the analysis raise it per image where the ice score says the
    // background really is radial. An engine nobody drives therefore never applies the correction.
    const auto radial = settings.GetBackgroundRadialCorrection();
    bkg_radial_auto = !radial.has_value();
    bkg_radial = radial.value_or(false);

    // The table spans zero growth up to whatever the widest reflection on this detector reaches, one
    // kernel per pixel of growth; with nothing elongated a single kernel is all there is, which is
    // the layout and the values of every build before the stencil existed. It is built only when the
    // correction can ever run - the rows are not cheap, and nothing may read them otherwise:
    // bkg_radial is raised after construction only by the auto mode (MXAnalysisWithoutFPGA), which
    // requires bkg_radial_auto, and the GPU allocates its curve buffers under the same condition.
    // n_kern is the largest row BraggStencilKernelIndex can select, plus one.
    r_max = std::hypot(std::max<double>(beam_x - plane_bounds.x0, plane_bounds.x1 - beam_x),
                       std::max<double>(beam_y - plane_bounds.y0, plane_bounds.y1 - beam_y));
    bkg_radial_built = bkg_radial || bkg_radial_auto;
    const float grow_max = bkg_radial_built ? BraggStencilMaxGrow_px(static_cast<float>(r_max), stencil)
                                            : 0.0f;
    n_kern = static_cast<int>(std::lround(grow_max)) + 1;
    // Every row must fit: the last one is built at grow = n_kern - 1, which rounding can put just
    // above grow_max.
    k_off = static_cast<int>(std::ceil(r3 + std::max<double>(grow_max, n_kern - 1))) + 1;
    k_len = 2 * k_off + 1;
    const RadialKernelKey kernel_key{r1_sq, r2, r3, n_kern, k_off, k_len};
    {
        const std::lock_guard lock(radial_kernel_mutex);
        if (const auto it = radial_kernel_cache.find(kernel_key); it != radial_kernel_cache.end())
            k_diff = it->second;
    }
    if (k_diff.empty()) {
        k_diff.reserve(static_cast<size_t>(n_kern) * k_len);
        for (int j = 0; j < n_kern; ++j)
            BuildRadialKernel(static_cast<float>(j));
        const std::lock_guard lock(radial_kernel_mutex);
        radial_kernel_cache.emplace(kernel_key, k_diff);
    }

    polarization = experiment.GetPolarizationFactor();
}

void BraggIntegrationEngine::BuildRadialKernel(float grow) {
    // Histogram the stencil over radial offset, averaged over azimuth so the kernel does not depend
    // on where the reflection sits. The average is over the SUB-PIXEL PHASE of the detector grid
    // against the radial direction, not over the stencil's own orientation: the stencil is built in
    // the reflection's frame at each azimuth, so an elongated one stays aligned with the radius, as
    // it is on the detector. k_diff is the annulus histogram minus the disk histogram, each
    // normalised, so dot(k_diff, B) is directly mean_annulus(B) - mean_disk(B).
    // The signal disk is a circle whatever the ring does, so its histogram is the same for every
    // kernel in the table - build it once.
    const bool first = hist_disk.empty();
    if (first)
        hist_disk.assign(k_len, 0.0);
    std::vector<double> hist_ann(k_len, 0.0);
    constexpr int n_phi = 512;
    const int span = static_cast<int>(std::ceil(r3 + grow)) + 1;
    const float si = r2 / (r2 + grow), so = r3 / (r3 + grow);
    const double q_in = 1.0 - static_cast<double>(si) * si;
    const double q_out = 1.0 - static_cast<double>(so) * so;
    for (int p = 0; p < n_phi; ++p) {
        const double phi = 2.0 * PI * p / n_phi, cp = std::cos(phi), sp = std::sin(phi);
        for (int dy = -span; dy <= span; ++dy)
            for (int dx = -span; dx <= span; ++dx) {
                const double d2 = static_cast<double>(dx) * dx + static_cast<double>(dy) * dy;
                const double rad = dx * cp + dy * sp;
                const int k = k_off + static_cast<int>(std::lround(rad));
                if (k < 0 || k >= k_len)
                    continue;
                const double rad2 = rad * rad;
                if (d2 < r1_sq) {
                    if (first) hist_disk[k] += 1.0;
                } else if (d2 - q_in * rad2 >= r2_sq && d2 - q_out * rad2 < r3_sq) {
                    hist_ann[k] += 1.0;
                }
            }
    }
    if (first) sum_disk = std::accumulate(hist_disk.begin(), hist_disk.end(), 0.0);
    const double sd = sum_disk;
    const double sa = std::accumulate(hist_ann.begin(), hist_ann.end(), 0.0);
    for (int k = 0; k < k_len; ++k)
        k_diff.push_back(static_cast<float>(hist_ann[k] / sa - hist_disk[k] / sd));
}

void BraggIntegrationEngine::RadialBackgroundFromRings(const std::vector<float> &ring_bkg, int32_t q_bins,
                                                       float low_q, float dq) {
    radial_curve.clear();
    const RingBackground rings(ring_bkg, q_bins, low_q, dq);
    if (rings.empty())
        return;
    // The same bins as the engines' own curve: one per pixel of distance from the beam, out to the far
    // corner.
    const int n = static_cast<int>(std::ceil(r_max)) + 2;
    const double two_k = 4.0 * PI / geom.GetWavelength_A();
    radial_curve.resize(n);
    for (int r = 0; r < n; ++r)
        radial_curve[r] = static_cast<float>(rings.At(two_k * std::sin(0.5 * std::atan(r / F_px))));
}

std::vector<BraggPanel> BraggIntegrationEngine::Panels(const std::vector<Reflection> &predicted,
                                                       size_t npredicted) const {
    std::vector<BraggPanel> out(npredicted);
    for (size_t i = 0; i < npredicted; ++i) {
        const float x = predicted[i].predicted_x, y = predicted[i].predicted_y;
        out[i] = MakeBraggPanel(modules, modules.ModuleOfImagePixel(x, y), static_cast<int64_t>(xpixel),
                                static_cast<int64_t>(ypixel), x, y);
    }
    return out;
}

std::vector<Reflection> BraggIntegrationEngine::Finalize(const std::vector<Reflection> &predicted,
                                                         size_t npredicted,
                                                         const std::vector<BraggPanel> &panels,
                                                         const std::vector<BraggFitResult> &results,
                                                         int64_t image_number) {
    // Sized to what is kept, not to what was predicted: this vector is the one every image holds to the
    // end of the pass, and a background ring lost to the neighbours drops a reflection here - on a
    // dense pattern more than half of them, which a reserve of npredicted would carry as dead capacity.
    size_t n_ok = 0;
    for (size_t i = 0; i < npredicted; ++i)
        n_ok += results[i].ok ? 1 : 0;
    std::vector<Reflection> out;
    out.reserve(n_ok);
    for (size_t i = 0; i < npredicted; ++i) {
        const auto &fr = results[i];
        if (!fr.ok)
            continue;
        Reflection refl = predicted[i];
        // In photons: the engine measured in ADU, with variances already taken at the gain.
        refl.I = fr.I / pixel_gain;
        refl.sigma = fr.sigma / pixel_gain;
        refl.bkg = fr.bkg / pixel_gain;
        refl.var_bkg = fr.var_bkg / (pixel_gain * pixel_gain);
        if (fr.has_observed) {
            refl.observed_x = fr.observed_x;
            refl.observed_y = fr.observed_y;
        }
        // has_observed is set only for a signal disk every pixel of which was readable.
        refl.clipped = !fr.has_observed;
        refl.overloaded = fr.overloaded;
        // A rotation merge drops the whole rocking event of an overloaded reflection. A still has no
        // rocking event: its overloaded observations are the fully recorded ones of the strongest
        // reflections, so dropping them leaves those reflections to their partial observations, whose
        // partiality correction then sets the merged intensity - and it reads low. Measured on clipped
        // CCD stills, the merged low-resolution intensities came out at 0.4-0.6x of the deposited ones.
        // So a still keeps the observation as a lower bound: each saturated pixel counted at the
        // saturation limit. (A fit to the unsaturated wings was tried and read lower still: the
        // profile learned from unsaturated spots does not describe the wings of a clipped one.)
        // That takes a pixel AT the limit: an observation overloaded only by pixels unreadable on
        // this frame has nothing to count at the limit, and counting those pixels as zero is no
        // lower bound of anything - it stays overloaded and is dropped. (A pixel stuck at the limit
        // is no clipped spot either; the pre-scan masks it, HotPixelFinder.)
        if (fr.overloaded && keep_overloads && fr.saturated > 0) {
            refl.I = fr.I + static_cast<float>(fr.saturated * (saturation_limit - fr.bkg));
            refl.sigma = std::sqrt(fr.var_bkg + std::max(refl.I, 1.0f));
            refl.overloaded = false;
            ++counts.overloads_kept;
        }
        refl.observed = true;
        if (polarization)
            refl.prescaling_corr /= geom.CalcAzIntPolarizationCorr(panels[i].plane_x, panels[i].plane_y,
                                                                   polarization.value());
        refl.image_scale_corr = refl.prescaling_corr * refl.qe_corr * refl.flight_corr / refl.partiality;
        refl.image_number = static_cast<float>(image_number);
        out.push_back(refl);
    }
    return out;
}