leonarski_fandClaude Opus 4.8
fccf9b83e7
Remove the [rsm] per-stage lap timing and the JFJOCH_RSM_NO_GPU / JFJOCH_RSM_CPU_COMBINE env gates now that the GPU-resident path is the validated default (it runs whenever a GPU is present, with the CPU loops as the bit-parity fallback; the diagnostic-dump path still uses the CPU combine). Honour a fixed (forced) mosaicity: SmoothMosaicityAndPartiality now overrides every frame with GetForcedMosaicity() when set, instead of always reading the per-frame integration value - so the caller can route the --mosaicity case through RotationScaleMerge (its partiality recompute makes it a natural fit) rather than a separate path. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
196 lines
11 KiB
C++
196 lines
11 KiB
C++
// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
|
|
// SPDX-License-Identifier: GPL-3.0-only
|
|
|
|
#pragma once
|
|
|
|
#include <cstdint>
|
|
#include <optional>
|
|
#include <vector>
|
|
|
|
#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 <memory>
|
|
#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). It does NOT support B-factor refinement, external-reference
|
|
// scaling, an absorption surface or wedge refinement - the caller rejects those combinations. Stills use
|
|
// the per-image ScaleOnTheFly (fixed partiality) instead.
|
|
class RotationScaleMerge {
|
|
public:
|
|
struct Result {
|
|
std::vector<MergedReflection> 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<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 = {});
|
|
|
|
// 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<char> &masked_ice_rings = {});
|
|
|
|
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 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<IntegrationOutcome> &partials_out; // written back at the end of scaling
|
|
std::optional<UnitCell> 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<double> d_min_limit;
|
|
bool merge_friedel = true;
|
|
double capture_uncertainty_coeff = 0.0;
|
|
double reject_nsigma = 0.0;
|
|
bool reject_outliers = false;
|
|
double rfree_fraction = 0.0;
|
|
int scaling_iter = 3;
|
|
bool scale_fulls = true;
|
|
double mosaicity_deg = 0.1;
|
|
float ice_half_width_q = 0.0f;
|
|
|
|
// Flat buffers, allocated once by Ingest() and reused across Run() calls.
|
|
std::vector<Obs> partials; // all per-frame partials, grouped by frame
|
|
std::vector<int32_t> frame_start, frame_count; // CSR ranges of `partials` per frame
|
|
std::vector<uint8_t> frame_cell_ok; // per-frame cell-consistency mask (1 = kept)
|
|
std::vector<uint8_t> finite_ok; // per-obs AcceptReflection finiteness (immutable; 1 = kept)
|
|
std::vector<double> 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<int32_t> perm;
|
|
std::vector<int32_t> rawrun_start, rawrun_count;
|
|
std::vector<int32_t> rawrun_h, rawrun_k, rawrun_l;
|
|
std::vector<float> rawrun_d; // representative resolution per raw hkl
|
|
std::vector<int32_t> rawrun_group; // dense ASU-group id per raw hkl (<0 = absent/out of range)
|
|
|
|
std::vector<Obs> fulls; // combined fulls (rebuilt each Run), sorted by frame
|
|
std::vector<int32_t> fulls_frame_start, fulls_frame_count; // CSR ranges of `fulls` per frame
|
|
std::vector<double> 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<uint8_t> 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<float> mos_smooth;
|
|
|
|
// Working per-group arrays (sized to the current group count; reused).
|
|
std::vector<int32_t> 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<RotationScaleMergeGPU> 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> &obs, int n_groups,
|
|
bool exclude_ice, const std::vector<char> &masked_ice_rings,
|
|
std::vector<double> &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> &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);
|
|
|
|
// corr = rlp / (partiality * G[frame]); leaves corr unchanged for frames that could not be fit.
|
|
void UpdateCorr(std::vector<Obs> &obs, const std::vector<double> &g,
|
|
const std::vector<uint8_t> &frame_scaled) const;
|
|
|
|
void SmoothG(std::vector<Obs> &obs, std::vector<double> &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<double> &g, int window,
|
|
std::vector<double> &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)
|
|
|
|
// 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<double> &partial_group_mean,
|
|
std::vector<double> &cc, std::vector<int64_t> &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<double> &cc, const std::vector<int64_t> &cc_n,
|
|
const std::vector<uint8_t> &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<char> &masked_ice_rings,
|
|
bool fulls_resident);
|
|
};
|