Jungfraujoch/tests/JFJochReceiverProcessingTest.cpp

// Copyright (2019-2024) Paul Scherrer Institute

#include <catch2/catch.hpp>

#include "../receiver/JFJochReceiverTest.h"
#include "../acquisition_device/HLSSimulatedDevice.h"
#include "../writer/HDF5Objects.h"
#include "../writer/HDF5ImagePusher.h"
#include "../receiver/JFJochReceiverService.h"
#include "../common/DiffractionGeometry.h"
#include "../export_images/WriteTIFF.h"
#include <fstream>

TEST_CASE("JFJochIntegrationTest_ZMQ_lysozyme_spot_and_index", "[JFJochReceiver]") {
    Logger logger("JFJochIntegrationTest_ZMQ_lysozyme_spot_and_index");

    RegisterHDF5Filter();

    const uint16_t nthreads = 4;

    DiffractionExperiment experiment(DetectorGeometry(8,2,8,36));
    experiment.ImagesPerTrigger(5).NumTriggers(1).DataFileCount(1).UseInternalPacketGenerator(true)
            .FilePrefix("lyso_test").ConversionOnFPGA(false)
            .DetectorDistance_mm(75).BeamY_pxl(1136).BeamX_pxl(1090).PhotonEnergy_keV(12.4)
            .SetUnitCell(UnitCell{.a = 36.9, .b = 78.95, .c = 78.95, .alpha =90, .beta = 90, .gamma = 90});

    // Load example image
    HDF5ReadOnlyFile data("../../tests/test_data/compression_benchmark.h5");
    HDF5DataSet dataset(data, "/entry/data/data");
    HDF5DataSpace file_space(dataset);

    REQUIRE(file_space.GetDimensions()[2] == experiment.GetXPixelsNum());
    REQUIRE(file_space.GetDimensions()[1] == experiment.GetYPixelsNum());
    std::vector<int16_t> image_conv (file_space.GetDimensions()[1] * file_space.GetDimensions()[2]);

    std::vector<hsize_t> start = {4,0,0};
    std::vector<hsize_t> file_size = {1, file_space.GetDimensions()[1], file_space.GetDimensions()[2]};
    dataset.ReadVector(image_conv, start, file_size);

    std::vector<int16_t> image_raw_geom(experiment.GetModulesNum() * RAW_MODULE_SIZE);
    ConvertedToRawGeometry(experiment, image_raw_geom.data(), image_conv.data());
    logger.Info("Loaded image");

    // Setup acquisition device
    AcquisitionDeviceGroup aq_devices;
    std::unique_ptr<HLSSimulatedDevice> test = std::make_unique<HLSSimulatedDevice>(0, 64);
    for (int m = 0; m < experiment.GetModulesNum(); m++)
        test->SetInternalGeneratorFrame((uint16_t *) image_raw_geom.data() + m * RAW_MODULE_SIZE, m);

    aq_devices.Add(std::move(test));

    JFCalibration calib(experiment);
    HDF5ImagePusher pusher;

    ZMQContext context;
    context.NumThreads(4);

    JFJochReceiverService service(aq_devices, logger, pusher);
    service.NumThreads(nthreads);

    SpotFindingSettings settings = DiffractionExperiment::DefaultDataProcessingSettings();
    settings.signal_to_noise_threshold = 2.5;
    settings.photon_count_threshold = 5;
    settings.min_pix_per_spot = 1;
    settings.max_pix_per_spot = 200;
    service.SetSpotFindingSettings(settings);

    service.Start(experiment, nullptr);
    auto receiver_out = service.Stop();

    std::string jpeg;
    PreviewJPEGSettings jpeg_settings{
            .saturation_value = 10,
            .jpeg_quality = 100,
            .show_spots = true
    };
    REQUIRE_NOTHROW(jpeg = service.GetJPEG(jpeg_settings));
    std::ofstream f("lyso_processing_test_min_pix_1.jpeg", std::ios::binary);
    f.write(jpeg.data(), jpeg.size());

    CHECK(receiver_out.efficiency == 1.0);
    CHECK(receiver_out.status.indexing_rate == 1.0);
    CHECK(receiver_out.status.images_sent == experiment.GetImageNum());
    CHECK(!receiver_out.status.cancelled);
}

