PostIndexingRefinement: Fixes

image_analysis/indexing/CMakeLists.txt

@@ -9,10 +9,11 @@ ADD_LIBRARY(JFJochIndexing STATIC
         AnalyzeIndexing.h
         FitProfileRadius.cpp
         FitProfileRadius.h
-        EigenRefine.h
+        PostIndexingRefinement.h
         FFTResult.h
         FFTIndexer.cpp
-        FFTIndexer.h)
+        FFTIndexer.h
+        PostIndexingRefinement.cpp)
 TARGET_LINK_LIBRARIES(JFJochIndexing JFJochCommon)
 
 IF (JFJOCH_CUDA_AVAILABLE)

image_analysis/indexing/EigenRefine.h

@@ -1,196 +0,0 @@
-// 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

image_analysis/indexing/FFBIDXIndexer.cpp

@@ -2,7 +2,7 @@
 // SPDX-License-Identifier: GPL-3.0-only
 
 #include "FFBIDXIndexer.h"
-#include "EigenRefine.h"
+#include "PostIndexingRefinement.h"
 
 void FFBIDXIndexer::SetupUnitCell(const std::optional<UnitCell> &cell) {
     if (!cell.has_value())

image_analysis/indexing/FFTIndexer.cpp

@@ -3,7 +3,7 @@
 
 #include "FFTIndexer.h"
 #include <Eigen/Eigen>
-#include "EigenRefine.h"
+#include "PostIndexingRefinement.h"
 
 FFTIndexer::FFTIndexer(const IndexingSettings &settings)
     : max_length_A(settings.GetFFT_MaxUnitCell_A()),

image_analysis/indexing/FFTIndexerCPU.cpp

@@ -2,7 +2,6 @@
 // SPDX-License-Identifier: GPL-3.0-only
 
 #include "FFTIndexerCPU.h"
-#include "EigenRefine.h"
 
 #include <cmath>
 #include <algorithm>

image_analysis/indexing/PostIndexingRefinement.cpp

@@ -0,0 +1,258 @@
+// SPDX-FileCopyrightText: 2025 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
+// SPDX-License-Identifier: GPL-3.0-only
+
+#include "PostIndexingRefinement.h"
+
+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;
+        const Eigen::MatrixX3<float> resid = CalculateResiduals(spots, cell);
+
+        std::vector<uint8_t> mask(spots.rows(), 0);
+        indexed_spot_count = 0;
+
+        for (int i = 0; i < spots.rows(); ++i) {
+            if (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;
+        } else {
+            if (row_norms.minCoeff() < p.min_length_A || row_norms.maxCoeff() > p.max_length_A)
+                continue;
+
+            float alpha = std::acos(cell_rows.row(1).normalized().dot(cell_rows.row(2).normalized())) * 180.0f / M_PI;
+            float beta = std::acos(cell_rows.row(0).normalized().dot(cell_rows.row(2).normalized())) * 180.0f / M_PI;
+            float gamma = std::acos(cell_rows.row(0).normalized().dot(cell_rows.row(1).normalized())) * 180.0f / M_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) {
+        bool too_similar = false;
+
+        for (const auto &selected: accepted) {
+            int64_t overlap = 0;
+            for (size_t i = 0; i < candidate.indexed_mask.size(); ++i) {
+                if (candidate.indexed_mask[i] && selected.indexed_mask[i])
+                    overlap++;
+            }
+
+            const int64_t max_set_size = std::max(candidate.indexed_spot_count, selected.indexed_spot_count);
+            if (overlap > static_cast<int64_t>(REFINE_CANDIDATE_OVERLAP_RATIO_THRESHOLD
+                                               * static_cast<float>(max_set_size))) {
+                too_similar = true;
+                break;
+            }
+        }
+
+        if (!too_similar)
+            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;
+}

image_analysis/indexing/PostIndexingRefinement.h

@@ -0,0 +1,36 @@
+// SPDX-FileCopyrightText: 2025 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
+// SPDX-License-Identifier: GPL-3.0-only
+
+#pragma once
+
+#include <vector>
+#include <optional>
+#include <algorithm>
+#include <cmath>
+#include <cstdint>
+#include <Eigen/Dense>
+
+#include "../common/CrystalLattice.h"
+
+constexpr float REFINE_CANDIDATE_SPOT_COUNT_RATIO_THRESHOLD = 0.9f;
+constexpr float REFINE_CANDIDATE_VOLUME_RATIO_THRESHOLD = 1.05f;
+constexpr float REFINE_CANDIDATE_OVERLAP_RATIO_THRESHOLD = 0.2f;
+constexpr float REFINE_MIN_VOLUME_EPSILON = 1e-12f;
+constexpr float REFINE_MIN_REFERENCE_LENGTH_EPSILON = 1e-6f;
+
+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;
+};
+
+std::vector<CrystalLattice> Refine(const std::vector<Coord> &in_spots,
+                                          size_t nspots,
+                                          Eigen::MatrixX3<float> &oCell,
+                                          Eigen::VectorX<float> &scores,
+                                          RefineParameters &p);