Jungfraujoch/image_analysis/bragg_integration/BraggStencil.h

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

#pragma once

// =============================================================================
// BraggStencil - the per-reflection signal/background geometry, shared by both engines
// =============================================================================
//
// The integrator classifies every pixel of a reflection's neighbourhood as signal (inside r1),
// background (the r2..r3 ring) or neither. Those three radii used to be plain circles of a fixed
// pixel size, which is wrong for a spot that is not round: a bandwidth streaks a reflection
// radially by sigma_bw = bw_sigma*Rpx, so at high resolution the r2..r3 ring sits on the
// reflection's own tails and measures them as background.
//
// So the RING is an ellipse, elongated along the beam->reflection direction by k_sigma times that
// streak, and left alone tangentially. The SIGNAL disk stays a circle: r1 governs n_inner, var_bkg,
// the box-sum acceptance gate (all-or-nothing `n_inner_valid == n_inner`, so growing it rejects every
// reflection with one bad pixel anywhere along a long streak) and the domain the profile WIDTH is
// learned over. In the default gaussian mode it does not set the intensity - the fit grid, ceil(r2),
// does.
//
// What the fixed circular disk loses is flux, and the loss is NOT a function of resolution alone:
// measured, it carries a directional component worth several A^2 with a definite principal axis. Nor
// is there anything to absorb it. There is no per-shell scale in the merge, and there cannot be a
// useful one: every scale is fitted against a reference built from a reflection's own symmetry
// equivalents, and equivalents share s^2 exactly, so any function of s^2 lies in the exact null space
// of the whole scaling model. The isotropic part is instead degenerate with the Wilson B and is
// silently reported as part of it - measured across a radius sweep, 0.5 A^2 on sharp strong data and
// up to ~10 A^2 on weak wide-spot data. Most of the loss is matched by a proportional sigma, so it
// moves no CC1/2 and no R_meas.
//
// The width is the BANDWIDTH STREAK ALONE, and deliberately not the profile's full radial variance
// (which also carries a sensor parallax term and a weak-spot capture term). Those two were measured
// on the rotation battery: they widen the ring on monochromatic data, where they are the only terms
// there are, and that neither helped the two crystals with clean high-resolution shells nor left the
// two weak ones alone. The bandwidth streak, by contrast, is a measured elongation of the recorded
// spot - principal axis along the radius to within a couple of degrees, azimuth-independent - and it
// is the term that carries the pink-beam case. Keeping only it also makes the whole feature exactly
// inert on a monochromatic beam, where bw_sigma is 0 and so is the growth.
//
// Rather than evaluate an ellipse, each pixel's squared distance has the radial part scaled down:
// with rad the pixel's radial offset and d2 its plain squared distance,
//
//     d2_inner = d2 - q_in *rad^2   tested against r2^2   (q_in  = 1 - (r2/(r2+grow))^2)
//     d2_outer = d2 - q_out*rad^2   tested against r3^2   (q_out = 1 - (r3/(r3+grow))^2)
//
// which is the ellipse (radial semi-axis r+grow, tangential semi-axis r) written so that grow = 0
// gives q = 0 and hence d2_inner == d2_outer == d2 EXACTLY, in floating point, not merely to within
// rounding. k_sigma = 0 therefore reproduces the old circular stencil bit for bit.
//
// This header is compiled by the host compiler, by nvcc and by the Metal compiler, so the geometry has
// exactly one definition: the engines classify by the same rule, built from the same inputs. The CUDA
// engine is not bit-identical to the CPU one - nvcc contracts a*b+c into an FMA and the host baseline
// has no FMA instruction to contract into - so a pixel within a rounding of a boundary can still fall
// either way, which is what its parity test's tolerances are for. The Metal engine is: neither the host
// nor the Metal compiler contracts anything, and with fast math off Metal rounds sqrt and division
// correctly, as the CPU does (BraggIntegrationEngineMetal.mm).
// =============================================================================

#ifndef __METAL_VERSION__
#include <cmath>
#include <cstdint>
#endif

#ifdef __CUDACC__
#define BRAGG_STENCIL_HD __host__ __device__ inline
#else
#define BRAGG_STENCIL_HD inline
#endif

// Metal keeps a program-scope constant in its constant address space, and a reference names the space
// it points into (here always the thread's own), so both say so.
#ifdef __METAL_VERSION__
#define BRAGG_STENCIL_CONSTANT constant
#define BRAGG_STENCIL_THREAD thread
#else
#define BRAGG_STENCIL_CONSTANT
#define BRAGG_STENCIL_THREAD
#endif

