139 lines
4.3 KiB
C++
139 lines
4.3 KiB
C++
// SPDX-FileCopyrightText: 2025 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
|
|
// SPDX-License-Identifier: GPL-3.0-only
|
|
|
|
#include "FFTIndexer.h"
|
|
#include "EigenRefine.h"
|
|
|
|
std::vector<Coord> FFTIndexer::FilterFFTResults() const {
|
|
std::multimap<float, FFTResult> fft_result_map;
|
|
|
|
// Order vectors by max FFT magnitude and select the 30 strongest ones
|
|
size_t max_vectors = 30;
|
|
|
|
for (int i = 0; i < direction_vectors.size(); i++)
|
|
fft_result_map.insert(std::make_pair(result_fft[i].magnitude, result_fft[i]));
|
|
|
|
std::vector<FFTResult> fft_result_filtered;
|
|
int count = 0;
|
|
for (auto it = fft_result_map.rbegin();
|
|
it != fft_result_map.rend() && count < max_vectors;
|
|
++it, ++count) {
|
|
fft_result_filtered.emplace_back(it->second);
|
|
}
|
|
|
|
std::vector<Coord> ret;
|
|
|
|
// Remove vectors less than 5 deg apart, as most likely these are colinear
|
|
constexpr float COS_5_DEG = 0.996194;
|
|
std::vector<bool> ignore(fft_result_filtered.size(), false);
|
|
|
|
for (int i = 0; i < fft_result_filtered.size(); i++) {
|
|
if (ignore[i])
|
|
continue;
|
|
Coord dir_i = direction_vectors.at(fft_result_filtered[i].direction);
|
|
for (int j = i + 1; j < fft_result_filtered.size(); j++) {
|
|
Coord dir_j = direction_vectors.at(fft_result_filtered[j].direction);
|
|
if (std::fabs(dir_i * dir_j) > COS_5_DEG)
|
|
ignore[j] = true;
|
|
}
|
|
ret.push_back(dir_i * fft_result_filtered[i].length);
|
|
}
|
|
return ret;
|
|
}
|
|
|
|
void Sort(Coord &A, Coord &B, Coord &C) {
|
|
// A is the smallest, C is the largest
|
|
|
|
if (A.Length() > B.Length())
|
|
std::swap(A, B);
|
|
|
|
if (B.Length() > C.Length())
|
|
std::swap(B, C);
|
|
|
|
if (A.Length() > B.Length())
|
|
std::swap(A, B);
|
|
}
|
|
|
|
std::vector<CrystalLattice> FFTIndexer::ReduceResults(const std::vector<Coord> &results) const {
|
|
if (results.size() < 3)
|
|
return {};
|
|
|
|
std::vector<CrystalLattice> candidates;
|
|
|
|
for (int i = 0; i < 3; i++) {
|
|
for (int j = 0; j < 3; j++) {
|
|
for (int k = 0; k < 3; k++) {
|
|
if (i + j + k + 2 >= results.size())
|
|
break;
|
|
|
|
Coord A = results[i];
|
|
Coord B = results[(i + j + 1)];
|
|
Coord C = results[(i + j + 1) + k + 1];
|
|
|
|
// sort vectors by length for reduction
|
|
Sort(A,B,C);
|
|
|
|
// Lattice reduction
|
|
B = B - std::round(B * A / (A * A)) * A;
|
|
C = C - std::round(C * A / (A * A)) * A;
|
|
C = C - std::round(C * B / (B * B)) * B;
|
|
|
|
// Check if values are OK after lattice reduction
|
|
if (A.Length() >= min_length_A && B.Length() >= min_length_A && C.Length() >= min_length_A)
|
|
candidates.emplace_back(A, B, C);
|
|
}
|
|
}
|
|
}
|
|
|
|
return candidates;
|
|
}
|
|
|
|
void FFTIndexer::SetupUnitCell(const std::optional<UnitCell> &cell) {
|
|
reference_unit_cell = cell;
|
|
}
|
|
|
|
std::vector<CrystalLattice> FFTIndexer::Run(const std::vector<Coord> &coord, size_t nspots) {
|
|
if (nspots > coord.size())
|
|
nspots = coord.size();
|
|
|
|
if (nspots <= viable_cell_min_spots)
|
|
return {};
|
|
|
|
assert(nspots <= FFT_MAX_SPOTS);
|
|
assert(coord.size() <= FFT_MAX_SPOTS);
|
|
|
|
ExecuteFFT(coord, nspots);
|
|
auto f = FilterFFTResults();
|
|
auto r = ReduceResults(f);
|
|
|
|
Eigen::MatrixX3<float> oCell(r.size() * 3u, 3u);
|
|
Eigen::VectorX<float> scores(r.size());
|
|
|
|
for (int i = 0; i < r.size(); i++) {
|
|
oCell(i * 3u, 0u) = r[i].Vec0().x;
|
|
oCell(i * 3u, 1u) = r[i].Vec0().y;
|
|
oCell(i * 3u, 2u) = r[i].Vec0().z;
|
|
|
|
oCell(i * 3u + 1, 0u) = r[i].Vec1().x;
|
|
oCell(i * 3u + 1, 1u) = r[i].Vec1().y;
|
|
oCell(i * 3u + 1, 2u) = r[i].Vec1().z;
|
|
|
|
oCell(i * 3u + 2, 0u) = r[i].Vec2().x;
|
|
oCell(i * 3u + 2, 1u) = r[i].Vec2().y;
|
|
oCell(i * 3u + 2, 2u) = r[i].Vec2().z;
|
|
|
|
// Bootstrap score
|
|
scores(i) = 0.2;
|
|
}
|
|
|
|
RefineParameters parameters{
|
|
.viable_cell_min_spots = viable_cell_min_spots,
|
|
.dist_tolerance_vs_reference = dist_tolerance_vs_reference,
|
|
.reference_unit_cell = reference_unit_cell,
|
|
.min_length_A = min_length_A,
|
|
.max_length_A = max_length_A,
|
|
.indexing_tolerance = indexing_tolerance
|
|
};
|
|
|
|
return Refine(coord, nspots, oCell, scores, parameters);
|
|
} |