Jungfraujoch/image_analysis/indexing/PostIndexingRefinement.cpp

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

#include "../../common/JFJochMath.h"
#include "PostIndexingRefinement.h"

#include <iostream>
#include <limits>
#include <future>
#include <optional>

namespace {
    struct config_ifssr final {
        float threshold_contraction = .8; // contract error threshold by this value in every iteration
        float max_distance = .00075; // max distance to reciprocal spots for inliers
        unsigned min_spots = 8; // minimum number of spots to fit against
        unsigned max_iter = 32; // max number of iterations
    };

    static std::pair<float, float> score_parts(float score) noexcept {
        float nsp = -std::floor(score);
        float s = score + nsp;
        return std::make_pair(nsp - 1, s);
    }

    struct RefinedCandidate {
        Eigen::Matrix3f cell;
        float score;
        float volume;
        int64_t indexed_spot_count;
        std::vector<uint8_t> indexed_mask;
    };

    static inline Eigen::MatrixX3<float> CalculateResiduals(
        const Eigen::Ref<const Eigen::MatrixX3<float>> &spots,
        const Eigen::Matrix3f &cell) {
        Eigen::MatrixX3<float> miller = (spots * cell).array().round().matrix();
        Eigen::MatrixX3<float> resid = miller * cell.inverse();
        resid -= spots;
        return resid;
    }

    static inline std::vector<uint8_t> ComputeIndexedMask(
     const Eigen::Ref<const Eigen::MatrixX3<float>> &spots,
     const Eigen::Matrix3f &cell,
     float indexing_tolerance,
     int64_t &indexed_spot_count) {
        const float indexing_tolerance_sq = indexing_tolerance * indexing_tolerance;

        // Compute fractional Miller indices. rint (round half to even) rather than round (round half
        // away from zero): without SSE4.1 Eigen has no vector round, so each element is a libm call,
        // while rint is a few inline instructions. Only the SQUARED residual is taken below and the two
        // rules can differ only at an exact .5, where either leaves |frac| = 0.5 - so the mask and the
        // count are the same. The refinement loop above keeps round: there the rounded value IS the
        // Miller index that goes into the residual and the QR solve, so its tie rule does matter.
        Eigen::MatrixX3<float> miller_frac = spots * cell;
        Eigen::MatrixX3<float> miller_int = miller_frac.array().rint().matrix();
        Eigen::MatrixX3<float> frac_resid = miller_frac - miller_int;

        std::vector<uint8_t> mask(spots.rows(), 0);
        indexed_spot_count = 0;

        for (int i = 0; i < spots.rows(); ++i) {
            if (frac_resid.row(i).squaredNorm() < indexing_tolerance_sq) {
                mask[i] = 1;
                indexed_spot_count++;
            }
        }

        return mask;
    }

    template<typename MatX3, typename VecX>
    static void RefineCandidateCells(const Eigen::Ref<const Eigen::MatrixX3<float>> &spots,
                                     Eigen::DenseBase<MatX3> &cells,
                                     Eigen::DenseBase<VecX> &scores,
                                     const config_ifssr &cifssr,
                                     unsigned block = 0, unsigned nblocks = 1) {
        using namespace Eigen;
        using Mx3 = MatrixX3<float>;
        using M3 = Matrix3<float>;

        const unsigned nspots = spots.rows();
        const unsigned ncells = scores.rows();
        VectorX<bool> below{nspots};
        Mx3 resid{nspots, 3u};
        Mx3 miller{nspots, 3u};
        M3 cell;

        const unsigned blocksize = (ncells + nblocks - 1u) / nblocks;
        const unsigned startcell = block * blocksize;
        const unsigned endcell = std::min(startcell + blocksize, ncells);

        for (unsigned j = startcell; j < endcell; j++) {
            if (nspots < cifssr.min_spots) {
                scores(j) = float{1.};
                continue;
            }

            cell = cells.block(3u * j, 0u, 3u, 3u).transpose(); // cell: col vectors
            const float scale = cell.colwise().norm().minCoeff();
            float threshold = score_parts(scores[j]).second / scale;

            for (unsigned niter = 1; niter < cifssr.max_iter && threshold > cifssr.max_distance; niter++) {
                miller = (spots * cell).array().round().matrix();
                resid = miller * cell.inverse();
                resid -= spots;

                below = (resid.rowwise().norm().array() < threshold);
                if (below.count() < cifssr.min_spots)
                    break;

                threshold *= cifssr.threshold_contraction;
                // Least squares spots * cell = miller over the spots below the threshold, by the normal
                // equations summed in spot order in double. A QR over all the spots reduces its column
                // norms in SIMD packets of the target's width, so its cell - and which candidate a
                // borderline first pass indexed with - depended on -march; these sums do not.
                Matrix3d ata = Matrix3d::Zero();
                Matrix3d atb = Matrix3d::Zero();
                for (unsigned i = 0; i < nspots; i++) {
                    if (!below[i])
                        continue;
                    for (int r = 0; r < 3; r++)
                        for (int c = 0; c < 3; c++) {
                            ata(r, c) += double(spots(i, r)) * double(spots(i, c));
                            atb(r, c) += double(spots(i, r)) * double(miller(i, c));
                        }
                }
                cell = (ata.inverse() * atb).cast<float>();
            }

            resid = CalculateResiduals(spots, cell);

            ArrayX<float> dist = resid.rowwise().norm();
            // A singular refined cell puts inf into cell.inverse() and 0*inf = NaN into the
            // residuals; a NaN key is UB in nth_element. Such a spot does not index at all,
            // which is exactly what +inf says - and inf, unlike NaN, orders consistently.
            dist = dist.isFinite().select(dist, std::numeric_limits<float>::infinity());
            auto nth = std::begin(dist) + (cifssr.min_spots - 1);
            std::nth_element(std::begin(dist), nth, std::end(dist));
            scores(j) = *nth;

            cells.block(3u * j, 0u, 3u, 3u) = cell.transpose();
        }
    }

