Jungfraujoch/image_analysis/bragg_integration/BraggIntegrate2D.cpp

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

#include "BraggIntegrate2D.h"

#include <algorithm>
#include <cmath>
#include <vector>

namespace {

void MarkReflectionMask(std::vector<uint8_t> &mask,
                        size_t xpixel, size_t ypixel,
                        const Reflection &r, float r_2, float r_2_sq) {

    int64_t x0 = std::floor(r.predicted_x - r_2 - 1.0f);
    int64_t x1 = std::ceil(r.predicted_x + r_2 + 1.0f);
    int64_t y0 = std::floor(r.predicted_y - r_2 - 1.0f);
    int64_t y1 = std::ceil(r.predicted_y + r_2 + 1.0f);

    if (x0 < 0)
        x0 = 0;
    if (y0 < 0)
        y0 = 0;
    if (x1 >= static_cast<int64_t>(xpixel))
        x1 = static_cast<int64_t>(xpixel) - 1;
    if (y1 >= static_cast<int64_t>(ypixel))
        y1 = static_cast<int64_t>(ypixel) - 1;

    for (int64_t y = y0; y <= y1; y++) {
        for (int64_t x = x0; x <= x1; x++) {
            const float dist_sq = (x - r.predicted_x) * (x - r.predicted_x)
                                + (y - r.predicted_y) * (y - r.predicted_y);
            if (dist_sq < r_2_sq)
                mask[y * xpixel + x] = 1;
        }
    }
}

std::vector<uint8_t> BuildReflectionMask(const std::vector<Reflection> &predicted,
                                         size_t npredicted,
                                         size_t xpixel, size_t ypixel,
                                         float r_2, float r_2_sq) {
    std::vector<uint8_t> mask(xpixel * ypixel, 0);

    for (size_t i = 0; i < npredicted; i++)
        MarkReflectionMask(mask, xpixel, ypixel, predicted.at(i), r_2, r_2_sq);

    return mask;
}

} // namespace

template<class T>
void IntegrateReflection(Reflection &r, const T *image, const std::vector<uint8_t> &reflection_mask,
                         size_t xpixel, size_t ypixel,
                         int64_t special_value, int64_t saturation,
                         float r_3, float r_1_sq, float r_2_sq, float r_3_sq,
                         float minimum_sigma_in_regards_to_i) {

    int64_t x0 = std::floor(r.predicted_x - r_3 - 1.0);
    int64_t x1 = std::ceil(r.predicted_x + r_3 + 1.0);
    int64_t y0 = std::floor(r.predicted_y - r_3 - 1.0);
    int64_t y1 = std::ceil(r.predicted_y + r_3 + 1.0);
    if (x0 < 0)
        x0 = 0;
    if (y0 < 0)
        y0 = 0;
    if (x1 >= xpixel)
        x1 = xpixel - 1;
    if (y1 >= ypixel)
        y1 = ypixel - 1;

    int64_t I_sum = 0;
    int64_t I_npixel_inner = 0;
    int64_t I_npixel_integrated = 0;
    int64_t I_sum_x = 0;
    int64_t I_sum_y = 0;

    std::vector<T> bkg_values;
    bkg_values.reserve(static_cast<size_t>((x1 - x0 + 1) * (y1 - y0 + 1)));

    for (int64_t y = y0; y <= y1; y++) {
        for (int64_t x = x0; x <= x1; x++) {
            const float dist_sq = (x - r.predicted_x) * (x - r.predicted_x) + (y - r.predicted_y) * (y - r.predicted_y);
            const auto pixel = image[y * xpixel + x];

            if (dist_sq < r_1_sq)
                I_npixel_inner++;

            // Masked/error pixels carry the type minimum and saturated the type maximum, but the lossy
            // codec can nudge them inward by one (masked observed as INT32_MIN+1), so reject the +/-1
            // band too - real calibrated counts never get anywhere near the type extremes.
            if (pixel == special_value || pixel == special_value + 1 || pixel == saturation || pixel == saturation - 1)
                continue;

            if (dist_sq < r_1_sq) {
                I_sum += pixel;
                I_sum_x += x * pixel;
                I_sum_y += y * pixel;
                I_npixel_integrated++;
            } else if (dist_sq >= r_2_sq && dist_sq < r_3_sq) {
                if (reflection_mask[y * xpixel + x] != 0)
                    continue;
                bkg_values.emplace_back(pixel);
            }
        }
    }

    if ((I_npixel_integrated == I_npixel_inner) && (bkg_values.size() > 5)) {
        // Mean, not median: the median of a right-skewed (Poisson) background sits below
        // the mean, so subtracting it under-subtracts and biases weak intensities
        // positive (fake <I/sig> in no-signal high-resolution shells).
        double bkg_sum = 0.0;
        for (const T v : bkg_values)
            bkg_sum += static_cast<double>(v);
        r.bkg = static_cast<float>(bkg_sum / static_cast<double>(bkg_values.size()));
        r.I = static_cast<float>(I_sum) - static_cast<float>(I_npixel_integrated) * r.bkg;
        if (I_sum > 0) {
            r.observed_x = static_cast<float>(I_sum_x) / static_cast<float>(I_sum);
            r.observed_y = static_cast<float>(I_sum_y) / static_cast<float>(I_sum);
        }

        // sigma is max of the:
        // - 1 (for zero photons)
        // - Poisson noise (sqrt(I_sum)) (for in between)
        // - minimum_sigma_in_regards_to_i of Intensity (for very large numbers)
        r.sigma = std::max(1.0f, r.I * minimum_sigma_in_regards_to_i);
        if (I_sum > 0)
            r.sigma = std::max(r.sigma, std::sqrt(static_cast<float>(I_sum)));
        r.observed = true;
    } else {
        r.I = 0;
        r.bkg = 0;
        r.sigma = NAN;
        r.observed = false;
    }
}