// The measured spot footprint (SpotFootprint.h): the radial and tangential standard deviations of a
// spot [px], tabulated against distance from the beam centre in bins of fp_bin_px. Held here, by value,
// because both engines build every stencil from these parameters - the GPU gets the table with the
// kernel arguments and needs no buffer of its own.
BRAGG_STENCIL_CONSTANT constexpr int BRAGG_FOOTPRINT_MAX_BINS = 16;
// A spot whose footprint reaches this many of its standard deviations past r1 has outgrown the disk, and
// gets the footprint stencil below.
BRAGG_STENCIL_CONSTANT constexpr float BRAGG_FOOTPRINT_NSIGMA = 3.0f;
// How far that stencil reaches, in the footprint's standard deviations: the summation ellipse ends and the
// background ring starts there, and the profile is fitted out to it. Further than the 3 sigma that decides
// whether a spot outgrew the disk, because a wide spot is not a Gaussian: mosaic streaks and diffuse halos
// put flux past 3 sigma, and a ring starting there reads it as background. Scored with XDS's own
// corrections applied to both integrations, so only the integration differs, 4 sigma closed about a third
// of the gap to XDS's 3D summation on the haloed small-molecule set and was neutral on the streaked one;
// 5 sigma was no better than 4.
BRAGG_STENCIL_CONSTANT constexpr float BRAGG_FOOTPRINT_REACH = 4.0f;

// Fixed per-experiment inputs to the stencil law (mirrors BraggIntegrationEngine's members).
struct BraggStencilParams {
    float beam_x = 0.0f, beam_y = 0.0f;
    float r1 = 4.0f;
    float r2 = 6.0f, r3 = 10.0f;
    float bw_sigma = 0.0f;    // radial streak per pixel of radius (bandwidth sigma, dimensionless)
    float k_sigma = 0.0f;     // radial sigmas to push the ring out by; 0 = the old circular stencil
    float max_grow = 0.0f;    // hard cap on the radial growth [px]; <= 0 means uncapped
    int   fp_n = 0;           // bins of the measured footprint; 0 = none measured
    float fp_bin_px = 0.0f;
    float fp_sigma_rad[BRAGG_FOOTPRINT_MAX_BINS] = {};
    float fp_sigma_tan[BRAGG_FOOTPRINT_MAX_BINS] = {};
};

// One reflection's stencil, in its own radial/tangential frame.
struct BraggStencil {
    float ux = 1.0f, uy = 0.0f;      // unit vector beam -> reflection
    float r0 = 0.0f;                 // distance from the beam centre [px]
    float grow = 0.0f;               // radial growth of the ring [px] (0 = circular)
    float grow_tan = 0.0f;           // tangential growth of the ring [px] (footprint only)
    float q_in = 0.0f, q_out = 0.0f; // radial shrink coefficients (0 = circular)
    float qt_in = 0.0f, qt_out = 0.0f; // tangential shrink coefficients (0 = circular)
    // The measured footprint where the spot outgrows the r1 disk (variances [px^2]), else 0: the
    // profile fit takes its Gaussian at least this wide.
    float fp_s2r = 0.0f, fp_s2t = 0.0f;
    // 1 / (BRAGG_FOOTPRINT_REACH * sigma)^2 along and across the radius: the footprint ellipse a wide
    // spot is SUMMED over, beside the r1 disk (BraggStencilInSignal).
    float core_r = 0.0f, core_t = 0.0f;
    float ex_in = 0.0f, ey_in = 0.0f;   // axis-aligned half-extent of the inner (r2) ellipse
    float ex_out = 0.0f, ey_out = 0.0f; // axis-aligned half-extent of the outer (r3) ellipse
};

// Radial growth of the ring at Rpx, in pixels: k_sigma times the bandwidth streak, capped. Shared
// so the stencil and the radial-background kernel table agree on what "grown by this much" means.
BRAGG_STENCIL_HD float BraggStencilGrow_px(float Rpx, const BRAGG_STENCIL_THREAD BraggStencilParams &p) {
    if (!(p.k_sigma > 0.0f)) return 0.0f;
    // The cap bounds what a mis-declared bandwidth can do to the bounding box. A caller that
    // forgets to set it gets the uncapped growth rather than silently no growth at all.
    const float grow = p.k_sigma * p.bw_sigma * Rpx;
    return (p.max_grow > 0.0f && grow > p.max_grow) ? p.max_grow : grow;
}

