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 "../../common/JFJochMath.h"
#include <cstdint>
#include <vector>
#include <cmath>
#include <algorithm>

#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>(PI) / 180.0f);
    }

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

    // Wrap to [-180, 180] (useful for residuals)
    inline float wrap_deg_pm180(float deg) {
        if (!std::isfinite(deg)) return std::numeric_limits<float>::quiet_NaN();; // or std::nullopt upstream

        deg = std::fmod(deg + 180.0f, 360.0f);
        if (!std::isfinite(deg))
            return std::numeric_limits<float>::quiet_NaN();
        if (deg < 0) deg += 360.0f;
        return deg - 180.0f;
    }

     // XDS convention: zeta = |m2 · e1| where e1 = (S × S0) / |S × S0|
    // This is the Lorentz factor component related to the rotation axis
    inline float calc_zeta(const Coord& S, const Coord& S0, const Coord& m2) {
        Coord S_cross_S0 = S % S0;
        float len = S_cross_S0.Length();
        if (len < 1e-12f) return 0.0f;
        Coord e1 = S_cross_S0 * (1.0f / len);
        return std::fabs(m2 * e1);
    }

    // XDS R(τ; σM/ζ) function - fraction of observed reflection intensity
    // τ = angular difference between reflection and Bragg maximum (radians)
    // delta_phi = oscillation range (radians)
    // sigma_M = mosaicity (radians)
    // zeta = |m2 · e1| Lorentz factor component
    inline float R_fraction(float tau, float delta_phi, float sigma_M, float zeta) {
        if (zeta < 1e-6f || sigma_M < 1e-9f)
            return 0.0f;

        const float sigma_eff = sigma_M / zeta;
        const float sqrt2_sigma = std::sqrt(2.0f) * sigma_eff;

        if (sqrt2_sigma < 1e-12f)
            return 0.0f;

        const float arg_plus = (tau + delta_phi / 2.0f) / sqrt2_sigma;
        const float arg_minus = (tau - delta_phi / 2.0f) / sqrt2_sigma;

        return 0.5f * (std::erf(arg_plus) - std::erf(arg_minus));
    }

    // Log-likelihood for a given sigma_M value
    // Returns sum of log(R) for all reflections
    inline double log_likelihood(const std::vector<float>& tau_values,
                                 const std::vector<float>& zeta_values,
                                 float delta_phi,
                                 float sigma_M) {
        double ll = 0.0;
        for (size_t i = 0; i < tau_values.size(); ++i) {
            float R = R_fraction(tau_values[i], delta_phi, sigma_M, zeta_values[i]);
            if (std::isfinite(R) && R > 1e-30f) {
                ll += std::log(static_cast<double>(R));
            } else {
                ll += -70.0;  // Large penalty for zero probability
            }
        }
        return ll;
    }

    // Golden section search for maximum likelihood sigma_M
    inline float find_sigma_M_mle(const std::vector<float>& tau_values,
                                  const std::vector<float>& zeta_values,
                                  float delta_phi,
                                  float sigma_min_deg = 0.001f,
                                  float sigma_max_deg = 2.0f) {
        const float golden = 0.618033988749895f;

        float a = deg_to_rad(sigma_min_deg);
        float b = deg_to_rad(sigma_max_deg);

        float c = b - golden * (b - a);
        float d = a + golden * (b - a);

        const float tol = 1e-6f;
        int iter = 0;
        while (std::fabs(b - a) > tol && iter++ < 100) {
            double fc = log_likelihood(tau_values, zeta_values, delta_phi, c);
            double fd = log_likelihood(tau_values, zeta_values, delta_phi, d);

            if (fc > fd) {
                b = d;
                d = c;
                c = b - golden * (b - a);
            } else {
                a = c;
                c = d;
                d = a + golden * (b - a);
            }
        }

        return (a + b) / 2.0f;
    }

    // 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 (!std::isfinite(rhs) || 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>(PI);
        auto shift_near = [&](float x) {
            if (!std::isfinite(x))
                return std::numeric_limits<float>::quiet_NaN();

            const float span_center = 0.5f * (phi0 + phi1);

            // Bring x close to the interval center using modulo 2π
            float shifted = x - span_center;
            shifted = std::fmod(shifted, two_pi);
            if (!std::isfinite(shifted))
                return std::numeric_limits<float>::quiet_NaN();

            // fmod can return negative values; normalize to [-π, π]
            if (shifted < -static_cast<float>(PI))
                shifted += two_pi;
            else if (shifted > static_cast<float>(PI))
                shifted -= two_pi;

            return shifted + span_center;
        };

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

    // XDS-style mosaicity calculation using maximum likelihood
    // Following Kabsch (2010) Acta Cryst. D66, 133-144
    std::optional<float> CalcMosaicityXDS(const DiffractionExperiment& experiment,
                                          const std::vector<SpotToSave> &spots,
                                          const Coord &astar, const Coord &bstar, const Coord &cstar) {
        const auto &axis_opt = experiment.GetGoniometer();
        if (!axis_opt.has_value())
            return std::nullopt;

        const GoniometerAxis& axis = *axis_opt;
        const Coord m2 = axis.GetAxis().Normalize();  // XDS notation: m2 is rotation axis
        const Coord S0 = experiment.GetScatteringVector();
        const float delta_phi_rad = deg_to_rad(axis.GetWedge_deg());

        std::vector<float> tau_values;
        std::vector<float> zeta_values;
        tau_values.reserve(spots.size());
        zeta_values.reserve(spots.size());

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

            // Find predicted phi angle. The search window must be wide enough to catch reflections
            // recorded at large rocking offset (|tau| up to ~mosaicity + dphi/2). Using ±wedge alone
            // clips the tau tail at the oscillation width, so the MLE then underestimates the mosaicity
            // ~2x (the tail reflections are exactly the ones that define the mosaic width). A generous
            // window (oscillation + ~0.8deg rocking allowance) lets the tail in; the MLE is insensitive
            // to making it wider still (it weights by the recorded fraction R(tau), which decays).
            const float window_deg = axis.GetWedge_deg() + 0.8f;
            const auto phi_pred_opt = predict_phi_deg_local(pstar, S0, m2, 0.0f, window_deg);
            if (!phi_pred_opt.has_value() || !std::isfinite(phi_pred_opt.value()))
                continue;

            // τ (tau) = angular deviation from Bragg position to center of oscillation range
            float tau_rad = deg_to_rad(wrap_deg_pm180(phi_pred_opt.value()));

            // Calculate diffracted beam direction S = S0 + p (at diffracting condition)
            // For zeta calculation, we need S at the predicted phi angle
            const float phi_pred_rad = deg_to_rad(phi_pred_opt.value());
            const float cos_phi = std::cos(phi_pred_rad);
            const float sin_phi = std::sin(phi_pred_rad);

            // Rotate pstar by predicted phi around m2 axis
            const float p_m2 = pstar * m2;
            const Coord p_parallel = m2 * p_m2;
            const Coord p_perp = pstar - p_parallel;
            const Coord m2_cross_p = m2 % pstar;

            const Coord p_rotated = p_parallel + p_perp * cos_phi + m2_cross_p * sin_phi;
            const Coord S = S0 + p_rotated;

            // Calculate zeta (XDS convention)
            float zeta = calc_zeta(S, S0, m2);

            // Filter out reflections with very small zeta (poorly determined)
            if (!std::isfinite(zeta) || !std::isfinite(tau_rad) || zeta < 0.1f)
                continue;

            tau_values.push_back(tau_rad);
            zeta_values.push_back(zeta);
        }

        if (tau_values.size() < 10)
            return std::nullopt;

        // Find sigma_M by maximizing log-likelihood
        float sigma_M_rad = find_sigma_M_mle(tau_values, zeta_values, delta_phi_rad);

        return rad_to_deg(sigma_M_rad);
    }
} // namespace

