Jungfraujoch/image_analysis/hot_pixels/HotPixels.cpp

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

#include "HotPixels.h"

#include <algorithm>
#include <cmath>
#include <cstring>
#include <limits>
#include <optional>

#include "../../common/CUDAWrapper.h"
#include "../../common/JFJochMath.h"
#include "../../common/ParallelFor.h"

#include <future>

namespace {

// The lower median of v[0..n), reordering it.
int32_t median_of(int32_t *v, size_t n) {
    std::nth_element(v, v + n / 2, v + n);
    return v[n / 2];
}

// A ring's background is a few counts, so its order statistics are read off a histogram of the values
// 0..RING_HIST_VALUES-1 - exactly, as long as the one asked for lands inside that range.
constexpr int32_t RING_HIST_VALUES = 1024;

// The rank-th smallest value, if it is inside the histogram; `below` counts the values under 0.
std::optional<int32_t> hist_order_statistic(const std::vector<uint32_t> &hist, size_t below, size_t rank) {
    if (rank < below)
        return std::nullopt;
    size_t seen = below;
    for (int32_t v = 0; v < RING_HIST_VALUES; v++) {
        seen += hist[v];
        if (seen > rank)
            return v;
    }
    return std::nullopt;
}

// The rank-th smallest |value - m|, if every value that close to m is inside the histogram (none
// below 0 anywhere, and m + d still in range).
std::optional<int32_t> hist_deviation_order_statistic(const std::vector<uint32_t> &hist, size_t below,
                                                      int32_t m, size_t rank) {
    if (below > 0)
        return std::nullopt;
    size_t seen = hist[m];
    for (int32_t d = 0; seen <= rank; ) {
        d++;
        if (m + d >= RING_HIST_VALUES)
            return std::nullopt;
        seen += hist[m + d] + (m - d >= 0 ? hist[m - d] : 0);
        if (seen > rank)
            return d;
    }
    return 0;
}

// The smallest k with npix * P(Binomial(n, q) >= k) below the family-wise rate, n + 1 if none.
int binomial_bound(int n, double q, double npix, double rate) {
    double tail = 0.0;
    int k = n + 1;
    for (int j = n; j >= 0; j--) {
        tail += std::exp(std::lgamma(n + 1.0) - std::lgamma(j + 1.0) - std::lgamma(n - j + 1.0)
                         + j * std::log(q) + (n - j) * std::log1p(-q));
        if (npix * tail >= rate) break;
        k = j;
    }
    return k;
}

}   // namespace

HotPixelFinder::HotPixelFinder(const DiffractionExperiment &experiment, const PixelMask &mask, size_t nthreads)
    : geometry(experiment.GetDiffractionGeometry()),
      axis(experiment.GetGoniometer() ? experiment.GetGoniometer()->GetAxis().Normalize() : Coord(0, 0, 0)),
      width(experiment.GetXPixelsNumConv()), height(experiment.GetYPixelsNumConv()),
      key(std::make_unique_for_overwrite<int32_t[]>(width * height)),
      n_lit(std::make_unique_for_overwrite<uint16_t[]>(width * height)),
      n_error(std::make_unique_for_overwrite<uint16_t[]>(width * height)),
      sum_value(std::make_unique_for_overwrite<int64_t[]>(width * height)),
      n_error_ring_ok(std::make_unique_for_overwrite<uint16_t[]>(width * height)),
      error_level_sum(std::make_unique_for_overwrite<int64_t[]>(width * height)) {
    const auto &m = mask.GetMask();
    // Rings of equal 2theta, RING_WIDTH_PX wide where the beam meets the detector.
    const float ring_rad = RING_WIDTH_PX * geometry.GetPixelSize_mm() / geometry.GetDetectorDistance_mm();
    ParallelChunks(static_cast<int>(height), nthreads, [&](int y0, int y1) {
        for (size_t y = y0; y < static_cast<size_t>(y1); y++)
            for (size_t x = 0; x < width; x++) {
                const size_t i = y * width + x;
                n_lit[i] = n_error[i] = n_error_ring_ok[i] = 0;
                sum_value[i] = error_level_sum[i] = 0;
                key[i] = -1;
                if (m[i] != 0) continue;
                const float fx = static_cast<float>(x), fy = static_cast<float>(y);
                const int ring = static_cast<int>(geometry.TwoTheta_rad(fx, fy) / ring_rad);
                const int sector = std::min(SECTORS - 1, static_cast<int>(geometry.Phi_rad(fx, fy)
                                                                          / static_cast<float>(2.0 * PI) * SECTORS));
                key[i] = ring * SECTORS + sector;
            }
    });
    for (size_t i = 0; i < width * height; i++)
        if (key[i] >= 0) {
            nrings = std::max(nrings, key[i] / SECTORS + 1);
            unmasked++;
        }

    key_frames.assign(static_cast<size_t>(nrings) * SECTORS, 0);
    key_level_sum.assign(static_cast<size_t>(nrings) * SECTORS, 0);
    key_local_level_sum.assign(static_cast<size_t>(nrings) * SECTORS, 0);
    key_begin.assign(static_cast<size_t>(nrings) * SECTORS + 1, 0);
    for (size_t i = 0; i < width * height; i++)
        if (key[i] >= 0) key_begin[key[i] + 1]++;
    for (size_t k = 1; k < key_begin.size(); k++)
        key_begin[k] += key_begin[k - 1];
}

