// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute // SPDX-License-Identifier: GPL-3.0-only #pragma once #include #include #include #include "../../common/DiffractionExperiment.h" #include "../../common/Logger.h" #include "../../common/Reflection.h" #include "../../common/UnitCell.h" #include "../IntegrationOutcome.h" #include "Merge.h" // MergedReflection, MergeStatistics #ifdef JFJOCH_USE_CUDA #include #include "RotationScaleMergeGPU.h" #endif // Dedicated, allocate-once scale+combine+merge for rotation data (the -P rot3d path): recompute the // per-frame partiality from the (smoothed) mosaicity, robustly fit a per-image scale G, 3D-combine each // rocking event's partials into fulls, refit a per-frame scale on the fulls (XDS order), and merge with // a global error model. // // The per-frame partial observations are ingested ONCE into flat vectors; the hkl->ASU grouping is // computed once per space group (by a sort, not a map) and reused across all scaling iterations; every // hot step is a flat loop over those vectors, so the whole pipeline maps onto CUDA kernels (segmented // reduction + per-frame solve) and runs GPU-resident when a GPU is present, with the CPU loops as the // bit-parity fallback. CC1/2 and the per-image CC are computed once at the end, not every iteration. // // Used only for the self-scaling rotation case with per-image G (Rotation partiality, a fixed/forced // mosaicity is honoured by the recompute). Post-scale-fulls correction stages (on by default via // ScalingSettings::CorrectionSurfaces): a global Debye-Waller decay and a goniometer-frame absorption // surface, both fitted on the host and pushed back to the resident (GPU) fulls before the merge. // External-reference scaling, the stills B-factor and wedge refinement are unsupported (caller rejects). // Stills use the per-image ScaleOnTheFly (fixed partiality) instead. class RotationScaleMerge { public: struct Result { std::vector merged; MergeStatistics statistics; double isa = 0.0; // 1/b of the fitted error model (0 if the model stayed at identity) }; // experiment: read live (its space group is changed by the caller between Run() calls). // partial_outcomes: the per-frame partials; the final per-frame scale (G, CC, mosaicity) is written // back onto them so the offline per-image scaling table is still exported. // reference_cell: the consensus cell (for the completeness count and the cell-consistency mask). RotationScaleMerge(const DiffractionExperiment &experiment, std::vector &partial_outcomes, std::optional reference_cell, int scaling_iterations, float ice_ring_half_width_q, size_t nthreads, Logger &logger, std::string observation_dump_path = {}); // Copy the per-frame partials into the flat buffers. Call once before the first Run(). void Ingest(); // Scale (per-frame G) -> smooth G -> 3D combine -> scale fulls -> merge -> error model -> statistics // for the space group currently set on the experiment, reusing the ingested buffers. // for_search: the de-novo P1 pass whose merged intensities feed the space-group search - ice-ring // reflections are dropped from the merge and the error model (kept otherwise, for completeness). // masked_ice_rings: rings (indices into ICE_RING_RES_A) to drop from the final merge; empty = none. Result Run(bool for_search, const std::vector &masked_ice_rings = {}); // Override the high-resolution cut for the next Run() - used to gate the de-novo P1 search pass at // >= 1 without cutting the final in-symmetry merge. Reset to the manual limit afterwards. void SetDMinLimit(std::optional d_min_A) { d_min_limit = d_min_A; } private: // One integrated observation - a per-frame partial during scaling/combine, or a combined full during // scale-fulls/merge. Flat (not nested per image); a POD so the arrays translate straight to CUDA. struct Obs { int32_t h, k, l; float I, sigma, d, rlp, partiality, zeta, delta_phi, bkg; float px = NAN, py = NAN; // predicted detector position (for the absorption surface; CPU path only) float image_number; // fractional frame position (for 3D-combine contiguity) int32_t frame; // index of the outcome whose per-frame scale G applies to this obs uint8_t on_ice; float corr; // image_scale_corr (working; updated by scaling) int32_t group; // dense ASU-group id for the current space group; <0 = never mergeable }; const DiffractionExperiment &x; std::vector &partials_out; // written back at the end of scaling std::optional reference_cell; size_t nthreads; Logger &logger; std::string observation_dump_path; // Fixed settings snapshot (read once in the ctor). int n_frames = 0; double min_partiality = 0.02; std::optional d_min_limit; bool merge_friedel = true; double capture_uncertainty_coeff = 0.0; double min_captured_fraction = 0.0; double reject_nsigma = 0.0; bool reject_outliers = false; double rfree_fraction = 0.0; int scaling_iter = 3; bool scale_fulls = true; bool refine_decay_b = false; // per-time-block Debye-Waller decay correction (radiation damage) int absorption_iter = 0; // >0: fit a goniometer-frame absorption surface over this many iterations double mosaicity_deg = 0.1; float ice_half_width_q = 0.0f; // Automatic high-resolution cutoff for the written reflections + reported shells (post-merge; the // scaling, combine and error model always run over the full range). Manual d_min_limit wins. ResolutionCutoffMethod resolution_cutoff_method = ResolutionCutoffMethod::Off; double resolution_cc_target = 0.30; int report_shell_count = 10; // Flat buffers, allocated once by Ingest() and reused across Run() calls. std::vector partials; // all per-frame partials, grouped by frame std::vector frame_start, frame_count; // CSR ranges of `partials` per frame std::vector frame_cell_ok; // per-frame cell-consistency mask (1 = kept) std::vector finite_ok; // per-obs AcceptReflection finiteness (immutable; 1 = kept) std::vector g_partial; // per-frame partial scale G // Raw-hkl ordering, built ONCE by Ingest and reused: `perm` lists partial indices sorted by // (raw h,k,l, image_number); each distinct raw hkl is a contiguous run [rawrun_start, +count) of it. // The expensive sort happens once here, so per-pass combine (event split) and ASU grouping are linear. std::vector perm; std::vector rawrun_start, rawrun_count; std::vector rawrun_h, rawrun_k, rawrun_l; std::vector rawrun_d; // representative resolution per raw hkl std::vector rawrun_group; // dense ASU-group id per raw hkl (<0 = absent/out of range) std::vector fulls; // combined fulls (rebuilt each Run), sorted by frame std::vector fulls_frame_start, fulls_frame_count; // CSR ranges of `fulls` per frame std::vector g_full; // per-frame scale on the fulls // Set by FitPerFrameG: which frames were fitted this call (so corr/G is updated only there). std::vector frame_scaled_scratch; // Per-frame mosaicity smoothed in frame order (deterministic); used to recompute partiality and // written back for the per-image scaling table. Empty if there is no per-frame mosaicity. std::vector mos_smooth; // Working per-group arrays (sized to the current group count; reused). std::vector group_h, group_k, group_l; #ifdef JFJOCH_USE_CUDA // GPU engine: the whole hot path (scaling, combine, scale-fulls, per-frame CC, smooth-G, merge + // error model) runs on the device, resident, when a GPU is present. Null / inactive otherwise, with // the CPU loops as the bit-parity fallback. Built in Ingest. std::unique_ptr gpu_; bool gpu_active_ = false; #endif // --- helpers (each a flat pass; see the .cpp) --- // Compute the dense ASU-group id for the current space group by grouping the (pre-sorted) raw-hkl // runs by their ASU key - one gemmi ASU reduction per distinct raw hkl, not per observation. Fills // rawrun_group, the group_h/k/l representative tables, and partials[].group; returns the group count. int ComputeAsuGroups(const HKLKeyGenerator &key_generator); // Inverse-variance per-group mean of I*corr over `obs` (the merge reference). exclude_ice/masked drop // those reflections (used for the error-model/merge means, not the scaling reference). void ReduceGroupMeans(const std::vector &obs, int n_groups, bool exclude_ice, const std::vector &masked_ice_rings, std::vector &out_mean) const; // Robust per-frame G fit (IRLS, Cauchy k=3), unity=false uses the rotation partiality, unity=true the // scale-fulls (partiality already folded in). Reads out_mean[group] as the reference intensity. void FitPerFrameG(std::vector &obs, const std::vector &fstart, const std::vector &fcount, const std::vector &group_mean_in, bool unity, std::vector &g); // corr = rlp / (partiality * G[frame]); leaves corr unchanged for frames that could not be fit. void UpdateCorr(std::vector &obs, const std::vector &g, const std::vector &frame_scaled) const; void SmoothG(std::vector &obs, std::vector &g, int window) const; // The windowed geometric mean of G over frames (the shared first half of SmoothG); the GPU path // applies the resulting ratio to the resident corr in a kernel instead of the host obs loop. void ComputeSmoothGWindow(const std::vector &g, int window, std::vector &g_smooth) const; // Smooth per-frame mosaicity in frame order and recompute each partial's partiality from it, so the // per-frame partials of one rocking event tile the curve consistently (they sum toward 1) before the // 3D combine. Deterministic (frame order); replaces the old arrival-order mosaicity moving average // that prediction applied. SG-independent, so done once in Ingest. void SmoothMosaicityAndPartiality(); void Combine(); // partials -> fulls (CPU) // Post-scale-fulls correction surfaces, each an alternating multiplicative fit of the host fulls' corr // against the merged reference (cheap host loops; the corrected corr is re-uploaded to the resident // fulls afterwards). Each is cross-validated (fit even frames, keep only if held-out odd equivalents // improve) so it is a no-op when its systematic is absent. RefineDecay fits a global Debye-Waller B // (resolution x time - radiation damage the resolution-flat per-frame G cannot capture; also gated on a // physical total-dB floor). RefineAbsorption fits a smooth factor over the diffracted-beam direction in // the goniometer frame (path-length / absorption; negligible at hard X-rays, matters at low energy). void RefineDecay(int n_groups); void RefineAbsorption(int n_iter, int n_groups); // Sort `fulls` by peak frame and (re)build fulls_frame_start/count (the per-frame CSR the scale-fulls // step slices). Shared by the CPU Combine tail and the GPU combine path. void SortFullsByFrame(); // Per-frame CC vs the partial merge reference (CPU; the GPU equivalent is gpu_->ComputePartialCC). void ComputePerFrameCC(const std::vector &partial_group_mean, std::vector &cc, std::vector &cc_n) const; // Write G/CC/mosaicity back onto the partials (once, at the end of partial scaling) from the given // per-frame cc/cc_n, so the offline per-image scaling table is still exported. void FinalizePerFrameScale(const std::vector &cc, const std::vector &cc_n, const std::vector &frame_scaled); // Error model + merge + statistics over the fulls (the last stage). n_groups is the fulls group count. // fulls_resident: the (scaled) fulls + their group CSR are still on the GPU, so the em-stats / samples // / merge-accumulate / R_meas reductions run there (only per-group + samples come back). Result MergeAndStats(int n_groups, bool for_search, const std::vector &masked_ice_rings, bool fulls_resident); };