Jungfraujoch/image_analysis/scale_merge/Merge.h

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

#pragma once

#include <algorithm>
#include <cmath>
#include <map>
#include <span>
#include <unordered_map>
#include <vector>

#include "../../common/Logger.h"
#include "../../common/DiffractionExperiment.h"
#include "../../common/Reflection.h"
#include "../IntegrationOutcome.h"

#include "AnisotropyAnalysis.h"
#include "HKLKey.h"

struct MergeStatisticsShell {
    float d_min = 0.0f;
    float d_max = 0.0f;
    float mean_one_over_d2 = 0;

    int total_observations = 0;
    int unique_reflections = 0;
    int possible_unique_reflections = 0;

    double mean_i_over_sigma = 0.0;

    double cc_half = 0.0f;
    double cc_ref = NAN;

    // Redundancy-independent merging R-factor (Diederichs & Karplus 1997), computed over the
    // observations that enter the merge: R_meas = sum_hkl sqrt(n/(n-1)) sum_i|I_i-<I>| / sum I_i.
    double r_meas = NAN;
    // The same with each observation weighted as the rotation merge weights it, 1/sigma^2 under the
    // error model (WeightedRmeasTerms); with equal sigmas it is r_meas. A frame kept at a small weight
    // then no longer sets the number with its noise. NaN where it is not computed (stills, per hand).
    double r_meas_weighted = NAN;
    // The other two merging R-factors, over the same observations and means as r_meas and differing
    // from it only in the per-reflection factor on sum_i|I_i-<I>|: 1 for R_merge, sqrt(1/(n-1)) for
    // the precision-indicating R_pim (Weiss (2001) J. Appl. Cryst. 34, 130-135). Reported for
    // deposition, which asks for all three.
    double r_merge = NAN;
    double r_pim = NAN;

    // Anomalous signal-to-noise (XDS "SigAno" / mmCIF pdbx_absDiff_over_sigma_anomalous):
    // <|I(+)-I(-)|> / <sigma(I(+)-I(-))> over acentric reflections measured in both hands. NaN when
    // there is no anomalous split (e.g. Friedel-merged with no mates, or the stills path).
    double abs_diff_over_sigma_anomalous = NAN;

    // CCanom (AIMLESS "CCanom", phenix "cc_anom"): the correlation of the anomalous difference
    // measured twice - I(+)-I(-) from one half of the observations against the same difference from
    // the other half, over acentric reflections whose BOTH mates split into two non-empty halves.
    // Unlike SigAno it is not a ratio against the error model, so a mis-estimated sigma cannot move
    // it. NaN when no reflection has an anomalous split (the stills path), which is not the same
    // statement as zero - zero means measured and absent.
    double cc_anom = NAN;

    // The serial-crystallography pair (stills only, NaN on rotation): every frame's observations go
    // whole to one half-set, as CrystFEL's partialator deals out crystals, and the two half-set means
    // give CC1/2 and R_split = (1/sqrt 2) sum|I_1 - I_2| / (1/2 sum (I_1 + I_2)) (White et al. (2012)
    // J. Appl. Cryst. 45, 335-341), over the reflections both halves measured (n_split_pairs). This
    // CC1/2 is not cc_half above, which deals each reflection's observations out evenly and is the
    // one the processing decides on.
    double r_split = NAN;
    double cc_half_split = NAN;
    int n_split_pairs = 0;
};

// Why a stretch of the sweep came out much weaker than the rest of the run. The codes are the field's
// own words - "crystal rotating out of the beam" (HKL-2000 manual), "loss of centring during crystal
// rotation" (autoPROC). What is DONE about a range is its disposition, below; the reason only names
// what was seen. Appended to, never renumbered: the value reaches the per-image array in the HDF5.
enum class SweepQualityReason {
    NoDiffraction,      // the range recorded essentially no diffraction from the indexed lattice
    CrystalOutOfBeam,   // frames were lost: the range gets a per-image scale far less often than the run
    WeakDiffraction,    // the frames all still index, but with much less intensity - cause not determined
    LossOfCentring,     // one cycle of modulation per revolution: the crystal is off the rotation axis
    RadiationDamage,    // the range runs to the end of a sweep whose quality was already decaying
    InconsistentWithMerge  // the frames diffract as the run does, but their intensities do not agree with it
};

