Jungfraujoch/tests/WriteReflectionsTest.cpp

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

#include <catch2/catch_all.hpp>

#include <algorithm>
#include <filesystem>
#include <fstream>
#include <string>
#include <cstring>

#include <gemmi/mtz.hpp>
#include <gemmi/read_cif.hpp>
#include <gemmi/numb.hpp>

#include "../common/DiffractionExperiment.h"
#include "../image_analysis/WriteReflections.h"
#include "../image_analysis/IntegrationOutcome.h"
#include "../image_analysis/scale_merge/Merge.h"
#include "../image_analysis/scale_merge/TwinningAnalysis.h"
#include <cmath>
#include "SyntheticMergedReflections.h"

namespace {
    // The wavelength CCP4's mtzlib substitutes when the data columns belong to dataset 0, which it
    // takes for the reserved HKL_base: Cu K-alpha, and the wrong edge for anything reading f'/f''
    // out of the file.
    constexpr double CU_KALPHA_A = 1.54187;

    constexpr UnitCell TETRAGONAL_CELL{47.0f, 47.0f, 63.0f, 90.0f, 90.0f, 90.0f};

    DiffractionExperiment TestExperiment() {
        DiffractionExperiment x;
        x.IncidentEnergy_keV(12.7f);         // ~0.976 A, nowhere near the Cu K-alpha default
        x.SpaceGroupNumber(96);              // P 43 21 2
        x.SetUnitCell(TETRAGONAL_CELL);
        return x;
    }
}

TEST_CASE("Merged MTZ: the data dataset is id 1 and carries the wavelength", "[write_reflections][portable]") {
    jfjoch_test::SyntheticMergeParams params;
    params.true_space_group = "P 43 21 2";
    params.twin_supergroup = "P 43 21 2";
    params.d_min_A = 5.0;
    const auto reflections = jfjoch_test::GenerateSyntheticMerged(params);
    REQUIRE(!reflections.empty());

    const auto experiment = TestExperiment();
    const auto path = (std::filesystem::temp_directory_path() / "rugnux_merged_wavelength.mtz").string();
    WriteMtzReflections(reflections, TETRAGONAL_CELL, experiment, OutputProvenance{}, path);
    const gemmi::Mtz mtz = gemmi::read_mtz_file(path);
    std::filesystem::remove(path);

    // HKL_base at id 0, the data at id 1. A data dataset written at id 0 occupies the id MTZ
    // reserves for the base, and mtzlib then reports CU_KALPHA_A instead of the real wavelength.
    REQUIRE(mtz.datasets.size() == 2);
    CHECK(mtz.datasets[0].id == 0);
    CHECK(mtz.datasets[0].dataset_name == "HKL_base");
    CHECK(mtz.datasets[1].id == 1);
    CHECK(mtz.datasets[1].wavelength == Catch::Approx(experiment.GetWavelength_A()).epsilon(1e-5));
    CHECK(mtz.datasets[1].wavelength != Catch::Approx(CU_KALPHA_A).epsilon(1e-3));

    // The wavelength is read off the dataset the data columns belong to, so they have to be on the
    // data dataset and not on the base.
    for (const char *label : {"IMEAN", "SIGIMEAN", "F", "SIGF", "FreeR_flag"}) {
        const gemmi::Mtz::Column *col = mtz.column_with_label(label);
        REQUIRE(col != nullptr);
        CHECK(col->dataset_id == 1);
    }

    // Cell and space group travel in the same header.
    CHECK(mtz.spacegroup != nullptr);
    CHECK(mtz.spacegroup->number == 96);
    CHECK(mtz.cell.a == Catch::Approx(TETRAGONAL_CELL.a).epsilon(1e-5));
    CHECK(mtz.cell.c == Catch::Approx(TETRAGONAL_CELL.c).epsilon(1e-5));
    CHECK(mtz.cell.gamma == Catch::Approx(90.0).epsilon(1e-5));
    CHECK(mtz.datasets[1].cell.a == Catch::Approx(TETRAGONAL_CELL.a).epsilon(1e-5));

    CHECK(mtz.nreflections == static_cast<int>(reflections.size()));
}

