Jungfraujoch/writer/HDF5DataFilePluginReflection.cpp

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

#include "HDF5DataFilePluginReflection.h"
#include "../include/spdlog/fmt/fmt.h"
#include "../common/Reflection.h"

void HDF5DataFilePluginReflection::OpenFile(HDF5File &data_file, const DataMessage &msg, size_t images_per_file) {
    reflection_group = std::make_unique<HDF5Group>(data_file, "/entry/reflections");
    reflection_group->NXClass("NXcollection");
}

void HDF5DataFilePluginReflection::Write(const DataMessage &msg, uint64_t image_number) {
    if (!reflection_group || msg.reflections.empty())
        return;

    std::vector<int32_t> h, k, l;
    std::vector<float> I, sigma, d;
    // The three corrections are stored as their reciprocals, in double so that the reciprocal of the
    // value read back is the float the reflection held.
    std::vector<double> lp, qe, flight;
    std::vector<float> image, phi, pred_x, pred_y, obs_x, obs_y, bkg, var_bkg, partiality, zeta,
            scale_factor, dist_ewald;
    std::vector<uint8_t> observed, clipped, overloaded;

    h.reserve(msg.reflections.size());
    k.reserve(msg.reflections.size());
    l.reserve(msg.reflections.size());
    I.reserve(msg.reflections.size());
    sigma.reserve(msg.reflections.size());
    d.reserve(msg.reflections.size());
    pred_x.reserve(msg.reflections.size());
    pred_y.reserve(msg.reflections.size());
    obs_x.reserve(msg.reflections.size());
    obs_y.reserve(msg.reflections.size());
    bkg.reserve(msg.reflections.size());
    var_bkg.reserve(msg.reflections.size());
    lp.reserve(msg.reflections.size());
    qe.reserve(msg.reflections.size());
    flight.reserve(msg.reflections.size());
    partiality.reserve(msg.reflections.size());
    image.reserve(msg.reflections.size());
    phi.reserve(msg.reflections.size());
    zeta.reserve(msg.reflections.size());
    scale_factor.reserve(msg.reflections.size());
    dist_ewald.reserve(msg.reflections.size());
    observed.reserve(msg.reflections.size());
    clipped.reserve(msg.reflections.size());
    overloaded.reserve(msg.reflections.size());

    for (const auto &refl : msg.reflections) {
        image.emplace_back(refl.image_number);
        h.emplace_back(refl.h);
        k.emplace_back(refl.k);
        l.emplace_back(refl.l);
        I.emplace_back(refl.I);
        sigma.emplace_back(refl.sigma);
        d.emplace_back(refl.d);
        pred_x.emplace_back(refl.predicted_x);
        pred_y.emplace_back(refl.predicted_y);
        obs_x.emplace_back(refl.observed_x);
        obs_y.emplace_back(refl.observed_y);
        bkg.emplace_back(refl.bkg);
        var_bkg.emplace_back(refl.var_bkg);
        lp.emplace_back(1.0/refl.prescaling_corr);
        qe.emplace_back(1.0/refl.qe_corr);
        flight.emplace_back(1.0/refl.flight_corr);
        partiality.emplace_back(refl.partiality);
        phi.emplace_back(refl.delta_phi_deg);
        zeta.emplace_back(refl.zeta);
        scale_factor.emplace_back(refl.image_scale_corr);
        dist_ewald.emplace_back(refl.dist_ewald);
        observed.emplace_back(refl.observed);
        clipped.emplace_back(refl.clipped);
        overloaded.emplace_back(refl.overloaded);
    }

    std::string image_group_name = fmt::format("image_{:06d}", image_number);
    HDF5Group image_group(*reflection_group, image_group_name);
    image_group.NXClass("NXreflections");
    image_group.SaveVector("h", h);
    image_group.SaveVector("k", k);
    image_group.SaveVector("l", l);
    image_group.SaveVector("d", d)->Units("Angstrom");
    image_group.SaveVector("delta_phi", phi);
    image_group.SaveVector("predicted_x", pred_x);
    image_group.SaveVector("predicted_y", pred_y);
    image_group.SaveVector("observed_x", obs_x);
    image_group.SaveVector("observed_y", obs_y);
    image_group.SaveVector("int_sum", I);
    image_group.SaveVector("int_err", sigma);
    image_group.SaveVector("background_mean", bkg);
    image_group.SaveVector("background_variance", var_bkg);
    image_group.SaveVector("observed_frame", image);
    image_group.SaveVector("lp", lp);
    image_group.SaveVector("qe", qe);
    image_group.SaveVector("flight", flight);
    image_group.SaveVector("partiality", partiality);
    image_group.SaveVector("zeta", zeta);
    image_group.SaveVector("image_scale_corr", scale_factor);
    image_group.SaveVector("excitation_error", dist_ewald)->Units("Angstrom^-1");
    image_group.SaveVector("observed", observed);
    image_group.SaveVector("clipped", clipped);
    image_group.SaveVector("overloaded", overloaded);
}

void HDF5DataFilePluginReflection::WriteFinal(HDF5File &data_file) {

}