leonarski_fandClaude Opus 4.8
87af0f419b
Merged reflections carried only intensities; a naive sqrt(max(I,0)) turns every weak or
negative measurement into a biased (or zero) amplitude. Add a French-Wilson step so the
merged reflections themselves carry proper Bayesian amplitudes.
New FrenchWilson.{h,cpp}: ApplyFrenchWilson() fills F and sigmaF on each MergedReflection
with the posterior mean |F| given I and its sigma under the Wilson prior:
- correct centric vs acentric prior (gemmi is_reflection_centric);
- epsilon (symmetry-enhancement) multiplicity: the shell mean is <I/eps> and the prior
mean for a reflection is eps * <I/eps>;
- numerically stable log-shift integration of the posterior;
- strong reflections (I > 4 sigma) short-circuit to sqrt(I) where the FW bias is negligible;
- unusable I/sigma falls back to sqrt(max(I,0)).
It runs as the last step of the merge routine (both MergeOnTheFly and RotationScaleMerge),
so F/sigmaF are part of the merged result and downstream consumers (file writing AND, later,
map calculation) all use the same amplitudes rather than each recomputing FW and risking
divergence. The writers just emit the fields; extended additively:
- MTZ: F (F) + SIGF (Q) columns;
- mmCIF: _refln.F_meas_au / _refln.F_meas_sigma_au;
- HKL text: F sigmaF appended (rfree flag kept in place for back-compat).
Existing intensity columns are unchanged, so current readers keep working while
phenix.refine (etc.) can now refine against the amplitudes.
Verified on lyso (1.3 A): all 173 negative-intensity reflections get F>0, strong reflections
have F/sqrt(I)=1.000, no NaN, and the mmCIF _refln loop is well-formed.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
1622 lines
82 KiB
C++
1622 lines
82 KiB
C++
// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
|
|
// SPDX-License-Identifier: GPL-3.0-only
|
|
|
|
#include "RotationScaleMerge.h"
|
|
|
|
#include <algorithm>
|
|
#include <atomic>
|
|
#include <cmath>
|
|
#include <cstdint>
|
|
#include <fstream>
|
|
#include <future>
|
|
#include <limits>
|
|
#include <random>
|
|
|
|
#include <gemmi/reciproc.hpp>
|
|
#include <gemmi/symmetry.hpp>
|
|
|
|
#include "HKLKey.h"
|
|
#include "RfreeFlags.h"
|
|
#include "FrenchWilson.h"
|
|
#include "ResolutionCutoff.h"
|
|
#include "../../common/CorrelationCoefficient.h"
|
|
#include "../../common/CrystalLattice.h"
|
|
#include "../../common/Definitions.h"
|
|
#include "../../common/JFJochException.h"
|
|
#include "../../common/ResolutionShells.h"
|
|
|
|
namespace {
|
|
// These mirror the per-image ScaleOnTheFly / Merge rocking-curve physics verbatim so this flat
|
|
// implementation is numerically identical - see the comments there for the details.
|
|
constexpr size_t MIN_REFLECTIONS = 20; // per-frame scale needs at least this many
|
|
constexpr double SCALE_ROBUST_K = 3.0; // Cauchy loss scale (sigma units) for the per-frame G fit
|
|
constexpr float MAX_FRAME_GAP = 2.0f; // a rocking event is a run of frames no more apart than this
|
|
constexpr double CHI2_1_MEDIAN = 0.454936;
|
|
// A post-scale-fulls correction surface (decay / absorption) is applied only if its held-out
|
|
// cross-validation improvement exceeds this fraction of the held-out scatter. A margin (not just >0)
|
|
// stops it acting on marginal, noise-level "improvements" that leave the merged quality unchanged or
|
|
// slightly worse - a correction should engage only when the generalizing signal is clear.
|
|
constexpr double CV_MIN_RELATIVE_GAIN = 0.02;
|
|
|
|
double SafeInv(double x, double fallback) {
|
|
if (!std::isfinite(x) || x == 0.0)
|
|
return fallback;
|
|
return 1.0 / x;
|
|
}
|
|
|
|
// Kabsch rotation partiality: the fraction of a reflection recorded in the sampled slice, from the
|
|
// erf of the rocking angle relative to the mosaic width. Identical to ScaleOnTheFly's RotationPartiality
|
|
// (and the predictor's), so recomputing here just swaps in the smoothed mosaicity.
|
|
float RotationPartiality(double delta_phi_deg, double zeta, double mosaicity_deg, double wedge_deg) {
|
|
const double half_wedge = wedge_deg / 2.0;
|
|
const double c1 = zeta / std::sqrt(2.0);
|
|
const double arg_plus = (delta_phi_deg + half_wedge) * c1 / mosaicity_deg;
|
|
const double arg_minus = (delta_phi_deg - half_wedge) * c1 / mosaicity_deg;
|
|
return static_cast<float>((std::erf(arg_plus) - std::erf(arg_minus)) / 2.0);
|
|
}
|
|
|
|
// Deterministic CC1/2 half from the frame's stable index (splitmix64), matching Merge.cpp.
|
|
int HalfForImage(int64_t image_id) {
|
|
uint64_t z = static_cast<uint64_t>(image_id) + 0x9e3779b97f4a7c15ULL;
|
|
z = (z ^ (z >> 30)) * 0xbf58476d1ce4e5b9ULL;
|
|
z = (z ^ (z >> 27)) * 0x94d049bb133111ebULL;
|
|
z = z ^ (z >> 31);
|
|
return static_cast<int>(z & 1ULL);
|
|
}
|
|
|
|
struct ScaleObs { double coeff, Iobs, weight; };
|
|
|
|
// Robust per-frame scale (linear in G): the exact objective ScaleOnTheFly::SolveScaleIRLS solves.
|
|
double SolveScaleIRLS(const std::vector<ScaleObs> &obs, double robust_k) {
|
|
auto weighted_scale = [&obs](auto robust_weight) {
|
|
double num = 0.0, den = 0.0;
|
|
for (const auto &o : obs) {
|
|
const double rw = robust_weight(o);
|
|
const double w2 = o.weight * o.weight;
|
|
num += rw * w2 * o.coeff * o.Iobs;
|
|
den += rw * w2 * o.coeff * o.coeff;
|
|
}
|
|
return den > 0.0 ? num / den : NAN;
|
|
};
|
|
|
|
double G = weighted_scale([](const ScaleObs &) { return 1.0; });
|
|
if (!std::isfinite(G))
|
|
return 1.0;
|
|
G = std::max(0.0, G);
|
|
|
|
const double k2 = robust_k * robust_k;
|
|
for (int iter = 0; iter < 30; ++iter) {
|
|
const double G_prev = G;
|
|
const double G_next = weighted_scale([&](const ScaleObs &o) {
|
|
const double res = o.weight * (G * o.coeff - o.Iobs);
|
|
return 1.0 / (1.0 + res * res / k2);
|
|
});
|
|
if (!std::isfinite(G_next))
|
|
break;
|
|
G = std::max(0.0, G_next);
|
|
if (std::abs(G - G_prev) <= 1e-7 * std::max(G, 1.0))
|
|
break;
|
|
}
|
|
return G;
|
|
}
|
|
|
|
// Run fn(i) for i in [0, n) over `nthreads` workers pulling from a shared atomic counter - the same
|
|
// self-load-balancing pattern the rest of the codebase uses (heavy frames don't stall light ones).
|
|
// Work-stealing per-item parallel: one atomic fetch per item. Use ONLY when the per-item work is
|
|
// heavy and uneven (e.g. per-frame fits) - the atomic amortises. For millions of tiny uniform items
|
|
// use ParallelChunks instead; a per-item atomic there is pure contention.
|
|
template <typename Fn>
|
|
void ParallelFor(int n, size_t nthreads, Fn fn) {
|
|
if (n <= 0) return;
|
|
if (nthreads <= 1 || n == 1) {
|
|
for (int i = 0; i < n; ++i) fn(i);
|
|
return;
|
|
}
|
|
const size_t local = std::min(nthreads, static_cast<size_t>(n));
|
|
std::atomic<int> next = 0;
|
|
std::vector<std::future<void>> futures;
|
|
futures.reserve(local);
|
|
for (size_t t = 0; t < local; ++t)
|
|
futures.emplace_back(std::async(std::launch::async, [&] {
|
|
for (int i = next.fetch_add(1); i < n; i = next.fetch_add(1))
|
|
fn(i);
|
|
}));
|
|
for (auto &f : futures) f.get();
|
|
}
|
|
|
|
// Chunked parallel: each worker gets one contiguous [lo, hi) range, no per-item synchronisation.
|
|
// Right for millions of cheap uniform items (the CPU stand-in for a flat CUDA grid-stride kernel).
|
|
template <typename Fn>
|
|
void ParallelChunks(int n, size_t nthreads, Fn fn) {
|
|
if (n <= 0) return;
|
|
const int nt = static_cast<int>(std::max<size_t>(1, std::min(nthreads, static_cast<size_t>(n))));
|
|
if (nt == 1) { fn(0, n); return; }
|
|
const int chunk = (n + nt - 1) / nt;
|
|
std::vector<std::future<void>> futures;
|
|
futures.reserve(nt);
|
|
for (int t = 0; t < nt; ++t) {
|
|
const int lo = t * chunk, hi = std::min(n, lo + chunk);
|
|
if (lo >= hi) break;
|
|
futures.emplace_back(std::async(std::launch::async, [&fn, lo, hi] { fn(lo, hi); }));
|
|
}
|
|
for (auto &f : futures) f.get();
|
|
}
|
|
|
|
double median_of(std::vector<double> &v) {
|
|
std::nth_element(v.begin(), v.begin() + v.size() / 2, v.end());
|
|
return v[v.size() / 2];
|
|
}
|
|
}
|
|
|
|
RotationScaleMerge::RotationScaleMerge(const DiffractionExperiment &experiment,
|
|
std::vector<IntegrationOutcome> &partial_outcomes,
|
|
std::optional<UnitCell> reference_cell,
|
|
int scaling_iterations, float ice_ring_half_width_q,
|
|
size_t nthreads, Logger &logger,
|
|
std::string observation_dump_path)
|
|
: x(experiment), partials_out(partial_outcomes), reference_cell(std::move(reference_cell)),
|
|
nthreads(nthreads == 0 ? std::thread::hardware_concurrency() : nthreads), logger(logger),
|
|
observation_dump_path(std::move(observation_dump_path)) {
|
|
const auto s = x.GetScalingSettings();
|
|
min_partiality = s.GetMinPartiality();
|
|
d_min_limit = s.GetHighResolutionLimit_A();
|
|
merge_friedel = s.GetMergeFriedel();
|
|
capture_uncertainty_coeff = s.GetCaptureUncertaintyCoeff();
|
|
min_captured_fraction = s.GetMinCapturedFraction();
|
|
reject_nsigma = s.GetOutlierRejectNsigma();
|
|
reject_outliers = reject_nsigma > 0.0;
|
|
rfree_fraction = s.GetRfreeFraction();
|
|
scale_fulls = s.GetScaleFulls();
|
|
// Decay + absorption correction surfaces are one master toggle (on by default; both cross-validated,
|
|
// so a no-op when their systematic is absent). Decoupled from the stills-only -B / RefineB flag.
|
|
refine_decay_b = s.GetCorrectionSurfaces();
|
|
absorption_iter = s.GetCorrectionSurfaces() ? s.GetAbsorptionIter() : 0;
|
|
scaling_iter = std::max(1, scaling_iterations);
|
|
resolution_cutoff_method = s.GetResolutionCutoff();
|
|
resolution_cc_target = s.GetResolutionCCTarget();
|
|
report_shell_count = s.GetReportShellCount();
|
|
if (const auto forced = s.GetForcedMosaicity(); forced.has_value())
|
|
mosaicity_deg = *forced;
|
|
else
|
|
mosaicity_deg = s.GetDefaultMosaicity();
|
|
ice_half_width_q = ice_ring_half_width_q;
|
|
}
|
|
|
|
void RotationScaleMerge::Ingest() {
|
|
n_frames = static_cast<int>(partials_out.size());
|
|
|
|
size_t total = 0;
|
|
for (const auto &o : partials_out) total += o.reflections.size();
|
|
|
|
partials.clear();
|
|
partials.reserve(total);
|
|
frame_start.assign(n_frames, 0);
|
|
frame_count.assign(n_frames, 0);
|
|
frame_cell_ok.assign(n_frames, 1);
|
|
g_partial.assign(n_frames, 1.0);
|
|
|
|
const float dist_tol = x.GetIndexingSettings().GetUnitCellDistTolerance();
|
|
const float ang_tol = x.GetIndexingSettings().GetUnitCellAngleTolerance_deg();
|
|
|
|
for (int o = 0; o < n_frames; ++o) {
|
|
frame_start[o] = static_cast<int32_t>(partials.size());
|
|
if (reference_cell) {
|
|
const auto cell = partials_out[o].latt.GetUnitCell();
|
|
frame_cell_ok[o] = cell.is_close(*reference_cell, dist_tol, ang_tol) ? 1 : 0;
|
|
}
|
|
for (const auto &r : partials_out[o].reflections) {
|
|
Obs obs{};
|
|
obs.h = r.h; obs.k = r.k; obs.l = r.l;
|
|
obs.I = r.I; obs.sigma = r.sigma; obs.d = r.d; obs.rlp = r.rlp;
|
|
obs.partiality = r.partiality; obs.zeta = r.zeta; obs.delta_phi = r.delta_phi_deg; obs.bkg = r.bkg;
|
|
obs.px = r.predicted_x; obs.py = r.predicted_y;
|
|
obs.image_number = r.image_number;
|
|
obs.frame = o;
|
|
obs.on_ice = r.on_ice_ring ? 1 : 0;
|
|
obs.corr = r.image_scale_corr;
|
|
obs.group = -1;
|
|
partials.push_back(obs);
|
|
}
|
|
frame_count[o] = static_cast<int32_t>(partials.size()) - frame_start[o];
|
|
}
|
|
|
|
// Per-obs AcceptReflection finiteness (immutable) - lets ComputeAsuGroups stamp the ASU-group id per
|
|
// obs from a flat 1-byte array instead of re-reading the fat Obs struct for every space group.
|
|
finite_ok.resize(partials.size());
|
|
for (size_t i = 0; i < partials.size(); ++i) {
|
|
const auto &o = partials[i];
|
|
finite_ok[i] = (std::isfinite(o.I) && std::isfinite(o.rlp) && o.rlp != 0.0f
|
|
&& std::isfinite(o.sigma) && o.sigma > 0.0f) ? 1 : 0;
|
|
}
|
|
|
|
// Sort ONCE by (raw h,k,l, image_number) and split into raw-hkl runs. This is the one expensive sort;
|
|
// both the 3D combine (event split) and the per-space-group ASU grouping reuse this order.
|
|
perm.resize(partials.size());
|
|
for (int i = 0; i < static_cast<int>(partials.size()); ++i) perm[i] = i;
|
|
std::sort(perm.begin(), perm.end(), [&](int32_t a, int32_t b) {
|
|
const auto &x1 = partials[a];
|
|
const auto &y1 = partials[b];
|
|
if (x1.h != y1.h) return x1.h < y1.h;
|
|
if (x1.k != y1.k) return x1.k < y1.k;
|
|
if (x1.l != y1.l) return x1.l < y1.l;
|
|
return x1.image_number < y1.image_number;
|
|
});
|
|
rawrun_start.clear(); rawrun_count.clear();
|
|
rawrun_h.clear(); rawrun_k.clear(); rawrun_l.clear(); rawrun_d.clear();
|
|
for (int i = 0; i < static_cast<int>(perm.size()); ) {
|
|
const auto &o0 = partials[perm[i]];
|
|
int j = i;
|
|
float d = NAN;
|
|
while (j < static_cast<int>(perm.size())) {
|
|
const auto &o = partials[perm[j]];
|
|
if (o.h != o0.h || o.k != o0.k || o.l != o0.l) break;
|
|
if (!std::isfinite(d) && std::isfinite(o.d) && o.d > 0.0f) d = o.d;
|
|
++j;
|
|
}
|
|
rawrun_start.push_back(i);
|
|
rawrun_count.push_back(j - i);
|
|
rawrun_h.push_back(o0.h); rawrun_k.push_back(o0.k); rawrun_l.push_back(o0.l);
|
|
rawrun_d.push_back(d);
|
|
i = j;
|
|
}
|
|
rawrun_group.assign(rawrun_start.size(), -1);
|
|
logger.Info("RotationScaleMerge: ingested {} partial observations from {} frames ({} distinct hkl)",
|
|
total, n_frames, rawrun_start.size());
|
|
|
|
SmoothMosaicityAndPartiality();
|
|
|
|
#ifdef JFJOCH_USE_CUDA
|
|
// Bring the partial-scaling loop onto the GPU when one is present. Upload the immutable per-obs
|
|
// fields once (corr lives on the device, refreshed each pass); the CPU keeps the sort/keying/combine.
|
|
gpu_ = std::make_unique<RotationScaleMergeGPU>();
|
|
gpu_active_ = gpu_->Available();
|
|
if (gpu_active_) {
|
|
const int n = static_cast<int>(partials.size());
|
|
std::vector<float> I(n), sigma(n), rlp(n), part(n), zeta(n), corr(n), bkg(n), img(n), dd(n),
|
|
px(n), py(n);
|
|
std::vector<uint8_t> onice(n);
|
|
std::vector<int32_t> frm(n);
|
|
for (int i = 0; i < n; ++i) {
|
|
const auto &o = partials[i];
|
|
I[i] = o.I; sigma[i] = o.sigma; rlp[i] = o.rlp; part[i] = o.partiality;
|
|
zeta[i] = o.zeta; onice[i] = o.on_ice; frm[i] = o.frame; corr[i] = o.corr;
|
|
bkg[i] = o.bkg; img[i] = o.image_number; dd[i] = o.d; px[i] = o.px; py[i] = o.py;
|
|
}
|
|
gpu_->SetPartials(n, n_frames, I.data(), sigma.data(), rlp.data(), part.data(), zeta.data(),
|
|
onice.data(), frm.data(), corr.data(), frame_start.data(), frame_count.data());
|
|
gpu_->SetCombineInputs(bkg.data(), img.data(), dd.data(), px.data(), py.data());
|
|
gpu_->SetRawRuns(static_cast<int>(rawrun_start.size()), static_cast<int>(perm.size()), perm.data(),
|
|
rawrun_start.data(), rawrun_count.data(),
|
|
rawrun_h.data(), rawrun_k.data(), rawrun_l.data());
|
|
gpu_->SetFrameCellOk(frame_cell_ok.data());
|
|
logger.Info("RotationScaleMerge: GPU scaling + combine + scale-fulls + merge active");
|
|
}
|
|
#endif
|
|
}
|
|
|
|
void RotationScaleMerge::SmoothMosaicityAndPartiality() {
|
|
// Per-frame mosaicity to recompute partiality from. A forced (fixed) mosaicity overrides every frame;
|
|
// otherwise use the per-frame value measured at integration (image-local, deterministic).
|
|
const auto forced_mosaicity = x.GetScalingSettings().GetForcedMosaicity();
|
|
std::vector<double> mos_raw(n_frames, NAN);
|
|
if (forced_mosaicity.has_value() && std::isfinite(*forced_mosaicity) && *forced_mosaicity > 0.0) {
|
|
for (int o = 0; o < n_frames; ++o) mos_raw[o] = *forced_mosaicity;
|
|
} else {
|
|
for (int o = 0; o < n_frames; ++o) {
|
|
const auto &m = partials_out[o].mosaicity_deg;
|
|
if (m && std::isfinite(*m) && *m > 0.0f) mos_raw[o] = *m;
|
|
}
|
|
}
|
|
|
|
// Frame-order moving average with the same window as smooth-G (a rotation range -> frame count).
|
|
// With smoothing off, fall back to the per-frame value (still deterministic, just unsmoothed).
|
|
const auto ss = x.GetScalingSettings();
|
|
const double smooth_deg = ss.GetSmoothGDegrees();
|
|
const auto gon = x.GetGoniometer();
|
|
const double osc = gon ? std::fabs(gon->GetIncrement_deg()) : 0.0;
|
|
mos_smooth.assign(n_frames, NAN);
|
|
if (smooth_deg > 0.0 && osc > 1e-6) {
|
|
int window = std::max(1, static_cast<int>(std::lround(smooth_deg / osc)));
|
|
if (window % 2 == 0) ++window;
|
|
const int half = window / 2;
|
|
for (int o = 0; o < n_frames; ++o) {
|
|
double sum = 0.0;
|
|
int cnt = 0;
|
|
for (int j = std::max(0, o - half); j <= std::min(n_frames - 1, o + half); ++j)
|
|
if (std::isfinite(mos_raw[j])) { sum += mos_raw[j]; ++cnt; }
|
|
if (cnt > 0) mos_smooth[o] = static_cast<float>(sum / cnt);
|
|
}
|
|
} else {
|
|
for (int o = 0; o < n_frames; ++o) mos_smooth[o] = static_cast<float>(mos_raw[o]);
|
|
}
|
|
|
|
// Recompute each partial's partiality from the smoothed mosaicity (same wedge the predictor used).
|
|
// Frames without a mosaicity keep the stored partiality.
|
|
const double wedge = gon ? std::fabs(gon->GetWedge_deg()) : 0.0;
|
|
ParallelChunks(static_cast<int>(partials.size()), nthreads, [&](int lo, int hi) {
|
|
for (int i = lo; i < hi; ++i) {
|
|
auto &o = partials[i];
|
|
const float mos = mos_smooth[o.frame];
|
|
if (std::isfinite(mos) && mos > 1e-6f && std::isfinite(o.zeta) && o.zeta > 0.0f
|
|
&& std::isfinite(o.delta_phi))
|
|
o.partiality = RotationPartiality(o.delta_phi, o.zeta, mos, wedge);
|
|
}
|
|
});
|
|
logger.Info("Recomputed partiality from frame-order-smoothed mosaicity");
|
|
}
|
|
|
|
int RotationScaleMerge::ComputeAsuGroups(const HKLKeyGenerator &keygen) {
|
|
// One ASU reduction per distinct raw hkl (not per observation): a raw hkl is eligible if it is not
|
|
// systematically absent and its resolution is in range. Group the eligible raw hkls by ASU key
|
|
// (sort the ~#distinct-hkl keys, not the millions of observations), then hand out dense ids.
|
|
const int n_run = static_cast<int>(rawrun_start.size());
|
|
std::vector<uint64_t> key(n_run);
|
|
std::vector<uint8_t> eligible(n_run, 0);
|
|
// The gemmi ASU reduction / absence test per raw hkl is the cost here and is independent per run
|
|
// (HKLKeyGenerator is const, so concurrent reads are safe) - compute keys in parallel chunks.
|
|
ParallelChunks(n_run, nthreads, [&](int lo, int hi) {
|
|
for (int r = lo; r < hi; ++r) {
|
|
rawrun_group[r] = -1;
|
|
if (keygen.IsSystematicallyAbsent(rawrun_h[r], rawrun_k[r], rawrun_l[r]))
|
|
continue;
|
|
const float d = rawrun_d[r]; // resolution is a per-raw-hkl property (all its partials share d)
|
|
if (!std::isfinite(d) || d <= 0.0f) continue;
|
|
if (d_min_limit && d < *d_min_limit) continue;
|
|
key[r] = keygen(rawrun_h[r], rawrun_k[r], rawrun_l[r]).pack();
|
|
eligible[r] = 1;
|
|
}
|
|
});
|
|
std::vector<int32_t> idx;
|
|
idx.reserve(n_run);
|
|
for (int r = 0; r < n_run; ++r)
|
|
if (eligible[r]) idx.push_back(r);
|
|
std::sort(idx.begin(), idx.end(), [&](int32_t a, int32_t b) { return key[a] < key[b]; });
|
|
|
|
group_h.clear(); group_k.clear(); group_l.clear();
|
|
int n_groups = 0;
|
|
for (size_t j = 0; j < idx.size(); ++j) {
|
|
if (j == 0 || key[idx[j]] != key[idx[j - 1]]) {
|
|
const int r = idx[j];
|
|
const auto hkl = keygen(rawrun_h[r], rawrun_k[r], rawrun_l[r]);
|
|
group_h.push_back(hkl.plus ? hkl.h : -hkl.h);
|
|
group_k.push_back(hkl.plus ? hkl.k : -hkl.k);
|
|
group_l.push_back(hkl.plus ? hkl.l : -hkl.l);
|
|
++n_groups;
|
|
}
|
|
rawrun_group[idx[j]] = n_groups - 1;
|
|
}
|
|
|
|
// Stamp the ASU-group id per obs from its raw hkl + the precomputed finiteness. For the GPU we build a
|
|
// flat group_ids array (fed to the reduction); the Obs.group field is written only when a CPU stage
|
|
// will read it - i.e. no GPU (with a GPU, scaling/CC/combine all read group_ids / rawrun_group, never
|
|
// partials.group) - so the default path skips a strided 6.3M pass over the fat Obs struct.
|
|
const bool need_obs_group =
|
|
#ifdef JFJOCH_USE_CUDA
|
|
!gpu_active_;
|
|
#else
|
|
true;
|
|
#endif
|
|
#ifdef JFJOCH_USE_CUDA
|
|
std::vector<int32_t> group_ids;
|
|
if (gpu_active_) group_ids.resize(partials.size());
|
|
#endif
|
|
// Parallel over raw-hkl runs: distinct runs own disjoint perm ranges, hence disjoint observations.
|
|
ParallelChunks(n_run, nthreads, [&](int rlo, int rhi) {
|
|
for (int r = rlo; r < rhi; ++r) {
|
|
const int g = rawrun_group[r];
|
|
const int lo = rawrun_start[r], hi = rawrun_start[r] + rawrun_count[r];
|
|
for (int p = lo; p < hi; ++p) {
|
|
const int i = perm[p];
|
|
const int gi = (g >= 0 && finite_ok[i]) ? g : -1;
|
|
#ifdef JFJOCH_USE_CUDA
|
|
if (gpu_active_) group_ids[i] = gi;
|
|
#endif
|
|
if (need_obs_group) partials[i].group = gi;
|
|
}
|
|
}
|
|
});
|
|
|
|
#ifdef JFJOCH_USE_CUDA
|
|
// Group-ordered permutation (obs bucketed by ASU group, obs-index order) + its CSR, so the GPU
|
|
// reduction is a deterministic segmented reduction (fixed order, no atomics). A stable counting sort
|
|
// by group, parallel via per-chunk histograms over the flat group_ids; bit-identical to a serial fill
|
|
// (chunk order == obs-index order, each chunk fills its groups sequentially).
|
|
if (gpu_active_) {
|
|
const int n = static_cast<int>(partials.size());
|
|
const int nt = static_cast<int>(std::max<size_t>(1, std::min(nthreads, static_cast<size_t>(std::max(1, n)))));
|
|
const int chunk = (n + nt - 1) / nt;
|
|
std::vector<std::vector<int32_t>> hist(nt, std::vector<int32_t>(n_groups, 0));
|
|
|
|
// Pass 1 (parallel): per-chunk group histogram over the flat group_ids.
|
|
std::vector<std::future<void>> f1;
|
|
for (int t = 0; t < nt; ++t) {
|
|
const int lo = t * chunk, hi = std::min(n, lo + chunk);
|
|
if (lo >= hi) break;
|
|
f1.emplace_back(std::async(std::launch::async, [&, t, lo, hi] {
|
|
auto &h = hist[t];
|
|
for (int i = lo; i < hi; ++i) { const int g = group_ids[i]; if (g >= 0) ++h[g]; }
|
|
}));
|
|
}
|
|
for (auto &f : f1) f.get();
|
|
|
|
// CSR starts + convert hist[t][g] into chunk t's write base for group g (exclusive prefix over t).
|
|
std::vector<int32_t> gstart(n_groups), gcount(n_groups);
|
|
int acc = 0;
|
|
for (int g = 0; g < n_groups; ++g) {
|
|
int base = acc;
|
|
gstart[g] = acc;
|
|
for (int t = 0; t < nt; ++t) { const int c = hist[t][g]; hist[t][g] = base; base += c; }
|
|
gcount[g] = base - acc;
|
|
acc = base;
|
|
}
|
|
|
|
// Pass 2 (parallel): each chunk fills its obs into gperm at its per-group base (stable).
|
|
std::vector<int32_t> gperm(acc);
|
|
std::vector<std::future<void>> f2;
|
|
for (int t = 0; t < nt; ++t) {
|
|
const int lo = t * chunk, hi = std::min(n, lo + chunk);
|
|
if (lo >= hi) break;
|
|
f2.emplace_back(std::async(std::launch::async, [&, t, lo, hi] {
|
|
std::vector<int32_t> fill = hist[t];
|
|
for (int i = lo; i < hi; ++i) { const int g = group_ids[i]; if (g >= 0) gperm[fill[g]++] = i; }
|
|
}));
|
|
}
|
|
for (auto &f : f2) f.get();
|
|
gpu_->SetGroups(n_groups, group_ids.data(), gperm.data(), acc, gstart.data(), gcount.data());
|
|
}
|
|
#endif
|
|
return n_groups;
|
|
}
|
|
|
|
void RotationScaleMerge::ReduceGroupMeans(const std::vector<Obs> &obs, int n_groups,
|
|
bool exclude_ice, const std::vector<char> &masked,
|
|
std::vector<double> &out_mean) const {
|
|
// Inverse-variance per-group mean of I*corr = the merge reference (a segmented reduction over the
|
|
// groups; the CPU stand-in for a CUDA reduce_by_key). No cell mask here: the scaling reference
|
|
// (MergeAll) is built without a reference cell - only the final merge applies it.
|
|
std::vector<double> sw(n_groups, 0.0), swI(n_groups, 0.0);
|
|
for (const auto &o : obs) {
|
|
if (o.group < 0) continue;
|
|
if (!(o.corr > 0.0f) || !std::isfinite(o.corr)) continue;
|
|
if (exclude_ice && o.on_ice) continue;
|
|
if (!masked.empty()) {
|
|
const int ring = IceRingIndex(o.d, ice_half_width_q);
|
|
if (ring >= 0 && ring < static_cast<int>(masked.size()) && masked[ring]) continue;
|
|
}
|
|
if (o.partiality < min_partiality) continue;
|
|
const float I_corr = o.I * o.corr;
|
|
const float sigma_corr = o.sigma * o.corr;
|
|
if (!std::isfinite(I_corr) || !std::isfinite(sigma_corr) || sigma_corr <= 0.0f) continue;
|
|
const double w = 1.0 / (static_cast<double>(sigma_corr) * sigma_corr);
|
|
sw[o.group] += w;
|
|
swI[o.group] += w * I_corr;
|
|
}
|
|
out_mean.assign(n_groups, NAN);
|
|
for (int g = 0; g < n_groups; ++g)
|
|
if (sw[g] > 0.0) out_mean[g] = swI[g] / sw[g];
|
|
}
|
|
|
|
void RotationScaleMerge::FitPerFrameG(std::vector<Obs> &obs, const std::vector<int32_t> &fstart,
|
|
const std::vector<int32_t> &fcount,
|
|
const std::vector<double> &group_mean_in,
|
|
bool unity, std::vector<double> &g) {
|
|
std::vector<uint8_t> scaled(fstart.size(), 0);
|
|
ParallelFor(static_cast<int>(fstart.size()), nthreads, [&](int f) {
|
|
std::vector<ScaleObs> so;
|
|
so.reserve(fcount[f]);
|
|
const int lo = fstart[f], hi = fstart[f] + fcount[f];
|
|
for (int i = lo; i < hi; ++i) {
|
|
const auto &o = obs[i];
|
|
if (o.group < 0) continue;
|
|
if (o.on_ice) continue;
|
|
const double mean = group_mean_in[o.group];
|
|
if (!std::isfinite(mean)) continue;
|
|
double coeff;
|
|
if (unity) {
|
|
coeff = mean; // partiality already folded into the full, rlp = 1
|
|
} else {
|
|
if (!(std::isfinite(o.zeta) && o.zeta > 0.0f)) continue; // Rotation model needs zeta > 0
|
|
coeff = o.partiality * SafeInv(o.rlp, 1.0) * mean;
|
|
}
|
|
so.push_back({coeff, static_cast<double>(o.I), SafeInv(o.sigma, 1.0)});
|
|
}
|
|
if (so.size() < MIN_REFLECTIONS) return; // leave g[f]/corr untouched (as ScaleOnTheFly does)
|
|
g[f] = SolveScaleIRLS(so, SCALE_ROBUST_K);
|
|
scaled[f] = 1;
|
|
});
|
|
// Remember which frames were fitted this call (so the caller updates corr only there).
|
|
frame_scaled_scratch = std::move(scaled);
|
|
}
|
|
|
|
void RotationScaleMerge::UpdateCorr(std::vector<Obs> &obs, const std::vector<double> &g,
|
|
const std::vector<uint8_t> &frame_scaled) const {
|
|
ParallelChunks(static_cast<int>(obs.size()), nthreads, [&](int lo, int hi) {
|
|
for (int i = lo; i < hi; ++i) {
|
|
auto &o = obs[i];
|
|
if (!frame_scaled[o.frame]) continue;
|
|
const double denom = static_cast<double>(o.partiality) * g[o.frame]; // B_term = 1 (no B refine)
|
|
if (std::isfinite(o.rlp) && std::isfinite(denom) && denom > 0.0)
|
|
o.corr = static_cast<float>(o.rlp / denom);
|
|
else
|
|
o.corr = NAN;
|
|
}
|
|
});
|
|
}
|
|
|
|
void RotationScaleMerge::RefineDecay(int n_groups) {
|
|
// Radiation damage weakens later frames more at higher resolution - a resolution x time (Debye-Waller)
|
|
// systematic the resolution-flat per-frame G leaves in. Model it as a single global relative-B rate:
|
|
// ln(I_ref / I_obs) = slope * (frame - frame_center) * s^2, s^2 = 1/(4 d^2), slope = 2 dB/dframe,
|
|
// fitted (inverse-variance-of-log weighted) against the merged reference.
|
|
//
|
|
// The fit is dominated by strong LOW-res reflections and EXTRAPOLATED to high s^2, so a slope fitted to
|
|
// noise / aggregate structure (not real per-reflection decay) would amplify into a high-res
|
|
// mis-correction and worsen the symmetry-equivalent scatter that sets ISa. An opt-in correction must
|
|
// never add noise where its systematic is absent (a low-dose reference set), so the slope is
|
|
// CROSS-VALIDATED: fit on even-frame fulls and score the held-out odd-frame equivalent chi^2, 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 chi^2 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;
|
|
};
|
|
// Mean studentized squared deviation over the subset {frame&1 == parity} when `slope` is applied,
|
|
// scored against that subset's own reference (no leakage). Lower = tighter equivalents.
|
|
auto subset_chi2 = [&](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 sum = 0.0; size_t n = 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);
|
|
const double sc = static_cast<double>(o.sigma) * o.corr * decay_factor(o, slope);
|
|
if (!std::isfinite(Iref) || !(sc > 0.0)) continue;
|
|
const double dd = (Is - Iref) / sc;
|
|
sum += dd * dd; ++n;
|
|
}
|
|
return n > 0 ? sum / static_cast<double>(n) : 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_chi2(1, 0.0) + subset_chi2(0, 0.0);
|
|
const double gain = base - (subset_chi2(1, fit_slope(0)) + subset_chi2(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);
|
|
}
|
|
|
|
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;
|
|
const int ncell = NB * NB;
|
|
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;
|
|
}
|
|
|
|
auto usable = [&](const Obs &o) {
|
|
return o.group >= 0 && o.corr > 0.0f && std::isfinite(o.corr) && o.partiality >= min_partiality;
|
|
};
|
|
// Fit the per-cell absorption 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 accumulated
|
|
// per-cell factor A[cell].
|
|
auto fit_surface = [&](int parity) -> std::vector<double> {
|
|
std::vector<double> A(ncell, 1.0);
|
|
for (int it = 0; it < n_iter; ++it) {
|
|
std::vector<double> sw(n_groups, 0.0), swI(n_groups, 0.0);
|
|
for (size_t i = 0; i < fulls.size(); ++i) {
|
|
const Obs &o = fulls[i];
|
|
if (!usable(o) || cell[i] < 0 || (parity >= 0 && (o.frame & 1) != parity)) continue;
|
|
const double a = A[cell[i]], sc = static_cast<double>(o.sigma) * o.corr * a, w = 1.0 / (sc * sc);
|
|
sw[o.group] += w; swI[o.group] += w * static_cast<double>(o.I) * o.corr * a;
|
|
}
|
|
std::vector<double> num(ncell, 0.0), den(ncell, 0.0);
|
|
for (size_t i = 0; i < fulls.size(); ++i) {
|
|
const Obs &o = fulls[i];
|
|
if (!usable(o) || cell[i] < 0 || (parity >= 0 && (o.frame & 1) != parity) || sw[o.group] <= 0.0)
|
|
continue;
|
|
const double Iref = swI[o.group] / sw[o.group], a = A[cell[i]];
|
|
const double Is = static_cast<double>(o.I) * o.corr * a, sc = static_cast<double>(o.sigma) * o.corr * a;
|
|
if (!std::isfinite(Iref) || Iref <= 0.0 || !(Is > 0.0) || !(sc > 0.0)) continue;
|
|
const double w = 1.0 / (sc * sc);
|
|
num[cell[i]] += w * Is * Iref; den[cell[i]] += w * Is * Is;
|
|
}
|
|
std::vector<double> dsorted = den;
|
|
std::nth_element(dsorted.begin(), dsorted.begin() + dsorted.size() / 2, dsorted.end());
|
|
const double lambda = 0.1 * std::max(1e-30, dsorted[dsorted.size() / 2]);
|
|
std::vector<double> upd(ncell, 1.0);
|
|
for (int c = 0; c < ncell; ++c) upd[c] = (num[c] + lambda) / (den[c] + lambda);
|
|
double logsum = 0.0, wsum = 0.0;
|
|
for (int c = 0; c < ncell; ++c) if (den[c] > 0.0) { logsum += den[c] * std::log(upd[c]); wsum += den[c]; }
|
|
const double gm = wsum > 0.0 ? std::exp(logsum / wsum) : 1.0;
|
|
for (int c = 0; c < ncell; ++c) A[c] = std::clamp(A[c] * upd[c] / gm, 0.25, 4.0);
|
|
}
|
|
return A;
|
|
};
|
|
// Mean studentized squared deviation over the subset {frame&1 == parity} with surface A applied,
|
|
// scored against that subset's own reference (no leakage). Lower = tighter equivalents.
|
|
auto score = [&](int parity, const std::vector<double> &A) -> double {
|
|
std::vector<double> sw(n_groups, 0.0), swI(n_groups, 0.0);
|
|
for (size_t i = 0; i < fulls.size(); ++i) {
|
|
const Obs &o = fulls[i];
|
|
if (!usable(o) || cell[i] < 0 || (o.frame & 1) != parity) continue;
|
|
const double a = A[cell[i]], Is = static_cast<double>(o.I) * o.corr * a;
|
|
const double sc = static_cast<double>(o.sigma) * o.corr * a, w = 1.0 / (sc * sc);
|
|
sw[o.group] += w; swI[o.group] += w * Is;
|
|
}
|
|
double sum = 0.0; size_t n = 0;
|
|
for (size_t i = 0; i < fulls.size(); ++i) {
|
|
const Obs &o = fulls[i];
|
|
if (!usable(o) || cell[i] < 0 || (o.frame & 1) != parity || sw[o.group] <= 0.0) continue;
|
|
const double a = A[cell[i]], Is = static_cast<double>(o.I) * o.corr * a;
|
|
const double sc = static_cast<double>(o.sigma) * o.corr * a;
|
|
const double Iref = swI[o.group] / sw[o.group];
|
|
if (!std::isfinite(Iref) || !(sc > 0.0)) continue;
|
|
const double dd = (Is - Iref) / sc;
|
|
sum += dd * dd; ++n;
|
|
}
|
|
return n > 0 ? sum / static_cast<double>(n) : 0.0;
|
|
};
|
|
|
|
// Cross-validate: fit the surface on even frames and score the held-out odd equivalents (and vice
|
|
// versa). A real absorption 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 or act on marginal noise.
|
|
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("Absorption correction: not cross-validated (held-out gain {:.1f}%, skipped)",
|
|
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("Absorption correction: goniometer-frame {}x{} surface (cross-validated, held-out gain {:.1f}%)",
|
|
NB, NB, 100.0 * gain / std::max(base, 1e-30));
|
|
}
|
|
|
|
void RotationScaleMerge::ComputeSmoothGWindow(const std::vector<double> &g, int window,
|
|
std::vector<double> &g_smooth) const {
|
|
const int n = static_cast<int>(g.size());
|
|
const int half = window / 2;
|
|
g_smooth.assign(n, NAN);
|
|
for (int o = 0; o < n; ++o) {
|
|
double sum_log = 0.0;
|
|
int count = 0;
|
|
for (int j = std::max(0, o - half); j <= std::min(n - 1, o + half); ++j) {
|
|
if (frame_scaled_scratch[j] && std::isfinite(g[j]) && g[j] > 0.0) {
|
|
sum_log += std::log(g[j]);
|
|
++count;
|
|
}
|
|
}
|
|
if (count > 0) g_smooth[o] = std::exp(sum_log / count);
|
|
}
|
|
}
|
|
|
|
void RotationScaleMerge::SmoothG(std::vector<Obs> &obs, std::vector<double> &g, int window) const {
|
|
const int n = static_cast<int>(g.size());
|
|
std::vector<double> g_smooth;
|
|
ComputeSmoothGWindow(g, window, g_smooth);
|
|
for (auto &o : obs) {
|
|
const int f = o.frame;
|
|
if (!frame_scaled_scratch[f] || !std::isfinite(g[f]) || g[f] <= 0.0 || !std::isfinite(g_smooth[f]))
|
|
continue;
|
|
if (std::isfinite(o.corr))
|
|
o.corr = static_cast<float>(o.corr * (g[f] / g_smooth[f]));
|
|
}
|
|
for (int f = 0; f < n; ++f)
|
|
if (frame_scaled_scratch[f] && std::isfinite(g[f]) && g[f] > 0.0 && std::isfinite(g_smooth[f]))
|
|
g[f] = g_smooth[f];
|
|
}
|
|
|
|
void RotationScaleMerge::Combine() {
|
|
fulls.clear();
|
|
g_full.assign(n_frames, 1.0);
|
|
|
|
// Combine one raw-hkl run into fulls, appended to `out`. Independent per run (the events of one hkl
|
|
// touch no shared state), so runs parallelise cleanly. Returns the number of usable partials seen.
|
|
// `dump` (serial path only) writes each emitted full for the diagnostic observation dump.
|
|
auto process_rawrun = [&](int r, std::vector<Obs> &out, std::ofstream *dump) -> size_t {
|
|
const int lo = rawrun_start[r], hi = rawrun_start[r] + rawrun_count[r];
|
|
std::vector<int32_t> ev; // usable perm-indices of this raw hkl, in image-number order
|
|
ev.reserve(hi - lo);
|
|
for (int p = lo; p < hi; ++p) {
|
|
const auto &o = partials[perm[p]];
|
|
if (!std::isfinite(o.corr) || o.corr <= 0.0f) continue;
|
|
if (!std::isfinite(o.I) || !std::isfinite(o.sigma) || o.sigma <= 0.0f) continue;
|
|
ev.push_back(perm[p]);
|
|
}
|
|
const int group = rawrun_group[r];
|
|
|
|
size_t i = 0;
|
|
while (i < ev.size()) {
|
|
size_t kk = i + 1;
|
|
float last_frame = partials[ev[i]].image_number;
|
|
while (kk < ev.size()) {
|
|
const float frame = partials[ev[kk]].image_number;
|
|
if (frame - last_frame > MAX_FRAME_GAP) break;
|
|
last_frame = frame;
|
|
++kk;
|
|
}
|
|
|
|
double pooled_bkg = 0.0;
|
|
int n_pool = 0;
|
|
for (size_t m = i; m < kk; ++m) {
|
|
const float b = partials[ev[m]].bkg;
|
|
if (std::isfinite(b)) { pooled_bkg += b; ++n_pool; }
|
|
}
|
|
pooled_bkg = n_pool > 0 ? pooled_bkg / n_pool : 0.0;
|
|
auto pooled_I = [&](const Obs &r2) {
|
|
const double n_bkg = std::max(0.0, static_cast<double>(r2.sigma) * r2.sigma - r2.I)
|
|
/ std::max(r2.bkg, 1.0f);
|
|
return static_cast<double>(r2.I) + n_bkg * (static_cast<double>(r2.bkg) - pooled_bkg);
|
|
};
|
|
|
|
double sum_w = 0.0, sum_wI = 0.0, sum_partiality = 0.0;
|
|
float d = NAN;
|
|
int peak_outcome = partials[ev[i]].frame;
|
|
float peak_frame = partials[ev[i]].image_number;
|
|
float peak_px = partials[ev[i]].px, peak_py = partials[ev[i]].py;
|
|
float peak_partiality = -1.0f;
|
|
const bool on_ice = partials[ev[i]].on_ice;
|
|
for (size_t m = i; m < kk; ++m) {
|
|
const auto &r2 = partials[ev[m]];
|
|
const double sigma_corr = static_cast<double>(r2.sigma) * r2.corr;
|
|
const double w = 1.0 / (sigma_corr * sigma_corr);
|
|
sum_w += w;
|
|
sum_wI += w * pooled_I(r2) * r2.corr;
|
|
sum_partiality += r2.partiality;
|
|
if (r2.partiality > peak_partiality) {
|
|
peak_partiality = r2.partiality;
|
|
peak_outcome = r2.frame;
|
|
peak_frame = r2.image_number;
|
|
peak_px = r2.px; peak_py = r2.py;
|
|
}
|
|
if (!std::isfinite(d) && std::isfinite(r2.d) && r2.d > 0.0f) d = r2.d;
|
|
}
|
|
double F = sum_wI / sum_w;
|
|
|
|
for (int iter = 0; iter < 3; ++iter) {
|
|
sum_w = 0.0; sum_wI = 0.0;
|
|
for (size_t m = i; m < kk; ++m) {
|
|
const auto &r2 = partials[ev[m]];
|
|
const double corr = r2.corr;
|
|
const double I_corr = pooled_I(r2) * corr;
|
|
const double sigma_corr = static_cast<double>(r2.sigma) * corr;
|
|
const double bkg_var = sigma_corr * sigma_corr - corr * I_corr;
|
|
double var = std::max(0.0, bkg_var) + corr * std::max(0.0, F);
|
|
if (!(var > 0.0)) var = sigma_corr * sigma_corr;
|
|
const double w = 1.0 / var;
|
|
sum_w += w;
|
|
sum_wI += w * I_corr;
|
|
}
|
|
F = sum_wI / sum_w;
|
|
}
|
|
|
|
const int n_frames_event = static_cast<int>(kk - i);
|
|
i = kk;
|
|
|
|
if (sum_w <= 0.0 || sum_partiality < min_partiality || sum_partiality < min_captured_fraction)
|
|
continue;
|
|
|
|
double sigma_full = 1.0 / std::sqrt(sum_w);
|
|
if (capture_uncertainty_coeff > 0.0) {
|
|
const double frac = std::min(1.0, sum_partiality);
|
|
const double extra = capture_uncertainty_coeff * (1.0 - frac) * std::max(0.0, F);
|
|
sigma_full = std::sqrt(sigma_full * sigma_full + extra * extra);
|
|
}
|
|
|
|
Obs full{};
|
|
full.h = rawrun_h[r]; full.k = rawrun_k[r]; full.l = rawrun_l[r];
|
|
full.I = static_cast<float>(F);
|
|
full.sigma = static_cast<float>(sigma_full);
|
|
full.d = d;
|
|
full.rlp = 1.0f;
|
|
full.partiality = 1.0f;
|
|
full.corr = 1.0f;
|
|
full.image_number = peak_frame;
|
|
full.px = peak_px; full.py = peak_py;
|
|
full.frame = peak_outcome;
|
|
full.on_ice = on_ice ? 1 : 0;
|
|
full.group = group; // the raw hkl's ASU group for the current space group (<0 = absent)
|
|
out.push_back(full);
|
|
|
|
if (dump != nullptr)
|
|
*dump << full.h << ' ' << full.k << ' ' << full.l << ' ' << full.I << ' '
|
|
<< full.sigma << ' ' << d << ' ' << n_frames_event << ' ' << sum_partiality << ' '
|
|
<< static_cast<int>(peak_frame) << '\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);
|
|
dump << "# h k l I sigma d n_frames captured_fraction\n";
|
|
}
|
|
fulls.reserve(n_run);
|
|
for (int r = 0; r < n_run; ++r)
|
|
n_used += process_rawrun(r, fulls, dump.is_open() ? &dump : nullptr);
|
|
} else {
|
|
// Parallel over contiguous rawrun chunks; concatenate the per-thread fulls in run order so the
|
|
// result is deterministic.
|
|
const int nt = static_cast<int>(std::min(nthreads, static_cast<size_t>(n_run)));
|
|
const int chunk = (n_run + nt - 1) / nt;
|
|
std::vector<std::vector<Obs>> part(nt);
|
|
std::vector<size_t> used(nt, 0);
|
|
std::vector<std::future<void>> futures;
|
|
futures.reserve(nt);
|
|
for (int t = 0; t < nt; ++t) {
|
|
const int r0 = t * chunk, r1 = std::min(n_run, r0 + chunk);
|
|
if (r0 >= r1) break;
|
|
futures.emplace_back(std::async(std::launch::async, [&, t, r0, r1] {
|
|
part[t].reserve(r1 - r0);
|
|
for (int r = r0; r < r1; ++r) used[t] += process_rawrun(r, part[t], nullptr);
|
|
}));
|
|
}
|
|
for (auto &f : futures) f.get();
|
|
size_t total = 0;
|
|
for (int t = 0; t < nt; ++t) { total += part[t].size(); n_used += used[t]; }
|
|
fulls.reserve(total);
|
|
for (int t = 0; t < nt; ++t)
|
|
fulls.insert(fulls.end(), part[t].begin(), part[t].end());
|
|
}
|
|
|
|
SortFullsByFrame();
|
|
logger.Info("3D combine: {} fulls from {} partials", fulls.size(), n_used);
|
|
}
|
|
|
|
void RotationScaleMerge::SortFullsByFrame() {
|
|
// Sort the fulls by their (peak) frame and build per-frame CSR ranges, so the scale-fulls step can
|
|
// fit a per-frame G by slicing contiguous ranges (the same layout the partials use).
|
|
std::sort(fulls.begin(), fulls.end(),
|
|
[](const Obs &a, const Obs &b) { return a.frame < b.frame; });
|
|
fulls_frame_start.assign(n_frames, 0);
|
|
fulls_frame_count.assign(n_frames, 0);
|
|
for (int i = 0; i < static_cast<int>(fulls.size()); ) {
|
|
const int f = fulls[i].frame;
|
|
fulls_frame_start[f] = i;
|
|
int j = i;
|
|
while (j < static_cast<int>(fulls.size()) && fulls[j].frame == f) ++j;
|
|
fulls_frame_count[f] = j - i;
|
|
i = j;
|
|
}
|
|
}
|
|
|
|
void RotationScaleMerge::ComputePerFrameCC(const std::vector<double> &partial_group_mean,
|
|
std::vector<double> &cc, std::vector<int64_t> &cc_n) const {
|
|
// Per-frame CC vs the merged reference (CalculateGlobalCC), computed once now (not every iteration).
|
|
cc.assign(n_frames, NAN);
|
|
cc_n.assign(n_frames, 0);
|
|
ParallelFor(n_frames, nthreads, [&](int f) {
|
|
double sx = 0, sy = 0, sx2 = 0, sy2 = 0, sxy = 0;
|
|
size_t n = 0;
|
|
const int lo = frame_start[f], hi = frame_start[f] + frame_count[f];
|
|
for (int i = lo; i < hi; ++i) {
|
|
const auto &o = partials[i];
|
|
if (o.on_ice) continue;
|
|
if (o.group < 0) continue;
|
|
if (o.partiality < min_partiality) continue;
|
|
if (!std::isfinite(o.I) || !std::isfinite(o.corr) || o.corr <= 0.0f) continue;
|
|
if (!std::isfinite(o.sigma) || o.sigma <= 0.0f) continue;
|
|
const double mean = partial_group_mean[o.group];
|
|
if (!std::isfinite(mean)) continue;
|
|
const double img = static_cast<double>(o.I) * o.corr;
|
|
sx += img; sy += mean; sx2 += img * img; sy2 += mean * mean; sxy += img * mean;
|
|
++n;
|
|
}
|
|
if (n < MIN_REFLECTIONS) return;
|
|
const double nd = static_cast<double>(n);
|
|
const double cov = sxy - sx * sy / nd;
|
|
const double vx = sx2 - sx * sx / nd;
|
|
const double vy = sy2 - sy * sy / nd;
|
|
if (vx > 0.0 && vy > 0.0) { cc[f] = cov / std::sqrt(vx * vy); cc_n[f] = static_cast<int64_t>(n); }
|
|
});
|
|
}
|
|
|
|
// Write the per-frame G / CC / mosaicity (from the given cc/cc_n) back onto the partials for the offline
|
|
// per-image scaling table. cc/cc_n are computed on the host (ComputePerFrameCC) or GPU (ComputePartialCC).
|
|
void RotationScaleMerge::FinalizePerFrameScale(const std::vector<double> &cc, const std::vector<int64_t> &cc_n,
|
|
const std::vector<uint8_t> &frame_scaled) {
|
|
for (int f = 0; f < n_frames; ++f) {
|
|
auto &o = partials_out[f];
|
|
if (frame_scaled[f]) {
|
|
o.image_scale_g = static_cast<float>(g_partial[f]);
|
|
o.mosaicity_deg = (f < static_cast<int>(mos_smooth.size()) && std::isfinite(mos_smooth[f]))
|
|
? mos_smooth[f] : static_cast<float>(mosaicity_deg);
|
|
if (std::isfinite(cc[f])) { o.image_scale_cc = static_cast<float>(cc[f]); o.image_scale_cc_n = cc_n[f]; }
|
|
else { o.image_scale_cc.reset(); o.image_scale_cc_n.reset(); }
|
|
} else {
|
|
o.image_scale_g.reset();
|
|
o.image_scale_cc.reset();
|
|
o.image_scale_cc_n.reset();
|
|
o.mosaicity_deg.reset();
|
|
}
|
|
o.image_scale_b_factor_Ang2.reset();
|
|
o.image_scale_wedge_deg.reset();
|
|
}
|
|
}
|
|
|
|
namespace {
|
|
// Possible unique reflections per shell for the completeness column - mirrors CalcPossibleReflections.
|
|
void PossiblePerShell(int space_group_number, const UnitCell &cell, double d_min, double d_max,
|
|
const ResolutionShells &shells, bool merge_friedel, std::vector<int> &possible) {
|
|
gemmi::UnitCell gemmi_cell = cell;
|
|
const gemmi::SpaceGroup *sg = gemmi::find_spacegroup_by_number(space_group_number);
|
|
if (sg == nullptr) return;
|
|
const std::vector<gemmi::Miller> hkls = gemmi::make_miller_vector(gemmi_cell, sg, d_min, d_max, true);
|
|
const gemmi::GroupOps gops = sg->operations();
|
|
CrystalLattice lattice(cell);
|
|
const auto astar = lattice.Astar(), bstar = lattice.Bstar(), cstar = lattice.Cstar();
|
|
for (const auto &hkl : hkls) {
|
|
const auto q = hkl[0] * astar + hkl[1] * bstar + hkl[2] * cstar;
|
|
const auto qlen = q.Length();
|
|
if (qlen < 1e-6) continue;
|
|
const auto shell = shells.GetShell(1.0 / qlen);
|
|
if (!shell.has_value()) continue;
|
|
const int s = *shell;
|
|
if (s >= 0 && s < static_cast<int>(possible.size()))
|
|
possible[s] += (merge_friedel || gops.is_reflection_centric(hkl)) ? 1 : 2;
|
|
}
|
|
}
|
|
}
|
|
|
|
RotationScaleMerge::Result RotationScaleMerge::MergeAndStats(int n_groups, bool for_search,
|
|
const std::vector<char> &masked,
|
|
bool fulls_resident) {
|
|
// A full is usable for the merge / error model if it passes AddImage's filters (with the current
|
|
// ice/masked-ring context). group >= 0 already encodes "not absent and passes AcceptReflection".
|
|
auto masked_ring = [&](const Obs &o) {
|
|
if (masked.empty()) return false;
|
|
const int ring = IceRingIndex(o.d, ice_half_width_q);
|
|
return ring >= 0 && ring < static_cast<int>(masked.size()) && masked[ring];
|
|
};
|
|
auto usable_merge = [&](const Obs &o) {
|
|
if (o.group < 0) return false;
|
|
if (!frame_cell_ok[o.frame]) return false;
|
|
if (!(o.corr > 0.0f) || !std::isfinite(o.corr)) return false;
|
|
if (for_search && o.on_ice) return false;
|
|
if (masked_ring(o)) return false;
|
|
if (o.partiality < min_partiality) return false;
|
|
const float I_corr = o.I * o.corr, sigma_corr = o.sigma * o.corr;
|
|
return std::isfinite(I_corr) && std::isfinite(sigma_corr) && sigma_corr > 0.0f;
|
|
};
|
|
|
|
// The em-stats / samples / merge-accumulate / R_meas reductions run on the resident, scaled fulls
|
|
// (their group CSR is still on the device from scale-fulls) when fulls_resident; the host keeps the
|
|
// I2-sort, the (a,b) fit, the export and the statistics. reject_outliers is excluded upstream.
|
|
bool use_gpu_merge = false;
|
|
#ifdef JFJOCH_USE_CUDA
|
|
use_gpu_merge = fulls_resident && !fulls.empty();
|
|
std::vector<uint8_t> gpu_masked(masked.begin(), masked.end());
|
|
#endif
|
|
|
|
// ---- Error model: fit dev2 = a*sigma^2 + b^2*<I>^2 from symmetry-equivalent scatter. ----
|
|
std::vector<double> em_mean(n_groups, NAN);
|
|
std::vector<float> reject_median(n_groups, NAN);
|
|
double error_model_a = 1.0, error_model_b = 0.0, error_model_chi2 = 0.0;
|
|
bool error_model_active = false;
|
|
{
|
|
struct Sample { double s2, I2, dev2; };
|
|
std::vector<Sample> samples;
|
|
std::vector<int32_t> cnt(n_groups, 0); // per-group usable count (both paths; feeds reject-median)
|
|
bool did_gpu = false;
|
|
#ifdef JFJOCH_USE_CUDA
|
|
if (use_gpu_merge) {
|
|
const int nf = static_cast<int>(fulls.size());
|
|
std::vector<double> gs2(nf), gI2(nf), gdev2(nf);
|
|
std::vector<uint8_t> gvalid(nf);
|
|
gpu_->MergeEmSamples(for_search, gpu_masked.data(), static_cast<int>(gpu_masked.size()),
|
|
ice_half_width_q, min_partiality, em_mean.data(), cnt.data(),
|
|
gs2.data(), gI2.data(), gdev2.data(), gvalid.data());
|
|
samples.reserve(nf);
|
|
for (int i = 0; i < nf; ++i)
|
|
if (gvalid[i]) samples.push_back({gs2[i], gI2[i], gdev2[i]});
|
|
did_gpu = true;
|
|
}
|
|
#endif
|
|
if (!did_gpu) {
|
|
// Per-group inverse-variance mean over usable fulls (>=2 obs), and the leverage-corrected samples.
|
|
std::vector<double> sw(n_groups, 0.0), swI(n_groups, 0.0);
|
|
for (const auto &o : fulls) {
|
|
if (!usable_merge(o)) continue;
|
|
const double sigma_corr = static_cast<double>(o.sigma) * o.corr;
|
|
const double w = 1.0 / (sigma_corr * sigma_corr);
|
|
sw[o.group] += w; swI[o.group] += w * (static_cast<double>(o.I) * o.corr); cnt[o.group]++;
|
|
}
|
|
for (int g = 0; g < n_groups; ++g)
|
|
if (cnt[g] >= 2 && sw[g] > 0.0) em_mean[g] = swI[g] / sw[g];
|
|
|
|
samples.reserve(fulls.size());
|
|
for (const auto &o : fulls) {
|
|
if (!usable_merge(o) || cnt[o.group] < 2) continue;
|
|
const double mean = em_mean[o.group];
|
|
if (!std::isfinite(mean)) continue;
|
|
const double sigma_corr = static_cast<double>(o.sigma) * o.corr;
|
|
const double s2 = sigma_corr * sigma_corr;
|
|
const double w = 1.0 / s2;
|
|
const double factor = 1.0 - w / sw[o.group];
|
|
if (factor < 0.05) continue;
|
|
const double resid = static_cast<double>(o.I) * o.corr - mean;
|
|
samples.push_back({s2, mean * mean, resid * resid / factor});
|
|
}
|
|
}
|
|
|
|
// Per-group outlier-rejection median of I*corr (host both paths - a per-group median is awkward on
|
|
// the GPU; cheap here, cnt >= 2 filter from the em pass). Fed to the merge accumulate.
|
|
if (reject_outliers) {
|
|
std::vector<std::vector<float>> iv(n_groups);
|
|
for (const auto &o : fulls)
|
|
if (usable_merge(o) && cnt[o.group] >= 2)
|
|
iv[o.group].push_back(o.I * o.corr);
|
|
for (int g = 0; g < n_groups; ++g)
|
|
if (!iv[g].empty()) {
|
|
std::nth_element(iv[g].begin(), iv[g].begin() + iv[g].size() / 2, iv[g].end());
|
|
reject_median[g] = iv[g][iv[g].size() / 2];
|
|
}
|
|
}
|
|
constexpr int n_bins = 16;
|
|
// Fit (a, b) from the intensity-binned median deviations. Factored into a lambda so it can be
|
|
// re-run on a misfit-free pool below; takes the samples by value (it sorts them in place).
|
|
auto fit_ab = [&](std::vector<Sample> smp) {
|
|
if (smp.size() < static_cast<size_t>(8 * n_bins))
|
|
return;
|
|
std::sort(smp.begin(), smp.end(), [](const Sample &a, const Sample &b) { return a.I2 < b.I2; });
|
|
std::vector<double> bs2, bI2, bd2;
|
|
const size_t per = smp.size() / n_bins;
|
|
for (int bin = 0; bin < n_bins; ++bin) {
|
|
const size_t lo = bin * per;
|
|
const size_t hi = (bin == n_bins - 1) ? smp.size() : lo + per;
|
|
std::vector<double> vs2, vI2, vd2;
|
|
for (size_t i = lo; i < hi; ++i) {
|
|
vs2.push_back(smp[i].s2); vI2.push_back(smp[i].I2); vd2.push_back(smp[i].dev2);
|
|
}
|
|
bs2.push_back(median_of(vs2));
|
|
bI2.push_back(median_of(vI2));
|
|
bd2.push_back(median_of(vd2) / CHI2_1_MEDIAN);
|
|
}
|
|
std::vector<double> bd2_sorted = bd2;
|
|
const double dev2_floor = std::max(1e-30, 1e-3 * median_of(bd2_sorted));
|
|
double Ass = 0, AsI = 0, AII = 0, Bs = 0, BI = 0;
|
|
for (int bin = 0; bin < n_bins; ++bin) {
|
|
const double s2 = bs2[bin], I2 = bI2[bin], d2 = bd2[bin];
|
|
const double d2w = std::max(d2, dev2_floor);
|
|
const double wgt = 1.0 / (d2w * d2w);
|
|
Ass += wgt * s2 * s2; AsI += wgt * s2 * I2; AII += wgt * I2 * I2;
|
|
Bs += wgt * s2 * d2; BI += wgt * I2 * d2;
|
|
}
|
|
const double det = Ass * AII - AsI * AsI;
|
|
if (det > 1e-10 * Ass * AII) {
|
|
error_model_a = std::clamp((Bs * AII - BI * AsI) / det, 0.25, 100.0);
|
|
const double b2 = std::max((Ass * BI - AsI * Bs) / det, 0.0);
|
|
error_model_b = std::sqrt(b2);
|
|
error_model_active = true;
|
|
std::vector<double> chi2;
|
|
chi2.reserve(smp.size());
|
|
for (const auto &s : smp) {
|
|
const double v = error_model_a * s.s2 + b2 * s.I2;
|
|
if (v > 0.0) chi2.push_back(s.dev2 / v);
|
|
}
|
|
error_model_chi2 = chi2.empty() ? 0.0 : median_of(chi2) / CHI2_1_MEDIAN;
|
|
}
|
|
};
|
|
fit_ab(samples);
|
|
// Refit on a misfit-free pool: the merge drops symmetry outliers (|I - median| > reject_nsigma *
|
|
// sigma) from the merged intensity, so drop the equivalent samples (dev2 > reject_nsigma^2 * model
|
|
// variance) from the error-model fit too, keeping the fitted sigmas consistent with the reflections
|
|
// that actually survive. Operates on the shared `samples`, so CPU and GPU stay bit-identical.
|
|
if (reject_outliers && error_model_active) {
|
|
const double ns2 = reject_nsigma * reject_nsigma, b2 = error_model_b * error_model_b;
|
|
std::vector<Sample> kept;
|
|
kept.reserve(samples.size());
|
|
for (const auto &s : samples) {
|
|
const double v = error_model_a * s.s2 + b2 * s.I2;
|
|
if (v > 0.0 && s.dev2 <= ns2 * v) kept.push_back(s);
|
|
}
|
|
if (kept.size() >= static_cast<size_t>(8 * n_bins) && kept.size() < samples.size())
|
|
fit_ab(std::move(kept));
|
|
}
|
|
}
|
|
if (error_model_active)
|
|
logger.Info("Error model: a={:.3f} b={:.3f} ISa={:.1f} chi2={:.2f}", error_model_a, error_model_b,
|
|
error_model_b > 0 ? 1.0 / error_model_b : 0.0, error_model_chi2);
|
|
|
|
auto corrected_sigma = [&](float I_corr, float sigma_corr, int g) -> float {
|
|
if (!error_model_active) return sigma_corr;
|
|
const double I_for_b = std::isfinite(em_mean[g]) ? em_mean[g] : I_corr;
|
|
const double v = error_model_a * static_cast<double>(sigma_corr) * sigma_corr
|
|
+ (error_model_b * I_for_b) * (error_model_b * I_for_b);
|
|
return v > 0.0 ? static_cast<float>(std::sqrt(v)) : sigma_corr;
|
|
};
|
|
// ---- Merge: per-group inverse-variance sums with corrected sigma + deterministic half sets. ----
|
|
struct Accum { double swI = 0, sw = 0, swIh[2] = {0, 0}, swh[2] = {0, 0}; size_t nh[2] = {0, 0}; float d = NAN; };
|
|
std::vector<Accum> acc(n_groups);
|
|
size_t reject_count = 0;
|
|
std::vector<char> rejected_obs(fulls.size(), 0); // per-full outlier-rejected flag (mirrors the GPU path)
|
|
bool did_gpu_acc = false;
|
|
#ifdef JFJOCH_USE_CUDA
|
|
if (use_gpu_merge) {
|
|
std::vector<double> aswI(n_groups), asw(n_groups), aswIh0(n_groups), aswIh1(n_groups),
|
|
aswh0(n_groups), aswh1(n_groups), ad(n_groups);
|
|
std::vector<int32_t> anh0(n_groups), anh1(n_groups), arej(n_groups);
|
|
gpu_->MergeAccum(error_model_a, error_model_b, error_model_active,
|
|
reject_outliers, reject_nsigma, reject_median.data(),
|
|
aswI.data(), asw.data(), aswIh0.data(), aswIh1.data(),
|
|
aswh0.data(), aswh1.data(), anh0.data(), anh1.data(), ad.data(), arej.data());
|
|
for (int g = 0; g < n_groups; ++g) {
|
|
Accum &a = acc[g];
|
|
a.swI = aswI[g]; a.sw = asw[g]; a.swIh[0] = aswIh0[g]; a.swIh[1] = aswIh1[g];
|
|
a.swh[0] = aswh0[g]; a.swh[1] = aswh1[g];
|
|
a.nh[0] = static_cast<size_t>(anh0[g]); a.nh[1] = static_cast<size_t>(anh1[g]);
|
|
a.d = static_cast<float>(ad[g]);
|
|
reject_count += static_cast<size_t>(arej[g]);
|
|
}
|
|
did_gpu_acc = true;
|
|
}
|
|
#endif
|
|
if (!did_gpu_acc)
|
|
for (const auto &o : fulls) {
|
|
if (!usable_merge(o)) continue;
|
|
const float I_corr = o.I * o.corr;
|
|
float sigma_corr = o.sigma * o.corr;
|
|
sigma_corr = corrected_sigma(I_corr, sigma_corr, o.group);
|
|
if (reject_outliers && error_model_active && std::isfinite(reject_median[o.group])
|
|
&& std::fabs(I_corr - reject_median[o.group]) > reject_nsigma * sigma_corr) {
|
|
++reject_count;
|
|
rejected_obs[&o - fulls.data()] = 1;
|
|
continue;
|
|
}
|
|
const double w = 1.0 / (static_cast<double>(sigma_corr) * sigma_corr);
|
|
const double wI = w * I_corr;
|
|
const int half = HalfForImage(o.frame);
|
|
auto &a = acc[o.group];
|
|
a.swI += wI; a.sw += w;
|
|
a.swIh[half] += wI; a.swh[half] += w; a.nh[half]++;
|
|
if (!std::isfinite(a.d) && std::isfinite(o.d) && o.d > 0.0f) a.d = o.d;
|
|
}
|
|
// ---- Export merged reflections (+ resolution-shell R-free flags). ----
|
|
Result result;
|
|
result.isa = error_model_b > 0 ? 1.0 / error_model_b : 0.0;
|
|
std::vector<double> merged_I(n_groups, NAN);
|
|
float d_min = std::numeric_limits<float>::max(), d_max = 0.0f;
|
|
for (int g = 0; g < n_groups; ++g) {
|
|
const auto &a = acc[g];
|
|
if (a.sw <= 0.0) continue;
|
|
MergedReflection mr{};
|
|
mr.h = group_h[g]; mr.k = group_k[g]; mr.l = group_l[g];
|
|
mr.I = static_cast<float>(a.swI / a.sw);
|
|
mr.sigma = SigmaWithSystematicFloor(1.0 / std::sqrt(a.sw), mr.I, error_model_b);
|
|
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);
|
|
}
|
|
}
|
|
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);
|
|
ApplyFrenchWilson(result.merged, x.GetSpaceGroupNumber().value_or(1));
|
|
|
|
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;
|
|
}
|
|
|
|
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;
|
|
}
|
|
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;
|
|
logger.Info("Merge complete ({} unique reflections)", result.merged.size());
|
|
return result;
|
|
}
|
|
|
|
RotationScaleMerge::Result RotationScaleMerge::Run(bool for_search,
|
|
const std::vector<char> &masked_ice_rings) {
|
|
const int sg_number = x.GetSpaceGroupNumber().value_or(1);
|
|
HKLKeyGenerator keygen(merge_friedel, sg_number);
|
|
|
|
// --- 1. Per-frame partial scaling (Rotation model, per-image G only). ---
|
|
const int n_groups = ComputeAsuGroups(keygen); // one ASU grouping, shared by partials and fulls
|
|
std::vector<double> partial_mean;
|
|
bool scaled_on_gpu = false;
|
|
#ifdef JFJOCH_USE_CUDA
|
|
if (gpu_active_) {
|
|
// Run the whole scaling loop on the GPU; corr stays RESIDENT across scaling -> smooth-G -> CC ->
|
|
// combine (and across passes, exactly as the old host round-trip did). Only the per-frame G/scaled
|
|
// come back (needed to compute smooth-G and the writeback).
|
|
gpu_->ScalePartials(scaling_iter, SCALE_ROBUST_K, min_partiality, d_min_limit.has_value());
|
|
frame_scaled_scratch.assign(n_frames, 0);
|
|
gpu_->GetG(g_partial.data(), frame_scaled_scratch.data());
|
|
scaled_on_gpu = true;
|
|
}
|
|
#endif
|
|
if (!scaled_on_gpu) {
|
|
for (int it = 0; it < scaling_iter; ++it) {
|
|
ReduceGroupMeans(partials, n_groups, false, {}, partial_mean);
|
|
FitPerFrameG(partials, frame_start, frame_count, partial_mean, /*unity=*/false, g_partial);
|
|
UpdateCorr(partials, g_partial, frame_scaled_scratch);
|
|
}
|
|
}
|
|
const std::vector<uint8_t> partial_scaled = frame_scaled_scratch;
|
|
|
|
// --- 2. Smooth G across frames (XDS DELPHI-like) before the combine. ---
|
|
const auto s = x.GetScalingSettings();
|
|
const double smooth_g_deg = s.GetSmoothGDegrees();
|
|
const auto gonio = x.GetGoniometer();
|
|
const double osc_deg = gonio ? std::fabs(gonio->GetIncrement_deg()) : 0.0;
|
|
if (smooth_g_deg > 0.0 && osc_deg > 1e-6) {
|
|
int window = std::max(1, static_cast<int>(std::lround(smooth_g_deg / osc_deg)));
|
|
if (window % 2 == 0) ++window;
|
|
bool smoothed_on_gpu = false;
|
|
#ifdef JFJOCH_USE_CUDA
|
|
if (gpu_active_) {
|
|
// Apply smooth-G to the resident corr: the host builds the per-frame ratio g/g_smooth (+ an
|
|
// apply flag), the kernel multiplies corr in place. Same guard / arithmetic as host SmoothG.
|
|
std::vector<double> g_smooth;
|
|
ComputeSmoothGWindow(g_partial, window, g_smooth);
|
|
std::vector<uint8_t> apply(n_frames, 0);
|
|
std::vector<double> ratio(n_frames, 1.0);
|
|
for (int f = 0; f < n_frames; ++f)
|
|
if (frame_scaled_scratch[f] && std::isfinite(g_partial[f]) && g_partial[f] > 0.0
|
|
&& std::isfinite(g_smooth[f])) {
|
|
apply[f] = 1;
|
|
ratio[f] = g_partial[f] / g_smooth[f];
|
|
}
|
|
gpu_->SmoothCorr(apply.data(), ratio.data());
|
|
for (int f = 0; f < n_frames; ++f) if (apply[f]) g_partial[f] = g_smooth[f];
|
|
smoothed_on_gpu = true;
|
|
}
|
|
#endif
|
|
if (!smoothed_on_gpu)
|
|
SmoothG(partials, g_partial, window);
|
|
}
|
|
|
|
#ifdef JFJOCH_USE_CUDA
|
|
// The GPU keeps corr resident through scaling + smooth-G; only the diagnostic dump falls back to the
|
|
// CPU combine, which reads host partials[].corr, so refresh it for that path. (The decay/absorption
|
|
// correction surfaces run on the GPU-combined fulls - see below - so they do NOT need the CPU combine.)
|
|
if (gpu_active_ && !observation_dump_path.empty()) {
|
|
std::vector<float> corr(partials.size());
|
|
gpu_->GetCorr(corr.data());
|
|
for (size_t i = 0; i < partials.size(); ++i) partials[i].corr = corr[i];
|
|
}
|
|
#endif
|
|
|
|
// Per-frame CC + write G/CC/mosaicity back onto the partials (once). On the GPU the group means +
|
|
// per-frame CC run on the resident (already smoothed) corr; only the tiny per-frame cc/cc_n come back.
|
|
std::vector<double> cc;
|
|
std::vector<int64_t> cc_n;
|
|
bool cc_on_gpu = false;
|
|
#ifdef JFJOCH_USE_CUDA
|
|
if (gpu_active_) {
|
|
cc.resize(n_frames);
|
|
cc_n.resize(n_frames);
|
|
gpu_->ComputePartialCC(min_partiality, cc.data(), cc_n.data());
|
|
cc_on_gpu = true;
|
|
}
|
|
#endif
|
|
if (!cc_on_gpu) {
|
|
ReduceGroupMeans(partials, n_groups, false, {}, partial_mean);
|
|
ComputePerFrameCC(partial_mean, cc, cc_n);
|
|
}
|
|
FinalizePerFrameScale(cc, cc_n, partial_scaled);
|
|
|
|
// --- 3. 3D combine of per-frame partials into fulls (fulls inherit their ASU group here). ---
|
|
bool combined_on_gpu = false;
|
|
bool scaled_fulls_on_gpu = false;
|
|
#ifdef JFJOCH_USE_CUDA
|
|
// GPU combine (+ scale-fulls) keeps the fulls resident on the device: combine, then build the frame /
|
|
// ASU-group CSRs on the host from just the small key arrays (a deterministic counting sort - no GPU
|
|
// stable-sort), scale the fulls in place, and download only once. Mirrors Combine() + the Unity
|
|
// scale-fulls loop below. Only the diagnostic dump (serial, one writer) falls back to CPU. The
|
|
// decay/absorption correction surfaces run on the downloaded fulls (px/py is carried through the
|
|
// combine) and their corrected corr is re-uploaded before the resident merge - so they stay on the GPU.
|
|
if (gpu_active_ && observation_dump_path.empty()) {
|
|
// The smoothed corr is already resident (scaling + smooth-G ran on the device, no round-trip).
|
|
const int nf = gpu_->Combine(rawrun_group.data(), min_partiality, capture_uncertainty_coeff,
|
|
min_captured_fraction);
|
|
g_full.assign(n_frames, 1.0);
|
|
|
|
if (scale_fulls && nf > 0) {
|
|
// Frame + group CSRs over the emit-ordered fulls, built by counting sort on the host (stable,
|
|
// deterministic). frame is always in [0, n_frames); group is <0 for absent/out-of-range fulls.
|
|
std::vector<int32_t> ff(nf), fg(nf);
|
|
gpu_->GetFullsKeys(ff.data(), fg.data());
|
|
std::vector<int32_t> f_start(n_frames, 0), f_count(n_frames, 0), f_perm(nf);
|
|
for (int i = 0; i < nf; ++i) ++f_count[ff[i]];
|
|
for (int f = 1; f < n_frames; ++f) f_start[f] = f_start[f - 1] + f_count[f - 1];
|
|
{ std::vector<int32_t> fill = f_start; for (int i = 0; i < nf; ++i) f_perm[fill[ff[i]]++] = i; }
|
|
gpu_->SetFullsFrameCSR(f_perm.data(), nf, f_start.data(), f_count.data());
|
|
|
|
std::vector<int32_t> g_count(n_groups, 0), g_start(n_groups, 0);
|
|
for (int i = 0; i < nf; ++i) if (fg[i] >= 0) ++g_count[fg[i]];
|
|
int acc = 0;
|
|
for (int g = 0; g < n_groups; ++g) { g_start[g] = acc; acc += g_count[g]; }
|
|
std::vector<int32_t> g_perm(acc);
|
|
{ std::vector<int32_t> fill = g_start; for (int i = 0; i < nf; ++i) if (fg[i] >= 0) g_perm[fill[fg[i]]++] = i; }
|
|
gpu_->SetFullsGroups(g_perm.data(), acc, g_start.data(), g_count.data());
|
|
|
|
gpu_->ScaleFulls(scaling_iter, SCALE_ROBUST_K, min_partiality);
|
|
scaled_fulls_on_gpu = true;
|
|
}
|
|
|
|
fulls.assign(nf, Obs{});
|
|
std::vector<int32_t> fh(nf), fk(nf), fl(nf), fframe(nf), fgroup(nf);
|
|
std::vector<float> fI(nf), fsig(nf), fd(nf), fimg(nf), fcorr(nf, 1.0f), fpx(nf), fpy(nf);
|
|
std::vector<uint8_t> fon(nf);
|
|
gpu_->GetFulls(fh.data(), fk.data(), fl.data(), fI.data(), fsig.data(), fd.data(),
|
|
fimg.data(), fframe.data(), fon.data(), fgroup.data());
|
|
gpu_->GetFullsPxPy(fpx.data(), fpy.data());
|
|
if (scaled_fulls_on_gpu) gpu_->GetFullsCorr(fcorr.data());
|
|
for (int i = 0; i < nf; ++i) {
|
|
Obs &o = fulls[i];
|
|
o.h = fh[i]; o.k = fk[i]; o.l = fl[i];
|
|
o.I = fI[i]; o.sigma = fsig[i]; o.d = fd[i];
|
|
o.rlp = 1.0f; o.partiality = 1.0f; o.corr = fcorr[i];
|
|
o.image_number = fimg[i]; o.frame = fframe[i]; o.px = fpx[i]; o.py = fpy[i];
|
|
o.on_ice = fon[i]; o.group = fgroup[i];
|
|
}
|
|
logger.Info("3D combine{} (GPU): {} fulls", scaled_fulls_on_gpu ? " + scale-fulls" : "", nf);
|
|
combined_on_gpu = true;
|
|
}
|
|
#endif
|
|
if (!combined_on_gpu)
|
|
Combine();
|
|
|
|
// --- 4. Scale the fulls (XDS order, Unity model). ---
|
|
if (scale_fulls && !scaled_fulls_on_gpu) {
|
|
std::vector<double> full_mean;
|
|
for (int it = 0; it < scaling_iter; ++it) {
|
|
ReduceGroupMeans(fulls, n_groups, false, {}, full_mean);
|
|
FitPerFrameG(fulls, fulls_frame_start, fulls_frame_count, full_mean, /*unity=*/true, g_full);
|
|
UpdateCorr(fulls, g_full, frame_scaled_scratch);
|
|
}
|
|
logger.Info("Scaled fulls (XDS order, Unity model)");
|
|
}
|
|
|
|
// --- 4b. Optional correction surfaces (decay = resolution x time; absorption = goniometer-frame
|
|
// diffracted-beam direction), each an alternating multiplicative fit of the fulls' corr against
|
|
// the merged reference. Cheap host loops over the downloaded fulls; applied only on the final
|
|
// in-symmetry merge, never the P1 space-group search pass (corrections there add risk and can
|
|
// perturb the symmetry determination). ---
|
|
[[maybe_unused]] const bool corrections = !for_search && (refine_decay_b || absorption_iter > 0);
|
|
if (!for_search && refine_decay_b)
|
|
RefineDecay(n_groups);
|
|
if (!for_search && absorption_iter > 0)
|
|
RefineAbsorption(absorption_iter, n_groups);
|
|
#ifdef JFJOCH_USE_CUDA
|
|
// The corrections mutate the host fulls' corr; when the merge runs on the resident (GPU) fulls, push
|
|
// the corrected corr back to the device so the merge reads it.
|
|
if (corrections && combined_on_gpu && scaled_fulls_on_gpu) {
|
|
std::vector<float> fcorr(fulls.size());
|
|
for (size_t i = 0; i < fulls.size(); ++i) fcorr[i] = fulls[i].corr;
|
|
gpu_->SetFullsCorr(fcorr.data());
|
|
}
|
|
#endif
|
|
|
|
// --- 5. Error model + merge + statistics. ---
|
|
auto r = MergeAndStats(n_groups, for_search, masked_ice_rings, combined_on_gpu && scaled_fulls_on_gpu);
|
|
return r;
|
|
}
|