const char *SweepQualityReasonCode(SweepQualityReason reason);   // machine-readable, e.g. "out_of_beam"
const char *SweepQualityReasonText(SweepQualityReason reason);   // for a sentence, e.g. "out of beam"

// What became of a frame's observations. The three codes partition the sweep - every processed image is
// exactly one of them - so the percentages add to 100 and "5% of the frames were useless" has a meaning.
enum class FrameDisposition {
    Merged,      // its observations are in the merged data at their own weight
    Downgraded,  // merged, but over a degraded stretch: carried at the reduced weight its own scale and
                 // sigmas give it, which for weak-but-consistent data is the only honest weight there is
    Rejected     // nothing of it reached the merge
};

const char *FrameDispositionCode(FrameDisposition disposition);   // "merged", "downgraded", "rejected"

// One reflection's share of R_meas when its observations carry weights v_i: given sum v|I_i - <I>|,
// sum v I_i, sum v and sum v^2, the numerator sqrt(n_eff/(n_eff-1)) sum v'|I_i - <I>| and the
// denominator sum v' I_i, with the weights rescaled to v' = v n_eff / sum v so they add up to Kish's
// effective count n_eff = (sum v)^2 / sum v^2. Equal weights give the ordinary R_meas terms. False when
// the reflection has no second effective observation (n_eff <= 1), which then counts in neither sum.
bool WeightedRmeasTerms(double sum_v_absdev, double sum_v_I, double sum_v, double sum_v2,
                        double &num, double &den);

struct SweepQualityRange {
    // Inclusive, in processed-image ordinals - the numbering of <prefix>_plot.txt and of every other
    // per-image array rugnux writes. With -s/--stride the source image is start + ordinal * stride.
    int first_image = 0;
    int last_image = 0;
    SweepQualityReason reason = SweepQualityReason::WeakDiffraction;
    // Fraction of the run's typical diffracting power missing over the range: 0 = as good as the run,
    // 1 = nothing at all. The rest are supporting numbers, relative to the run unless stated.
    float severity = 0.0f;
    float rotation_deg = 0.0f;         // width of the range
    // Gauged on the run's precision-weighted typical frame (TypicalFrameScale), not on its median
    // frame, which on a sweep that spent most of its turn out of the beam is itself a dead one.
    float mean_relative_scale = 1.0f;  // <per-image scale> / typical frame scale
    float mean_relative_cc = 1.0f;     // <per-image CC to merge> / run median
    float indexed_fraction = 1.0f;     // frames in the range that got a per-image scale at all
    float relative_b = NAN;            // mean of the radiation-damage monitor's per-batch B over the range
    // What was done with the range, and what the merged data say about keeping it: delta-CC1/2 is the
    // overall CC1/2 with the range minus the CC1/2 without it, over the reflections the range touches
    // (negative = keeping it makes the merged intensities worse), beside the standard error of a CC1/2
    // on that many reflections. NaN where the range carries too few reflections to be measured.
    FrameDisposition disposition = FrameDisposition::Downgraded;
    float delta_cc_half = NAN;
    float delta_cc_half_se = NAN;
};

// Sweep-quality diagnostic (rotation only). `measured` separates "the run is clean" from "this did not
// run": the range list is empty in the first case and in the second alike.
struct SweepQuality {
    bool measured = false;
    float sweep_deg = 0.0f;
    float flux_peak_to_trough = 1.0f;        // the incident-flux proxy, over the whole run
    float modulation_peak_to_trough = 1.0f;  // depth of a DIAGNOSED once-per-revolution modulation (1 = none)
    std::vector<SweepQualityRange> ranges;

