Jungfraujoch/docs/GPU_DATA_ANALYSIS.md

GPU data analysis

Jungfraujoch runs most of its crystallographic analysis on NVIDIA GPUs where one is present: the per-image loop of jfjoch_broker, jfjoch_viewer and rugnux, and most of rugnux's dataset-level steps - the pre-scan, scaling and merging, model validation and molecular replacement. This page lists the GPU engines (all of them in .cu files under image_analysis/), with what each computes. The methods themselves are described once, in the CPU/GPU data analysis reference and its siblings, which each entry links to.

Every engine has a CPU counterpart that computes the same thing, and the CPU-only build (JFJOCH_USE_CUDA=OFF) runs the whole pipeline on it; the one exception is the ffbidx indexer, which needs a GPU (the CPU build indexes with fftw). How closely the two agree is stated engine by engine below and in the source: where the arithmetic is integer, or the order of every sum is fixed to the host's, the result is the same to the bit; where the GPU uses float distances, cuFFT for FFTW, or a different (fixed) summation order, the two agree to rounding. Every engine is deterministic on its own - no floating-point atomics decide a result - so the same input gives the same output on every run of a GPU build.

Getting the image onto the card

• Bitshuffle+LZ4 decoder (BSLZ4DecoderGPU.cu) - decompresses the stored chunk on the device, so only the compressed bytes cross PCIe, and un-transposes it straight into the preprocessing step. The same pixels as the host decoder. The approach is Jon Wright's (ESRF). §0
• CBF byte-offset decoder (CBFDecoderGPU.cu) - decodes a miniCBF x-CBF_BYTE_OFFSET stream on the device into the same int32 pixels as the host decoder, refusing a stream that ends before the image does. §0
• marCCD frames (ImagePreprocessorGPU.cu) - the stored 16-bit pixels are uploaded as they are and widened to int32 on the card, inside preprocessing. §0
• Preprocessing (ImagePreprocessorGPU.cu, ImagePreprocessorBufferGPU.cu) - applies the pixel mask and flags error and saturated pixels, counting each, and keeps the preprocessed image on the device for every engine below; on request it gathers the pixels of the strong spots, or small squares around given positions, back to the host. Integer, so identical to the host. §1.6

Per image: profiles and spots

• Azimuthal integration (AzIntEngineGPU.cu) - bins every pixel into its resolution (and optionally azimuth) bin with the geometric corrections applied and sums them, as a radial profile. Fixed-point sums, so the profile does not depend on the order the threads add in. §2
• ROI integration (ROIIntegrationGPU.cu) - sums the pixels of user-defined detector regions of interest (box, circle, azimuthal), each pixel feeding any subset of them, on the EIGER/CPU path where the FPGA does not. Integer sums. FPGA: ROI integration
• Local-statistics spot finder (ImageSpotFinderGPU.cu) - flags strong pixels by their signal-to-noise against a local box. §3.1
• Adaptive spot finder (AdaptiveSpotFinderGPU.cu) - reduces the image into resolution rings once, and from that one pass gives both the azimuthal profile and each ring's robust background, which sets the self-calibrating detection threshold; the result is intersected with the local test above. §3.2
• Spot extractor (SpotExtractorGPU.cu) - joins the flagged pixels into connected components and sums each into a spot descriptor on the device, so only the spot list comes back. The same components, in the same order, as the host labelling (the sparse formulation of the ACTS/traccc project). §3.4

Pre-scan accumulators

These read a subset of the frames before processing starts, on every card at once, and are added together across cards; their sums are integers, so the split of the frames between the cards changes nothing.

• Beam-stop projection and mask (ShadowAccumulatorGPU.cu, ShadowMaskGPU.cu) - accumulates the per-pixel maximum, sum and count of the raw counts straight from the compressed frames, then draws the beam-stop shadow mask from them, step for step as the host does. §1.5
• Hot pixels (HotPixelsGPU.cu) - per frame, the ring-sector and ring order statistics, the levels made from them and the per-pixel counts of lit frames, from which persistently hot pixels are masked. Exact order statistics and integer sums: the host's result to the bit. §1.6
• Beam centre from the background (BeamCenterBackgroundGPU.cu) - bins the usable pixels about a trial centre and sums each cell, the two passes of the background-isotropy fit; the fit itself stays on the host. §1.4
• Beam centre by FFT (BeamCenterFFTGPU.cu) - the self-convolutions that score every candidate centre at once, their combination into one score surface, and the shortlist of its peaks, with only the shortlist brought back. Runs on FFTW on the host instead where it would need more than half of the card's total memory, decided before the card is touched and logged. §1.4

Indexing

• FFT indexer (FFTIndexerGPU.cu) - projects the spots' reciprocal-space vectors onto a dense set of directions, histograms each projection and transforms it with cuFFT, 4096 directions at a time, and picks each direction's most prominent lattice peak; the candidate vectors then go to the host's lattice reduction. cuFFT gives the same bits at any batch size, so the batching changes nothing but the memory held. A rotation run gives its indexers' memory back to the card once a pass has done its last indexing. §5.1, §5.2
• Fast-feedback indexer (ffbidx) - third-party, from the fast-feedback-indexer project (PSI); a known-cell indexer for serial stills, used as fetched. §4

Prediction and integration

• Still prediction (BraggPredictionGPU.cu) - the reflections of a lattice within the excitation-error cutoff, with their detector positions and per-reflection corrections. §8.2
• Rotation prediction (BraggPredictionRotGPU.cu) - the reflections crossing the Ewald sphere within an image's rotation range, with their rotation angle and partiality; they are ordered, and under a cap the best-recorded ones selected, on the device, in exactly the order the host gives. §8.3
• Bragg integration (BraggIntegrationEngineGPU.cu) - for every predicted reflection: the three-ring regions and neighbour ownership, box summation, the per-shell profile learned from the strong spots and rebuilt per reflection, the radial background-curvature correction, and the profile fit. The profile fit is in single precision on both engines, in the same arithmetic, so the two differ only in summation order. §9.2, §9.3

