leonarski_f
dd0bffb283
Build Packages / Unit tests (push) Skipped
Build Packages / build:windows:nocuda (push) Successful in 11m6s
Build Packages / build:rpm (rocky8_nocuda) (push) Successful in 10m27s
Build Packages / build:rpm (rocky9_nocuda) (push) Successful in 10m54s
Build Packages / build:rpm (ubuntu2204_nocuda) (push) Successful in 9m25s
Build Packages / build:rpm (ubuntu2404_nocuda) (push) Successful in 10m5s
Build Packages / build:rpm (rocky8_sls9) (push) Successful in 11m33s
Build Packages / build:rpm (rocky9_sls9) (push) Successful in 11m19s
Build Packages / build:rpm (rocky8) (push) Successful in 12m23s
Build Packages / build:rpm (rocky9) (push) Successful in 13m21s
Build Packages / build:rpm (ubuntu2204) (push) Successful in 12m30s
Build Packages / build:rpm (ubuntu2404) (push) Successful in 11m55s
Build Packages / DIALS test (push) Successful in 13m42s
Build Packages / XDS test (durin plugin) (push) Successful in 9m26s
Build Packages / XDS test (JFJoch plugin) (push) Successful in 6m41s
Build Packages / XDS test (neggia plugin) (push) Successful in 6m12s
Build Packages / Generate python client (push) Successful in 19s
Build Packages / Build documentation (push) Successful in 52s
Build Packages / Create release (push) Skipped
Build Packages / build:viewer-tgz:cpu (push) Successful in 5m29s
Build Packages / build:viewer-tgz:cuda (push) Successful in 6m12s
Build Packages / build:windows:cuda (push) Successful in 18m36s
This is an UNSTABLE release. It includes many experimental features, as well as many AI generated fixes. We recommend using rc.152 for production use. * rugnux: Add `--model model.pdb` - score the merged data against an atomic model and compute initial maps. It reports R-work/R-free (scaling the model to the observed amplitudes with an overall scale, an anisotropic B and a flat bulk solvent - the standard few-parameter model, so a batch of maps stays directly comparable) and writes 2Fo-Fc / Fo-Fc electron-density maps (CCP4) plus a map-coefficient MTZ. The structure itself is not refined; the model is only re-fractionalised into the data cell. * rugnux: The merged reflection output now carries French-Wilson amplitudes (|F| and its sigma) next to the intensities - MTZ `F`/`SIGF`, mmCIF `_refln.F_meas_au`, and the text HKL - computed with the correct centric/acentric Wilson prior and epsilon multiplicity, so a downstream program (e.g. phenix.refine) can refine against amplitudes. The intensity columns are unchanged. * rugnux: R-free test-set flags are now assigned deterministically and consistently across symmetry - a Bijvoet pair I(+)/I(-) is never split between the work and free sets, and the assignment is a reproducible per-hkl hash that depends only on the reflection index, so every dataset of one crystal form gets the same ~5% free set (what a multi-dataset campaign such as PanDDA needs). On small data the fraction is floored so the test set stays large enough for a stable R-free (~500 reflections, capped at 10%); it stays flat at 5% on ordinary data. When a reference MTZ carries a `FreeR_flag` column its test set is imported instead, letting a whole campaign inherit one shared free set. * rugnux: A reference MTZ (`--reference-mtz`) can now fix the space group and cell for rotation data too (previously rejected), without being used to scale - the rotation merge stays self-consistent. When the crystal has an indexing (merohedral) ambiguity - a lattice symmetry higher than its Laue symmetry, e.g. P3/P4/P6/C2 - the reference also resolves it: each candidate reindexing (identity plus the twin-law cosets of the metric symmetry) is scored by its intensity correlation against the reference and the data are re-merged in the best-correlating one. This is a metric-preserving relabelling of hkl (the cell is unchanged) and a no-op for a holohedral crystal such as lysozyme. * rugnux: `--model` validation now aligns the data to the model before scoring - the observed reflections are reindexed into the model's enantiomorph when the two differ only by hand (indistinguishable from merged intensities). A merohedral indexing ambiguity is resolved against the reference MTZ when one is given (so a whole campaign shares one indexing convention); only with a model and no reference does validation fall back to fitting each candidate reindexing and keeping the lowest R-free. * rugnux: De-novo symmetry - recover a genuine high-symmetry group whose data are imperfectly scaled. Such a merge's within-orbit chi² lands just past the self-consistency bound (each real symmetry step adds a little systematic scatter), right where a merohedral twin also lands, so the chi² ratio alone cannot separate them. The candidate is now rescued when the extra intensity-proportional systematic error it invokes stays small relative to the confirmed subgroup - a genuine symmetry step gains multiplicity without inflating the merge error model's b, whereas a twin forces non-equivalent reflections together and b balloons. Fixes cubic insulin (I23 instead of I222) with no change to any other crystal in the test battery, including the twins that must stay in their lower symmetry. * Docs: Document the French-Wilson amplitude estimation, R-free flagging, reference-based space-group/ambiguity resolution, and model-based validation/maps in CPU_DATA_ANALYSIS.md. * Frontend: The status-bar pill now shows a progress bar during detector calibration (previously only during measurement), and the calibration state and its button are labelled "Calibration"/"CALIBRATE" (the internal `Pedestal` state name is unchanged for back-compatibility).Reviewed-on: #69 Co-authored-by: Filip Leonarski <filip.leonarski@psi.ch>
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;
|
|
}
|