    // The disposition of the whole sweep. The three counts add up to the frame count; `rejected_deg` is
    // the same rejection in degrees, because a percentage of FRAMES can be moved by re-slicing the same
    // experiment and a percentage of the rotation cannot.
    int frames_merged = 0;
    int frames_downgraded = 0;
    int frames_rejected = 0;
    float rejected_deg = 0.0f;
    std::vector<uint8_t> frame_disposition;   // FrameDisposition per processed image

    // The per-batch delta-CC1/2 curve the disposition was decided on, measured on the run as it stood
    // before anything was removed (NaN where a batch carries too few reflections to be judged). Same
    // batches as the radiation-damage monitor's relative-B curve, so the two read side by side.
    std::vector<float> delta_cc_half_batch;
    float delta_cc_half_batch_deg = 0.0f;
};

// The same merge counted with Friedel's law OFF: an acentric reflection is two rows, one per Bijvoet
// hand, and a centric is one. XDS's FRIEDEL'S_LAW=FALSE table, and what an anomalous CORRECT.LP is to
// be read against. Formed from the same per-hand accumulators as SigAno and the I(+)/I(-) export, so
// it is a second PRESENTATION of one merge and not a second merge: the intensities, the scaling, the
// error model and every decision behind it are the Friedel-merged ones. `measured` is false when the
// merge carried no anomalous split to count (the stills path, or a pass that writes nothing).
struct AnomalousMergeStatistics {
    std::vector<MergeStatisticsShell> shells;
    MergeStatisticsShell overall;
    bool measured = false;
};

// An observation the Wilson outlier test removed from the merge (see WilsonOutliers.h): its own
// indices, its normalised intensity, and where it was recorded, so a recurring detector position shows.
struct WilsonRejectedObservation {
    int32_t h = 0, k = 0, l = 0;
    float d = NAN;
    float e2 = NAN;
    float image = NAN;       // fractional image number
    float x = NAN, y = NAN;  // predicted detector position (pixels)
};

struct MergeStatistics {
    std::vector<MergeStatisticsShell> shells;
    MergeStatisticsShell overall;
    AnomalousMergeStatistics per_hand;

    // Dataset-wide isotropic Wilson B-factor estimate (A^2) from the log-linear fit of the shell-mean
    // merged intensity against 1/d^2 (CalcGlobalWilsonB) - the analogue of XDS's "WILSON LINE ... B=".
    // Diagnostic only; not used in scaling. NaN when not determined.
    double wilson_b = NAN;
    double wilson_b_correlation = NAN;
    double wilson_b_d_low_A = NAN;     // the window the fit was made over
    double wilson_b_d_high_A = NAN;

    // Radiation-damage monitor (rotation only): the relative Debye-Waller B change measured from the first
    // to the last frame of the run (A^2; positive = high-resolution intensity fades with dose = damage) and
    // the per-batch relative-B curve it was derived from. Measured before any decay/relative-B correction is
    // applied, so it reports how much radiation damage was present. Diagnostic; NaN / empty for stills or
    // when not determined. batch_deg is the rotation width per batch of the curve. A batch the data cannot
    // measure is NaN in the curve, and delta_b is NaN when the curve is not a trend a single number
    // summarises - damage is progressive, so a curve that is not is telling the user about something else.
    // Observations the outlier rejection dropped before merging. Rejected observations are excluded
    // from the merge AND from R_meas and the CC(1/2) half-sets, which is the XDS convention and is
    // right - those statistics describe the data as merged. But it means an over-rejecting run looks
    // BETTER by every number it reports, so the count has to be visible or the failure is silent.
    // Zero when rejection is off.
    size_t n_observations_rejected = 0;
    // Rocking events (rotation) left out of the merge because a pixel of their spot was saturated
    // somewhere along the rocking curve - the measurement lacks its brightest part (XDS's OVERLOAD).
    // Not part of n_observations_rejected, which counts what the outlier tests took.
    size_t n_reflections_rejected_overloaded = 0;
    // The part of those the Wilson test removed (see WilsonOutliers.h), each listed, in the written range.
    std::vector<WilsonRejectedObservation> wilson_rejected;
    double radiation_damage_delta_b = NAN;
    std::vector<float> radiation_damage_b_batch;
    double radiation_damage_batch_deg = 0.0;

