Jungfraujoch/image_analysis/MXAnalysisWithoutFPGA.cpp

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

#include "MXAnalysisWithoutFPGA.h"

#include <algorithm>
#include <spdlog/spdlog.h>

#include "spot_finding/StrongPixelSet.h"
#include "../compression/JFJochDecompress.h"
#include "../common/CUDAWrapper.h"

#include "spot_finding/SpotUtils.h"
#include "bragg_prediction/BraggPredictionFactory.h"
#include "image_preprocessing/ImagePreprocessorCPU.h"

#include "azint/AzIntEngineCPU.h"
#include "roi/ROIIntegrationCPU.h"
#include "spot_finding/ImageSpotFinderCPU.h"
#include "spot_finding/AdaptiveSpotFinderCPU.h"
#include "bragg_integration/BraggIntegrationEngineCPU.h"
#ifdef JFJOCH_USE_CUDA
#include "azint/AzIntEngineGPU.h"
#include "roi/ROIIntegrationGPU.h"
#include "spot_finding/ImageSpotFinderGPU.h"
#include "spot_finding/AdaptiveSpotFinderGPU.h"
#include "image_preprocessing/ImagePreprocessorGPU.h"
#include "image_preprocessing/ImagePreprocessorBufferGPU.h"
#include "bragg_integration/BraggIntegrationEngineGPU.h"
#include "../common/CUDAWrapper.h"
#endif


MXAnalysisWithoutFPGA::MXAnalysisWithoutFPGA(const DiffractionExperiment &in_experiment,
                                             const AzimuthalIntegrationMapping &in_integration,
                                             const PixelMask &in_mask,
                                             IndexAndRefine &in_indexer,
                                             bool in_enable_fused_adaptive_gpu)
    : experiment(in_experiment),
      integration(in_integration),
      enable_fused_adaptive_gpu(in_enable_fused_adaptive_gpu),
      npixels(experiment.GetPixelsNum()),
      xpixels(experiment.GetXPixelsNum()),
      indexer(in_indexer),
      prediction(CreateBraggPrediction(experiment.IsRotationIndexing())),
      mask(in_mask),
      mask_high_res(-1),
      mask_low_res(-1) {
#ifdef JFJOCH_USE_CUDA
    if (get_gpu_count() == 0) {
#endif
        preprocessor_buffer = std::make_unique<ImagePreprocessorBuffer>(experiment.GetPixelsNum());
        auto cpu_preprocessor = std::make_unique<ImagePreprocessorCPU>(in_experiment, in_mask);
        preprocessor_cpu = cpu_preprocessor.get();
        preprocessor = std::move(cpu_preprocessor);
        bragg_engine = std::make_unique<BraggIntegrationEngineCPU>(in_experiment, in_mask);
        if (experiment.ROI().size() >= 1)
            roi = std::make_unique<ROIIntegrationCPU>(experiment);
#ifdef JFJOCH_USE_CUDA
    } else {
        stream = std::make_shared<CudaStream>();
        // The host copy of the preprocessed image is only read when a CPU engine wants it, which is
        // the same condition that drives copy_image_to_host below. Skipping it also skips page-locking
        // 4 bytes per pixel per worker.
        preprocessor_buffer = std::make_unique<ImagePreprocessorBufferGPU>(
                experiment.GetPixelsNum(), /*host_mirror=*/!enable_fused_adaptive_gpu);
        // The preprocessed image only has to come back to the host if a CPU engine reads it. Every
        // engine built below runs on the GPU, except the CPU adaptive finder that is kept when the fused
        // GPU engine is off - so that is the one case that needs the copy. Every caller currently passes
        // enable_fused_adaptive_gpu = true, so on the GPU path the copy is off in practice.
        preprocessor = std::make_unique<ImagePreprocessorGPU>(in_experiment, in_mask, stream,
                                                             /*copy_image_to_host=*/!enable_fused_adaptive_gpu);
        bragg_engine = std::make_unique<BraggIntegrationEngineGPU>(in_experiment, stream, in_mask);
        if (experiment.ROI().size() >= 1)
            roi = std::make_unique<ROIIntegrationGPU>(experiment, stream);
        if (enable_fused_adaptive_gpu) {
            // One GPU engine that computes the azimuthal profile and the adaptive spot mask in a single
            // image pass. fused_adaptive aliases it so Analyze() can lift the profile out of it.
            auto fused = std::make_unique<AdaptiveSpotFinderGPU>(integration, stream);
            fused_adaptive = fused.get();
            adaptiveSpotFinder = std::move(fused);
        }
    }
#endif
    if (!adaptiveSpotFinder) {
        auto cpu = std::make_unique<AdaptiveSpotFinderCPU>(integration);
        // Where the azimuthal integration is the CPU engine too, the finder's first ring pass is the
        // same pass over the same pixels, and gives the profile (see AdaptiveSpotFinderCPU).
        bool cpu_azint = true;
#ifdef JFJOCH_USE_CUDA
        cpu_azint = !stream;
#endif
        if (cpu_azint) {
            cpu->FuseAzimuthalIntegration(true);
            fused_adaptive_cpu = cpu.get();
        }
        adaptiveSpotFinder = std::move(cpu);
    }
}

