Jungfraujoch/image_analysis/indexing/AnalyzeIndexing.cpp

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

#include <cstdint>
#include <vector>
#include "AnalyzeIndexing.h"

#include "FitProfileRadius.h"

namespace {
    inline bool ok(float x) {
        if (!std::isfinite(x))
            return false;
        if (x < 0.0)
            return false;
        return true;
    }

    inline float deg_to_rad(float deg) {
        return deg * (static_cast<float>(M_PI) / 180.0f);
    }

    inline float rad_to_deg(float rad) {
        return rad * (180.0f / static_cast<float>(M_PI));
    }

    // Wrap to [-180, 180] (useful for residuals)
    inline float wrap_deg_pm180(float deg) {
        while (deg > 180.0f) deg -= 360.0f;
        while (deg < -180.0f) deg += 360.0f;
        return deg;
    }

    // Solve A cos(phi) + B sin(phi) + D = 0, return solutions in [phi0, phi1] (radians)
    inline int solve_trig(float A, float B, float D,
                          float phi0, float phi1,
                          float out_phi[2]) {
        const float R = std::sqrt(A * A + B * B);
        if (!(R > 0.0f))
            return 0;

        const float rhs = -D / R;
        if (rhs < -1.0f || rhs > 1.0f)
            return 0;

        const float phi_ref = std::atan2(B, A);
        const float delta = std::acos(rhs);

        float s1 = phi_ref + delta;
        float s2 = phi_ref - delta;

        const float two_pi = 2.0f * static_cast<float>(M_PI);
        auto shift_near = [&](float x) {
            while (x < phi0 - two_pi) x += two_pi;
            while (x > phi1 + two_pi) x -= two_pi;
            return x;
        };

        s1 = shift_near(s1);
        s2 = shift_near(s2);

        int n = 0;
        if (s1 >= phi0 && s1 <= phi1) out_phi[n++] = s1;
        if (s2 >= phi0 && s2 <= phi1) {
            if (n == 0 || std::fabs(s2 - out_phi[0]) > 1e-6f) out_phi[n++] = s2;
        }
        return n;
    }

    // Find predicted phi (deg) for given g0 around phi_obs (deg) within +/- half_window_deg.
    // Returns nullopt if no solution in the local window.
    inline std::optional<float> predict_phi_deg_local(const Coord &g0,
                                                      const Coord &S0,
                                                      const Coord &w_unit,
                                                      float phi_obs_deg,
                                                      float half_window_deg) {
        const float phi0 = deg_to_rad(phi_obs_deg - half_window_deg);
        const float phi1 = deg_to_rad(phi_obs_deg + half_window_deg);

        // Decompose g0 into parallel/perp to w
        const float g_par_s = g0 * w_unit;
        const Coord g_par = w_unit * g_par_s;
        const Coord g_perp = g0 - g_par;

        const float g_perp2 = g_perp * g_perp;
        if (g_perp2 < 1e-12f)
            return std::nullopt;

        const float k2 = (S0 * S0); // |S0|^2 = (1/lambda)^2

        // Equation: |S0 + g(phi)|^2 = |S0|^2
        const Coord p = S0 + g_par;
        const Coord w_x_gperp = w_unit % g_perp;

        const float A = 2.0f * (p * g_perp);
        const float B = 2.0f * (p * w_x_gperp);
        const float D = (p * p) + g_perp2 - k2;

        float sols[2]{};
        const int nsol = solve_trig(A, B, D, phi0, phi1, sols);
        if (nsol == 0)
            return std::nullopt;

        // Pick the solution closest to phi_obs
        const float phi_obs = deg_to_rad(phi_obs_deg);
        float best_phi = sols[0];
        float best_err = std::fabs(sols[0] - phi_obs);

        if (nsol == 2) {
            const float err2 = std::fabs(sols[1] - phi_obs);
            if (err2 < best_err) {
                best_err = err2;
                best_phi = sols[1];
            }
        }

        return rad_to_deg(best_phi);
    }

