AnalyzeIndexing: Log-likelihood approach for mosaicity calculation (like XDS)

image_analysis/indexing/AnalyzeIndexing.cpp

@@ -3,6 +3,9 @@
 
 #include <cstdint>
 #include <vector>
+#include <cmath>
+#include <algorithm>
+
 #include "AnalyzeIndexing.h"
 
 #include "FitProfileRadius.h"
@@ -31,6 +34,89 @@ namespace {
         return deg;
     }
 
+     // 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 (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;
+
+        while (std::fabs(b - a) > tol) {
+            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,
@@ -117,6 +203,76 @@ namespace {
         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
+            const auto phi_pred_opt = predict_phi_deg_local(pstar, S0, m2, 0.0f, axis.GetWedge_deg());
+            if (!phi_pred_opt.has_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 (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);
+    }
+
     std::optional<float> CalcMosaicity(const DiffractionExperiment& experiment,
                                        const std::vector<SpotToSave> &spots,
                                        const Coord &astar, const Coord &bstar, const Coord &cstar) {
@@ -221,7 +377,7 @@ bool AnalyzeIndexing(DataMessage &message,
         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);
+        message.mosaicity_deg = CalcMosaicityXDS(experiment, message.spots, astar, bstar, cstar);
 
         return true;
     }