Jungfraujoch/image_analysis/bragg_integration/SpotFootprint.cpp

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

#include "SpotFootprint.h"

#include <algorithm>
#include <cmath>
#include <map>
#include <tuple>

namespace {

// The window around a spot is BRAGG_FOOTPRINT_NSIGMA-like: three of its standard deviations, at least
// a few pixels and at most this many, which is wider than any spot the integrator could hold.
constexpr float WINDOW_NSIGMA = 3.0f;
constexpr float WINDOW_MIN_PX = 3.0f;
constexpr float WINDOW_MAX_PX = 24.0f;
// The background is the median of an elliptical annulus between these multiples of the window.
constexpr float BKG_INNER = 1.5f;
constexpr float BKG_OUTER = 2.2f;
constexpr int ITERATIONS = 8;

inline bool valid(int32_t v) { return v != INT32_MIN && v != INT32_MAX; }

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

}   // namespace

FootprintOutcome MeasureFootprintSpot(const FootprintTile &tile, int width, int height, float beam_x, float beam_y,
                                      float x, float y, FootprintSpot &spot) {
    const int half = static_cast<int>(std::ceil(BKG_OUTER * WINDOW_MAX_PX)) + 1;
    std::vector<float> ring;
    const float rx = x - beam_x, ry = y - beam_y;
    const float r = std::sqrt(rx * rx + ry * ry);
    if (!(r > 1.0f)) return FootprintOutcome::LeftOut;
    const float ux = rx / r, uy = ry / r;
    const int ix = static_cast<int>(std::lround(x)), iy = static_cast<int>(std::lround(y));
    if (ix - half < 0 || iy - half < 0 || ix + half >= width || iy + half >= height) return FootprintOutcome::LeftOut;

    // Start from a compact spot at the prediction; each round re-centres on the signal and takes
    // the window to three of the widths just measured.
    float cx = x, cy = y;
    float s2r = 1.0f, s2t = 1.0f;
    for (int it = 0; it < ITERATIONS; ++it) {
        const float wr = std::clamp(WINDOW_NSIGMA * std::sqrt(s2r), WINDOW_MIN_PX, WINDOW_MAX_PX);
        const float wt = std::clamp(WINDOW_NSIGMA * std::sqrt(s2t), WINDOW_MIN_PX, WINDOW_MAX_PX);
        const float reach = BKG_OUTER * std::max(wr, wt);
        const int x0 = static_cast<int>(std::floor(cx - reach)), x1 = static_cast<int>(std::ceil(cx + reach));
        const int y0 = static_cast<int>(std::floor(cy - reach)), y1 = static_cast<int>(std::ceil(cy + reach));
        if (x0 < 0 || y0 < 0 || x1 >= width || y1 >= height) return FootprintOutcome::LeftOut;
        // The re-centring is not bounded, so the window can walk off a tile the frame still holds.
        if (x0 < tile.x0 || y0 < tile.y0 || x1 >= tile.x0 + tile.w || y1 >= tile.y0 + tile.h)
            return FootprintOutcome::OffTile;
        const auto pixel = [&](int px, int py) {
            return tile.pixels[static_cast<size_t>(py - tile.y0) * tile.w + (px - tile.x0)];
        };

        ring.clear();
        for (int py = y0; py <= y1; ++py)
            for (int px = x0; px <= x1; ++px) {
                const float dx = px - cx, dy = py - cy;
                const float rad = dx * ux + dy * uy, tn = -dx * uy + dy * ux;
                const float e = rad * rad / (wr * wr) + tn * tn / (wt * wt);
                const int32_t v = pixel(px, py);
                if (e >= BKG_INNER * BKG_INNER && e < BKG_OUTER * BKG_OUTER && valid(v))
                    ring.push_back(static_cast<float>(v));
            }
        if (ring.size() < 10) return FootprintOutcome::LeftOut;
        const double bkg = median_of(ring);

        double w = 0.0, mr = 0.0, mt = 0.0, m2r = 0.0, m2t = 0.0;
        for (int py = y0; py <= y1; ++py)
            for (int px = x0; px <= x1; ++px) {
                const float dx = px - cx, dy = py - cy;
                const float rad = dx * ux + dy * uy, tn = -dx * uy + dy * ux;
                if (rad * rad / (wr * wr) + tn * tn / (wt * wt) >= 1.0f) continue;
                const int32_t v = pixel(px, py);
                if (!valid(v)) return FootprintOutcome::LeftOut;
                const double net = static_cast<double>(v) - bkg;
                w += net;
                mr += net * rad;
                mt += net * tn;
                m2r += net * rad * rad;
                m2t += net * tn * tn;
            }
        if (!(w > 0.0)) return FootprintOutcome::LeftOut;
        const double cr = mr / w, ct = mt / w;
        s2r = static_cast<float>(std::max(0.25, m2r / w - cr * cr));
        s2t = static_cast<float>(std::max(0.25, m2t / w - ct * ct));
        cx += static_cast<float>(cr * ux - ct * uy);
        cy += static_cast<float>(cr * uy + ct * ux);
    }
    spot = {r, std::sqrt(s2r), std::sqrt(s2t)};
    return FootprintOutcome::Measured;
}

