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](CPU_DATA_ANALYSIS.md) 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 `ffbidx` indexer included (an in-tree
CPU port of the library's kernels). 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](CPU_DATA_ANALYSIS_IMAGE.md#0-getting-the-image-onto-the-gpu-device-side-decoding)
- **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](CPU_DATA_ANALYSIS_IMAGE.md#0-getting-the-image-onto-the-gpu-device-side-decoding)
- **marCCD frames** (`ImagePreprocessorGPU.cu`) - the stored 16-bit pixels are uploaded as they are
  and widened to int32 on the card, inside preprocessing.
  [§0](CPU_DATA_ANALYSIS_IMAGE.md#0-getting-the-image-onto-the-gpu-device-side-decoding)
- **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](CPU_DATA_ANALYSIS_IMAGE.md#16-defective-pixels)

## 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](CPU_DATA_ANALYSIS_IMAGE.md#2-azimuthal-integration-radial-profiles)
- **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](FPGA_DATA_ANALYSIS.md#region-of-interest-roi-integration)
- **Local-statistics spot finder** (`ImageSpotFinderGPU.cu`) - flags strong pixels by their
  signal-to-noise against a local box. [§3.1](CPU_DATA_ANALYSIS_IMAGE.md#31-strong-pixel-detection-by-local-statistics)
- **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](CPU_DATA_ANALYSIS_IMAGE.md#32-adaptive-self-calibrating-detection)
- **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](CPU_DATA_ANALYSIS_IMAGE.md#34-connected-component-labeling-ccl)

## 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](CPU_DATA_ANALYSIS_IMAGE.md#15-finding-the-beam-stop)
- **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](CPU_DATA_ANALYSIS_IMAGE.md#16-defective-pixels)
- **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](CPU_DATA_ANALYSIS_IMAGE.md#14-the-beam-centre)
- **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](CPU_DATA_ANALYSIS_IMAGE.md#14-the-beam-centre)

## 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](CPU_DATA_ANALYSIS_INDEXING.md#51-directional-projections-and-histograms),
  [§5.2](CPU_DATA_ANALYSIS_INDEXING.md#52-fft-peak-picking-and-candidate-vectors)
- **Fast-feedback indexer** (`ffbidx`) - third-party, from the fast-feedback-indexer project (PSI); a
  known-cell indexer for serial stills, used as fetched. Its CPU counterpart is an in-tree port with
  the same samples and scores, agreeing to rounding.
  [§4.3](CPU_DATA_ANALYSIS_INDEXING.md#43-fast-feedback-indexing-known-unit-cell)

## 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](CPU_DATA_ANALYSIS_INTEGRATION.md#82-still-prediction-excitation-error-cutoff)
- **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](CPU_DATA_ANALYSIS_INTEGRATION.md#83-rotation-prediction-laue-equation--partiality-model)
- **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](CPU_DATA_ANALYSIS_INTEGRATION.md#92-box-summation-seed-and-fallback),
  [§9.3](CPU_DATA_ANALYSIS_INTEGRATION.md#93-profile-fitted-extraction-default)

## 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](CPU_DATA_ANALYSIS_INTEGRATION.md#102-partiality-models),
  [§10.6](CPU_DATA_ANALYSIS_INTEGRATION.md#106-rotation-datasets-combining-partials-into-fulls-3d-integration)
- **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](CPU_DATA_ANALYSIS_INTEGRATION.md#101-observation-model),
  [§10.3](CPU_DATA_ANALYSIS_INTEGRATION.md#103-smoothing-of-per-frame-scales)
- **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](CPU_DATA_ANALYSIS_INTEGRATION.md#106-rotation-datasets-combining-partials-into-fulls-3d-integration)
- **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](CPU_DATA_ANALYSIS_INTEGRATION.md#106-rotation-datasets-combining-partials-into-fulls-3d-integration)
- **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](CPU_DATA_ANALYSIS_INTEGRATION.md#104-merging-estimator),
  [§10.5](CPU_DATA_ANALYSIS_INTEGRATION.md#105-merging-statistics)
- **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](CPU_DATA_ANALYSIS_INTEGRATION.md#105-merging-statistics)
- **Wilson outliers** (`WilsonOutliersGPU.cu`) - the per-observation Wilson-statistics rejection,
  the host's result to the bit. [§13.3](CPU_DATA_ANALYSIS_DECISIONS.md#133-outlier-rejection)
- **Pair sort** (`PairSortGPU.cu`) - sorts the (deviation, weight) pairs of the outlier test's
  weighted median; the same sequence as `std::sort`.
  [§13.3](CPU_DATA_ANALYSIS_DECISIONS.md#133-outlier-rejection)

## 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](CPU_DATA_ANALYSIS_DECISIONS.md#148-rigid-body-placement-of-the-model)); 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). On a machine with several cards, the rigid-body engines
the first card has no room for go on the others, each against its own budget.

- **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](CPU_DATA_ANALYSIS_DECISIONS.md#141-model-structure-factors)
- **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](CPU_DATA_ANALYSIS_DECISIONS.md#142-bulk-solvent-and-scaling)
- **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](CPU_DATA_ANALYSIS_DECISIONS.md#141-model-structure-factors),
  [§14.4](CPU_DATA_ANALYSIS_DECISIONS.md#144-electron-density-maps)
- **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](CPU_DATA_ANALYSIS_DECISIONS.md#142-bulk-solvent-and-scaling)
- **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](CPU_DATA_ANALYSIS_DECISIONS.md#148-rigid-body-placement-of-the-model),
  [§14.5](CPU_DATA_ANALYSIS_DECISIONS.md#145-does-the-model-fit-the-null-it-is-scored-against)

## 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](CPU_DATA_ANALYSIS_DECISIONS.md#1410-molecular-replacement-placing-a-model-the-rigid-body-cannot)).
On a machine with several cards, every other card holds a copy of the molecular transform where it
fits, the rotation and translation functions are split over the cards by orientation, and the
rigid-body engines are spread over them as in model validation; the result does not depend on how
many cards there are.

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](CPU_DATA_ANALYSIS_DECISIONS.md#1410-molecular-replacement-placing-a-model-the-rigid-body-cannot), its keys and decision rules in
[Molecular replacement](RUGNUX_REPORT.md#molecular-replacement).

## Apple GPUs (Metal)

A macOS build has no CUDA, and three engines run on the Apple GPU instead (none under
`rugnux --no-gpu`; where the GPU fails, the CPU takes over with a warning):

- **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](CPU_DATA_ANALYSIS_INTEGRATION.md#92-box-summation-seed-and-fallback)
- **Model density** (`ModelDensityMetal.mm`, `.metal`) - the density of the model's atoms on a grid,
  for model validation and molecular replacement, as a gather over bricks of the grid as under CUDA;
  the same grid as the CPU's **to the bit**.
  [§14.1](CPU_DATA_ANALYSIS_DECISIONS.md#141-model-structure-factors)
- **Molecular-replacement rotation function** (`MolecularReplacementMetal.mm`, `.metal`) - the
  rotation function's trial orientations scored from a copy of the molecular transform; equal to the
  CPU's to rounding.
  [§14.10](CPU_DATA_ANALYSIS_DECISIONS.md#1410-molecular-replacement-placing-a-model-the-rigid-body-cannot)