void HotPixelFinder::AddImage(const int32_t *image, std::vector<int32_t> &scratch) {
    const size_t npix = width * height;
    const size_t nkeys = static_cast<size_t>(nrings) * SECTORS;

    // The frame's valid counts laid out by ring and sector. A saturated pixel has no count to add.
    scratch.resize(npix);
    std::vector<uint32_t> count(nkeys, 0);
    for (size_t i = 0; i < npix; i++) {
        const int32_t k = key[i], v = image[i];
        if (k < 0 || v == INT32_MIN || v == INT32_MAX) continue;
        scratch[key_begin[k] + count[k]++] = v;
    }

    std::vector<int32_t> sector_level(nkeys, 0);
    for (size_t k = 0; k < nkeys; k++)
        if (count[k] >= MIN_SECTOR_PIXELS)
            sector_level[k] = median_of(scratch.data() + key_begin[k], count[k]);

    // Each ring's median and its robust spread about that median: off the histogram of its values
    // where both land inside it, otherwise with the ring's sectors packed together - forward into the
    // space the ring owns, so nothing is overwritten before it is moved - and selected.
    std::vector<int32_t> ring_level(nrings, 0);
    std::vector<float> ring_spread(nrings, 0.0f);
    std::vector<char> ring_ok(nrings, 0);
    std::vector<uint32_t> hist(RING_HIST_VALUES);
    for (int r = 0; r < nrings; r++) {
        size_t n = 0;
        for (int s = 0; s < SECTORS; s++)
            n += count[r * SECTORS + s];
        if (n < MIN_RING_PIXELS) continue;
        ring_ok[r] = 1;

        std::fill(hist.begin(), hist.end(), 0);
        size_t below = 0;
        for (int s = 0; s < SECTORS; s++) {
            const size_t k = r * SECTORS + s;
            for (size_t j = 0; j < count[k]; j++) {
                const int32_t v = scratch[key_begin[k] + j];
                if (v < 0) below++;
                else if (v < RING_HIST_VALUES) hist[v]++;
            }
        }
        const auto median = hist_order_statistic(hist, below, n / 2);
        const auto mad = median ? hist_deviation_order_statistic(hist, below, *median, n / 2) : std::nullopt;
        if (median && mad) {
            ring_level[r] = *median;
            ring_spread[r] = 1.4826f * static_cast<float>(*mad);
            continue;
        }

        int32_t *ring = scratch.data() + key_begin[r * SECTORS];
        size_t packed = 0;
        for (int s = 0; s < SECTORS; s++) {
            const size_t k = r * SECTORS + s;
            std::memmove(ring + packed, scratch.data() + key_begin[k], count[k] * sizeof(int32_t));
            packed += count[k];
        }
        ring_level[r] = median_of(ring, n);
        for (size_t j = 0; j < n; j++)
            ring[j] = std::abs(ring[j] - ring_level[r]);
        ring_spread[r] = 1.4826f * static_cast<float>(median_of(ring, n));
    }

    // The level a pixel is judged dark against is its own sector's - the ring's where the sector has
    // too few pixels for one.
    std::vector<int32_t> local_level(nkeys);
    for (size_t k = 0; k < nkeys; k++)
        local_level[k] = count[k] >= MIN_SECTOR_PIXELS ? sector_level[k] : ring_level[k / SECTORS];

    std::vector<int32_t> level;
    std::vector<float> threshold;
    AddLevels(sector_level, ring_level, ring_spread, ring_ok, local_level, level, threshold);

    const size_t rows_per_band = (height + BANDS - 1) / BANDS;
    const size_t first = next_band.fetch_add(1);
    for (size_t b = 0; b < BANDS; b++) {
        const size_t band = (first + b) % BANDS;
        const size_t begin = std::min(npix, band * rows_per_band * width);
        const size_t end = std::min(npix, (band + 1) * rows_per_band * width);
        std::lock_guard lock(band_mutex[band]);
        for (size_t i = begin; i < end; i++) {
            const int32_t k = key[i], v = image[i];
            if (k < 0) continue;
            if (v == INT32_MIN) {
                n_error[i]++;
                if (ring_ok[k / SECTORS]) {   // counted in its ring-sector's frames, but not valid here
                    n_error_ring_ok[i]++;
                    error_level_sum[i] += level[k];
                }
                continue;
            }
            if (!ring_ok[k / SECTORS]) continue;
            sum_value[i] += v;
            if (v == INT32_MAX || static_cast<float>(v) > threshold[k])
                n_lit[i]++;
        }
    }
}

