Jungfraujoch/image_analysis/scale_merge/ReindexAmbiguity.cpp

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

#include "ReindexAmbiguity.h"

#include <cmath>
#include <algorithm>
#include <unordered_map>
#include <unordered_set>
#include <utility>

#include "gemmi/twin.hpp"

#include "HKLKey.h"
#include "CrystalSetting.h"

namespace {
    constexpr size_t MIN_REFLECTIONS = 20;

    struct BestReindex {
        gemmi::Op op = gemmi::Op::identity();
        bool is_identity = true;
        double score = 0.0;
        double identity_score = 0.0;
    };

    // Among identity (the baseline, scored `identity_score`) and `ops`, return the operator with the
    // highest `score_op(op)`. A non-finite score never wins (too little overlap to decide). Shared by the
    // stills per-image resolver and the rotation post-merge ChooseReindex.
    BestReindex PickBestReindex(double identity_score, const std::vector<gemmi::Op> &ops,
                                const std::function<double(const gemmi::Op &)> &score_op) {
        BestReindex best;
        best.identity_score = identity_score;
        best.score = identity_score;
        for (const auto &op : ops) {
            const double sc = score_op(op);
            if (std::isfinite(sc) && (!std::isfinite(best.score) || sc > best.score)) {
                best.score = sc;
                best.op = op;
                best.is_identity = false;
            }
        }
        return best;
    }
}

std::vector<gemmi::Op> ReindexAmbiguityOperators(const UnitCell &cell, const gemmi::SpaceGroup &space_group,
                                                 double max_obliquity_deg) {
    return gemmi::find_twin_laws(static_cast<gemmi::UnitCell>(cell), &space_group, max_obliquity_deg,
                                 /*all_ops=*/false);
}

CrystalLattice ReindexLattice(const CrystalLattice &latt, const gemmi::Op &op) {
    const gemmi::Op inv = op.inverse();
    const Coord old_star[3] = {latt.Astar(), latt.Bstar(), latt.Cstar()};
    Coord rs[3];
    for (int i = 0; i < 3; ++i)
        rs[i] = old_star[0] * (static_cast<float>(inv.rot[i][0]) / gemmi::Op::DEN)
              + old_star[1] * (static_cast<float>(inv.rot[i][1]) / gemmi::Op::DEN)
              + old_star[2] * (static_cast<float>(inv.rot[i][2]) / gemmi::Op::DEN);
    const float v_star = rs[0] * (rs[1] % rs[2]);
    return {(rs[1] % rs[2]) / v_star, (rs[2] % rs[0]) / v_star, (rs[0] % rs[1]) / v_star};
}

std::vector<MergedReflection> ReindexReflections(const std::vector<MergedReflection> &merged,
                                                 const gemmi::Op &op) {
    std::vector<MergedReflection> out = merged;
    for (auto &r : out) {
        const gemmi::Op::Miller h = op.apply_to_hkl({{static_cast<int>(r.h),
                                                      static_cast<int>(r.k),
                                                      static_cast<int>(r.l)}});
        r.h = h[0];
        r.k = h[1];
        r.l = h[2];
    }
    return out;
}

std::vector<MergedReflection> ReindexMergedIntoAsu(const std::vector<MergedReflection> &merged,
                                                   const gemmi::Op &op, const gemmi::SpaceGroup &space_group,
                                                   bool merge_friedel) {
    // merge_friedel=false so the generator reports which side of the Friedel pair the ASU index was
    // reached from; that sign is what the two hands have to follow.
    const HKLKeyGenerator key_gen(/*merge_friedel=*/false, space_group);

    std::vector<MergedReflection> out = merged;
    for (auto &r : out) {
        const gemmi::Op::Miller h = op.apply_to_hkl({{r.h, r.k, r.l}});
        const HKLKey key = key_gen(h[0], h[1], h[2]);
        // The hands follow a CHANGE of hand, not the hand the new index happens to land on. An
        // anomalous merge stores the minus mate at -hkl_asu (Merge.cpp), so a row can already be on
        // the minus side before the operator is applied, while I_plus/I_minus are attached in the
        // plus convention on both mates alike (RotationScaleMerge.cpp keys the lookup at plus=true).
        // Since the operator is rotation-type, op(-x) == -op(x), and the minus mate's key.plus is
        // therefore the negation of the plus mate's - testing key.plus alone swaps exactly the rows
        // that must not move, and leaves the ones that must. With merge_friedel every row is stored
        // at the ASU representative, so was_plus is always true and this reduces to the old test.
        const bool was_plus = key_gen(r.h, r.k, r.l).plus;
        if (was_plus != key.plus) {
            std::swap(r.I_plus, r.I_minus);
            std::swap(r.sigma_plus, r.sigma_minus);
            std::swap(r.F_plus, r.F_minus);
            std::swap(r.sigmaF_plus, r.sigmaF_minus);
        }
        const bool at_asu_index = merge_friedel || key.plus;
        r.h = at_asu_index ? key.h : -key.h;
        r.k = at_asu_index ? key.k : -key.k;
        r.l = at_asu_index ? key.l : -key.l;
    }
    return out;
}

