Jungfraujoch/common/AzimuthalIntegrationMapping.cpp

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

#include "JFJochMath.h"
#include <algorithm>
#include <bit>
#include <cmath>
#include <thread>
#include <future>

#include "AzimuthalIntegrationMapping.h"
#include "JFJochException.h"
#include "DiffractionGeometry.h"
#include "RawToConvertedGeometry.h"
#include "TableChecksum.h"
#include "ParallelFor.h"

AzimuthalIntegrationMapping::AzimuthalIntegrationMapping(const DiffractionExperiment &experiment,
                                                         const PixelMask &mask,
                                                         size_t in_nthreads)
    : AzimuthalIntegrationMapping(experiment, mask, nullptr, in_nthreads) {}

AzimuthalIntegrationMapping::AzimuthalIntegrationMapping(const DiffractionExperiment &experiment,
                                                         const PixelMask &mask,
                                                         AzimuthalIntegrationGeometryCache &cache,
                                                         size_t in_nthreads)
    : AzimuthalIntegrationMapping(experiment, mask, &cache, in_nthreads) {}

AzimuthalIntegrationMapping::AzimuthalIntegrationMapping(const DiffractionExperiment &experiment,
                                           const PixelMask& mask,
                                           AzimuthalIntegrationGeometryCache *cache,
                                           size_t in_nthreads)
    : settings(experiment.GetAzimuthalIntegrationSettings()),
      wavelength(experiment.GetWavelength_A()),
      // The dimensions of the image this mapping is built for: converted when the geometry is
      // transformed, raw when it is not. They have to match pixel_to_bin, which is sized per mode
      // below - the adaptive spot finders walk the image with these and index pixel_to_bin with it.
      width(experiment.GetXPixelsNum()),
      height(experiment.GetYPixelsNum()) {

    if (width <= 0)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Detector width must be above 0");

    if (height <= 0)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Detector height must be above 0");

    if (settings.GetBinCount() >= UINT16_MAX)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Cannot handle more than 65534 az. int. bins");

    polarization_factor = experiment.GetPolarizationFactor();

    if (in_nthreads == 0)
        nthreads = AvailableCpus();
    else
        nthreads = in_nthreads;

    nthreads = std::clamp<size_t>(nthreads, 1, 64);

    if (!experiment.IsGeometryTransformed())
        SetupRawGeom(experiment, mask.GetMaskRaw());
    else if (cache)
        SetupConvGeomCached(experiment.GetDiffractionGeometry(), mask.GetMask(), *cache);
    else
        SetupConvGeom(experiment.GetDiffractionGeometry(), &mask.GetMask());

    UpdateMaxBinNumber();

    pixel_to_bin_checksum = TableChecksum(pixel_to_bin.data(), pixel_to_bin.size() * sizeof(uint16_t));
    corrections_checksum = TableChecksum(corrections.data(), corrections.size() * sizeof(float));
}

void AzimuthalIntegrationMapping::SetupConvGeomRows(const DiffractionGeometry &geom, const std::vector<uint32_t> *mask,
    size_t row0, size_t row_end) {
    for (size_t row = row0; row < row_end && row < height; row++) {
        for (size_t col = 0; col < width; col++)
            SetupPixel(geom, mask, row * width + col, col, row);
    }
}

void AzimuthalIntegrationMapping::SetupConvGeom(const DiffractionGeometry &geom, const std::vector<uint32_t> *mask) {
    pixel_to_bin.resize(width * height, UINT16_MAX);
    pixel_resolution.resize(width * height, 0);
    corrections.resize(width * height, 0);

    if (mask && mask->size() != width * height)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Mask size invalid");

    if (nthreads <= 1) {
        SetupConvGeomRows(geom, mask, 0, height);
    } else {
        auto local_nthreads = std::min(nthreads, height);
        std::vector<std::future<void>> futures;
        futures.reserve(local_nthreads);

        for (size_t t = 0; t < local_nthreads; ++t)
            futures.emplace_back(std::async(std::launch::async,
                &AzimuthalIntegrationMapping::SetupConvGeomRows,
                this, std::cref(geom), mask,
                t * height / local_nthreads,
                (t + 1) * height / local_nthreads));

        for (auto &f: futures)
            f.get();
    }
}