// The footprint at Rpx: linear between bin centres, flat beyond the first and the last.
BRAGG_STENCIL_HD void BraggFootprintAt(float Rpx, const BRAGG_STENCIL_THREAD BraggStencilParams &p,
                                        BRAGG_STENCIL_THREAD float &s_rad, BRAGG_STENCIL_THREAD float &s_tan) {
    const float t = Rpx / p.fp_bin_px - 0.5f;
    int i = (int) floorf(t);
    float f = t - (float) i;
    if (i < 0) { i = 0; f = 0.0f; }
    if (i >= p.fp_n - 1) { i = p.fp_n - 1; f = 0.0f; }
    const int j = i + 1 < p.fp_n ? i + 1 : i;
    s_rad = p.fp_sigma_rad[i] + f * (p.fp_sigma_rad[j] - p.fp_sigma_rad[i]);
    s_tan = p.fp_sigma_tan[i] + f * (p.fp_sigma_tan[j] - p.fp_sigma_tan[i]);
}

// Growth of the ring beyond r2 that the footprint asks for along one axis: none while the spot fits
// the r1 disk - there the disk resolves the spot and the stencil stays what it always was - and
// otherwise out to BRAGG_FOOTPRINT_REACH of it, capped as the bandwidth growth is.
BRAGG_STENCIL_HD float BraggFootprintGrow_px(float sigma, const BRAGG_STENCIL_THREAD BraggStencilParams &p) {
    const float g = BRAGG_FOOTPRINT_REACH * sigma - p.r2;
    if (!(g > 0.0f)) return 0.0f;
    return (p.max_grow > 0.0f && g > p.max_grow) ? p.max_grow : g;
}

// The largest growth any reflection on a detector reaching r_max can get, radially - what the
// buffers that hold a ring are sized by.
BRAGG_STENCIL_HD float BraggStencilMaxGrow_px(float r_max, const BRAGG_STENCIL_THREAD BraggStencilParams &p) {
    float g = BraggStencilGrow_px(r_max, p);
    for (int i = 0; i < p.fp_n; ++i) {
        const float s = fmaxf(p.fp_sigma_rad[i], p.fp_sigma_tan[i]);
        if (BRAGG_FOOTPRINT_NSIGMA * s > p.r1) g = fmaxf(g, BraggFootprintGrow_px(s, p));
    }
    return g;
}

// The panel a reflection is on (DetectorModules), as the stencil needs it. Everything radial about a
// reflection - its distance from the beam, the direction the ring is elongated in, the radial
// background - is a property of the detector PLANE, while the pixels it is walked over are the image's.
// A panel's image pixels lie on the plane at a fixed turn (and maybe mirror), so the radial direction is
// taken on the plane and expressed in the panel's own pixel axes; a panel is rigid, so a distance in its
// pixels is the same distance on the plane. A pixel outside the panel's box is not the reflection's
// neighbour on the plane, whatever is stored there: the engines read it as a gap.
// Without panels the image is the plane: the whole image, plane == image, and every expression below
// is the one it always was.
struct BraggPanel {
    float plane_x, plane_y;     // the reflection on the plane
    float fx, fy, sx, sy;       // the plane step of one image pixel along x, and of one along y
    int32_t x0, y0, x1, y1;     // the panel's pixels in the image, inclusive
};
// The Metal engine hands both to its kernels as they lie in host memory, and this header is compiled on
// both sides: a layout the two compilers disagree on fails here, on whichever side it is.
static_assert(sizeof(BraggStencilParams) == 168, "BraggStencilParams layout shared with the Metal kernels");
static_assert(sizeof(BraggPanel) == 40, "BraggPanel layout shared with the Metal kernels");

BRAGG_STENCIL_HD bool BraggPanelHolds(const BRAGG_STENCIL_THREAD BraggPanel &pn, int x, int y) {
    return x >= pn.x0 && x <= pn.x1 && y >= pn.y0 && y <= pn.y1;
}