void MXAnalysisWithoutFPGA::Analyze(DataMessage &output,
                                    AzimuthalIntegrationProfile &profile,
                                    const SpotFindingSettings &spot_finding_settings) {
    if ((output.image.GetWidth() != xpixels)
        || (output.image.GetWidth() * output.image.GetHeight() != npixels))
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Mismatch in pixel size");

    // Decompress on the device where the preprocessor can, so only the compressed chunk crosses PCIe
    // and the host does no decompression at all. AnalyzeCompressed says whether it took the image;
    // when it declines (a CPU preprocessor, or an algorithm with no device decoder) fall through to
    // the host route unchanged. The two produce the same preprocessed image.
    const auto compression_start_time = std::chrono::steady_clock::now();
    ImageStatistics ret{};
    bool decoded_on_device = false;
    try {
        decoded_on_device = preprocessor->AnalyzeCompressed(*preprocessor_buffer, output.image, ret);
    } catch (const JFJochException &e) {
        // The device route must never be the reason a frame fails: whatever it could not handle, the
        // host decoder gets its turn. If the data really is bad the host throws too and the caller
        // sees the same error it saw before any of this existed - but a GPU-side problem costs speed
        // rather than the acquisition. Unless the context is lost, which the host route cannot survive.
        cuda_throw_if_context_lost();
        spdlog::warn("Device decoding failed ({}), falling back to host decompression", e.what());
        decoded_on_device = false;
        // Handled here, so it must not stay behind as the thread's last error: the host route below
        // runs kernels of its own and checks cudaGetLastError() after them, which would report this
        // failure again - fatally - over an image that decoded perfectly well on the host.
        cuda_clear_error();
    }
    const auto compression_end_time = std::chrono::steady_clock::now();

    // The fused GPU engine (rugnux offline, GPU, adaptive detection) produces the azimuthal profile as
    // a byproduct of spot finding, so the separate azint pass is skipped in that case and the profile is
    // lifted out of the finder below.
    const bool fused = enable_fused_adaptive_gpu && spot_finding_settings.enable
                    && spot_finding_settings.adaptive_threshold && fused_adaptive != nullptr;
    const bool fused_cpu = spot_finding_settings.enable && spot_finding_settings.adaptive_threshold
                           && fused_adaptive_cpu != nullptr;

    if (preprocessor_cpu) {
        const auto preprocessing_start_time = std::chrono::steady_clock::now();
        ret = PreprocessCPU(output.image, fused_cpu);
        const auto preprocessing_end_time = std::chrono::steady_clock::now();
        output.preprocessing_time_s = std::chrono::duration<float>(preprocessing_end_time - preprocessing_start_time).count();
        // Decoded inside the preprocessing pass, whose time includes it.
        if (output.image.GetCompressionAlgorithm() != CompressionAlgorithm::NO_COMPRESSION)
            output.compression_time_s = 0.0f;
    } else if (!decoded_on_device) {
        const uint8_t *image_ptr = Decompress(output.image);
        const auto decompressed_time = std::chrono::steady_clock::now();
        if (output.image.GetCompressionAlgorithm() != CompressionAlgorithm::NO_COMPRESSION)
            output.compression_time_s = std::chrono::duration<float>(decompressed_time - compression_start_time).count();

        const auto preprocessing_start_time = std::chrono::steady_clock::now();
        ret = preprocessor->Analyze(*preprocessor_buffer, image_ptr, output.image.GetMode());
        const auto preprocessing_end_time = std::chrono::steady_clock::now();
        output.preprocessing_time_s = std::chrono::duration<float>(preprocessing_end_time - preprocessing_start_time).count();
    } else {
        // Decode and preprocess are one device operation here, but the decompression is still a real,
        // separately measurable cost - the decoder brackets it with CUDA events - so it is still
        // reported as one. Leaving compression_time_s unset instead would blank the broker's
        // "compression" plot trace and fill /entry/profiling/compressionTime with NaN.
        const float total_s = std::chrono::duration<float>(compression_end_time - compression_start_time).count();
        const float decompress_s = std::min(preprocessor->GetLastDecompressionTime_s(), total_s);
        output.compression_time_s = decompress_s;
        output.preprocessing_time_s = total_s - decompress_s;
    }

    if (!fused && !fused_cpu) {
        const auto azint_start_time = std::chrono::steady_clock::now();
        AzInt().Run(*preprocessor_buffer, profile);
        const auto azint_end_time = std::chrono::steady_clock::now();
        output.azint_time_s = std::chrono::duration<float>(azint_end_time - azint_start_time).count();
    }

    if (roi)
        roi->Run(*preprocessor_buffer, output.roi);

    if (spot_finding_settings.enable) {
        // Update resolution mask
        if (mask_high_res != spot_finding_settings.high_resolution_limit
            || mask_low_res != spot_finding_settings.low_resolution_limit)
            UpdateMaskResolution(spot_finding_settings);

        ImageSpotFinder &finder = spot_finding_settings.adaptive_threshold
                                      ? static_cast<ImageSpotFinder &>(*adaptiveSpotFinder)
                                      : FixedThresholdFinder();
        const auto integrate_fn = [this](const std::vector<Reflection> &predicted, size_t npredicted,
                                         int64_t image_number) {
            if (count_crowding)
                crowding_engine->CountNeighbourCrowding(predicted, npredicted);
            return bragg_engine->Run(*preprocessor_buffer, predicted, npredicted, image_number);
        };
        // The supercell probe integrates on the same engine, but what the engine counts is the run's
        // measurement of its own stencil, and the probe's reflections are not the run's: keep them out.
        const auto probe_integrate_fn = [this](const std::vector<Reflection> &predicted, size_t npredicted,
                                               int64_t image_number) {
            const BraggIntegrationCounts before = bragg_engine->Counts();
            auto ret = bragg_engine->Run(*preprocessor_buffer, predicted, npredicted, image_number);
            const BraggIntegrationCounts after = bragg_engine->Counts();
            probe_counts.predicted += after.predicted - before.predicted;
            probe_counts.bkg_starved += after.bkg_starved - before.bkg_starved;
            probe_counts.bkg_starved_by_neighbour += after.bkg_starved_by_neighbour - before.bkg_starved_by_neighbour;
            probe_counts.profile_fallback += after.profile_fallback - before.profile_fallback;
            probe_counts.overlapped += after.overlapped - before.overlapped;
            probe_counts.overlapped_major += after.overlapped_major - before.overlapped_major;
            return ret;
        };

        // The radial background correction has to be decided BEFORE this image is integrated, so the
        // ice score is taken here rather than with the other per-image quantities at the end of the
        // function. It needs the peak-excluded per-ring background, which the adaptive finder has as
        // soon as it has detected - so this must be called after detection and before integration.
        // Where no such background exists (no adaptive finder), auto leaves the correction off: the
        // plain profile carries the Bragg peaks and cannot support an absolute threshold.
        const auto decide_radial_background = [this, &spot_finding_settings, &output]() {
            if (!bragg_engine->IsBackgroundRadialAuto())
                return;
            // Same condition as the score at the end of this function: the adaptive finder holds its
            // ring background from whenever it last ran, so requiring that it ran for THIS image is
            // what keeps a stale curve out.
            if (!spot_finding_settings.adaptive_threshold)
                return;
            const std::vector<float> &ring_bkg = adaptiveSpotFinder->GetRingBackground();
            if (ring_bkg.empty())
                return;
            output.ice_ring_score = AzimuthalIntegrationProfile::IceRingScore(
                    ring_bkg, integration.GetQBinCount(), integration.Settings(),
                    spot_finding_settings.ice_ring_width_Q_recipA);
            bragg_engine->BackgroundRadial(*output.ice_ring_score
                                           >= experiment.GetScalingSettings().GetIceMinScore());
        };

        // A missing min-pix (std::nullopt) means "choose it per image". This applies only to the stills
        // indexing path (each frame is indexed independently); rotation indexing builds one lattice from
        // all frames, so it keeps the fixed min-pix and the single-pass finder.
        const bool adaptive_min_pix = !spot_finding_settings.min_pix_per_spot.has_value()
                                   && spot_finding_settings.indexing
                                   && !experiment.IsRotationIndexing();
        if (adaptive_min_pix) {
            // Choose the per-image min-pix adaptively instead of a fixed one. min-pix filters connected
            // components AFTER detection, so BOTH the detection (the expensive per-pixel pass) and the
            // connected-component search run ONCE and only the filter is repeated; the azimuthal
            // profile is the one Detect() computed.
            // Index at 3/2/1 (index-only, no integration/accumulation), keep whichever maximises
            // n_indexed^2 / n_total (indexed count weighted by indexed fraction) together with its spot
            // list, and integrate that one.
            const auto detect_start_time = std::chrono::steady_clock::now();
            finder.Detect(*preprocessor_buffer, spot_finding_settings);
            const auto &components = finder.ExtractComponents(*preprocessor_buffer, spot_finding_settings);
            float spot_finding_time_s =
                    std::chrono::duration<float>(std::chrono::steady_clock::now() - detect_start_time).count();
            float indexing_time_s = 0.0f;

            SpotFindingSettings s = spot_finding_settings;
            std::vector<DiffractionSpot> best_spots;
            int best_mp = 0;
            double best_score = -1.0;
            for (int mp : {3, 2, 1}) {
                s.min_pix_per_spot = mp;
                const auto extract_start_time = std::chrono::steady_clock::now();
                std::vector<DiffractionSpot> spots = ImageSpotFinder::Filter(components, s);
                spot_finding_time_s +=
                        std::chrono::duration<float>(std::chrono::steady_clock::now() - extract_start_time).count();
                SpotAnalyze(experiment, s, spots, output);
                const bool indexed = indexer.IndexFrameOnly(output, s);
                indexing_time_s += output.indexing_time_s.value_or(0.0f);
                if (indexed) {
                    const double n_idx = static_cast<double>(output.spot_count_indexed.value_or(0));
                    const double n_tot = static_cast<double>(std::max<int64_t>(1, output.spot_count.value_or(1)));
                    const double score = n_idx * n_idx / n_tot;
                    if (score > best_score) {
                        best_score = score;
                        best_mp = mp;
                        best_spots = std::move(spots);
                    }
                }
            }
            if (best_mp != 0) {
                // Index and integrate the winning spot list; no spot finding left to do.
                s.min_pix_per_spot = best_mp;
                SpotAnalyze(experiment, s, best_spots, output);
                decide_radial_background();
                indexer.ProcessImage(output, s, *prediction, integrate_fn, probe_integrate_fn);
                indexing_time_s += output.indexing_time_s.value_or(0.0f);
            }
            // Each indexer call reports only its own time, so the escalation's total is summed here.
            output.spot_finding_time_s = spot_finding_time_s;
            output.indexing_time_s = indexing_time_s;
        } else {
            const auto spot_finding_start_time = std::chrono::steady_clock::now();
            const std::vector<DiffractionSpot> spots = finder.Run(*preprocessor_buffer, spot_finding_settings);
            SpotAnalyze(experiment, spot_finding_settings, spots, output);
            output.spot_finding_time_s = std::chrono::duration<float>(std::chrono::steady_clock::now() - spot_finding_start_time).count();
            decide_radial_background();
            if (spot_finding_settings.indexing)
                indexer.ProcessImage(output, spot_finding_settings, *prediction, integrate_fn, probe_integrate_fn);
        }

        // Recorded whichever way the frame went. A frame holding StrongPixelLimit of them is given up
        // on and reports no spots at all, and this is the only thing that tells such a frame from one
        // that did not diffract.
        output.strong_pixel_count = finder.StrongPixelCount();

#ifdef JFJOCH_USE_CUDA
        if (fused) {
            // Lift the azimuthal profile the fused engine computed in the same detection pass; its azint
            // cost is folded into spot_finding_time_s above.
            profile.Clear(integration);
            profile += fused_adaptive->GetProfile();
            output.azint_time_s = 0.0f;
        }
#endif
        if (fused_cpu) {
            fused_adaptive_cpu->GetProfile(profile);
            output.azint_time_s = 0.0f;
        }
    }

    output.max_viable_pixel_value = ret.max_value;
    output.min_viable_pixel_value = ret.min_value;
    output.error_pixel_count = ret.error_pixel_count;
    output.saturated_pixel_count = ret.saturated_pixel_count;
    output.az_int_profile = profile.GetResult();
    output.az_int_profile_count = profile.GetPixelCount();
    output.az_int_profile_std = profile.GetStd();

    output.bkg_estimate = profile.GetBkgEstimate(integration.Settings());

    // The ice score wants a radial profile with the Bragg peaks taken OUT of it. The azimuthal profile
    // is a plain per-ring mean, so a strong low-resolution reflection landing in a ring's bin is
    // indistinguishable from ice sitting there - measured, that alone lifts clean crystals to a score
    // of 1.5-4.2, right into the range real ice occupies. The adaptive spot finder already computes
    // exactly what is wanted: a sigma-clipped per-ring background, in the same bins, from which the
    // peaks have been removed (an ice ring is azimuthally smooth, so it survives the clip). It is in
    // raw counts rather than corrected ones, which the score does not care about - it is a ratio to the
    // background interpolated under the ring, and the corrections are smooth in radius.
    const std::vector<float> &ring_bkg = adaptiveSpotFinder->GetRingBackground();
    const bool have_ring_bkg = spot_finding_settings.enable && spot_finding_settings.adaptive_threshold
                            && !ring_bkg.empty();
    output.ice_ring_score = AzimuthalIntegrationProfile::IceRingScore(
            have_ring_bkg ? ring_bkg : profile.GetResult1D(), integration.GetQBinCount(),
            integration.Settings(), spot_finding_settings.ice_ring_width_Q_recipA);
}

