Jungfraujoch/image_analysis/geom_refinement/BeamCenterFromBackground.cpp

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

#include "BeamCenterFromBackground.h"
#include "BackgroundBand.h"

#include <algorithm>
#include <cmath>
#include <thread>

#include "../../common/CompressedImage.h"
#include "../../common/JFJochMath.h"
#include "../../common/ParallelFor.h"
#ifdef JFJOCH_USE_CUDA
#include "../../common/CUDAWrapper.h"
#include "BeamCenterBackgroundGPU.h"
#endif

namespace {

// The band the background is fitted over. The low-resolution end sits outside the beam stop
// and its penumbra, the high-resolution end where the solvent ring has died away.
constexpr float BAND_LOW_RES_A = 12.0f;
constexpr float BAND_HIGH_RES_A = 2.2f;

constexpr int SECTORS = BackgroundBand::SECTORS;
constexpr int RADIAL_BINS = BackgroundBand::RADIAL_BINS;

// A cell with fewer pixels than this has no usable mean.
constexpr int MIN_PIXELS_PER_CELL = 20;

// A radial bin missing more azimuth than this is a partial ring - it leaves the detector, or a
// module gap eats it - and a partial ring biases the profile it is compared against.
constexpr float MIN_SECTOR_COVERAGE = 0.85f;

// Fractions of a bin's pixels are Bragg peaks. Two rounds of clipping at the upper 2 sigma take
// the mean back to the background without needing the pixel values a second time.
constexpr float CLIP_SIGMA = 2.0f;
constexpr int CLIP_ROUNDS = 2;

// The profile is rebuilt at the trial centre every iteration, so a centre that is off smears the
// solvent ring and flattens g', which over-estimates the shift. Half steps damp that; the fixed
// point is unchanged, only the path to it.
constexpr float DAMPING = 0.5f;
constexpr float CONVERGED_PXL = 0.02f;

// The step is compared with the one before it, and a walk whose step CANCELS its predecessor twice
// running has crossed its fixed point rather than walked toward it: with DAMPING = 0.5 that happens
// where the local gain exceeds 2, which is a period-2 limit cycle the walk will not leave.
// Instrumented, a sweep that spent 96 of its 100 iterations in one oscillated by 0.08 px about a
// centre it knew to 1.7 px. Cancelling is the whole of it, and it is not the same as merely turning
// by more than a right angle: a travelling walk wobbles, and on one sweep a single 0.56 px step
// between 2 px ones turned twice in passing and ended the fit 118 px short of the centre the same
// fit reaches from the other side - in one build and not another, because the sign of a dot product
// between two near-orthogonal steps is set by the last ulp and not by the data.
constexpr int REVERSALS_AT_THE_FIXED_POINT = 2;

// A travel budget, not a convergence criterion. The shift a sector's regression can report is
// bounded by the width of the features it reads - a sector whose profile has moved a long way is
// not g + d*g' for any d, and the least-squares projection onto g' returns far less than the true
// d - so the walk advances by a bounded distance per iteration, of the order of twenty pixels,
// however far it still has to go. Ten iterations therefore cap the fit at about two hundred
// pixels of travel and a centre further out than that is left part way there, still walking, with
// the per-iteration precision reported as though it had arrived. The count is set to cross a
// detector and then converge on it; a centre that is already close leaves on CONVERGED_PXL, or on
// the reversal test above, after a handful. Measured on a sweep whose centre is 170 px out: the
// walk crawls at 1-2 px an iteration for its first forty, accelerates across the gap, and reaches
// CONVERGED_PXL at 107-115 - so a budget of 100 stopped it in the last few pixels of its approach
// and left where it stopped to the build.
constexpr int MAX_ITERATIONS = 300;

// The two passes over the pixels are split into this many row blocks, each accumulating into cells
// of its own, and the blocks are folded in block order. The split is a property of the image and
// not of the machine, so the sums are grouped the same way however many threads ran them.
constexpr int BLOCKS = 64;

// Below these the fit has not seen enough of the detector to be believed at all.
constexpr int MIN_USABLE_SECTORS = SECTORS * 3 / 5;
constexpr int MIN_USABLE_RADIAL_BINS = 15;

float median_of(std::vector<float> &v) {
    const size_t half = v.size() / 2;
    std::nth_element(v.begin(), v.begin() + half, v.end());
    return v[half];
}

}   // namespace