BRAGG_STENCIL_HD BraggStencil MakeBraggStencil(const BRAGG_STENCIL_THREAD BraggPanel &pn,
                                               const BRAGG_STENCIL_THREAD BraggStencilParams &p) {
    BraggStencil s;
    const float rx = pn.plane_x - p.beam_x, ry = pn.plane_y - p.beam_y;
    const float r0 = sqrtf(rx * rx + ry * ry);
    s.r0 = r0;
    if (r0 > 1e-6f) {
        s.ux = rx / r0;
        s.uy = ry / r0;
    }
    // The radial direction in the panel's pixel axes. A pixel step (dx, dy) is f dx + s dy on the plane,
    // so its radial part is dx (u.f) + dy (u.s). The identity of a flat detector leaves u as it is.
    const float ux = s.ux, uy = s.uy;
    s.ux = ux * pn.fx + uy * pn.fy;
    s.uy = ux * pn.sx + uy * pn.sy;

    s.grow = BraggStencilGrow_px(r0, p);
    // A spot wider than the r1 disk can resolve keeps its background ring clear of itself: the ring
    // starts BRAGG_FOOTPRINT_REACH of the measured footprint out, radially and tangentially apart.
    if (p.fp_n > 0) {
        float s_rad, s_tan;
        BraggFootprintAt(r0, p, s_rad, s_tan);
        if (BRAGG_FOOTPRINT_NSIGMA * fmaxf(s_rad, s_tan) > p.r1) {
            s.fp_s2r = s_rad * s_rad;
            s.fp_s2t = s_tan * s_tan;
            s.core_r = 1.0f / (BRAGG_FOOTPRINT_REACH * BRAGG_FOOTPRINT_REACH * s.fp_s2r);
            s.core_t = 1.0f / (BRAGG_FOOTPRINT_REACH * BRAGG_FOOTPRINT_REACH * s.fp_s2t);
            s.grow = fmaxf(s.grow, BraggFootprintGrow_px(s_rad, p));
            s.grow_tan = BraggFootprintGrow_px(s_tan, p);
        }
    }
    const float grow = s.grow, grow_t = s.grow_tan;
    if (grow == 0.0f && grow_t == 0.0f) {   // the common case, and the only one the online path can reach
        s.ex_in = p.r2; s.ey_in = p.r2;
        s.ex_out = p.r3; s.ey_out = p.r3;
        return s;
    }

    const float a_in = p.r2 + grow, a_out = p.r3 + grow;          // semi-axes along the radius
    const float b_in = p.r2 + grow_t, b_out = p.r3 + grow_t;      // and across it
    const float si = p.r2 / a_in, so = p.r3 / a_out;
    s.q_in = 1.0f - si * si;
    s.q_out = 1.0f - so * so;
    if (grow_t > 0.0f) {
        const float ti = p.r2 / b_in, to = p.r3 / b_out;
        s.qt_in = 1.0f - ti * ti;
        s.qt_out = 1.0f - to * to;
    }

    // Axis-aligned half-extents of an ellipse with semi-axes (a along u, b across it).
    const float ux2 = s.ux * s.ux, uy2 = s.uy * s.uy;
    s.ex_in = sqrtf(a_in * a_in * ux2 + b_in * b_in * uy2);
    s.ey_in = sqrtf(a_in * a_in * uy2 + b_in * b_in * ux2);
    s.ex_out = sqrtf(a_out * a_out * ux2 + b_out * b_out * uy2);
    s.ey_out = sqrtf(a_out * a_out * uy2 + b_out * b_out * ux2);
    return s;
}

// The cell of a learned profile (ProfileEmpirical) for grid offset (dx, dy) in a panel's pixels, on a grid
// of half-size R. The learned profile is kept in the plane's axes, so that reflections on panels turned
// against each other pool one shape rather than four orientations of it. A grid of whole pixels can only
// follow a quarter turn (and a mirror), so each panel axis is taken as the plane axis nearest to it; the
// small tilt of a placed panel is left out. Without panels, the cell of (dx, dy) itself.
BRAGG_STENCIL_HD int BraggProfileCell(const BRAGG_STENCIL_THREAD BraggPanel &pn, int dx, int dy, int R) {
    const bool f_on_x = pn.fx * pn.fx >= pn.fy * pn.fy, s_on_x = pn.sx * pn.sx > pn.sy * pn.sy;
    const int f = f_on_x ? (pn.fx > 0.0f ? 1 : -1) : (pn.fy > 0.0f ? 1 : -1);
    const int s = s_on_x ? (pn.sx > 0.0f ? 1 : -1) : (pn.sy > 0.0f ? 1 : -1);
    const int px = (f_on_x ? f * dx : 0) + (s_on_x ? s * dy : 0);
    const int py = (f_on_x ? 0 : f * dx) + (s_on_x ? 0 : s * dy);
    return (py + R) * (2 * R + 1) + (px + R);
}

// A reflection at (x, y) of a flat detector, where the image is the plane.
BRAGG_STENCIL_HD BraggStencil MakeBraggStencil(float x, float y, const BRAGG_STENCIL_THREAD BraggStencilParams &p) {
    const BraggPanel pn = {x, y, 1.0f, 0.0f, 0.0f, 1.0f, 0, 0, 0, 0};
    return MakeBraggStencil(pn, p);
}