ImageSpotFinder &MXAnalysisWithoutFPGA::FixedThresholdFinder() {
    if (!spotFinder) {
#ifdef JFJOCH_USE_CUDA
        if (stream)
            spotFinder = std::make_unique<ImageSpotFinderGPU>(experiment.GetXPixelsNum(),
                                                              experiment.GetYPixelsNum(), stream);
        else
#endif
            spotFinder = std::make_unique<ImageSpotFinderCPU>(experiment.GetXPixelsNum(),
                                                              experiment.GetYPixelsNum());
        // It missed every mask update that happened before it existed, so it takes the current one
        // now. Without this it would find spots outside the resolution limits.
        if (mask_resolution)
            spotFinder->SetResolutionMaskBits(*mask_resolution);
    }
    return *spotFinder;
}

AzIntEngine &MXAnalysisWithoutFPGA::AzInt() {
    if (!azint) {
#ifdef JFJOCH_USE_CUDA
        if (stream)
            azint = std::make_unique<AzIntEngineGPU>(integration, stream);
        else
#endif
            azint = std::make_unique<AzIntEngineCPU>(integration);
    }
    return *azint;
}

void MXAnalysisWithoutFPGA::RebuildROI() {
    if (experiment.ROI().empty()) {
        roi.reset();
        return;
    }
#ifdef JFJOCH_USE_CUDA
    if (stream) {
        roi = std::make_unique<ROIIntegrationGPU>(experiment, stream);
        return;
    }
#endif
    roi = std::make_unique<ROIIntegrationCPU>(experiment);
}