TEST_CASE("JFJochIntegrationTest_ZMQ_lysozyme_spot_and_index_min_pix_2", "[JFJochReceiver]") {
    Logger logger("JFJochIntegrationTest_ZMQ_lysozyme_spot_and_index_min_pix_2");

    RegisterHDF5Filter();

    const uint16_t nthreads = 4;

    DiffractionExperiment experiment(DetectorGeometry(8,2,8,36));
    experiment.ImagesPerTrigger(5).NumTriggers(1).DataFileCount(1).UseInternalPacketGenerator(true)
            .FilePrefix("lyso_test_min_pix_2").ConversionOnFPGA(false)
            .DetectorDistance_mm(75).BeamY_pxl(1136).BeamX_pxl(1090).PhotonEnergy_keV(12.4)
            .SetUnitCell(UnitCell{.a = 36.9, .b = 78.95, .c = 78.95, .alpha =90, .beta = 90, .gamma = 90});

    // Load example image
    HDF5ReadOnlyFile data("../../tests/test_data/compression_benchmark.h5");
    HDF5DataSet dataset(data, "/entry/data/data");
    HDF5DataSpace file_space(dataset);

    REQUIRE(file_space.GetDimensions()[2] == experiment.GetXPixelsNum());
    REQUIRE(file_space.GetDimensions()[1] == experiment.GetYPixelsNum());
    std::vector<int16_t> image_conv (file_space.GetDimensions()[1] * file_space.GetDimensions()[2]);

    std::vector<hsize_t> start = {2,0,0};
    std::vector<hsize_t> file_size = {1, file_space.GetDimensions()[1], file_space.GetDimensions()[2]};
    dataset.ReadVector(image_conv, start, file_size);

    std::vector<int16_t> image_raw_geom(experiment.GetModulesNum() * RAW_MODULE_SIZE);
    ConvertedToRawGeometry(experiment, image_raw_geom.data(), image_conv.data());
    logger.Info("Loaded image");

    // Setup acquisition device
    AcquisitionDeviceGroup aq_devices;
    std::unique_ptr<HLSSimulatedDevice> test = std::make_unique<HLSSimulatedDevice>(0, 64);
    for (int m = 0; m < experiment.GetModulesNum(); m++)
        test->SetInternalGeneratorFrame((uint16_t *) image_raw_geom.data() + m * RAW_MODULE_SIZE, m);
    aq_devices.Add(std::move(test));

    JFCalibration calib(experiment);
    HDF5ImagePusher pusher;

    ZMQContext context;
    context.NumThreads(4);

    JFJochReceiverService service(aq_devices, logger, pusher);
    service.NumThreads(nthreads);

    SpotFindingSettings settings = DiffractionExperiment::DefaultDataProcessingSettings();
    settings.signal_to_noise_threshold = 2.5;
    settings.photon_count_threshold = 3;
    settings.min_pix_per_spot = 2;
    settings.max_pix_per_spot = 200;
    service.SetSpotFindingSettings(settings);

    service.Start(experiment, nullptr);
    auto receiver_out = service.Stop();

    std::string jpeg;
    PreviewJPEGSettings jpeg_settings{
        .saturation_value = 10,
        .jpeg_quality = 100,
        .show_spots = true
    };
    REQUIRE_NOTHROW(jpeg = service.GetJPEG(jpeg_settings));
    std::ofstream f("lyso_processing_test_min_pix_2.jpeg", std::ios::binary);
    f.write(jpeg.data(), jpeg.size());

    CHECK(receiver_out.efficiency == 1.0);
    CHECK(receiver_out.status.indexing_rate == 1.0);
    CHECK(receiver_out.status.images_sent == experiment.GetImageNum());
    CHECK(!receiver_out.status.cancelled);
}