TEST_CASE("Unmerged MTZ: LP is Lorentz-polarization, QE the sensor efficiency, FLIGHT the flight path",
          "[write_reflections][portable]") {
    // The whole point of the split: LP must mean what XDS and DIALS mean by it, and the raw count
    // sum must still be recoverable from the file alone, as I / LP * QE * FLIGHT.
    auto experiment = TestExperiment();
    experiment.Goniometer(GoniometerAxis("omega", 0.0f, 0.1f, Coord(-1, 0, 0), {}));

    IntegrationOutcome outcome;
    const float raw[3] = {1000.0f, 250.0f, 40.0f};
    const float lp[3]  = {1.75f, 2.50f, 0.90f};    // Lorentz x polarization, and nothing else
    const float qe[3]  = {0.9375f, 0.8125f, 1.0f};    // 1.0 = the sensor said nothing to correct
    const float fl[3]  = {1.0125f, 1.0400f, 1.0f};    // >= 1: an oblique reflection crossed more air
    for (int i = 0; i < 3; ++i) {
        Reflection r{};
        r.h = 4 + i; r.k = 2; r.l = 6;
        r.image_number = static_cast<float>(i);
        r.d = 5.0f + i;
        r.I = raw[i];                                 // the writer is what applies the factor
        r.sigma = std::sqrt(raw[i]);
        r.prescaling_corr = lp[i];
        r.qe_corr = qe[i];
        r.flight_corr = fl[i];
        r.partiality = 1.0f;
        r.predicted_x = 100.0f + i; r.predicted_y = 200.0f + i;
        r.observed_x = NAN; r.observed_y = NAN;
        outcome.reflections.push_back(r);
    }

    const auto path = (std::filesystem::temp_directory_path() / "rugnux_unmerged_qe.mtz").string();
    WriteUnmergedMtzReflections({outcome}, TETRAGONAL_CELL, experiment, false, OutputProvenance{}, path);
    const gemmi::Mtz mtz = gemmi::read_mtz_file(path);
    std::filesystem::remove(path);

    const gemmi::Mtz::Column *c_I = mtz.column_with_label("I");
    const gemmi::Mtz::Column *c_lp = mtz.column_with_label("LP");
    const gemmi::Mtz::Column *c_qe = mtz.column_with_label("QE");
    const gemmi::Mtz::Column *c_fl = mtz.column_with_label("FLIGHT");
    REQUIRE(c_I != nullptr);
    REQUIRE(c_lp != nullptr);
    REQUIRE(c_qe != nullptr);          // DIALS writes this column even when there is nothing in it
    REQUIRE(c_fl != nullptr);
    REQUIRE(mtz.nreflections == 3);

    for (int i = 0; i < 3; ++i) {
        const float I = mtz.data[i * mtz.columns.size() + c_I->idx];
        const float LP = mtz.data[i * mtz.columns.size() + c_lp->idx];
        const float QE = mtz.data[i * mtz.columns.size() + c_qe->idx];
        const float FL = mtz.data[i * mtz.columns.size() + c_fl->idx];
        INFO("row " << i);
        // LP holds Lorentz x polarization alone: the sensor term was never inside it.
        CHECK(LP == Catch::Approx(lp[i]).epsilon(1e-5));
        // QE is a divisor normalised to 1 at normal incidence, so it never drops below 1.
        CHECK(QE == Catch::Approx(1.0f / qe[i]).epsilon(1e-5));
        CHECK(QE >= 1.0f);
        // FLIGHT is a divisor in the same convention, and runs the other way: the sensor favours an
        // oblique reflection, the medium attenuates it, so this one never rises above 1.
        CHECK(FL == Catch::Approx(1.0f / fl[i]).epsilon(1e-5));
        CHECK(FL <= 1.0f);
        // ... and the three together put the raw counts back.
        CHECK(I / LP * QE * FL == Catch::Approx(raw[i]).epsilon(1e-4));
        // The intensity itself is the fully corrected value - all three applied.
        CHECK(I == Catch::Approx(raw[i] * lp[i] * qe[i] * fl[i]).epsilon(1e-5));
    }
}