void AzimuthalIntegrationMapping::SetupConvGeomCached(const DiffractionGeometry &geom,
                                                      const std::vector<uint32_t> &mask,
                                                      AzimuthalIntegrationGeometryCache &cache) {
    if (mask.size() != width * height)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Mask size invalid");

    // Everything SetupPixel reads, compared bit for bit.
    std::vector<uint32_t> key;
    for (const float f: {geom.GetBeamX_pxl(), geom.GetBeamY_pxl(), geom.GetDetectorDistance_mm(),
                         geom.GetPixelSize_mm(), geom.GetWavelength_A(),
                         settings.GetLowQ_recipA(), settings.GetHighQ_recipA(), settings.GetQSpacing_recipA(),
                         polarization_factor.value_or(0.0f)})
        key.push_back(std::bit_cast<uint32_t>(f));
    for (const float f: geom.GetDetectorMatrix().arr())
        key.push_back(std::bit_cast<uint32_t>(f));
    key.insert(key.end(), {static_cast<uint32_t>(settings.GetAzimuthalBinCount()),
                           static_cast<uint32_t>(settings.IsSolidAngleCorrection()),
                           static_cast<uint32_t>(settings.IsPolarizationCorrection()),
                           static_cast<uint32_t>(polarization_factor.has_value()),
                           static_cast<uint32_t>(width), static_cast<uint32_t>(height)});

    std::lock_guard lock(cache.m);
    auto &entries = cache.entries;
    auto it = std::find_if(entries.begin(), entries.end(), [&](const auto &e) { return e.key == key; });
    if (it == entries.end()) {
        SetupConvGeom(geom, nullptr);
        if (entries.size() == AzimuthalIntegrationGeometryCache::MAX_ENTRIES)
            entries.erase(entries.begin());
        entries.push_back({key, pixel_to_bin, pixel_resolution, corrections});
    } else {
        std::rotate(it, it + 1, entries.end());
        pixel_to_bin = entries.back().pixel_to_bin;
        pixel_resolution = entries.back().pixel_resolution;
        corrections = entries.back().corrections;
    }

    // What SetupPixel leaves at a masked pixel.
    ParallelChunks(static_cast<int>(mask.size()), nthreads, [&](int lo, int hi) {
        for (int i = lo; i < hi; i++)
            if (mask[i] != 0) {
                pixel_to_bin[i] = UINT16_MAX;
                pixel_resolution[i] = 0;
                corrections[i] = 0;
            }
    });
}

void AzimuthalIntegrationMapping::SetupRawGeom(const DiffractionExperiment &experiment,
                                               const std::vector<uint32_t> &mask) {
    if (mask.size() != RAW_MODULE_SIZE * experiment.GetModulesNum())
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Mask size invalid");

    pixel_to_bin.resize(RAW_MODULE_SIZE * experiment.GetModulesNum(), UINT16_MAX);
    pixel_resolution.resize(RAW_MODULE_SIZE * experiment.GetModulesNum(), 0);
    corrections.resize(RAW_MODULE_SIZE * experiment.GetModulesNum(), 0);

    auto geom = experiment.GetDiffractionGeometry();

    if (nthreads <= 1) {
        for (int m = 0; m < experiment.GetModulesNum(); m++) {
            for (int pxl = 0; pxl < RAW_MODULE_SIZE; pxl++) {
                auto [x,y] = RawToConvertedCoordinate(experiment, m, pxl);
                SetupPixel(geom, &mask, m * RAW_MODULE_SIZE + pxl, x, y);
            }
        }
    } else {
        auto local_nthreads = std::min<size_t>(nthreads, experiment.GetModulesNum());
        std::vector<std::future<void>> futures;
        futures.reserve(local_nthreads);

        for (size_t t = 0; t < local_nthreads; ++t) {
            const size_t module_begin = t * experiment.GetModulesNum() / local_nthreads;
            const size_t module_end = (t + 1) * experiment.GetModulesNum() / local_nthreads;

            futures.emplace_back(std::async(std::launch::async, [&, module_begin, module_end] {
                for (size_t m = module_begin; m < module_end; ++m) {
                    for (int pxl = 0; pxl < RAW_MODULE_SIZE; ++pxl) {
                        auto [x, y] = RawToConvertedCoordinate(experiment, m, pxl);
                        SetupPixel(geom, &mask, m * RAW_MODULE_SIZE + pxl, x, y);
                    }
                }
            }));
        }

        for (auto &f: futures)
            f.get();
    }
}

void AzimuthalIntegrationMapping::SetupPixel(const DiffractionGeometry &geom,
                                      const std::vector<uint32_t> *mask,
                                      uint32_t pxl, uint32_t col, uint32_t row) {
    if (mask && (*mask)[pxl] != 0)
        return;

    auto x = static_cast<float>(col);
    auto y = static_cast<float>(row);
    float d = geom.PxlToRes(x, y);
    float phi_rad = geom.Phi_rad(x, y);

    pixel_resolution[pxl] = d;

    float corr = 1.0;
    if (settings.IsSolidAngleCorrection())
        corr /= geom.CalcAzIntSolidAngleCorr(x, y);
    if (settings.IsPolarizationCorrection() && polarization_factor)
        corr /= geom.CalcAzIntPolarizationCorr(x, y, polarization_factor.value());
    corrections[pxl] = corr;

    if (d > 0) {
        float q = 2.0f * static_cast<float>(PI) / d;
        pixel_to_bin[pxl] = settings.GetBin(q, phi_rad * 180.0 / PI);
    }
}