TEST_CASE("GenerateResolutionMap") {
    DiffractionExperiment experiment(DetectorGeometry(8,2,8,36));
    experiment.ImagesPerTrigger(5).NumTriggers(1).DataFileCount(1).UseInternalPacketGenerator(true)
            .FilePrefix("lyso_test").ConversionOnFPGA(false)
            .DetectorDistance_mm(75).BeamY_pxl(1136).BeamX_pxl(1090).PhotonEnergy_keV(12.4)
            .SetUnitCell(UnitCell{.a = 36.9, .b = 78.95, .c = 78.95, .alpha =90, .beta = 90, .gamma = 90});

    std::vector<float> spot_finder_resolution_map(experiment.GetModulesNum() * RAW_MODULE_SIZE, 1.0);

    for (int m = 0; m < experiment.GetModulesNum(); m++)
        CalcSpotFinderResolutionMap(spot_finder_resolution_map.data() + m * RAW_MODULE_SIZE, experiment, m);

    std::vector<uint32_t> spot_finder_resolution_map_int(spot_finder_resolution_map.size());
    for (int i = 0; i < spot_finder_resolution_map.size(); i++)
        spot_finder_resolution_map_int[i] = static_cast<uint32_t>(spot_finder_resolution_map[i] * 100);
    std::vector<uint32_t> spot_finder_resolution_map_int_conv(experiment.GetPixelsNum(), 0);

    RawToConvertedGeometry(experiment, spot_finder_resolution_map_int_conv.data(), spot_finder_resolution_map_int.data());

    WriteTIFFToFile("ResolutionMap.tiff", spot_finder_resolution_map_int_conv.data(), experiment.GetXPixelsNum(),
                    experiment.GetYPixelsNum(), 4);
}

TEST_CASE("JFJochIntegrationTest_ZMQ_lysozyme_resolution", "[JFJochReceiver]") {
    Logger logger("JFJochIntegrationTest_ZMQ_lysozyme_resolution");

    RegisterHDF5Filter();

    const uint16_t nthreads = 4;

    DiffractionExperiment experiment(DetectorGeometry(8,2,8,36));
    experiment.ImagesPerTrigger(5).NumTriggers(1).DataFileCount(1).UseInternalPacketGenerator(true)
            .FilePrefix("lyso_test_resolution").ConversionOnFPGA(false)
            .DetectorDistance_mm(75).BeamY_pxl(1136).BeamX_pxl(1090).PhotonEnergy_keV(12.4)
            .SetUnitCell(UnitCell{.a = 36.9, .b = 78.95, .c = 78.95, .alpha =90, .beta = 90, .gamma = 90})
            .NeuralNetModelPath("../../resonet/traced_resnet_model.pt");

    // Load example image
    HDF5ReadOnlyFile data("../../tests/test_data/compression_benchmark.h5");
    HDF5DataSet dataset(data, "/entry/data/data");
    HDF5DataSpace file_space(dataset);

    REQUIRE(file_space.GetDimensions()[2] == experiment.GetXPixelsNum());
    REQUIRE(file_space.GetDimensions()[1] == experiment.GetYPixelsNum());
    std::vector<int16_t> image_conv (file_space.GetDimensions()[1] * file_space.GetDimensions()[2]);

    std::vector<hsize_t> start = {4,0,0};
    std::vector<hsize_t> file_size = {1, file_space.GetDimensions()[1], file_space.GetDimensions()[2]};
    dataset.ReadVector(image_conv, start, file_size);

    std::vector<int16_t> image_raw_geom(experiment.GetModulesNum() * RAW_MODULE_SIZE);
    ConvertedToRawGeometry(experiment, image_raw_geom.data(), image_conv.data());
    logger.Info("Loaded image");

    // Setup acquisition device
    AcquisitionDeviceGroup aq_devices;
    std::unique_ptr<HLSSimulatedDevice> test = std::make_unique<HLSSimulatedDevice>(0, 64);
    for (int m = 0; m < experiment.GetModulesNum(); m++)
        test->SetInternalGeneratorFrame((uint16_t *) image_raw_geom.data() + m * RAW_MODULE_SIZE, m);
    aq_devices.Add(std::move(test));

    JFCalibration calib(experiment);
    HDF5ImagePusher pusher;

    ZMQContext context;
    context.NumThreads(4);

    JFJochReceiverService service(aq_devices, logger, pusher);
    service.NumThreads(nthreads);

    SpotFindingSettings settings = DiffractionExperiment::DefaultDataProcessingSettings();
    settings.signal_to_noise_threshold = 2.5;
    settings.photon_count_threshold = 5;
    settings.min_pix_per_spot = 1;
    settings.max_pix_per_spot = 200;
    service.SetSpotFindingSettings(settings);

    service.Start(experiment, nullptr);
    auto receiver_out = service.Stop();

    CHECK(receiver_out.efficiency == 1.0);
    CHECK(receiver_out.status.indexing_rate == 1.0);
    CHECK(receiver_out.status.images_sent == experiment.GetImageNum());
    CHECK(!receiver_out.status.cancelled);
}