TEST_CASE("Unmerged MTZ: built on several workers, the same rows as on one", "[write_reflections][portable]") {
    // A sweep large enough that the sort, the event sums and the row fill are all split into
    // pieces, with events of several parts, events cut short by the sweep's ends, and parts that
    // share (h,k,l) and image_number - the tie the sort has to break the same way every time.
    auto experiment = TestExperiment();
    experiment.Goniometer(GoniometerAxis("omega", 0.0f, 0.1f, Coord(-1, 0, 0), {}));

    uint32_t seed = 12345;
    const auto next = [&seed] { seed = seed * 1664525u + 1013904223u; return (seed >> 8) / 16777216.0f; };

    std::vector<IntegrationOutcome> outcomes(200);
    for (int f = 0; f < 200; ++f) {
        for (int q = 0; q < 8000; ++q) {
            if ((f + q) % 7 >= 3)
                continue;   // each reflection is seen on three consecutive images out of seven
            Reflection r{};
            r.h = q % 20 - 10; r.k = (q / 20) % 20 - 10; r.l = q / 400 + 1;
            r.image_number = static_cast<float>(f);
            r.I = 1000.0f * next();
            r.sigma = 1.0f + 10.0f * next();
            r.bkg = next();
            r.var_bkg = next();
            r.prescaling_corr = 1.0f + next();
            r.partiality = 0.2f + 0.2f * next();
            r.predicted_x = 1000.0f * next(); r.predicted_y = 1000.0f * next();
            r.observed_x = NAN; r.observed_y = NAN;
            outcomes[f].reflections.push_back(r);
            if (q % 97 == 0) {
                r.I = 1000.0f * next();
                outcomes[f].reflections.push_back(r);
            }
        }
    }

    for (bool sum_partials : {true, false}) {
        INFO("sum_partials " << sum_partials);
        const MtzTable serial = UnmergedMtzTable(outcomes, TETRAGONAL_CELL, experiment, sum_partials,
                                                 OutputProvenance{}, 1);
        const MtzTable parallel = UnmergedMtzTable(outcomes, TETRAGONAL_CELL, experiment, sum_partials,
                                                   OutputProvenance{}, 8);
        REQUIRE(serial.columns[0].values.size() > 100000);
        REQUIRE(parallel.columns.size() == serial.columns.size());
        // Bitwise, so the NaN of a missing value compares equal to itself.
        for (size_t c = 0; c < serial.columns.size(); c++) {
            REQUIRE(parallel.columns[c].values.size() == serial.columns[c].values.size());
            CHECK(std::memcmp(parallel.columns[c].values.data(), serial.columns[c].values.data(),
                              serial.columns[c].values.size() * sizeof(float)) == 0);
        }
        REQUIRE(parallel.batches.size() == serial.batches.size());
        CHECK(parallel.batches.front().number == serial.batches.front().number);
        CHECK(parallel.batches.back().number == serial.batches.back().number);
    }
}

TEST_CASE("SHELX .hkl of unmerged fulls: fixed 3I4,2F8.2, scaled to fit, 0 0 0 terminator",
          "[write_reflections][portable]") {
    // Two fulls of the same reflection and a Friedel mate: written as measured, nothing averaged. The
    // largest |I| or sigma is put at 9999.00; an unusable full (sigma not positive) is left out.
    const std::vector<ScaledFull> fulls{
            {1, 2, 3, 200000.0f, 2000.0f},
            {-1, -2, -3, 100000.0f, 1200.0f},
            {1, 2, 3, -50.0f, 20.0f},
            {4, 5, 6, 10.0f, 0.0f},
    };
    const auto path = (std::filesystem::temp_directory_path() / "jfjoch_scaled_fulls_test.hkl").string();
    WriteShelxHklReflections(fulls, path, 2);

    std::vector<std::string> lines;
    {
        // Closed before the remove: Windows refuses to delete a file that is still open.
        std::ifstream in(path);
        for (std::string line; std::getline(in, line);)
            lines.push_back(line);
    }
    std::filesystem::remove(path);

    REQUIRE(lines.size() == 4);
    for (const auto &line : lines)
        CHECK(line.size() == 28);
    CHECK(lines[0] == "   1   2   3 9999.00   99.99");
    CHECK(lines[1] == "  -1  -2  -3 4999.50   59.99");
    CHECK(lines[2] == "   1   2   3   -2.50    1.00");
    CHECK(lines[3] == "   0   0   0    0.00    0.00");
}

TEST_CASE("CC* from CC1/2", "[write_reflections][portable]") {
    CHECK(CCStar(1.0) == Catch::Approx(1.0));
    CHECK(CCStar(0.5) == Catch::Approx(std::sqrt(2.0 / 3.0)));
    CHECK(std::isnan(CCStar(0.0)));
    CHECK(std::isnan(CCStar(-0.1)));
    CHECK(std::isnan(CCStar(NAN)));
}

