Jungfraujoch/image_analysis/LoadFCalcFromMtz.cpp

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

#include "LoadFCalcFromMtz.h"

#include <cmath>
#include <cstdint>
#include <cstring>
#include <limits>
#include <sstream>
#include <stdexcept>
#include <string>
#include <vector>

#include <gemmi/cif2mtz.hpp>
#include <gemmi/mtz.hpp>
#include <gemmi/read_cif.hpp>
#include <gemmi/refln.hpp>
#include <gemmi/symmetry.hpp>

namespace {

// Reference intensities can come from a calculated structure factor F-model (squared to an
// intensity), or from a merged/observed mean-intensity column (used directly - this is what lets
// reference-based scaling be self-seeded from the data's own previous merge). column_with_one_of_labels would
// hide which label matched, so the priority list is walked explicitly to record the choice.
const gemmi::Mtz::Column* SelectDefaultColumn(const gemmi::Mtz& mtz, bool& square) {
    if (const auto* col = mtz.column_with_label("F-model", nullptr, 'F')) {
        square = true;
        return col;
    }
    for (const char* label : {"IMEAN", "I", "IOBS", "Iobs", "I-obs"}) {
        if (const auto* col = mtz.column_with_label(label, nullptr, 'J')) {
            square = false;
            return col;
        }
    }
    for (const auto& c : mtz.columns)
        if (c.type == 'J') {
            square = false;
            return &c;
        }
    // Fall back to a plain structure-factor amplitude (type F): a deposited/observed FP or FOBS, or a
    // lone FC, squared to an intensity like F-model. Lets a reference MTZ that carries only amplitudes
    // (no F-model or IMEAN/I - e.g. a deposition that has only FP) still seed CCref
    // without an explicit --reference-column. Only these named observed/calculated labels are matched:
    // a bare "first type-F column" catch-all would silently pick up map coefficients (FWT, DELFWT,
    // FC_ALL) and seed CCref from the wrong quantity - require --reference-column in that case instead.
    for (const char* label : {"FP", "FOBS", "F", "FC", "Fobs", "F-obs"}) {
        if (const auto* col = mtz.column_with_label(label, nullptr, 'F')) {
            square = true;
            return col;
        }
    }
    return nullptr;
}

// GEMMI's SF-mmCIF -> MTZ conversion, with one change to where FreeR_flag comes from. _refln.status
// ('f' = free, 'o' = working, anything else no flag) is the PDB's own record of the test set, so it
// is preferred to _refln.pdbx_r_free_flag, whose convention varies by refinement program. A status
// column with no 'f' at all says nothing about the test set and is not used. Either way FreeR_flag
// ends up in the CCP4 convention the MTZ path already reads (status: f -> 0, o -> 1).
std::vector<std::string> SfCifSpec(const gemmi::cif::Loop& loop) {
    bool status_has_free = false;
    const int status = loop.find_tag("_refln.status");
    if (status >= 0)
        for (std::size_t i = status; i < loop.values.size(); i += loop.width())
            if (gemmi::cif::as_string(loop.values[i]) == "f") {
                status_has_free = true;
                break;
            }

    std::vector<std::string> spec;
    for (const char** line = gemmi::CifToMtz::default_spec(true); *line != nullptr; ++line) {
        const std::string l = *line;
        if (status_has_free && l.starts_with("pdbx_r_free_flag "))
            continue;
        if (!status_has_free && l.starts_with("status "))
            continue;
        spec.push_back(l);
    }
    return spec;
}

// A deposited structure-factor file can hold several reflection blocks: the main merged data first,
// then possibly further merged sets (another wavelength, anomalous) and unmerged data (_diffrn_refln).
// The first merged block with the requested column - or, with none requested, a column the default
// choice accepts - is used; unmerged blocks and anomalous-only blocks (only I(+)/I(-) or F(+)/F(-))
// are passed over.
gemmi::Mtz MtzFromSfCif(gemmi::cif::Document&& doc, const std::optional<std::string>& column,
                        std::string& source) {
    std::vector<gemmi::ReflnBlock> blocks = gemmi::as_refln_blocks(std::move(doc.blocks));
    gemmi::Logger logger;   // no callback: GEMMI's conversion notes are not shown
    for (auto& rb : blocks) {
        if (rb.refln_loop == nullptr)
            continue;
        gemmi::CifToMtz converter;
        converter.spec_lines = SfCifSpec(*rb.refln_loop);
        gemmi::Mtz mtz = converter.convert_block_to_mtz(rb, logger);
        bool square;
        if (column ? mtz.column_with_label(*column) != nullptr : SelectDefaultColumn(mtz, square) != nullptr) {
            source = "SF-mmCIF block " + rb.block.name;
            return mtz;
        }
    }
    throw std::runtime_error(column ? "no merged reflection block has column '" + *column + "'"
                                    : "no merged reflection block has an intensity or amplitude column to use as reference");
}

// The reference is an MTZ or a structure-factor mmCIF (SF-mmCIF, as deposited in the PDB), gzipped or
// not (by the .gz extension, as for every GEMMI read). The format is recognised by content: an MTZ
// starts with "MTZ ", anything else is parsed as CIF and converted to an MTZ in memory.
gemmi::Mtz ReadReferenceAsMtz(const std::string& path, const std::optional<std::string>& column,
                              std::string& source) {
    gemmi::CharArray mem = gemmi::read_into_buffer_gz(path);
    if (mem.size() >= 4 && std::memcmp(mem.data(), "MTZ ", 4) == 0) {
        gemmi::Mtz mtz;
        gemmi::MemoryStream stream(mem.data(), mem.size());
        mtz.read_stream(stream, true);
        source = "MTZ";
        return mtz;
    }
    return MtzFromSfCif(gemmi::read_cif_from_memory(mem.data(), mem.size(), path.c_str()), column, source);
}

} // namespace