    // The reference test, on a candidate's three rows: sorted lengths within the relative tolerance and
    // sorted angles within REFINE_ANGLE_TOLERANCE_VS_REFERENCE_DEG of any of the reference cells. Each
    // angle is folded to its acute complement (min(x,180-x)) so the obtuse/acute setting choice is
    // irrelevant. Guards against a right-edges/wrong-angle cell (a pseudo-symmetric near-metric, e.g. a
    // monoclinic beta refined to the wrong value) passing on lengths.
    bool MatchesReference(const Eigen::Matrix3f &cell_rows, const std::vector<UnitCell> &reference_cells,
                          float dist_tolerance) {
        std::array<float, 3> obs = {cell_rows.row(0).norm(), cell_rows.row(1).norm(), cell_rows.row(2).norm()};
        std::sort(obs.begin(), obs.end());

        auto fold = [](float deg) { return std::min(deg, 180.0f - deg); };
        auto row_angle = [&](int i, int j) {
            return std::acos(std::clamp(cell_rows.row(i).normalized().dot(cell_rows.row(j).normalized()),
                                        -1.0f, 1.0f)) * 180.0f / PI;
        };
        std::array<float, 3> obs_ang = {fold(row_angle(1, 2)), fold(row_angle(0, 2)), fold(row_angle(0, 1))};
        std::sort(obs_ang.begin(), obs_ang.end());

        auto matches = [&](const UnitCell &reference) {
            std::array<float, 3> ref = {reference.a, reference.b, reference.c};
            std::sort(ref.begin(), ref.end());
            for (int k = 0; k < 3; ++k) {
                const float denom = std::max(ref[k], REFINE_MIN_REFERENCE_LENGTH_EPSILON);
                if (std::abs(obs[k] - ref[k]) / denom > dist_tolerance)
                    return false;
            }

            std::array<float, 3> ref_ang = {fold(reference.alpha), fold(reference.beta), fold(reference.gamma)};
            std::sort(ref_ang.begin(), ref_ang.end());
            for (int k = 0; k < 3; ++k) {
                if (std::abs(obs_ang[k] - ref_ang[k]) > REFINE_ANGLE_TOLERANCE_VS_REFERENCE_DEG)
                    return false;
            }
            return true;
        };
        return std::any_of(reference_cells.begin(), reference_cells.end(), matches);
    }