    std::optional<float> CalcMosaicity(const DiffractionExperiment& experiment,
                                       const std::vector<SpotToSave> &spots,
                                       const Coord &astar, const Coord &bstar, const Coord &cstar) {
        const auto &axis = experiment.GetGoniometer();
        if (axis.has_value()) {
            const Coord w = axis->GetAxis().Normalize();
            const Coord S0 = experiment.GetScatteringVector();
            const float wedge_deg = axis->GetWedge_deg();

            double sum_sq = 0.0;
            int count = 0;

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

                const Coord pstar = astar * static_cast<float>(s.h) + bstar * static_cast<float>(s.k) + cstar * static_cast<float>(s.l);

                // Local solve window: +/- 1 wedge (easy/robust first try)
                const auto phi_pred_deg_opt = predict_phi_deg_local(pstar, S0, w, 0.0, wedge_deg);
                if (!phi_pred_deg_opt.has_value())
                    continue;

                float dphi = wrap_deg_pm180(phi_pred_deg_opt.value());
                sum_sq += static_cast<double>(dphi) * static_cast<double>(dphi);
                count++;
            }

            if (count > 0) {
                return static_cast<float>(std::sqrt(sum_sq / static_cast<double>(count)));
            }
        }
        return std::nullopt;
    }

} // namespace

bool AnalyzeIndexing(DataMessage &message,
                     const DiffractionExperiment &experiment,
                     const CrystalLattice &latt) {
    std::vector<uint8_t> indexed_spots(message.spots.size());

    // Check spots
    const Coord a = latt.Vec0();
    const Coord b = latt.Vec1();
    const Coord c = latt.Vec2();

    const Coord astar = latt.Astar();
    const Coord bstar = latt.Bstar();
    const Coord cstar = latt.Cstar();

    const bool index_ice_ring = experiment.GetIndexingSettings().GetIndexIceRings();
    const auto geom = experiment.GetDiffractionGeometry();
    const auto indexing_tolerance = experiment.GetIndexingSettings().GetTolerance();
    const auto indexing_tolerance_sq = indexing_tolerance * indexing_tolerance;
    const auto viable_cell_min_spots = experiment.GetIndexingSettings().GetViableCellMinSpots();

    size_t nspots_ref = 0;
    size_t nspots_indexed = 0;

    // identify indexed spots
    for (int i = 0; i < message.spots.size(); i++) {
        auto recip = message.spots[i].ReciprocalCoord(geom);

        float h_fp = recip * a;
        float k_fp = recip * b;
        float l_fp = recip * c;

        float h_frac = h_fp - std::round(h_fp);
        float k_frac = k_fp - std::round(k_fp);
        float l_frac = l_fp - std::round(l_fp);

        float norm_sq = h_frac * h_frac + k_frac * k_frac + l_frac * l_frac;

        Coord recip_pred = std::round(h_fp) * astar + std::round(k_fp) * bstar + std::round(l_fp) * cstar;

        // See indexing_peak_check() in peaks.c in CrystFEL
        if (norm_sq < indexing_tolerance_sq) {
            if (index_ice_ring || !message.spots[i].ice_ring)
                nspots_indexed++;
            indexed_spots[i] = 1;
            message.spots[i].dist_ewald_sphere = geom.DistFromEwaldSphere(recip_pred);
            message.spots[i].h = std::lround(h_fp);
            message.spots[i].k = std::lround(k_fp);
            message.spots[i].l = std::lround(l_fp);
        }
        if (index_ice_ring || !message.spots[i].ice_ring)
            nspots_ref++;
    }

    if (nspots_indexed >= viable_cell_min_spots && nspots_indexed >= std::lround(min_percentage_spots * nspots_ref)) {
        auto uc = latt.GetUnitCell();
        if (!ok(uc.a) || !ok(uc.b) || !ok(uc.c) || !ok(uc.alpha) || !ok(uc.beta) || !ok(uc.gamma))
            return false;

        message.indexing_result = true;
        assert(indexed_spots.size() == message.spots.size());
        for (int i = 0; i < message.spots.size(); i++)
            message.spots[i].indexed = indexed_spots[i];

        message.profile_radius = FitProfileRadius(message.spots);
        message.spot_count_indexed = nspots_indexed;
        message.indexing_lattice = latt;
        message.indexing_unit_cell = latt.GetUnitCell();
        message.mosaicity_deg = CalcMosaicity(experiment, message.spots, astar, bstar, cstar);

        return true;
    }

    message.indexing_result = false;
    return false;
}