template<class T>
std::vector<Reflection> IntegrateInternal(const DiffractionExperiment &experiment,
                                          const CompressedImage &image,
                                          const std::vector<Reflection> &predicted,
                                          size_t npredicted,
                                          int64_t special_value,
                                          int64_t saturation,
                                          int64_t image_number) {

    std::vector<Reflection> ret;
    ret.reserve(npredicted);

    auto settings = experiment.GetBraggIntegrationSettings();
    auto geom = experiment.GetDiffractionGeometry();

    std::vector<uint8_t> buffer;
    auto ptr = reinterpret_cast<const T *>(image.GetUncompressedPtr(buffer));

    const float r_3 = settings.GetR3();
    const float r_1_sq = settings.GetR1() * settings.GetR1();
    const float r_2 = settings.GetR2();
    const float r_2_sq = settings.GetR2() * settings.GetR2();
    const float r_3_sq = settings.GetR3() * settings.GetR3();

    const float minimum_sigma_in_regards_to_i = settings.GetMinimumSigmaInRegardsToI();
    const auto reflection_mask = BuildReflectionMask(predicted, npredicted,
                                                     image.GetWidth(), image.GetHeight(),
                                                     r_2, r_2_sq);

    for (int i = 0; i < npredicted; i++) {
        auto r = predicted.at(i);
        IntegrateReflection(r, ptr, reflection_mask, image.GetWidth(), image.GetHeight(), special_value, saturation,
                            r_3, r_1_sq, r_2_sq, r_3_sq, minimum_sigma_in_regards_to_i);
        if (r.observed) {
            if (experiment.GetPolarizationFactor())
                r.rlp /= geom.CalcAzIntPolarizationCorr(r.predicted_x, r.predicted_y, experiment.GetPolarizationFactor().value());
            r.image_scale_corr = r.rlp / r.partiality;
            r.image_number = static_cast<float>(image_number);
            ret.emplace_back(r);
        }
    }
    return ret;
}

std::vector<Reflection> BraggIntegrate2D(const DiffractionExperiment &experiment,
                                         const CompressedImage &image,
                                         const std::vector<Reflection> &predicted,
                                         size_t npredicted,
                                         int64_t image_number) {
    if (image.GetCompressedSize() == 0 || predicted.empty())
        return {};

    switch (image.GetMode()) {
        case CompressedImageMode::Int8:
           return IntegrateInternal<int8_t>(experiment, image, predicted, npredicted, INT8_MIN, INT8_MAX, image_number);
        case CompressedImageMode::Int16:
            return IntegrateInternal<int16_t>(experiment, image, predicted, npredicted, INT16_MIN, INT16_MAX, image_number);
        case CompressedImageMode::Int32:
            return IntegrateInternal<int32_t>(experiment, image, predicted, npredicted, INT32_MIN, INT32_MAX, image_number);
        case CompressedImageMode::Uint8:
            return IntegrateInternal<uint8_t>(experiment, image, predicted, npredicted, UINT8_MAX, UINT8_MAX, image_number);
        case CompressedImageMode::Uint16:
            return IntegrateInternal<uint16_t>(experiment, image, predicted, npredicted, UINT16_MAX, UINT16_MAX, image_number);
        case CompressedImageMode::Uint32:
            return IntegrateInternal<uint32_t>(experiment, image, predicted, npredicted, UINT32_MAX, UINT32_MAX, image_number);
        default:
            throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                                  "Image mode not supported");
    }
}