Jungfraujoch/image_analysis/indexing/FFTIndexer.h

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

#pragma once

#include "Indexer.h"
#include "../../common/UnitCell.h"
#include "../../common/Coord.h"
#include "../../common/CrystalLattice.h"
#include "../../common/IndexingSettings.h"
#include "FFTResult.h"
#include "SpindleBlindFraction.h"

#define FFT_MAX_SPOTS (64*1024)

class FFTIndexer : public Indexer {

protected:
    const float min_length_A;
    const float max_length_A;
    const float min_angle_deg;
    const float max_angle_deg;

    const int nDirections;
    const unsigned refine_threads;
    float histogram_spacing;
    int64_t histogram_size;
    std::vector<Coord> direction_vectors;

    size_t input_size;
    size_t output_size;

    // Per-direction best peaks (host)
    std::vector<FFTResult> result_fft;

    void SetupUnitCell(const std::optional<UnitCell> &cell) override;
    // widen=false: combine only the shortest few filtered vectors (fast, correct for compact cells).
    // widen=true: anchor the two short axes but let the third range over all vectors (reaches the long
    // axis of a large/elongated cell). Used only as a fallback when the standard reduction indexes poorly.
    std::vector<CrystalLattice> ReduceResults(const std::vector<Coord> &results, bool widen) const;
    // `magnitudes`, when given, receives the FFT peak prominence of each returned row, in the
    // same order. The spindle severity needs it to tell a crystal's real rows from the short
    // spurious ones a long cell produces; nothing else does, so it is optional.
    std::vector<Coord> FilterFFTResults(size_t max_vectors,
                                        std::vector<float> *magnitudes = nullptr) const;
    // Replace each row whose FFT length is an alias of a longer one by that longer row (see the definition).
    void DeAlias(const std::vector<Coord> &coord, size_t nspots, std::vector<Coord> &rows) const;
    std::vector<CrystalLattice> ReduceAndRefine(const std::vector<Coord> &coord, size_t nspots,
                                                const std::vector<Coord> &filtered, bool widen);
    float IndexedFraction(const CrystalLattice &latt, const std::vector<Coord> &coord, size_t nspots) const;
    void SetupDirectionVectors();
    // Replace the search directions with a small cap around `axis` (same count, so the transform and
    // its plan are unchanged) and return the strongest peak found in it.
    std::optional<Coord> SearchCap(const std::vector<Coord> &coord, size_t nspots,
                                   const Coord &axis, float half_angle_deg);
    // Direction perpendicular to every vector of `filtered`, when they all lie in one plane and so
    // cannot close a cell; empty otherwise.
    std::optional<Coord> DegeneratePlaneNormal(const std::vector<Coord> &filtered) const;
    // Filled by RunInternal from the shortlist it already computes; read out by Indexer::Run.
    std::optional<SpindleSeverity> spindle_severity;
    std::optional<SpindleSeverity> GetSpindleSeverity() const override { return spindle_severity; }
    std::optional<SpindleSeverity> RunSeverityOnly(const std::vector<Coord> &coord) override;

    virtual void ExecuteFFT(const std::vector<Coord> &coord, size_t nspots) = 0;
    // Called after direction_vectors is rewritten, for implementations that keep a copy of it.
    virtual void DirectionsChanged() {}
public:

    explicit FFTIndexer(const IndexingSettings& settings);
    ~FFTIndexer() override = default;

    std::vector<CrystalLattice> RunInternal(const std::vector<Coord> &coord, size_t nspots) override;
};