scale_merge: harden R-free flag assignment (Bijvoet-safe, deterministic, stratified)

The R-free flags were drawn with a per-reflection std::bernoulli_distribution off a
shared std::mt19937(12345). Two problems: in anomalous mode I(+) and I(-) are separate
merged rows, so a Bijvoet pair could be split across the work and free sets (biasing
R-free); and the draw order depended on merge/call order, so it was not reproducible or
parallel-safe. The resolution-shell grouping was cosmetic — it still drew an independent
Bernoulli per row, so it neither tied mates together nor guaranteed a per-shell quota.

New AssignRfreeFlags() (own file RfreeFlags.{h,cpp}) replaces both call sites (stills
MergeOnTheFly, rotation RotationScaleMerge). It:
  - keys each reflection by its Friedel-merged (Laue) ASU index, so all symmetry- and
    Friedel-equivalent reflections (including a Bijvoet pair) share one flag and are never
    split;
  - selects the free set by a deterministic splitmix64 hash of that key (same idiom as the
    CC1/2 half-set split), so it is reproducible run-to-run and order-independent;
  - stratifies by resolution: exactly round(fraction * n) of the distinct reflections in
    each of 20 shells is flagged free.

Verified on an anomalous lyso merge (29536 refl, 5% free): 13548 Bijvoet pairs, 0 split;
free fraction 4.6-5.4% in every resolution shell.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
2026-07-12 21:27:44 +02:00
co-authored by Claude Opus 4.8
parent 5addb4483c
commit 4c95a289b6
5 changed files with 103 additions and 57 deletions
@@ -13,6 +13,8 @@ ADD_LIBRARY(JFJochScaleMerge
ResolutionCutoff.h
HKLKey.cpp
HKLKey.h
RfreeFlags.cpp
RfreeFlags.h
ScalingResult.h
ScalingResult.cpp)
TARGET_LINK_LIBRARIES(JFJochScaleMerge Ceres::ceres Eigen3::Eigen JFJochCommon)
+2 -36
View File
@@ -17,6 +17,7 @@
#include "../../common/ResolutionShells.h"
#include "../../common/Definitions.h"
#include "HKLKey.h"
#include "RfreeFlags.h"
namespace {
// Deterministic CC1/2 half-set assignment: a splitmix64 bit-mix of the image's stable index.
@@ -438,9 +439,6 @@ bool MergeOnTheFly::Mask(const IntegrationOutcome &outcome, bool cc_mask) {
std::vector<MergedReflection> MergeOnTheFly::ExportReflections() {
std::unique_lock ul(merged_mutex);
float d_min = std::numeric_limits<float>::max();
float d_max = 0.0f;
std::vector<MergedReflection> out;
out.reserve(accumulator.size());
for (const auto &accum: accumulator | std::views::values) {
@@ -470,42 +468,10 @@ std::vector<MergedReflection> MergeOnTheFly::ExportReflections() {
if (!std::isfinite(accum.d) || accum.d <= 0.0f)
continue;
d_min = std::min(d_min, accum.d);
d_max = std::max(d_max, accum.d);
out.emplace_back(mr);
}
const double rfree_fraction = scaling_settings.GetRfreeFraction();
if (rfree_fraction > 0.0 && !out.empty()) {
if (d_min < d_max && d_min > 0.0f) {
constexpr int n_shells = 20;
const float d_min_pad = d_min * 0.999f;
const float d_max_pad = d_max * 1.001f;
ResolutionShells shells(d_min_pad, d_max_pad, n_shells);
std::vector<std::vector<size_t>> shell_groups(n_shells);
for (size_t i = 0; i < out.size(); ++i) {
const auto shell = shells.GetShell(out[i].d);
if (!shell.has_value())
continue;
const int s = *shell;
if (s >= 0 && s < n_shells)
shell_groups[s].push_back(i);
}
std::mt19937 rfree_rng(12345u);
std::bernoulli_distribution rfree_dist(rfree_fraction);
for (const auto &group: shell_groups) {
for (const size_t idx: group)
out[idx].rfree_flag = rfree_dist(rfree_rng);
}
}
}
AssignRfreeFlags(out, space_group_number, scaling_settings.GetRfreeFraction());
return out;
}
+79
View File
@@ -0,0 +1,79 @@
// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only
#include "RfreeFlags.h"
#include <algorithm>
#include <cmath>
#include <limits>
#include <unordered_set>
#include "HKLKey.h"
#include "../../common/ResolutionShells.h"
namespace {
// splitmix64 bit-mix of a key -> uniform double in [0, 1). Same key -> same value, so all
// mates of a reflection (which share the Laue-ASU key) get the same draw. Same idiom as the
// CC1/2 half-set split (HalfForImage in Merge.cpp).
double UniformFromKey(uint64_t key) {
uint64_t z = key + 0x9e3779b97f4a7c15ULL;
z = (z ^ (z >> 30)) * 0xbf58476d1ce4e5b9ULL;
z = (z ^ (z >> 27)) * 0x94d049bb133111ebULL;
z = z ^ (z >> 31);
return static_cast<double>(z >> 11) * (1.0 / 9007199254740992.0);
}
}
void AssignRfreeFlags(std::vector<MergedReflection> &merged, int32_t space_group_number,
double rfree_fraction) {
for (auto &r : merged)
r.rfree_flag = false;
if (rfree_fraction <= 0.0 || merged.empty())
return;
// Friedel-merged (Laue) ASU key: I(+)/I(-) and every symmetry mate collapse to one key, so a
// reflection's free/work status is a pure function of this key -> mates never split.
const HKLKeyGenerator laue_key(/*merge_friedel=*/true, space_group_number);
float d_min = std::numeric_limits<float>::max(), d_max = 0.0f;
for (const auto &r : merged)
if (std::isfinite(r.d) && r.d > 0.0f) {
d_min = std::min(d_min, r.d);
d_max = std::max(d_max, r.d);
}
// No usable resolution: fall back to a global deterministic hash (still mate-consistent).
if (!(d_min < d_max && d_min > 0.0f)) {
for (auto &r : merged)
r.rfree_flag = UniformFromKey(laue_key(r).pack()) < rfree_fraction;
return;
}
// Stratify by resolution: within each shell take exactly round(fraction * n) of the distinct
// Laue reflections as the free set, chosen by hash rank. This guarantees ~fraction of the data
// is free in every shell, deterministically and without splitting any Bijvoet/symmetry set.
constexpr int n_shells = 20;
ResolutionShells shells(d_min * 0.999f, d_max * 1.001f, n_shells);
std::vector<std::vector<std::pair<double, uint64_t>>> shell_keys(n_shells); // (hash, key)
std::unordered_set<uint64_t> seen;
seen.reserve(merged.size());
for (const auto &r : merged) {
const uint64_t key = laue_key(r).pack();
if (!seen.insert(key).second) // count each Laue set once (Bijvoet mates share a key)
continue;
const auto shell = shells.GetShell(r.d);
if (shell)
shell_keys[*shell].push_back({UniformFromKey(key), key});
}
std::unordered_set<uint64_t> free_keys;
free_keys.reserve(static_cast<size_t>(rfree_fraction * seen.size()) + 1);
for (auto &keys : shell_keys) {
std::sort(keys.begin(), keys.end());
const auto n_free = static_cast<size_t>(std::llround(rfree_fraction * keys.size()));
for (size_t i = 0; i < n_free; ++i)
free_keys.insert(keys[i].second);
}
for (auto &r : merged)
r.rfree_flag = free_keys.count(laue_key(r).pack()) > 0;
}
+18
View File
@@ -0,0 +1,18 @@
// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only
#pragma once
#include <cstdint>
#include <vector>
#include "../../common/Reflection.h"
// Assign R-free (test-set) flags to merged reflections. Two guarantees:
// - All symmetry- and Friedel-equivalent reflections share one flag (the flag is a pure function
// of the Friedel-merged Laue-ASU key), so a Bijvoet pair I(+)/I(-) is never split across sets,
// and the assignment is reproducible run-to-run and independent of merge/call order.
// - The free set is stratified across resolution: ~rfree_fraction of the distinct reflections in
// every resolution shell is flagged free.
void AssignRfreeFlags(std::vector<MergedReflection> &merged, int32_t space_group_number,
double rfree_fraction);
@@ -16,6 +16,7 @@
#include <gemmi/symmetry.hpp>
#include "HKLKey.h"
#include "RfreeFlags.h"
#include "ResolutionCutoff.h"
#include "../../common/CorrelationCoefficient.h"
#include "../../common/CrystalLattice.h"
@@ -1297,28 +1298,8 @@ RotationScaleMerge::Result RotationScaleMerge::MergeAndStats(int n_groups, bool
// the space-group search still sees the full range.
const std::optional<double> effective_d_min = ApplyResolutionCutoff(
result.merged, d_min_limit, resolution_cutoff_method, resolution_cc_target, for_search, logger);
if (effective_d_min) {
// Recompute the merged resolution span for the R-free binning below over the trimmed set.
d_min = std::numeric_limits<float>::max(); d_max = 0.0f;
for (const auto &m : result.merged) {
if (!std::isfinite(m.d) || m.d <= 0.0f) continue;
d_min = std::min(d_min, m.d); d_max = std::max(d_max, m.d);
}
}
if (rfree_fraction > 0.0 && !result.merged.empty() && d_min < d_max && d_min > 0.0f) {
constexpr int n_shells = 20;
ResolutionShells shells(d_min * 0.999f, d_max * 1.001f, n_shells);
std::vector<std::vector<size_t>> shell_groups(n_shells);
for (size_t i = 0; i < result.merged.size(); ++i) {
const auto shell = shells.GetShell(result.merged[i].d);
if (shell && *shell >= 0 && *shell < n_shells) shell_groups[*shell].push_back(i);
}
std::mt19937 rng(12345u);
std::bernoulli_distribution dist(rfree_fraction);
for (const auto &grp : shell_groups)
for (const size_t i : grp) result.merged[i].rfree_flag = dist(rng);
}
AssignRfreeFlags(result.merged, x.GetSpaceGroupNumber().value_or(1), rfree_fraction);
if (reject_count > 0)
logger.Info("Merge outlier rejection: dropped {} observations", reject_count);