leonarski_fandClaude Opus 4.8
c5cfdcef7e
Rotation scaling modelled dose-dependent scattering power with a single global relative-B decay slope only (RefineDecay); per-frame refinement was scale-only. Add a per-batch relative Debye-Waller B, both as an opt-in correction and as an always-on measurement. - FitRelativeBCurve / SolveCurvatureSmoothedB: per rotation-range batch, fit the s^2 slope of ln(Iref/Iobs) that the resolution-flat per-frame G leaves, regularised by a second-difference (curvature) penalty so a genuine smooth relative-B is kept while resolution-correlated noise (which an unconstrained per-batch fit would chase into an oscillating curve) is suppressed. - Correction (--relative-b[=deg], default off): applied after RefineDecay, gated by a physical peak-to-peak floor and cross-validated by ASU-group parity (a per-batch parameter owns whole frames, so it is scored on held-out equivalents, not held-out frames). On a battery re-scaling identical stored reflections it is a no-op on well-behaved data and a small gain on radiation-damaged rotation data; never a regression. - Radiation-damage monitor (always on for rotation, report-only): measure the relative-B change from the pristine low-dose start of the run to its end (a whole-run-mean reference inverts on damaged data) plus the per-batch curve, and report them on the command line (in the radiation-damage report and the merge statistics) and in the mmCIF (_reflns.jfjoch_radiation_damage_ relative_B and a per-batch loop). A large magnitude flags radiation damage, complementing the existing per-image CC / mosaicity read. - --scale now reproduces the de-novo rotation merge (it previously omitted scale-fulls and smooth-G, silently using a weaker model), and honours --relative-b, so an offline re-scale matches the full pipeline. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2024 lines
108 KiB
C++
2024 lines
108 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; // 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
|
|
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();
|
|
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). Decoupled from the stills-only -B / RefineB flag.
|
|
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);
|
|
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] > 0.0) {
|
|
sum_log += std::log(g[j]);
|
|
++count;
|
|
}
|
|
}
|
|
if (count > 0) g_smooth[o] = std::exp(sum_log / count);
|
|
}
|
|
}
|
|
|
|
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];
|
|
}
|
|
|
|
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_b_factor_Ang2.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;
|
|
}
|
|
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);
|
|
|
|
// --- 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)");
|
|
}
|
|
|
|
// --- 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 && 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;
|
|
}
|