Jungfraujoch/tools/jfjoch_stills_gpu_replay.cpp

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

// Replays the frames of a stored serial collection through the broker's GPU analysis of converted
// JUNGFRAU frames (StillsGPUAnalysis): the pre-pass on the first images, then every frame analysed with the
// beam centre and mask it decided. Reports the pre-pass, the indexed fraction and the sustained rate, and
// writes the mask the frames were analysed with (<prefix>_mask.u32, uint32, row by row).

#include <atomic>
#include <chrono>
#include <cstring>
#include <fstream>
#include <iostream>
#include <thread>
#include <getopt.h>

#include "../common/ImageBuffer.h"
#include "../common/Logger.h"
#include "../image_analysis/indexing/IndexerThreadPool.h"
#include "../reader/JFJochHDF5Reader.h"
#include "../receiver/StillsGPUAnalysis.h"
#include "../writer/HDF5Objects.h"

void print_usage() {
    std::cout << "Usage: jfjoch_stills_gpu_replay [options] <input master>" << std::endl;
    std::cout << "   -o <prefix>  output prefix (default stills_gpu_replay)" << std::endl;
    std::cout << "   -s <image>   first image (default 0)" << std::endl;
    std::cout << "   -n <images>  images to replay (default all)" << std::endl;
    std::cout << "   -r <Hz>      images given per second (default 0 = as fast as they are read)" << std::endl;
    std::cout << "   -N <num>     analysis threads (default 16)" << std::endl;
    std::cout << "   -b <MiB>     image buffer the frames wait in (default 8192)" << std::endl;
    std::cout << "   -x <px>, -y <px>  beam centre, overriding the file" << std::endl;
    std::cout << "   -C <a,b,c,alpha,beta,gamma>  unit cell of the collection" << std::endl;
    std::cout << "   -F           the broker's default fixed-threshold spot finder (default: adaptive, as the receiver)" << std::endl;
    std::cout << "   -c           leave the file's unit cell out of the collection" << std::endl;
}