void HotPixelFinder::AddLevels(const std::vector<int32_t> &sector_level, const std::vector<int32_t> &ring_level,
                               const std::vector<float> &ring_spread, const std::vector<char> &ring_ok,
                               const std::vector<int32_t> &local_level,
                               std::vector<int32_t> &level, std::vector<float> &threshold) {
    const size_t nkeys = static_cast<size_t>(nrings) * SECTORS;
    level.assign(nkeys, 0);
    threshold.assign(nkeys, 0.0f);
    for (size_t k = 0; k < nkeys; k++) {
        const int r = static_cast<int>(k / SECTORS);
        level[k] = std::max(ring_level[r], sector_level[k]);
        const float noise = std::max(std::sqrt(static_cast<float>(std::max(level[k], 0))), ring_spread[r]);
        // One rounding for level + nsigma * noise and one for the offset, written out so that every
        // compiler takes the same two - the device takes them too (HotPixelsGPU.cu).
        threshold[k] = std::fma(LIT_NSIGMA, noise, static_cast<float>(level[k])) + LIT_OFFSET;
    }

    // Every sum is an integer, so the result does not depend on the order the frames arrive in.
    std::lock_guard lock(m);
    frames++;
    for (size_t k = 0; k < nkeys; k++)
        if (ring_ok[k / SECTORS]) {
            key_frames[k]++;
            key_level_sum[k] += level[k];
            key_local_level_sum[k] += local_level[k];
        }
}

HotPixelFinder::~HotPixelFinder() {
#ifdef JFJOCH_USE_CUDA
    // Each device half goes with its own card current, where its buffers are.
    for (auto &g : gpus)
        if (g) {
            ScopedGpuPin pin(g->GetDevice());
            g.reset();
        }
#endif
}

#ifdef JFJOCH_USE_CUDA
void HotPixelFinder::PrepareDevice() {
    std::lock_guard lock(m);
    if (!gpus.empty())
        return;
    std::vector<std::future<std::unique_ptr<HotPixelFinderGPU>>> built;
    for (int32_t c = 0; c < get_gpu_count(); c++)
        built.push_back(std::async(std::launch::async, [this, c] {
            set_gpu(c);
            return std::make_unique<HotPixelFinderGPU>(key.get(), width * height, key_begin, nrings, SECTORS,
                                                       HotPixelLevelRules{MIN_SECTOR_PIXELS, MIN_RING_PIXELS,
                                                                          LIT_NSIGMA, LIT_OFFSET});
        }));
    for (auto &b : built)
        gpus.push_back(b.get());
}

void HotPixelFinder::AddDeviceImage(const int32_t *device_image, HotPixelFinderGPU::Frame &frame) {
    PrepareDevice();
    int device = 0;
    cudaGetDevice(&device);
    gpus[device]->Add(device_image, frame);
    std::lock_guard lock(m);
    frames++;
}