void MeasureFootprintSpots(const int32_t *img, int width, int height, float beam_x, float beam_y,
                           const std::vector<float> &x, const std::vector<float> &y,
                           std::vector<FootprintSpot> &out) {
    const FootprintTile frame{img, 0, 0, width, height};
    for (size_t s = 0; s < x.size(); ++s) {
        FootprintSpot spot;
        if (MeasureFootprintSpot(frame, width, height, beam_x, beam_y, x[s], y[s], spot) == FootprintOutcome::Measured)
            out.push_back(spot);
    }
}

SpotFootprint FootprintFromSpots(const std::vector<FootprintSpot> &spots, float r_max) {
    SpotFootprint fp;
    if (!(r_max > 0.0f)) return fp;
    const float bin = r_max / FOOTPRINT_BINS;
    std::vector<std::vector<float>> rad(FOOTPRINT_BINS), tan(FOOTPRINT_BINS);
    for (const auto &s : spots) {
        const int b = std::clamp(static_cast<int>(s.r_px / bin), 0, FOOTPRINT_BINS - 1);
        rad[b].push_back(s.sigma_rad);
        tan[b].push_back(s.sigma_tan);
    }
    std::vector<int> filled;
    std::vector<float> mr(FOOTPRINT_BINS, 0.0f), mt(FOOTPRINT_BINS, 0.0f);
    for (int b = 0; b < FOOTPRINT_BINS; ++b)
        if (static_cast<int>(rad[b].size()) >= FOOTPRINT_MIN_SPOTS_PER_BIN) {
            mr[b] = median_of(rad[b]);
            mt[b] = median_of(tan[b]);
            filled.push_back(b);
        }
    if (filled.empty()) return fp;
    fp.bin_px = bin;
    for (int b = 0; b < FOOTPRINT_BINS; ++b) {
        // The nearest bin that has enough spots; the inner one on a tie.
        int best = filled.front();
        for (int f : filled)
            if (std::abs(f - b) < std::abs(best - b)) best = f;
        fp.sigma_rad.push_back(mr[best]);
        fp.sigma_tan.push_back(mt[best]);
    }
    return fp;
}

void MeasureFootprintOffsets(const std::vector<SpotToSave> &spots, const std::vector<Reflection> &reflections,
                             float beam_x, float beam_y, std::vector<FootprintOffset> &out) {
    std::map<std::tuple<int, int, int>, const Reflection *> predicted;
    for (const auto &r : reflections)
        predicted[{r.h, r.k, r.l}] = &r;
    for (const auto &s : spots) {
        if (!s.indexed || s.lattice != 0) continue;
        const auto it = predicted.find({static_cast<int>(s.h), static_cast<int>(s.k), static_cast<int>(s.l)});
        if (it == predicted.end()) continue;
        const float px = it->second->predicted_x, py = it->second->predicted_y;
        const float rx = px - beam_x, ry = py - beam_y;
        const float r = std::sqrt(rx * rx + ry * ry);
        if (!(r > 1.0f)) continue;
        const float ux = rx / r, uy = ry / r;
        const float dx = s.x - px, dy = s.y - py;
        out.push_back({r, dx * ux + dy * uy, -dx * uy + dy * ux});
    }
}

SpotFootprint FootprintWithOffsets(const SpotFootprint &widths, std::vector<FootprintOffset> offsets) {
    if (widths.empty()) return widths;
    // Sorted, so the sums below do not depend on the order the frames were measured in.
    std::sort(offsets.begin(), offsets.end());
    const int n = static_cast<int>(widths.sigma_rad.size());
    std::vector<double> s2r(n, 0.0), s2t(n, 0.0);
    std::vector<int> cnt(n, 0);
    for (const auto &o : offsets) {
        const int b = std::clamp(static_cast<int>(o.r_px / widths.bin_px), 0, n - 1);
        s2r[b] += static_cast<double>(o.off_rad) * o.off_rad;
        s2t[b] += static_cast<double>(o.off_tan) * o.off_tan;
        ++cnt[b];
    }
    std::vector<int> filled;
    for (int b = 0; b < n; ++b)
        if (cnt[b] >= FOOTPRINT_MIN_SPOTS_PER_BIN) filled.push_back(b);
    if (filled.empty()) return widths;
    SpotFootprint fp = widths;
    for (int b = 0; b < n; ++b) {
        int best = filled.front();
        for (int f : filled)
            if (std::abs(f - b) < std::abs(best - b)) best = f;
        fp.sigma_rad[b] = static_cast<float>(std::sqrt(widths.sigma_rad[b] * widths.sigma_rad[b] + s2r[best] / cnt[best]));
        fp.sigma_tan[b] = static_cast<float>(std::sqrt(widths.sigma_tan[b] * widths.sigma_tan[b] + s2t[best] / cnt[best]));
    }
    return fp;
}