uint16_t AzimuthalIntegrationMapping::GetBinNumber() const {
    return settings.GetBinCount();
}

const std::vector<uint16_t> &AzimuthalIntegrationMapping::GetPixelToBin() const {
    return pixel_to_bin;
}

const std::vector<float> &AzimuthalIntegrationMapping::GetBinToQ() const {
    return bin_to_q;
}

const std::vector<float> &AzimuthalIntegrationMapping::GetBinToD() const {
    return bin_to_d;
}

const std::vector<float> &AzimuthalIntegrationMapping::GetBinToTwoTheta() const {
    return bin_to_2theta;
}

const std::vector<float> &AzimuthalIntegrationMapping::GetBinToPhi() const {
    return bin_to_phi;
}

uint16_t AzimuthalIntegrationMapping::QToBin(float q) const {
    return settings.QToBin(q);
}

void AzimuthalIntegrationMapping::UpdateMaxBinNumber() {
    bin_to_q.resize(settings.GetBinCount());
    bin_to_d.resize(settings.GetBinCount());
    bin_to_2theta.resize(settings.GetBinCount());
    bin_to_phi.resize(settings.GetBinCount());

    for (int j = 0; j < settings.GetAzimuthalBinCount(); j++) {
        for (int i = 0; i < settings.GetQBinCount(); i++) {
            bin_to_q[j * settings.GetQBinCount() + i] = static_cast<float>(settings.GetQSpacing_recipA() * (i + 0.5) + settings.GetLowQ_recipA());
            bin_to_d[j * settings.GetQBinCount() + i] = 2.0f * static_cast<float>(PI) / bin_to_q[j * settings.GetQBinCount() + i];
            bin_to_2theta[j * settings.GetQBinCount() + i] = 2.0f * asinf(bin_to_q[i] * wavelength / (4.0f * static_cast<float>(PI))) * 180.0f /
                               static_cast<float>(PI);
            bin_to_phi[j * settings.GetQBinCount() + i] = static_cast<float>(j) * 360.0f / static_cast<float>(settings.GetAzimuthalBinCount());
        }
    }
}

const std::vector<float> &AzimuthalIntegrationMapping::Corrections() const {
    return corrections;
}

const std::vector<float> &AzimuthalIntegrationMapping::Resolution() const {
    return pixel_resolution;
}

uint64_t AzimuthalIntegrationMapping::GetPixelToBinChecksum() const {
    return pixel_to_bin_checksum;
}

uint64_t AzimuthalIntegrationMapping::GetCorrectionsChecksum() const {
    return corrections_checksum;
}

std::shared_ptr<const std::vector<uint32_t>>
AzimuthalIntegrationMapping::ResolutionMaskBits(std::optional<float> high_res,
                                                std::optional<float> low_res) const {
    const std::lock_guard lock(res_mask_mutex);
    if (res_mask_bits && res_mask_high == high_res && res_mask_low == low_res)
        return res_mask_bits;

    // An unset limit masks nothing at that end. At the high-resolution end 0 does that on its own - no
    // pixel has d < 0, and the detector's own edge is where the pixels stop anyway; at the
    // low-resolution end every pixel lies above any finite stand-in, so it takes an infinite one.
    const float high = high_res.value_or(0.0f);
    const float low = low_res.value_or(INFINITY);

    const size_t npixel = pixel_resolution.size();
    auto bits = std::make_shared<std::vector<uint32_t>>(npixel / 32 + (npixel % 32 != 0 ? 1 : 0), 0);
    for (size_t i = 0; i < npixel; i++)
        if (pixel_resolution[i] > low || pixel_resolution[i] < high)
            (*bits)[i / 32] |= 1u << (i % 32);

    res_mask_high = high_res;
    res_mask_low = low_res;
    res_mask_bits = bits;
    return bits;
}

const AzimuthalIntegrationSettings &AzimuthalIntegrationMapping::Settings() const {
    return settings;
}

size_t AzimuthalIntegrationMapping::GetWidth() const {
    return width;
}

size_t AzimuthalIntegrationMapping::GetHeight() const {
    return height;
}

int32_t AzimuthalIntegrationMapping::GetAzimuthalBinCount() const {
    return settings.GetAzimuthalBinCount();
}

int32_t AzimuthalIntegrationMapping::GetQBinCount() const {
    return settings.GetQBinCount();
}

size_t AzimuthalIntegrationMapping::GetNThreads() const {
    return nthreads;
}