bool AnalyzeIndexing(DataMessage &message,
                     const DiffractionExperiment &experiment,
                     const CrystalLattice &latt,
                     const std::vector<CrystalLattice> &extra_lattices) {
    auto start_time = std::chrono::steady_clock::now();

    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;

    // Reciprocal radius squared (|s|^2 = (2 sin(theta)/lambda)^2) per spot, and the largest among the
    // indexed spots - i.e. the highest resolution at which this lattice actually diffracts.
    std::vector<float> spot_q_sq(message.spots.size(), 0.0f);
    float indexed_q_sq_max = 0.0f;

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

        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_q_sq_max = std::max(indexed_q_sq_max, spot_q_sq[i]);
            }
            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);
        }
    }

    // Reference count for the indexed-fraction test = spots within the resolution range that this
    // lattice actually diffracts to (out to the highest-resolution indexed spot). Spots beyond that
    // limit are noise: these weakly-diffracting crystals reach only ~4 A while the detector spans
    // ~1.5 A, so most found spots are unindexable high-resolution background. Counting them in the
    // denominator makes the 20% floor unreachable and rejects every frame. This can only shrink
    // nspots_ref versus counting all spots, so it never rejects a frame that passes today.
    for (int i = 0; i < message.spots.size(); i++) {
        if ((index_ice_ring || !message.spots[i].ice_ring) && spot_q_sq[i] <= indexed_q_sq_max)
            nspots_ref++;
    }

    int64_t indexing_lattice_count = 0;
    bool outcome = false;
    // Minimum fraction of the in-resolution spots a candidate lattice must index to be accepted.
    // Lowering it admits weaker/sparser crystals (more real ones on flooded XFEL frames, but also more
    // spurious lattices that a downstream merge-consistency gate must remove).
    constexpr float min_frac = 0.20f;
    if (nspots_indexed >= viable_cell_min_spots && nspots_indexed >= std::lround(min_frac * 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)) {
            message.indexing_result = true;
            indexing_lattice_count++;

            assert(indexed_spots.size() == message.spots.size());
            for (int i = 0; i < message.spots.size(); i++) {
                message.spots[i].indexed = indexed_spots[i];
                message.spots[i].lattice = indexed_spots[i] ? 0 : -1;
            }
            message.profile_radius = FitProfileRadius(message.spots,
                                                      experiment.GetBandwidthFWHM().value_or(0.0f) / 2.3548f,
                                                      experiment.GetWavelength_A());
            message.spot_count_indexed = nspots_indexed;
            message.indexing_lattice = latt;
            message.indexing_unit_cell = latt.GetUnitCell();
            message.mosaicity_deg = CalcMosaicityXDS(experiment, message.spots, astar, bstar, cstar);

            // Assign remaining (unindexed) spots to extra lattices, in order.
            // Spots already assigned to the main lattice (lattice == 0) are never
            // overwritten. Each extra lattice gets index 1, 2, 3, ...
            const size_t n_extra = std::min<size_t>(extra_lattices.size(), experiment.GetIndexingSettings().GetMaxExtraLattices());
            message.indexing_extra_lattices.clear();
            message.indexing_extra_lattices.reserve(n_extra);

            for (size_t li = 0; li < n_extra; li++) {
                const CrystalLattice &el = extra_lattices[li];

                const Coord ea = el.Vec0();
                const Coord eb = el.Vec1();
                const Coord ec = el.Vec2();
                const Coord east = el.Astar();
                const Coord ebst = el.Bstar();
                const Coord ecst = el.Cstar();

                const int64_t lattice_id = static_cast<int64_t>(li) + 1;

                for (int i = 0; i < message.spots.size(); i++) {
                    // Do not overwrite spots already assigned to a lattice
                    if (message.spots[i].lattice >= 0)
                        continue;

                    auto recip = message.spots[i].ReciprocalCoord(geom);

                    float h_fp = recip * ea;
                    float k_fp = recip * eb;
                    float l_fp = recip * ec;

                    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;

                    if (norm_sq < indexing_tolerance_sq) {
                        Coord recip_pred = std::round(h_fp) * east
                                         + std::round(k_fp) * ebst
                                         + std::round(l_fp) * ecst;
                        message.spots[i].indexed = true;
                        message.spots[i].lattice = lattice_id;
                        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);
                    }
                }

                message.indexing_extra_lattices.push_back(el);
                indexing_lattice_count++;
            }
            outcome = true;
        }
    }

    auto end_time = std::chrono::steady_clock::now();
    message.index_analysis_time_s = std::chrono::duration<float>(end_time - start_time).count();
    message.indexing_lattice_count = indexing_lattice_count;
    message.indexing_result = outcome;
    return outcome;
}