// Which of the n_kern radial-background kernels this reflection uses: one per whole pixel of growth.
// Taken from the stencil, so the kernel and the ring it corrects are built from the same radius -
// deriving it from a separately computed r0 let the two engines round differently and pick
// different rows.
BRAGG_STENCIL_HD int BraggStencilKernelIndex(const BRAGG_STENCIL_THREAD BraggStencil &s, int n_kern) {
    if (n_kern <= 1) return 0;
    int idx = (int) lroundf(s.grow);
    if (idx < 0) idx = 0;
    if (idx >= n_kern) idx = n_kern - 1;
    return idx;
}

// The three squared distances a pixel at offset (ddx, ddy) is tested by, plus its radial offset
// (which the radial-background curve bins on). `signal` is the plain circular distance.
struct BraggStencilDist {
    float signal, inner, outer, rad;
};

// --- Ownership of a shared signal region ---------------------------------------------------------
// Two reflections whose centres are closer than 2*r1 have signal disks that intersect, and a pixel in
// the intersection carries both. It belongs to the NEARER centre. The union mask above cannot say
// that (it also marks a reflection's own core), so the decision is made once per image into an owner
// map: every reflection writes (quantised distance << 24) | index over its claim disk with an atomic
// minimum, so the nearest centre wins whatever order the writes arrive in and the lowest index breaks
// a tie. Eight bits of distance over the claim radius is 1/40 px at the default r2 = 6 px, far finer
// than the prediction itself. A pixel nobody claimed is nobody's neighbour, hence the owner's.
BRAGG_STENCIL_CONSTANT constexpr uint32_t BRAGG_OWNER_NONE = 0xffffffffu;
BRAGG_STENCIL_CONSTANT constexpr uint32_t BRAGG_OWNER_IDX_MASK = 0x00ffffffu;

BRAGG_STENCIL_HD int BraggOwnerQuant(float d, float inv_claim) {
    int q = (int) (d * inv_claim * 255.0f);
    return q > 255 ? 255 : q;
}

BRAGG_STENCIL_HD uint32_t BraggOwnerKey(float d, float inv_claim, int index) {
    return ((uint32_t) BraggOwnerQuant(d, inv_claim) << 24)
           | ((uint32_t) index & BRAGG_OWNER_IDX_MASK);
}

BRAGG_STENCIL_HD bool BraggOwnedBy(uint32_t owner, int index) {
    return owner == BRAGG_OWNER_NONE
           || (owner & BRAGG_OWNER_IDX_MASK) == ((uint32_t) index & BRAGG_OWNER_IDX_MASK);
}

// Widening that split - keeping a pixel only where no other centre is within its distance PLUS a
// margin - was built and measured, and it is worse, monotonically: on the crowded crystal the residual
// bias of the pixels that were kept grew from +0.072 at no margin to +0.087, +0.144 and +0.209 in ln
// intensity at 1, 2 and 3 px, and the merge fell apart with it. What the margin removes is the
// reflection's own profile, not the neighbour's tail, so the plain nearest-centre split is the rule.

BRAGG_STENCIL_HD BraggStencilDist BraggStencilDistances(const BRAGG_STENCIL_THREAD BraggStencil &s, float ddx, float ddy) {
    BraggStencilDist d;
    d.rad = ddx * s.ux + ddy * s.uy;
    d.signal = ddx * ddx + ddy * ddy;
    const float rad2 = d.rad * d.rad;
    d.inner = d.signal - s.q_in * rad2;
    d.outer = d.signal - s.q_out * rad2;
    if (s.qt_in != 0.0f) {   // grown tangentially as well (the footprint); tan^2 = d2 - rad^2
        const float tan2 = d.signal - rad2;
        d.inner -= s.qt_in * tan2;
        d.outer -= s.qt_out * tan2;
    }
    return d;
}

// The signal region: the r1 disk, and for a spot that outgrows it the whole footprint ellipse too, so
// the summation - the seed of the profile fit, and what the fit falls back to - holds the spot rather
// than its core. Without a footprint this is the plain disk test, bit for bit.
BRAGG_STENCIL_HD bool BraggStencilInSignal(const BRAGG_STENCIL_THREAD BraggStencil &s, float ddx, float ddy, float d2, float r1_sq) {
    if (d2 < r1_sq) return true;
    if (!(s.fp_s2t > 0.0f)) return false;
    const float rad = ddx * s.ux + ddy * s.uy, tn = -ddx * s.uy + ddy * s.ux;
    return rad * rad * s.core_r + tn * tn * s.core_t < 1.0f;
}