ReindexChoice ChooseReindex(const std::vector<MergedReflection> &merged,
                            const UnitCell &cell, const gemmi::SpaceGroup &space_group,
                            const std::function<double(const std::vector<MergedReflection> &)> &score,
                            double max_obliquity_deg) {
    const auto ops = ReindexAmbiguityOperators(cell, space_group, max_obliquity_deg);
    const BestReindex best = PickBestReindex(
        score(merged), ops,
        [&](const gemmi::Op &op) { return score(ReindexReflections(merged, op)); });

    ReindexChoice choice;
    choice.op = best.op;
    choice.is_identity = best.is_identity;
    choice.score = best.score;
    choice.identity_score = best.identity_score;
    choice.n_candidates = 1 + static_cast<int>(ops.size());
    return choice;
}

double ReferenceIntensityCC(const std::vector<MergedReflection> &merged,
                            const std::vector<MergedReflection> &reference,
                            const gemmi::SpaceGroup &space_group) {
    const HKLKeyGenerator key(/*merge_friedel=*/true, space_group);

    std::unordered_map<uint64_t, double> ref;
    ref.reserve(reference.size());
    for (const auto &r : reference)
        if (std::isfinite(r.I))
            ref[key(r).pack()] = r.I;

    double sx = 0, sy = 0, sxx = 0, syy = 0, sxy = 0;
    int n = 0;
    for (const auto &m : merged) {
        if (!std::isfinite(m.I))
            continue;
        const auto it = ref.find(key(m).pack());
        if (it == ref.end())
            continue;
        const double x = m.I, y = it->second;
        sx += x; sy += y; sxx += x * x; syy += y * y; sxy += x * y;
        ++n;
    }
    // NaN, not 0: PickBestReindex discards only non-finite scores, so a 0 here is a SCORE - an operator
    // that mapped almost everything outside the reference's coverage would beat a measured identity that
    // merely correlates a little negatively, and the whole run would be reindexed on no data. The same
    // condition already returns NaN in the per-image resolver's ReferenceCC, against the same minimum.
    if (n < static_cast<int>(MIN_REFLECTIONS))
        return NAN;
    const double cov = n * sxy - sx * sy;
    const double vx = n * sxx - sx * sx;
    const double vy = n * syy - sy * sy;
    return (vx > 0 && vy > 0) ? cov / std::sqrt(vx * vy) : NAN;
}

ReferenceSettingChoice ChooseReferenceSetting(const std::vector<MergedReflection> &merged,
                                              const UnitCell &data_cell,
                                              const gemmi::SpaceGroup &data_space_group,
                                              const std::vector<MergedReflection> &reference,
                                              const std::optional<UnitCell> &reference_cell,
                                              const gemmi::SpaceGroup &reference_space_group) {
    // Every candidate as the operator a Miller index moves by: the twin law first (it acts in the
    // data's frame), then the change of basis. The identity leads, so it wins a tie.
    std::vector<gemmi::Op> hkl_ops{gemmi::Op::identity()};
    std::vector<gemmi::Op> twins{gemmi::Op::identity()};
    for (const auto &law : ReindexAmbiguityOperators(data_cell, data_space_group))
        twins.push_back(law);
    if (reference_cell)
        for (const gemmi::Op &cob : CellMappingOperators(data_cell, *reference_cell))
            for (const gemmi::Op &twin : twins)
                hkl_ops.push_back(twin.combine(HklOperator(cob)));

    ReferenceSettingChoice choice;
    choice.n_candidates = static_cast<int>(hkl_ops.size());
    choice.identity_cc = ReferenceIntensityCC(merged, reference, reference_space_group);
    choice.cc = choice.identity_cc;
    for (size_t i = 1; i < hkl_ops.size(); i++) {
        const double cc = ReferenceIntensityCC(ReindexReflections(merged, hkl_ops[i]), reference,
                                               reference_space_group);
        if (std::isfinite(cc) && (!std::isfinite(choice.cc) || cc > choice.cc)) {
            choice.cc = cc;
            choice.cob = HklOperator(hkl_ops[i]);   // the inverse: the change of basis that moves indices by it
        }
    }
    return choice;
}

