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"
#include "scale_merge/ScaleOnTheFly.h"

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

    mosaicity.resize(x.GetImageNum(), NAN);
    scale_cc.resize(x.GetImageNum(), 0);
    unit_cells.resize(x.GetImageNum());
}

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

    auto result = rotation_indexer->ProcessImage(msg.number, msg.spots);
    if (result) {
        // For rotation indexing, indexing rate is calculated only for frames, where "global" rotation indexing solution was found
        msg.indexing_result = false;

        // get rotated lattice
        auto gon = result->axis;
        if (gon) {
            const float angle_deg = gon->GetAngle_deg(msg.number) + gon->GetWedge_deg() / 2.0f;
            outcome.lattice_candidate = result->lattice.Multiply(gon->GetTransformationAngle(-angle_deg));
        }

        outcome.experiment.BeamX_pxl(result->geom.GetBeamX_pxl())
                .BeamY_pxl(result->geom.GetBeamY_pxl())
                .DetectorDistance_mm(result->geom.GetDetectorDistance_mm())
                .PoniRot1_rad(result->geom.GetPoniRot1_rad())
                .PoniRot2_rad(result->geom.GetPoniRot2_rad())
        .Goniometer(result->axis);
        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;
}

IndexAndRefine::IndexingOutcome IndexAndRefine::DetermineLatticeAndSymmetry(DataMessage &msg) {
    auto indexing_start_time = std::chrono::steady_clock::now();

    IndexingOutcome outcome(experiment);

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

    if (indexer_result.executed)
        msg.indexing_result = false;

    if (!indexer_result.lattice.empty()) {
        auto latt = indexer_result.lattice[0];
        if (latt.CalcVolume() > 1.0) {
            auto sg = experiment.GetGemmiSpaceGroup();

            // If space group and cell provided => always enforce symmetry in refinement
            // If space group not provided => guess symmetry
            if (sg && 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;
            }
        }
    }

    auto indexing_end_time = std::chrono::steady_clock::now();
    msg.indexing_time_s = std::chrono::duration<float>(indexing_end_time - indexing_start_time).count();

    return outcome;
}

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

    auto start_time = std::chrono::steady_clock::now();

    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,
        .refine_detector_angles = false,
        .refine_unit_cell = !experiment.IsRotationIndexing(),
        .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.symmetry.crystal_system == gemmi::CrystalSystem::Monoclinic)
        outcome.lattice_candidate->ReorderMonoclinic();

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

    auto end_time = std::chrono::steady_clock::now();
    msg.refinement_time_s = std::chrono::duration_cast<std::chrono::duration<double>>(end_time - start_time).count();
}

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


    CrystalLattice latt = outcome.lattice_candidate.value();

    if (rotation_indexer) {
        // Use moving average for mosaicity and profile_radius (also add beam center later)
        if (msg.mosaicity_deg)
            msg.mosaicity_deg = rotation_parameters.Mosaicity(msg.mosaicity_deg.value());
        if (msg.profile_radius) {
            msg.profile_radius = rotation_parameters.ProfileRadius(msg.profile_radius.value());
        }
    }

    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;

    float wedge_deg = 0.0f;
    float mos_deg = 0.1f;

    if (experiment.GetGoniometer().has_value()) {
        wedge_deg = experiment.GetGoniometer()->GetWedge_deg() / 2.0;

        if (msg.mosaicity_deg) {
            mos_deg = msg.mosaicity_deg.value();
            mosaicity[msg.number] = mos_deg;
        }
    }

    IntegrationOutcome i_outcome{
        .geom = outcome.experiment.GetDiffractionGeometry(),
        .latt = latt,
        .mosaicity_deg = mos_deg,
        .image_scale_b_factor_Ang2 = msg.image_scale_b_factor,
        .image_scale_cc = msg.image_scale_cc,
    };

    const BraggPredictionSettings settings_prediction{
        .high_res_A = experiment.GetBraggIntegrationSettings().GetDMinLimit_A(),
        .ewald_dist_cutoff = ewald_dist_cutoff,
        .max_hkl = 100,
        .centering = outcome.symmetry.centering,
        .wedge_deg = std::fabs(wedge_deg),
        .mosaicity_deg = std::fabs(mos_deg)
    };

    auto pred_start_time = std::chrono::steady_clock::now();
    auto nrefl = prediction.Calc(outcome.experiment, latt, settings_prediction);
    auto pred_end_time = std::chrono::steady_clock::now();
    msg.bragg_prediction_time_s = std::chrono::duration<float>(pred_end_time - pred_start_time).count();

    auto integration_start_time = std::chrono::steady_clock::now();
    i_outcome.reflections = BraggIntegrate2D(outcome.experiment, image, prediction.GetReflections(), nrefl, msg.number);
    msg.integrated_reflections = i_outcome.reflections.size();

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

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

    CalcISigma(msg, i_outcome.reflections);
    CalcWilsonBFactor(msg, i_outcome.reflections);

    auto integration_end_time = std::chrono::steady_clock::now();
    msg.integration_time_s = std::chrono::duration<float>(integration_end_time - integration_start_time).count();

    ScaleImage(msg, i_outcome);

    // Copy reflections to outgoing message
    msg.reflections = i_outcome.reflections;

    {
        const std::unique_lock ul(reflections_mutex);
        integration_outcome[msg.number] = std::move(i_outcome);
    }
}

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

    IndexingOutcome outcome(experiment);
  
    if (rotation_indexer)
        outcome = DetermineLatticeAndSymmetryRotation(msg);
    else
        outcome = DetermineLatticeAndSymmetry(msg);

    if (!outcome.lattice_candidate)
        return;

    if (experiment.GetIndexingSettings().GetGeomRefinementAlgorithm() != GeomRefinementAlgorithmEnum::None)
        RefineGeometryIfNeeded(msg, outcome);

    if (!outcome.lattice_candidate.has_value())
        return;

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

    {
        std::unique_lock ul(reflections_mutex);
        unit_cells[msg.number] = outcome.lattice_candidate->GetUnitCell();
    }
    msg.lattice_type = outcome.symmetry;

    if (spot_finding_settings.quick_integration)
        QuickPredictAndIntegrate(msg, spot_finding_settings, image, prediction, outcome);
}

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

