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 <sstream>
#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;
    std::cout << "   -S <group>   space group of the collection (the merge's)" << std::endl;
    std::cout << "   -M           no merge (StillsGPUAnalysis merge = false)" << std::endl;
    std::cout << "   -e <images>  the first images given as empty frames (no diffraction)" << 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;
    std::optional<std::string> space_group;
    bool merge = true;
    int64_t empty_frames = 0;

    int opt;
    while ((opt = getopt(argc, argv, "o:s:n:r:N:b:x:y:C:cFS:Me:")) != -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 'S': space_group = optarg; break;
            case 'M': merge = false; break;
            case 'e': empty_frames = atoll(optarg); 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);
    if (space_group)
        x.SetSpaceGroup(*gemmi::find_spacegroup_by_name(*space_group));
    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();
    {
        // As the receiver: half the buffer for the frames waiting, the other half for those sent.
        StillsGPUAnalysis gpu(x, mask, start, &pool, nthreads, buffer.GetStatus().total_slots / 2, 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; }, {}, nullptr, merge);

        JFJochReaderRawImage raw;
        std::vector<uint8_t> empty;
        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);
            if (i < empty_frames) {
                // An empty frame of the same size and pixel type: no compression, every pixel zero.
                empty.assign(raw.image.GetUncompressedSize(), 0);
                raw.image = CompressedImage(empty.data(), empty.size(), raw.image.GetWidth(), raw.image.GetHeight(),
                                            raw.image.GetMode());
            }
            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();
        if (const auto c = gpu.GetConsensusUnitCell())
            logger.Info("Consensus cell {:.3f} {:.3f} {:.3f} {:.2f} {:.2f} {:.2f}", c->a, c->b, c->c, c->alpha, c->beta,
                        c->gamma);
        if (const auto &m = gpu.GetMerged()) {
            const auto &o = m->merge.statistics.overall;
            logger.Info("Merge: {} crystals, {} unique to {:.2f} A, CC1/2 {:.4f}, R_split {:.4f}, <I/sigma> {:.2f}, "
                        "{:.2f} s after the last frame was analysed", m->images, o.unique_reflections, o.d_min,
                        o.cc_half, o.r_split, o.mean_i_over_sigma, m->seconds);
            std::ostringstream os;
            os << m->merge.statistics;
            logger.Info("Merging statistics:\n{}", os.str());
        }
    }
    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;
}