    // Stretches of the sweep over which the crystal delivered much less than the rest of the run
    // (MeasureSweepQuality), and what became of every frame of the sweep (MeasureBatchDeltaCCHalf).
    SweepQuality sweep_quality;

    // Diffraction anisotropy (AnalyzeAnisotropy): the anisotropy tensor, the diffraction limit along
    // each of its principal directions, and whether either is established above this dataset's own
    // systematic error. Report-only - nothing is corrected and no reflection is removed. Empty
    // (n_reflections = 0) when the diagnostic did not run.
    AnisotropyResult anisotropy;
};


// CC* = sqrt(2 CC1/2 / (1 + CC1/2)), the correlation of the merged data with the noise-free signal
// (Karplus & Diederichs (2012) Science 336, 1030-1033). NaN where CC1/2 is not positive.
inline double CCStar(double cc_half) {
    return cc_half > 0.0 ? std::sqrt(2.0 * cc_half / (1.0 + cc_half)) : NAN;
}

std::ostream &operator<<(std::ostream &output, const MergeStatisticsShell &in);
std::ostream &operator<<(std::ostream &output, const MergeStatistics &in);
// The header and the rule of the table operator<< prints, for a table with the same columns rendered
// elsewhere (the report's reference-range table).
std::string ShellTableHeader();
std::string ShellTableRule();
// The serial table of R_split and the frame-split CC1/2, per shell and overall; empty where they were
// not computed (rotation).
std::string SerialShellTable(const MergeStatistics &in);

struct MergeAccum {
    int32_t h = 0;
    int32_t k = 0;
    int32_t l = 0;
    float d = NAN;
    bool on_ice_ring = false;
    double sum_wI = 0.0;
    double sum_w = 0.0;
    double sum_wI_half[2] = {0.0, 0.0};
    double sum_w_half[2] = {0.0, 0.0};
    size_t n_half[2] = {0, 0};
    // The same sums with each frame's observations kept whole in the frame's half (R_split).
    double sum_wI_split[2] = {0.0, 0.0};
    double sum_w_split[2] = {0.0, 0.0};
};

// XDS's error-model convention (Diederichs, Acta Cryst. D66 (2010) 733) is sigma^2 = a*(sigma0^2 +
// b*I^2), reported with ISa = 1/sqrt(a*b). Jungfraujoch fits sigma^2 = a*sigma0^2 + (b*<I>)^2 - the
// same `a`, but a `b` that is a FRACTION of the intensity - so b_xds = b^2/a, and the two ISa
// expressions are the same number: 1/sqrt(a * b^2/a) = 1/b. Only `b` needs converting, and only
// where it is reported: doing it here rather than in the fit leaves every merge weight untouched.
//
// NOTE this is NOT the `b` of SearchSpaceGroup's merge_systematic_b, which is a third, unrelated
// quantity (a fraction of I, fitted with no `a` at all) whose gate constants are calibrated in that
// convention. Do not "make them consistent".
struct XdsErrorModel {
    double a = 1.0;
    double b = 0.0;
    double isa = 0.0;
};

inline XdsErrorModel ToXdsErrorModel(double a, double b) {
    if (!(a > 0.0) || !(b > 0.0))
        return {a, 0.0, 0.0};
    return {a, b * b / a, 1.0 / b};
}

class MergeOnTheFly {
    mutable std::mutex merged_mutex;

    const gemmi::SpaceGroup space_group;

    ScalingSettings scaling_settings;
    IndexingSettings indexing_settings;

