Jungfraujoch/image_analysis/bragg_integration/BraggPassB.h

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

#pragma once

// Pass B of BraggIntegrationEngineCPU, the profile fit, as the data it hands on. The CPU works out in
// double what one reflection's fit needs (BraggFitSetup), the fit itself is single precision
// (BraggFitOut), and the CPU finishes it in double. The Metal engine runs the middle on the GPU, so this
// is compiled by the host and by the Metal compiler (no default member initialisers, as in BraggPassA.h).

#ifndef __METAL_VERSION__
#include <cstdint>

#include "BraggExp.h"   // the Metal engine puts it in front
#endif

struct BraggFitSetup {
    float px, py;                  // predicted position
    int32_t cx, cy;                // profile grid centre
    int32_t Rf;                    // grid half-size
    int32_t shell;                 // the learned profile of this shell, or -1: the Gaussian below
    uint32_t cells;                // where the Metal engine keeps the Gaussian, Gf^2 floats
    int32_t fit;                   // 0: not fitted (no pass A result, or overloaded)
    int32_t footprint;             // the stencil has a measured footprint: it decides the signal disk
    float fx, fy;                  // sub-pixel offset of the prediction from the grid centre
    float ux, uy;                  // radial direction
    float inv_2s2r, inv_2s2t;      // 1 / (2 sigma^2), radially and tangentially
    float bkg;                     // background per pixel
    float B, v_min;                // the variance model's floor, and the least variance a weight takes
    float I0;                      // the box sum, where the fit starts
};

struct BraggFitOut {
    float p_grid, p_valid, p_own;  // profile mass on the detector, readable, and this reflection's own
    float p_peak, p_lost_peak;     // the largest profile value, and the largest one unreadable
    float m_all, m_read;           // mass in the signal disk, and the part the box sum read
    float I, den, wsum;            // the fit; den = 0 when it did not run or failed
};

// What the Metal engine's fit kernel is told about the frame, beside the stencil.
struct BraggPassBFrame {
    int32_t width, height;
    float r1_sq;
    int32_t R;                     // the learned profiles' grid half-size
    int32_t use_owner, exclude;    // as in BraggPassAFrame
    uint32_t n;                    // reflections
};

// Compiled on both sides, so a layout the host and the Metal compiler disagree on fails to compile.
static_assert(sizeof(BraggFitSetup) == 76, "BraggFitSetup layout shared with the Metal kernels");
static_assert(sizeof(BraggFitOut) == 40, "BraggFitOut layout shared with the Metal kernels");
static_assert(sizeof(BraggPassBFrame) == 28, "BraggPassBFrame layout shared with the Metal kernels");

// The Gaussian profile's (unnormalised) value at grid offset (dx, dy).
inline float BraggGaussian(int dx, int dy, float fx, float fy, float ux, float uy, float inv_2s2r,
                           float inv_2s2t) {
    const float ex = dx - fx, ey = dy - fy;
    const float rad = ex * ux + ey * uy, tn = -ex * uy + ey * ux;
    return BraggExp(-rad * rad * inv_2s2r - tn * tn * inv_2s2t);
}