    // A candidate that is an integer SUPERCELL of the reference lattice: the FFT builds one from rows it
    // reports at a multiple of their length, and on a twin by reticular merohedry the cell indexing both
    // domains is their coincidence lattice, a supercell of either. Every node of the candidate is then a
    // node of the reference lattice, so each reference basis vector is a combination of the candidate's
    // rows with coefficients in (1/n)Z, n the volume ratio. Those combinations are enumerated, and of the
    // triples that match the reference, span its volume and contain the candidate's lattice, the one
    // that indexes the most spots is returned (on a twin, the stronger domain).
    std::optional<Eigen::Matrix3f> StepDownToReference(const Eigen::Matrix3f &cell_rows,
                                                       const std::vector<UnitCell> &reference_cells,
                                                       float dist_tolerance,
                                                       const Eigen::Ref<const Eigen::MatrixX3<float>> &spots,
                                                       float indexing_tolerance) {
        constexpr int MAX_INDEX = 10;
        const float vol = std::abs(cell_rows.determinant());
        std::optional<Eigen::Matrix3f> best;
        int64_t best_count = -1;
        for (const auto &ref: reference_cells) {
            const float ref_vol = CrystalLattice(ref).CalcVolume();
            if (!(ref_vol > 0.0f))
                continue;
            const float ratio = vol / ref_vol;
            const int n = static_cast<int>(std::lround(ratio));
            if (n < 2 || n > MAX_INDEX || std::abs(ratio - static_cast<float>(n)) > 0.15f * static_cast<float>(n))
                continue;

            const std::array<float, 3> lengths = {ref.a, ref.b, ref.c};
            // Of v and -v only the one met first is kept: either spans the same lattice with the same
            // count, and the first triple with the most spots is then one the full enumeration would
            // have returned too - in a loop an eighth of the length.
            std::array<std::vector<Eigen::RowVector3f>, 3> near;
            for (int i = -n; i <= n; i++)
                for (int j = -n; j <= n; j++)
                    for (int k = -n; k <= n; k++) {
                        if (!(i < 0 || (i == 0 && (j < 0 || (j == 0 && k < 0)))))
                            continue;
                        const Eigen::RowVector3f v = (static_cast<float>(i) * cell_rows.row(0)
                                                      + static_cast<float>(j) * cell_rows.row(1)
                                                      + static_cast<float>(k) * cell_rows.row(2))
                                                     / static_cast<float>(n);
                        for (int a = 0; a < 3; a++)
                            if (std::abs(v.norm() - lengths[a]) <= dist_tolerance * lengths[a])
                                near[a].push_back(v);
                    }

            // The volume test rejects nearly every triple, and the volume is |(v0 x v1) . v2|: one dot
            // product per triple once the cross product is taken per pair. That is asked first, with a
            // wider bound than the test itself so that its different rounding cannot drop a triple
            // the test would pass.
            for (const auto &v0: near[0])
                for (const auto &v1: near[1]) {
                    const Eigen::RowVector3f v01 = v0.cross(v1);
                    for (const auto &v2: near[2]) {
                        if (std::abs(std::abs(v01.dot(v2)) - ref_vol) > 0.16f * ref_vol)
                            continue;
                        Eigen::Matrix3f sub;
                        sub << v0, v1, v2;
                        if (std::abs(std::abs(sub.determinant()) - ref_vol) > 0.15f * ref_vol)
                            continue;
                        const Eigen::Matrix3f m = cell_rows * sub.inverse();
                        if ((m - m.array().round().matrix()).cwiseAbs().maxCoeff() > 0.1f)
                            continue;
                        if (!MatchesReference(sub, {ref}, dist_tolerance))
                            continue;
                        int64_t count = 0;
                        ComputeIndexedMask(spots, sub.transpose(), indexing_tolerance, count);
                        if (count > best_count) {
                            best_count = count;
                            best = sub;
                        }
                    }
                }
        }
        return best;
    }
}

// One lattice contains the other: the rows of the larger cell are integer combinations of the rows
// of the smaller, and the volumes differ by an integer factor of at least two.
bool IntegerRelated(const Eigen::Matrix3f &a, const Eigen::Matrix3f &b) {
    const bool a_larger = std::abs(a.determinant()) > std::abs(b.determinant());
    const Eigen::Matrix3f &large = a_larger ? a : b, &small = a_larger ? b : a;
    if (std::abs(large.determinant()) < 1.5f * std::abs(small.determinant()))
        return false;
    const Eigen::Matrix3f m = large * small.inverse();
    return (m - m.array().round().matrix()).cwiseAbs().maxCoeff() < 0.1f;
}

