leonarski_fandClaude Opus 5
4895dc1018
The flag was accepted on rotation data and did nothing - it is read only by the stills merge (Merge.cpp), and the CLI warned about that rather than fixing it. Meanwhile RotationScaleMerge already COMPUTES a per-frame correlation against the merged reference and writes it to the per-image table; nothing acted on it. Wire the two together. A rejected frame has its partials' corr set to 0, which is how a frame already leaves the pipeline - every consumer requires corr > 0, so the combine, the merge and the error model all drop it together. The GPU path reuses the SmoothCorr kernel with a ratio of 0, so one implementation covers both. Off by default (0), and verified bit-identical to the previous binary when off. What it catches, on the two rotation datasets that have a population to catch: a two-lattice crystal - two lattices in two physical AREAS of the sample, so the sweep passes from one to the other and whole blocks of frames measure a different crystal from the one being merged (frames 500-700 index perfectly well at a per-frame CC of 0.22 against 0.47-0.56 either side, in 11 contiguous runs). R_meas 28.6 -> 24.6%, CC1/2 93.6 -> 95.1, high-shell CC 23.4 -> 38.3. a second dataset with 9.5% of frames below CC 0.30: R_meas 24.3 -> 23.4%, CC1/2 92.6 -> 93.4. The criterion is "this frame disagrees with the merged reference", NOT "this frame is off-crystal". It happens to catch both, because a frame that measures nothing and a frame that measures a DIFFERENT crystal fail the same test, and it does not need to know which. For the two-area case that is a workaround, not a treatment: it recovers one crystal by discarding the other, where processing the two as separate sweeps would keep both. The frame-block structure is clean enough that such a split could be detected automatically. WHY THERE IS NO DEFAULT. The per-frame CC is not comparable between datasets - it is as much a measure of data quality as of frame validity. Measured medians across the battery run from 0.30 to 0.81, so one absolute bound removes 13 frames from one dataset and 584 of 1800 from another: battery at --min-image-cc 30, 33 crystals: no point group changed (30/33), four crystals clearly better (one +5.4 CC1/2 points, the two-lattice case above, and ISa gains of 1.3-4.6 on three others) - and one healthy crystal lost a third of its frames and with them its high-resolution shell (CC1/2_hi 26.2 -> 2.0). This is the same trap as an absolute bound on any per-operator or per-frame agreement statistic, and the same one the per-frame scale guard avoids by measuring against the run's own median. A principled version would cut on the SHAPE of the per-frame CC distribution - a dataset with a bad subpopulation is bimodal, a uniformly weak one is not - rather than on an absolute value. Until that exists this stays opt-in, and the per-image CC it keys on is already in the _image.dat table for anyone choosing a value. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
2134 lines
114 KiB
C++
2134 lines
114 KiB
C++
// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
|
|
// SPDX-License-Identifier: GPL-3.0-only
|
|
|
|
#include "RotationScaleMerge.h"
|
|
|
|
#include <algorithm>
|
|
#include <atomic>
|
|
#include <cmath>
|
|
#include <cstdint>
|
|
#include <fstream>
|
|
#include <future>
|
|
#include <limits>
|
|
#include <random>
|
|
#include <unordered_map>
|
|
|
|
#include <gemmi/reciproc.hpp>
|
|
#include <gemmi/symmetry.hpp>
|
|
|
|
#include "HKLKey.h"
|
|
#include "RfreeFlags.h"
|
|
#include "FrenchWilson.h"
|
|
#include "ResolutionCutoff.h"
|
|
#include "../../common/CorrelationCoefficient.h"
|
|
#include "../../common/CrystalLattice.h"
|
|
#include "../../common/Definitions.h"
|
|
#include "../../common/JFJochException.h"
|
|
#include "../../common/ResolutionShells.h"
|
|
|
|
namespace {
|
|
// These mirror the per-image ScaleOnTheFly / Merge rocking-curve physics verbatim so this flat
|
|
// implementation is numerically identical - see the comments there for the details.
|
|
constexpr size_t MIN_REFLECTIONS = 20;
|
|
constexpr int64_t MIN_REFLECTIONS_FOR_IMAGE_CC = 20; // below this a frame's CC means nothing // per-frame scale needs at least this many
|
|
constexpr double SCALE_ROBUST_K = 3.0; // Cauchy loss scale (sigma units) for the per-frame G fit
|
|
|
|
// A fitted per-frame scale below this fraction of the run's median is not a measurement of anything:
|
|
// it would mean the frame received 2% of its neighbours' dose while still producing indexable spots.
|
|
// It matters because the scale enters as 1/G, so a G that collapses toward zero multiplies every
|
|
// intensity on that frame without bound - and sigma by the identical factor, which is why no
|
|
// sigma-based outlier test can see it. Measured over 12 rotation crystals in the default
|
|
// configuration, the smallest LEGITIMATE min(G)/median(G) is 0.070; the case this bound exists for
|
|
// sat at 0.0017 (a 546x amplification that put 99% of sum(I^2) into 0.66% of the observations and
|
|
// took the merged CC1/2 from 93.7% to 17.2%). 0.02 is 3.5x below anything real and 12x above the
|
|
// failure.
|
|
constexpr double MIN_CREDIBLE_SCALE_RATIO = 0.02;
|
|
constexpr float MAX_FRAME_GAP = 2.0f; // a rocking event is a run of frames no more apart than this
|
|
constexpr double CHI2_1_MEDIAN = 0.454936;
|
|
// A post-scale-fulls correction surface (decay / absorption) is applied only if its held-out
|
|
// cross-validation improvement exceeds this fraction of the held-out scatter. A margin (not just >0)
|
|
// stops it acting on marginal, noise-level "improvements" that leave the merged quality unchanged or
|
|
// slightly worse - a correction should engage only when the generalizing signal is clear.
|
|
constexpr double CV_MIN_RELATIVE_GAIN = 0.02;
|
|
|
|
double SafeInv(double x, double fallback) {
|
|
if (!std::isfinite(x) || x == 0.0)
|
|
return fallback;
|
|
return 1.0 / x;
|
|
}
|
|
|
|
// Kabsch rotation partiality: the fraction of a reflection recorded in the sampled slice, from the
|
|
// erf of the rocking angle relative to the mosaic width. Identical to ScaleOnTheFly's RotationPartiality
|
|
// (and the predictor's), so recomputing here just swaps in the smoothed mosaicity.
|
|
float RotationPartiality(double delta_phi_deg, double zeta, double mosaicity_deg, double wedge_deg) {
|
|
const double half_wedge = wedge_deg / 2.0;
|
|
const double c1 = zeta / std::sqrt(2.0);
|
|
const double arg_plus = (delta_phi_deg + half_wedge) * c1 / mosaicity_deg;
|
|
const double arg_minus = (delta_phi_deg - half_wedge) * c1 / mosaicity_deg;
|
|
return static_cast<float>((std::erf(arg_plus) - std::erf(arg_minus)) / 2.0);
|
|
}
|
|
|
|
// Deterministic CC1/2 half from the frame's stable index (splitmix64), matching Merge.cpp.
|
|
int HalfForImage(int64_t image_id) {
|
|
uint64_t z = static_cast<uint64_t>(image_id) + 0x9e3779b97f4a7c15ULL;
|
|
z = (z ^ (z >> 30)) * 0xbf58476d1ce4e5b9ULL;
|
|
z = (z ^ (z >> 27)) * 0x94d049bb133111ebULL;
|
|
z = z ^ (z >> 31);
|
|
return static_cast<int>(z & 1ULL);
|
|
}
|
|
|
|
struct ScaleObs { double coeff, Iobs, weight; };
|
|
|
|
// Robust per-frame scale (linear in G): the exact objective ScaleOnTheFly::SolveScaleIRLS solves.
|
|
double SolveScaleIRLS(const std::vector<ScaleObs> &obs, double robust_k) {
|
|
auto weighted_scale = [&obs](auto robust_weight) {
|
|
double num = 0.0, den = 0.0;
|
|
for (const auto &o : obs) {
|
|
const double rw = robust_weight(o);
|
|
const double w2 = o.weight * o.weight;
|
|
num += rw * w2 * o.coeff * o.Iobs;
|
|
den += rw * w2 * o.coeff * o.coeff;
|
|
}
|
|
return den > 0.0 ? num / den : NAN;
|
|
};
|
|
|
|
double G = weighted_scale([](const ScaleObs &) { return 1.0; });
|
|
if (!std::isfinite(G))
|
|
return 1.0;
|
|
G = std::max(0.0, G);
|
|
|
|
const double k2 = robust_k * robust_k;
|
|
for (int iter = 0; iter < 30; ++iter) {
|
|
const double G_prev = G;
|
|
const double G_next = weighted_scale([&](const ScaleObs &o) {
|
|
const double res = o.weight * (G * o.coeff - o.Iobs);
|
|
return 1.0 / (1.0 + res * res / k2);
|
|
});
|
|
if (!std::isfinite(G_next))
|
|
break;
|
|
G = std::max(0.0, G_next);
|
|
if (std::abs(G - G_prev) <= 1e-7 * std::max(G, 1.0))
|
|
break;
|
|
}
|
|
return G;
|
|
}
|
|
|
|
// Run fn(i) for i in [0, n) over `nthreads` workers pulling from a shared atomic counter - the same
|
|
// self-load-balancing pattern the rest of the codebase uses (heavy frames don't stall light ones).
|
|
// Work-stealing per-item parallel: one atomic fetch per item. Use ONLY when the per-item work is
|
|
// heavy and uneven (e.g. per-frame fits) - the atomic amortises. For millions of tiny uniform items
|
|
// use ParallelChunks instead; a per-item atomic there is pure contention.
|
|
template <typename Fn>
|
|
void ParallelFor(int n, size_t nthreads, Fn fn) {
|
|
if (n <= 0) return;
|
|
if (nthreads <= 1 || n == 1) {
|
|
for (int i = 0; i < n; ++i) fn(i);
|
|
return;
|
|
}
|
|
const size_t local = std::min(nthreads, static_cast<size_t>(n));
|
|
std::atomic<int> next = 0;
|
|
std::vector<std::future<void>> futures;
|
|
futures.reserve(local);
|
|
for (size_t t = 0; t < local; ++t)
|
|
futures.emplace_back(std::async(std::launch::async, [&] {
|
|
for (int i = next.fetch_add(1); i < n; i = next.fetch_add(1))
|
|
fn(i);
|
|
}));
|
|
for (auto &f : futures) f.get();
|
|
}
|
|
|
|
// Chunked parallel: each worker gets one contiguous [lo, hi) range, no per-item synchronisation.
|
|
// Right for millions of cheap uniform items (the CPU stand-in for a flat CUDA grid-stride kernel).
|
|
template <typename Fn>
|
|
void ParallelChunks(int n, size_t nthreads, Fn fn) {
|
|
if (n <= 0) return;
|
|
const int nt = static_cast<int>(std::max<size_t>(1, std::min(nthreads, static_cast<size_t>(n))));
|
|
if (nt == 1) { fn(0, n); return; }
|
|
const int chunk = (n + nt - 1) / nt;
|
|
std::vector<std::future<void>> futures;
|
|
futures.reserve(nt);
|
|
for (int t = 0; t < nt; ++t) {
|
|
const int lo = t * chunk, hi = std::min(n, lo + chunk);
|
|
if (lo >= hi) break;
|
|
futures.emplace_back(std::async(std::launch::async, [&fn, lo, hi] { fn(lo, hi); }));
|
|
}
|
|
for (auto &f : futures) f.get();
|
|
}
|
|
|
|
double median_of(std::vector<double> &v) {
|
|
std::nth_element(v.begin(), v.begin() + v.size() / 2, v.end());
|
|
return v[v.size() / 2];
|
|
}
|
|
}
|
|
|
|
RotationScaleMerge::RotationScaleMerge(const DiffractionExperiment &experiment,
|
|
std::vector<IntegrationOutcome> &partial_outcomes,
|
|
std::optional<UnitCell> reference_cell,
|
|
int scaling_iterations, float ice_ring_half_width_q,
|
|
size_t nthreads, Logger &logger,
|
|
std::string observation_dump_path)
|
|
: x(experiment), partials_out(partial_outcomes), reference_cell(std::move(reference_cell)),
|
|
nthreads(nthreads == 0 ? std::thread::hardware_concurrency() : nthreads), logger(logger),
|
|
observation_dump_path(std::move(observation_dump_path)) {
|
|
const auto s = x.GetScalingSettings();
|
|
min_partiality = s.GetMinPartiality();
|
|
d_min_limit = s.GetHighResolutionLimit_A();
|
|
merge_friedel = s.GetMergeFriedel();
|
|
capture_uncertainty_coeff = s.GetCaptureUncertaintyCoeff();
|
|
min_captured_fraction = s.GetMinCapturedFraction();
|
|
min_cc_for_image = s.GetMinCCForImage();
|
|
reject_nsigma = s.GetOutlierRejectNsigma();
|
|
reject_outliers = reject_nsigma > 0.0;
|
|
rfree_fraction = s.GetRfreeFraction();
|
|
scale_fulls = s.GetScaleFulls();
|
|
// Decay + absorption correction surfaces are one master toggle (on by default; both cross-validated,
|
|
// so a no-op when their systematic is absent).
|
|
refine_decay_b = s.GetCorrectionSurfaces();
|
|
absorption_iter = s.GetCorrectionSurfaces() ? s.GetAbsorptionIter() : 0;
|
|
modulation_iter = s.GetCorrectionSurfaces() ? s.GetAbsorptionIter() : 0;
|
|
relative_b_deg = s.GetRelativeBDegrees(); // opt-in per-batch relative-B (0 = off)
|
|
scaling_iter = std::max(1, scaling_iterations);
|
|
resolution_cutoff_method = s.GetResolutionCutoff();
|
|
resolution_cc_target = s.GetResolutionCCTarget();
|
|
report_shell_count = s.GetReportShellCount();
|
|
if (const auto forced = s.GetForcedMosaicity(); forced.has_value())
|
|
mosaicity_deg = *forced;
|
|
else
|
|
mosaicity_deg = s.GetDefaultMosaicity();
|
|
ice_half_width_q = ice_ring_half_width_q;
|
|
}
|
|
|
|
void RotationScaleMerge::Ingest() {
|
|
n_frames = static_cast<int>(partials_out.size());
|
|
|
|
size_t total = 0;
|
|
for (const auto &o : partials_out) total += o.reflections.size();
|
|
|
|
partials.clear();
|
|
partials.reserve(total);
|
|
frame_start.assign(n_frames, 0);
|
|
frame_count.assign(n_frames, 0);
|
|
frame_cell_ok.assign(n_frames, 1);
|
|
g_partial.assign(n_frames, 1.0);
|
|
|
|
const float dist_tol = x.GetIndexingSettings().GetUnitCellDistTolerance();
|
|
const float ang_tol = x.GetIndexingSettings().GetUnitCellAngleTolerance_deg();
|
|
|
|
for (int o = 0; o < n_frames; ++o) {
|
|
frame_start[o] = static_cast<int32_t>(partials.size());
|
|
if (reference_cell) {
|
|
const auto cell = partials_out[o].latt.GetUnitCell();
|
|
frame_cell_ok[o] = cell.is_close(*reference_cell, dist_tol, ang_tol) ? 1 : 0;
|
|
}
|
|
for (const auto &r : partials_out[o].reflections) {
|
|
Obs obs{};
|
|
obs.h = r.h; obs.k = r.k; obs.l = r.l;
|
|
obs.I = r.I; obs.sigma = r.sigma; obs.d = r.d; obs.rlp = r.rlp;
|
|
obs.partiality = r.partiality; obs.zeta = r.zeta; obs.delta_phi = r.delta_phi_deg; obs.bkg = r.bkg;
|
|
obs.px = r.predicted_x; obs.py = r.predicted_y;
|
|
obs.image_number = r.image_number;
|
|
obs.frame = o;
|
|
obs.on_ice = r.on_ice_ring ? 1 : 0;
|
|
obs.corr = r.image_scale_corr;
|
|
obs.group = -1;
|
|
partials.push_back(obs);
|
|
}
|
|
frame_count[o] = static_cast<int32_t>(partials.size()) - frame_start[o];
|
|
}
|
|
|
|
// Per-obs AcceptReflection finiteness (immutable) - lets ComputeAsuGroups stamp the ASU-group id per
|
|
// obs from a flat 1-byte array instead of re-reading the fat Obs struct for every space group.
|
|
finite_ok.resize(partials.size());
|
|
for (size_t i = 0; i < partials.size(); ++i) {
|
|
const auto &o = partials[i];
|
|
finite_ok[i] = (std::isfinite(o.I) && std::isfinite(o.rlp) && o.rlp != 0.0f
|
|
&& std::isfinite(o.sigma) && o.sigma > 0.0f) ? 1 : 0;
|
|
}
|
|
|
|
// Sort ONCE by (raw h,k,l, image_number) and split into raw-hkl runs. This is the one expensive sort;
|
|
// both the 3D combine (event split) and the per-space-group ASU grouping reuse this order.
|
|
perm.resize(partials.size());
|
|
for (int i = 0; i < static_cast<int>(partials.size()); ++i) perm[i] = i;
|
|
std::sort(perm.begin(), perm.end(), [&](int32_t a, int32_t b) {
|
|
const auto &x1 = partials[a];
|
|
const auto &y1 = partials[b];
|
|
if (x1.h != y1.h) return x1.h < y1.h;
|
|
if (x1.k != y1.k) return x1.k < y1.k;
|
|
if (x1.l != y1.l) return x1.l < y1.l;
|
|
return x1.image_number < y1.image_number;
|
|
});
|
|
rawrun_start.clear(); rawrun_count.clear();
|
|
rawrun_h.clear(); rawrun_k.clear(); rawrun_l.clear(); rawrun_d.clear();
|
|
for (int i = 0; i < static_cast<int>(perm.size()); ) {
|
|
const auto &o0 = partials[perm[i]];
|
|
int j = i;
|
|
float d = NAN;
|
|
while (j < static_cast<int>(perm.size())) {
|
|
const auto &o = partials[perm[j]];
|
|
if (o.h != o0.h || o.k != o0.k || o.l != o0.l) break;
|
|
if (!std::isfinite(d) && std::isfinite(o.d) && o.d > 0.0f) d = o.d;
|
|
++j;
|
|
}
|
|
rawrun_start.push_back(i);
|
|
rawrun_count.push_back(j - i);
|
|
rawrun_h.push_back(o0.h); rawrun_k.push_back(o0.k); rawrun_l.push_back(o0.l);
|
|
rawrun_d.push_back(d);
|
|
i = j;
|
|
}
|
|
rawrun_group.assign(rawrun_start.size(), -1);
|
|
logger.Info("RotationScaleMerge: ingested {} partial observations from {} frames ({} distinct hkl)",
|
|
total, n_frames, rawrun_start.size());
|
|
|
|
SmoothMosaicityAndPartiality();
|
|
|
|
#ifdef JFJOCH_USE_CUDA
|
|
// Bring the partial-scaling loop onto the GPU when one is present. Upload the immutable per-obs
|
|
// fields once (corr lives on the device, refreshed each pass); the CPU keeps the sort/keying/combine.
|
|
gpu_ = std::make_unique<RotationScaleMergeGPU>();
|
|
gpu_active_ = gpu_->Available();
|
|
if (gpu_active_) {
|
|
const int n = static_cast<int>(partials.size());
|
|
std::vector<float> I(n), sigma(n), rlp(n), part(n), zeta(n), corr(n), bkg(n), img(n), dd(n),
|
|
px(n), py(n);
|
|
std::vector<uint8_t> onice(n);
|
|
std::vector<int32_t> frm(n);
|
|
for (int i = 0; i < n; ++i) {
|
|
const auto &o = partials[i];
|
|
I[i] = o.I; sigma[i] = o.sigma; rlp[i] = o.rlp; part[i] = o.partiality;
|
|
zeta[i] = o.zeta; onice[i] = o.on_ice; frm[i] = o.frame; corr[i] = o.corr;
|
|
bkg[i] = o.bkg; img[i] = o.image_number; dd[i] = o.d; px[i] = o.px; py[i] = o.py;
|
|
}
|
|
gpu_->SetPartials(n, n_frames, I.data(), sigma.data(), rlp.data(), part.data(), zeta.data(),
|
|
onice.data(), frm.data(), corr.data(), frame_start.data(), frame_count.data());
|
|
gpu_->SetCombineInputs(bkg.data(), img.data(), dd.data(), px.data(), py.data());
|
|
gpu_->SetRawRuns(static_cast<int>(rawrun_start.size()), static_cast<int>(perm.size()), perm.data(),
|
|
rawrun_start.data(), rawrun_count.data(),
|
|
rawrun_h.data(), rawrun_k.data(), rawrun_l.data());
|
|
gpu_->SetFrameCellOk(frame_cell_ok.data());
|
|
logger.Info("RotationScaleMerge: GPU scaling + combine + scale-fulls + merge active");
|
|
}
|
|
#endif
|
|
}
|
|
|
|
void RotationScaleMerge::SmoothMosaicityAndPartiality() {
|
|
// Per-frame mosaicity to recompute partiality from. A forced (fixed) mosaicity overrides every frame;
|
|
// otherwise use the per-frame value measured at integration (image-local, deterministic).
|
|
const auto forced_mosaicity = x.GetScalingSettings().GetForcedMosaicity();
|
|
std::vector<double> mos_raw(n_frames, NAN);
|
|
if (forced_mosaicity.has_value() && std::isfinite(*forced_mosaicity) && *forced_mosaicity > 0.0) {
|
|
for (int o = 0; o < n_frames; ++o) mos_raw[o] = *forced_mosaicity;
|
|
} else {
|
|
for (int o = 0; o < n_frames; ++o) {
|
|
const auto &m = partials_out[o].mosaicity_deg;
|
|
if (m && std::isfinite(*m) && *m > 0.0f) mos_raw[o] = *m;
|
|
}
|
|
}
|
|
|
|
// Frame-order moving average with the same window as smooth-G (a rotation range -> frame count).
|
|
// With smoothing off, fall back to the per-frame value (still deterministic, just unsmoothed).
|
|
const auto ss = x.GetScalingSettings();
|
|
const double smooth_deg = ss.GetSmoothGDegrees();
|
|
const auto gon = x.GetGoniometer();
|
|
const double osc = gon ? std::fabs(gon->GetIncrement_deg()) : 0.0;
|
|
mos_smooth.assign(n_frames, NAN);
|
|
if (smooth_deg > 0.0 && osc > 1e-6) {
|
|
int window = std::max(1, static_cast<int>(std::lround(smooth_deg / osc)));
|
|
if (window % 2 == 0) ++window;
|
|
const int half = window / 2;
|
|
for (int o = 0; o < n_frames; ++o) {
|
|
double sum = 0.0;
|
|
int cnt = 0;
|
|
for (int j = std::max(0, o - half); j <= std::min(n_frames - 1, o + half); ++j)
|
|
if (std::isfinite(mos_raw[j])) { sum += mos_raw[j]; ++cnt; }
|
|
if (cnt > 0) mos_smooth[o] = static_cast<float>(sum / cnt);
|
|
}
|
|
} else {
|
|
for (int o = 0; o < n_frames; ++o) mos_smooth[o] = static_cast<float>(mos_raw[o]);
|
|
}
|
|
|
|
// Recompute each partial's partiality from the smoothed mosaicity (same wedge the predictor used).
|
|
// Frames without a mosaicity keep the stored partiality.
|
|
const double wedge = gon ? std::fabs(gon->GetWedge_deg()) : 0.0;
|
|
ParallelChunks(static_cast<int>(partials.size()), nthreads, [&](int lo, int hi) {
|
|
for (int i = lo; i < hi; ++i) {
|
|
auto &o = partials[i];
|
|
const float mos = mos_smooth[o.frame];
|
|
if (std::isfinite(mos) && mos > 1e-6f && std::isfinite(o.zeta) && o.zeta > 0.0f
|
|
&& std::isfinite(o.delta_phi))
|
|
o.partiality = RotationPartiality(o.delta_phi, o.zeta, mos, wedge);
|
|
}
|
|
});
|
|
logger.Info("Recomputed partiality from frame-order-smoothed mosaicity");
|
|
}
|
|
|
|
int RotationScaleMerge::ComputeAsuGroups(const HKLKeyGenerator &keygen) {
|
|
// One ASU reduction per distinct raw hkl (not per observation): a raw hkl is eligible if it is not
|
|
// systematically absent and its resolution is in range. Group the eligible raw hkls by ASU key
|
|
// (sort the ~#distinct-hkl keys, not the millions of observations), then hand out dense ids.
|
|
const int n_run = static_cast<int>(rawrun_start.size());
|
|
std::vector<uint64_t> key(n_run);
|
|
std::vector<uint8_t> eligible(n_run, 0);
|
|
// The gemmi ASU reduction / absence test per raw hkl is the cost here and is independent per run
|
|
// (HKLKeyGenerator is const, so concurrent reads are safe) - compute keys in parallel chunks.
|
|
ParallelChunks(n_run, nthreads, [&](int lo, int hi) {
|
|
for (int r = lo; r < hi; ++r) {
|
|
rawrun_group[r] = -1;
|
|
if (keygen.IsSystematicallyAbsent(rawrun_h[r], rawrun_k[r], rawrun_l[r]))
|
|
continue;
|
|
const float d = rawrun_d[r]; // resolution is a per-raw-hkl property (all its partials share d)
|
|
if (!std::isfinite(d) || d <= 0.0f) continue;
|
|
if (d_min_limit && d < *d_min_limit) continue;
|
|
key[r] = keygen(rawrun_h[r], rawrun_k[r], rawrun_l[r]).pack();
|
|
eligible[r] = 1;
|
|
}
|
|
});
|
|
std::vector<int32_t> idx;
|
|
idx.reserve(n_run);
|
|
for (int r = 0; r < n_run; ++r)
|
|
if (eligible[r]) idx.push_back(r);
|
|
std::sort(idx.begin(), idx.end(), [&](int32_t a, int32_t b) { return key[a] < key[b]; });
|
|
|
|
group_h.clear(); group_k.clear(); group_l.clear();
|
|
int n_groups = 0;
|
|
for (size_t j = 0; j < idx.size(); ++j) {
|
|
if (j == 0 || key[idx[j]] != key[idx[j - 1]]) {
|
|
const int r = idx[j];
|
|
const auto hkl = keygen(rawrun_h[r], rawrun_k[r], rawrun_l[r]);
|
|
group_h.push_back(hkl.plus ? hkl.h : -hkl.h);
|
|
group_k.push_back(hkl.plus ? hkl.k : -hkl.k);
|
|
group_l.push_back(hkl.plus ? hkl.l : -hkl.l);
|
|
++n_groups;
|
|
}
|
|
rawrun_group[idx[j]] = n_groups - 1;
|
|
}
|
|
|
|
// Stamp the ASU-group id per obs from its raw hkl + the precomputed finiteness. For the GPU we build a
|
|
// flat group_ids array (fed to the reduction); the Obs.group field is written only when a CPU stage
|
|
// will read it - i.e. no GPU (with a GPU, scaling/CC/combine all read group_ids / rawrun_group, never
|
|
// partials.group) - so the default path skips a strided 6.3M pass over the fat Obs struct.
|
|
const bool need_obs_group =
|
|
#ifdef JFJOCH_USE_CUDA
|
|
!gpu_active_;
|
|
#else
|
|
true;
|
|
#endif
|
|
#ifdef JFJOCH_USE_CUDA
|
|
std::vector<int32_t> group_ids;
|
|
if (gpu_active_) group_ids.resize(partials.size());
|
|
#endif
|
|
// Parallel over raw-hkl runs: distinct runs own disjoint perm ranges, hence disjoint observations.
|
|
ParallelChunks(n_run, nthreads, [&](int rlo, int rhi) {
|
|
for (int r = rlo; r < rhi; ++r) {
|
|
const int g = rawrun_group[r];
|
|
const int lo = rawrun_start[r], hi = rawrun_start[r] + rawrun_count[r];
|
|
for (int p = lo; p < hi; ++p) {
|
|
const int i = perm[p];
|
|
const int gi = (g >= 0 && finite_ok[i]) ? g : -1;
|
|
#ifdef JFJOCH_USE_CUDA
|
|
if (gpu_active_) group_ids[i] = gi;
|
|
#endif
|
|
if (need_obs_group) partials[i].group = gi;
|
|
}
|
|
}
|
|
});
|
|
|
|
#ifdef JFJOCH_USE_CUDA
|
|
// Group-ordered permutation (obs bucketed by ASU group, obs-index order) + its CSR, so the GPU
|
|
// reduction is a deterministic segmented reduction (fixed order, no atomics). A stable counting sort
|
|
// by group, parallel via per-chunk histograms over the flat group_ids; bit-identical to a serial fill
|
|
// (chunk order == obs-index order, each chunk fills its groups sequentially).
|
|
if (gpu_active_) {
|
|
const int n = static_cast<int>(partials.size());
|
|
const int nt = static_cast<int>(std::max<size_t>(1, std::min(nthreads, static_cast<size_t>(std::max(1, n)))));
|
|
const int chunk = (n + nt - 1) / nt;
|
|
std::vector<std::vector<int32_t>> hist(nt, std::vector<int32_t>(n_groups, 0));
|
|
|
|
// Pass 1 (parallel): per-chunk group histogram over the flat group_ids.
|
|
std::vector<std::future<void>> f1;
|
|
for (int t = 0; t < nt; ++t) {
|
|
const int lo = t * chunk, hi = std::min(n, lo + chunk);
|
|
if (lo >= hi) break;
|
|
f1.emplace_back(std::async(std::launch::async, [&, t, lo, hi] {
|
|
auto &h = hist[t];
|
|
for (int i = lo; i < hi; ++i) { const int g = group_ids[i]; if (g >= 0) ++h[g]; }
|
|
}));
|
|
}
|
|
for (auto &f : f1) f.get();
|
|
|
|
// CSR starts + convert hist[t][g] into chunk t's write base for group g (exclusive prefix over t).
|
|
std::vector<int32_t> gstart(n_groups), gcount(n_groups);
|
|
int acc = 0;
|
|
for (int g = 0; g < n_groups; ++g) {
|
|
int base = acc;
|
|
gstart[g] = acc;
|
|
for (int t = 0; t < nt; ++t) { const int c = hist[t][g]; hist[t][g] = base; base += c; }
|
|
gcount[g] = base - acc;
|
|
acc = base;
|
|
}
|
|
|
|
// Pass 2 (parallel): each chunk fills its obs into gperm at its per-group base (stable).
|
|
std::vector<int32_t> gperm(acc);
|
|
std::vector<std::future<void>> f2;
|
|
for (int t = 0; t < nt; ++t) {
|
|
const int lo = t * chunk, hi = std::min(n, lo + chunk);
|
|
if (lo >= hi) break;
|
|
f2.emplace_back(std::async(std::launch::async, [&, t, lo, hi] {
|
|
std::vector<int32_t> fill = hist[t];
|
|
for (int i = lo; i < hi; ++i) { const int g = group_ids[i]; if (g >= 0) gperm[fill[g]++] = i; }
|
|
}));
|
|
}
|
|
for (auto &f : f2) f.get();
|
|
gpu_->SetGroups(n_groups, group_ids.data(), gperm.data(), acc, gstart.data(), gcount.data());
|
|
}
|
|
#endif
|
|
return n_groups;
|
|
}
|
|
|
|
void RotationScaleMerge::ReduceGroupMeans(const std::vector<Obs> &obs, int n_groups,
|
|
bool exclude_ice, const std::vector<char> &masked,
|
|
std::vector<double> &out_mean) const {
|
|
// Inverse-variance per-group mean of I*corr = the merge reference (a segmented reduction over the
|
|
// groups; the CPU stand-in for a CUDA reduce_by_key). No cell mask here: the scaling reference
|
|
// (MergeAll) is built without a reference cell - only the final merge applies it.
|
|
std::vector<double> sw(n_groups, 0.0), swI(n_groups, 0.0);
|
|
for (const auto &o : obs) {
|
|
if (o.group < 0) continue;
|
|
if (!(o.corr > 0.0f) || !std::isfinite(o.corr)) continue;
|
|
if (exclude_ice && o.on_ice) continue;
|
|
if (!masked.empty()) {
|
|
const int ring = IceRingIndex(o.d, ice_half_width_q);
|
|
if (ring >= 0 && ring < static_cast<int>(masked.size()) && masked[ring]) continue;
|
|
}
|
|
if (o.partiality < min_partiality) continue;
|
|
const float I_corr = o.I * o.corr;
|
|
const float sigma_corr = o.sigma * o.corr;
|
|
if (!std::isfinite(I_corr) || !std::isfinite(sigma_corr) || sigma_corr <= 0.0f) continue;
|
|
const double w = 1.0 / (static_cast<double>(sigma_corr) * sigma_corr);
|
|
sw[o.group] += w;
|
|
swI[o.group] += w * I_corr;
|
|
}
|
|
out_mean.assign(n_groups, NAN);
|
|
for (int g = 0; g < n_groups; ++g)
|
|
if (sw[g] > 0.0) out_mean[g] = swI[g] / sw[g];
|
|
}
|
|
|
|
void RotationScaleMerge::FitPerFrameG(std::vector<Obs> &obs, const std::vector<int32_t> &fstart,
|
|
const std::vector<int32_t> &fcount,
|
|
const std::vector<double> &group_mean_in,
|
|
bool unity, std::vector<double> &g) {
|
|
std::vector<uint8_t> scaled(fstart.size(), 0);
|
|
ParallelFor(static_cast<int>(fstart.size()), nthreads, [&](int f) {
|
|
std::vector<ScaleObs> so;
|
|
so.reserve(fcount[f]);
|
|
const int lo = fstart[f], hi = fstart[f] + fcount[f];
|
|
for (int i = lo; i < hi; ++i) {
|
|
const auto &o = obs[i];
|
|
if (o.group < 0) continue;
|
|
if (o.on_ice) continue;
|
|
const double mean = group_mean_in[o.group];
|
|
if (!std::isfinite(mean)) continue;
|
|
double coeff;
|
|
if (unity) {
|
|
coeff = mean; // partiality already folded into the full, rlp = 1
|
|
} else {
|
|
if (!(std::isfinite(o.zeta) && o.zeta > 0.0f)) continue; // Rotation model needs zeta > 0
|
|
coeff = o.partiality * SafeInv(o.rlp, 1.0) * mean;
|
|
}
|
|
so.push_back({coeff, static_cast<double>(o.I), SafeInv(o.sigma, 1.0)});
|
|
}
|
|
if (so.size() < MIN_REFLECTIONS) return; // leave g[f]/corr untouched (as ScaleOnTheFly does)
|
|
g[f] = SolveScaleIRLS(so, SCALE_ROBUST_K);
|
|
scaled[f] = 1;
|
|
});
|
|
// Remember which frames were fitted this call (so the caller updates corr only there).
|
|
frame_scaled_scratch = std::move(scaled);
|
|
}
|
|
|
|
void RotationScaleMerge::UpdateCorr(std::vector<Obs> &obs, const std::vector<double> &g,
|
|
const std::vector<uint8_t> &frame_scaled) const {
|
|
ParallelChunks(static_cast<int>(obs.size()), nthreads, [&](int lo, int hi) {
|
|
for (int i = lo; i < hi; ++i) {
|
|
auto &o = obs[i];
|
|
if (!frame_scaled[o.frame]) continue;
|
|
const double denom = static_cast<double>(o.partiality) * g[o.frame]; // B_term = 1 (no B refine)
|
|
if (std::isfinite(o.rlp) && std::isfinite(denom) && denom > 0.0)
|
|
o.corr = static_cast<float>(o.rlp / denom);
|
|
else
|
|
o.corr = NAN;
|
|
}
|
|
});
|
|
}
|
|
|
|
void RotationScaleMerge::RefineDecay(int n_groups) {
|
|
// Radiation damage weakens later frames more at higher resolution - a resolution x time (Debye-Waller)
|
|
// systematic the resolution-flat per-frame G leaves in. Model it as a single global relative-B rate:
|
|
// ln(I_ref / I_obs) = slope * (frame - frame_center) * s^2, s^2 = 1/(4 d^2), slope = 2 dB/dframe,
|
|
// fitted (inverse-variance-of-log weighted) against the merged reference.
|
|
//
|
|
// The fit is dominated by strong LOW-res reflections and EXTRAPOLATED to high s^2, so a slope fitted to
|
|
// noise / aggregate structure (not real per-reflection decay) would amplify into a high-res
|
|
// mis-correction and worsen the symmetry-equivalent scatter that sets ISa. An opt-in correction must
|
|
// never add noise where its systematic is absent (a low-dose reference set), so the slope is
|
|
// CROSS-VALIDATED: fit on even-frame fulls and score the held-out odd-frame equivalent agreement (a
|
|
// sigma-independent Rmeas-like metric), and vice versa. Real decay generalizes across the split (both
|
|
// held-out scores improve); a noise slope does not. Apply the full-data slope only if the mean held-out
|
|
// agreement strictly improves; else it is a no-op.
|
|
if (fulls.empty() || n_frames <= 0)
|
|
return;
|
|
const double fcenter = 0.5 * n_frames;
|
|
auto s2_of = [](float d) { return d > 0.0f ? 1.0 / (4.0 * static_cast<double>(d) * d) : 0.0; };
|
|
auto usable = [&](const Obs &o) {
|
|
return o.group >= 0 && o.corr > 0.0f && std::isfinite(o.corr) && (o.d > 0.0f)
|
|
&& o.partiality >= min_partiality;
|
|
};
|
|
auto decay_factor = [&](const Obs &o, double slope) {
|
|
return std::exp(slope * (o.image_number - fcenter) * s2_of(o.d));
|
|
};
|
|
|
|
// Fit the global slope over the subset {frame&1 == parity} (parity < 0 = all fulls), using an
|
|
// inverse-variance reference built from that same subset. slope*factor applied to the current corr.
|
|
auto fit_slope = [&](int parity) -> double {
|
|
std::vector<double> sw(n_groups, 0.0), swI(n_groups, 0.0);
|
|
for (const auto &o : fulls) {
|
|
if (!usable(o) || (parity >= 0 && (o.frame & 1) != parity)) continue;
|
|
const double sc = static_cast<double>(o.sigma) * o.corr, w = 1.0 / (sc * sc);
|
|
sw[o.group] += w; swI[o.group] += w * static_cast<double>(o.I) * o.corr;
|
|
}
|
|
double Sw = 0, Sx = 0, Sy = 0, Sxx = 0, Sxy = 0;
|
|
for (const auto &o : fulls) {
|
|
if (!usable(o) || (parity >= 0 && (o.frame & 1) != parity) || sw[o.group] <= 0.0) continue;
|
|
const double Iref = swI[o.group] / sw[o.group];
|
|
const double Is = static_cast<double>(o.I) * o.corr, sc = static_cast<double>(o.sigma) * o.corr;
|
|
if (!std::isfinite(Iref) || Iref <= 0.0 || !(Is > 0.0)) continue;
|
|
const double w = (Is / sc) * (Is / sc);
|
|
const double x = (o.image_number - fcenter) * s2_of(o.d), y = std::log(Iref / Is);
|
|
Sw += w; Sx += w * x; Sy += w * y; Sxx += w * x * x; Sxy += w * x * y;
|
|
}
|
|
const double var = (Sw > 0.0) ? Sxx - Sx * Sx / Sw : 0.0;
|
|
// Clamp only to guard against a near-collinear (var ~ 0) blow-up; generous enough to reach very
|
|
// strong damage (slope = 2 dB/dframe, so +-1.0 admits total relative-B up to ~n_frames/2 A^2).
|
|
return (var > 0.0) ? std::clamp((Sxy - Sx * Sy / Sw) / var, -1.0, 1.0) : 0.0;
|
|
};
|
|
// Rmeas-like (sigma-INDEPENDENT) disagreement of the equivalents over the subset {frame&1 == parity}
|
|
// when `slope` is applied, scored against that subset's own reference (no leakage): sum|Is-Iref|/sum|Iref|.
|
|
// Scoring on the studentized deviation instead would let the slope "improve" the held-out chi^2 by
|
|
// reshaping sigma - it scales sigma up on late / high-angle frames, down-weighting them - without
|
|
// actually tightening the intensities; a fractional metric cannot be gamed that way. This mirrors the
|
|
// absorption / modulation gate in ApplyCellSurface (which was switched off studentized chi^2 for the
|
|
// same reason); the reference is still the inverse-variance mean.
|
|
auto subset_disagreement = [&](int parity, double slope) -> double {
|
|
std::vector<double> sw(n_groups, 0.0), swI(n_groups, 0.0);
|
|
for (const auto &o : fulls) {
|
|
if (!usable(o) || (o.frame & 1) != parity) continue;
|
|
const double Is = static_cast<double>(o.I) * o.corr * decay_factor(o, slope);
|
|
const double sc = static_cast<double>(o.sigma) * o.corr * decay_factor(o, slope);
|
|
const double w = 1.0 / (sc * sc);
|
|
sw[o.group] += w; swI[o.group] += w * Is;
|
|
}
|
|
double num = 0.0, den = 0.0;
|
|
for (const auto &o : fulls) {
|
|
if (!usable(o) || (o.frame & 1) != parity || sw[o.group] <= 0.0) continue;
|
|
const double Iref = swI[o.group] / sw[o.group];
|
|
const double Is = static_cast<double>(o.I) * o.corr * decay_factor(o, slope);
|
|
if (!std::isfinite(Iref) || Iref <= 0.0) continue;
|
|
num += std::abs(Is - Iref); den += Iref;
|
|
}
|
|
return den > 0.0 ? num / den : 0.0;
|
|
};
|
|
|
|
// A meaningful radiation-damage relative-B accrues several A^2 across a dataset. Below a physical floor
|
|
// the decay is negligible and "correcting" it only adds per-equivalent perturbation noise (equivalents
|
|
// sit at the same s^2 but different frames, so the correction spreads them) - net harmful. slope =
|
|
// 2 dB/dframe, so the total relative-B change across the run is |slope/2 * n_frames|.
|
|
constexpr double DECAY_MIN_DELTA_B = 2.0; // A^2, minimum total relative-B over the run to engage
|
|
const double slope = fit_slope(-1);
|
|
const double total_delta_B = std::fabs(0.5 * slope * n_frames);
|
|
if (total_delta_B < DECAY_MIN_DELTA_B) {
|
|
logger.Info("Decay correction: negligible radiation damage (total dB = {:.2f} A^2 < {:.1f}, skipped)",
|
|
total_delta_B, DECAY_MIN_DELTA_B);
|
|
return;
|
|
}
|
|
// And require the slope to cross-validate by a clear margin: fit on even frames, score the held-out
|
|
// odd equivalents (and vice versa). Real damage generalizes; a fluke does not.
|
|
const double base = subset_disagreement(1, 0.0) + subset_disagreement(0, 0.0);
|
|
const double gain = base - (subset_disagreement(1, fit_slope(0)) + subset_disagreement(0, fit_slope(1)));
|
|
if (!(gain > CV_MIN_RELATIVE_GAIN * base)) {
|
|
logger.Info("Decay correction: not cross-validated (dB = {:.2f} A^2, held-out gain {:.1f}%, skipped)",
|
|
total_delta_B, 100.0 * gain / std::max(base, 1e-30));
|
|
return;
|
|
}
|
|
for (auto &o : fulls)
|
|
if (usable(o))
|
|
o.corr = static_cast<float>(o.corr * decay_factor(o, slope));
|
|
logger.Info("Decay correction: total relative-B = {:.2f} A^2 over run (dB/dframe = {:.2e}, cross-validated)",
|
|
0.5 * slope * n_frames, 0.5 * slope);
|
|
}
|
|
|
|
std::vector<double> RotationScaleMerge::SolveCurvatureSmoothedB(const std::vector<double> &num,
|
|
const std::vector<double> &den) const {
|
|
// Solve for a SMOOTH per-batch relative-B from the per-batch normal-equation data (num_c = sum w s^2 y,
|
|
// den_c = sum w s^4). A per-batch regression alone (b_c = num_c/den_c) over-fits: a genuine relative-B is
|
|
// a low-curvature function of dose / orientation, but an unconstrained fit chases resolution-correlated
|
|
// noise into a rough, oscillating curve. So minimise data-fidelity PLUS a second-difference (curvature)
|
|
// penalty sum_c den_c (b_c - num_c/den_c)^2 + mu * sum_c (b_{c-1} - 2 b_c + b_{c+1})^2, which leaves a
|
|
// straight RAMP untouched but suppresses oscillation, solving the small SPD normal system A b = num,
|
|
// A = diag(den + lambda) + mu * D^T D, by Gauss-Seidel (n_batch <= a few tens). Not anchored - callers set
|
|
// the gauge (zero-mean for the correction, low-dose reference for the monitor).
|
|
constexpr double RELATIVE_B_MAX = 50.0; // A^2, clamp a single batch's relative-B (generous guard)
|
|
constexpr double RELATIVE_B_CURVATURE_K = 8.0; // curvature penalty weight (x the median batch leverage)
|
|
const int n_batch = static_cast<int>(num.size());
|
|
std::vector<double> dsorted(den);
|
|
std::sort(dsorted.begin(), dsorted.end());
|
|
const double med_den = dsorted.empty() ? 0.0 : dsorted[dsorted.size() / 2];
|
|
const double lambda = 0.1 * med_den; // ridge ~ 10% of the median batch leverage
|
|
const double mu = RELATIVE_B_CURVATURE_K * med_den; // curvature penalty
|
|
std::vector<double> A(static_cast<size_t>(n_batch) * n_batch, 0.0);
|
|
auto Aat = [&](int i, int j) -> double & { return A[static_cast<size_t>(i) * n_batch + j]; };
|
|
for (int c = 0; c < n_batch; ++c) Aat(c, c) = den[c] + lambda;
|
|
constexpr int st[3] = {-1, 0, 1};
|
|
constexpr double sv[3] = {1.0, -2.0, 1.0}; // second-difference stencil
|
|
for (int c = 1; c < n_batch - 1; ++c)
|
|
for (int a = 0; a < 3; ++a)
|
|
for (int b2 = 0; b2 < 3; ++b2)
|
|
Aat(c + st[a], c + st[b2]) += mu * sv[a] * sv[b2];
|
|
std::vector<double> b(n_batch, 0.0);
|
|
for (int it = 0; it < 100; ++it)
|
|
for (int c = 0; c < n_batch; ++c) {
|
|
if (!(Aat(c, c) > 0.0)) continue;
|
|
double s = num[c];
|
|
for (int j = 0; j < n_batch; ++j)
|
|
if (j != c) s -= Aat(c, j) * b[j];
|
|
b[c] = std::clamp(s / Aat(c, c), -RELATIVE_B_MAX, RELATIVE_B_MAX);
|
|
}
|
|
return b;
|
|
}
|
|
|
|
std::vector<double> RotationScaleMerge::FitRelativeBCurve(int n_groups, int n_batch, int frames_per_batch,
|
|
int gparity) const {
|
|
// Fit one relative-B per batch over the ASU-group subset {group&1 == gparity} (gparity < 0 = all fulls),
|
|
// against an inverse-variance reference built from that same subset (no leakage). Per batch the model is
|
|
// y = ln(I_ref / I_obs) = b * s^2 through the origin (the resolution-flat intercept is the per-frame G,
|
|
// already applied), so the data term is b_c = sum(w s^2 y) / sum(w s^2 s^2); SolveCurvatureSmoothedB then
|
|
// smooths it. Finally subtract the leverage-weighted mean: a constant B is a global Wilson-B degenerate
|
|
// with overall scale, and cannot change equivalent agreement (within one ASU group every equivalent sits
|
|
// at the same s^2, so only b VARYING between batches moves the merge). Drives the per-batch correction.
|
|
auto s2_of = [](float d) { return d > 0.0f ? 1.0 / (4.0 * static_cast<double>(d) * d) : 0.0; };
|
|
auto usable = [&](const Obs &o) {
|
|
return o.group >= 0 && o.corr > 0.0f && std::isfinite(o.corr) && (o.d > 0.0f)
|
|
&& o.partiality >= min_partiality;
|
|
};
|
|
auto batch_of = [&](const Obs &o) { return std::min(n_batch - 1, o.frame / frames_per_batch); };
|
|
|
|
std::vector<double> sw(n_groups, 0.0), swI(n_groups, 0.0);
|
|
for (const auto &o : fulls) {
|
|
if (!usable(o) || (gparity >= 0 && (o.group & 1) != gparity)) continue;
|
|
const double sc = static_cast<double>(o.sigma) * o.corr, w = 1.0 / (sc * sc);
|
|
sw[o.group] += w; swI[o.group] += w * static_cast<double>(o.I) * o.corr;
|
|
}
|
|
std::vector<double> num(n_batch, 0.0), den(n_batch, 0.0);
|
|
for (const auto &o : fulls) {
|
|
if (!usable(o) || (gparity >= 0 && (o.group & 1) != gparity) || sw[o.group] <= 0.0) continue;
|
|
const double Iref = swI[o.group] / sw[o.group];
|
|
const double Is = static_cast<double>(o.I) * o.corr, sc = static_cast<double>(o.sigma) * o.corr;
|
|
if (!std::isfinite(Iref) || Iref <= 0.0 || !(Is > 0.0)) continue;
|
|
const double w = (Is / sc) * (Is / sc);
|
|
const double s2 = s2_of(o.d), y = std::log(Iref / Is);
|
|
num[batch_of(o)] += w * s2 * y; den[batch_of(o)] += w * s2 * s2;
|
|
}
|
|
std::vector<double> b = SolveCurvatureSmoothedB(num, den);
|
|
double dw = 0.0, dbw = 0.0;
|
|
for (int c = 0; c < n_batch; ++c) { dw += den[c]; dbw += den[c] * b[c]; }
|
|
const double bmean = dw > 0.0 ? dbw / dw : 0.0;
|
|
for (auto &v : b) v -= bmean;
|
|
return b;
|
|
}
|
|
|
|
void RotationScaleMerge::MeasureRadiationDamageB(int n_groups) {
|
|
// Radiation-damage MONITOR (report-only): measure the per-batch relative Debye-Waller B on the scaled
|
|
// fulls BEFORE any correction, so it captures how the high-resolution scattering power fades with
|
|
// accumulated dose. Crucially the reference is the PRISTINE LOW-DOSE crystal (first ~10% of frames), not
|
|
// the whole-run mean: a group measured both early and late then contributes ln(I_early / I_late) vs s^2,
|
|
// whose slope is the ADDED B - an absolute, correctly-signed damage read (positive = high-res fades).
|
|
// A whole-run-mean reference inverts on heavily damaged data (the surviving late high-res reflections
|
|
// look strong against a mean dominated by the frames where they had already faded). Stores the first->last
|
|
// relative-B change (the headline "B-factor difference") and the per-batch curve; never touches corr.
|
|
rad_damage_delta_b = std::numeric_limits<double>::quiet_NaN();
|
|
rad_damage_b_batch.clear();
|
|
rad_damage_batch_deg = 0.0;
|
|
if (fulls.empty() || n_frames <= 0)
|
|
return;
|
|
const auto gon = x.GetGoniometer();
|
|
const double osc_deg = gon ? std::fabs(gon->GetIncrement_deg()) : 0.0;
|
|
if (!(osc_deg > 1e-6))
|
|
return;
|
|
constexpr double MONITOR_BATCH_DEG = 10.0; // rotation width per batch for the damage curve
|
|
const int frames_per_batch = std::max(1, static_cast<int>(std::lround(MONITOR_BATCH_DEG / osc_deg)));
|
|
const int n_batch = std::max(1, (n_frames + frames_per_batch - 1) / frames_per_batch);
|
|
if (n_batch < 2) // a single batch is a global Wilson-B (degenerate with scale) - nothing to monitor
|
|
return;
|
|
auto s2_of = [](float d) { return d > 0.0f ? 1.0 / (4.0 * static_cast<double>(d) * d) : 0.0; };
|
|
auto usable = [&](const Obs &o) {
|
|
return o.group >= 0 && o.corr > 0.0f && std::isfinite(o.corr) && (o.d > 0.0f)
|
|
&& o.partiality >= min_partiality;
|
|
};
|
|
auto batch_of = [&](const Obs &o) { return std::min(n_batch - 1, o.frame / frames_per_batch); };
|
|
|
|
// Low-dose reference: inverse-variance group means over the first ~10% of frames (>= 1 batch).
|
|
const int ref_frames = std::max(frames_per_batch, n_frames / 10);
|
|
std::vector<double> sw0(n_groups, 0.0), swI0(n_groups, 0.0);
|
|
for (const auto &o : fulls) {
|
|
if (!usable(o) || o.frame >= ref_frames) continue;
|
|
const double sc = static_cast<double>(o.sigma) * o.corr, w = 1.0 / (sc * sc);
|
|
sw0[o.group] += w; swI0[o.group] += w * static_cast<double>(o.I) * o.corr;
|
|
}
|
|
// Per batch: slope of ln(I_ref0 / I_obs) vs s^2, over groups also seen in the low-dose reference.
|
|
std::vector<double> num(n_batch, 0.0), den(n_batch, 0.0);
|
|
for (const auto &o : fulls) {
|
|
if (!usable(o) || sw0[o.group] <= 0.0) continue;
|
|
const double Iref = swI0[o.group] / sw0[o.group];
|
|
const double Is = static_cast<double>(o.I) * o.corr, sc = static_cast<double>(o.sigma) * o.corr;
|
|
if (!std::isfinite(Iref) || Iref <= 0.0 || !(Is > 0.0)) continue;
|
|
const double w = (Is / sc) * (Is / sc);
|
|
const double s2 = s2_of(o.d), y = std::log(Iref / Is);
|
|
num[batch_of(o)] += w * s2 * y; den[batch_of(o)] += w * s2 * s2;
|
|
}
|
|
std::vector<double> b = SolveCurvatureSmoothedB(num, den);
|
|
// Anchor so the reference batches read ~0 (b = added B relative to the low-dose start).
|
|
const int n_ref_batch = std::max(1, ref_frames / frames_per_batch);
|
|
double ref_mean = 0.0;
|
|
for (int c = 0; c < n_ref_batch && c < n_batch; ++c) ref_mean += b[c];
|
|
ref_mean /= std::min(n_ref_batch, n_batch);
|
|
for (auto &v : b) v -= ref_mean;
|
|
// Headline: the first->last relative-B change from a least-squares linear trend of the (smoothed) curve -
|
|
// robust to a single noisy end batch, unlike raw endpoints. Positive = high-res fades with dose = damage.
|
|
double sx = 0, sy = 0, sxx = 0, sxy = 0;
|
|
for (int c = 0; c < n_batch; ++c) { sx += c; sy += b[c]; sxx += double(c) * c; sxy += double(c) * b[c]; }
|
|
const double det = n_batch * sxx - sx * sx;
|
|
const double slope = det > 0.0 ? (n_batch * sxy - sx * sy) / det : 0.0;
|
|
rad_damage_delta_b = slope * (n_batch - 1);
|
|
rad_damage_batch_deg = MONITOR_BATCH_DEG;
|
|
rad_damage_b_batch.resize(n_batch);
|
|
for (int c = 0; c < n_batch; ++c) rad_damage_b_batch[c] = static_cast<float>(b[c]);
|
|
}
|
|
|
|
void RotationScaleMerge::RefineRelativeB(int n_groups) {
|
|
// RefineDecay removes the AVERAGE radiation-damage falloff as a single global relative-B slope, but the
|
|
// relative scattering power drifts NON-monotonically across a run (absorption path as the crystal
|
|
// rotates, slippage, dose bursts, a partly-shadowed wedge). XDS/AIMLESS model this with a per-BATCH
|
|
// relative-B on top of the per-batch scale. Here, on the fulls after RefineDecay: fit the smoothed
|
|
// per-batch relative-B (FitRelativeBCurve) and apply exp(b_batch * s^2). Same anti-overfit posture as
|
|
// RefineDecay - a physical spread floor and a cross-validated held-out gain - but the CV split is by ASU
|
|
// GROUP parity, not frame parity: a per-batch parameter owns whole frames, so it cannot be scored on a
|
|
// held-out FRAME; splitting the equivalents instead tests whether a batch's B generalises to reflections
|
|
// it was not fit on, which is exactly the over-fit risk of a many-parameter model.
|
|
if (fulls.empty() || n_frames <= 0 || relative_b_deg <= 0.0)
|
|
return;
|
|
const auto gon = x.GetGoniometer();
|
|
const double osc_deg = gon ? std::fabs(gon->GetIncrement_deg()) : 0.0;
|
|
if (!(osc_deg > 1e-6))
|
|
return;
|
|
const int frames_per_batch = std::max(1, static_cast<int>(std::lround(relative_b_deg / osc_deg)));
|
|
const int n_batch = std::max(1, (n_frames + frames_per_batch - 1) / frames_per_batch);
|
|
if (n_batch < 2) // a single batch is a global Wilson-B (degenerate with overall scale) - nothing relative
|
|
return;
|
|
auto s2_of = [](float d) { return d > 0.0f ? 1.0 / (4.0 * static_cast<double>(d) * d) : 0.0; };
|
|
auto usable = [&](const Obs &o) {
|
|
return o.group >= 0 && o.corr > 0.0f && std::isfinite(o.corr) && (o.d > 0.0f)
|
|
&& o.partiality >= min_partiality;
|
|
};
|
|
auto batch_of = [&](const Obs &o) { return std::min(n_batch - 1, o.frame / frames_per_batch); };
|
|
auto b_factor = [&](const Obs &o, const std::vector<double> &b) {
|
|
return std::exp(b[batch_of(o)] * s2_of(o.d));
|
|
};
|
|
// Held-out disagreement over the group subset {group&1 == gparity} with `b` applied: the same
|
|
// sigma-INDEPENDENT Rmeas-like metric as RefineDecay (sum|Is-Iref|/sum|Iref|), so a relative-B cannot
|
|
// "win" by reshaping sigma. Reference built from the scored subset (no leakage).
|
|
auto subset_disagreement = [&](int gparity, const std::vector<double> &b) -> double {
|
|
std::vector<double> sw(n_groups, 0.0), swI(n_groups, 0.0);
|
|
for (const auto &o : fulls) {
|
|
if (!usable(o) || (o.group & 1) != gparity) continue;
|
|
const double f = b_factor(o, b);
|
|
const double Is = static_cast<double>(o.I) * o.corr * f, sc = static_cast<double>(o.sigma) * o.corr * f;
|
|
const double w = 1.0 / (sc * sc);
|
|
sw[o.group] += w; swI[o.group] += w * Is;
|
|
}
|
|
double num = 0.0, den = 0.0;
|
|
for (const auto &o : fulls) {
|
|
if (!usable(o) || (o.group & 1) != gparity || sw[o.group] <= 0.0) continue;
|
|
const double Iref = swI[o.group] / sw[o.group];
|
|
const double Is = static_cast<double>(o.I) * o.corr * b_factor(o, b);
|
|
if (!std::isfinite(Iref) || Iref <= 0.0) continue;
|
|
num += std::abs(Is - Iref); den += Iref;
|
|
}
|
|
return den > 0.0 ? num / den : 0.0;
|
|
};
|
|
|
|
// A relative-B whose peak-to-peak is below a couple A^2 is negligible; "correcting" it only spreads
|
|
// equivalents (same s^2, different batch) into extra scatter - net harmful, exactly like RefineDecay's
|
|
// sub-floor slope. (RELATIVE_B_MIN_SPREAD ~ DECAY_MIN_DELTA_B.)
|
|
constexpr double RELATIVE_B_MIN_SPREAD = 2.0; // A^2, minimum peak-to-peak relative-B to engage
|
|
const std::vector<double> b_all = FitRelativeBCurve(n_groups, n_batch, frames_per_batch, -1);
|
|
const double bmin = *std::min_element(b_all.begin(), b_all.end());
|
|
const double bmax = *std::max_element(b_all.begin(), b_all.end());
|
|
if (bmax - bmin < RELATIVE_B_MIN_SPREAD) {
|
|
logger.Info("Relative-B: negligible variation ({} batches, peak-to-peak {:.2f} A^2 < {:.1f}, skipped)",
|
|
n_batch, bmax - bmin, RELATIVE_B_MIN_SPREAD);
|
|
return;
|
|
}
|
|
// Cross-validate: fit on even ASU groups, score the held-out odd equivalents (and vice-versa). Real
|
|
// per-batch structure generalises across the equivalent split; over-fit per-batch noise does not.
|
|
const std::vector<double> zero(n_batch, 0.0);
|
|
const double base = subset_disagreement(0, zero) + subset_disagreement(1, zero);
|
|
const std::vector<double> b_even = FitRelativeBCurve(n_groups, n_batch, frames_per_batch, 0);
|
|
const std::vector<double> b_odd = FitRelativeBCurve(n_groups, n_batch, frames_per_batch, 1);
|
|
const double gain = base - (subset_disagreement(1, b_even) + subset_disagreement(0, b_odd));
|
|
if (!(gain > CV_MIN_RELATIVE_GAIN * base)) {
|
|
logger.Info("Relative-B: not cross-validated ({} batches, peak-to-peak {:.2f} A^2, held-out gain {:.1f}%, skipped)",
|
|
n_batch, bmax - bmin, 100.0 * gain / std::max(base, 1e-30));
|
|
return;
|
|
}
|
|
for (auto &o : fulls)
|
|
if (usable(o))
|
|
o.corr = static_cast<float>(o.corr * b_factor(o, b_all));
|
|
logger.Info("Relative-B: {} batches of ~{:.1f} deg, peak-to-peak {:.2f} A^2, held-out gain {:.1f}% (cross-validated)",
|
|
n_batch, relative_b_deg, bmax - bmin, 100.0 * gain / std::max(base, 1e-30));
|
|
}
|
|
|
|
void RotationScaleMerge::RefineAbsorption(int n_iter, int n_groups) {
|
|
// Absorption / path-length: a smooth multiplicative factor over the diffracted-beam direction in the
|
|
// goniometer (crystal) frame. Each full's lab diffracted direction (from its predicted detector
|
|
// position, flat-detector-perpendicular approximation) is de-rotated by the spindle into a
|
|
// goniometer-fixed frame, so a fixed crystal-frame direction is sampled at many spindle angles and its
|
|
// cell is over-determined - robust, unlike a raw detector-position x time grid. Negligible at hard
|
|
// X-rays / thin crystals (the factor stays ~1); the point is low-energy data where absorption is large.
|
|
const auto gon_opt = x.GetGoniometer();
|
|
if (!gon_opt.has_value() || fulls.empty())
|
|
return;
|
|
const GoniometerAxis &gon = *gon_opt;
|
|
const float beam_x = x.GetBeamX_pxl(), beam_y = x.GetBeamY_pxl();
|
|
const float F = x.GetDetectorDistance_mm() / x.GetPixelSize_mm();
|
|
if (!(F > 0.0f))
|
|
return;
|
|
|
|
constexpr int NB = 8;
|
|
std::vector<int32_t> cell(fulls.size(), -1);
|
|
for (size_t i = 0; i < fulls.size(); ++i) {
|
|
const Obs &o = fulls[i];
|
|
if (!std::isfinite(o.px) || !std::isfinite(o.py))
|
|
continue;
|
|
const Coord s1 = Coord((o.px - beam_x) / F, (o.py - beam_y) / F, 1.0f).Normalize();
|
|
const Coord u = gon.GetTransformationAngle(gon.GetAngle_deg(o.image_number)).transpose() * s1;
|
|
const int ix = std::clamp(static_cast<int>((u.x + 1.0f) * 0.5f * NB), 0, NB - 1);
|
|
const int iy = std::clamp(static_cast<int>((u.y + 1.0f) * 0.5f * NB), 0, NB - 1);
|
|
cell[i] = ix * NB + iy;
|
|
}
|
|
ApplyCellSurface(cell, NB * NB, n_iter, n_groups, "Absorption (goniometer-frame 8x8)");
|
|
}
|
|
|
|
void RotationScaleMerge::RefineModulation(int n_iter, int n_groups) {
|
|
// Detector-plane modulation (flat-field): a smooth multiplicative factor over the position where a
|
|
// reflection lands on the detector (predicted px, py). It absorbs detector-response and geometric
|
|
// systematics that vary across the detector plane; symmetry-equivalents of one reflection land at
|
|
// different positions as the crystal rotates, over-determining the surface. Because it is a property of
|
|
// the detector + beam (not the rotation), the same correction concept transfers to stills.
|
|
if (fulls.empty())
|
|
return;
|
|
float pxmin = std::numeric_limits<float>::infinity(), pxmax = -pxmin, pymin = pxmin, pymax = -pxmin;
|
|
for (const Obs &o : fulls) {
|
|
if (!std::isfinite(o.px) || !std::isfinite(o.py)) continue;
|
|
pxmin = std::min(pxmin, o.px); pxmax = std::max(pxmax, o.px);
|
|
pymin = std::min(pymin, o.py); pymax = std::max(pymax, o.py);
|
|
}
|
|
if (!(pxmax > pxmin) || !(pymax > pymin))
|
|
return;
|
|
constexpr int NB = 16;
|
|
const float sx = NB / (pxmax - pxmin), sy = NB / (pymax - pymin);
|
|
std::vector<int32_t> cell(fulls.size(), -1);
|
|
for (size_t i = 0; i < fulls.size(); ++i) {
|
|
const Obs &o = fulls[i];
|
|
if (!std::isfinite(o.px) || !std::isfinite(o.py)) continue;
|
|
const int ix = std::clamp(static_cast<int>((o.px - pxmin) * sx), 0, NB - 1);
|
|
const int iy = std::clamp(static_cast<int>((o.py - pymin) * sy), 0, NB - 1);
|
|
cell[i] = ix * NB + iy;
|
|
}
|
|
ApplyCellSurface(cell, NB * NB, n_iter, n_groups, "Modulation (detector-frame 16x16)");
|
|
}
|
|
|
|
void RotationScaleMerge::ApplyCellSurface(const std::vector<int32_t> &cell, int ncell, int n_iter,
|
|
int n_groups, const char *name) {
|
|
auto usable = [&](const Obs &o) {
|
|
return o.group >= 0 && o.corr > 0.0f && std::isfinite(o.corr) && o.partiality >= min_partiality;
|
|
};
|
|
// Fit the per-cell factor over the subset {frame&1 == parity} (parity < 0 = all fulls), n_iter
|
|
// alternating rounds against that subset's own reference (Tikhonov pull to 1, gauge-fixed to a
|
|
// den-weighted geometric mean of 1 so it never drifts the overall scale). Returns the per-cell factor.
|
|
auto fit_surface = [&](int parity) -> std::vector<double> {
|
|
std::vector<double> A(ncell, 1.0);
|
|
for (int it = 0; it < n_iter; ++it) {
|
|
std::vector<double> sw(n_groups, 0.0), swI(n_groups, 0.0);
|
|
for (size_t i = 0; i < fulls.size(); ++i) {
|
|
const Obs &o = fulls[i];
|
|
if (!usable(o) || cell[i] < 0 || (parity >= 0 && (o.frame & 1) != parity)) continue;
|
|
const double a = A[cell[i]], sc = static_cast<double>(o.sigma) * o.corr * a, w = 1.0 / (sc * sc);
|
|
sw[o.group] += w; swI[o.group] += w * static_cast<double>(o.I) * o.corr * a;
|
|
}
|
|
std::vector<double> num(ncell, 0.0), den(ncell, 0.0);
|
|
for (size_t i = 0; i < fulls.size(); ++i) {
|
|
const Obs &o = fulls[i];
|
|
if (!usable(o) || cell[i] < 0 || (parity >= 0 && (o.frame & 1) != parity) || sw[o.group] <= 0.0)
|
|
continue;
|
|
const double Iref = swI[o.group] / sw[o.group], a = A[cell[i]];
|
|
const double Is = static_cast<double>(o.I) * o.corr * a, sc = static_cast<double>(o.sigma) * o.corr * a;
|
|
if (!std::isfinite(Iref) || Iref <= 0.0 || !(Is > 0.0) || !(sc > 0.0)) continue;
|
|
const double w = 1.0 / (sc * sc);
|
|
num[cell[i]] += w * Is * Iref; den[cell[i]] += w * Is * Is;
|
|
}
|
|
std::vector<double> dsorted = den;
|
|
std::nth_element(dsorted.begin(), dsorted.begin() + dsorted.size() / 2, dsorted.end());
|
|
const double lambda = 0.1 * std::max(1e-30, dsorted[dsorted.size() / 2]);
|
|
std::vector<double> upd(ncell, 1.0);
|
|
for (int c = 0; c < ncell; ++c) upd[c] = (num[c] + lambda) / (den[c] + lambda);
|
|
double logsum = 0.0, wsum = 0.0;
|
|
for (int c = 0; c < ncell; ++c) if (den[c] > 0.0) { logsum += den[c] * std::log(upd[c]); wsum += den[c]; }
|
|
const double gm = wsum > 0.0 ? std::exp(logsum / wsum) : 1.0;
|
|
for (int c = 0; c < ncell; ++c) A[c] = std::clamp(A[c] * upd[c] / gm, 0.25, 4.0);
|
|
}
|
|
return A;
|
|
};
|
|
// Mean studentized squared deviation over the subset {frame&1 == parity} with surface A applied,
|
|
// scored against that subset's own reference (no leakage). Lower = tighter equivalents.
|
|
auto score = [&](int parity, const std::vector<double> &A) -> double {
|
|
std::vector<double> sw(n_groups, 0.0), swI(n_groups, 0.0);
|
|
for (size_t i = 0; i < fulls.size(); ++i) {
|
|
const Obs &o = fulls[i];
|
|
if (!usable(o) || cell[i] < 0 || (o.frame & 1) != parity) continue;
|
|
const double a = A[cell[i]], Is = static_cast<double>(o.I) * o.corr * a;
|
|
const double sc = static_cast<double>(o.sigma) * o.corr * a, w = 1.0 / (sc * sc);
|
|
sw[o.group] += w; swI[o.group] += w * Is;
|
|
}
|
|
// Rmeas-like (sigma-INDEPENDENT) agreement of the held-out equivalents: sum|Is - Iref| / sum|Iref|.
|
|
// Scoring on the studentized deviation instead lets a surface "improve" the held-out chi^2 by
|
|
// reshaping sigma (via corr) without tightening the actual intensities - which on mis-indexed / bad
|
|
// data passes cross-validation yet worsens Rmeas. A fractional metric cannot be gamed that way.
|
|
double num = 0.0, den = 0.0;
|
|
for (size_t i = 0; i < fulls.size(); ++i) {
|
|
const Obs &o = fulls[i];
|
|
if (!usable(o) || cell[i] < 0 || (o.frame & 1) != parity || sw[o.group] <= 0.0) continue;
|
|
const double a = A[cell[i]], Is = static_cast<double>(o.I) * o.corr * a;
|
|
const double Iref = swI[o.group] / sw[o.group];
|
|
if (!std::isfinite(Iref) || Iref <= 0.0) continue;
|
|
num += std::abs(Is - Iref); den += Iref;
|
|
}
|
|
return den > 0.0 ? num / den : 0.0;
|
|
};
|
|
|
|
// Cross-validate: fit the surface on even frames and score the held-out odd equivalents (and vice
|
|
// versa). A real surface generalizes; an over-fit one (few obs / cell) does not. Apply only if the
|
|
// held-out equivalent chi^2 improves by a clear margin - so the correction can never worsen the data.
|
|
const std::vector<double> ident(ncell, 1.0);
|
|
const std::vector<double> A_even = fit_surface(0), A_odd = fit_surface(1);
|
|
const double base = score(1, ident) + score(0, ident);
|
|
const double gain = base - (score(1, A_even) + score(0, A_odd));
|
|
if (!(gain > CV_MIN_RELATIVE_GAIN * base)) {
|
|
logger.Info("{} correction: not cross-validated (held-out gain {:.1f}%, skipped)",
|
|
name, 100.0 * gain / std::max(base, 1e-30));
|
|
return;
|
|
}
|
|
const std::vector<double> A = fit_surface(-1);
|
|
for (size_t i = 0; i < fulls.size(); ++i)
|
|
if (cell[i] >= 0)
|
|
fulls[i].corr = static_cast<float>(fulls[i].corr * A[cell[i]]);
|
|
logger.Info("{} correction: cross-validated, held-out gain {:.1f}%",
|
|
name, 100.0 * gain / std::max(base, 1e-30));
|
|
}
|
|
|
|
void RotationScaleMerge::ComputeSmoothGWindow(const std::vector<double> &g, int window,
|
|
std::vector<double> &g_smooth) const {
|
|
const int n = static_cast<int>(g.size());
|
|
const int half = window / 2;
|
|
g_smooth.assign(n, NAN);
|
|
|
|
// The run's typical scale, and the floor below which a fitted scale is not credible
|
|
// (MIN_CREDIBLE_SCALE_RATIO). The window average is a GEOMETRIC mean, so a scale collapsing toward
|
|
// zero does not just corrupt its own frame - its logarithm drags the whole window down. Where a run
|
|
// has a stretch of frames too sparse to fit, the only fitted frames in a window can BE the collapsed
|
|
// ones, and the mean then averages the fault with itself and makes it permanent.
|
|
std::vector<double> fitted;
|
|
fitted.reserve(n);
|
|
for (int j = 0; j < n; ++j)
|
|
if (frame_scaled_scratch[j] && std::isfinite(g[j]) && g[j] > 0.0)
|
|
fitted.push_back(g[j]);
|
|
double g_typ = NAN;
|
|
if (!fitted.empty()) {
|
|
const size_t mid = fitted.size() / 2;
|
|
std::nth_element(fitted.begin(), fitted.begin() + mid, fitted.end());
|
|
g_typ = fitted[mid];
|
|
}
|
|
const double g_floor = std::isfinite(g_typ) ? g_typ * MIN_CREDIBLE_SCALE_RATIO : 0.0;
|
|
|
|
for (int o = 0; o < n; ++o) {
|
|
double sum_log = 0.0;
|
|
int count = 0;
|
|
for (int j = std::max(0, o - half); j <= std::min(n - 1, o + half); ++j) {
|
|
if (frame_scaled_scratch[j] && std::isfinite(g[j]) && g[j] >= g_floor && g[j] > 0.0) {
|
|
sum_log += std::log(g[j]);
|
|
++count;
|
|
}
|
|
}
|
|
if (count > 0) g_smooth[o] = std::exp(sum_log / count);
|
|
// A frame whose own scale is not credible takes the neighbourhood's, or - when the neighbourhood
|
|
// holds nothing credible either - the run's typical scale. Callers divide g by g_smooth, so this
|
|
// replaces the collapsed scale rather than layering another correction on top of it.
|
|
if (frame_scaled_scratch[o] && std::isfinite(g[o]) && g[o] > 0.0 && g[o] < g_floor
|
|
&& !std::isfinite(g_smooth[o]))
|
|
g_smooth[o] = g_typ;
|
|
}
|
|
}
|
|
|
|
void RotationScaleMerge::SmoothG(std::vector<Obs> &obs, std::vector<double> &g, int window) const {
|
|
const int n = static_cast<int>(g.size());
|
|
std::vector<double> g_smooth;
|
|
ComputeSmoothGWindow(g, window, g_smooth);
|
|
for (auto &o : obs) {
|
|
const int f = o.frame;
|
|
if (!frame_scaled_scratch[f] || !std::isfinite(g[f]) || g[f] <= 0.0 || !std::isfinite(g_smooth[f]))
|
|
continue;
|
|
if (std::isfinite(o.corr))
|
|
o.corr = static_cast<float>(o.corr * (g[f] / g_smooth[f]));
|
|
}
|
|
for (int f = 0; f < n; ++f)
|
|
if (frame_scaled_scratch[f] && std::isfinite(g[f]) && g[f] > 0.0 && std::isfinite(g_smooth[f]))
|
|
g[f] = g_smooth[f];
|
|
}
|
|
|
|
bool RotationScaleMerge::RejectCollapsedFullScales() {
|
|
// The Unity model leaves corr = 1/G, and every full of a frame carries that frame's G.
|
|
std::vector<double> g_frame(n_frames, NAN);
|
|
for (const auto &o : fulls)
|
|
if (o.frame >= 0 && o.frame < n_frames && std::isfinite(o.corr) && o.corr > 0.0f)
|
|
g_frame[o.frame] = 1.0 / static_cast<double>(o.corr);
|
|
|
|
std::vector<double> fitted;
|
|
fitted.reserve(g_frame.size());
|
|
for (const double gf : g_frame)
|
|
if (std::isfinite(gf) && gf > 0.0)
|
|
fitted.push_back(gf);
|
|
if (fitted.empty())
|
|
return false;
|
|
const size_t mid = fitted.size() / 2;
|
|
std::nth_element(fitted.begin(), fitted.begin() + mid, fitted.end());
|
|
const double g_floor = fitted[mid] * MIN_CREDIBLE_SCALE_RATIO;
|
|
|
|
int n_rejected = 0;
|
|
double worst = 1.0;
|
|
for (auto &o : fulls) {
|
|
if (o.frame < 0 || o.frame >= n_frames)
|
|
continue;
|
|
const double gf = g_frame[o.frame];
|
|
if (!std::isfinite(gf) || gf >= g_floor)
|
|
continue;
|
|
worst = std::max(worst, 1.0 / gf);
|
|
o.corr = 1.0f; // as if the frame's scale had never been fitted
|
|
++n_rejected;
|
|
}
|
|
if (n_rejected > 0)
|
|
logger.Warning("Rejected the fitted scale of {} full(s) whose frame scaled to less than {:.0f}x "
|
|
"below the run median (worst {:.0f}x amplification); those frames merge unscaled",
|
|
n_rejected, 1.0 / MIN_CREDIBLE_SCALE_RATIO, worst / fitted[mid]);
|
|
return n_rejected > 0;
|
|
}
|
|
|
|
void RotationScaleMerge::Combine() {
|
|
fulls.clear();
|
|
g_full.assign(n_frames, 1.0);
|
|
|
|
// Combine one raw-hkl run into fulls, appended to `out`. Independent per run (the events of one hkl
|
|
// touch no shared state), so runs parallelise cleanly. Returns the number of usable partials seen.
|
|
// `dump` (serial path only) writes each emitted full for the diagnostic observation dump.
|
|
auto process_rawrun = [&](int r, std::vector<Obs> &out, std::ofstream *dump) -> size_t {
|
|
const int lo = rawrun_start[r], hi = rawrun_start[r] + rawrun_count[r];
|
|
std::vector<int32_t> ev; // usable perm-indices of this raw hkl, in image-number order
|
|
ev.reserve(hi - lo);
|
|
for (int p = lo; p < hi; ++p) {
|
|
const auto &o = partials[perm[p]];
|
|
if (!std::isfinite(o.corr) || o.corr <= 0.0f) continue;
|
|
if (!std::isfinite(o.I) || !std::isfinite(o.sigma) || o.sigma <= 0.0f) continue;
|
|
ev.push_back(perm[p]);
|
|
}
|
|
const int group = rawrun_group[r];
|
|
|
|
size_t i = 0;
|
|
while (i < ev.size()) {
|
|
size_t kk = i + 1;
|
|
float last_frame = partials[ev[i]].image_number;
|
|
while (kk < ev.size()) {
|
|
const float frame = partials[ev[kk]].image_number;
|
|
if (frame - last_frame > MAX_FRAME_GAP) break;
|
|
last_frame = frame;
|
|
++kk;
|
|
}
|
|
|
|
double pooled_bkg = 0.0;
|
|
int n_pool = 0;
|
|
for (size_t m = i; m < kk; ++m) {
|
|
const float b = partials[ev[m]].bkg;
|
|
if (std::isfinite(b)) { pooled_bkg += b; ++n_pool; }
|
|
}
|
|
pooled_bkg = n_pool > 0 ? pooled_bkg / n_pool : 0.0;
|
|
auto pooled_I = [&](const Obs &r2) {
|
|
const double n_bkg = std::max(0.0, static_cast<double>(r2.sigma) * r2.sigma - r2.I)
|
|
/ std::max(r2.bkg, 1.0f);
|
|
return static_cast<double>(r2.I) + n_bkg * (static_cast<double>(r2.bkg) - pooled_bkg);
|
|
};
|
|
|
|
double sum_w = 0.0, sum_wI = 0.0, sum_partiality = 0.0;
|
|
float d = NAN;
|
|
int peak_outcome = partials[ev[i]].frame;
|
|
float peak_frame = partials[ev[i]].image_number;
|
|
float peak_px = partials[ev[i]].px, peak_py = partials[ev[i]].py;
|
|
float peak_partiality = -1.0f;
|
|
const bool on_ice = partials[ev[i]].on_ice;
|
|
for (size_t m = i; m < kk; ++m) {
|
|
const auto &r2 = partials[ev[m]];
|
|
const double sigma_corr = static_cast<double>(r2.sigma) * r2.corr;
|
|
const double w = 1.0 / (sigma_corr * sigma_corr);
|
|
sum_w += w;
|
|
sum_wI += w * pooled_I(r2) * r2.corr;
|
|
sum_partiality += r2.partiality;
|
|
if (r2.partiality > peak_partiality) {
|
|
peak_partiality = r2.partiality;
|
|
peak_outcome = r2.frame;
|
|
peak_frame = r2.image_number;
|
|
peak_px = r2.px; peak_py = r2.py;
|
|
}
|
|
if (!std::isfinite(d) && std::isfinite(r2.d) && r2.d > 0.0f) d = r2.d;
|
|
}
|
|
double F = sum_wI / sum_w;
|
|
|
|
for (int iter = 0; iter < 3; ++iter) {
|
|
sum_w = 0.0; sum_wI = 0.0;
|
|
for (size_t m = i; m < kk; ++m) {
|
|
const auto &r2 = partials[ev[m]];
|
|
const double corr = r2.corr;
|
|
const double I_corr = pooled_I(r2) * corr;
|
|
const double sigma_corr = static_cast<double>(r2.sigma) * corr;
|
|
const double bkg_var = sigma_corr * sigma_corr - corr * I_corr;
|
|
double var = std::max(0.0, bkg_var) + corr * std::max(0.0, F);
|
|
if (!(var > 0.0)) var = sigma_corr * sigma_corr;
|
|
const double w = 1.0 / var;
|
|
sum_w += w;
|
|
sum_wI += w * I_corr;
|
|
}
|
|
F = sum_wI / sum_w;
|
|
}
|
|
|
|
const int n_frames_event = static_cast<int>(kk - i);
|
|
i = kk;
|
|
|
|
if (sum_w <= 0.0 || sum_partiality < min_partiality || sum_partiality < min_captured_fraction)
|
|
continue;
|
|
|
|
double sigma_full = 1.0 / std::sqrt(sum_w);
|
|
if (capture_uncertainty_coeff > 0.0) {
|
|
const double frac = std::min(1.0, sum_partiality);
|
|
const double extra = capture_uncertainty_coeff * (1.0 - frac) * std::max(0.0, F);
|
|
sigma_full = std::sqrt(sigma_full * sigma_full + extra * extra);
|
|
}
|
|
|
|
Obs full{};
|
|
full.h = rawrun_h[r]; full.k = rawrun_k[r]; full.l = rawrun_l[r];
|
|
full.I = static_cast<float>(F);
|
|
full.sigma = static_cast<float>(sigma_full);
|
|
full.d = d;
|
|
full.rlp = 1.0f;
|
|
full.partiality = 1.0f;
|
|
full.corr = 1.0f;
|
|
full.image_number = peak_frame;
|
|
full.px = peak_px; full.py = peak_py;
|
|
full.frame = peak_outcome;
|
|
full.on_ice = on_ice ? 1 : 0;
|
|
full.group = group; // the raw hkl's ASU group for the current space group (<0 = absent)
|
|
out.push_back(full);
|
|
|
|
if (dump != nullptr)
|
|
*dump << full.h << ' ' << full.k << ' ' << full.l << ' ' << full.I << ' '
|
|
<< full.sigma << ' ' << d << ' ' << n_frames_event << ' ' << sum_partiality << ' '
|
|
<< static_cast<int>(peak_frame) << ' ' << full.px << ' ' << full.py << '\n';
|
|
}
|
|
return ev.size();
|
|
};
|
|
|
|
const int n_run = static_cast<int>(rawrun_start.size());
|
|
size_t n_used = 0;
|
|
|
|
if (!observation_dump_path.empty() || nthreads <= 1) {
|
|
// Serial (diagnostic dump needs a single writer).
|
|
std::ofstream dump;
|
|
if (!observation_dump_path.empty()) {
|
|
dump.open(observation_dump_path);
|
|
// Columns are enough to rescale/merge the combined fulls in an external program: peak_frame
|
|
// is a per-full batch (ZD) and px/py the predicted detector position (XD/YD).
|
|
dump << "# h k l I sigma d n_frames captured_fraction peak_frame px py\n";
|
|
}
|
|
fulls.reserve(n_run);
|
|
for (int r = 0; r < n_run; ++r)
|
|
n_used += process_rawrun(r, fulls, dump.is_open() ? &dump : nullptr);
|
|
} else {
|
|
// Parallel over contiguous rawrun chunks; concatenate the per-thread fulls in run order so the
|
|
// result is deterministic.
|
|
const int nt = static_cast<int>(std::min(nthreads, static_cast<size_t>(n_run)));
|
|
const int chunk = (n_run + nt - 1) / nt;
|
|
std::vector<std::vector<Obs>> part(nt);
|
|
std::vector<size_t> used(nt, 0);
|
|
std::vector<std::future<void>> futures;
|
|
futures.reserve(nt);
|
|
for (int t = 0; t < nt; ++t) {
|
|
const int r0 = t * chunk, r1 = std::min(n_run, r0 + chunk);
|
|
if (r0 >= r1) break;
|
|
futures.emplace_back(std::async(std::launch::async, [&, t, r0, r1] {
|
|
part[t].reserve(r1 - r0);
|
|
for (int r = r0; r < r1; ++r) used[t] += process_rawrun(r, part[t], nullptr);
|
|
}));
|
|
}
|
|
for (auto &f : futures) f.get();
|
|
size_t total = 0;
|
|
for (int t = 0; t < nt; ++t) { total += part[t].size(); n_used += used[t]; }
|
|
fulls.reserve(total);
|
|
for (int t = 0; t < nt; ++t)
|
|
fulls.insert(fulls.end(), part[t].begin(), part[t].end());
|
|
}
|
|
|
|
SortFullsByFrame();
|
|
logger.Info("3D combine: {} fulls from {} partials", fulls.size(), n_used);
|
|
}
|
|
|
|
void RotationScaleMerge::SortFullsByFrame() {
|
|
// Sort the fulls by their (peak) frame and build per-frame CSR ranges, so the scale-fulls step can
|
|
// fit a per-frame G by slicing contiguous ranges (the same layout the partials use).
|
|
std::sort(fulls.begin(), fulls.end(),
|
|
[](const Obs &a, const Obs &b) { return a.frame < b.frame; });
|
|
fulls_frame_start.assign(n_frames, 0);
|
|
fulls_frame_count.assign(n_frames, 0);
|
|
for (int i = 0; i < static_cast<int>(fulls.size()); ) {
|
|
const int f = fulls[i].frame;
|
|
fulls_frame_start[f] = i;
|
|
int j = i;
|
|
while (j < static_cast<int>(fulls.size()) && fulls[j].frame == f) ++j;
|
|
fulls_frame_count[f] = j - i;
|
|
i = j;
|
|
}
|
|
}
|
|
|
|
void RotationScaleMerge::ComputePerFrameCC(const std::vector<double> &partial_group_mean,
|
|
std::vector<double> &cc, std::vector<int64_t> &cc_n) const {
|
|
// Per-frame CC vs the merged reference (CalculateGlobalCC), computed once now (not every iteration).
|
|
cc.assign(n_frames, NAN);
|
|
cc_n.assign(n_frames, 0);
|
|
ParallelFor(n_frames, nthreads, [&](int f) {
|
|
double sx = 0, sy = 0, sx2 = 0, sy2 = 0, sxy = 0;
|
|
size_t n = 0;
|
|
const int lo = frame_start[f], hi = frame_start[f] + frame_count[f];
|
|
for (int i = lo; i < hi; ++i) {
|
|
const auto &o = partials[i];
|
|
if (o.on_ice) continue;
|
|
if (o.group < 0) continue;
|
|
if (o.partiality < min_partiality) continue;
|
|
if (!std::isfinite(o.I) || !std::isfinite(o.corr) || o.corr <= 0.0f) continue;
|
|
if (!std::isfinite(o.sigma) || o.sigma <= 0.0f) continue;
|
|
const double mean = partial_group_mean[o.group];
|
|
if (!std::isfinite(mean)) continue;
|
|
const double img = static_cast<double>(o.I) * o.corr;
|
|
sx += img; sy += mean; sx2 += img * img; sy2 += mean * mean; sxy += img * mean;
|
|
++n;
|
|
}
|
|
if (n < MIN_REFLECTIONS) return;
|
|
const double nd = static_cast<double>(n);
|
|
const double cov = sxy - sx * sy / nd;
|
|
const double vx = sx2 - sx * sx / nd;
|
|
const double vy = sy2 - sy * sy / nd;
|
|
if (vx > 0.0 && vy > 0.0) { cc[f] = cov / std::sqrt(vx * vy); cc_n[f] = static_cast<int64_t>(n); }
|
|
});
|
|
}
|
|
|
|
// Write the per-frame G / CC / mosaicity (from the given cc/cc_n) back onto the partials for the offline
|
|
// per-image scaling table. cc/cc_n are computed on the host (ComputePerFrameCC) or GPU (ComputePartialCC).
|
|
void RotationScaleMerge::FinalizePerFrameScale(const std::vector<double> &cc, const std::vector<int64_t> &cc_n,
|
|
const std::vector<uint8_t> &frame_scaled) {
|
|
for (int f = 0; f < n_frames; ++f) {
|
|
auto &o = partials_out[f];
|
|
if (frame_scaled[f]) {
|
|
o.image_scale_g = static_cast<float>(g_partial[f]);
|
|
o.mosaicity_deg = (f < static_cast<int>(mos_smooth.size()) && std::isfinite(mos_smooth[f]))
|
|
? mos_smooth[f] : static_cast<float>(mosaicity_deg);
|
|
if (std::isfinite(cc[f])) { o.image_scale_cc = static_cast<float>(cc[f]); o.image_scale_cc_n = cc_n[f]; }
|
|
else { o.image_scale_cc.reset(); o.image_scale_cc_n.reset(); }
|
|
} else {
|
|
o.image_scale_g.reset();
|
|
o.image_scale_cc.reset();
|
|
o.image_scale_cc_n.reset();
|
|
o.mosaicity_deg.reset();
|
|
}
|
|
o.image_scale_wedge_deg.reset();
|
|
}
|
|
}
|
|
|
|
namespace {
|
|
// Possible unique reflections per shell for the completeness column - mirrors CalcPossibleReflections.
|
|
void PossiblePerShell(int space_group_number, const UnitCell &cell, double d_min, double d_max,
|
|
const ResolutionShells &shells, bool merge_friedel, std::vector<int> &possible) {
|
|
gemmi::UnitCell gemmi_cell = cell;
|
|
const gemmi::SpaceGroup *sg = gemmi::find_spacegroup_by_number(space_group_number);
|
|
if (sg == nullptr) return;
|
|
const std::vector<gemmi::Miller> hkls = gemmi::make_miller_vector(gemmi_cell, sg, d_min, d_max, true);
|
|
const gemmi::GroupOps gops = sg->operations();
|
|
CrystalLattice lattice(cell);
|
|
const auto astar = lattice.Astar(), bstar = lattice.Bstar(), cstar = lattice.Cstar();
|
|
for (const auto &hkl : hkls) {
|
|
const auto q = hkl[0] * astar + hkl[1] * bstar + hkl[2] * cstar;
|
|
const auto qlen = q.Length();
|
|
if (qlen < 1e-6) continue;
|
|
const auto shell = shells.GetShell(1.0 / qlen);
|
|
if (!shell.has_value()) continue;
|
|
const int s = *shell;
|
|
if (s >= 0 && s < static_cast<int>(possible.size()))
|
|
possible[s] += (merge_friedel || gops.is_reflection_centric(hkl)) ? 1 : 2;
|
|
}
|
|
}
|
|
}
|
|
|
|
RotationScaleMerge::Result RotationScaleMerge::MergeAndStats(int n_groups, bool for_search,
|
|
const std::vector<char> &masked,
|
|
bool fulls_resident) {
|
|
// A full is usable for the merge / error model if it passes AddImage's filters (with the current
|
|
// ice/masked-ring context). group >= 0 already encodes "not absent and passes AcceptReflection".
|
|
auto masked_ring = [&](const Obs &o) {
|
|
if (masked.empty()) return false;
|
|
const int ring = IceRingIndex(o.d, ice_half_width_q);
|
|
return ring >= 0 && ring < static_cast<int>(masked.size()) && masked[ring];
|
|
};
|
|
auto usable_merge = [&](const Obs &o) {
|
|
if (o.group < 0) return false;
|
|
if (!frame_cell_ok[o.frame]) return false;
|
|
if (!(o.corr > 0.0f) || !std::isfinite(o.corr)) return false;
|
|
if (for_search && o.on_ice) return false;
|
|
if (masked_ring(o)) return false;
|
|
if (o.partiality < min_partiality) return false;
|
|
const float I_corr = o.I * o.corr, sigma_corr = o.sigma * o.corr;
|
|
return std::isfinite(I_corr) && std::isfinite(sigma_corr) && sigma_corr > 0.0f;
|
|
};
|
|
|
|
// The em-stats / samples / merge-accumulate / R_meas reductions run on the resident, scaled fulls
|
|
// (their group CSR is still on the device from scale-fulls) when fulls_resident; the host keeps the
|
|
// I2-sort, the (a,b) fit, the export and the statistics. reject_outliers is excluded upstream.
|
|
bool use_gpu_merge = false;
|
|
#ifdef JFJOCH_USE_CUDA
|
|
use_gpu_merge = fulls_resident && !fulls.empty();
|
|
std::vector<uint8_t> gpu_masked(masked.begin(), masked.end());
|
|
#endif
|
|
|
|
// ---- Error model: fit dev2 = a*sigma^2 + b^2*<I>^2 from symmetry-equivalent scatter. ----
|
|
std::vector<double> em_mean(n_groups, NAN);
|
|
std::vector<float> reject_median(n_groups, NAN);
|
|
double error_model_a = 1.0, error_model_b = 0.0, error_model_chi2 = 0.0;
|
|
bool error_model_active = false;
|
|
{
|
|
struct Sample { double s2, I2, dev2; };
|
|
std::vector<Sample> samples;
|
|
std::vector<int32_t> cnt(n_groups, 0); // per-group usable count (both paths; feeds reject-median)
|
|
bool did_gpu = false;
|
|
#ifdef JFJOCH_USE_CUDA
|
|
if (use_gpu_merge) {
|
|
const int nf = static_cast<int>(fulls.size());
|
|
std::vector<double> gs2(nf), gI2(nf), gdev2(nf);
|
|
std::vector<uint8_t> gvalid(nf);
|
|
gpu_->MergeEmSamples(for_search, gpu_masked.data(), static_cast<int>(gpu_masked.size()),
|
|
ice_half_width_q, min_partiality, em_mean.data(), cnt.data(),
|
|
gs2.data(), gI2.data(), gdev2.data(), gvalid.data());
|
|
samples.reserve(nf);
|
|
for (int i = 0; i < nf; ++i)
|
|
if (gvalid[i]) samples.push_back({gs2[i], gI2[i], gdev2[i]});
|
|
did_gpu = true;
|
|
}
|
|
#endif
|
|
if (!did_gpu) {
|
|
// Per-group inverse-variance mean over usable fulls (>=2 obs), and the leverage-corrected samples.
|
|
std::vector<double> sw(n_groups, 0.0), swI(n_groups, 0.0);
|
|
for (const auto &o : fulls) {
|
|
if (!usable_merge(o)) continue;
|
|
const double sigma_corr = static_cast<double>(o.sigma) * o.corr;
|
|
const double w = 1.0 / (sigma_corr * sigma_corr);
|
|
sw[o.group] += w; swI[o.group] += w * (static_cast<double>(o.I) * o.corr); cnt[o.group]++;
|
|
}
|
|
for (int g = 0; g < n_groups; ++g)
|
|
if (cnt[g] >= 2 && sw[g] > 0.0) em_mean[g] = swI[g] / sw[g];
|
|
|
|
samples.reserve(fulls.size());
|
|
for (const auto &o : fulls) {
|
|
if (!usable_merge(o) || cnt[o.group] < 2) continue;
|
|
const double mean = em_mean[o.group];
|
|
if (!std::isfinite(mean)) continue;
|
|
const double sigma_corr = static_cast<double>(o.sigma) * o.corr;
|
|
const double s2 = sigma_corr * sigma_corr;
|
|
const double w = 1.0 / s2;
|
|
const double factor = 1.0 - w / sw[o.group];
|
|
if (factor < 0.05) continue;
|
|
const double resid = static_cast<double>(o.I) * o.corr - mean;
|
|
samples.push_back({s2, mean * mean, resid * resid / factor});
|
|
}
|
|
}
|
|
|
|
// Per-group outlier-rejection median of I*corr (host both paths - a per-group median is awkward on
|
|
// the GPU; cheap here, cnt >= 2 filter from the em pass). Fed to the merge accumulate.
|
|
if (reject_outliers) {
|
|
std::vector<std::vector<float>> iv(n_groups);
|
|
for (const auto &o : fulls)
|
|
if (usable_merge(o) && cnt[o.group] >= 2)
|
|
iv[o.group].push_back(o.I * o.corr);
|
|
for (int g = 0; g < n_groups; ++g)
|
|
if (!iv[g].empty()) {
|
|
std::nth_element(iv[g].begin(), iv[g].begin() + iv[g].size() / 2, iv[g].end());
|
|
reject_median[g] = iv[g][iv[g].size() / 2];
|
|
}
|
|
}
|
|
constexpr int n_bins = 16;
|
|
// Fit (a, b) from the intensity-binned median deviations. Factored into a lambda so it can be
|
|
// re-run on a misfit-free pool below; takes the samples by value (it sorts them in place).
|
|
auto fit_ab = [&](std::vector<Sample> smp) {
|
|
if (smp.size() < static_cast<size_t>(8 * n_bins))
|
|
return;
|
|
std::sort(smp.begin(), smp.end(), [](const Sample &a, const Sample &b) { return a.I2 < b.I2; });
|
|
std::vector<double> bs2, bI2, bd2;
|
|
const size_t per = smp.size() / n_bins;
|
|
for (int bin = 0; bin < n_bins; ++bin) {
|
|
const size_t lo = bin * per;
|
|
const size_t hi = (bin == n_bins - 1) ? smp.size() : lo + per;
|
|
std::vector<double> vs2, vI2, vd2;
|
|
for (size_t i = lo; i < hi; ++i) {
|
|
vs2.push_back(smp[i].s2); vI2.push_back(smp[i].I2); vd2.push_back(smp[i].dev2);
|
|
}
|
|
bs2.push_back(median_of(vs2));
|
|
bI2.push_back(median_of(vI2));
|
|
bd2.push_back(median_of(vd2) / CHI2_1_MEDIAN);
|
|
}
|
|
std::vector<double> bd2_sorted = bd2;
|
|
const double dev2_floor = std::max(1e-30, 1e-3 * median_of(bd2_sorted));
|
|
double Ass = 0, AsI = 0, AII = 0, Bs = 0, BI = 0;
|
|
for (int bin = 0; bin < n_bins; ++bin) {
|
|
const double s2 = bs2[bin], I2 = bI2[bin], d2 = bd2[bin];
|
|
const double d2w = std::max(d2, dev2_floor);
|
|
const double wgt = 1.0 / (d2w * d2w);
|
|
Ass += wgt * s2 * s2; AsI += wgt * s2 * I2; AII += wgt * I2 * I2;
|
|
Bs += wgt * s2 * d2; BI += wgt * I2 * d2;
|
|
}
|
|
const double det = Ass * AII - AsI * AsI;
|
|
if (det > 1e-10 * Ass * AII) {
|
|
error_model_a = std::clamp((Bs * AII - BI * AsI) / det, 0.25, 100.0);
|
|
const double b2 = std::max((Ass * BI - AsI * Bs) / det, 0.0);
|
|
error_model_b = std::sqrt(b2);
|
|
error_model_active = true;
|
|
std::vector<double> chi2;
|
|
chi2.reserve(smp.size());
|
|
for (const auto &s : smp) {
|
|
const double v = error_model_a * s.s2 + b2 * s.I2;
|
|
if (v > 0.0) chi2.push_back(s.dev2 / v);
|
|
}
|
|
error_model_chi2 = chi2.empty() ? 0.0 : median_of(chi2) / CHI2_1_MEDIAN;
|
|
}
|
|
};
|
|
fit_ab(samples);
|
|
// Refit on a misfit-free pool: the merge drops symmetry outliers (|I - median| > reject_nsigma *
|
|
// sigma) from the merged intensity, so drop the equivalent samples (dev2 > reject_nsigma^2 * model
|
|
// variance) from the error-model fit too, keeping the fitted sigmas consistent with the reflections
|
|
// that actually survive. Operates on the shared `samples`, so CPU and GPU stay bit-identical.
|
|
if (reject_outliers && error_model_active) {
|
|
const double ns2 = reject_nsigma * reject_nsigma, b2 = error_model_b * error_model_b;
|
|
std::vector<Sample> kept;
|
|
kept.reserve(samples.size());
|
|
for (const auto &s : samples) {
|
|
const double v = error_model_a * s.s2 + b2 * s.I2;
|
|
if (v > 0.0 && s.dev2 <= ns2 * v) kept.push_back(s);
|
|
}
|
|
if (kept.size() >= static_cast<size_t>(8 * n_bins) && kept.size() < samples.size())
|
|
fit_ab(std::move(kept));
|
|
}
|
|
}
|
|
// Asymptotic I/sigma. ISa is by definition the I -> infinity limit of the signal-to-noise, i.e. the
|
|
// reproducibility of the strongest reflections (Diederichs, Acta Cryst. D66 (2010), 733-740). The
|
|
// (a, b) fit above spans the whole intensity range, and a mild excess of scatter at intermediate
|
|
// intensity raises its systematic term `b`, so 1/b understates that limit. Read it instead directly
|
|
// from the strong equivalents: for each well-measured reflection group the counting-subtracted
|
|
// fractional scatter of its symmetry mates estimates the systematic term, and the robust median over
|
|
// strong groups is the asymptote. The I/sigma threshold that selects "strong" is relaxed on weak /
|
|
// radiation-damaged data that has few strong reflections (rather than fall back to the higher
|
|
// whole-range b); it falls back only when even the relaxed set is too small. Host-side over the
|
|
// merged fulls, so CPU and GPU agree.
|
|
double error_model_b_asymptotic = error_model_b;
|
|
if (error_model_active) {
|
|
struct GroupScatter { double sum = 0, sum_sq = 0, sum_var = 0; int n = 0; };
|
|
std::vector<GroupScatter> gs(n_groups);
|
|
for (const auto &o : fulls) {
|
|
if (!usable_merge(o)) continue;
|
|
const double I_corr = static_cast<double>(o.I) * o.corr;
|
|
const double sigma_corr = static_cast<double>(o.sigma) * o.corr;
|
|
auto &g = gs[o.group];
|
|
g.sum += I_corr; g.sum_sq += I_corr * I_corr; g.sum_var += sigma_corr * sigma_corr; ++g.n;
|
|
}
|
|
// Per-group counting-subtracted fractional systematic variance, paired with the group's I/sigma.
|
|
std::vector<std::pair<double, double>> group_scatter; // (systematic b^2, I/sigma)
|
|
for (const auto &g : gs) {
|
|
if (g.n < 5) continue;
|
|
const double mean = g.sum / g.n;
|
|
const double counting = g.sum_var / g.n;
|
|
if (mean <= 0.0 || counting <= 0.0) continue;
|
|
const double variance = (g.sum_sq - g.sum * g.sum / g.n) / (g.n - 1);
|
|
group_scatter.push_back({std::max(variance - counting, 0.0) / (mean * mean),
|
|
mean / std::sqrt(counting)});
|
|
}
|
|
auto asymptote_above = [&](double snr_min, size_t min_groups) -> double {
|
|
std::vector<double> b2;
|
|
for (const auto &[b2_value, snr] : group_scatter)
|
|
if (snr >= snr_min) b2.push_back(b2_value);
|
|
return b2.size() >= min_groups ? std::sqrt(median_of(b2)) : 0.0;
|
|
};
|
|
double b_asy = asymptote_above(20.0, 100); // tight threshold on data that supports it
|
|
if (b_asy <= 0.0) b_asy = asymptote_above(10.0, 50); // relaxed for weak / damaged data
|
|
if (b_asy > 0.0) error_model_b_asymptotic = b_asy;
|
|
}
|
|
// Guard a degenerate low-multiplicity fit: with too few symmetry equivalents both the (a, b) fit and
|
|
// the per-group scatter collapse toward zero, and 1/b then reports an impossibly high asymptotic
|
|
// I/sigma. Real macromolecular data does not exceed ISa ~50; past a generous cap treat the asymptote
|
|
// as unmeasured (result.isa left undetermined) rather than emit a spurious extreme.
|
|
constexpr double MIN_ASYMPTOTIC_B = 0.01; // ISa cap 100
|
|
if (error_model_b_asymptotic < MIN_ASYMPTOTIC_B) error_model_b_asymptotic = 0.0;
|
|
if (error_model_active)
|
|
logger.Info("Error model: a={:.3f} b={:.3f} ISa={:.1f} chi2={:.2f}", error_model_a, error_model_b,
|
|
error_model_b_asymptotic > 0 ? 1.0 / error_model_b_asymptotic : 0.0, error_model_chi2);
|
|
|
|
auto corrected_sigma = [&](float I_corr, float sigma_corr, int g) -> float {
|
|
if (!error_model_active) return sigma_corr;
|
|
const double I_for_b = std::isfinite(em_mean[g]) ? em_mean[g] : I_corr;
|
|
const double v = error_model_a * static_cast<double>(sigma_corr) * sigma_corr
|
|
+ (error_model_b * I_for_b) * (error_model_b * I_for_b);
|
|
return v > 0.0 ? static_cast<float>(std::sqrt(v)) : sigma_corr;
|
|
};
|
|
// ---- Merge: per-group inverse-variance sums with corrected sigma + deterministic half sets. ----
|
|
struct Accum { double swI = 0, sw = 0, swIh[2] = {0, 0}, swh[2] = {0, 0}; size_t nh[2] = {0, 0}; float d = NAN; };
|
|
std::vector<Accum> acc(n_groups);
|
|
size_t reject_count = 0;
|
|
std::vector<char> rejected_obs(fulls.size(), 0); // per-full outlier-rejected flag (mirrors the GPU path)
|
|
bool did_gpu_acc = false;
|
|
#ifdef JFJOCH_USE_CUDA
|
|
if (use_gpu_merge) {
|
|
std::vector<double> aswI(n_groups), asw(n_groups), aswIh0(n_groups), aswIh1(n_groups),
|
|
aswh0(n_groups), aswh1(n_groups), ad(n_groups);
|
|
std::vector<int32_t> anh0(n_groups), anh1(n_groups), arej(n_groups);
|
|
gpu_->MergeAccum(error_model_a, error_model_b, error_model_active,
|
|
reject_outliers, reject_nsigma, reject_median.data(),
|
|
aswI.data(), asw.data(), aswIh0.data(), aswIh1.data(),
|
|
aswh0.data(), aswh1.data(), anh0.data(), anh1.data(), ad.data(), arej.data());
|
|
for (int g = 0; g < n_groups; ++g) {
|
|
Accum &a = acc[g];
|
|
a.swI = aswI[g]; a.sw = asw[g]; a.swIh[0] = aswIh0[g]; a.swIh[1] = aswIh1[g];
|
|
a.swh[0] = aswh0[g]; a.swh[1] = aswh1[g];
|
|
a.nh[0] = static_cast<size_t>(anh0[g]); a.nh[1] = static_cast<size_t>(anh1[g]);
|
|
a.d = static_cast<float>(ad[g]);
|
|
reject_count += static_cast<size_t>(arej[g]);
|
|
}
|
|
did_gpu_acc = true;
|
|
}
|
|
#endif
|
|
if (!did_gpu_acc)
|
|
for (const auto &o : fulls) {
|
|
if (!usable_merge(o)) continue;
|
|
const float I_corr = o.I * o.corr;
|
|
float sigma_corr = o.sigma * o.corr;
|
|
sigma_corr = corrected_sigma(I_corr, sigma_corr, o.group);
|
|
if (reject_outliers && error_model_active && std::isfinite(reject_median[o.group])
|
|
&& std::fabs(I_corr - reject_median[o.group]) > reject_nsigma * sigma_corr) {
|
|
++reject_count;
|
|
rejected_obs[&o - fulls.data()] = 1;
|
|
continue;
|
|
}
|
|
const double w = 1.0 / (static_cast<double>(sigma_corr) * sigma_corr);
|
|
const double wI = w * I_corr;
|
|
const int half = HalfForImage(o.frame);
|
|
auto &a = acc[o.group];
|
|
a.swI += wI; a.sw += w;
|
|
a.swIh[half] += wI; a.swh[half] += w; a.nh[half]++;
|
|
if (!std::isfinite(a.d) && std::isfinite(o.d) && o.d > 0.0f) a.d = o.d;
|
|
}
|
|
// ---- Export merged reflections (+ resolution-shell R-free flags). ----
|
|
Result result;
|
|
result.isa = error_model_b_asymptotic > 0 ? 1.0 / error_model_b_asymptotic : 0.0;
|
|
std::vector<double> merged_I(n_groups, NAN);
|
|
float d_min = std::numeric_limits<float>::max(), d_max = 0.0f;
|
|
for (int g = 0; g < n_groups; ++g) {
|
|
const auto &a = acc[g];
|
|
if (a.sw <= 0.0) continue;
|
|
MergedReflection mr{};
|
|
mr.h = group_h[g]; mr.k = group_k[g]; mr.l = group_l[g];
|
|
mr.I = static_cast<float>(a.swI / a.sw);
|
|
// The systematic floor caps a high-multiplicity merged sigma at the reproducibility that the
|
|
// strongest reflections actually reach, i.e. the asymptotic term, so merged I/sigma approaches ISa.
|
|
mr.sigma = SigmaWithSystematicFloor(1.0 / std::sqrt(a.sw), mr.I, error_model_b_asymptotic);
|
|
mr.I_half[0] = mr.I_half[1] = NAN;
|
|
mr.sigma_half[0] = mr.sigma_half[1] = NAN;
|
|
mr.d = a.d;
|
|
if (a.nh[0] + a.nh[1] > 0 && a.swh[0] > 0.0 && a.swh[1] > 0.0) {
|
|
for (int i = 0; i < 2; ++i) {
|
|
mr.I_half[i] = static_cast<float>(a.swIh[i] / a.swh[i]);
|
|
mr.sigma_half[i] = SigmaWithSystematicFloor(1.0 / std::sqrt(a.swh[i]), mr.I_half[i], error_model_b_asymptotic);
|
|
}
|
|
}
|
|
if (!std::isfinite(a.d) || a.d <= 0.0f) continue;
|
|
d_min = std::min(d_min, a.d);
|
|
d_max = std::max(d_max, a.d);
|
|
merged_I[g] = mr.I;
|
|
result.merged.push_back(mr);
|
|
}
|
|
|
|
// Automatic high-resolution cutoff (post-merge): trim the written reflections + reported shells to
|
|
// the CC1/2 fall-off. The scaling, combine and error model above already ran over the full range,
|
|
// and the per-image _process.h5 is written elsewhere from the partials, so no data is lost. A manual
|
|
// --scaling-high-resolution (d_min_limit) wins; the P1 search merge (for_search) is never cut, so
|
|
// 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);
|
|
|
|
AssignRfreeFlags(result.merged, x.GetSpaceGroupNumber().value_or(1), rfree_fraction);
|
|
// French-Wilson (F, and F(+)/F(-) from the anomalous split) is deferred until after the anomalous
|
|
// accumulator below has attached I(+)/I(-), so the two hands get their amplitudes in one pass.
|
|
|
|
if (reject_count > 0)
|
|
logger.Info("Merge outlier rejection: dropped {} observations", reject_count);
|
|
// ---- Statistics (report_shell_count shells): completeness, multiplicity, <I/sigma>, R_meas, CC1/2. ----
|
|
const int n_shells = report_shell_count;
|
|
float sd_min = std::numeric_limits<float>::max(), sd_max = 0.0f;
|
|
for (const auto &m : result.merged) {
|
|
if (!std::isfinite(m.d) || m.d <= 0.0f) continue;
|
|
if (effective_d_min && m.d < *effective_d_min) continue;
|
|
sd_min = std::min(sd_min, m.d); sd_max = std::max(sd_max, m.d);
|
|
}
|
|
if (!(sd_min < sd_max && sd_min > 0.0f))
|
|
throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
|
|
"RotationScaleMerge: resolution calculation failed");
|
|
const float d_min_pad = sd_min * 0.999f, d_max_pad = sd_max * 1.001f;
|
|
ResolutionShells shells(d_min_pad, d_max_pad, n_shells);
|
|
const auto shell_mean_1_d2 = shells.GetShellMeanOneOverResSq();
|
|
const auto shell_min_res = shells.GetShellMinRes();
|
|
|
|
struct ShellAcc {
|
|
int unique = 0, total_obs = 0, possible = 0;
|
|
double sum_i_over_sigma = 0.0; int n_i_over_sigma = 0;
|
|
CorrelationCoefficient cc_half;
|
|
};
|
|
std::vector<ShellAcc> sa(n_shells);
|
|
std::vector<int> possible(n_shells, 0);
|
|
if (reference_cell)
|
|
PossiblePerShell(x.GetSpaceGroupNumber().value_or(1), *reference_cell, d_min_pad, d_max_pad,
|
|
shells, merge_friedel, possible);
|
|
for (int s = 0; s < n_shells; ++s) sa[s].possible = possible[s];
|
|
|
|
CorrelationCoefficient cc_half_overall;
|
|
for (const auto &m : result.merged) {
|
|
const auto shell = shells.GetShell(m.d);
|
|
if (!shell || *shell < 0 || *shell >= n_shells) continue;
|
|
if (std::isfinite(m.I) && std::isfinite(m.sigma) && m.sigma > 0.0) {
|
|
auto &s = sa[*shell];
|
|
s.unique++;
|
|
s.sum_i_over_sigma += m.I / m.sigma; s.n_i_over_sigma++;
|
|
if (std::isfinite(m.I_half[0]) && std::isfinite(m.I_half[1])) {
|
|
s.cc_half.Add(m.I_half[0], m.I_half[1]);
|
|
cc_half_overall.Add(m.I_half[0], m.I_half[1]);
|
|
}
|
|
}
|
|
}
|
|
|
|
// R_meas: re-walk the fulls (Mask = cell only; no ice / masked-ring / error-model), accumulate
|
|
// |I_i - <I>| per reflection.
|
|
struct RmeasObs { double sum_abs_dev = 0, sum_I = 0; int n = 0, shell = -1; };
|
|
std::vector<RmeasObs> rmeas(n_groups);
|
|
bool did_gpu_rmeas = false;
|
|
#ifdef JFJOCH_USE_CUDA
|
|
if (use_gpu_merge) {
|
|
// Per-group R_meas + usable count on the GPU; the shell is assigned per group (its fulls share d).
|
|
std::vector<double> rabsdev(n_groups), rsumI(n_groups);
|
|
std::vector<int32_t> rn(n_groups), rnusable(n_groups);
|
|
gpu_->MergeRmeas(merged_I.data(), rabsdev.data(), rsumI.data(), rn.data(), rnusable.data());
|
|
for (int g = 0; g < n_groups; ++g) {
|
|
if (rnusable[g] == 0) continue;
|
|
const auto shell = shells.GetShell(acc[g].d);
|
|
if (!shell || *shell < 0 || *shell >= n_shells) continue;
|
|
sa[*shell].total_obs += rnusable[g];
|
|
if (std::isfinite(merged_I[g]) && rn[g] > 0) {
|
|
auto &r = rmeas[g];
|
|
r.sum_abs_dev = rabsdev[g]; r.sum_I = rsumI[g]; r.n = rn[g]; r.shell = *shell;
|
|
}
|
|
}
|
|
did_gpu_rmeas = true;
|
|
}
|
|
#endif
|
|
if (!did_gpu_rmeas)
|
|
for (const auto &o : fulls) {
|
|
if (o.group < 0) continue;
|
|
if (rejected_obs[&o - fulls.data()]) continue; // outlier-rejected in the merge -> also out of R_meas
|
|
if (!frame_cell_ok[o.frame]) continue;
|
|
if (!(o.corr > 0.0f) || !std::isfinite(o.corr)) continue;
|
|
if (o.partiality < min_partiality) continue;
|
|
const float I_corr = o.I * o.corr, sigma_corr = o.sigma * o.corr;
|
|
if (!std::isfinite(I_corr) || !std::isfinite(sigma_corr) || sigma_corr <= 0.0f) continue;
|
|
const auto shell = shells.GetShell(o.d);
|
|
if (!shell || *shell < 0 || *shell >= n_shells) continue;
|
|
sa[*shell].total_obs++;
|
|
if (std::isfinite(merged_I[o.group])) {
|
|
auto &r = rmeas[o.group];
|
|
r.sum_abs_dev += std::fabs(static_cast<double>(I_corr) - merged_I[o.group]);
|
|
r.sum_I += I_corr; r.n++; r.shell = *shell;
|
|
}
|
|
}
|
|
std::vector<double> rmeas_num(n_shells, 0.0), rmeas_den(n_shells, 0.0);
|
|
double rmeas_num_all = 0.0, rmeas_den_all = 0.0;
|
|
for (const auto &r : rmeas) {
|
|
if (r.n < 2 || r.shell < 0 || r.shell >= n_shells) continue;
|
|
const double factor = std::sqrt(static_cast<double>(r.n) / (r.n - 1));
|
|
rmeas_num[r.shell] += factor * r.sum_abs_dev; rmeas_den[r.shell] += r.sum_I;
|
|
rmeas_num_all += factor * r.sum_abs_dev; rmeas_den_all += r.sum_I;
|
|
}
|
|
|
|
// ---- Anomalous split (always, even when the merge is Friedel-averaged): for each acentric
|
|
// reflection, the inverse-variance I(+)/I(-) from the SAME scaled fulls (so it never touches the
|
|
// Friedel-merged IMEAN, scaling or error model above). Lets I(+)/I(-) and F(+)/F(-) be written by
|
|
// default without scaling anomalously; a reflection with only one mate, or a centric, is left
|
|
// without the split. Skipped for the P1 search pass, which has no use for it. ----
|
|
struct AnomExport { float Ip = NAN, sIp = NAN, Im = NAN, sIm = NAN; };
|
|
std::unordered_map<uint64_t, AnomExport> anom_export;
|
|
// SigAno = <|I(+)-I(-)|>/<sigma(I(+)-I(-))> accumulated on the standard report shells (num=sum|dI|,
|
|
// den=sum sigma(dI)); mmCIF pdbx_absDiff_over_sigma_anomalous. NaN'd below when there is no split.
|
|
std::vector<double> sig_num(n_shells, 0.0), sig_den(n_shells, 0.0);
|
|
double sig_num_all = 0.0, sig_den_all = 0.0;
|
|
size_t sig_n = 0;
|
|
if (!for_search) {
|
|
const int sg_num = x.GetSpaceGroupNumber().value_or(1);
|
|
const HKLKeyGenerator anom_keygen(/*merge_friedel=*/false, sg_num);
|
|
const gemmi::GroupOps gops = gemmi::find_spacegroup_by_number(sg_num)->operations();
|
|
struct AnomAcc { double swI[2] = {}; double sw[2] = {}; int32_t h = 0, k = 0, l = 0; float d = NAN; }; // [hand] 0=I(+) 1=I(-)
|
|
std::unordered_map<uint64_t, AnomAcc> anom;
|
|
anom.reserve(result.merged.size() * 2 + 1);
|
|
for (const auto &o : fulls) {
|
|
if (!usable_merge(o)) continue;
|
|
if (rejected_obs[&o - fulls.data()]) continue; // outlier-rejected in the merge (CPU path)
|
|
const HKLKey ak = anom_keygen(o.h, o.k, o.l);
|
|
const int hand = ak.plus ? 0 : 1;
|
|
const float I_corr = o.I * o.corr;
|
|
const float sigma_corr = corrected_sigma(I_corr, o.sigma * o.corr, o.group);
|
|
if (!(sigma_corr > 0.0f) || !std::isfinite(sigma_corr)) continue;
|
|
const double w = 1.0 / (static_cast<double>(sigma_corr) * sigma_corr);
|
|
AnomAcc &a = anom[HKLKey{ak.h, ak.k, ak.l, true}.pack()];
|
|
a.h = ak.h; a.k = ak.k; a.l = ak.l; a.d = o.d;
|
|
a.swI[hand] += w * static_cast<double>(I_corr); a.sw[hand] += w;
|
|
}
|
|
for (const auto &[fkey, a] : anom) {
|
|
// Centrics have I(+)=I(-) by symmetry; leave them without an anomalous split.
|
|
if (gops.is_reflection_centric(gemmi::Op::Miller{a.h, a.k, a.l})) continue;
|
|
AnomExport ex;
|
|
if (a.sw[0] > 0.0) { ex.Ip = static_cast<float>(a.swI[0] / a.sw[0]); ex.sIp = static_cast<float>(1.0 / std::sqrt(a.sw[0])); }
|
|
if (a.sw[1] > 0.0) { ex.Im = static_cast<float>(a.swI[1] / a.sw[1]); ex.sIm = static_cast<float>(1.0 / std::sqrt(a.sw[1])); }
|
|
if (std::isfinite(ex.Ip) || std::isfinite(ex.Im)) anom_export[fkey] = ex;
|
|
if (a.sw[0] > 0.0 && a.sw[1] > 0.0 && a.d > 0.0f) {
|
|
const double abs_di = std::fabs(a.swI[0] / a.sw[0] - a.swI[1] / a.sw[1]);
|
|
const double sig_di = std::sqrt(1.0 / a.sw[0] + 1.0 / a.sw[1]);
|
|
const auto shell = shells.GetShell(a.d);
|
|
if (shell && *shell >= 0 && *shell < n_shells) { sig_num[*shell] += abs_di; sig_den[*shell] += sig_di; }
|
|
sig_num_all += abs_di; sig_den_all += sig_di; ++sig_n;
|
|
}
|
|
}
|
|
}
|
|
|
|
MergeStatistics &out = result.statistics;
|
|
out.shells.resize(n_shells);
|
|
for (int s = 0; s < n_shells; ++s) {
|
|
auto &ss = out.shells[s];
|
|
ss.mean_one_over_d2 = shell_mean_1_d2[s];
|
|
ss.d_min = shell_min_res[s];
|
|
ss.d_max = s == 0 ? d_max_pad : shell_min_res[s - 1];
|
|
ss.total_observations = sa[s].total_obs;
|
|
ss.unique_reflections = sa[s].unique;
|
|
ss.possible_unique_reflections = sa[s].possible;
|
|
ss.mean_i_over_sigma = sa[s].n_i_over_sigma > 0 ? sa[s].sum_i_over_sigma / sa[s].n_i_over_sigma : 0.0;
|
|
ss.cc_half = sa[s].cc_half.GetCC();
|
|
ss.cc_ref = NAN;
|
|
ss.r_meas = rmeas_den[s] > 0.0 ? rmeas_num[s] / rmeas_den[s] : NAN;
|
|
ss.abs_diff_over_sigma_anomalous = sig_den[s] > 0.0 ? sig_num[s] / sig_den[s] : NAN;
|
|
}
|
|
auto &overall = out.overall;
|
|
overall.d_min = sd_min; overall.d_max = sd_max;
|
|
double sum_ios = 0.0; int n_ios = 0;
|
|
for (int s = 0; s < n_shells; ++s) {
|
|
overall.total_observations += sa[s].total_obs;
|
|
overall.unique_reflections += sa[s].unique;
|
|
overall.possible_unique_reflections += sa[s].possible;
|
|
sum_ios += sa[s].sum_i_over_sigma; n_ios += sa[s].n_i_over_sigma;
|
|
}
|
|
overall.mean_i_over_sigma = n_ios > 0 ? sum_ios / n_ios : 0.0;
|
|
overall.cc_half = cc_half_overall.GetCC();
|
|
overall.cc_ref = NAN;
|
|
overall.r_meas = rmeas_den_all > 0.0 ? rmeas_num_all / rmeas_den_all : NAN;
|
|
overall.abs_diff_over_sigma_anomalous = sig_den_all > 0.0 ? sig_num_all / sig_den_all : NAN;
|
|
if (std::isfinite(overall.abs_diff_over_sigma_anomalous))
|
|
logger.Info("Anomalous signal SigAno = {:.2f} ({} acentric pairs)",
|
|
overall.abs_diff_over_sigma_anomalous, sig_n);
|
|
|
|
// Radiation-damage monitor (measured before any correction by MeasureRadiationDamageB): carry the
|
|
// first->last relative-B change and the per-batch curve into the reported statistics / mmCIF.
|
|
out.radiation_damage_delta_b = rad_damage_delta_b;
|
|
out.radiation_damage_b_batch = rad_damage_b_batch;
|
|
out.radiation_damage_batch_deg = rad_damage_batch_deg;
|
|
|
|
// Attach the per-reflection anomalous split so the writer can emit I(+)/I(-) by default (each merged
|
|
// reflection maps to its Friedel-ASU key; in an anomalous merge both mates map to the same key).
|
|
if (!anom_export.empty()) {
|
|
const HKLKeyGenerator anom_keygen(/*merge_friedel=*/false, x.GetSpaceGroupNumber().value_or(1));
|
|
for (auto &r : result.merged) {
|
|
const HKLKey ak = anom_keygen(r.h, r.k, r.l);
|
|
const auto it = anom_export.find(HKLKey{ak.h, ak.k, ak.l, true}.pack());
|
|
if (it == anom_export.end()) continue;
|
|
r.I_plus = it->second.Ip; r.sigma_plus = it->second.sIp;
|
|
r.I_minus = it->second.Im; r.sigma_minus = it->second.sIm;
|
|
}
|
|
}
|
|
|
|
// French-Wilson amplitudes for IMEAN and (now that they are attached) each Bijvoet hand.
|
|
ApplyFrenchWilson(result.merged, x.GetSpaceGroupNumber().value_or(1));
|
|
|
|
logger.Info("Merge complete ({} unique reflections)", result.merged.size());
|
|
return result;
|
|
}
|
|
|
|
RotationScaleMerge::Result RotationScaleMerge::Run(bool for_search,
|
|
const std::vector<char> &masked_ice_rings) {
|
|
const int sg_number = x.GetSpaceGroupNumber().value_or(1);
|
|
HKLKeyGenerator keygen(merge_friedel, sg_number);
|
|
|
|
// --- 1. Per-frame partial scaling (Rotation model, per-image G only). ---
|
|
const int n_groups = ComputeAsuGroups(keygen); // one ASU grouping, shared by partials and fulls
|
|
std::vector<double> partial_mean;
|
|
bool scaled_on_gpu = false;
|
|
#ifdef JFJOCH_USE_CUDA
|
|
if (gpu_active_) {
|
|
// Run the whole scaling loop on the GPU; corr stays RESIDENT across scaling -> smooth-G -> CC ->
|
|
// combine (and across passes, exactly as the old host round-trip did). Only the per-frame G/scaled
|
|
// come back (needed to compute smooth-G and the writeback).
|
|
gpu_->ScalePartials(scaling_iter, SCALE_ROBUST_K, min_partiality, d_min_limit.has_value());
|
|
frame_scaled_scratch.assign(n_frames, 0);
|
|
gpu_->GetG(g_partial.data(), frame_scaled_scratch.data());
|
|
scaled_on_gpu = true;
|
|
}
|
|
#endif
|
|
if (!scaled_on_gpu) {
|
|
for (int it = 0; it < scaling_iter; ++it) {
|
|
ReduceGroupMeans(partials, n_groups, false, {}, partial_mean);
|
|
FitPerFrameG(partials, frame_start, frame_count, partial_mean, /*unity=*/false, g_partial);
|
|
UpdateCorr(partials, g_partial, frame_scaled_scratch);
|
|
}
|
|
}
|
|
const std::vector<uint8_t> partial_scaled = frame_scaled_scratch;
|
|
|
|
// --- 2. Smooth G across frames (XDS DELPHI-like) before the combine. ---
|
|
const auto s = x.GetScalingSettings();
|
|
const double smooth_g_deg = s.GetSmoothGDegrees();
|
|
const auto gonio = x.GetGoniometer();
|
|
const double osc_deg = gonio ? std::fabs(gonio->GetIncrement_deg()) : 0.0;
|
|
if (smooth_g_deg > 0.0 && osc_deg > 1e-6) {
|
|
int window = std::max(1, static_cast<int>(std::lround(smooth_g_deg / osc_deg)));
|
|
if (window % 2 == 0) ++window;
|
|
bool smoothed_on_gpu = false;
|
|
#ifdef JFJOCH_USE_CUDA
|
|
if (gpu_active_) {
|
|
// Apply smooth-G to the resident corr: the host builds the per-frame ratio g/g_smooth (+ an
|
|
// apply flag), the kernel multiplies corr in place. Same guard / arithmetic as host SmoothG.
|
|
std::vector<double> g_smooth;
|
|
ComputeSmoothGWindow(g_partial, window, g_smooth);
|
|
std::vector<uint8_t> apply(n_frames, 0);
|
|
std::vector<double> ratio(n_frames, 1.0);
|
|
for (int f = 0; f < n_frames; ++f)
|
|
if (frame_scaled_scratch[f] && std::isfinite(g_partial[f]) && g_partial[f] > 0.0
|
|
&& std::isfinite(g_smooth[f])) {
|
|
apply[f] = 1;
|
|
ratio[f] = g_partial[f] / g_smooth[f];
|
|
}
|
|
gpu_->SmoothCorr(apply.data(), ratio.data());
|
|
for (int f = 0; f < n_frames; ++f) if (apply[f]) g_partial[f] = g_smooth[f];
|
|
smoothed_on_gpu = true;
|
|
}
|
|
#endif
|
|
if (!smoothed_on_gpu)
|
|
SmoothG(partials, g_partial, window);
|
|
}
|
|
|
|
#ifdef JFJOCH_USE_CUDA
|
|
// The GPU keeps corr resident through scaling + smooth-G; only the diagnostic dump falls back to the
|
|
// CPU combine, which reads host partials[].corr, so refresh it for that path. (The decay/absorption
|
|
// correction surfaces run on the GPU-combined fulls - see below - so they do NOT need the CPU combine.)
|
|
if (gpu_active_ && !observation_dump_path.empty()) {
|
|
std::vector<float> corr(partials.size());
|
|
gpu_->GetCorr(corr.data());
|
|
for (size_t i = 0; i < partials.size(); ++i) partials[i].corr = corr[i];
|
|
}
|
|
#endif
|
|
|
|
// Per-frame CC + write G/CC/mosaicity back onto the partials (once). On the GPU the group means +
|
|
// per-frame CC run on the resident (already smoothed) corr; only the tiny per-frame cc/cc_n come back.
|
|
std::vector<double> cc;
|
|
std::vector<int64_t> cc_n;
|
|
bool cc_on_gpu = false;
|
|
#ifdef JFJOCH_USE_CUDA
|
|
if (gpu_active_) {
|
|
cc.resize(n_frames);
|
|
cc_n.resize(n_frames);
|
|
gpu_->ComputePartialCC(min_partiality, cc.data(), cc_n.data());
|
|
cc_on_gpu = true;
|
|
}
|
|
#endif
|
|
if (!cc_on_gpu) {
|
|
ReduceGroupMeans(partials, n_groups, false, {}, partial_mean);
|
|
ComputePerFrameCC(partial_mean, cc, cc_n);
|
|
}
|
|
FinalizePerFrameScale(cc, cc_n, partial_scaled);
|
|
|
|
// --- 2b. Drop frames that do not agree with the merged reference (--min-image-cc). ---
|
|
if (min_cc_for_image > 0.0) {
|
|
std::vector<uint8_t> reject(n_frames, 0);
|
|
int n_rejected = 0;
|
|
for (int f = 0; f < n_frames; ++f)
|
|
if (std::isfinite(cc[f]) && cc_n[f] >= MIN_REFLECTIONS_FOR_IMAGE_CC && cc[f] < min_cc_for_image) {
|
|
reject[f] = 1;
|
|
++n_rejected;
|
|
}
|
|
if (n_rejected > 0) {
|
|
// Zeroing corr is how a frame leaves the pipeline: every consumer requires corr > 0, so the
|
|
// frame's partials stop being usable for the combine, the merge and the error model alike.
|
|
bool rejected_on_gpu = false;
|
|
#ifdef JFJOCH_USE_CUDA
|
|
if (gpu_active_) {
|
|
std::vector<double> ratio(n_frames, 1.0);
|
|
for (int f = 0; f < n_frames; ++f) if (reject[f]) ratio[f] = 0.0;
|
|
gpu_->SmoothCorr(reject.data(), ratio.data());
|
|
rejected_on_gpu = true;
|
|
}
|
|
#endif
|
|
if (!rejected_on_gpu)
|
|
for (auto &o : partials)
|
|
if (reject[o.frame]) o.corr = 0.0f;
|
|
logger.Info("Rejected {} of {} frames correlating below {:.2f} with the merged reference",
|
|
n_rejected, n_frames, min_cc_for_image);
|
|
}
|
|
}
|
|
|
|
// --- 3. 3D combine of per-frame partials into fulls (fulls inherit their ASU group here). ---
|
|
bool combined_on_gpu = false;
|
|
bool scaled_fulls_on_gpu = false;
|
|
#ifdef JFJOCH_USE_CUDA
|
|
// GPU combine (+ scale-fulls) keeps the fulls resident on the device: combine, then build the frame /
|
|
// ASU-group CSRs on the host from just the small key arrays (a deterministic counting sort - no GPU
|
|
// stable-sort), scale the fulls in place, and download only once. Mirrors Combine() + the Unity
|
|
// scale-fulls loop below. Only the diagnostic dump (serial, one writer) falls back to CPU. The
|
|
// decay/absorption correction surfaces run on the downloaded fulls (px/py is carried through the
|
|
// combine) and their corrected corr is re-uploaded before the resident merge - so they stay on the GPU.
|
|
if (gpu_active_ && observation_dump_path.empty()) {
|
|
// The smoothed corr is already resident (scaling + smooth-G ran on the device, no round-trip).
|
|
const int nf = gpu_->Combine(rawrun_group.data(), min_partiality, capture_uncertainty_coeff,
|
|
min_captured_fraction);
|
|
g_full.assign(n_frames, 1.0);
|
|
|
|
if (scale_fulls && nf > 0) {
|
|
// Frame + group CSRs over the emit-ordered fulls, built by counting sort on the host (stable,
|
|
// deterministic). frame is always in [0, n_frames); group is <0 for absent/out-of-range fulls.
|
|
std::vector<int32_t> ff(nf), fg(nf);
|
|
gpu_->GetFullsKeys(ff.data(), fg.data());
|
|
std::vector<int32_t> f_start(n_frames, 0), f_count(n_frames, 0), f_perm(nf);
|
|
for (int i = 0; i < nf; ++i) ++f_count[ff[i]];
|
|
for (int f = 1; f < n_frames; ++f) f_start[f] = f_start[f - 1] + f_count[f - 1];
|
|
{ std::vector<int32_t> fill = f_start; for (int i = 0; i < nf; ++i) f_perm[fill[ff[i]]++] = i; }
|
|
gpu_->SetFullsFrameCSR(f_perm.data(), nf, f_start.data(), f_count.data());
|
|
|
|
std::vector<int32_t> g_count(n_groups, 0), g_start(n_groups, 0);
|
|
for (int i = 0; i < nf; ++i) if (fg[i] >= 0) ++g_count[fg[i]];
|
|
int acc = 0;
|
|
for (int g = 0; g < n_groups; ++g) { g_start[g] = acc; acc += g_count[g]; }
|
|
std::vector<int32_t> g_perm(acc);
|
|
{ std::vector<int32_t> fill = g_start; for (int i = 0; i < nf; ++i) if (fg[i] >= 0) g_perm[fill[fg[i]]++] = i; }
|
|
gpu_->SetFullsGroups(g_perm.data(), acc, g_start.data(), g_count.data());
|
|
|
|
gpu_->ScaleFulls(scaling_iter, SCALE_ROBUST_K, min_partiality);
|
|
scaled_fulls_on_gpu = true;
|
|
}
|
|
|
|
fulls.assign(nf, Obs{});
|
|
std::vector<int32_t> fh(nf), fk(nf), fl(nf), fframe(nf), fgroup(nf);
|
|
std::vector<float> fI(nf), fsig(nf), fd(nf), fimg(nf), fcorr(nf, 1.0f), fpx(nf), fpy(nf);
|
|
std::vector<uint8_t> fon(nf);
|
|
gpu_->GetFulls(fh.data(), fk.data(), fl.data(), fI.data(), fsig.data(), fd.data(),
|
|
fimg.data(), fframe.data(), fon.data(), fgroup.data());
|
|
gpu_->GetFullsPxPy(fpx.data(), fpy.data());
|
|
if (scaled_fulls_on_gpu) gpu_->GetFullsCorr(fcorr.data());
|
|
for (int i = 0; i < nf; ++i) {
|
|
Obs &o = fulls[i];
|
|
o.h = fh[i]; o.k = fk[i]; o.l = fl[i];
|
|
o.I = fI[i]; o.sigma = fsig[i]; o.d = fd[i];
|
|
o.rlp = 1.0f; o.partiality = 1.0f; o.corr = fcorr[i];
|
|
o.image_number = fimg[i]; o.frame = fframe[i]; o.px = fpx[i]; o.py = fpy[i];
|
|
o.on_ice = fon[i]; o.group = fgroup[i];
|
|
}
|
|
logger.Info("3D combine{} (GPU): {} fulls", scaled_fulls_on_gpu ? " + scale-fulls" : "", nf);
|
|
combined_on_gpu = true;
|
|
}
|
|
#endif
|
|
if (!combined_on_gpu)
|
|
Combine();
|
|
|
|
// --- 4. Scale the fulls (XDS order, Unity model). ---
|
|
if (scale_fulls && !scaled_fulls_on_gpu) {
|
|
std::vector<double> full_mean;
|
|
for (int it = 0; it < scaling_iter; ++it) {
|
|
ReduceGroupMeans(fulls, n_groups, false, {}, full_mean);
|
|
FitPerFrameG(fulls, fulls_frame_start, fulls_frame_count, full_mean, /*unity=*/true, g_full);
|
|
UpdateCorr(fulls, g_full, frame_scaled_scratch);
|
|
}
|
|
logger.Info("Scaled fulls (XDS order, Unity model)");
|
|
}
|
|
const bool rejected_full_scales = scale_fulls && RejectCollapsedFullScales();
|
|
|
|
// --- 4b. Optional correction surfaces (decay = resolution x time; absorption = goniometer-frame
|
|
// diffracted-beam direction), each an alternating multiplicative fit of the fulls' corr against
|
|
// the merged reference. Cheap host loops over the downloaded fulls; applied only on the final
|
|
// in-symmetry merge, never the P1 space-group search pass (corrections there add risk and can
|
|
// perturb the symmetry determination). ---
|
|
[[maybe_unused]] const bool corrections =
|
|
!for_search && (refine_decay_b || absorption_iter > 0 || modulation_iter > 0 || relative_b_deg > 0.0);
|
|
// Radiation-damage monitor: measure the per-batch relative-B on the scaled fulls BEFORE any correction
|
|
// (report-only; captures the full damage signature, not a residual). Skipped on the P1 search pass.
|
|
if (!for_search)
|
|
MeasureRadiationDamageB(n_groups);
|
|
if (!for_search && refine_decay_b)
|
|
RefineDecay(n_groups);
|
|
if (!for_search && relative_b_deg > 0.0)
|
|
RefineRelativeB(n_groups); // per-batch relative-B on top of the single decay slope
|
|
if (!for_search && absorption_iter > 0)
|
|
RefineAbsorption(absorption_iter, n_groups);
|
|
if (!for_search && modulation_iter > 0)
|
|
RefineModulation(modulation_iter, n_groups);
|
|
#ifdef JFJOCH_USE_CUDA
|
|
// The corrections mutate the host fulls' corr; when the merge runs on the resident (GPU) fulls, push
|
|
// the corrected corr back to the device so the merge reads it.
|
|
if ((corrections || rejected_full_scales) && combined_on_gpu && scaled_fulls_on_gpu) {
|
|
std::vector<float> fcorr(fulls.size());
|
|
for (size_t i = 0; i < fulls.size(); ++i) fcorr[i] = fulls[i].corr;
|
|
gpu_->SetFullsCorr(fcorr.data());
|
|
}
|
|
#endif
|
|
|
|
// --- 5. Error model + merge + statistics. ---
|
|
auto r = MergeAndStats(n_groups, for_search, masked_ice_rings, combined_on_gpu && scaled_fulls_on_gpu);
|
|
return r;
|
|
}
|