void MXAnalysisWithoutFPGA::AnalyzeROIOnly(DataMessage &output) {
    if ((output.image.GetWidth() != xpixels)
        || (output.image.GetWidth() * output.image.GetHeight() != npixels))
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Mismatch in pixel size");

    const uint8_t *image_ptr = Decompress(output.image);
    preprocessor->Analyze(*preprocessor_buffer, image_ptr, output.image.GetMode());
    RunROIOnly(output);
}

const uint8_t *MXAnalysisWithoutFPGA::Decompress(const CompressedImage &image) {
    // An uncompressed image is read straight out of the message and never touches decompression_buffer,
    // so what the upload reads is the reader's own bytes: page-lock those instead. Pageable memory is
    // staged by the driver a chunk at a time, which on a 4-byte 16 Mpx frame is 12.9 ms against 4.0 ms
    // pinned - and a miniCBF sweep takes this path for every image.
    if (image.GetCompressionAlgorithm() == CompressionAlgorithm::NO_COMPRESSION)
        preprocessor->PinInputRegion(image.GetCompressed(), image.GetUncompressedSize());
    else
        preprocessor->PinInputBuffer(decompression_buffer, image.GetUncompressedSize());
    return image.GetUncompressedPtr(decompression_buffer);
}

ImageStatistics MXAnalysisWithoutFPGA::PreprocessCPU(const CompressedImage &image, bool ring_pass) {
    ImageStatistics ret{};
    if (ring_pass)
        fused_adaptive_cpu->BeginRings();
    const auto process_block = [&](size_t first, const uint8_t *data, size_t n) {
        preprocessor_cpu->AnalyzeBlock(*preprocessor_buffer, first, data, n, image.GetMode(), ret);
        if (ring_pass)
            fused_adaptive_cpu->AccumulateRingsBlock(*preprocessor_buffer, first, n);
    };
    if (JFJochIsBitshuffle(image.GetCompressionAlgorithm()))
        JFJochDecompressBlocks(image.GetCompressionAlgorithm(), image.GetCompressed(), image.GetCompressedSize(),
                               image.GetWidth() * image.GetHeight() * image.GetNumChannels(), image.GetByteDepth(),
                               process_block);
    else
        process_block(0, Decompress(image), npixels);
    return ret;
}

