Jungfraujoch/image_analysis/bragg_prediction/BraggPrediction.h

// SPDX-FileCopyrightText: 2025 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only

#pragma once

#include <vector>

#include "../../common/CrystalLattice.h"
#include "../../common/DiffractionExperiment.h"
#include "../../common/Reflection.h"

struct BraggPredictionSettings {
    float high_res_A = 1.5;
    float ewald_dist_cutoff = 0.0005;
    // Per-index half-widths of the box the predictor walks: h runs -max_h..+max_h, and so on. One limit
    // per axis rather than one cube, because each index is bounded by its OWN axis (|h| <= a/d_min), so
    // a cube sized for the longest axis walks the short ones far past anything the resolution cut can
    // keep - on a 149/83/226 A cell that is ~16x the candidates a per-axis box generates.
    int max_h = 100;
    int max_k = 100;
    int max_l = 100;
    char centering = 'P';
    float wedge_deg = 0.1f;
    float mosaicity_deg = 0.2f;
    float min_zeta = 0.05;
    float mosaicity_multiplier = 4.0;
    // Relative X-ray bandwidth Δλ/λ expressed as a Gaussian sigma (0 = monochromatic).
    // Stills: the Ewald-shell acceptance is thickened radially per reflection by
    // σ_bw = |recip_z|·bandwidth_sigma (= bλ/2d²), so the 1/d² pink-beam smear no
    // longer clips high-resolution reflections.
    // Rotation: differentiating Bragg's law at fixed d gives an extra rocking width
    // Δθ = bandwidth_sigma·tan(θ_B), a spread in the same glancing angle the mosaic spread
    // smears, so it adds to σ_M in quadrature. It is NOT divided by ζ: rotating the crystal
    // by Δφ changes θ by ζ·Δφ, so the 1/ζ that turns an angular width into a rotation width is
    // already the one the partiality applies to σ_M. CalcMosaicityXDS deconvolves the same term
    // out of the fitted σ_M, so it is not counted twice.
    float bandwidth_sigma = 0.0f;
};

class BraggPrediction {
protected:
    // Not const: the buffer grows to fit a frame that predicts more than it holds.
    int max_reflections;
    std::vector<Reflection> reflections;

    // Make room for `count` reflections. Overridden where device buffers have to follow. Called only
    // when a frame predicted more than the current capacity, so a run pays for it a handful of times.
    // Every Calc grows rather than stops: the CPU ones double the buffer when the walk fills it, the GPU
    // ones re-run the kernel at the counted total. A CPU Calc that stopped at the capacity instead kept
    // whichever reflections the h/k/l walk reached first, not the best-recorded ones, so a large cell
    // merged a different reflection set on a CPU-only build than on a GPU one. The one cap on what
    // leaves Calc is output_limit, applied by TruncateToOutput on both.
    virtual void GrowCapacity(int count);

    // Deterministically cap Calc's output at output_limit: if more were predicted, keep the best-recorded
    // ones - largest partiality on the rotation path, and, since partiality is 1 for every still, smallest
    // excitation error on the still path - with hkl breaking what is left. Returns the kept count. Below
    // the cap it is a no-op. Call at the end of every Calc override.
    int TruncateToOutput(int count);

    // Put the first `count` predicted reflections in an order that depends only on the reflections
    // themselves. The GPU predictors append at an atomic counter, so a reflection's POSITION in the
    // array is decided by the order the blocks happened to finish - and that position is not private
    // to the predictor: BraggOwnerKey packs it into the owner map as the tie-break between two centres
    // equidistant from a shared pixel, and every downstream sort that is not a total order (the
    // ingest and post-refine bucket sorts) resolves its ties by the order it receives. Two runs of the
    // same binary on the same image therefore integrated a different set of reflections. hkl is a
    // property of the reflection; delta_phi separates the two rocking solutions one hkl can have.
    // Call before TruncateToOutput, whose own pick is then reproducible as well.
    void OrderOutput(int count);

    // Scratch for OrderOutput. A Reflection is ~88 bytes and a large cell predicts tens of thousands
    // of them per frame, so sorting the structs themselves moves several megabytes an image; sorting
    // a 20-byte key and gathering once is the same order for a third of the traffic. Members rather
    // than locals so the two allocations happen once per engine, not once per image.
    struct OrderKey { int32_t h, k, l; float delta_phi_deg; int32_t index; };
    std::vector<OrderKey> order_keys;
    std::vector<Reflection> order_scratch;
public:
    // Calc returns at most output_limit, the number that flows downstream and is serialized - kept low so
    // the per-image reflection list stays within the frame transport headroom.
    // kPredictionCapacity is the prediction buffer's starting size. The buffer grows to whatever a frame
    // actually predicts (GrowCapacity), so a large cell is not truncated. It used to be a hard cap, and
    // overflowing it was lossy - and on the GPU NON-DETERMINISTIC, since the kernels claim slots with an
    // atomicAdd, so which reflections survived depended on block scheduling and changed between runs.
    static constexpr int kPredictionCapacity = 20000;
    // How many reflections may flow downstream per image, offline. Sized for a large unit cell: a
    // ~2.8e6 A^3 cell predicts up to ~44000 per frame at 2.4 A. Truncating below what the frame really
    // has costs more than it saves - the selection keeps the best-recorded reflections, and the rotation
    // combine rebuilds a full FROM the partials it drops (measured on such a crystal: CC1/2 98 -> 60,
    // ISa 8.6 -> 1.8). DiffractionExperiment's image-buffer headroom is derived from this, so the
    // transport can carry what the analysis produces.
    static constexpr int kPredictionOutput = 65536;
    // What the ONLINE path may carry per image. The acquisition transports every reflection list through
    // a fixed-size image-buffer slot, and the slot size divides a fixed total - so sizing the slot for
    // kPredictionOutput would cut the number of slots, and with it the receiver's ability to absorb a
    // burst, by about three. Online keeps the transport-sized limit it has always had; offline, which
    // has no such budget, keeps the full one. DiffractionExperiment's buffer headroom derives from THIS.
    static constexpr int kOnlineMaxReflections = 10000;

    // How many reflections Calc may return. A caller with a tighter cap than the offline one - online,
    // whose transport slot holds kOnlineMaxReflections - sets its own here, so the surplus is dropped
    // before it is integrated rather than integrated and then thrown away.
    int output_limit = kPredictionOutput;

    explicit BraggPrediction(int max_reflections = kPredictionCapacity);

    virtual ~BraggPrediction() = default;
    virtual int Calc(const DiffractionExperiment &experiment, const CrystalLattice &lattice,
        const BraggPredictionSettings &settings);
    const std::vector<Reflection> &GetReflections() const;
};