HotPixelFinderGPU *HotPixelFinder::FoldedGpu() {
    if (gpus.empty())
        return nullptr;
    for (size_t c = 1; c < gpus.size(); c++) {
        {
            ScopedGpuPin pin(gpus[0]->GetDevice());
            gpus[0]->Fold(*gpus[c]);
        }
        ScopedGpuPin pin(gpus[c]->GetDevice());
        gpus[c].reset();
    }
    gpus.resize(1);
    return gpus[0].get();
}
#endif

HotPixelFinder::Result HotPixelFinder::GetMask(double oscillation_deg, double spacing_deg, size_t nthreads) {
    std::lock_guard lock(m);
    Result ret;
    ret.frames = frames;
    ret.mask.assign(width * height, 0);
    const int n = static_cast<int>(frames);

    // The chance rate per ring, from the pixels lit on no more than half of their frames: whatever
    // lights those - reflections, zingers, noise above the bound - lights a defect-free pixel too.
    // Counted in integers by blocks of rows in parallel, so the totals do not depend on the split.
    std::vector<double> lit(nrings, 0.0), seen(nrings, 0.0);
#ifdef JFJOCH_USE_CUDA
    HotPixelFinderGPU *gpu = FoldedGpu();
    std::optional<ScopedGpuPin> pin;
    if (gpu)
        pin.emplace(gpu->GetDevice());
    if (gpu) {
        std::vector<int64_t> device_lit, device_seen;
        gpu->ChanceCounts(device_lit, device_seen);
        for (int r = 0; r < nrings; r++) {
            lit[r] = static_cast<double>(device_lit[r]);
            seen[r] = static_cast<double>(device_seen[r]);
        }
    } else
#endif
    {
        std::vector<std::vector<int64_t>> block_lit(BANDS), block_seen(BANDS);
        const size_t rows_per_band = (height + BANDS - 1) / BANDS;
        ParallelFor(static_cast<int>(BANDS), nthreads, [&](int b) {
            block_lit[b].assign(nrings, 0);
            block_seen[b].assign(nrings, 0);
            const size_t begin = std::min(width * height, b * rows_per_band * width);
            const size_t end = std::min(width * height, (b + 1) * rows_per_band * width);
            for (size_t i = begin; i < end; i++)
                if (key[i] >= 0 && n_valid(i) > 0 && 2 * n_lit[i] <= n_valid(i)) {
                    block_lit[b][key[i] / SECTORS] += n_lit[i];
                    block_seen[b][key[i] / SECTORS] += n_valid(i);
                }
        });
        for (int r = 0; r < nrings; r++)
            for (size_t b = 0; b < BANDS; b++) {
                lit[r] += static_cast<double>(block_lit[b][r]);
                seen[r] += static_cast<double>(block_seen[b][r]);
            }
    }
    std::vector<int> k_chance(nrings, n + 1);
    for (int r = 0; r < nrings; r++)
        if (seen[r] > 0.0)
            k_chance[r] = binomial_bound(n, std::clamp(lit[r] / seen[r], 1e-6, 0.999),
                                         static_cast<double>(unmasked), FAMILY_WISE_RATE);

    // Persistent: lit on more frames than one reflection or chance explains.
    const int min_valid = std::max(10, n / 2);
    // Whether the pixel's own sums are on the host. A sum of zero reads as dark, so a pixel whose sums
    // stayed on the device must not be judged by the dead rule.
    std::vector<uint8_t> sums_known(width * height, 1);
#ifdef JFJOCH_USE_CUDA
    // With a GPU the per-pixel sums stay there. Only the pixels that can be masked come back - those the
    // tests below could pass (see GetCandidates) - into the host arrays, which are zero everywhere else,
    // so the tests below run on them unchanged.
    if (gpu) {
        const auto c = gpu->GetCandidates(frames, min_valid, spacing_deg > 0.0, k_chance);
        std::fill(sums_known.begin(), sums_known.end(), 0);
        for (size_t j = 0; j < c.index.size(); j++) {
            const size_t i = c.index[j];
            sums_known[i] = 1;
            n_lit[i] = c.n_lit[j];
            n_error[i] = c.n_error[j];
            n_error_ring_ok[i] = c.n_error_ring_ok[j];
            sum_value[i] = c.sum_value[j];
            error_level_sum[i] = c.error_level_sum[j];
        }
        for (size_t k = 0; k < key_frames.size(); k++) {
            key_frames[k] = static_cast<uint16_t>(c.key_frames[k]);
            key_level_sum[k] = c.key_level_sum[k];
            key_local_level_sum[k] = c.key_local_level_sum[k];
        }
    }
#endif
    std::vector<uint8_t> persistent(width * height, 0);
    ParallelChunks(static_cast<int>(height), nthreads, [&](int y0, int y1) {
        for (size_t y = y0; y < static_cast<size_t>(y1); y++)
            for (size_t x = 0; x < width; x++) {
                const size_t i = y * width + x;
                if (key[i] < 0 || n_valid(i) < min_valid || n_lit[i] == 0 || !(spacing_deg > 0.0)) continue;
                // One reflection's stay on this pixel, in sampled frames: 1 / |zeta| turns the rocking
                // width into rotation, and zeta = |axis . (s1 x s0)| with s0 along z.
                const Coord s1 = geometry.LabCoord(static_cast<float>(x), static_cast<float>(y)).Normalize();
                const double zeta = std::max(1e-4, static_cast<double>(std::fabs(axis * (s1 % Coord(0, 0, 1)))));
                const double stay = (oscillation_deg + ROCKING_WIDTH_DEG) / zeta;
                const double k_bragg = 1.0 + std::ceil(stay / spacing_deg);
                const double k = std::min<double>(std::max<double>(k_bragg, k_chance[key[i] / SECTORS]),
                                                  n_valid(i));
                persistent[i] = n_lit[i] >= k;
            }
    });

    // Masked: a persistent pixel whose mean is STRONG_RATIO times its ring's and whose excess is above
    // the Poisson bound - a weaker one cannot make an outlier. Alone or in a patch: a patch lit on frame
    // after frame is stationary in the lab and so no reflection of a rotating crystal, whatever made it,
    // and every reflection that crosses it is integrated on top of it. Measured: a ring of about twenty
    // pixels at 1e5-3e5 counts on every frame (a dead centre in the middle) merged the reflections
    // crossing it to 90 times the strongest real one; the space-group search, which merges without the
    // equivalents an outlier test needs, read the crystal's four-fold off them as unrelated.
    ParallelChunks(static_cast<int>(height), nthreads, [&](int y0, int y1) {
        for (size_t y = y0; y < static_cast<size_t>(y1); y++)
            for (size_t x = 0; x < width; x++) {
                const size_t i = y * width + x;
                if (key[i] < 0) continue;
                if (2 * n_error[i] > n) {
                    ret.mask[i] = 2;
                    continue;
                }
                // Dead: the mirror of hot. A pixel reading at most half of its own sector's level, by
                // more than the bound a hot pixel's excess has to clear, is not recording the light that
                // reaches it, and every reflection whose disk covers it is integrated short by that
                // deficit on every frame - a dead row along a detector edge merged reflections there
                // to large negative intensities. Nothing that is not a defect keeps a pixel dark on
                // frame after frame: a reflection only adds counts. A pixel that also held the error
                // value on some frames is left to the rule above.
                if (sums_known[i] && n_error[i] == 0 && n_valid(i) >= min_valid
                    && 2 * sum_value[i] <= key_local_level_sum[key[i]]) {
                    const double mean_local = static_cast<double>(key_local_level_sum[key[i]]) / n_valid(i);
                    const double deficit = mean_local - static_cast<double>(sum_value[i]) / n_valid(i);
                    if (deficit >= LIT_NSIGMA * std::sqrt(mean_local) + LIT_OFFSET) {
                        ret.mask[i] = 3;
                        continue;
                    }
                }
                if (!persistent[i]) continue;
                const double mean_level = static_cast<double>(sum_level(i)) / n_valid(i);
                const double excess = static_cast<double>(sum_value[i] - sum_level(i)) / n_valid(i);
                if (static_cast<double>(sum_value[i]) >= STRONG_RATIO * static_cast<double>(sum_level(i))
                    && excess >= LIT_NSIGMA * std::sqrt(std::max(mean_level, 0.0)) + LIT_OFFSET)
                    ret.mask[i] = 1;
            }
    });
    for (const uint32_t v : ret.mask) {
        ret.hot += v == 1;
        ret.error += v == 2;
        ret.dead += v == 3;
    }
    return ret;
}