Jungfraujoch/writer/HDF5DataFilePluginMX.cpp

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

#include "HDF5DataFilePluginMX.h"
#define RESERVE_IMAGES 1000

inline std::string to_bravais_code(const std::optional<LatticeMessage> &lm_opt) {
    if (!lm_opt.has_value()) return "";
    const auto &lm = lm_opt.value();
    const char C = lm.centering;

    switch (lm.crystal_system) {
        case gemmi::CrystalSystem::Triclinic:
            if (C == 'P') return "aP";
            break;
        case gemmi::CrystalSystem::Monoclinic:
            if (C == 'P') return "mP";
            if (C == 'A') return "mA";
            if (C == 'B') return "mB";
            if (C == 'C') return "mC";
            break;
        case gemmi::CrystalSystem::Orthorhombic:
            if (C == 'P') return "oP";
            if (C == 'A') return "oA";
            if (C == 'B') return "oB";
            if (C == 'C') return "oC";
            if (C == 'F') return "oF";
            if (C == 'I') return "oI";
            break;
        case gemmi::CrystalSystem::Tetragonal:
            if (C == 'P') return "tP";
            if (C == 'I') return "tI";
            break;
        case gemmi::CrystalSystem::Trigonal:
            if (C == 'R') return "hR"; // rhombohedral in hex setting
            break;
        case gemmi::CrystalSystem::Hexagonal:
            if (C == 'P') return "hP";
            break;
        case gemmi::CrystalSystem::Cubic:
            if (C == 'P') return "cP";
            if (C == 'F') return "cF";
            if (C == 'I') return "cI";
            break;
    }
    return "";
}

HDF5DataFilePluginMX::HDF5DataFilePluginMX(const StartMessage &msg)
    : max_spots(msg.max_spot_count), indexing(msg.indexing_algorithm != IndexingAlgorithmEnum::None) {
}

void HDF5DataFilePluginMX::OpenFile(HDF5File &data_file, const DataMessage &msg) {
    bkg_estimate.reserve(RESERVE_IMAGES);

    if (max_spots == 0)
        return;

    spot_x.reserve(max_spots * RESERVE_IMAGES);
    spot_y.reserve(max_spots * RESERVE_IMAGES);
    spot_int.reserve(max_spots * RESERVE_IMAGES);
    spot_indexed.reserve(max_spots * RESERVE_IMAGES);
    spot_ice_ring.reserve(max_spots * RESERVE_IMAGES);
    spot_h.reserve(max_spots * RESERVE_IMAGES);
    spot_k.reserve(max_spots * RESERVE_IMAGES);
    spot_l.reserve(max_spots * RESERVE_IMAGES);
    spot_dist_ewald.reserve(max_spots * RESERVE_IMAGES);

    beam_corr_x.reserve(RESERVE_IMAGES);
    beam_corr_y.reserve(RESERVE_IMAGES);

    npeaks.reserve(RESERVE_IMAGES);
    strong_pixel_count.reserve(RESERVE_IMAGES);
    indexed.reserve(RESERVE_IMAGES);
    profile_radius.reserve(RESERVE_IMAGES);
    mosaicity_deg.reserve(RESERVE_IMAGES);
    b_factor.reserve(RESERVE_IMAGES);
    indexed_lattice.reserve(9 * RESERVE_IMAGES);
    resolution_estimate.reserve(RESERVE_IMAGES);

    spot_count_total.reserve(RESERVE_IMAGES);
    spot_count_ice.reserve(RESERVE_IMAGES);
    spot_count_indexed.reserve(RESERVE_IMAGES);
    spot_count_low_res.reserve(RESERVE_IMAGES);
}

void HDF5DataFilePluginMX::Write(const DataMessage &msg, uint64_t image_number) {
    if (msg.bkg_estimate.has_value())
        bkg_estimate[image_number] = msg.bkg_estimate.value();

    if (max_spots == 0)
        return;

    if (image_number >= max_image_number) {
        max_image_number = image_number;

        spot_x.resize(max_spots * (max_image_number + 1));
        spot_y.resize(max_spots * (max_image_number + 1));
        spot_int.resize(max_spots * (max_image_number + 1));
        spot_ice_ring.resize(max_spots * (max_image_number + 1));
        spot_h.resize(max_spots * (max_image_number + 1));
        spot_k.resize(max_spots * (max_image_number + 1));
        spot_l.resize(max_spots * (max_image_number + 1));
        spot_dist_ewald.resize(max_spots * (max_image_number + 1));

        if (indexing)
            spot_indexed.resize(max_spots * (max_image_number + 1));
        if (indexing)
            indexed_lattice.resize((max_image_number + 1) * 9);
    }

    uint32_t spot_cnt = std::min(msg.spots.size(), max_spots);

    for (int i = 0; i < spot_cnt; i++) {
        spot_x[max_spots * image_number + i] = msg.spots[i].x;
        spot_y[max_spots * image_number + i] = msg.spots[i].y;
        spot_int[max_spots * image_number + i] = msg.spots[i].intensity;
        if (indexing) {
            spot_indexed[max_spots * image_number + i] = msg.spots[i].indexed;
            spot_h[max_spots * image_number + i] = msg.spots[i].h;
            spot_k[max_spots * image_number + i] = msg.spots[i].k;
            spot_l[max_spots * image_number + i] = msg.spots[i].l;
            spot_dist_ewald[max_spots * image_number + i] = msg.spots[i].dist_ewald_sphere;
        }
        spot_ice_ring[max_spots * image_number + i] = msg.spots[i].ice_ring;
    }

    npeaks[image_number] = spot_cnt;
    strong_pixel_count[image_number] = msg.strong_pixel_count.value_or(0);
    spot_count_total[image_number] = msg.spot_count.value_or(0);
    spot_count_ice[image_number] = msg.spot_count_ice_rings.value_or(0);
    spot_count_low_res[image_number] = msg.spot_count_low_res.value_or(0);

    if (indexing) {
        indexed[image_number] = msg.indexing_result.value_or(0);
        profile_radius[image_number] = msg.profile_radius.value_or(NAN);
        mosaicity_deg[image_number] = msg.mosaicity_deg.value_or(NAN);
        b_factor[image_number] = msg.b_factor.value_or(NAN);
        resolution_estimate[image_number] = msg.resolution_estimate.value_or(NAN);
        beam_corr_x[image_number] = msg.beam_corr_x.value_or(NAN);
        beam_corr_y[image_number] = msg.beam_corr_y.value_or(NAN);
        spot_count_indexed[image_number] = msg.spot_count_indexed.value_or(0);

        if (msg.indexing_lattice) {
            auto tmp = msg.indexing_lattice->GetVector();

            for (int i = 0; i < 9; i++)
                indexed_lattice[image_number * 9 + i] = tmp[i];
        } else {
            for (int i = 0; i < 9; i++)
                indexed_lattice[image_number * 9 + i] = NAN;
        }

        if (msg.lattice_type) {
            bravais_lattice[image_number] = to_bravais_code(msg.lattice_type);
            niggli_class[image_number] = msg.lattice_type->niggli_class;
        } else {
            bravais_lattice[image_number] = "";
            niggli_class[image_number] = 0;
        }
    }
}