ReferenceMtzData LoadReferenceMtz(const std::string& path,
                                  const std::optional<std::string>& column) {
    ReferenceMtzData out;
    const gemmi::Mtz mtz = ReadReferenceAsMtz(path, column, out.source);

    // Columns the caller may pick as reference intensities/structure factors.
    for (const auto& c : mtz.columns)
        if (c.type == 'J' || c.type == 'F')
            out.candidate_columns.push_back({c.label, c.type});

    const gemmi::Mtz::Column* col = nullptr;
    if (column) {
        col = mtz.column_with_label(*column, nullptr);
        if (col == nullptr)
            throw std::runtime_error("Reference MTZ has no column '" + *column + "'");
        out.squared = (col->type == 'F');
        out.default_column = false;
    } else {
        col = SelectDefaultColumn(mtz, out.squared);
        if (col == nullptr)
            throw std::runtime_error("MTZ has no F-model, intensity (J) or amplitude (F) column to use as reference");
        out.default_column = true;
    }
    out.used_column = col->label;
    out.used_column_type = col->type;

    // Optional cross-validation flags: a campaign shares ONE test set, so if the reference carries a
    // FreeR flag we import it and every dataset can inherit the same free reflections.
    const gemmi::Mtz::Column* free_col = nullptr;
    for (const char* label : {"FreeR_flag", "FREE", "FREER", "RFREE", "R-free-flags", "FreeRflag"})
        if ((free_col = mtz.column_with_label(label, nullptr, 'I')) != nullptr)
            break;

    // Cell and space group the reference was recorded in, for display and the consistency check.
    const gemmi::UnitCell& gc = mtz.get_cell();
    const bool cell_valid = gc.a > 0 && gc.b > 0 && gc.c > 0;
    if (cell_valid)
        out.cell = UnitCell{static_cast<float>(gc.a), static_cast<float>(gc.b), static_cast<float>(gc.c),
                            static_cast<float>(gc.alpha), static_cast<float>(gc.beta),
                            static_cast<float>(gc.gamma)};
    if (mtz.spacegroup != nullptr) {
        out.space_group_number = mtz.spacegroup->number;
        out.space_group_name = mtz.spacegroup->short_name();
        out.space_group_xhm = mtz.spacegroup->xhm();
        out.point_group = mtz.spacegroup->point_group_hm();
    }

    out.reflections.reserve(static_cast<std::size_t>(mtz.nreflections));
    std::vector<int> raw_free;                        // raw FreeR value, aligned with out.reflections
    if (free_col != nullptr)
        raw_free.reserve(static_cast<std::size_t>(mtz.nreflections));

    const std::size_t stride = mtz.columns.size();
    double d_min = std::numeric_limits<double>::max();
    double d_max = 0.0;

    for (int i = 0; i < mtz.nreflections; ++i) {
        const float v = (*col)[static_cast<std::size_t>(i)];
        if (std::isnan(v))
            continue;

        const std::size_t row = static_cast<std::size_t>(i) * stride;
        MergedReflection r;
        r.h = static_cast<int32_t>(mtz.data[row + 0]);
        r.k = static_cast<int32_t>(mtz.data[row + 1]);
        r.l = static_cast<int32_t>(mtz.data[row + 2]);
        r.I = out.squared ? v * v : v;
        r.sigma = NAN;

        if (cell_valid) {
            r.d = static_cast<float>(gc.calculate_d({{r.h, r.k, r.l}}));
            if (std::isfinite(r.d) && r.d > 0.0f) {
                d_min = std::min<double>(d_min, r.d);
                d_max = std::max<double>(d_max, r.d);
            }
        }

        out.reflections.emplace_back(r);
        if (free_col != nullptr) {
            const float f = (*free_col)[static_cast<std::size_t>(i)];
            raw_free.push_back(std::isnan(f) ? -1 : static_cast<int>(std::lround(f)));
        }
    }

    if (d_max > 0.0) {
        out.d_min = d_min;
        out.d_max = d_max;
    }

    // Set the free flag from the raw column. The CCP4/refmac convention is that the test set is
    // flag 0 (this also handles the historical 0-19 thin-shell format, where 0 is ~5%). If flag 0
    // would instead be the majority (a phenix-style file where 1 marks free), take the complement.
    if (free_col != nullptr && !raw_free.empty()) {
        std::size_t zeros = 0;
        for (int f : raw_free)
            if (f == 0)
                ++zeros;
        const bool free_is_zero = zeros * 2 <= raw_free.size();
        for (std::size_t j = 0; j < out.reflections.size(); ++j) {
            const bool is_free = (raw_free[j] == 0) == free_is_zero;
            out.reflections[j].rfree_flag = is_free;
            if (is_free)
                ++out.n_free;
        }
        out.has_free_flags = true;
        out.free_column = free_col->label;
    }

    return out;
}