std::optional<BeamCenterEstimate>
FindBeamCenterFromBackground(const DiffractionExperiment &experiment, const PixelMask &mask,
                             const std::vector<float> &mean, size_t nthreads,
                             std::optional<std::pair<float, float>> start, bool allow_device) {
    if (nthreads == 0)
        nthreads = std::max(1u, std::thread::hardware_concurrency());
    const auto W = static_cast<int>(experiment.GetXPixelsNumConv());
    const auto H = static_cast<int>(experiment.GetYPixelsNumConv());
    const size_t n_pixels = static_cast<size_t>(W) * H;
    if (mean.size() != n_pixels)
        return {};

    const auto &pixel_mask = mask.GetMask(experiment);
    auto geom = experiment.GetDiffractionGeometry();

    const float wavelength = geom.GetWavelength_A();
    const float sin_high = wavelength / (2.0f * BAND_HIGH_RES_A);
    if (sin_high >= 1.0f)
        return {};
    const float tt_lo = 2.0f * std::asin(wavelength / (2.0f * BAND_LOW_RES_A));
    const float tt_hi = 2.0f * std::asin(sin_high);
    const float d_tt = (tt_hi - tt_lo) / RADIAL_BINS;

    const auto rot = geom.GetDetectorMatrix().arr();   // row major
    const float pixel_size = geom.GetPixelSize_mm();
    const float distance = geom.GetDetectorDistance_mm();

    float beam_x = start ? start->first : geom.GetBeamX_pxl();
    float beam_y = start ? start->second : geom.GetBeamY_pxl();

    constexpr int n_cells = RADIAL_BINS * SECTORS;
    std::vector<double> sum(n_cells), sum_sq(n_cells), sum_jx(n_cells), sum_jy(n_cells);
    std::vector<int32_t> count(n_cells), count_all(n_cells);
    std::vector<float> profile(RADIAL_BINS), d_profile(RADIAL_BINS), clip_limit(n_cells);
    std::vector<char> radial_ok(RADIAL_BINS);

    // One set of cells per block, allocated once and reused by every iteration.
    std::vector<int> block_row(BLOCKS + 1);
    for (int b = 0; b <= BLOCKS; b++)
        block_row[b] = static_cast<int>(static_cast<int64_t>(b) * H / BLOCKS);
    std::vector<double> block_sum(static_cast<size_t>(BLOCKS) * n_cells);
    std::vector<double> block_sum_sq(static_cast<size_t>(BLOCKS) * n_cells);
    std::vector<double> block_jx(static_cast<size_t>(BLOCKS) * n_cells);
    std::vector<double> block_jy(static_cast<size_t>(BLOCKS) * n_cells);
    std::vector<int32_t> block_count(static_cast<size_t>(BLOCKS) * n_cells);

    // Whether a pixel can take part at all, which does not depend on the centre.
    std::vector<char, NoInitAllocator<char>> usable(n_pixels);
    ParallelFor(BLOCKS, nthreads, [&](int b) {
        for (size_t i = static_cast<size_t>(block_row[b]) * W; i < static_cast<size_t>(block_row[b + 1]) * W; i++)
            usable[i] = pixel_mask[i] == 0 && std::isfinite(mean[i]);
    });
    // The pixels of each block that fall in the band at the current centre, with their cell and value,
    // in pixel order from the block's first pixel on. The clipping rounds read these instead of the
    // whole detector, in the same order.
    std::vector<int32_t, NoInitAllocator<int32_t>> band_cell(n_pixels);
    std::vector<float, NoInitAllocator<float>> band_value(n_pixels);
    std::vector<size_t> band_pixels(BLOCKS);

    // The blocks' cells folded in block order. Each cell is folded on its own, so the cells are split
    // over the threads and every cell is still summed in the same order.
    const auto fold = [&](bool with_jacobian) {
        ParallelChunks(n_cells, nthreads, [&](int c0, int c1) {
            for (int c = c0; c < c1; c++) {
                double s = 0, ss = 0, jx = 0, jy = 0;
                int32_t n = 0;
                for (int b = 0; b < BLOCKS; b++) {
                    const size_t k = static_cast<size_t>(b) * n_cells + c;
                    s += block_sum[k]; ss += block_sum_sq[k];
                    if (with_jacobian) { jx += block_jx[k]; jy += block_jy[k]; }
                    n += block_count[k];
                }
                sum[c] = s; sum_sq[c] = ss; count[c] = n;
                if (with_jacobian) { sum_jx[c] = jx; sum_jy[c] = jy; }
            }
        });
    };

    // Most pixels lie outside the band. Those clearly outside it in tan(2theta) = rho / lz - by a
    // margin far above float rounding - skip the square root and both atan2 below; every pixel the exact
    // test would keep still reaches it.
    const double tan_lo = std::tan(static_cast<double>(tt_lo)) * (1.0 - 1e-3);
    const double tan_hi = tt_hi < PI / 2 ? std::tan(static_cast<double>(tt_hi)) * (1.0 + 1e-3) : INFINITY;

    BackgroundBand band{};
    for (int k = 0; k < 9; k++)
        band.rot[k] = rot[k];
    band.pixel_size = pixel_size;
    band.distance = distance;
    band.tt_lo = tt_lo;
    band.tt_hi = tt_hi;
    band.d_tt = d_tt;
    band.tan_lo = tan_lo;
    band.tan_hi = tan_hi;

#ifndef JFJOCH_USE_CUDA
    (void) allow_device;
#else
    // With a GPU the two passes over the pixels run there, and only the cells come back.
    std::unique_ptr<BeamCenterBackgroundGPU> gpu;
    if (allow_device && get_gpu_count() > 0)
        gpu = std::make_unique<BeamCenterBackgroundGPU>(W, H, block_row, usable.data(), mean.data());
#endif

    float step_x = 0.0f, step_y = 0.0f, sigma_x = 0.0f, sigma_y = 0.0f;
    float previous_x = 0.0f, previous_y = 0.0f;
    int reversals = 0;
    // The binning pass about the current centre, then the clip rounds over the pixels it binned.
    const auto bin_cpu = [&] {
        ParallelFor(BLOCKS, nthreads, [&](int b) {
            double *b_sum = block_sum.data() + static_cast<size_t>(b) * n_cells;
            double *b_sum_sq = block_sum_sq.data() + static_cast<size_t>(b) * n_cells;
            double *b_jx = block_jx.data() + static_cast<size_t>(b) * n_cells;
            double *b_jy = block_jy.data() + static_cast<size_t>(b) * n_cells;
            int32_t *b_count = block_count.data() + static_cast<size_t>(b) * n_cells;
            std::fill(b_sum, b_sum + n_cells, 0.0);
            std::fill(b_sum_sq, b_sum_sq + n_cells, 0.0);
            std::fill(b_jx, b_jx + n_cells, 0.0);
            std::fill(b_jy, b_jy + n_cells, 0.0);
            std::fill(b_count, b_count + n_cells, 0);
            int32_t *cells = band_cell.data() + static_cast<size_t>(block_row[b]) * W;
            float *values = band_value.data() + static_cast<size_t>(block_row[b]) * W;
            size_t n_band = 0;
            for (int y = block_row[b]; y < block_row[b + 1]; y++) {
                for (int x = 0; x < W; x++) {
                    const size_t i = static_cast<size_t>(y) * W + x;
                    if (!usable[i])
                        continue;
                    float jac_x, jac_y;
                    const int cell = BackgroundBandCell(band, x, y, beam_x, beam_y, jac_x, jac_y);
                    if (cell < 0)
                        continue;
                    cells[n_band] = cell;
                    values[n_band] = mean[i];
                    n_band++;
                    b_count[cell]++;
                    b_sum[cell] += mean[i];
                    b_sum_sq[cell] += static_cast<double>(mean[i]) * mean[i];
                    b_jx[cell] += jac_x;
                    b_jy[cell] += jac_y;
                }
            }
            band_pixels[b] = n_band;
        });
        fold(true);
    };
    const auto clip_cpu = [&] {
        ParallelFor(BLOCKS, nthreads, [&](int b) {
            double *b_sum = block_sum.data() + static_cast<size_t>(b) * n_cells;
            double *b_sum_sq = block_sum_sq.data() + static_cast<size_t>(b) * n_cells;
            int32_t *b_count = block_count.data() + static_cast<size_t>(b) * n_cells;
            std::fill(b_sum, b_sum + n_cells, 0.0);
            std::fill(b_sum_sq, b_sum_sq + n_cells, 0.0);
            std::fill(b_count, b_count + n_cells, 0);
            const int32_t *cells = band_cell.data() + static_cast<size_t>(block_row[b]) * W;
            const float *values = band_value.data() + static_cast<size_t>(block_row[b]) * W;
            for (size_t j = 0; j < band_pixels[b]; j++) {
                const int32_t c = cells[j];
                const float value = values[j];
                if (clip_limit[c] < 0.0f || value > clip_limit[c])
                    continue;
                b_count[c]++;
                b_sum[c] += value;
                b_sum_sq[c] += static_cast<double>(value) * value;
            }
        });
        fold(false);
    };

    for (int iteration = 0; iteration < MAX_ITERATIONS; iteration++) {
#ifdef JFJOCH_USE_CUDA
        if (gpu)
            gpu->Bin(band, beam_x, beam_y, sum, sum_sq, sum_jx, sum_jy, count);
        else
#endif
            bin_cpu();

        count_all = count;   // the Jacobian sums belong to the unclipped pixel set

        for (int round = 0; round < CLIP_ROUNDS; round++) {
            for (int c = 0; c < n_cells; c++) {
                if (count[c] < MIN_PIXELS_PER_CELL) { clip_limit[c] = -1.0f; continue; }
                const double m = sum[c] / count[c];
                const double variance = std::max(sum_sq[c] / count[c] - m * m, 0.0);
                clip_limit[c] = static_cast<float>(m + CLIP_SIGMA * std::sqrt(variance));
            }
#ifdef JFJOCH_USE_CUDA
            if (gpu)
                gpu->Clip(clip_limit, sum, sum_sq, count);
            else
#endif
                clip_cpu();
        }

        // Radial profile: the median over the sectors that have a mean, on rings that are
        // almost fully covered.
        int usable_radial = 0;
        for (int r = 0; r < RADIAL_BINS; r++) {
            std::vector<float> present;
            for (int s = 0; s < SECTORS; s++)
                if (count[r * SECTORS + s] >= MIN_PIXELS_PER_CELL)
                    present.push_back(static_cast<float>(sum[r * SECTORS + s] / count[r * SECTORS + s]));
            radial_ok[r] = static_cast<float>(present.size()) >= MIN_SECTOR_COVERAGE * SECTORS;
            profile[r] = radial_ok[r] ? median_of(present) : 0.0f;
            usable_radial += radial_ok[r];
        }
        if (usable_radial < MIN_USABLE_RADIAL_BINS)
            return {};
        // Central difference, so a bin next to a gap in the profile drops out with it. The test
        // reads the ring BEFORE it, so it has to read the covered/not-covered flags as they were,
        // not as this same loop has already rewritten them.
        const std::vector<char> covered = radial_ok;
        for (int r = 0; r < RADIAL_BINS; r++) {
            const bool have = r > 0 && r + 1 < RADIAL_BINS && covered[r - 1] && covered[r] && covered[r + 1];
            d_profile[r] = have ? (profile[r + 1] - profile[r - 1]) / (2 * d_tt) : 0.0f;
            radial_ok[r] = have;
        }

        // Per sector: regress (profile of the sector - common profile) on {g, g'}. The first
        // coefficient is the sector's amplitude, the second its radial shift; only the shift
        // is carried on.
        std::vector<float> shift, weight, jacobian_x, jacobian_y;
        for (int s = 0; s < SECTORS; s++) {
            double a11 = 0, a12 = 0, a22 = 0, b1 = 0, b2 = 0;
            double jx = 0, jy = 0;
            int n = 0;
            for (int r = 0; r < RADIAL_BINS; r++) {
                const int c = r * SECTORS + s;
                if (!radial_ok[r] || count[c] < MIN_PIXELS_PER_CELL)
                    continue;
                const double g = profile[r], dg = d_profile[r];
                const double y = sum[c] / count[c] - profile[r];
                a11 += g * g; a12 += g * dg; a22 += dg * dg;
                b1 += g * y; b2 += dg * y;
                jx += sum_jx[c] / count_all[c];
                jy += sum_jy[c] / count_all[c];
                n++;
            }
            const double det = a11 * a22 - a12 * a12;
            if (n < MIN_USABLE_RADIAL_BINS || det <= 0)
                continue;
            const double amplitude = (a22 * b1 - a12 * b2) / det;
            const double this_shift = (a11 * b2 - a12 * b1) / det;
            // Residual sum of squares from the normal equations, without a second pass.
            double residual = 0;
            for (int r = 0; r < RADIAL_BINS; r++) {
                const int c = r * SECTORS + s;
                if (!radial_ok[r] || count[c] < MIN_PIXELS_PER_CELL)
                    continue;
                const double e = sum[c] / count[c] - profile[r] - amplitude * profile[r] - this_shift * d_profile[r];
                residual += e * e;
            }
            const double variance = residual / (n - 2) * (a11 / det);
            if (!(variance > 0))
                continue;
            shift.push_back(static_cast<float>(this_shift));
            weight.push_back(static_cast<float>(1.0 / variance));
            jacobian_x.push_back(static_cast<float>(jx / n));
            jacobian_y.push_back(static_cast<float>(jy / n));
        }
        if (static_cast<int>(shift.size()) < MIN_USABLE_SECTORS)
            return {};

        // shift_k = Jx_k dx + Jy_k dy, robustified so one bad sector cannot carry the answer.
        std::vector<float> w = weight;
        double c11 = 0, c12 = 0, c22 = 0;
        for (int round = 0; round < 3; round++) {
            c11 = c12 = c22 = 0;
            double r1 = 0, r2 = 0;
            for (size_t k = 0; k < shift.size(); k++) {
                c11 += w[k] * jacobian_x[k] * jacobian_x[k];
                c12 += w[k] * jacobian_x[k] * jacobian_y[k];
                c22 += w[k] * jacobian_y[k] * jacobian_y[k];
                r1 += w[k] * jacobian_x[k] * shift[k];
                r2 += w[k] * jacobian_y[k] * shift[k];
            }
            const double det = c11 * c22 - c12 * c12;
            if (det <= 0)
                return {};
            step_x = static_cast<float>((c22 * r1 - c12 * r2) / det);
            step_y = static_cast<float>((c11 * r2 - c12 * r1) / det);
            std::vector<float> residual(shift.size());
            for (size_t k = 0; k < shift.size(); k++)
                residual[k] = shift[k] - jacobian_x[k] * step_x - jacobian_y[k] * step_y;
            std::vector<float> absolute(residual.size());
            for (size_t k = 0; k < residual.size(); k++) absolute[k] = std::abs(residual[k]);
            const float scale = 1.4826f * median_of(absolute) + 1e-30f;
            for (size_t k = 0; k < shift.size(); k++) {
                const float t = residual[k] / (3 * scale);
                w[k] = weight[k] / (1.0f + t * t);
            }
        }

        double chi2 = 0;
        for (size_t k = 0; k < shift.size(); k++) {
            const double e = shift[k] - jacobian_x[k] * step_x - jacobian_y[k] * step_y;
            chi2 += w[k] * e * e;
        }
        chi2 = std::max(chi2 / (shift.size() - 2), 1.0);
        const double det = c11 * c22 - c12 * c12;
        sigma_x = static_cast<float>(std::sqrt(c22 / det * chi2));
        sigma_y = static_cast<float>(std::sqrt(c11 / det * chi2));

        const bool cancels = std::hypot(step_x + previous_x, step_y + previous_y)
                             < std::min(std::hypot(step_x, step_y), std::hypot(previous_x, previous_y));
        reversals = cancels ? reversals + 1 : 0;
        previous_x = step_x;
        previous_y = step_y;

        beam_x += DAMPING * step_x;
        beam_y += DAMPING * step_y;
        if (std::hypot(step_x, step_y) < CONVERGED_PXL || reversals >= REVERSALS_AT_THE_FIXED_POINT)
            break;
    }

    // A fit that leaves on the iteration cap has not converged - it was still walking when it ran
    // out - and the precision of its last step is not what it knows the centre to. The step it
    // still wanted to take is a floor under what is left, so report that instead: it turns a
    // confidently wrong answer into one the caller's sigma gate refuses. A fit that stopped on
    // CONVERGED_PXL or at its fixed point takes a step far under any sigma worth reporting, so
    // this never touches those.
    return BeamCenterEstimate{beam_x, beam_y,
                              std::max({sigma_x, sigma_y, std::hypot(step_x, step_y)})};
}

