leonarski_f
f9eee954de
Build Packages / Unit tests (push) Failing after 3m9s
Build Packages / build:rpm (rocky8) (push) Failing after 4m37s
Build Packages / build:rpm (rocky9_nocuda) (push) Failing after 5m3s
Build Packages / build:rpm (rocky8_sls9) (push) Failing after 4m57s
Build Packages / build:rpm (ubuntu2404_nocuda) (push) Failing after 5m22s
Build Packages / build:rpm (rocky9_sls9) (push) Failing after 5m57s
Build Packages / build:rpm (rocky8_nocuda) (push) Failing after 6m17s
Build Packages / build:rpm (ubuntu2204_nocuda) (push) Failing after 6m56s
Build Packages / Generate python client (push) Successful in 30s
Build Packages / build:rpm (rocky9) (push) Failing after 4m39s
Build Packages / Create release (push) Skipped
Build Packages / Build documentation (push) Successful in 36s
Build Packages / DIALS test (push) Failing after 3m34s
Build Packages / build:rpm (ubuntu2404) (push) Failing after 4m4s
Build Packages / XDS test (JFJoch plugin) (push) Failing after 3m7s
Build Packages / XDS test (durin plugin) (push) Failing after 3m47s
Build Packages / build:rpm (ubuntu2204) (push) Failing after 4m34s
Build Packages / XDS test (neggia plugin) (push) Failing after 3m49s
430 lines
16 KiB
C++
430 lines
16 KiB
C++
// SPDX-FileCopyrightText: 2025 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
|
||
// SPDX-License-Identifier: GPL-3.0-only
|
||
|
||
#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>(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) {
|
||
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>(M_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>(M_PI))
|
||
shifted += two_pi;
|
||
else if (shifted > static_cast<float>(M_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
|
||
const auto phi_pred_opt = predict_phi_deg_local(pstar, S0, m2, 0.0f, axis.GetWedge_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;
|
||
|
||
// 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++;
|
||
}
|
||
|
||
bool outcome = false;
|
||
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)) {
|
||
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.spots[i].lattice = indexed_spots[i] ? 0 : -1;
|
||
}
|
||
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 = 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);
|
||
}
|
||
|
||
message.indexing_lattice_count = static_cast<int64_t>(1 + message.indexing_extra_lattices.size());
|
||
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_result = outcome;
|
||
return outcome;
|
||
}
|