    std::optional<UnitCell> reference_cell;
    std::optional<double> high_resolution_limit;
    std::optional<double> low_resolution_limit;
    std::optional<double> image_cc_limit;
    // Apply image_cc_limit in Mask(). One flag for the whole engine, not a per-call argument, so the
    // merge, the error model and MergeStats can never disagree about which images are in.
    bool filter_by_image_cc = false;
    double min_partiality = 0.02;
    // The share of a partial the model could not see, (1 - p), carried in its variance as
    // (capture_coeff * (1 - p) * <I>)^2 - the rotation combine's capture term, per observation.
    double capture_coeff = 0.0;

    // When set, ice-ring-flagged reflections are left out of this merge. Used for the P1 pass whose
    // merged intensities feed the space-group search and the error model - those model fits must not
    // see the ice-contaminated intensities. The final in-symmetry merge keeps them (for completeness).
    bool exclude_ice_rings = false;
    size_t nthreads = 1;

    HKLKeyGenerator generator;

    std::unordered_map<uint64_t, MergeAccum> accumulator;   // exported in key order

    // Global error model (XDS form): sigma_corr^2 = a*sigma^2 + (b*<I>)^2. a rescales the
    // (under-estimated) counting variance; the (b*<I>)^2 term adds the intensity-
    // proportional systematic error that counting statistics miss, so strong reflections
    // are no longer over-weighted. ISa = 1/b is the asymptotic I/sigma. Refined from the
    // scatter of symmetry equivalents (RefineErrorModel); identity until then.
    bool   error_model_active = false;
    double error_model_a = 1.0;
    double error_model_b = 0.0;
    double error_model_chi2 = 0.0;   // achieved median reduced chi^2 (~1.0 = honestly calibrated sigmas)
    // The (b*I)^2 term uses the reflection's *mean* intensity (constant over its
    // observations), so it inflates sigma without biasing the inverse-variance weights -
    // using the per-observation I_i instead would over-weight down-fluctuated points.
    std::unordered_map<uint64_t, float> error_model_mean_I;
    [[nodiscard]] float CorrectedSigma(float I_corr, float sigma_corr, float image_scale_corr,
                                       float var_bkg, float partiality,
                                       uint64_t hkl_key) const;

    // Optional per-observation outlier rejection: drop observations whose corrected
    // intensity lies more than reject_nsigma error-model sigmas from the reflection's
    // *median* (a robust centre). The error-model sigma already captures the genuine
    // (e.g. partiality) scatter, so this removes only the tail beyond it - zingers,
    // overlaps, mis-indexed frames - not good partials. Populated by RefineErrorModel.
    bool   reject_outliers = false;
    double reject_nsigma = 6.0;
    std::unordered_map<uint64_t, float> reject_median_I;
    // The asymmetric band of OutlierBand.h: exp(+-n s) - 1 about the reflection mean, s the spread of
    // ln(I / median) measured on this merge's observations.
    double reject_up = 0.0, reject_down = 0.0;
    size_t reject_count = 0;

    bool Mask(const IntegrationOutcome &outcome) const;

    // One accepted observation, reduced to what the accumulator adds for it.
    struct Contribution {
        uint64_t key;
        int32_t h, k, l;   // the reflection's own index, as the accumulator records it
        float w, wI, d;
        bool on_ice_ring;
    };
    // The observations of one image that reach the merge; returns how many the outlier test dropped.
    size_t Contributions(const IntegrationOutcome &outcome, std::vector<Contribution> &out) const;
    void Accumulate(const std::vector<Contribution> &contributions, int64_t image_id);
public:
    MergeOnTheFly(const DiffractionExperiment &x);
    MergeOnTheFly& ReferenceCell(const std::optional<UnitCell> &cell);
    MergeOnTheFly& ExcludeIceRings(bool input) { exclude_ice_rings = input; return *this; }
    MergeOnTheFly& FilterByImageCC(bool input) { filter_by_image_cc = input; return *this; }
    // Workers for the passes over every observation (AddImages, RefineErrorModel, MergeStats); the
    // result does not depend on the number.
    MergeOnTheFly& Threads(size_t input) { nthreads = std::max<size_t>(1, input); return *this; }