void MXAnalysisWithoutFPGA::RunROIOnly(DataMessage &output) {
    output.roi.clear();
    if (roi)
        roi->Run(*preprocessor_buffer, output.roi);
}

BraggIntegrationCounts MXAnalysisWithoutFPGA::BraggCounts() const {
    if (!bragg_engine)
        return {};
    BraggIntegrationCounts c = bragg_engine->Counts();
    c.predicted -= probe_counts.predicted;
    c.bkg_starved -= probe_counts.bkg_starved;
    c.bkg_starved_by_neighbour -= probe_counts.bkg_starved_by_neighbour;
    c.profile_fallback -= probe_counts.profile_fallback;
    c.overlapped -= probe_counts.overlapped;
    c.overlapped_major -= probe_counts.overlapped_major;
    return c;
}

void MXAnalysisWithoutFPGA::CountCrowdingAt(const BraggIntegrationSettings &settings) {
    DiffractionExperiment x = experiment;
    x.ImportBraggIntegrationSettings(settings);
    crowding_engine = std::make_unique<BraggIntegrationEngineCPU>(x, mask);
}

BraggIntegrationCounts MXAnalysisWithoutFPGA::CrowdingCounts() const {
    return crowding_engine ? crowding_engine->Counts() : BraggIntegrationCounts{};
}

void MXAnalysisWithoutFPGA::UpdateMaskResolution(const SpotFindingSettings &settings) {
    mask_low_res = settings.low_resolution_limit;
    mask_high_res = settings.high_resolution_limit;
    // The mask is a pure function of the resolution map and the two limits, so the mapping builds it -
    // once for all the workers, which otherwise each walked every pixel of the detector to arrive at
    // the same bits.
    mask_resolution = integration.ResolutionMaskBits(mask_high_res, mask_low_res);

    // The finders keep their own copy (the GPU ones a bit-packed device copy), so the mask is handed
    // over here - when the limits change - rather than with every image.
    if (spotFinder)
        spotFinder->SetResolutionMaskBits(*mask_resolution);
    adaptiveSpotFinder->SetResolutionMaskBits(*mask_resolution);
}