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
// and the all-or-nothing `n_inner_valid == n_inner` acceptance gate, so growing it rejects every
// reflection with one bad pixel anywhere along a long streak. The flux it loses is a function of
// resolution alone and the per-shell scale absorbs it.
//
// 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 both the host compiler and nvcc, so the geometry has exactly one
// definition: the two engines classify by the same rule, built from the same inputs. They are not
// bit-identical to each other - 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 the parity test's tolerances are for.
// =============================================================================

#include <cmath>
#include <cstdint>

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

// Fixed per-experiment inputs to the stencil law (mirrors BraggIntegrationEngine's members).
struct BraggStencilParams {
    float beam_x = 0.0f, beam_y = 0.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
};

// 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 q_in = 0.0f, q_out = 0.0f; // radial shrink coefficients (0 = circular)
    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 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;
}

BRAGG_STENCIL_HD BraggStencil MakeBraggStencil(float px_x, float px_y, const BraggStencilParams &p) {
    BraggStencil s;
    const float rx = px_x - p.beam_x, ry = px_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;
    }

    s.grow = BraggStencilGrow_px(r0, p);
    const float grow = s.grow;
    if (grow == 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;
    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;

    // 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 + p.r2 * p.r2 * uy2);
    s.ey_in = sqrtf(a_in * a_in * uy2 + p.r2 * p.r2 * ux2);
    s.ex_out = sqrtf(a_out * a_out * ux2 + p.r3 * p.r3 * uy2);
    s.ey_out = sqrtf(a_out * a_out * uy2 + p.r3 * p.r3 * ux2);
    return s;
}

// 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 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.
constexpr uint32_t BRAGG_OWNER_NONE = 0xffffffffu;
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 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;
    return d;
}