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 <filesystem>
#include <string>
#include <cstring>

#include <gemmi/mtz.hpp>

#include "../common/DiffractionExperiment.h"
#include "../image_analysis/WriteReflections.h"
#include "../image_analysis/IntegrationOutcome.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, 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, 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 gemmi::Mtz serial = UnmergedMtz(outcomes, TETRAGONAL_CELL, experiment, sum_partials, 1);
        const gemmi::Mtz parallel = UnmergedMtz(outcomes, TETRAGONAL_CELL, experiment, sum_partials, 8);
        REQUIRE(serial.nreflections > 100000);
        REQUIRE(parallel.nreflections == serial.nreflections);
        CHECK(std::memcmp(parallel.data.data(), serial.data.data(), serial.data.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);
    }
}