Jungfraujoch/image_analysis/IndexAndRefine.cpp

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

#include "IndexAndRefine.h"

#include "bragg_integration/BraggIntegrate2D.h"
#include "bragg_integration/CalcISigma.h"
#include "geom_refinement/XtalOptimizer.h"
#include "indexing/AnalyzeIndexing.h"
#include "indexing/FFTIndexer.h"
#include "lattice_search/LatticeSearch.h"

IndexAndRefine::IndexAndRefine(const DiffractionExperiment &x, IndexerThreadPool *indexer)
    : index_ice_rings(x.GetIndexingSettings().GetIndexIceRings()),
      experiment(x),
      geom_(x.GetDiffractionGeometry()),
      indexer_(indexer) {
    if (indexer && x.GetGoniometer() && x.GetIndexingSettings().GetRotationIndexing())
        rotation_indexer = std::make_unique<RotationIndexer>(x, *indexer);
}

void IndexAndRefine::SetLattice(const CrystalLattice &lattice) {
    if (rotation_indexer)
        rotation_indexer->SetLattice(lattice);
}

IndexAndRefine::IndexingOutcome IndexAndRefine::DetermineLatticeAndSymmetry(DataMessage &msg) {
    IndexingOutcome outcome(experiment);

    if (rotation_indexer) {
        auto result = rotation_indexer->ProcessImage(msg.number, msg.spots);
        if (result) {
            outcome.lattice_candidate = result->lattice;

            outcome.experiment.BeamX_pxl(result->geom.GetBeamX_pxl())
                                  .BeamY_pxl(result->geom.GetBeamY_pxl())
                                  .DetectorDistance_mm(result->geom.GetDetectorDistance_mm());

            outcome.symmetry.centering = result->search_result.centering;
            outcome.symmetry.niggli_class = result->search_result.niggli_class;
            outcome.symmetry.crystal_system = result->search_result.system;
        }
        return outcome;
    }

    // Convert input spots to reciprocal space
    std::vector<Coord> recip;
    recip.reserve(msg.spots.size());
    for (const auto &i: msg.spots) {
        if (index_ice_rings || !i.ice_ring)
            recip.push_back(i.ReciprocalCoord(geom_));
    }

    auto indexer_result = indexer_->Run(experiment, recip).get();
    msg.indexing_time_s = indexer_result.indexing_time_s;

    if (indexer_result.lattice.empty())
        return outcome;

    auto latt = indexer_result.lattice[0];
    auto sg = experiment.GetGemmiSpaceGroup();

    // If space group provided => always enforce symmetry in refinement
    // If space group not provided => guess symmetry
    if (sg) {
        // If space group provided but no unit cell fixed, it is better to keep triclinic for now
        if (experiment.GetUnitCell()) {
            outcome.symmetry = LatticeMessage{
                .centering = sg->centring_type(),
                .niggli_class = 0,
                .crystal_system = sg->crystal_system()
            };
        }
        outcome.lattice_candidate = latt;
    } else {
        auto sym_result = LatticeSearch(latt);
        outcome.symmetry = LatticeMessage{
            .centering = sym_result.centering,
            .niggli_class = sym_result.niggli_class,
            .crystal_system = sym_result.system
        };
        outcome.lattice_candidate = sym_result.conventional;
    }

    return outcome;
}

void IndexAndRefine::RefineGeometryIfNeeded(DataMessage &msg, IndexAndRefine::IndexingOutcome &outcome) {
    if (!outcome.lattice_candidate)
        return;

    XtalOptimizerData data{
        .geom = outcome.experiment.GetDiffractionGeometry(),
        .latt = *outcome.lattice_candidate,
        .crystal_system = outcome.symmetry.crystal_system,
        .min_spots = experiment.GetIndexingSettings().GetViableCellMinSpots(),
        .refine_beam_center = true,
        .refine_distance_mm = false,
        .max_time = 0.04 // 40 ms is max allowed time for the operation
    };

    if (outcome.symmetry.crystal_system == gemmi::CrystalSystem::Trigonal)
        data.crystal_system = gemmi::CrystalSystem::Hexagonal;

    switch (experiment.GetIndexingSettings().GetGeomRefinementAlgorithm()) {
        case GeomRefinementAlgorithmEnum::None:
            break;
        case GeomRefinementAlgorithmEnum::BeamCenter:
            if (XtalOptimizer(data, msg.spots)) {
                outcome.experiment.BeamX_pxl(data.geom.GetBeamX_pxl())
                                      .BeamY_pxl(data.geom.GetBeamY_pxl());
                outcome.beam_center_updated = true;
            }
            break;
    }

    outcome.lattice_candidate = data.latt;

    if (outcome.beam_center_updated) {
        msg.beam_corr_x = data.beam_corr_x;
        msg.beam_corr_y = data.beam_corr_y;
    }
}