    // Fit the global error model from the spread of symmetry-equivalent observations.
    // Call once before merging; AddImage then applies it.
    void RefineErrorModel(const std::vector<IntegrationOutcome> &outcomes);
    [[nodiscard]] bool   ErrorModelActive() const { return error_model_active; }
    [[nodiscard]] double ErrorModelA() const { return error_model_a; }
    [[nodiscard]] double ErrorModelB() const { return error_model_b; }
    [[nodiscard]] double ErrorModelChi2() const { return error_model_chi2; }

    // Outlier rejection (driven by ScalingSettings::GetOutlierRejectNsigma) reports its count.
    [[nodiscard]] size_t RejectedCount() const { return reject_count; }

    // image_id is the image's stable identity (its index in the outcomes vector). It seeds the CC1/2
    // half-set PREFERENCE, which the balancing rule in AddImage then overrides where a reflection's
    // two halves would come out uneven - so the half an observation lands in is a function of the
    // order AddImage is called in. CALL IT SERIALLY, in image order: the mutex makes concurrent calls
    // safe, not reproducible. The rotation path ranks instead (RotationScaleMerge::AssignHalvesByRank)
    // because it holds every observation at once; a streaming accumulator cannot, and there is no
    // second, device-side walk of the stills observations for this to disagree with.
    void AddImage(const IntegrationOutcome& outcome, int64_t image_id);
    // AddImage for every outcome, the image id being its index: the same merge, bit for bit, with the
    // per-observation work spread over the workers and the accumulation kept in image order.
    void AddImages(const std::vector<IntegrationOutcome> &outcomes);

    // d_min_override, when set, is the effective high-resolution limit for the shell table (used for
    // the automatic resolution cutoff computed by the caller); otherwise the manual
    // ScalingSettings high-resolution limit stands. The number of shells is ScalingSettings::ReportShellCount.
    // declared_range: bin over exactly this range instead - its bounds are the shell grid's and the
    // completeness denominator's whether or not any reflection reaches them, and the limits above
    // are ignored. The reference-range table of --report-resolution.
    MergeStatistics MergeStats(const std::vector<MergedReflection> &merged,
                               const std::vector<IntegrationOutcome> &reflections,
                               const std::vector<MergedReflection> &reference = {},
                               std::optional<double> d_min_override = std::nullopt,
                               std::optional<ReportResolutionRange> declared_range = std::nullopt);

    std::vector<MergedReflection> ExportReflections();
    // The merged intensities alone - no free-R flags, no French-Wilson amplitudes - which is all a
    // scaling reference reads.
    std::vector<MergedReflection> ExportIntensities();
};

// The merged intensities of the images whose cell is close to reference_cell (all of them without one):
// the reference a per-image scale is fitted against.
std::vector<MergedReflection> MergeAll(const DiffractionExperiment &x,
                                       const std::vector<IntegrationOutcome> &reflections,
                                       size_t nthreads,
                                       const std::optional<UnitCell> &reference_cell);

// Pearson CC between one image's corrected intensities (I * image_scale_corr) and a reference set of
// full intensities, over the reflections that would enter the merge (non-ice, within the resolution
// limit, partiality above the floor, finite). {NAN, n} when fewer than 20 reflections qualify.
// This is the per-image image_scale_cc: ScaleOnTheFly sets it, and StillsPartialityRefine recomputes it
// after refining the partiality model, so the reported CC always describes the corrections that will be
// merged - which matters because --min-image-cc drops images by it.
std::pair<double, size_t> ImageReferenceCC(std::span<const Reflection> reflections,
                                           const std::map<HKLKey, double> &reference,
                                           const HKLKeyGenerator &generator,
                                           std::optional<double> d_min_limit,
                                           std::optional<double> d_max_limit,
                                           double min_partiality);