Scaling and merging (rotation)

One engine (RotationScaleMergeGPU.cu) keeps every observation resident on the card for the whole scaling run; each per-group sum is a segmented reduction in a fixed order, so it is deterministic. Whether the scaling runs there at all is decided before any observation is on the card, from the card's total memory: a set whose sort alone, or whose scaling (estimated from its observation and distinct-hkl counts once the sort has counted them), would take more than 80% of it scales on the CPU instead, with a warning in the log (RotationScaleMerge: ... scaling on the CPU instead). Judged on the total and not on what is free, the same input on the same card always takes the same path.

• Ingest - the integrated reflections uploaded frame by frame, twice: first only what the sort and the resolution cut need, then, once the scaling is known to fit, the observations the cut kept, written straight to their place (the predicted detector positions stay on the host); the frame flux applied; the sort into raw-hkl runs and the resolution cut; the recomputed rotation angles (delta-phi), the rocking events relaid, and the partiality recomputed from the smoothed mosaicity. The same sort and cut as the host, to the bit. §10.2, §10.6
• Per-frame scaling (ReduceGroupMeans, FitPerFrameG, UpdateCorr, SmoothCorr, the zeta and frame filters, PerFrameCC) - the reference intensity of every symmetry group, the per-frame scale fitted against it, its smoothing over the sweep, and the per-frame correlation behind --min-image-cc. §10.1, §10.3
• Combine - each reflection's partials summed into a full, with its capture fraction and variance; then the fulls' own per-frame scale and its smoothing. §10.6
• Correction surfaces - the decay, absorption and modulation surfaces' terms, their reference and their least-squares sums over the fulls, on held-out subsets. §10.6
• Merge - the inverse-variance merge (I(+)/I(-) and the mean), the half-sets for CC1/2, R_meas, the per-group scatter, the anomalous (Bijvoet) excess and the error-model samples. §10.4, §10.5
• Error model (ErrorModelGPU.cu) - the binned fit of the error model's a and b, the bin medians identical to the host's; a and b can differ from the host fit in their last bits (summation order). §10.5
• Wilson outliers (WilsonOutliersGPU.cu) - the per-observation Wilson-statistics rejection, the host's result to the bit. §13.3
• Pair sort (PairSortGPU.cu) - sorts the (deviation, weight) pairs of the outlier test's weighted median; the same sequence as std::sort. §13.3

Model validation

The model's structure factors and the rigid body run on the card where there is one, every part placed on it or on the CPU up front against one budget, 80% of the card's total memory (§14.8); the GPU and CPU agree to rounding (float distances, cuFFT for FFTW, the crystal's symmetry composed in reciprocal space instead of by symmetrizing the grid).

• Model density (ModelDensityGPU.cu) - the density of one copy of the model's atoms on a grid, as a gather over bricks of the grid, bit-identical from run to run. §14.1
• Bulk-solvent mask (ModelMaskGPU.cu) - the solvent mask from every symmetry image of every atom, with island removal and the shrink, as gemmi computes it. §14.2
• Structure factors and maps (ModelStructureFactorsGPU.cu) - F_calc and F_mask from the two grids by cuFFT, composed over the group's operators and read off at the data's reflections; and the inverse transform of map coefficients into maps. §14.1, §14.4
• Scale fit (ModelScaleGPU.cu) - the sums of the overall-scale, anisotropic-B and bulk-solvent fit over the reflections, in double and a fixed order, for many fits at once (the solvent grid, the null replicates); the Levenberg-Marquardt control stays on the host. Used where its engine fits the validation's budget, the CPU fit otherwise. §14.2
• Rigid body (RigidBodyGPU.cu) - the rigid-body target for a placement: the model moved, its F_calc and mask, the scale, the residuals and their Jacobian; used for the placement itself and for every replicate of the null. §14.8, §14.5

Molecular replacement

• Molecular transform (MolecularReplacementGPU.cu, MapToFPhiGPU) - the Fourier transform of the search model's density in a large box, by cuFFT, where its buffers fit the search's budget (80% of the card, the search holding nothing else yet), kept only as far out as it is read.
• Rotation function (MolecularReplacementGPU.cu) - for every trial orientation, the model's intensities summed over the copies, normalised by resolution bin and correlated with the data's normalised intensities.
• Translation function (MolecularReplacementGPU.cu) - for each orientation, the function over every position in the cell, beside the copies already placed where there are any, and its mean, spread and highest point.

What of the search goes on the card is planned once per merge on the same budget as model validation's, before the search starts (§14.10).

The placements are then refined, scored and checked against their null with the rigid body and the structure factors above; the likelihood scoring runs on the host. The search is described in §14.10, its keys and decision rules in Molecular replacement.

Apple GPUs (Metal)

A macOS build has no CUDA, and one engine runs on the Apple GPU instead:

• Bragg integration (BraggIntegrationEngineMetal.mm, .metal) - the CPU engine with its pixel work moved to the GPU: the neighbour mask, the ownership of shared pixels, the box summation's walk over every reflection, the background clip and the single-precision profile fit. Everything computed in double and the profile learning stay on the CPU, as Apple GPUs have no double precision. The fit's Gaussian uses an exp of our own (BraggExp.h) on both, as no GPU exp matches the CPU's. The same result as the CPU engine to the bit, which a macOS CI step checks on every push; rugnux --no-gpu keeps it off the GPU, and so does --background-trim, whose trimmed mean has no Metal kernel. §9.2