void IndexAndRefine::QuickPredictAndIntegrate(DataMessage &msg,
                                              const SpotFindingSettings &spot_finding_settings,
                                              const CompressedImage &image,
                                              BraggPrediction &prediction,
                                              const IndexAndRefine::IndexingOutcome &outcome) {
    if (!spot_finding_settings.quick_integration)
        return;
    if (!outcome.lattice_candidate)
        return;

    float ewald_dist_cutoff = 0.001f;

    if (msg.profile_radius)
        ewald_dist_cutoff = msg.profile_radius.value() * 2.0f;
    if (experiment.GetBraggIntegrationSettings().GetFixedProfileRadius_recipA())
        ewald_dist_cutoff = experiment.GetBraggIntegrationSettings().GetFixedProfileRadius_recipA().value() * 3.0f;

    BraggPredictionSettings settings_prediction{
        .high_res_A = experiment.GetBraggIntegrationSettings().GetDMinLimit_A(),
        .ewald_dist_cutoff = ewald_dist_cutoff,
        .max_hkl = 100,
        .centering = outcome.symmetry.centering
    };

    prediction.Calc(experiment, *outcome.lattice_candidate, settings_prediction);

    auto refl = prediction.GetReflections();

    BraggIntegrate2D(outcome.experiment, image, refl, msg.number);

    std::vector<Reflection> refl_ret;
    refl_ret.reserve(refl.size());

    for (const auto& r : refl) {
        if (r.observed)
            refl_ret.push_back(r);
    }

    constexpr size_t kMaxReflections = 20000;
    if (refl_ret.size() > kMaxReflections) {
        // Keep only smallest d (highest resolution)
        std::nth_element(refl_ret.begin(),
                         refl_ret.begin() + static_cast<long>(kMaxReflections),
                         refl_ret.end(),
                         [](const Reflection& a, const Reflection& b) {
                             return a.d < b.d;
                         });
        refl_ret.resize(kMaxReflections);

        // Optional: make output ordered by d (nice for downstream / debugging)
        std::sort(refl_ret.begin(), refl_ret.end(),
                  [](const Reflection& a, const Reflection& b) { return a.d < b.d; });
    }

    msg.reflections = std::move(refl_ret);

    CalcISigma(msg);
    CalcWilsonBFactor(msg);
}

void IndexAndRefine::ProcessImage(DataMessage &msg,
                                  const SpotFindingSettings &spot_finding_settings,
                                  const CompressedImage &image,
                                  BraggPrediction &prediction) {
    if (!indexer_ || !spot_finding_settings.indexing)
        return;

    msg.indexing_result = false;

    IndexingOutcome outcome = DetermineLatticeAndSymmetry(msg);
    if (!outcome.lattice_candidate)
        return;

    RefineGeometryIfNeeded(msg, outcome);
    if (!outcome.lattice_candidate)
        return;

    if (!AnalyzeIndexing(msg, outcome.experiment, *outcome.lattice_candidate, outcome.experiment.GetGoniometer()))
        return;

    msg.lattice_type = outcome.symmetry;

    QuickPredictAndIntegrate(msg, spot_finding_settings, image, prediction, outcome);
}

std::optional<RotationIndexerResult> IndexAndRefine::Finalize() {
    if (rotation_indexer)
        return rotation_indexer->GetLattice();
    return {};
}