Jungfraujoch/image_analysis/indexing/EigenRefine.h

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

#ifndef JFJOCH_EIGENREFINE_H
#define JFJOCH_EIGENREFINE_H

#include <vector>
#include <optional>
#include <Eigen/Dense>

#include "../common/CrystalLattice.h"

struct RefineParameters {
    int64_t viable_cell_min_spots;
    float dist_tolerance_vs_reference;
    std::optional<UnitCell> reference_unit_cell;
    float min_length_A;
    float max_length_A;
    float min_angle_deg;
    float max_angle_deg;
    float indexing_tolerance;
};


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

template<typename MatX3, typename VecX>
static void refine(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 = round((spots * cell).array());
            resid = miller * cell.inverse(); // reciprocal spots induced by <cell>
            resid -= spots; // distance between induced and given 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));
        } {
            // calc score
            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();
    }
}

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

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

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

    // Select cell that explains most spots
    int64_t max_indexed_spot_count = 0;
    float min_score = -1;
    int64_t id = -1;

    const float indexing_tolerance_sq = p.indexing_tolerance * p.indexing_tolerance;

    for (int i = 0; i < scores.size(); i++) {
        // Get cell vectors
        auto cell = oCell.block(3u * i, 0u, 3u, 3u);

        Eigen::Vector3f row_norms = cell.rowwise().norm();

        // Check for distance vs. reference unit cell
        if (p.reference_unit_cell) {
            // Compare edge lengths up to 5% deviation, permutation-invariant
            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) {
                // Guard against zero/near-zero reference values
                const float denom = std::max(ref[k], 1e-6f);
                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;
        } else {
            // Check lengths (A, B, C)
            if (row_norms.minCoeff() < p.min_length_A || row_norms.maxCoeff() > p.max_length_A)
                continue;

            // Calculate angles (alpha, beta, gamma) in degrees
            float alpha = std::acos(cell.row(1).normalized().dot(cell.row(2).normalized())) * 180.0f / M_PI;
            float beta = std::acos(cell.row(0).normalized().dot(cell.row(2).normalized())) * 180.0f / M_PI;
            float gamma = std::acos(cell.row(0).normalized().dot(cell.row(1).normalized())) * 180.0f / M_PI;

            // Check if angles are within allowed range
            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;
        }

        M3x resid = spots.topRows(nspots) * cell.transpose();
        const M3x miller = round(resid.array());
        resid -= miller;

        int64_t indexed_spot_count = (resid.rowwise().squaredNorm().array() < indexing_tolerance_sq).count();
        if (indexed_spot_count > max_indexed_spot_count) {
            max_indexed_spot_count = indexed_spot_count;
            min_score = scores(i);
            id = i;
        } if (indexed_spot_count == max_indexed_spot_count && scores(i) < min_score) {
            min_score = scores(i);
            id = i;
        }
    }

    if (id == -1)
        return {};

    auto cell = oCell.block(3u * id, 0u, 3u, 3u);

    if (cell.determinant() < .0f)
        cell = -cell;

    return { CrystalLattice(
            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))
    )};
}



#endif //JFJOCH_EIGENREFINE_H