std::optional<BeamCenterEstimate>
FindBeamCenter(const DiffractionExperiment &experiment, const PixelMask &mask,
               const std::vector<float> &mean, size_t nthreads, BeamCenterFFTResult *capture) {
    // The image the capture scores: the projection with everything the mask excludes taken out.
    // The beam stop is in that mask by the time this runs, and it is the part that matters - an
    // opaque region on one side of the beam is centrosymmetric about ITS own centre, and where it
    // is large enough that preference beats the background's.
    std::vector<float> scored = mean;
    const auto &pixel_mask = mask.GetMask(experiment);
    if (pixel_mask.size() == scored.size())
        for (size_t i = 0; i < scored.size(); i++)
            if (pixel_mask[i] != 0)
                scored[i] = NAN;

    const BeamCenterFFTResult fft = FindBeamCenterFFT(experiment, scored);
    if (capture)
        *capture = fft;
    if (fft.point.empty())
        return FindBeamCenterFromBackground(experiment, mask, mean, nthreads);

    const std::pair<float, float> seed{fft.point[0].beam_x_pxl, fft.point[0].beam_y_pxl};
    if (auto refined = FindBeamCenterFromBackground(experiment, mask, mean, nthreads, seed))
        return refined;

    // The walk has nothing to fit AT THE CAPTURE. That is a statement about the capture as much as
    // about the background: a centre far enough out leaves too few fully covered rings inside the
    // fitted band for the profile to exist at all. Measured on one small-molecule set whose capture
    // lands 536 px out - the walk from the file's centre answers there to 2.9 px, and taking the
    // capture instead would have been the composition's only regression on the corpus. So ask the
    // walk again where it used to be asked, and keep the capture only where neither start works.
    if (auto from_file = FindBeamCenterFromBackground(experiment, mask, mean, nthreads))
        return from_file;
    return BeamCenterEstimate{seed.first, seed.second, BEAM_CENTER_CAPTURE_SIGMA_PXL};
}