std::vector<CrystalLattice> Refine(const std::vector<Coord> &in_spots,
                                   size_t nspots,
                                   Eigen::MatrixX3<float> &oCell,
                                   Eigen::VectorX<float> &scores,
                                   RefineParameters &p) {
    std::vector<CrystalLattice> ret;

    Eigen::MatrixX3<float> spots(in_spots.size(), 3u);

    for (int i = 0; i < in_spots.size(); i++) {
        spots(i, 0u) = in_spots[i].x;
        spots(i, 1u) = in_spots[i].y;
        spots(i, 2u) = in_spots[i].z;
    }

    config_ifssr cifssr{
        .min_spots = static_cast<uint32_t>(p.viable_cell_min_spots)
    };

    // Candidate cells refine independently - a block touches only its own scores(j) and cells rows, and
    // holds its own scratch - so splitting them across threads gives the same numbers as one thread.
    // Only worth it where few indexer threads run (the rotation first pass uses two, one per scheme,
    // and leaves the rest of the machine idle); refine_threads stays 1 everywhere else.
    const unsigned ncells = static_cast<unsigned>(scores.rows());
    const unsigned nblocks = std::max(1u, std::min(p.refine_threads, ncells));
    if (nblocks == 1) {
        RefineCandidateCells(spots.topRows(nspots), oCell, scores, cifssr);
    } else {
        // Futures, not bare threads: an exception in a block (out of memory) then reaches the caller
        // instead of calling std::terminate.
        std::vector<std::future<void>> workers;
        workers.reserve(nblocks - 1);
        for (unsigned b = 1; b < nblocks; b++)
            workers.push_back(std::async(std::launch::async, [&, b] {
                RefineCandidateCells(spots.topRows(nspots), oCell, scores, cifssr, b, nblocks);
            }));
        RefineCandidateCells(spots.topRows(nspots), oCell, scores, cifssr, 0, nblocks);
        for (auto &w : workers)
            w.get();
    }

    std::vector<RefinedCandidate> candidates;

    // Angle bounds as cosines, once, for the per-candidate test below.
    const float cos_min_angle = std::cos(p.min_angle_deg * PI / 180.0f);
    const float cos_max_angle = std::cos(p.max_angle_deg * PI / 180.0f);

    // A reference cell is typed the way a deposit or a paper states it, which for a centred lattice
    // is the CONVENTIONAL cell - and the candidates it is compared against are Niggli-reduced
    // PRIMITIVE cells, whose edges a C, I, F or R description does not have. Compared as typed, a
    // centred reference therefore rejects the true candidate: it survives only where the transform
    // happened to build the centred cell itself as well - a cell of a sublattice, on which the run
    // then proceeds - and otherwise no lattice is returned at all, so the user who quotes the deposit
    // is exactly the one the option fails. The reference is therefore expanded into the primitive
    // lattices its six numbers could stand for, one per centring, each reduced the way a candidate
    // is; a candidate matching any of them is kept. The cell as typed stays in the set because
    // ffbidx hands back the basis it was given, unreduced.
    std::vector<UnitCell> reference_cells;
    if (p.reference_unit_cell) {
        reference_cells.push_back(*p.reference_unit_cell);
        const CrystalLattice reference(*p.reference_unit_cell);
        for (const char centering: {'A', 'B', 'C', 'I', 'F', 'R'})
            reference_cells.push_back(reference.ToPrimitive(centering).NiggliReduce().GetUnitCell());
        reference_cells.push_back(reference.NiggliReduce().GetUnitCell());
    }

   for (int i = 0; i < scores.size(); i++) {
        Eigen::Matrix3f cell_rows = oCell.block(3u * i, 0u, 3u, 3u);
        Eigen::Matrix3f cell_cols = cell_rows.transpose();
        Eigen::Vector3f row_norms = cell_rows.rowwise().norm();

        if (!reference_cells.empty()) {
            if (!MatchesReference(cell_rows, reference_cells, p.dist_tolerance_vs_reference)) {
                if (!p.step_down_to_reference)
                    continue;
                const auto sub = StepDownToReference(cell_rows, reference_cells, p.dist_tolerance_vs_reference,
                                                     spots.topRows(nspots), p.indexing_tolerance);
                if (!sub)
                    continue;
                cell_rows = *sub;
                cell_cols = cell_rows.transpose();
                row_norms = cell_rows.rowwise().norm();
            }
        } else {
            if (row_norms.minCoeff() < p.min_length_A || row_norms.maxCoeff() > p.max_length_A)
                continue;
        }

        // Filter for wrong angles. Compared as COSINES, not angles: acos is strictly decreasing on
        // [-1, 1], so "angle outside [min_angle, max_angle]" is exactly "cosine outside
        // [cos(max_angle), cos(min_angle)]" with the ends swapped - and the three acos calls the
        // comparison needed disappear. They were not cheap: this runs per candidate cell per image,
        // and on a serial-stills run acos was 41% of the whole process.
        const float cos_alpha = cell_rows.row(1).normalized().dot(cell_rows.row(2).normalized());
        const float cos_beta = cell_rows.row(0).normalized().dot(cell_rows.row(2).normalized());
        const float cos_gamma = cell_rows.row(0).normalized().dot(cell_rows.row(1).normalized());

        if (cos_alpha > cos_min_angle || cos_alpha < cos_max_angle ||
            cos_beta > cos_min_angle || cos_beta < cos_max_angle ||
            cos_gamma > cos_min_angle || cos_gamma < cos_max_angle)
            continue;

        // Three coplanar rows pass both tests above - 120 deg apart in one plane is inside any angle
        // window - and the refinement can walk a candidate there. Such a cell has no reciprocal basis,
        // so it is not a lattice: the same volume test as CrystalLattice::VolumeFraction().
        if (std::abs(cell_rows.determinant()) < MIN_BASIS_VOLUME_FRACTION * row_norms.prod())
            continue;

        int64_t indexed_spot_count = 0;
        auto indexed_mask = ComputeIndexedMask(spots.topRows(nspots), cell_cols, p.indexing_tolerance, indexed_spot_count);

        if (indexed_spot_count < p.viable_cell_min_spots)
            continue;

        candidates.emplace_back(RefinedCandidate{
            .cell = cell_rows,
            .score = scores(i),
            .volume = std::abs(cell_rows.determinant()),
            .indexed_spot_count = indexed_spot_count,
            .indexed_mask = std::move(indexed_mask)
        });
    }

    std::sort(candidates.begin(), candidates.end(),
              [](const RefinedCandidate &a, const RefinedCandidate &b) {
                  const auto max_spots = std::max(a.indexed_spot_count, b.indexed_spot_count);
                  const auto min_spots = std::min(a.indexed_spot_count, b.indexed_spot_count);
                  const bool spot_counts_close = (max_spots > 0)
                                                 && (static_cast<float>(min_spots) / static_cast<float>(max_spots)
                                                     >= REFINE_CANDIDATE_SPOT_COUNT_RATIO_THRESHOLD);

                  if (!spot_counts_close)
                      return a.indexed_spot_count > b.indexed_spot_count;

                  const float max_volume = std::max(a.volume, b.volume);
                  const float min_volume = std::max(std::min(a.volume, b.volume), REFINE_MIN_VOLUME_EPSILON);
                  const bool volume_differs = (max_volume / min_volume) > REFINE_CANDIDATE_VOLUME_RATIO_THRESHOLD;

                  if (volume_differs)
                      return a.volume < b.volume;

                  if (a.score != b.score)
                      return a.score < b.score;

                  return a.indexed_spot_count > b.indexed_spot_count;
              });

    std::vector<RefinedCandidate> accepted;
    // Candidates the overlap test drops only for their overlap with a cell they are an integer multiple
    // or sub-cell of: one of those two lattices contains the other, so that overlap is total by
    // construction and says nothing about which of them is the crystal's. They are kept, in rank order,
    // after every cell the test accepts - which come first and in the order they always had - for the
    // caller to decide on.
    std::vector<RefinedCandidate> multiples;

    for (const auto &candidate: candidates) {
        int64_t overlap = 0, overlap_unrelated = 0;

        // Check all already selected lattices and see how many spots are already indexed for the candidate
        // If the overlap is more than 40% of indexed spots - we assume the lattice doesn't bring anything new
        for (const auto &selected: accepted) {
            int64_t n = 0;
            for (size_t i = 0; i < candidate.indexed_mask.size(); ++i) {
                if (candidate.indexed_mask[i] && selected.indexed_mask[i])
                    n++;
            }
            overlap += n;
            if (!IntegerRelated(candidate.cell, selected.cell))
                overlap_unrelated += n;
        }

        const auto bound = static_cast<int64_t>(REFINE_CANDIDATE_OVERLAP_RATIO_THRESHOLD
                                                * static_cast<float>(candidate.indexed_spot_count));
        if (overlap < bound)
            accepted.emplace_back(candidate);
        else if (overlap_unrelated < bound)
            multiples.emplace_back(candidate);
    }
    for (auto &m : multiples)
        accepted.push_back(std::move(m));

    ret.reserve(accepted.size());

    for (auto &candidate: accepted) {
        auto cell = candidate.cell;
        if (cell.determinant() < .0f)
            cell = -cell;

        ret.emplace_back(
            Coord(cell(0, 0), cell(0, 1), cell(0, 2)),
            Coord(cell(1, 0), cell(1, 1), cell(1, 2)),
            Coord(cell(2, 0), cell(2, 1), cell(2, 2))
        );
    }

    return ret;
}