ReferenceMatch MatchReference(const std::vector<MergedReflection> &merged,
                              const std::vector<MergedReflection> &reference,
                              const gemmi::SpaceGroup &space_group) {
    ReferenceMatch out;
    out.cc = ReferenceIntensityCC(merged, reference, space_group);
    const HKLKeyGenerator key(/*merge_friedel=*/true, space_group);
    std::unordered_set<uint64_t> ref;
    float ref_d_min = INFINITY;
    for (const auto &r : reference) {
        ref.insert(key(r).pack());
        if (std::isfinite(r.d) && r.d > 0.0f)
            ref_d_min = std::min(ref_d_min, r.d);
    }
    size_t n = 0, matched = 0;
    for (const auto &m : merged) {
        if (std::isfinite(ref_d_min) && !(m.d >= ref_d_min))
            continue;
        ++n;
        matched += ref.count(key(m).pack());
    }
    out.matched_fraction = n > 0 ? static_cast<double>(matched) / static_cast<double>(n) : 0.0;
    return out;
}

ReindexAmbiguityResolver::ReindexAmbiguityResolver(const DiffractionExperiment &x,
                                                   const std::vector<MergedReflection> &reference)
    : s(x.GetScalingSettings()),
      hkl_key_generator(s.GetMergeFriedel(), x.GetSpaceGroupOrP1()) {
    for (const auto &r : reference)
        reference_data[hkl_key_generator(r)] = r.I;
    if (x.GetUnitCell().has_value() && x.GetGemmiSpaceGroup().has_value())
        ops = ReindexAmbiguityOperators(*x.GetUnitCell(), x.GetSpaceGroupOrP1());
}

bool ReindexAmbiguityResolver::Accept(const Reflection &r) const {
    if (r.on_ice_ring)   // ice-contaminated intensity would bias the correlation; keep it out
        return false;
    return AcceptReflection(r, s.GetHighResolutionLimit_A(), s.GetLowResolutionLimit_A());
}

double ReindexAmbiguityResolver::ReferenceCC(std::span<const Reflection> reflections,
                                             const gemmi::Op &op) const {
    double sx = 0, sy = 0, sxx = 0, syy = 0, sxy = 0;
    size_t n = 0;
    for (const auto &r : reflections) {
        if (!Accept(r) || r.partiality < s.GetMinPartiality())
            continue;
        if (!std::isfinite(r.I) || !std::isfinite(r.sigma) || r.sigma <= 0.0f || r.partiality <= 0.0f)
            continue;
        const gemmi::Op::Miller h = op.apply_to_hkl({{r.h, r.k, r.l}});
        const auto it = reference_data.find(hkl_key_generator(h[0], h[1], h[2]));
        if (it == reference_data.end())
            continue;
        const double x = static_cast<double>(r.I) * (r.prescaling_corr * r.qe_corr * r.flight_corr) / r.partiality;
        const double y = it->second;
        if (!std::isfinite(x) || !std::isfinite(y))
            continue;
        sx += x; sy += y; sxx += x * x; syy += y * y; sxy += x * y;
        ++n;
    }
    if (n < MIN_REFLECTIONS)
        return NAN;
    const double nd = static_cast<double>(n);
    const double cov = sxy - sx * sy / nd;
    const double vx = sxx - sx * sx / nd;
    const double vy = syy - sy * sy / nd;
    return (vx > 0 && vy > 0) ? cov / std::sqrt(vx * vy) : NAN;
}

// Serial stills index each crystal in one of the merohedrally-equivalent hands at random; pick, for this
// image alone, the reindexing whose intensities correlate best with the external reference and apply it.
std::vector<double> ReindexAmbiguityResolver::Resolve(std::span<Reflection> reflections, CrystalLattice &latt) const {
    if (ops.empty())
        return {};

    std::vector<double> cc{ReferenceCC(reflections, gemmi::Op::identity())};
    const BestReindex best = PickBestReindex(cc[0], ops, [&](const gemmi::Op &op) {
        cc.push_back(ReferenceCC(reflections, op));
        return cc.back();
    });
    if (!best.is_identity)
        ReindexImage(reflections, latt, best.op);
    return cc;
}

void ReindexImage(std::span<Reflection> reflections, CrystalLattice &latt, const gemmi::Op &op) {
    for (auto &r : reflections) {
        const gemmi::Op::Miller h = op.apply_to_hkl({{r.h, r.k, r.l}});
        r.h = h[0]; r.k = h[1]; r.l = h[2];
    }
    // Everything after this - the sublattice relabel, the partiality post-refinement - predicts a spot
    // from its index and the lattice, so the lattice has to describe the new indices. Left on the old
    // axes, a frame relabelled by an operator that is only a near-symmetry of the lattice was put
    // through the sublattice relabel a second time, and its reflections came out at indices finer
    // than anything measured.
    latt = ReindexLattice(latt, op);
}