IndexAndRefine &IndexAndRefine::ReferenceIntensities(std::vector<MergedReflection> &reference) {
    scaling_engine = std::make_unique<ScaleOnTheFly>(experiment, reference);
    return *this;
}

void IndexAndRefine::ScaleImage(DataMessage &msg, IntegrationOutcome& outcome) {
    if (!scaling_engine)
        return;

    auto scaling_start_time = std::chrono::steady_clock::now();
    scaling_engine->Scale(outcome);
    auto scaling_end_time = std::chrono::steady_clock::now();

    scale_cc[msg.number] = outcome.image_scale_cc.value_or(NAN);
    msg.image_scale_b_factor = outcome.image_scale_b_factor_Ang2;
    msg.image_scale_factor = outcome.image_scale_g;
    msg.image_scale_mosaicity = outcome.mosaicity_deg;
    msg.image_scale_cc = outcome.image_scale_cc;
    msg.image_scale_time_s = std::chrono::duration<float>(scaling_end_time - scaling_start_time).count();
}

ScalingResult IndexAndRefine::ScaleAllImages(const std::vector<MergedReflection> &reference, size_t nthreads) {
    ScaleOnTheFly scaling(experiment, reference);
    scaling.Scale(integration_outcome, nthreads);
    scale_cc.resize(integration_outcome.size());

    for (int i = 0; i < integration_outcome.size(); i++)
        scale_cc.at(i) = integration_outcome[i].image_scale_cc.value_or(NAN);

    return ScalingResult(integration_outcome);
}

const std::vector<float> &IndexAndRefine::GetImageCC() const {
    return scale_cc;
}

const std::vector<std::optional<UnitCell> > & IndexAndRefine::GetUnitCells() const {
    return unit_cells;
}

std::optional<UnitCell> IndexAndRefine::GetConsensusUnitCell() const {
    const auto dist_tolerance = experiment.GetIndexingSettings().GetUnitCellDistTolerance();
    const auto angle_tolerance = experiment.GetIndexingSettings().GetUnitCellAngleTolerance_deg();

    if (rotation_indexer) {
        auto result = rotation_indexer->GetLattice();
        if (!result)
            return {};
        return result->lattice.GetUnitCell();
    }

    std::vector<UnitCell> cells;
    {
        std::unique_lock ul(reflections_mutex);
        cells.reserve(unit_cells.size());
        for (const auto &cell: unit_cells) {
            if (cell && cell->is_finite())
                cells.emplace_back(*cell);
        }
    }

    if (cells.empty())
        return {};

    if (experiment.GetUnitCell()) {
        std::vector<UnitCell> accepted;
        accepted.reserve(cells.size());

        for (const auto &cell: cells) {
            if (cell.is_close(*experiment.GetUnitCell(), dist_tolerance, angle_tolerance))
                accepted.emplace_back(cell);
        }

        return MeanUnitCell(accepted);
    }

    size_t best_count = 0;
    UnitCell best_reference{};

    for (const auto &ref: cells) {
        size_t count = 0;
        for (const auto &cell: cells) {
            if (cell.is_close(ref, dist_tolerance, angle_tolerance))
                ++count;
        }

        if (count > best_count) {
            best_count = count;
            best_reference = ref;
        }
    }

    if (best_count == 0)
        return {};

    std::vector<UnitCell> accepted;
    accepted.reserve(best_count);

    for (const auto &cell: cells) {
        if (cell.is_close(best_reference, dist_tolerance, angle_tolerance))
            accepted.emplace_back(cell);
    }

    return MeanUnitCell(accepted);
}

std::vector<IntegrationOutcome> &IndexAndRefine::GetIntegrationOutcome() {
    return integration_outcome;
}

const std::vector<IntegrationOutcome> &IndexAndRefine::GetIntegrationOutcome() const {
    return integration_outcome;
}