TEST_CASE("Merged mmCIF: PDBx keys, merging R-factors and provenance",
          "[write_reflections][portable]") {
    jfjoch_test::SyntheticMergeParams params;
    params.true_space_group = "P 43 21 2";
    params.twin_supergroup = "P 43 21 2";
    params.d_min_A = 5.0;
    const auto reflections = jfjoch_test::GenerateSyntheticMerged(params);
    REQUIRE(!reflections.empty());

    auto experiment = TestExperiment();
    experiment.SampleTemperature_K(100.0f);

    MergeStatistics statistics;
    MergeStatisticsShell shell;
    shell.d_min = 5.0f; shell.d_max = 30.0f;
    shell.total_observations = 400; shell.unique_reflections = 100; shell.possible_unique_reflections = 110;
    shell.mean_i_over_sigma = 10.0;
    shell.cc_half = 0.98;
    shell.r_merge = 0.05; shell.r_meas = 0.06; shell.r_pim = 0.03;
    statistics.shells = {shell};
    statistics.overall = shell;

    OutputProvenance provenance;
    provenance.collection_date = "2026-07-13T13:48:52.413Z";
    provenance.command_line = "rugnux -o 'my run' /data/x_master.h5";
    provenance.input_file = "/data/x_master.h5";
    provenance.input_identity = "1234 bytes, modified 2026-07-13T14:00:00.000Z";

    const auto path = (std::filesystem::temp_directory_path() / "rugnux_merged_test.cif").string();
    WriteMmcifReflections(reflections, TETRAGONAL_CELL, experiment, statistics, ErrorModelReport{},
                          TwinningAnalysisResult{}, provenance, path, 2);
    gemmi::cif::Document doc = gemmi::read_cif_gz(path);
    std::filesystem::remove(path);

    REQUIRE(doc.blocks.size() == 1);
    gemmi::cif::Block &merged = doc.blocks[0];
    CHECK(merged.name == "sample");
    CHECK(*merged.find_value("_reflns.pdbx_ordinal") == "1");
    CHECK(*merged.find_value("_reflns.pdbx_diffrn_id") == "1");
    CHECK(*merged.find_value("_reflns.pdbx_Rmerge_I_obs") == "0.0500");
    CHECK(*merged.find_value("_reflns.pdbx_Rpim_I_all") == "0.0300");
    CHECK(gemmi::cif::as_number(*merged.find_value("_reflns.pdbx_CC_star"))
          == Catch::Approx(std::sqrt(2 * 0.98 / 1.98)).margin(1e-4));
    CHECK(*merged.find_value("_diffrn.ambient_temp") == "100.0");
    CHECK(*merged.find_value("_diffrn_detector.pdbx_collection_date") == "2026-07-13");
    // Provenance: the command line survives its own quotes, as a text field.
    CHECK(gemmi::cif::as_string(*merged.find_value("_jfjoch_provenance.command_line"))
          == provenance.command_line);
    CHECK(gemmi::cif::as_string(*merged.find_value("_jfjoch_provenance.input_file")) == provenance.input_file);
    CHECK(gemmi::cif::as_string(*merged.find_value("_jfjoch_provenance.input_identity"))
          == provenance.input_identity);

    gemmi::cif::Table software = merged.find("_software.", {"pdbx_ordinal", "classification"});
    REQUIRE(software.length() == 2);
    CHECK(gemmi::cif::as_string(software[0][1]) == "data reduction");
    CHECK(gemmi::cif::as_string(software[1][1]) == "data scaling");

    gemmi::cif::Table shells = merged.find("_reflns_shell.", {"pdbx_ordinal", "Rmerge_I_obs",
                                                                     "pdbx_Rpim_I_all", "pdbx_CC_star"});
    REQUIRE(shells.length() == 1);
    CHECK(shells[0][1] == "0.0500");
    CHECK(shells[0][2] == "0.0300");

}

TEST_CASE("MTZ history: command line and input, in records of at most 80 characters",
          "[write_reflections][portable]") {
    jfjoch_test::SyntheticMergeParams params;
    params.true_space_group = "P 43 21 2";
    params.twin_supergroup = "P 43 21 2";
    params.d_min_A = 5.0;
    const auto reflections = jfjoch_test::GenerateSyntheticMerged(params);

    OutputProvenance provenance;
    // A space falls exactly on the first record's last character, which a reader would strip.
    provenance.command_line = "rugnux -o out " + std::string(56, 'x') + " " + std::string(150, 'x') + "_master.h5";
    provenance.input_file = "/data/" + std::string(150, 'x') + "_master.h5";
    provenance.input_identity = "1234 bytes, modified 2026-07-13T14:00:00.000Z";

    const auto path = (std::filesystem::temp_directory_path() / "rugnux_history.mtz").string();
    WriteMtzReflections(reflections, TETRAGONAL_CELL, TestExperiment(), provenance, path);
    const gemmi::Mtz mtz = gemmi::read_mtz_file(path);
    std::filesystem::remove(path);

    REQUIRE(!mtz.history.empty());
    CHECK(mtz.history.front().rfind("From Rugnux ", 0) == 0);
    std::string joined;
    for (const auto &line : mtz.history) {
        CHECK(line.size() <= 80);
        joined += line;
    }
    // Each value is all there, run on over records that start with '+'.
    std::string command;
    for (size_t i = 0; i < mtz.history.size(); ++i)
        if (mtz.history[i].rfind("Command: ", 0) == 0) {
            command = mtz.history[i].substr(9);
            for (size_t j = i + 1; j < mtz.history.size() && mtz.history[j].rfind("+", 0) == 0; ++j)
                command += mtz.history[j].substr(1);
        }
    CHECK(command == provenance.command_line);
    CHECK(joined.find("Input file: " + provenance.input_identity) != std::string::npos);
    CHECK(mtz.history.size() <= 30);
}

