Jungfraujoch/image_analysis/spot_finding/SpotUtils.cpp

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

#include "../../common/JFJochMath.h"
#include "SpotUtils.h"
#include "../../common/ResolutionShells.h"

void CountSpots(DataMessage &msg,
                const std::vector<SpotToSave> &spots,
                float d_min_A) {
    int64_t low_res = 0;
    int64_t ice_ring = 0;
    for (auto &s: spots) {
        if (s.ice_ring)
            ice_ring++;

        if (s.d_A > d_min_A)
            low_res++;
    }
    msg.spot_count = spots.size();
    msg.spot_count_low_res = low_res;
    msg.spot_count_ice_rings = ice_ring;
}

void MarkIceRings(std::vector<SpotToSave> &spots, float tolerance_q_recipA) {
    std::vector<float> ice_rings_q;

    for (const auto &i: ICE_RING_RES_A)
        ice_rings_q.push_back(2 * PI / i);

    for (auto &s: spots) {
        auto spot_q = 2 * PI / s.d_A;
        bool tmp = false;
        for (const auto &q: ice_rings_q)
            tmp |= (fabs(spot_q - q) < tolerance_q_recipA);
        s.ice_ring = tmp;
    }
}

void FilterSpotsByCount(std::vector<SpotToSave> &input, int64_t count, bool by_weight) {
    size_t output_size = std::min<size_t>(input.size(), count);

    std::ranges::partial_sort(input, input.begin() + output_size,
                              std::ranges::less{}, // comparator on the projected key
                              [by_weight](const SpotToSave &s) {
                                  // projection: key to compare by. non-ice first (false < true), then
                                  // by soft quality weight (higher first) when requested -- so a loose
                                  // detector's bright junk cannot evict faint clean Bragg -- else by
                                  // raw intensity. Intensity is the tie-breaker under the weight.
                                  if (by_weight)
                                      return std::tuple{s.ice_ring, -s.weight, -s.intensity};
                                  return std::tuple{s.ice_ring, 0.0f, -s.intensity};
                              });
    input.resize(output_size);
}

void FilterSpuriousHighResolutionSpots(std::vector<SpotToSave> &spots, float threshold) {
    std::ranges::sort(spots, [](SpotToSave &a, SpotToSave &b) {
        return a.d_A > b.d_A;
    });

    // Apply 1/d gap threshold: find first gap in q = 1/d exceeding dist_threshold and ignore spots after it
    if (spots.size() >= 2 && threshold > 0.0f) {
        size_t cut_index = spots.size(); // default: keep all
        // d_A sorted descending → q = 1/d_A sorted ascending
        // We check consecutive q gaps: Δq_i = (1/d_i) - (1/d_{i+1})
        for (size_t i = 0; i + 1 < spots.size(); ++i) {
            float d1 = spots[i].d_A;
            float d2 = spots[i + 1].d_A;
            // Avoid division by zero; d_A should be > 0 in valid data
            if (d1 <= 0.0f || d2 <= 0.0f)
                continue;
            float q1 = 2 * PI / d1;
            float q2 = 2 * PI / d2;
            float dq = q2 - q1; // should be >= 0 due to sorting
            if (dq > threshold) {
                cut_index = i + 1; // keep up to i inclusive
                break;
            }
        }
        if (cut_index < spots.size())
            spots.resize(cut_index);
    }
}

std::optional<float> GetResolution(const std::vector<SpotToSave> &spots) {
    std::vector<float> resolutions;
    resolutions.reserve(spots.size());

    for (const auto &spot: spots) {
        if (!spot.ice_ring)
            resolutions.push_back(spot.d_A);
    }

    std::ranges::sort(resolutions);

    if (resolutions.size() < 4)
        return std::nullopt;
    if (resolutions.size() < 20)
        return resolutions[2];
    return resolutions[static_cast<size_t>(resolutions.size() * 0.05)];
}

void GenerateSpotPlot(DataMessage &msg, const std::vector<SpotToSave> &spots, float d_min_A) {
    const int nshells = 20;
    ResolutionShells shells(d_min_A, 50.0, nshells);

    std::vector<float> intensity(nshells);
    std::vector<float> count(nshells);

    for (const auto &s: spots) {
        if (s.ice_ring)
            continue;

        if (auto shell = shells.GetShell(s.d_A)) {
            intensity[*shell] += s.intensity;
            count[*shell] += 1.0f;
        }
    }

    std::vector<float> result(nshells);
    for (int i = 0; i < nshells; ++i) {
        if (count[i] > 0)
            result[i] = intensity[i] / count[i];
        else
            result[i] = 0.0f;
    }

    msg.spot_plot_one_over_d_square = shells.GetShellMeanOneOverResSq();
    msg.spot_plot_intensity = result;
    msg.spot_plot_count = count;
}

void SpotAnalyze(const DiffractionExperiment &experiment,
                 const SpotFindingSettings &spot_finding_settings,
                 const std::vector<DiffractionSpot> &spots,
                 DataMessage &output) {
    auto geom = experiment.GetDiffractionGeometry();

    std::vector<SpotToSave> spots_out;

    for (const auto &spot: spots) {
        if (auto s = spot.Export(geom, output.number); s.has_value())
            spots_out.push_back(s.value());
    }

    if (spot_finding_settings.high_res_gap_Q_recipA.has_value())
        FilterSpuriousHighResolutionSpots(spots_out, spot_finding_settings.high_res_gap_Q_recipA.value());

    if (experiment.GetDatasetSettings().IsDetectIceRings() && spot_finding_settings.ice_ring_width_Q_recipA > 0.0f)
        MarkIceRings(spots_out, spot_finding_settings.ice_ring_width_Q_recipA);

    CountSpots(output, spots_out, spot_finding_settings.cutoff_spot_count_low_res);

    GenerateSpotPlot(output, spots_out, spot_finding_settings.high_resolution_limit);

    output.resolution_estimate = GetResolution(spots_out);

    FilterSpotsByCount(spots_out, experiment.GetMaxSpotCount(), spot_finding_settings.soft_weight);

    output.spots = spots_out;
}