leonarski_fandClaude Opus 5
52f0e58cae
SmoothGeometry de-rotates each frame's lattice to a common reference, averages in frame order and rotates back. A frame that never indexed keeps a default-constructed CrystalLattice whose vectors are all ZERO - and zero is finite, so the isfinite guard let it through. Those zero vectors were averaged into their neighbours' smoothed orientation, pulling it toward the origin, and they were scored in the leave-one-out cross-validation that picks the smoothing window. On a crystal where 374 of 900 frames fail to index, the effect on the window choice is not subtle. Measured: before n_scored 900 (only 526 indexed) CV score ~504-542 A^2 window +-12 after n_scored 516-526 CV score 0.160-0.175 window +-2 The score was inflated ~3000x and the choice among windows was noise. It settled on +-12 frames - 9.6 degrees of goniometer rotation - on a crystal whose orientation genuinely drifts by ~8 degrees over the sweep, so every partial's delta_phi was recomputed from an orientation averaged across that drift. Require a real cell. Exactly inert when every frame indexes, and no threshold is touched. The crystal that exposed it goes P1 -> P2_1, observations 60107 -> 77021, completeness 64.1% -> 93.0%, multiplicity 1.10 -> 2.0, CC1/2 70.0% -> 84.9%, R_meas low shell 37.3% -> 22.1%, and its 2-fold operator CC 0.330 -> 0.669, comfortably clear of the 0.5 gate. Battery over 37 crystals: space groups 33/37 -> 34/37, and that crystal is the ONLY flip - no losses. Another crystal is rescued from near-total collapse (4402 -> 139213 observations) because the two-pass "going back to the header geometry" fallback stops firing. Anomalous peak height +0.043 +- 0.022 sigma over 7 crystals, so the background clip's gain is intact. Merged quality is otherwise neutral (CC1/2 6 better/6 worse, R_meas_lo 9/6) with observations up on 18 crystals. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
2462 lines
130 KiB
C++
2462 lines
130 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 <array>
|
|
#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 int64_t MIN_REFLECTIONS_FOR_IMAGE_CC = 20; // below this a frame's CC means nothing
|
|
constexpr double SCALE_ROBUST_K = 3.0; // Cauchy loss scale (sigma units) for the per-frame G fit
|
|
|
|
// A fitted per-frame scale below this fraction of the run's median is not a measurement of anything:
|
|
// it would mean the frame received 2% of its neighbours' dose while still producing indexable spots.
|
|
// It matters because the scale enters as 1/G, so a G that collapses toward zero multiplies every
|
|
// intensity on that frame without bound - and sigma by the identical factor, which is why no
|
|
// sigma-based outlier test can see it. Measured over 12 rotation crystals in the default
|
|
// configuration, the smallest LEGITIMATE min(G)/median(G) is 0.070; the case this bound exists for
|
|
// sat at 0.0017 (a 546x amplification that put 99% of sum(I^2) into 0.66% of the observations and
|
|
// took the merged CC1/2 from 93.7% to 17.2%). 0.02 is 3.5x below anything real and 12x above the
|
|
// failure.
|
|
constexpr double MIN_CREDIBLE_SCALE_RATIO = 0.02;
|
|
constexpr float MAX_FRAME_GAP = 2.0f; // a rocking event is a run of frames no more apart than this
|
|
constexpr double CHI2_1_MEDIAN = 0.454936;
|
|
// A post-scale-fulls correction surface (decay / absorption) is applied only if its held-out
|
|
// cross-validation improvement exceeds this fraction of the held-out scatter. A margin (not just >0)
|
|
// stops it acting on marginal, noise-level "improvements" that leave the merged quality unchanged or
|
|
// slightly worse - a correction should engage only when the generalizing signal is clear.
|
|
constexpr double CV_MIN_RELATIVE_GAIN = 0.02;
|
|
|
|
double SafeInv(double x, double fallback) {
|
|
if (!std::isfinite(x) || x == 0.0)
|
|
return fallback;
|
|
return 1.0 / x;
|
|
}
|
|
|
|
// Kabsch rotation partiality: the fraction of a reflection recorded in the sampled slice, from the
|
|
// erf of the rocking angle relative to the mosaic width. Identical to ScaleOnTheFly's RotationPartiality
|
|
// (and the predictor's), so recomputing here just swaps in the smoothed mosaicity.
|
|
float RotationPartiality(double delta_phi_deg, double zeta, double mosaicity_deg, double wedge_deg) {
|
|
const double half_wedge = wedge_deg / 2.0;
|
|
const double c1 = zeta / std::sqrt(2.0);
|
|
const double arg_plus = (delta_phi_deg + half_wedge) * c1 / mosaicity_deg;
|
|
const double arg_minus = (delta_phi_deg - half_wedge) * c1 / mosaicity_deg;
|
|
return static_cast<float>((std::erf(arg_plus) - std::erf(arg_minus)) / 2.0);
|
|
}
|
|
|
|
// Deterministic CC1/2 half from the frame's stable index (splitmix64), matching Merge.cpp.
|
|
int HalfForImage(int64_t image_id) {
|
|
uint64_t z = static_cast<uint64_t>(image_id) + 0x9e3779b97f4a7c15ULL;
|
|
z = (z ^ (z >> 30)) * 0xbf58476d1ce4e5b9ULL;
|
|
z = (z ^ (z >> 27)) * 0x94d049bb133111ebULL;
|
|
z = z ^ (z >> 31);
|
|
return static_cast<int>(z & 1ULL);
|
|
}
|
|
|
|
struct ScaleObs { double coeff, Iobs, weight; };
|
|
|
|
// Robust per-frame scale (linear in G): the exact objective ScaleOnTheFly::SolveScaleIRLS solves.
|
|
double SolveScaleIRLS(const std::vector<ScaleObs> &obs, double robust_k) {
|
|
auto weighted_scale = [&obs](auto robust_weight) {
|
|
double num = 0.0, den = 0.0;
|
|
for (const auto &o : obs) {
|
|
const double rw = robust_weight(o);
|
|
const double w2 = o.weight * o.weight;
|
|
num += rw * w2 * o.coeff * o.Iobs;
|
|
den += rw * w2 * o.coeff * o.coeff;
|
|
}
|
|
return den > 0.0 ? num / den : NAN;
|
|
};
|
|
|
|
double G = weighted_scale([](const ScaleObs &) { return 1.0; });
|
|
if (!std::isfinite(G))
|
|
return 1.0;
|
|
G = std::max(0.0, G);
|
|
|
|
const double k2 = robust_k * robust_k;
|
|
for (int iter = 0; iter < 30; ++iter) {
|
|
const double G_prev = G;
|
|
const double G_next = weighted_scale([&](const ScaleObs &o) {
|
|
const double res = o.weight * (G * o.coeff - o.Iobs);
|
|
return 1.0 / (1.0 + res * res / k2);
|
|
});
|
|
if (!std::isfinite(G_next))
|
|
break;
|
|
G = std::max(0.0, G_next);
|
|
if (std::abs(G - G_prev) <= 1e-7 * std::max(G, 1.0))
|
|
break;
|
|
}
|
|
return G;
|
|
}
|
|
|
|
// Run fn(i) for i in [0, n) over `nthreads` workers pulling from a shared atomic counter - the same
|
|
// self-load-balancing pattern the rest of the codebase uses (heavy frames don't stall light ones).
|
|
// Work-stealing per-item parallel: one atomic fetch per item. Use ONLY when the per-item work is
|
|
// heavy and uneven (e.g. per-frame fits) - the atomic amortises. For millions of tiny uniform items
|
|
// use ParallelChunks instead; a per-item atomic there is pure contention.
|
|
template <typename Fn>
|
|
void ParallelFor(int n, size_t nthreads, Fn fn) {
|
|
if (n <= 0) return;
|
|
if (nthreads <= 1 || n == 1) {
|
|
for (int i = 0; i < n; ++i) fn(i);
|
|
return;
|
|
}
|
|
const size_t local = std::min(nthreads, static_cast<size_t>(n));
|
|
std::atomic<int> next = 0;
|
|
std::vector<std::future<void>> futures;
|
|
futures.reserve(local);
|
|
for (size_t t = 0; t < local; ++t)
|
|
futures.emplace_back(std::async(std::launch::async, [&] {
|
|
for (int i = next.fetch_add(1); i < n; i = next.fetch_add(1))
|
|
fn(i);
|
|
}));
|
|
for (auto &f : futures) f.get();
|
|
}
|
|
|
|
// Chunked parallel: each worker gets one contiguous [lo, hi) range, no per-item synchronisation.
|
|
// Right for millions of cheap uniform items (the CPU stand-in for a flat CUDA grid-stride kernel).
|
|
template <typename Fn>
|
|
void ParallelChunks(int n, size_t nthreads, Fn fn) {
|
|
if (n <= 0) return;
|
|
const int nt = static_cast<int>(std::max<size_t>(1, std::min(nthreads, static_cast<size_t>(n))));
|
|
if (nt == 1) { fn(0, n); return; }
|
|
const int chunk = (n + nt - 1) / nt;
|
|
std::vector<std::future<void>> futures;
|
|
futures.reserve(nt);
|
|
for (int t = 0; t < nt; ++t) {
|
|
const int lo = t * chunk, hi = std::min(n, lo + chunk);
|
|
if (lo >= hi) break;
|
|
futures.emplace_back(std::async(std::launch::async, [&fn, lo, hi] { fn(lo, hi); }));
|
|
}
|
|
for (auto &f : futures) f.get();
|
|
}
|
|
|
|
double median_of(std::vector<double> &v) {
|
|
std::nth_element(v.begin(), v.begin() + v.size() / 2, v.end());
|
|
return v[v.size() / 2];
|
|
}
|
|
}
|
|
|
|
RotationScaleMerge::RotationScaleMerge(const DiffractionExperiment &experiment,
|
|
std::vector<IntegrationOutcome> &partial_outcomes,
|
|
std::optional<UnitCell> reference_cell,
|
|
int scaling_iterations, float ice_ring_half_width_q,
|
|
size_t nthreads, Logger &logger,
|
|
std::string observation_dump_path)
|
|
: x(experiment), partials_out(partial_outcomes), reference_cell(std::move(reference_cell)),
|
|
nthreads(nthreads == 0 ? std::thread::hardware_concurrency() : nthreads), logger(logger),
|
|
observation_dump_path(std::move(observation_dump_path)) {
|
|
const auto s = x.GetScalingSettings();
|
|
min_partiality = s.GetMinPartiality();
|
|
d_min_limit = s.GetHighResolutionLimit_A();
|
|
merge_friedel = s.GetMergeFriedel();
|
|
capture_uncertainty_coeff = s.GetCaptureUncertaintyCoeff();
|
|
min_captured_fraction = s.GetMinCapturedFraction();
|
|
min_cc_for_image = s.GetMinCCForImage();
|
|
search_min_zeta = s.GetSearchMinZeta();
|
|
reject_nsigma = s.GetOutlierRejectNsigma();
|
|
reject_outliers = reject_nsigma > 0.0;
|
|
rfree_fraction = s.GetRfreeFraction();
|
|
scale_fulls = s.GetScaleFulls();
|
|
// Decay + absorption correction surfaces are one master toggle (on by default; both cross-validated,
|
|
// so a no-op when their systematic is absent).
|
|
refine_decay_b = s.GetCorrectionSurfaces();
|
|
absorption_iter = s.GetCorrectionSurfaces() ? s.GetAbsorptionIter() : 0;
|
|
modulation_iter = s.GetCorrectionSurfaces() ? s.GetAbsorptionIter() : 0;
|
|
relative_b_deg = s.GetRelativeBDegrees(); // opt-in per-batch relative-B (0 = off)
|
|
scaling_iter = std::max(1, scaling_iterations);
|
|
resolution_cutoff_method = s.GetResolutionCutoff();
|
|
resolution_cc_target = s.GetResolutionCCTarget();
|
|
report_shell_count = s.GetReportShellCount();
|
|
if (const auto forced = s.GetForcedMosaicity(); forced.has_value())
|
|
mosaicity_deg = *forced;
|
|
else
|
|
mosaicity_deg = s.GetDefaultMosaicity();
|
|
ice_half_width_q = ice_ring_half_width_q;
|
|
}
|
|
|
|
void RotationScaleMerge::Ingest() {
|
|
n_frames = static_cast<int>(partials_out.size());
|
|
|
|
size_t total = 0;
|
|
for (const auto &o : partials_out) total += o.reflections.size();
|
|
|
|
partials.clear();
|
|
partials.reserve(total);
|
|
frame_start.assign(n_frames, 0);
|
|
frame_count.assign(n_frames, 0);
|
|
frame_cell_ok.assign(n_frames, 1);
|
|
g_partial.assign(n_frames, 1.0);
|
|
|
|
const float dist_tol = x.GetIndexingSettings().GetUnitCellDistTolerance();
|
|
const float ang_tol = x.GetIndexingSettings().GetUnitCellAngleTolerance_deg();
|
|
|
|
for (int o = 0; o < n_frames; ++o) {
|
|
frame_start[o] = static_cast<int32_t>(partials.size());
|
|
if (reference_cell) {
|
|
const auto cell = partials_out[o].latt.GetUnitCell();
|
|
frame_cell_ok[o] = cell.is_close(*reference_cell, dist_tol, ang_tol) ? 1 : 0;
|
|
}
|
|
for (const auto &r : partials_out[o].reflections) {
|
|
Obs obs{};
|
|
obs.h = r.h; obs.k = r.k; obs.l = r.l;
|
|
obs.I = r.I; obs.sigma = r.sigma; obs.d = r.d; obs.rlp = r.rlp;
|
|
obs.partiality = r.partiality; obs.zeta = r.zeta; obs.delta_phi = r.delta_phi_deg; obs.bkg = r.bkg;
|
|
obs.px = r.predicted_x; obs.py = r.predicted_y;
|
|
obs.image_number = r.image_number;
|
|
obs.frame = o;
|
|
obs.on_ice = r.on_ice_ring ? 1 : 0;
|
|
obs.corr = r.image_scale_corr;
|
|
obs.group = -1;
|
|
partials.push_back(obs);
|
|
}
|
|
frame_count[o] = static_cast<int32_t>(partials.size()) - frame_start[o];
|
|
}
|
|
|
|
// Per-obs AcceptReflection finiteness (immutable) - lets ComputeAsuGroups stamp the ASU-group id per
|
|
// obs from a flat 1-byte array instead of re-reading the fat Obs struct for every space group.
|
|
finite_ok.resize(partials.size());
|
|
for (size_t i = 0; i < partials.size(); ++i) {
|
|
const auto &o = partials[i];
|
|
finite_ok[i] = (std::isfinite(o.I) && std::isfinite(o.rlp) && o.rlp != 0.0f
|
|
&& std::isfinite(o.sigma) && o.sigma > 0.0f) ? 1 : 0;
|
|
}
|
|
|
|
// Sort ONCE by (raw h,k,l, image_number) and split into raw-hkl runs. This is the one expensive sort;
|
|
// both the 3D combine (event split) and the per-space-group ASU grouping reuse this order.
|
|
// Sorting an index array whose comparator dereferences the 72-byte Obs is a cache miss per
|
|
// comparison over a multi-GB array, so sort a packed copy of the key fields instead. Started from
|
|
// the identity order with the same comparisons in the same order, std::sort takes exactly the same
|
|
// branches and produces exactly the same permutation.
|
|
struct SortKey { int32_t h, k, l; float image_number; int32_t idx; };
|
|
perm.resize(partials.size());
|
|
rawrun_start.clear(); rawrun_count.clear();
|
|
rawrun_h.clear(); rawrun_k.clear(); rawrun_l.clear(); rawrun_d.clear();
|
|
{
|
|
std::vector<SortKey> keys(partials.size());
|
|
for (size_t i = 0; i < partials.size(); ++i) {
|
|
const auto &o = partials[i];
|
|
keys[i] = SortKey{o.h, o.k, o.l, o.image_number, static_cast<int32_t>(i)};
|
|
}
|
|
std::sort(keys.begin(), keys.end(), [](const SortKey &a, const SortKey &b) {
|
|
if (a.h != b.h) return a.h < b.h;
|
|
if (a.k != b.k) return a.k < b.k;
|
|
if (a.l != b.l) return a.l < b.l;
|
|
return a.image_number < b.image_number;
|
|
});
|
|
for (size_t i = 0; i < keys.size(); ++i) perm[i] = keys[i].idx;
|
|
for (int i = 0; i < static_cast<int>(keys.size()); ) {
|
|
const SortKey k0 = keys[i];
|
|
int j = i;
|
|
float d = NAN;
|
|
while (j < static_cast<int>(keys.size())) {
|
|
if (keys[j].h != k0.h || keys[j].k != k0.k || keys[j].l != k0.l) break;
|
|
if (!std::isfinite(d)) {
|
|
const float dj = partials[keys[j].idx].d; // resolution: only until one is usable
|
|
if (std::isfinite(dj) && dj > 0.0f) d = dj;
|
|
}
|
|
++j;
|
|
}
|
|
rawrun_start.push_back(i);
|
|
rawrun_count.push_back(j - i);
|
|
rawrun_h.push_back(k0.h); rawrun_k.push_back(k0.k); rawrun_l.push_back(k0.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
|
|
}
|
|
|
|
namespace {
|
|
// Rotate v about a unit axis by angle (Rodrigues).
|
|
Coord RotateAbout(const Coord &v, const Coord &axis, double angle) {
|
|
const double c = std::cos(angle), s = std::sin(angle);
|
|
return v * static_cast<float>(c) + (axis % v) * static_cast<float>(s)
|
|
+ axis * static_cast<float>((axis * v) * (1.0 - c));
|
|
}
|
|
|
|
// The rotation angle (rad) at which reciprocal vector p0 meets the Ewald sphere, in the frame's
|
|
// own reference. Same solve as BraggPredictionRot, kept to the branch nearest phi_near so the
|
|
// reflection keeps the rocking-curve position it was integrated at.
|
|
std::optional<float> SolvePhi(const Coord &p0, const Coord &S0, const Coord &m1, const Coord &m2,
|
|
const Coord &m3, float m2_S0, float m3_S0, float phi_near) {
|
|
const float p0_sq = p0 * p0;
|
|
if (!(p0_sq > 0.0f))
|
|
return std::nullopt;
|
|
const float p0_m1 = p0 * m1, p0_m2 = p0 * m2, p0_m3 = p0 * m3;
|
|
const float rho_sq = p0_sq - p0_m2 * p0_m2;
|
|
const float p_m3 = (-p0_sq / 2 - p0_m2 * m2_S0) / m3_S0;
|
|
if (!(rho_sq > p_m3 * p_m3))
|
|
return std::nullopt;
|
|
const float root = std::sqrt(rho_sq - p_m3 * p_m3);
|
|
|
|
std::optional<float> best;
|
|
for (const float p_m1 : {root, -root}) {
|
|
const float cosphi = (p_m1 * p0_m1 + p_m3 * p0_m3) / rho_sq;
|
|
const float sinphi = (p_m1 * p0_m3 - p_m3 * p0_m1) / rho_sq;
|
|
const float phi = -std::atan2(sinphi, cosphi);
|
|
if (!best.has_value() || std::fabs(phi - phi_near) < std::fabs(*best - phi_near))
|
|
best = phi;
|
|
}
|
|
return best;
|
|
}
|
|
} // namespace
|
|
|
|
// Smooth the per-frame crystal orientation in frame order, then recompute every partial's delta_phi
|
|
// from the smoothed lattice.
|
|
//
|
|
// Geometry is re-refined independently on EVERY frame, against that frame's spots alone - as few as
|
|
// a dozen on a sparse crystal. Measured over ten datasets the per-frame orientation carries two
|
|
// components: a slow drift that is real (rugnux and XDS agree to R^2 0.83-0.88 on the two crystals
|
|
// that genuinely slip, by 1.5 and 0.54 deg), and a fast jitter that is fit noise, scaling with
|
|
// spots-per-frame at exponent -0.79 where counting noise alone would give -0.5. The jitter is worth
|
|
// 1-8% on merged intensities, 24% on the sparsest crystal.
|
|
//
|
|
// It cannot simply be refined less. Batteries: no per-image refinement at all loses six space groups
|
|
// and a whole crystal; dropping the beam-centre terms holds the space groups but is worse on 31 of
|
|
// 37 crystals. The freedom is earning its keep, so the fix is to keep it and suppress only the
|
|
// frequency band that cannot be physical - a crystal does not re-orient and snap back frame to frame.
|
|
//
|
|
// The window is capped: cross-validation scores how well the neighbours PREDICT a frame's
|
|
// orientation, which on a crystal that barely drifts keeps improving with width, but the per-frame
|
|
// fit is also absorbing a real per-frame systematic (dropping its beam-centre terms costs 31 of 37
|
|
// crystals), and smoothing too wide destroys that. The cap keeps the absorption local.
|
|
//
|
|
// The window is chosen per dataset by leave-one-out cross-validation rather than fixed, because the
|
|
// two components vary hugely between crystals (drift 0.018 to 1.288 deg, jitter 0.005 to 0.221): one
|
|
// window would over-smooth a fast-slipping crystal and under-smooth a sparse one. The frame's own
|
|
// value is excluded when scoring, so a window is only preferred if it PREDICTS that frame better,
|
|
// which is exactly the noise-versus-drift trade.
|
|
void RotationScaleMerge::SmoothGeometry() {
|
|
const auto gon = x.GetGoniometer();
|
|
if (!gon || n_frames < 5)
|
|
return;
|
|
|
|
const Coord axis = gon->GetAxis().Normalize();
|
|
const double increment_rad = gon->GetIncrement_deg() * PI / 180.0;
|
|
if (!(std::fabs(increment_rad) > 0.0))
|
|
return;
|
|
|
|
// De-rotate each frame's lattice back to a common reference, so what is left is the crystal
|
|
// orientation itself rather than the goniometer turning. Sign checked against the data: undoing
|
|
// the turn drops the median frame-to-frame change of a lattice vector 0.170 -> 0.0066 A on a
|
|
// well-conditioned crystal (a factor of 26); getting it backwards doubles it instead.
|
|
std::vector<std::array<Coord, 3>> orient(n_frames);
|
|
std::vector<uint8_t> have(n_frames, 0);
|
|
for (int f = 0; f < n_frames; ++f) {
|
|
const auto &l = partials_out[f].latt;
|
|
const Coord v[3] = {l.Vec0(), l.Vec1(), l.Vec2()};
|
|
if (!std::isfinite(v[0].x) || !std::isfinite(v[1].x) || !std::isfinite(v[2].x))
|
|
continue;
|
|
// A frame that never indexed keeps a default-constructed lattice, whose vectors are all zero
|
|
// - finite, so the test above lets it through. Averaging that in pulls its neighbours'
|
|
// orientation towards the origin, and scoring it makes the cross-validation below compare
|
|
// windows on noise. Require a real cell.
|
|
if (!(l.CalcVolume() > 0.0f))
|
|
continue;
|
|
for (int j = 0; j < 3; ++j)
|
|
orient[f][j] = RotateAbout(v[j], axis, f * increment_rad);
|
|
have[f] = 1;
|
|
}
|
|
|
|
auto window_mean = [&](int f, int half, bool exclude_self) -> std::optional<std::array<Coord, 3>> {
|
|
std::array<Coord, 3> sum{Coord{0, 0, 0}, Coord{0, 0, 0}, Coord{0, 0, 0}};
|
|
int n = 0;
|
|
for (int g = std::max(0, f - half); g <= std::min(n_frames - 1, f + half); ++g) {
|
|
if (!have[g] || (exclude_self && g == f))
|
|
continue;
|
|
for (int j = 0; j < 3; ++j)
|
|
sum[j] = sum[j] + orient[g][j];
|
|
++n;
|
|
}
|
|
if (n == 0)
|
|
return std::nullopt;
|
|
for (int j = 0; j < 3; ++j)
|
|
sum[j] = sum[j] * (1.0f / n);
|
|
return sum;
|
|
};
|
|
|
|
// Pick the half-window that best predicts each frame's own orientation from its neighbours.
|
|
int best_half = 0;
|
|
double best_score = std::numeric_limits<double>::infinity();
|
|
for (const int half : {1, 2, 4, 7, 12, 20}) {
|
|
double score = 0.0;
|
|
int n = 0;
|
|
for (int f = 0; f < n_frames; ++f) {
|
|
if (!have[f])
|
|
continue;
|
|
const auto pred = window_mean(f, half, true);
|
|
if (!pred)
|
|
continue;
|
|
for (int j = 0; j < 3; ++j) {
|
|
const Coord d = orient[f][j] - (*pred)[j];
|
|
score += d * d;
|
|
}
|
|
++n;
|
|
}
|
|
if (n > 0 && score / n < best_score) {
|
|
best_score = score / n;
|
|
best_half = half;
|
|
}
|
|
}
|
|
if (best_half == 0)
|
|
return;
|
|
|
|
std::vector<std::array<Coord, 3>> smoothed(n_frames);
|
|
std::vector<uint8_t> ok(n_frames, 0);
|
|
for (int f = 0; f < n_frames; ++f) {
|
|
if (!have[f])
|
|
continue;
|
|
if (const auto m = window_mean(f, best_half, false)) {
|
|
for (int j = 0; j < 3; ++j)
|
|
smoothed[f][j] = RotateAbout((*m)[j], axis, -f * increment_rad);
|
|
ok[f] = 1;
|
|
}
|
|
}
|
|
|
|
// Recompute delta_phi from the smoothed lattice. Everything downstream - partiality, the 3D
|
|
// combine, the per-frame scale - is rebuilt from it, and this runs before the GPU upload, so the
|
|
// device path picks it up with no separate kernel.
|
|
const Coord S0 = x.GetScatteringVector();
|
|
const Coord m2 = axis;
|
|
const Coord m1 = (m2 % S0).Normalize();
|
|
const Coord m3 = (m1 % m2).Normalize();
|
|
const float m2_S0 = m2 * S0, m3_S0 = m3 * S0;
|
|
|
|
std::atomic<int64_t> changed{0};
|
|
ParallelChunks(static_cast<int>(partials.size()), nthreads, [&](int lo, int hi) {
|
|
int64_t local = 0;
|
|
for (int i = lo; i < hi; ++i) {
|
|
auto &o = partials[i];
|
|
if (o.frame < 0 || o.frame >= n_frames || !ok[o.frame])
|
|
continue;
|
|
const CrystalLattice sl(smoothed[o.frame][0], smoothed[o.frame][1], smoothed[o.frame][2]);
|
|
const Coord p0 = sl.Astar() * static_cast<float>(o.h) + sl.Bstar() * static_cast<float>(o.k)
|
|
+ sl.Cstar() * static_cast<float>(o.l);
|
|
const float phi_near = o.delta_phi * static_cast<float>(PI) / 180.0f;
|
|
const auto phi = SolvePhi(p0, S0, m1, m2, m3, m2_S0, m3_S0, phi_near);
|
|
if (!phi || !std::isfinite(*phi))
|
|
continue;
|
|
o.delta_phi = *phi * 180.0f / static_cast<float>(PI);
|
|
++local;
|
|
}
|
|
changed += local;
|
|
});
|
|
|
|
logger.Info("Smoothed per-frame geometry over +-{} frames (chosen by cross-validation); "
|
|
"recomputed delta_phi for {} of {} partials", best_half, changed.load(),
|
|
partials.size());
|
|
}
|
|
|
|
void RotationScaleMerge::SmoothMosaicityAndPartiality() {
|
|
SmoothGeometry();
|
|
|
|
// 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;
|
|
};
|
|
// The observations any of the passes below can use, in fulls order, split by frame parity. Every
|
|
// pass then walks one list instead of re-testing (and re-reading) the whole of fulls - the
|
|
// cross-validated halves cost half a pass rather than a whole one with a skip. The lists keep the
|
|
// fulls order, so each sum below is formed from exactly the same terms in exactly the same order.
|
|
std::vector<int32_t> idx_all, idx_even, idx_odd;
|
|
idx_all.reserve(fulls.size());
|
|
for (size_t i = 0; i < fulls.size(); ++i) {
|
|
if (!usable(fulls[i]) || cell[i] < 0) continue;
|
|
idx_all.push_back(static_cast<int32_t>(i));
|
|
(fulls[i].frame & 1 ? idx_odd : idx_even).push_back(static_cast<int32_t>(i));
|
|
}
|
|
auto subset = [&](int parity) -> const std::vector<int32_t> & {
|
|
return parity < 0 ? idx_all : (parity ? idx_odd : idx_even);
|
|
};
|
|
// Per-group reference accumulators, reused by every pass (n_groups is large, so re-allocating them
|
|
// ~22x per surface is ~22x the page faults for nothing).
|
|
std::vector<double> sw(n_groups), swI(n_groups);
|
|
|
|
// 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> {
|
|
const std::vector<int32_t> &sel = subset(parity);
|
|
std::vector<double> A(ncell, 1.0);
|
|
for (int it = 0; it < n_iter; ++it) {
|
|
std::fill(sw.begin(), sw.end(), 0.0);
|
|
std::fill(swI.begin(), swI.end(), 0.0);
|
|
for (const int32_t i : sel) {
|
|
const Obs &o = fulls[i];
|
|
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 (const int32_t i : sel) {
|
|
const Obs &o = fulls[i];
|
|
if (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 {
|
|
const std::vector<int32_t> &sel = subset(parity);
|
|
std::fill(sw.begin(), sw.end(), 0.0);
|
|
std::fill(swI.begin(), swI.end(), 0.0);
|
|
for (const int32_t i : sel) {
|
|
const Obs &o = fulls[i];
|
|
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 (const int32_t i : sel) {
|
|
const Obs &o = fulls[i];
|
|
if (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);
|
|
|
|
// No collapsed-scale handling here: ReplaceCollapsedScales has already run on g, so every fitted
|
|
// entry is credible. That matters because the window average is a GEOMETRIC mean - a scale
|
|
// collapsing toward zero would not just corrupt its own frame, its logarithm would drag the whole
|
|
// window down, and where the only fitted frames in a window ARE the collapsed ones the mean would
|
|
// average the fault with itself and make it permanent.
|
|
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);
|
|
}
|
|
}
|
|
|
|
// Replace any per-frame scale that fitted far below the run's median with that median, and report the
|
|
// per-frame ratio by which the frame's corr has to be rescaled to match (1 where nothing changed).
|
|
// Returns false when there was nothing to replace.
|
|
//
|
|
// A collapsed G multiplies the frame's intensities by 1/G AND its sigmas by the same factor, so no
|
|
// downstream test can recognise it: the merge's n-sigma outlier cut scales with the very number that is
|
|
// wrong. Only a total collapse gives itself away, by overflowing corr to infinity.
|
|
bool RotationScaleMerge::ReplaceCollapsedScales(const std::vector<uint8_t> &fitted_mask,
|
|
std::vector<double> &g,
|
|
std::vector<uint8_t> &apply,
|
|
std::vector<double> &ratio) const {
|
|
const int n = static_cast<int>(g.size());
|
|
apply.assign(n, 0);
|
|
ratio.assign(n, 1.0);
|
|
|
|
std::vector<double> fitted;
|
|
fitted.reserve(n);
|
|
for (int f = 0; f < n; ++f)
|
|
if (fitted_mask[f] && std::isfinite(g[f]) && g[f] > 0.0)
|
|
fitted.push_back(g[f]);
|
|
if (fitted.empty())
|
|
return false;
|
|
|
|
const size_t mid = fitted.size() / 2;
|
|
std::nth_element(fitted.begin(), fitted.begin() + mid, fitted.end());
|
|
const double g_typ = fitted[mid];
|
|
const double g_floor = g_typ * MIN_CREDIBLE_SCALE_RATIO;
|
|
|
|
int n_replaced = 0;
|
|
double worst = 1.0;
|
|
for (int f = 0; f < n; ++f) {
|
|
if (!fitted_mask[f] || !std::isfinite(g[f]) || g[f] <= 0.0 || g[f] >= g_floor)
|
|
continue;
|
|
// corr = rlp / (partiality * G), so swapping G for g_typ scales this frame's corr by G / g_typ.
|
|
apply[f] = 1;
|
|
ratio[f] = g[f] / g_typ;
|
|
worst = std::max(worst, g_typ / g[f]);
|
|
g[f] = g_typ;
|
|
++n_replaced;
|
|
}
|
|
if (n_replaced > 0)
|
|
logger.Warning("Replaced the fitted scale of {} frame(s) that came out more than {:.0f}x below "
|
|
"the run median (worst {:.0f}x) - a scale that small amplifies the frame's "
|
|
"intensities, and its sigmas with them, by the same factor",
|
|
n_replaced, 1.0 / MIN_CREDIBLE_SCALE_RATIO, worst);
|
|
return n_replaced > 0;
|
|
}
|
|
|
|
void RotationScaleMerge::SmoothG(std::vector<Obs> &obs, std::vector<double> &g, int window) const {
|
|
const int n = static_cast<int>(g.size());
|
|
std::vector<double> g_smooth;
|
|
ComputeSmoothGWindow(g, window, g_smooth);
|
|
for (auto &o : obs) {
|
|
const int f = o.frame;
|
|
if (!frame_scaled_scratch[f] || !std::isfinite(g[f]) || g[f] <= 0.0 || !std::isfinite(g_smooth[f]))
|
|
continue;
|
|
if (std::isfinite(o.corr))
|
|
o.corr = static_cast<float>(o.corr * (g[f] / g_smooth[f]));
|
|
}
|
|
for (int f = 0; f < n; ++f)
|
|
if (frame_scaled_scratch[f] && std::isfinite(g[f]) && g[f] > 0.0 && std::isfinite(g_smooth[f]))
|
|
g[f] = g_smooth[f];
|
|
}
|
|
|
|
bool RotationScaleMerge::RejectCollapsedFullScales() {
|
|
// The Unity model leaves corr = 1/G, and every full of a frame carries that frame's G. A full whose
|
|
// frame was NOT fitted still carries the 1.0 the combine gave it, on both the host and the device
|
|
// path - those are not measurements of anything and must stay out of the median, or a run with many
|
|
// unfitted frames drags the median toward 1 and the floor with it.
|
|
std::vector<double> g_frame(n_frames, NAN);
|
|
for (const auto &o : fulls)
|
|
if (o.frame >= 0 && o.frame < n_frames && std::isfinite(o.corr) && o.corr > 0.0f
|
|
&& o.corr != 1.0f)
|
|
g_frame[o.frame] = 1.0 / static_cast<double>(o.corr);
|
|
|
|
std::vector<double> fitted;
|
|
fitted.reserve(g_frame.size());
|
|
for (const double gf : g_frame)
|
|
if (std::isfinite(gf) && gf > 0.0)
|
|
fitted.push_back(gf);
|
|
if (fitted.empty())
|
|
return false;
|
|
const size_t mid = fitted.size() / 2;
|
|
std::nth_element(fitted.begin(), fitted.begin() + mid, fitted.end());
|
|
const double fitted_median = fitted[mid];
|
|
const double g_floor = fitted_median * MIN_CREDIBLE_SCALE_RATIO;
|
|
|
|
int n_rejected = 0;
|
|
double worst = 1.0;
|
|
for (auto &o : fulls) {
|
|
if (o.frame < 0 || o.frame >= n_frames)
|
|
continue;
|
|
const double gf = g_frame[o.frame];
|
|
if (!std::isfinite(gf) || gf >= g_floor)
|
|
continue;
|
|
worst = std::max(worst, fitted_median / gf);
|
|
o.corr = 1.0f; // as if the frame's scale had never been fitted
|
|
++n_rejected;
|
|
}
|
|
if (n_rejected > 0)
|
|
logger.Warning("Rejected the fitted scale of {} full(s) whose frame scaled to less than {:.0f}x "
|
|
"below the run median (worst {:.0f}x amplification); those frames merge unscaled",
|
|
n_rejected, 1.0 / MIN_CREDIBLE_SCALE_RATIO, worst);
|
|
return n_rejected > 0;
|
|
}
|
|
|
|
void RotationScaleMerge::Combine() {
|
|
fulls.clear();
|
|
g_full.assign(n_frames, 1.0);
|
|
|
|
// Combine one raw-hkl run into fulls, appended to `out`. Independent per run (the events of one hkl
|
|
// touch no shared state), so runs parallelise cleanly. Returns the number of usable partials seen.
|
|
// `dump` (serial path only) writes each emitted full for the diagnostic observation dump.
|
|
auto process_rawrun = [&](int r, std::vector<Obs> &out, std::ofstream *dump) -> size_t {
|
|
const int lo = rawrun_start[r], hi = rawrun_start[r] + rawrun_count[r];
|
|
std::vector<int32_t> ev; // usable perm-indices of this raw hkl, in image-number order
|
|
ev.reserve(hi - lo);
|
|
for (int p = lo; p < hi; ++p) {
|
|
const auto &o = partials[perm[p]];
|
|
if (!std::isfinite(o.corr) || o.corr <= 0.0f) continue;
|
|
if (!std::isfinite(o.I) || !std::isfinite(o.sigma) || o.sigma <= 0.0f) continue;
|
|
ev.push_back(perm[p]);
|
|
}
|
|
const int group = rawrun_group[r];
|
|
|
|
size_t i = 0;
|
|
while (i < ev.size()) {
|
|
size_t kk = i + 1;
|
|
float last_frame = partials[ev[i]].image_number;
|
|
while (kk < ev.size()) {
|
|
const float frame = partials[ev[kk]].image_number;
|
|
if (frame - last_frame > MAX_FRAME_GAP) break;
|
|
last_frame = frame;
|
|
++kk;
|
|
}
|
|
|
|
double pooled_bkg = 0.0;
|
|
int n_pool = 0;
|
|
for (size_t m = i; m < kk; ++m) {
|
|
const float b = partials[ev[m]].bkg;
|
|
if (std::isfinite(b)) { pooled_bkg += b; ++n_pool; }
|
|
}
|
|
pooled_bkg = n_pool > 0 ? pooled_bkg / n_pool : 0.0;
|
|
auto pooled_I = [&](const Obs &r2) {
|
|
const double n_bkg = std::max(0.0, static_cast<double>(r2.sigma) * r2.sigma - r2.I)
|
|
/ std::max(r2.bkg, 1.0f);
|
|
return static_cast<double>(r2.I) + n_bkg * (static_cast<double>(r2.bkg) - pooled_bkg);
|
|
};
|
|
|
|
double sum_w = 0.0, sum_wI = 0.0, sum_partiality = 0.0;
|
|
float d = NAN;
|
|
int peak_outcome = partials[ev[i]].frame;
|
|
float peak_frame = partials[ev[i]].image_number;
|
|
float peak_px = partials[ev[i]].px, peak_py = partials[ev[i]].py;
|
|
float peak_partiality = -1.0f;
|
|
const bool on_ice = partials[ev[i]].on_ice;
|
|
for (size_t m = i; m < kk; ++m) {
|
|
const auto &r2 = partials[ev[m]];
|
|
const double sigma_corr = static_cast<double>(r2.sigma) * r2.corr;
|
|
const double w = 1.0 / (sigma_corr * sigma_corr);
|
|
sum_w += w;
|
|
sum_wI += w * pooled_I(r2) * r2.corr;
|
|
sum_partiality += r2.partiality;
|
|
if (r2.partiality > peak_partiality) {
|
|
peak_partiality = r2.partiality;
|
|
peak_outcome = r2.frame;
|
|
peak_frame = r2.image_number;
|
|
peak_px = r2.px; peak_py = r2.py;
|
|
}
|
|
if (!std::isfinite(d) && std::isfinite(r2.d) && r2.d > 0.0f) d = r2.d;
|
|
}
|
|
double F = sum_wI / sum_w;
|
|
|
|
for (int iter = 0; iter < 3; ++iter) {
|
|
sum_w = 0.0; sum_wI = 0.0;
|
|
for (size_t m = i; m < kk; ++m) {
|
|
const auto &r2 = partials[ev[m]];
|
|
const double corr = r2.corr;
|
|
const double I_corr = pooled_I(r2) * corr;
|
|
const double sigma_corr = static_cast<double>(r2.sigma) * corr;
|
|
// max(0, I): a down-fluctuated partial carries no Poisson signal to remove, and
|
|
// removing a negative one inflates the background part instead of leaving it alone.
|
|
const double bkg_var = sigma_corr * sigma_corr - corr * std::max(0.0, I_corr);
|
|
double var = std::max(0.0, bkg_var) + corr * std::max(0.0, F);
|
|
if (!(var > 0.0)) var = sigma_corr * sigma_corr;
|
|
const double w = 1.0 / var;
|
|
sum_w += w;
|
|
sum_wI += w * I_corr;
|
|
}
|
|
F = sum_wI / sum_w;
|
|
}
|
|
|
|
const int n_frames_event = static_cast<int>(kk - i);
|
|
i = kk;
|
|
|
|
if (sum_w <= 0.0 || sum_partiality < min_partiality || sum_partiality < min_captured_fraction)
|
|
continue;
|
|
|
|
double sigma_full = 1.0 / std::sqrt(sum_w);
|
|
if (capture_uncertainty_coeff > 0.0) {
|
|
const double frac = std::min(1.0, sum_partiality);
|
|
const double extra = capture_uncertainty_coeff * (1.0 - frac) * std::max(0.0, F);
|
|
sigma_full = std::sqrt(sigma_full * sigma_full + extra * extra);
|
|
}
|
|
|
|
Obs full{};
|
|
full.h = rawrun_h[r]; full.k = rawrun_k[r]; full.l = rawrun_l[r];
|
|
full.I = static_cast<float>(F);
|
|
full.sigma = static_cast<float>(sigma_full);
|
|
full.d = d;
|
|
full.rlp = 1.0f;
|
|
full.partiality = 1.0f;
|
|
full.corr = 1.0f;
|
|
full.image_number = peak_frame;
|
|
full.px = peak_px; full.py = peak_py;
|
|
full.frame = peak_outcome;
|
|
full.on_ice = on_ice ? 1 : 0;
|
|
full.group = group; // the raw hkl's ASU group for the current space group (<0 = absent)
|
|
out.push_back(full);
|
|
|
|
if (dump != nullptr)
|
|
*dump << full.h << ' ' << full.k << ' ' << full.l << ' ' << full.I << ' '
|
|
<< full.sigma << ' ' << d << ' ' << n_frames_event << ' ' << sum_partiality << ' '
|
|
<< static_cast<int>(peak_frame) << ' ' << full.px << ' ' << full.py << '\n';
|
|
}
|
|
return ev.size();
|
|
};
|
|
|
|
const int n_run = static_cast<int>(rawrun_start.size());
|
|
size_t n_used = 0;
|
|
|
|
if (!observation_dump_path.empty() || nthreads <= 1) {
|
|
// Serial (diagnostic dump needs a single writer).
|
|
std::ofstream dump;
|
|
if (!observation_dump_path.empty()) {
|
|
dump.open(observation_dump_path);
|
|
// Columns are enough to rescale/merge the combined fulls in an external program: peak_frame
|
|
// is a per-full batch (ZD) and px/py the predicted detector position (XD/YD).
|
|
dump << "# h k l I sigma d n_frames captured_fraction peak_frame px py\n";
|
|
}
|
|
fulls.reserve(n_run);
|
|
for (int r = 0; r < n_run; ++r)
|
|
n_used += process_rawrun(r, fulls, dump.is_open() ? &dump : nullptr);
|
|
} else {
|
|
// Parallel over contiguous rawrun chunks; concatenate the per-thread fulls in run order so the
|
|
// result is deterministic.
|
|
const int nt = static_cast<int>(std::min(nthreads, static_cast<size_t>(n_run)));
|
|
const int chunk = (n_run + nt - 1) / nt;
|
|
std::vector<std::vector<Obs>> part(nt);
|
|
std::vector<size_t> used(nt, 0);
|
|
std::vector<std::future<void>> futures;
|
|
futures.reserve(nt);
|
|
for (int t = 0; t < nt; ++t) {
|
|
const int r0 = t * chunk, r1 = std::min(n_run, r0 + chunk);
|
|
if (r0 >= r1) break;
|
|
futures.emplace_back(std::async(std::launch::async, [&, t, r0, r1] {
|
|
part[t].reserve(r1 - r0);
|
|
for (int r = r0; r < r1; ++r) used[t] += process_rawrun(r, part[t], nullptr);
|
|
}));
|
|
}
|
|
for (auto &f : futures) f.get();
|
|
size_t total = 0;
|
|
for (int t = 0; t < nt; ++t) { total += part[t].size(); n_used += used[t]; }
|
|
fulls.reserve(total);
|
|
for (int t = 0; t < nt; ++t)
|
|
fulls.insert(fulls.end(), part[t].begin(), part[t].end());
|
|
}
|
|
|
|
SortFullsByFrame();
|
|
logger.Info("3D combine: {} fulls from {} partials", fulls.size(), n_used);
|
|
}
|
|
|
|
void RotationScaleMerge::SortFullsByFrame() {
|
|
// Sort the fulls by their (peak) frame and build per-frame CSR ranges, so the scale-fulls step can
|
|
// fit a per-frame G by slicing contiguous ranges (the same layout the partials use).
|
|
std::sort(fulls.begin(), fulls.end(),
|
|
[](const Obs &a, const Obs &b) { return a.frame < b.frame; });
|
|
fulls_frame_start.assign(n_frames, 0);
|
|
fulls_frame_count.assign(n_frames, 0);
|
|
for (int i = 0; i < static_cast<int>(fulls.size()); ) {
|
|
const int f = fulls[i].frame;
|
|
fulls_frame_start[f] = i;
|
|
int j = i;
|
|
while (j < static_cast<int>(fulls.size()) && fulls[j].frame == f) ++j;
|
|
fulls_frame_count[f] = j - i;
|
|
i = j;
|
|
}
|
|
}
|
|
|
|
void RotationScaleMerge::ComputePerFrameCC(const std::vector<double> &partial_group_mean,
|
|
std::vector<double> &cc, std::vector<int64_t> &cc_n) const {
|
|
// Per-frame CC vs the merged reference (CalculateGlobalCC), computed once now (not every iteration).
|
|
cc.assign(n_frames, NAN);
|
|
cc_n.assign(n_frames, 0);
|
|
ParallelFor(n_frames, nthreads, [&](int f) {
|
|
double sx = 0, sy = 0, sx2 = 0, sy2 = 0, sxy = 0;
|
|
size_t n = 0;
|
|
const int lo = frame_start[f], hi = frame_start[f] + frame_count[f];
|
|
for (int i = lo; i < hi; ++i) {
|
|
const auto &o = partials[i];
|
|
if (o.on_ice) continue;
|
|
if (o.group < 0) continue;
|
|
if (o.partiality < min_partiality) continue;
|
|
if (!std::isfinite(o.I) || !std::isfinite(o.corr) || o.corr <= 0.0f) continue;
|
|
if (!std::isfinite(o.sigma) || o.sigma <= 0.0f) continue;
|
|
const double mean = partial_group_mean[o.group];
|
|
if (!std::isfinite(mean)) continue;
|
|
const double img = static_cast<double>(o.I) * o.corr;
|
|
sx += img; sy += mean; sx2 += img * img; sy2 += mean * mean; sxy += img * mean;
|
|
++n;
|
|
}
|
|
if (n < MIN_REFLECTIONS) return;
|
|
const double nd = static_cast<double>(n);
|
|
const double cov = sxy - sx * sy / nd;
|
|
const double vx = sx2 - sx * sx / nd;
|
|
const double vy = sy2 - sy * sy / nd;
|
|
if (vx > 0.0 && vy > 0.0) { cc[f] = cov / std::sqrt(vx * vy); cc_n[f] = static_cast<int64_t>(n); }
|
|
});
|
|
}
|
|
|
|
// Write the per-frame G / CC / mosaicity (from the given cc/cc_n) back onto the partials for the offline
|
|
// per-image scaling table. cc/cc_n are computed on the host (ComputePerFrameCC) or GPU (ComputePartialCC).
|
|
void RotationScaleMerge::FinalizePerFrameScale(const std::vector<double> &cc, const std::vector<int64_t> &cc_n,
|
|
const std::vector<uint8_t> &frame_scaled) {
|
|
for (int f = 0; f < n_frames; ++f) {
|
|
auto &o = partials_out[f];
|
|
if (frame_scaled[f]) {
|
|
o.image_scale_g = static_cast<float>(g_partial[f]);
|
|
o.mosaicity_deg = (f < static_cast<int>(mos_smooth.size()) && std::isfinite(mos_smooth[f]))
|
|
? mos_smooth[f] : static_cast<float>(mosaicity_deg);
|
|
if (std::isfinite(cc[f])) { o.image_scale_cc = static_cast<float>(cc[f]); o.image_scale_cc_n = cc_n[f]; }
|
|
else { o.image_scale_cc.reset(); o.image_scale_cc_n.reset(); }
|
|
} else {
|
|
o.image_scale_g.reset();
|
|
o.image_scale_cc.reset();
|
|
o.image_scale_cc_n.reset();
|
|
o.mosaicity_deg.reset();
|
|
}
|
|
o.image_scale_wedge_deg.reset();
|
|
}
|
|
}
|
|
|
|
namespace {
|
|
// Possible unique reflections per shell for the completeness column - mirrors CalcPossibleReflections.
|
|
void PossiblePerShell(int space_group_number, const UnitCell &cell, double d_min, double d_max,
|
|
const ResolutionShells &shells, bool merge_friedel, std::vector<int> &possible) {
|
|
gemmi::UnitCell gemmi_cell = cell;
|
|
const gemmi::SpaceGroup *sg = gemmi::find_spacegroup_by_number(space_group_number);
|
|
if (sg == nullptr) return;
|
|
const std::vector<gemmi::Miller> hkls = gemmi::make_miller_vector(gemmi_cell, sg, d_min, d_max, true);
|
|
const gemmi::GroupOps gops = sg->operations();
|
|
CrystalLattice lattice(cell);
|
|
const auto astar = lattice.Astar(), bstar = lattice.Bstar(), cstar = lattice.Cstar();
|
|
for (const auto &hkl : hkls) {
|
|
const auto q = hkl[0] * astar + hkl[1] * bstar + hkl[2] * cstar;
|
|
const auto qlen = q.Length();
|
|
if (qlen < 1e-6) continue;
|
|
const auto shell = shells.GetShell(1.0 / qlen);
|
|
if (!shell.has_value()) continue;
|
|
const int s = *shell;
|
|
if (s >= 0 && s < static_cast<int>(possible.size()))
|
|
possible[s] += (merge_friedel || gops.is_reflection_centric(hkl)) ? 1 : 2;
|
|
}
|
|
}
|
|
}
|
|
|
|
RotationScaleMerge::Result RotationScaleMerge::MergeAndStats(int n_groups, bool for_search,
|
|
const std::vector<char> &masked,
|
|
bool fulls_resident) {
|
|
// A full is usable for the merge / error model if it passes AddImage's filters (with the current
|
|
// ice/masked-ring context). group >= 0 already encodes "not absent and passes AcceptReflection".
|
|
auto masked_ring = [&](const Obs &o) {
|
|
if (masked.empty()) return false;
|
|
const int ring = IceRingIndex(o.d, ice_half_width_q);
|
|
return ring >= 0 && ring < static_cast<int>(masked.size()) && masked[ring];
|
|
};
|
|
auto usable_merge = [&](const Obs &o) {
|
|
if (o.group < 0) return false;
|
|
if (!frame_cell_ok[o.frame]) return false;
|
|
if (!(o.corr > 0.0f) || !std::isfinite(o.corr)) return false;
|
|
if (for_search && o.on_ice) return false;
|
|
if (masked_ring(o)) return false;
|
|
if (o.partiality < min_partiality) return false;
|
|
const float I_corr = o.I * o.corr, sigma_corr = o.sigma * o.corr;
|
|
return std::isfinite(I_corr) && std::isfinite(sigma_corr) && sigma_corr > 0.0f;
|
|
};
|
|
|
|
// The em-stats / samples / merge-accumulate / R_meas reductions run on the resident, scaled fulls
|
|
// (their group CSR is still on the device from scale-fulls) when fulls_resident; the host keeps the
|
|
// I2-sort, the (a,b) fit, the export and the statistics. Outlier rejection runs on the device too,
|
|
// and MergeAccum hands the per-full flags back so the host-side reductions see the same rejections.
|
|
bool use_gpu_merge = false;
|
|
#ifdef JFJOCH_USE_CUDA
|
|
use_gpu_merge = fulls_resident && !fulls.empty();
|
|
std::vector<uint8_t> gpu_masked(masked.begin(), masked.end());
|
|
#endif
|
|
|
|
// ---- Error model: fit dev2 = a*sigma^2 + b^2*<I>^2 from symmetry-equivalent scatter. ----
|
|
std::vector<double> em_mean(n_groups, NAN);
|
|
std::vector<float> reject_median(n_groups, NAN);
|
|
double error_model_a = 1.0, error_model_b = 0.0, error_model_chi2 = 0.0;
|
|
bool error_model_active = false;
|
|
{
|
|
struct Sample { double s2, I2, dev2; };
|
|
std::vector<Sample> samples;
|
|
std::vector<int32_t> cnt(n_groups, 0); // per-group usable count (both paths; feeds reject-median)
|
|
bool did_gpu = false;
|
|
#ifdef JFJOCH_USE_CUDA
|
|
if (use_gpu_merge) {
|
|
const int nf = static_cast<int>(fulls.size());
|
|
std::vector<double> gs2(nf), gI2(nf), gdev2(nf);
|
|
std::vector<uint8_t> gvalid(nf);
|
|
gpu_->MergeEmSamples(for_search, gpu_masked.data(), static_cast<int>(gpu_masked.size()),
|
|
ice_half_width_q, min_partiality, em_mean.data(), cnt.data(),
|
|
gs2.data(), gI2.data(), gdev2.data(), gvalid.data());
|
|
samples.reserve(nf);
|
|
for (int i = 0; i < nf; ++i)
|
|
if (gvalid[i]) samples.push_back({gs2[i], gI2[i], gdev2[i]});
|
|
did_gpu = true;
|
|
}
|
|
#endif
|
|
if (!did_gpu) {
|
|
// Per-group inverse-variance mean over usable fulls (>=2 obs), and the leverage-corrected samples.
|
|
std::vector<double> sw(n_groups, 0.0), swI(n_groups, 0.0);
|
|
for (const auto &o : fulls) {
|
|
if (!usable_merge(o)) continue;
|
|
const double sigma_corr = static_cast<double>(o.sigma) * o.corr;
|
|
const double w = 1.0 / (sigma_corr * sigma_corr);
|
|
sw[o.group] += w; swI[o.group] += w * (static_cast<double>(o.I) * o.corr); cnt[o.group]++;
|
|
}
|
|
for (int g = 0; g < n_groups; ++g)
|
|
if (cnt[g] >= 2 && sw[g] > 0.0) em_mean[g] = swI[g] / sw[g];
|
|
|
|
samples.reserve(fulls.size());
|
|
for (const auto &o : fulls) {
|
|
if (!usable_merge(o) || cnt[o.group] < 2) continue;
|
|
const double mean = em_mean[o.group];
|
|
if (!std::isfinite(mean)) continue;
|
|
const double sigma_corr = static_cast<double>(o.sigma) * o.corr;
|
|
const double s2 = sigma_corr * sigma_corr;
|
|
const double w = 1.0 / s2;
|
|
const double factor = 1.0 - w / sw[o.group];
|
|
if (factor < 0.05) continue;
|
|
const double resid = static_cast<double>(o.I) * o.corr - mean;
|
|
samples.push_back({s2, mean * mean, resid * resid / factor});
|
|
}
|
|
}
|
|
|
|
// Per-group outlier-rejection median of I*corr (host both paths - a per-group median is awkward on
|
|
// the GPU; cheap here, cnt >= 2 filter from the em pass). Fed to the merge accumulate.
|
|
if (reject_outliers) {
|
|
std::vector<std::vector<float>> iv(n_groups);
|
|
for (const auto &o : fulls)
|
|
if (usable_merge(o) && cnt[o.group] >= 2)
|
|
iv[o.group].push_back(o.I * o.corr);
|
|
for (int g = 0; g < n_groups; ++g)
|
|
if (!iv[g].empty()) {
|
|
std::nth_element(iv[g].begin(), iv[g].begin() + iv[g].size() / 2, iv[g].end());
|
|
reject_median[g] = iv[g][iv[g].size() / 2];
|
|
}
|
|
}
|
|
constexpr int n_bins = 16;
|
|
// Fit (a, b) from the intensity-binned median deviations. Factored into a lambda so it can be
|
|
// re-run on a misfit-free pool below; takes the samples by value (it sorts them in place).
|
|
auto fit_ab = [&](std::vector<Sample> smp) {
|
|
if (smp.size() < static_cast<size_t>(8 * n_bins))
|
|
return;
|
|
std::sort(smp.begin(), smp.end(), [](const Sample &a, const Sample &b) { return a.I2 < b.I2; });
|
|
std::vector<double> bs2, bI2, bd2;
|
|
const size_t per = smp.size() / n_bins;
|
|
for (int bin = 0; bin < n_bins; ++bin) {
|
|
const size_t lo = bin * per;
|
|
const size_t hi = (bin == n_bins - 1) ? smp.size() : lo + per;
|
|
std::vector<double> vs2, vI2, vd2;
|
|
for (size_t i = lo; i < hi; ++i) {
|
|
vs2.push_back(smp[i].s2); vI2.push_back(smp[i].I2); vd2.push_back(smp[i].dev2);
|
|
}
|
|
bs2.push_back(median_of(vs2));
|
|
bI2.push_back(median_of(vI2));
|
|
bd2.push_back(median_of(vd2) / CHI2_1_MEDIAN);
|
|
}
|
|
std::vector<double> bd2_sorted = bd2;
|
|
const double dev2_floor = std::max(1e-30, 1e-3 * median_of(bd2_sorted));
|
|
double Ass = 0, AsI = 0, AII = 0, Bs = 0, BI = 0;
|
|
for (int bin = 0; bin < n_bins; ++bin) {
|
|
const double s2 = bs2[bin], I2 = bI2[bin], d2 = bd2[bin];
|
|
const double d2w = std::max(d2, dev2_floor);
|
|
const double wgt = 1.0 / (d2w * d2w);
|
|
Ass += wgt * s2 * s2; AsI += wgt * s2 * I2; AII += wgt * I2 * I2;
|
|
Bs += wgt * s2 * d2; BI += wgt * I2 * d2;
|
|
}
|
|
const double det = Ass * AII - AsI * AsI;
|
|
if (det > 1e-10 * Ass * AII) {
|
|
error_model_a = std::clamp((Bs * AII - BI * AsI) / det, 0.25, 100.0);
|
|
const double b2 = std::max((Ass * BI - AsI * Bs) / det, 0.0);
|
|
error_model_b = std::sqrt(b2);
|
|
error_model_active = true;
|
|
std::vector<double> chi2;
|
|
chi2.reserve(smp.size());
|
|
for (const auto &s : smp) {
|
|
const double v = error_model_a * s.s2 + b2 * s.I2;
|
|
if (v > 0.0) chi2.push_back(s.dev2 / v);
|
|
}
|
|
error_model_chi2 = chi2.empty() ? 0.0 : median_of(chi2) / CHI2_1_MEDIAN;
|
|
}
|
|
};
|
|
fit_ab(samples);
|
|
// Refit on a misfit-free pool: the merge drops symmetry outliers (|I - median| > reject_nsigma *
|
|
// sigma) from the merged intensity, so drop the equivalent samples (dev2 > reject_nsigma^2 * model
|
|
// variance) from the error-model fit too, keeping the fitted sigmas consistent with the reflections
|
|
// that actually survive. Operates on the shared `samples`, so CPU and GPU stay bit-identical.
|
|
if (reject_outliers && error_model_active) {
|
|
const double ns2 = reject_nsigma * reject_nsigma, b2 = error_model_b * error_model_b;
|
|
std::vector<Sample> kept;
|
|
kept.reserve(samples.size());
|
|
for (const auto &s : samples) {
|
|
const double v = error_model_a * s.s2 + b2 * s.I2;
|
|
if (v > 0.0 && s.dev2 <= ns2 * v) kept.push_back(s);
|
|
}
|
|
if (kept.size() >= static_cast<size_t>(8 * n_bins) && kept.size() < samples.size())
|
|
fit_ab(std::move(kept));
|
|
}
|
|
}
|
|
// Asymptotic I/sigma. ISa is by definition the I -> infinity limit of the signal-to-noise, i.e. the
|
|
// reproducibility of the strongest reflections (Diederichs, Acta Cryst. D66 (2010), 733-740). The
|
|
// (a, b) fit above spans the whole intensity range, and a mild excess of scatter at intermediate
|
|
// intensity raises its systematic term `b`, so 1/b understates that limit. Read it instead directly
|
|
// from the strong equivalents: for each well-measured reflection group the counting-subtracted
|
|
// fractional scatter of its symmetry mates estimates the systematic term, and the robust median over
|
|
// strong groups is the asymptote. The I/sigma threshold that selects "strong" is relaxed on weak /
|
|
// radiation-damaged data that has few strong reflections (rather than fall back to the higher
|
|
// whole-range b); it falls back only when even the relaxed set is too small. Host-side over the
|
|
// merged fulls, so CPU and GPU agree.
|
|
double error_model_b_asymptotic = error_model_b;
|
|
if (error_model_active) {
|
|
struct GroupScatter { double sum = 0, sum_sq = 0, sum_var = 0; int n = 0; };
|
|
std::vector<GroupScatter> gs(n_groups);
|
|
for (const auto &o : fulls) {
|
|
if (!usable_merge(o)) continue;
|
|
const double I_corr = static_cast<double>(o.I) * o.corr;
|
|
const double sigma_corr = static_cast<double>(o.sigma) * o.corr;
|
|
auto &g = gs[o.group];
|
|
g.sum += I_corr; g.sum_sq += I_corr * I_corr; g.sum_var += sigma_corr * sigma_corr; ++g.n;
|
|
}
|
|
// Per-group counting-subtracted fractional systematic variance, paired with the group's I/sigma.
|
|
std::vector<std::pair<double, double>> group_scatter; // (systematic b^2, I/sigma)
|
|
for (const auto &g : gs) {
|
|
if (g.n < 5) continue;
|
|
const double mean = g.sum / g.n;
|
|
const double counting = g.sum_var / g.n;
|
|
if (mean <= 0.0 || counting <= 0.0) continue;
|
|
const double variance = (g.sum_sq - g.sum * g.sum / g.n) / (g.n - 1);
|
|
group_scatter.push_back({std::max(variance - counting, 0.0) / (mean * mean),
|
|
mean / std::sqrt(counting)});
|
|
}
|
|
auto asymptote_above = [&](double snr_min, size_t min_groups) -> double {
|
|
std::vector<double> b2;
|
|
for (const auto &[b2_value, snr] : group_scatter)
|
|
if (snr >= snr_min) b2.push_back(b2_value);
|
|
return b2.size() >= min_groups ? std::sqrt(median_of(b2)) : 0.0;
|
|
};
|
|
double b_asy = asymptote_above(20.0, 100); // tight threshold on data that supports it
|
|
if (b_asy <= 0.0) b_asy = asymptote_above(10.0, 50); // relaxed for weak / damaged data
|
|
if (b_asy > 0.0) error_model_b_asymptotic = b_asy;
|
|
}
|
|
// Guard a degenerate low-multiplicity fit: with too few symmetry equivalents both the (a, b) fit and
|
|
// the per-group scatter collapse toward zero, and 1/b then reports an impossibly high asymptotic
|
|
// I/sigma. Real macromolecular data does not exceed ISa ~50; past a generous cap report the asymptote
|
|
// as unmeasured rather than emit a spurious extreme. Only the REPORT is dropped: the sigma floor
|
|
// below keeps using the fitted value, because a degenerate fit is precisely where the raw
|
|
// 1/sqrt(sum_w) is least trustworthy and leaving the merged sigma uncapped would be the opposite of
|
|
// what this guard is for.
|
|
constexpr double MIN_ASYMPTOTIC_B = 0.01; // ISa cap 100
|
|
const double isa_reported = error_model_b_asymptotic >= MIN_ASYMPTOTIC_B
|
|
? 1.0 / error_model_b_asymptotic : 0.0;
|
|
if (error_model_active)
|
|
logger.Info("Error model: a={:.3f} b={:.3f} ISa={:.1f} chi2={:.2f}", error_model_a, error_model_b,
|
|
isa_reported, 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<uint8_t> rejected_obs(fulls.size(), 0); // per-full outlier-rejected flag (both paths)
|
|
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(),
|
|
rejected_obs.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 = isa_reported;
|
|
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
|
|
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.
|
|
FrenchWilsonOptions fw_opts;
|
|
fw_opts.num_threads = static_cast<int>(nthreads);
|
|
ApplyFrenchWilson(result.merged, x.GetSpaceGroupNumber().value_or(1), fw_opts);
|
|
|
|
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);
|
|
|
|
// Put back the observations the previous pass filtered out of its own merge (see where this is
|
|
// filled). This pass decides for itself which ones to drop, and the final merge drops none.
|
|
if (!corr_before_pass_filters.empty()) {
|
|
for (size_t i = 0; i < partials.size(); ++i)
|
|
partials[i].corr = corr_before_pass_filters[i];
|
|
#ifdef JFJOCH_USE_CUDA
|
|
if (gpu_active_)
|
|
gpu_->SetCorr(corr_before_pass_filters.data());
|
|
#endif
|
|
corr_before_pass_filters.clear();
|
|
corr_before_pass_filters.shrink_to_fit();
|
|
}
|
|
|
|
// --- 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;
|
|
|
|
// --- 1b. Guard the per-frame partial scales. Unconditional: the smooth-G window below used to be
|
|
// where this lived, so --smooth-g 0, a dataset with no oscillation width, and every caller
|
|
// that never sets a smoothing range at all (the viewer among them) had no guard whatsoever.
|
|
{
|
|
std::vector<uint8_t> apply;
|
|
std::vector<double> ratio;
|
|
if (ReplaceCollapsedScales(frame_scaled_scratch, g_partial, apply, ratio)) {
|
|
bool applied_on_gpu = false;
|
|
#ifdef JFJOCH_USE_CUDA
|
|
if (gpu_active_) {
|
|
// Same kernel smooth-G uses: multiply the resident corr by a per-frame ratio in place.
|
|
gpu_->SmoothCorr(apply.data(), ratio.data());
|
|
applied_on_gpu = true;
|
|
}
|
|
#endif
|
|
if (!applied_on_gpu)
|
|
for (auto &o : partials)
|
|
if (apply[o.frame] && std::isfinite(o.corr))
|
|
o.corr = static_cast<float>(o.corr * ratio[o.frame]);
|
|
}
|
|
}
|
|
|
|
// --- 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, so bring it back to the host now: the two
|
|
// filters below decide which observations leave the merge and they run on the host, and the CPU
|
|
// combine used by the diagnostic dump reads the same array. Writing the host copy to the device
|
|
// without this step would overwrite the scaling that just ran with the values ingested before it.
|
|
// (The decay/absorption correction surfaces run on the GPU-combined fulls - see below - so they do
|
|
// NOT need the CPU combine.)
|
|
if (gpu_active_) {
|
|
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);
|
|
|
|
// Both filters below remove observations by zeroing corr, which is what takes an observation out of
|
|
// the 3D combine, the merge and the error model alike (excluding them from the ASU grouping is NOT
|
|
// enough - the combine selects on corr, so their intensity would still reach the fulls). Each is a
|
|
// decision of THIS pass and the next pass re-makes it from scratch, so keep what corr was and put it
|
|
// back when that pass starts: nothing else restores it. The scaling loop rewrites corr only for
|
|
// frames it managed to fit, so a frame left with too few well-measured reflections would otherwise
|
|
// carry its dropped observations at zero into every later pass, the final production merge included.
|
|
const bool pass_filters = (for_search && search_min_zeta > 0.0) || min_cc_for_image > 0.0;
|
|
if (pass_filters) {
|
|
corr_before_pass_filters.resize(partials.size());
|
|
for (size_t i = 0; i < partials.size(); ++i)
|
|
corr_before_pass_filters[i] = partials[i].corr;
|
|
}
|
|
bool corr_filtered = false;
|
|
|
|
// --- 2a. On the de-novo search pass only, drop observations whose Lorentz geometry is poor. ---
|
|
if (for_search && search_min_zeta > 0.0) {
|
|
int64_t n_dropped = 0;
|
|
for (auto &o : partials)
|
|
if (!(std::isfinite(o.zeta) && o.zeta >= search_min_zeta)) {
|
|
if (std::isfinite(o.corr) && o.corr > 0.0f) ++n_dropped;
|
|
o.corr = 0.0f;
|
|
}
|
|
corr_filtered = corr_filtered || n_dropped > 0;
|
|
if (n_dropped > 0)
|
|
logger.Info("Space-group search: ignoring {} observations with |zeta| < {:.2f} "
|
|
"(they cross the Ewald sphere near-tangentially and are measured worst)",
|
|
n_dropped, search_min_zeta);
|
|
}
|
|
|
|
// --- 2b. Drop frames that do not agree with the merged reference (--min-image-cc). ---
|
|
if (min_cc_for_image > 0.0) {
|
|
std::vector<uint8_t> reject(n_frames, 0);
|
|
int n_rejected = 0;
|
|
for (int f = 0; f < n_frames; ++f)
|
|
if (std::isfinite(cc[f]) && cc_n[f] >= MIN_REFLECTIONS_FOR_IMAGE_CC && cc[f] < min_cc_for_image) {
|
|
reject[f] = 1;
|
|
++n_rejected;
|
|
}
|
|
if (n_rejected > 0) {
|
|
for (auto &o : partials)
|
|
if (reject[o.frame]) o.corr = 0.0f;
|
|
corr_filtered = true;
|
|
logger.Info("Rejected {} of {} frames correlating below {:.2f} with the merged reference",
|
|
n_rejected, n_frames, min_cc_for_image);
|
|
}
|
|
}
|
|
|
|
#ifdef JFJOCH_USE_CUDA
|
|
// Both filters ran on the host, so hand the result back to the device the combine reads from. One
|
|
// upload for the two of them - the frame rejection used to go to the device on its own path, which
|
|
// left it doing nothing whenever the CPU combine was in use (the diagnostic dump).
|
|
if (gpu_active_ && corr_filtered) {
|
|
std::vector<float> corr(partials.size());
|
|
for (size_t i = 0; i < partials.size(); ++i) corr[i] = partials[i].corr;
|
|
gpu_->SetCorr(corr.data());
|
|
}
|
|
#endif
|
|
|
|
// --- 3. 3D combine of per-frame partials into fulls (fulls inherit their ASU group here). ---
|
|
bool combined_on_gpu = false;
|
|
bool scaled_fulls_on_gpu = false;
|
|
#ifdef JFJOCH_USE_CUDA
|
|
// GPU combine (+ scale-fulls) keeps the fulls resident on the device: combine, then build the frame /
|
|
// ASU-group CSRs on the host from just the small key arrays (a deterministic counting sort - no GPU
|
|
// stable-sort), scale the fulls in place, and download only once. Mirrors Combine() + the Unity
|
|
// scale-fulls loop below. Only the diagnostic dump (serial, one writer) falls back to CPU. The
|
|
// decay/absorption correction surfaces run on the downloaded fulls (px/py is carried through the
|
|
// combine) and their corrected corr is re-uploaded before the resident merge - so they stay on the GPU.
|
|
if (gpu_active_ && observation_dump_path.empty()) {
|
|
// The smoothed corr is already resident (scaling + smooth-G ran on the device, no round-trip).
|
|
const int nf = gpu_->Combine(rawrun_group.data(), min_partiality, capture_uncertainty_coeff,
|
|
min_captured_fraction);
|
|
g_full.assign(n_frames, 1.0);
|
|
|
|
if (scale_fulls && nf > 0) {
|
|
// Frame + group CSRs over the emit-ordered fulls, built by counting sort on the host (stable,
|
|
// deterministic). frame is always in [0, n_frames); group is <0 for absent/out-of-range fulls.
|
|
std::vector<int32_t> ff(nf), fg(nf);
|
|
gpu_->GetFullsKeys(ff.data(), fg.data());
|
|
std::vector<int32_t> f_start(n_frames, 0), f_count(n_frames, 0), f_perm(nf);
|
|
for (int i = 0; i < nf; ++i) ++f_count[ff[i]];
|
|
for (int f = 1; f < n_frames; ++f) f_start[f] = f_start[f - 1] + f_count[f - 1];
|
|
{ std::vector<int32_t> fill = f_start; for (int i = 0; i < nf; ++i) f_perm[fill[ff[i]]++] = i; }
|
|
gpu_->SetFullsFrameCSR(f_perm.data(), nf, f_start.data(), f_count.data());
|
|
|
|
std::vector<int32_t> g_count(n_groups, 0), g_start(n_groups, 0);
|
|
for (int i = 0; i < nf; ++i) if (fg[i] >= 0) ++g_count[fg[i]];
|
|
int acc = 0;
|
|
for (int g = 0; g < n_groups; ++g) { g_start[g] = acc; acc += g_count[g]; }
|
|
std::vector<int32_t> g_perm(acc);
|
|
{ std::vector<int32_t> fill = g_start; for (int i = 0; i < nf; ++i) if (fg[i] >= 0) g_perm[fill[fg[i]]++] = i; }
|
|
gpu_->SetFullsGroups(g_perm.data(), acc, g_start.data(), g_count.data());
|
|
|
|
gpu_->ScaleFulls(scaling_iter, SCALE_ROBUST_K, min_partiality);
|
|
scaled_fulls_on_gpu = true;
|
|
}
|
|
|
|
fulls.assign(nf, Obs{});
|
|
std::vector<int32_t> fh(nf), fk(nf), fl(nf), fframe(nf), fgroup(nf);
|
|
std::vector<float> fI(nf), fsig(nf), fd(nf), fimg(nf), fcorr(nf, 1.0f), fpx(nf), fpy(nf);
|
|
std::vector<uint8_t> fon(nf);
|
|
gpu_->GetFulls(fh.data(), fk.data(), fl.data(), fI.data(), fsig.data(), fd.data(),
|
|
fimg.data(), fframe.data(), fon.data(), fgroup.data());
|
|
gpu_->GetFullsPxPy(fpx.data(), fpy.data());
|
|
if (scaled_fulls_on_gpu) gpu_->GetFullsCorr(fcorr.data());
|
|
for (int i = 0; i < nf; ++i) {
|
|
Obs &o = fulls[i];
|
|
o.h = fh[i]; o.k = fk[i]; o.l = fl[i];
|
|
o.I = fI[i]; o.sigma = fsig[i]; o.d = fd[i];
|
|
o.rlp = 1.0f; o.partiality = 1.0f; o.corr = fcorr[i];
|
|
o.image_number = fimg[i]; o.frame = fframe[i]; o.px = fpx[i]; o.py = fpy[i];
|
|
o.on_ice = fon[i]; o.group = fgroup[i];
|
|
}
|
|
logger.Info("3D combine{} (GPU): {} fulls", scaled_fulls_on_gpu ? " + scale-fulls" : "", nf);
|
|
combined_on_gpu = true;
|
|
}
|
|
#endif
|
|
if (!combined_on_gpu)
|
|
Combine();
|
|
|
|
// --- 4. Scale the fulls (XDS order, Unity model). ---
|
|
if (scale_fulls && !scaled_fulls_on_gpu) {
|
|
std::vector<double> full_mean;
|
|
for (int it = 0; it < scaling_iter; ++it) {
|
|
ReduceGroupMeans(fulls, n_groups, false, {}, full_mean);
|
|
FitPerFrameG(fulls, fulls_frame_start, fulls_frame_count, full_mean, /*unity=*/true, g_full);
|
|
UpdateCorr(fulls, g_full, frame_scaled_scratch);
|
|
}
|
|
logger.Info("Scaled fulls (XDS order, Unity model)");
|
|
}
|
|
const bool rejected_full_scales = scale_fulls && RejectCollapsedFullScales();
|
|
|
|
// --- 4b. Optional correction surfaces (decay = resolution x time; absorption = goniometer-frame
|
|
// diffracted-beam direction), each an alternating multiplicative fit of the fulls' corr against
|
|
// the merged reference. Cheap host loops over the downloaded fulls; applied only on the final
|
|
// in-symmetry merge, never the P1 space-group search pass (corrections there add risk and can
|
|
// perturb the symmetry determination). ---
|
|
[[maybe_unused]] const bool corrections =
|
|
!for_search && (refine_decay_b || absorption_iter > 0 || modulation_iter > 0 || relative_b_deg > 0.0);
|
|
// Radiation-damage monitor: measure the per-batch relative-B on the scaled fulls BEFORE any correction
|
|
// (report-only; captures the full damage signature, not a residual). Skipped on the P1 search pass.
|
|
if (!for_search)
|
|
MeasureRadiationDamageB(n_groups);
|
|
if (!for_search && refine_decay_b)
|
|
RefineDecay(n_groups);
|
|
if (!for_search && relative_b_deg > 0.0)
|
|
RefineRelativeB(n_groups); // per-batch relative-B on top of the single decay slope
|
|
if (!for_search && absorption_iter > 0)
|
|
RefineAbsorption(absorption_iter, n_groups);
|
|
if (!for_search && modulation_iter > 0)
|
|
RefineModulation(modulation_iter, n_groups);
|
|
#ifdef JFJOCH_USE_CUDA
|
|
// The corrections mutate the host fulls' corr; when the merge runs on the resident (GPU) fulls, push
|
|
// the corrected corr back to the device so the merge reads it.
|
|
if ((corrections || rejected_full_scales) && combined_on_gpu && scaled_fulls_on_gpu) {
|
|
std::vector<float> fcorr(fulls.size());
|
|
for (size_t i = 0; i < fulls.size(); ++i) fcorr[i] = fulls[i].corr;
|
|
gpu_->SetFullsCorr(fcorr.data());
|
|
}
|
|
#endif
|
|
|
|
// --- 5. Error model + merge + statistics. ---
|
|
auto r = MergeAndStats(n_groups, for_search, masked_ice_rings, combined_on_gpu && scaled_fulls_on_gpu);
|
|
return r;
|
|
}
|