int main(int argc, char **argv) {
    Logger logger("jfjoch_stills_gpu_replay");
    RegisterHDF5Filter();

    std::string prefix = "stills_gpu_replay";
    int64_t first = 0;
    std::optional<int64_t> images;
    double rate = 0.0;
    int nthreads = 16;
    size_t buffer_MiB = 8192;
    std::optional<float> beam_x, beam_y;
    std::optional<UnitCell> cell;
    bool no_cell = false;
    bool fixed_threshold = false;

    int opt;
    while ((opt = getopt(argc, argv, "o:s:n:r:N:b:x:y:C:cF")) != -1) {
        switch (opt) {
            case 'o': prefix = optarg; break;
            case 's': first = atoll(optarg); break;
            case 'n': images = atoll(optarg); break;
            case 'r': rate = atof(optarg); break;
            case 'N': nthreads = atoi(optarg); break;
            case 'b': buffer_MiB = atoll(optarg); break;
            case 'x': beam_x = atof(optarg); break;
            case 'y': beam_y = atof(optarg); break;
            case 'c': no_cell = true; break;
            case 'F': fixed_threshold = true; break;
            case 'C': {
                UnitCell c{};
                if (sscanf(optarg, "%f,%f,%f,%f,%f,%f", &c.a, &c.b, &c.c, &c.alpha, &c.beta, &c.gamma) != 6) {
                    print_usage();
                    return EXIT_FAILURE;
                }
                cell = c;
                break;
            }
            default: print_usage(); return EXIT_FAILURE;
        }
    }
    if (optind + 1 != argc) {
        print_usage();
        return EXIT_FAILURE;
    }

    JFJochHDF5Reader reader;
    reader.ReadFile(argv[optind]);
    const auto dataset = reader.GetDataset();
    const int64_t n = std::min<int64_t>(images.value_or(reader.GetNumberOfImages() - first),
                                        reader.GetNumberOfImages() - first);
    const auto stored = reader.GetStoredPixelFormat();

    DiffractionExperiment x(dataset->experiment);
    x.ImagesPerTrigger(n).NumTriggers(1);
    if (beam_x)
        x.BeamX_pxl(*beam_x);
    if (beam_y)
        x.BeamY_pxl(*beam_y);
    if (no_cell) {
        // Indexed de novo, as the broker does where no cell is stated.
        x.SetUnitCell(std::nullopt);
        IndexingSettings indexing = x.GetIndexingSettings();
        indexing.Algorithm(IndexingAlgorithmEnum::Auto);
        x.ImportIndexingSettings(indexing);
    }
    if (cell)
        x.SetUnitCell(cell);
    const PixelMask &mask = *dataset->pixel_mask;

    StartMessage start;
    x.FillMessage(start);
    start.number_of_images = n;
    start.bit_depth_image = stored.bit_depth;
    start.pixel_signed = stored.is_signed;
    start.pixel_mask["default"] = mask.GetMask(x);

    SpotFindingSettings spot_finding = DiffractionExperiment::DefaultDataProcessingSettings();
    // As JFJochReceiverFPGA analyses converted frames.
    if (!fixed_threshold) {
        spot_finding.adaptive_threshold = true;
        spot_finding.min_pix_per_spot = std::nullopt;
    }

    IndexerThreadPool pool(x.GetIndexingSettings());

    // The frames wait in an image buffer, as in the receiver; here every slot can hold one.
    ImageBuffer buffer(buffer_MiB * 1024 * 1024);
    buffer.StartMeasurement(x);

    std::atomic<int64_t> analysed = 0, indexed = 0, spots = 0;
    std::atomic<double> processing_s = 0.0;
    std::chrono::steady_clock::time_point t_ready;
    const auto t0 = std::chrono::steady_clock::now();
    {
        StillsGPUAnalysis gpu(x, mask, start, &pool, nthreads, buffer.GetStatus().total_slots, logger,
                              [&](const StillsGPUAnalysis::Setup &s) {
                                  t_ready = std::chrono::steady_clock::now();
                                  const auto &m = s.mask.GetMask();
                                  std::ofstream f(prefix + "_mask.u32", std::ios::binary);
                                  f.write(reinterpret_cast<const char *>(m.data()), m.size() * sizeof(uint32_t));
                                  logger.Info("Ready after {:.1f} s: beam centre {:.2f} {:.2f}",
                                              std::chrono::duration<double>(t_ready - t0).count(),
                                              s.experiment.GetBeamX_pxl(), s.experiment.GetBeamY_pxl());
                              },
                              [&](DataMessage &m, const AzimuthalIntegrationProfile &) {
                                  ++analysed;
                                  spots += static_cast<int64_t>(m.spots.size());
                                  processing_s += m.processing_time_s.value_or(0.0f);
                                  if (m.indexing_result.value_or(false))
                                      ++indexed;
                              },
                              [&] { return spot_finding; });

        JFJochReaderRawImage raw;
        for (int64_t i = 0; i < n; i++) {
            if (rate > 0)
                std::this_thread::sleep_until(t0 + std::chrono::duration<double>(static_cast<double>(i) / rate));
            reader.ReadRawImage(first + i, raw);
            auto held = std::make_shared<HeldImage>();
            // Unlike the receiver, the replay waits for a slot rather than lose the frame.
            while ((held->slot = buffer.GetImageSlot()) == nullptr)
                std::this_thread::sleep_for(std::chrono::milliseconds(1));
            memcpy(held->slot->GetImage(), raw.image.GetCompressed(), raw.image.GetCompressedSize());
            held->message.number = i;
            held->message.image = CompressedImage(held->slot->GetImage(), raw.image.GetCompressedSize(),
                                                  raw.image.GetWidth(), raw.image.GetHeight(), raw.image.GetMode(),
                                                  raw.image.GetCompressionAlgorithm());
            held->message.image_collection_efficiency = 1.0f;
            while (!gpu.Hold(held))
                std::this_thread::sleep_for(std::chrono::milliseconds(1));
        }
        logger.Info("Gave {} images in {:.1f} s", n,
                    std::chrono::duration<double>(std::chrono::steady_clock::now() - t0).count());
        gpu.Finish();
    }
    const auto t1 = std::chrono::steady_clock::now();
    const double after_ready = std::chrono::duration<double>(t1 - t_ready).count();
    logger.Info("Analysed {} images, {} indexed ({:.1f}%); {:.1f} s after the pre-pass = {:.1f} Hz; total {:.1f} s",
                analysed.load(), indexed.load(), 100.0 * indexed / std::max<int64_t>(1, analysed), after_ready,
                analysed / after_ready, std::chrono::duration<double>(t1 - t0).count());
    logger.Info("Per image: {:.1f} ms analysis (one thread), {:.1f} spots",
                1e3 * processing_s / std::max<int64_t>(1, analysed),
                static_cast<double>(spots) / std::max<int64_t>(1, analysed));
    return EXIT_SUCCESS;
}