TEST_CASE("Unmerged MTZ: on stills, only the crystals on the run's cell", "[write_reflections][portable]") {
    // A still indexed with its c axis doubled names other reflections by the same indices, so its rows
    // would be read as the wrong reflections; it is left out, the crystal on the run's cell is written.
    const auto experiment = TestExperiment();   // no goniometer: stills
    std::vector<IntegrationOutcome> outcomes(2);
    outcomes[0].latt = CrystalLattice(TETRAGONAL_CELL);
    outcomes[1].latt = CrystalLattice(UnitCell{47.0f, 47.0f, 126.0f, 90.0f, 90.0f, 90.0f});
    for (int f = 0; f < 2; ++f)
        for (int l = 1; l <= 5; ++l) {
            Reflection r{};
            r.h = 1; r.k = 2; r.l = l;
            r.image_number = static_cast<float>(f);
            r.I = 100.0f;
            r.sigma = 10.0f;
            r.prescaling_corr = 1.0f;
            r.partiality = 1.0f;
            outcomes[f].reflections.push_back(r);
        }
    const MtzTable mtz = UnmergedMtzTable(outcomes, TETRAGONAL_CELL, experiment, false, OutputProvenance{}, 1);
    const auto batch = std::find_if(mtz.columns.begin(), mtz.columns.end(),
                                    [](const MtzTable::Column &c) { return c.label == "BATCH"; });
    REQUIRE(batch != mtz.columns.end());
    CHECK(batch->values.size() == 5);
    for (const float b : batch->values)
        CHECK(b == 1.0f);
}

TEST_CASE("Unmerged MTZ: the merge's scaled fulls replace the summed events, those it did not keep left out",
          "[write_reflections][portable]") {
    auto experiment = TestExperiment();
    experiment.Goniometer(GoniometerAxis("omega", 0.0f, 0.1f, Coord(-1, 0, 0), {}));

    // (1,2,3) crosses twice, on frames 10-12 and 100-102; (2,3,4) once, on frames 50-51
    std::vector<IntegrationOutcome> outcomes(110);
    const auto add = [&](int frame, int h, int k, int l) {
        Reflection r{};
        r.h = h; r.k = k; r.l = l;
        r.image_number = static_cast<float>(frame);
        r.I = 100.0f;
        r.sigma = 10.0f;
        r.prescaling_corr = 1.0f;
        r.partiality = 0.4f;
        r.observed_x = NAN; r.observed_y = NAN;
        outcomes[frame].reflections.push_back(r);
    };
    for (int f : {10, 11, 12, 100, 101, 102})
        add(f, 1, 2, 3);
    for (int f : {50, 51})
        add(f, 2, 3, 4);

    const MtzTable unscaled = UnmergedMtzTable(outcomes, TETRAGONAL_CELL, experiment, true, OutputProvenance{}, 2);
    REQUIRE(unscaled.columns[0].values.size() == 3);

    // the merge kept the first crossing of (1,2,3) and that of (2,3,4), and rejected the second crossing
    const std::vector<ScaledFull> fulls{{1, 2, 3, 5000.0f, 50.0f, 11.0f}, {2, 3, 4, 7000.0f, 70.0f, 50.0f}};
    const MtzTable scaled = UnmergedMtzTable(outcomes, TETRAGONAL_CELL, experiment, true, OutputProvenance{}, 2,
                                             &fulls);
    const auto column = [&](const char *label) -> const std::vector<float> & {
        for (const auto &c : scaled.columns)
            if (c.label == label)
                return c.values;
        FAIL("no column " << label);
        return scaled.columns[0].values;
    };
    REQUIRE(column("I").size() == 2);
    std::vector<std::pair<float, float>> rows;
    for (size_t i = 0; i < 2; ++i)
        rows.emplace_back(column("I")[i], column("SIGI")[i]);
    std::sort(rows.begin(), rows.end());
    CHECK(rows[0] == std::pair<float, float>(5000.0f, 50.0f));
    CHECK(rows[1] == std::pair<float, float>(7000.0f, 70.0f));
    for (float batch : column("BATCH"))
        CHECK(batch < 60.0f);   // not the second crossing, at batch 102
}