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>

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
        Eigen::MatrixX3<float> miller_frac = spots * cell;
        Eigen::MatrixX3<float> miller_int = miller_frac.array().round().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};
        MatrixX3<bool> sel{nspots, 3u};
        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;
                sel.colwise() = below;
                HouseholderQR<Mx3> qr{sel.select(spots, .0f)};
                cell = qr.solve(sel.select(miller, .0f));
            }

            resid = CalculateResiduals(spots, cell);

            ArrayX<float> dist = resid.rowwise().norm();
            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();
        }
    }
}

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)
    };

    RefineCandidateCells(spots.topRows(nspots), oCell, scores, cifssr);

    std::vector<RefinedCandidate> candidates;

   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 (p.reference_unit_cell) {
            std::array<float, 3> obs = {row_norms(0), row_norms(1), row_norms(2)};
            std::array<float, 3> ref = {
                static_cast<float>(p.reference_unit_cell->a),
                static_cast<float>(p.reference_unit_cell->b),
                static_cast<float>(p.reference_unit_cell->c)
            };
            std::sort(obs.begin(), obs.end());
            std::sort(ref.begin(), ref.end());

            bool lengths_ok = true;
            for (int k = 0; k < 3; ++k) {
                const float denom = std::max(ref[k], REFINE_MIN_REFERENCE_LENGTH_EPSILON);
                const float rel_dev = std::abs(obs[k] - ref[k]) / denom;
                if (rel_dev > p.dist_tolerance_vs_reference) {
                    lengths_ok = false;
                    break;
                }
            }
            if (!lengths_ok)
                continue;

            // Also require the angles to match the reference. Fold each to its acute complement
            // (min(x,180-x)) so the obtuse/acute setting choice is irrelevant, then compare the
            // sorted triples. 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.
            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::array<float, 3> ref_ang = {
                fold(static_cast<float>(p.reference_unit_cell->alpha)),
                fold(static_cast<float>(p.reference_unit_cell->beta)),
                fold(static_cast<float>(p.reference_unit_cell->gamma))
            };
            std::sort(obs_ang.begin(), obs_ang.end());
            std::sort(ref_ang.begin(), ref_ang.end());

            bool angles_ok = true;
            for (int k = 0; k < 3; ++k) {
                if (std::abs(obs_ang[k] - ref_ang[k]) > REFINE_ANGLE_TOLERANCE_VS_REFERENCE_DEG) {
                    angles_ok = false;
                    break;
                }
            }
            if (!angles_ok)
                continue;
        } else {
            if (row_norms.minCoeff() < p.min_length_A || row_norms.maxCoeff() > p.max_length_A)
                continue;
        }

        // Filter for wrong angles
        float alpha = std::acos(cell_rows.row(1).normalized().dot(cell_rows.row(2).normalized())) * 180.0f / PI;
        float beta = std::acos(cell_rows.row(0).normalized().dot(cell_rows.row(2).normalized())) * 180.0f / PI;
        float gamma = std::acos(cell_rows.row(0).normalized().dot(cell_rows.row(1).normalized())) * 180.0f / PI;

        if (alpha < p.min_angle_deg || alpha > p.max_angle_deg ||
            beta < p.min_angle_deg || beta > p.max_angle_deg ||
            gamma < p.min_angle_deg || gamma > p.max_angle_deg)
            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;

    for (const auto &candidate: candidates) {
        int64_t overlap = 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) {
            for (size_t i = 0; i < candidate.indexed_mask.size(); ++i) {
                if (candidate.indexed_mask[i] && selected.indexed_mask[i])
                    overlap++;
            }
        }

        if (overlap < static_cast<int64_t>(REFINE_CANDIDATE_OVERLAP_RATIO_THRESHOLD
                                           * static_cast<float>(candidate.indexed_spot_count))) {
            accepted.emplace_back(candidate);
        }
    }

    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;
}