// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute // SPDX-License-Identifier: GPL-3.0-only #include "ReferenceFreeReindex.h" #include #include #include #include #include #include #include "HKLKey.h" #include "ReindexAmbiguity.h" #include "../../common/CrystalLattice.h" #include "../../common/ParallelFor.h" namespace { constexpr size_t MIN_REFLECTIONS = 20; constexpr int RESOLUTION_BINS = 20; constexpr int RESTARTS = 4; constexpr int MAX_SWEEPS = 100; constexpr double DECIDED_SIGMA = 3.0; // Two fits closer than this (a 0.1% difference of one axis, squared) are the same fit. constexpr double METRIC_TIE = 1e-6; double CellDistance(const UnitCell &u, const UnitCell &ref) { const auto rel = [](double v, double r) { return (v - r) / r; }; const auto ang = [](double v, double r) { return (v - r) * std::numbers::pi / 180.0; }; return rel(u.a, ref.a) * rel(u.a, ref.a) + rel(u.b, ref.b) * rel(u.b, ref.b) + rel(u.c, ref.c) * rel(u.c, ref.c) + ang(u.alpha, ref.alpha) * ang(u.alpha, ref.alpha) + ang(u.beta, ref.beta) * ang(u.beta, ref.beta) + ang(u.gamma, ref.gamma) * ang(u.gamma, ref.gamma); } double Median(std::vector v) { if (v.empty()) return NAN; std::nth_element(v.begin(), v.begin() + v.size() / 2, v.end()); return v[v.size() / 2]; } // Any consistent answer relabelled as a whole by one of the indexings is as consistent: take the // one that leaves the most crystals where they are - unless the cells decide an indexing, in which // case the crystals they anchor fix the whole. Two indexings are the same for the merge when they // send a general reflection to the same reflection of the asymmetric unit. Returns, for each // indexing g, the indexing a crystal in g is in after that relabel. choice: each crystal's indexing. std::vector GlobalRelabel(const std::vector &indexings, const std::vector &choice, const HKLKeyGenerator &key_gen, bool anchored) { const int K = static_cast(indexings.size()); const std::array general{2, 5, 11}; const auto same = [&](const gemmi::Op::Miller &a, const gemmi::Op::Miller &b) { return key_gen(a[0], a[1], a[2]) == key_gen(b[0], b[1], b[2]); }; const auto composed = [&](int global, int g) { // indexing g, then the overall relabel return indexings[global].apply_to_hkl(indexings[g].apply_to_hkl(general)); }; int global = 0; size_t global_kept = 0; for (int h = 0; h < (anchored ? 1 : K); ++h) { size_t kept = 0; for (const int c : choice) kept += same(composed(h, c), general); if (kept > global_kept) { global_kept = kept; global = h; } } std::vector becomes(K); for (int c = 0; c < K; ++c) { becomes[c] = c; for (int g = 0; g < K; ++g) if (same(composed(global, c), indexings[g].apply_to_hkl(general))) { becomes[c] = g; break; } } return becomes; } } MetricAnchors AnchorOnCells(const std::vector &outcomes, const UnitCell &cell, const std::vector &ops) { std::vector indexings{gemmi::Op::identity()}; indexings.insert(indexings.end(), ops.begin(), ops.end()); const size_t K = indexings.size(); // How much better an image's own refined cell fits the run's cell as indexed than after indexing g. const auto cell_preference = [&](const IntegrationOutcome &o, size_t g) { return CellDistance(ReindexLattice(o.latt, indexings[g]).GetUnitCell(), cell) - CellDistance(o.latt.GetUnitCell(), cell); }; // Which indexings the cells decide: those the indexed images' cells prefer the identity over (a tie // counting half) for at least METRIC_DECIDED of them. MetricAnchors out; std::vector decided(K, false); for (size_t g = 1; g < K; ++g) { double prefer = 0, n = 0; for (const auto &o : outcomes) { if (o.reflections.empty()) continue; const double p = cell_preference(o, g); prefer += std::fabs(p) < METRIC_TIE ? 0.5 : p > 0 ? 1.0 : 0.0; n += 1; } out.metric_preference.push_back(n > 0 ? prefer / n : 0.5); decided[g] = out.metric_preference.back() >= METRIC_DECIDED; } // Where the cells decide an indexing, an image whose own cell prefers the identity stays as indexed. out.allowed.assign(outcomes.size() * K, 1); for (size_t o = 0; o < outcomes.size(); ++o) for (size_t g = 1; g < K; ++g) if (decided[g] && !outcomes[o].reflections.empty() && cell_preference(outcomes[o], g) >= METRIC_TIE) out.allowed[o * K + g] = 0; return out; } ReferenceAnchorResult AnchorReferenceReindex(std::vector &outcomes, const UnitCell &cell, const std::vector &ops) { std::vector indexings{gemmi::Op::identity()}; indexings.insert(indexings.end(), ops.begin(), ops.end()); const size_t K = indexings.size(); // The indexing best correlated with the reference among those open: the one as indexed unless another // beats it; a score that is not finite never wins (the rule of the per-image resolver). const auto best = [&](const std::vector &cc, const uint8_t *open) { size_t b = 0; for (size_t g = 1; g < K; ++g) if (open[g] && std::isfinite(cc[g]) && (!std::isfinite(cc[b]) || cc[g] > cc[b])) b = g; return b; }; const std::vector all_open(K, 1); // Each image back in the indexing it was indexed in. std::vector by_reference(outcomes.size(), 0); for (size_t o = 0; o < outcomes.size(); ++o) { auto &out = outcomes[o]; if (out.reference_cc.size() != K) continue; by_reference[o] = best(out.reference_cc, all_open.data()); if (by_reference[o] != 0) ReindexImage(out.reflections, out.latt, indexings[by_reference[o]].inverse()); } const MetricAnchors anchors = AnchorOnCells(outcomes, cell, ops); ReferenceAnchorResult result; result.metric_preference = anchors.metric_preference; for (size_t o = 0; o < outcomes.size(); ++o) { auto &out = outcomes[o]; if (out.reference_cc.size() != K) continue; const size_t g = best(out.reference_cc, &anchors.allowed[o * K]); result.overruled += g != by_reference[o]; if (g != 0) ReindexImage(out.reflections, out.latt, indexings[g]); out.reference_cc.clear(); // the scores were for the indexing as indexed, which is no longer known } return result; } ReferenceFreeReindexResult ResolveReindexReferenceFree(std::vector &outcomes, const DiffractionExperiment &x, const UnitCell &cell, const std::vector &all_ops, size_t nthreads, bool apply) { ReferenceFreeReindexResult out; out.choice.assign(outcomes.size(), -1); if (all_ops.empty()) return out; const auto &s = x.GetScalingSettings(); const gemmi::SpaceGroup &sg = x.GetSpaceGroupOrP1(); std::vector indexings{gemmi::Op::identity()}; indexings.insert(indexings.end(), all_ops.begin(), all_ops.end()); const int K = static_cast(indexings.size()); const MetricAnchors anchors = AnchorOnCells(outcomes, cell, all_ops); out.metric_preference = anchors.metric_preference; const auto accept = [&](const Reflection &r) { return !r.on_ice_ring && AcceptReflection(r, s.GetHighResolutionLimit_A(), s.GetLowResolutionLimit_A()) && r.partiality >= s.GetMinPartiality() && r.partiality > 0.0f && std::isfinite(r.I) && std::isfinite(r.sigma) && r.sigma > 0.0f; }; const auto corrected = [](const Reflection &r) { return static_cast(r.I) * (r.prescaling_corr * r.qe_corr * r.flight_corr) / r.partiality; }; // The crystals are compared on the square root of the intensity, its sign kept. Intensities follow // the Wilson distribution, whose long tail lets a few very strong reflections dominate every product // of two crystals: then the strongest crystals settle among themselves, and a group of them can hold // one another in the wrong indexing however much the weaker majority disagrees. Amplitudes have no // such tail. const auto amplitude = [&](const Reflection &r) { const double i = corrected(r); return std::copysign(std::sqrt(std::fabs(i)), i); }; const auto amplitude_sigma = [&](const Reflection &r) { const double sigma = r.sigma * (r.prescaling_corr * r.qe_corr * r.flight_corr) / r.partiality; return sigma / (2.0 * std::sqrt(std::fabs(corrected(r)) + sigma)); }; // What each image has to offer: how many reflections, how fine, and how far its indices reach in // any of the indexings (which sizes the index table below). const int n_out = static_cast(outcomes.size()); std::vector count(n_out, 0); std::vector> reach(n_out, {0, 0, 0}); std::vector s3_max(n_out, 0.0); ParallelFor(n_out, nthreads, [&](int o) { for (const auto &r : outcomes[o].reflections) { if (!accept(r)) continue; ++count[o]; s3_max[o] = std::max(s3_max[o], 1.0 / (static_cast(r.d) * r.d * r.d)); for (const auto &op : indexings) { const auto h = op.apply_to_hkl({{r.h, r.k, r.l}}); for (int i = 0; i < 3; ++i) reach[o][i] = std::max(reach[o][i], std::abs(h[i])); } } }); // The crystals: images with enough reflections, in image order. std::vector crystal_outcome; std::vector begin{0}; std::array bound{0, 0, 0}; double s3_top = 0.0; for (int o = 0; o < n_out; ++o) { if (count[o] < MIN_REFLECTIONS) continue; crystal_outcome.push_back(o); begin.push_back(begin.back() + count[o]); for (int i = 0; i < 3; ++i) bound[i] = std::max(bound[i], reach[o][i]); s3_top = std::max(s3_top, s3_max[o]); } const int N = static_cast(crystal_outcome.size()); if (N < 2) return out; // The indexings each crystal may take (AnchorOnCells): a crystal its own cell anchors is one of those // the others are decided against. std::vector allowed(static_cast(N) * K); std::vector free(N, false); for (int c = 0; c < N; ++c) for (int g = 0; g < K; ++g) { allowed[c * K + g] = anchors.allowed[crystal_outcome[c] * K + g]; free[c] = free[c] || (g > 0 && allowed[c * K + g]); } out.crystals = std::count(free.begin(), free.end(), true); if (out.crystals == 0) return out; const size_t nnz = begin.back(); // The mean and spread of the amplitudes in shells of equal reciprocal volume over the whole run, so // that a crystal's amplitudes are compared after the fall-off with resolution, which every indexing // shares, is out. Past the resolution the run reaches, a shell holds only noise and its mean may be // zero or negative, so a shell is put on its mean and spread rather than divided by its mean. const auto bin_of = [&](const Reflection &r) { const double s3 = 1.0 / (static_cast(r.d) * r.d * r.d); return std::min(RESOLUTION_BINS - 1, static_cast(RESOLUTION_BINS * s3 / s3_top)); }; std::vector> bin_part(N); ParallelFor(N, nthreads, [&](int c) { auto &b = bin_part[c]; b.fill(0.0); for (const auto &r : outcomes[crystal_outcome[c]].reflections) if (accept(r)) { const double a = amplitude(r); b[bin_of(r)] += 1.0; b[RESOLUTION_BINS + bin_of(r)] += a; b[2 * RESOLUTION_BINS + bin_of(r)] += a * a; } }); std::array bin_mean{}, bin_sd{}; { std::array total{}; for (const auto &b : bin_part) // in crystal order: the same sum at any thread count for (int i = 0; i < 3 * RESOLUTION_BINS; ++i) total[i] += b[i]; for (int i = 0; i < RESOLUTION_BINS; ++i) { const double n = total[i]; bin_mean[i] = n > 0 ? total[RESOLUTION_BINS + i] / n : 0.0; const double var = n > 0 ? total[2 * RESOLUTION_BINS + i] / n - bin_mean[i] * bin_mean[i] : 0.0; bin_sd[i] = var > 0 ? std::sqrt(var) : 1.0; } } // Every index any indexing can produce, mapped to a dense number of its reflection in the // asymmetric unit (Friedel mates together: the anomalous difference would only add noise here). const HKLKeyGenerator key_gen(/*merge_friedel=*/true, sg); const int dh = 2 * bound[0] + 1, dk = 2 * bound[1] + 1, dl = 2 * bound[2] + 1; const auto entry = [&](int h, int k, int l) { return (static_cast(h + bound[0]) * dk + (k + bound[1])) * dl + (l + bound[2]); }; std::vector table_key(static_cast(dh) * dk * dl); ParallelFor(dh, nthreads, [&](int ih) { for (int k = -bound[1]; k <= bound[1]; ++k) for (int l = -bound[2]; l <= bound[2]; ++l) table_key[entry(ih - bound[0], k, l)] = key_gen(ih - bound[0], k, l).pack(); }); std::vector table(table_key.size()); int32_t n_keys = 0; { std::unordered_map dense; for (size_t i = 0; i < table_key.size(); ++i) { const auto [it, inserted] = dense.try_emplace(table_key[i], n_keys); n_keys += inserted; table[i] = it->second; } } table_key = {}; // Each crystal's amplitudes, put on the shell's mean and spread and standardised within the crystal // (the unknown scale and offset of an image drop out), and where each lands in every indexing. Each is // weighted by how much of it is signal: a standardised amplitude has unit variance, of which its // measurement error takes (sigma / sd)^2, so a weak or barely recorded reflection counts for little. std::vector z(nnz); std::vector> key(K, std::vector(nnz)); ParallelFor(N, nthreads, [&](int c) { std::vector err; size_t j = begin[c]; double sum = 0, sum2 = 0; for (const auto &r : outcomes[crystal_outcome[c]].reflections) { if (!accept(r)) continue; const double v = (amplitude(r) - bin_mean[bin_of(r)]) / bin_sd[bin_of(r)]; z[j] = static_cast(v); err.push_back(amplitude_sigma(r) / bin_sd[bin_of(r)]); sum += v; sum2 += v * v; for (int g = 0; g < K; ++g) { const auto h = indexings[g].apply_to_hkl({{r.h, r.k, r.l}}); key[g][j] = table[entry(h[0], h[1], h[2])]; } ++j; } const double n = static_cast(begin[c + 1] - begin[c]); const double mean = sum / n; const double sd = std::sqrt(std::max(0.0, sum2 / n - mean * mean)); for (size_t i = begin[c]; i < begin[c + 1]; ++i) { const double e = err[i - begin[c]] / sd; z[i] = sd > 0 ? static_cast((z[i] - mean) / sd / (1.0 + e * e)) : 0.0f; } }); table = {}; // Following Brehm & Diederichs (2014) Acta Cryst. D70, 101-109: each crystal in turn takes the // indexing in which it agrees best with all the others, until none changes. Agreement is the sum // over all pairs of crystals of the products of their standardised amplitudes on common // reflections, which is the dot product of a crystal with the sum of all the others - so a pass over // the crystals costs one pass over the reflections, not one per pair of crystals. The sum of all // pairs only ever grows, so it ends; a few starts are made and the one that ends highest is kept. std::vector M(n_keys); const auto add = [&](int c, int g, double sign) { for (size_t i = begin[c]; i < begin[c + 1]; ++i) M[key[g][i]] += sign * z[i]; }; const auto dot = [&](int c, int g) { double d = 0; for (size_t i = begin[c]; i < begin[c + 1]; ++i) d += z[i] * M[key[g][i]]; return d; }; std::vector best_choice; double best_objective = -1.0; for (int restart = 0; restart < RESTARTS; ++restart) { std::vector choice(N, 0); std::mt19937 rng(restart); if (restart > 0) for (int c = 0; c < N; ++c) { const int g = static_cast(rng() % K); choice[c] = allowed[c * K + g] ? g : 0; } std::fill(M.begin(), M.end(), 0.0); for (int c = 0; c < N; ++c) add(c, choice[c], 1.0); int sweeps = 0; // Run until a pass moves no more than one crystal in a thousand: what still moves then are crystals // too weak to be decided either way. for (int changed = N; changed > N / 1000 && sweeps < MAX_SWEEPS; ++sweeps) { changed = 0; for (int c = 0; c < N; ++c) { add(c, choice[c], -1.0); int best = choice[c]; double best_dot = dot(c, best); for (int g = 0; g < K; ++g) { if (g == choice[c] || !allowed[c * K + g]) continue; const double d = dot(c, g); if (d > best_dot) { best_dot = d; best = g; } } changed += best != choice[c]; choice[c] = best; add(c, best, 1.0); } } double objective = 0; for (double m : M) objective += m * m; if (objective > best_objective) { best_objective = objective; best_choice = choice; out.sweeps = sweeps; } } // How clearly each crystal is decided: its CC with the mean of all the other crystals, in the chosen // indexing and in the best other one, and the difference of the two in sigma (Fisher's z; each CC // on n reflections has an error of 1/sqrt(n - 3), and the two are nearly independent). std::fill(M.begin(), M.end(), 0.0); std::vector C(n_keys, 0.0); for (int c = 0; c < N; ++c) for (size_t i = begin[c]; i < begin[c + 1]; ++i) { M[key[best_choice[c]][i]] += z[i]; C[key[best_choice[c]][i]] += 1.0; } std::vector sigma(N, NAN), cc_chosen(N, NAN), cc_other(N, NAN); ParallelFor(N, nthreads, [&](int c) { std::vector cc(K, NAN); std::vector n_used(K, 0.0); for (int g = 0; g < K; ++g) { double sx = 0, sy = 0, sxx = 0, syy = 0, sxy = 0, n = 0; for (size_t i = begin[c]; i < begin[c + 1]; ++i) { const int32_t k = key[g][i]; const double own = g == best_choice[c] ? 1.0 : 0.0; const double cnt = C[k] - own; if (cnt <= 0) continue; const double xv = z[i], yv = (M[k] - own * z[i]) / cnt; sx += xv; sy += yv; sxx += xv * xv; syy += yv * yv; sxy += xv * yv; n += 1; } const double vx = sxx - sx * sx / n, vy = syy - sy * sy / n; if (n > 3 && vx > 0 && vy > 0) { cc[g] = std::clamp((sxy - sx * sy / n) / std::sqrt(vx * vy), -0.999, 0.999); n_used[g] = n; } } const int g0 = best_choice[c]; int other = -1; for (int g = 0; g < K; ++g) if (g != g0 && allowed[c * K + g] && std::isfinite(cc[g]) && (other < 0 || cc[g] > cc[other])) other = g; if (!std::isfinite(cc[g0]) || other < 0) return; cc_chosen[c] = cc[g0]; cc_other[c] = cc[other]; sigma[c] = (std::atanh(cc[g0]) - std::atanh(cc[other])) * std::sqrt((std::min(n_used[g0], n_used[other]) - 3.0) / 2.0); }); std::vector sig, ccc, cco; for (int c = 0; c < N; ++c) if (free[c] && std::isfinite(sigma[c])) { sig.push_back(sigma[c]); ccc.push_back(cc_chosen[c]); cco.push_back(cc_other[c]); } out.median_sigma = Median(sig); out.cc_chosen = Median(ccc); out.cc_other = Median(cco); if (!sig.empty()) out.decided_fraction = static_cast(std::count_if(sig.begin(), sig.end(), [](double v) { return v > DECIDED_SIGMA; })) / static_cast(sig.size()); const bool anchored = std::any_of(out.metric_preference.begin(), out.metric_preference.end(), [](double p) { return p >= METRIC_DECIDED; }); const std::vector becomes = GlobalRelabel(indexings, best_choice, key_gen, anchored); for (int c = 0; c < N; ++c) { const int final_g = becomes[best_choice[c]]; out.choice[crystal_outcome[c]] = final_g; if (final_g == 0) continue; ++out.reindexed; if (apply) { auto &o = outcomes[crystal_outcome[c]]; ReindexImage(o.reflections, o.latt, indexings[final_g]); } } return out; } int ReindexOnReference(IntegrationOutcome &outcome, const ReindexAmbiguityResolver &resolver, const UnitCell &cell, const std::vector &ops, const std::vector &metric_preference) { if (ops.empty() || outcome.reflections.empty()) return 0; const std::vector cc = resolver.Scores(outcome.reflections); const auto open = [&](size_t g) { if (g == 0 || metric_preference.size() < g || metric_preference[g - 1] < METRIC_DECIDED) return true; const double p = CellDistance(ReindexLattice(outcome.latt, ops[g - 1]).GetUnitCell(), cell) - CellDistance(outcome.latt.GetUnitCell(), cell); return p < METRIC_TIE; }; // The rule of AnchorReferenceReindex: as indexed unless an open indexing beats it. size_t best = 0; for (size_t g = 1; g <= ops.size(); ++g) if (open(g) && std::isfinite(cc[g]) && (!std::isfinite(cc[best]) || cc[g] > cc[best])) best = g; if (best != 0) ReindexImage(outcome.reflections, outcome.latt, ops[best - 1]); return static_cast(best); } size_t RelabelRunIndexing(std::vector &outcomes, std::vector &choice, const DiffractionExperiment &x, const std::vector &ops, const std::vector &metric_preference) { if (ops.empty()) return 0; std::vector indexings{gemmi::Op::identity()}; indexings.insert(indexings.end(), ops.begin(), ops.end()); const HKLKeyGenerator key_gen(/*merge_friedel=*/true, x.GetSpaceGroupOrP1()); const bool anchored = std::any_of(metric_preference.begin(), metric_preference.end(), [](double p) { return p >= METRIC_DECIDED; }); std::vector taking_part; for (size_t o = 0; o < outcomes.size(); ++o) if (choice[o] >= 0 && !outcomes[o].reflections.empty()) taking_part.push_back(choice[o]); const std::vector becomes = GlobalRelabel(indexings, taking_part, key_gen, anchored); size_t moved = 0; for (size_t o = 0; o < outcomes.size(); ++o) { if (choice[o] < 0 || becomes[choice[o]] == choice[o]) continue; auto &out = outcomes[o]; ReindexImage(out.reflections, out.latt, indexings[choice[o]].inverse()); choice[o] = becomes[choice[o]]; ReindexImage(out.reflections, out.latt, indexings[choice[o]]); ++moved; } return moved; }