std::string ReferenceConsistencyWarning(const ReferenceMtzData& reference,
                                        const std::optional<UnitCell>& data_cell,
                                        std::optional<int> data_space_group_number) {
    std::vector<std::string> issues;

    if (reference.cell && data_cell) {
        const auto& r = *reference.cell;
        const auto& d = *data_cell;
        auto rel = [](float x, float y) { return y != 0.0f ? std::abs(x - y) / std::abs(y) : 0.0f; };
        const bool len_off = rel(r.a, d.a) > 0.02f || rel(r.b, d.b) > 0.02f || rel(r.c, d.c) > 0.02f;
        const bool ang_off = std::abs(r.alpha - d.alpha) > 1.5f || std::abs(r.beta - d.beta) > 1.5f
                          || std::abs(r.gamma - d.gamma) > 1.5f;
        if (len_off || ang_off) {
            std::ostringstream ss;
            ss.setf(std::ios::fixed);
            ss.precision(2);
            ss << "unit cell differs (reference " << r.a << " " << r.b << " " << r.c << " "
               << r.alpha << " " << r.beta << " " << r.gamma << ", data " << d.a << " " << d.b
               << " " << d.c << " " << d.alpha << " " << d.beta << " " << d.gamma << ")";
            issues.push_back(ss.str());
        }
    }

    if (!reference.point_group.empty() && data_space_group_number) {
        const auto* sg = gemmi::find_spacegroup_by_number(*data_space_group_number);
        if (sg != nullptr && reference.point_group != sg->point_group_hm())
            issues.push_back("point group differs (reference " + reference.point_group + ", data "
                             + std::string(sg->point_group_hm()) + ")");
    }

    if (issues.empty())
        return {};

    std::string out = "reference may not match the data: ";
    for (std::size_t i = 0; i < issues.size(); ++i)
        out += (i ? "; " : "") + issues[i];
    return out;
}