void HDF5DataFilePluginMX::WriteFinal(HDF5File &data_file) {
    HDF5Group(data_file, "/entry/MX").NXClass("NXcollection");

    if (!spot_x.empty()) {
        data_file.SaveVector("/entry/MX/peakXPosRaw", spot_x, {(hsize_t) (max_image_number + 1), max_spots});
        data_file.SaveVector("/entry/MX/peakYPosRaw", spot_y, {(hsize_t) (max_image_number + 1), max_spots});
        data_file.SaveVector("/entry/MX/peakTotalIntensity", spot_int, {(hsize_t) (max_image_number + 1), max_spots});
        data_file.SaveVector("/entry/MX/peakIceRingRes", spot_ice_ring, {(hsize_t) (max_image_number + 1), max_spots});
        data_file.SaveVector("/entry/MX/nPeaks", npeaks.vec());

        if (indexing) {
            data_file.SaveVector("/entry/MX/peakH", spot_h, {(hsize_t) (max_image_number + 1), max_spots});
            data_file.SaveVector("/entry/MX/peakK", spot_k, {(hsize_t) (max_image_number + 1), max_spots});
            data_file.SaveVector("/entry/MX/peakL", spot_l, {(hsize_t) (max_image_number + 1), max_spots});
            data_file.SaveVector("/entry/MX/peakDistEwaldSphere", spot_dist_ewald,
                                 {(hsize_t) (max_image_number + 1), max_spots});
            data_file.SaveVector("/entry/MX/peakIndexed", spot_indexed, {(hsize_t) (max_image_number + 1), max_spots});
        }
    }


    if (!strong_pixel_count.empty())
        data_file.SaveVector("/entry/MX/strongPixels", strong_pixel_count.vec());

    if (!spot_count_ice.empty())
        data_file.SaveVector("/entry/MX/peakCountIceRingRes", spot_count_ice.vec());
    if (!spot_count_indexed.empty())
        data_file.SaveVector("/entry/MX/peakCountIndexed", spot_count_indexed.vec());
    if (!spot_count_low_res.empty())
        data_file.SaveVector("/entry/MX/peakCountLowRes", spot_count_low_res.vec());
    if (!spot_count_total.empty())
        data_file.SaveVector("/entry/MX/peakCountUnfiltered", spot_count_total.vec());

    if (!indexed.empty())
        data_file.SaveVector("/entry/MX/imageIndexed", indexed.vec());
    if (!indexed_lattice.empty())
        data_file.SaveVector("/entry/MX/latticeIndexed", indexed_lattice, {(hsize_t) (max_image_number + 1), 9})
                ->Units("Angstrom");
    if (!bkg_estimate.empty())
        data_file.SaveVector("/entry/MX/bkgEstimate", bkg_estimate.vec());
    if (!profile_radius.empty())
        data_file.SaveVector("/entry/MX/profileRadius", profile_radius.vec())->Units("Angstrom^-1");
    if (!mosaicity_deg.empty())
        data_file.SaveVector("/entry/MX/mosaicity", mosaicity_deg.vec())->Units("deg");
    if (!b_factor.empty())
        data_file.SaveVector("/entry/MX/bFactor", b_factor.vec())->Units("Angstrom^2");

    if (!beam_corr_x.empty())
        data_file.SaveVector("/entry/MX/beam_corr_x", beam_corr_x.vec())->Units("pixel");
    if (!beam_corr_y.empty())
        data_file.SaveVector("/entry/MX/beam_corr_y", beam_corr_y.vec())->Units("pixel");
    if (!niggli_class.empty())
        data_file.SaveVector("/entry/MX/niggli_class", niggli_class.vec());
    if (!bravais_lattice.empty())
        data_file.SaveVector("/entry/MX/bravais_lattice", bravais_lattice.vec());
    if (!resolution_estimate.empty())
        data_file.SaveVector("/entry/MX/resolutionEstimate", resolution_estimate.vec())->Units("Angstrom");
}