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 <map>
#include <random>
#include <unordered_map>
#include <vector>

#include "../../common/Logger.h"
#include "../../common/DiffractionExperiment.h"
#include "../../common/Reflection.h"
#include "../IntegrationOutcome.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;
};

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


std::ostream &operator<<(std::ostream &output, const MergeStatisticsShell &in);
std::ostream &operator<<(std::ostream &output, const MergeStatistics &in);

struct MergeAccum {
    int32_t h = 0;
    int32_t k = 0;
    int32_t l = 0;
    float d = NAN;
    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};
};

class MergeOnTheFly {
    mutable std::mutex merged_mutex;

    const int space_group_number = 1;

    ScalingSettings scaling_settings;
    IndexingSettings indexing_settings;

    std::optional<UnitCell> reference_cell;
    std::optional<double> high_resolution_limit;
    std::optional<double> image_cc_limit;
    double min_partiality = 0.02;

    HKLKeyGenerator generator;

    // To select images for half-datasets to calculate CC1/2, I use a random number generator with a fixed seed.
    // This makes sure that images are selected randomly, but in a fully reproducible manner (at least for the same binary)
    std::mt19937 rng{123456789u};
    std::bernoulli_distribution half_dist{0.5};

    std::map<uint64_t, MergeAccum> accumulator;

    // 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;
    // 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, 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;
    size_t reject_count = 0;

    bool Mask(const IntegrationOutcome &outcome, bool cc_mask);
public:
    MergeOnTheFly(const DiffractionExperiment &x);
    MergeOnTheFly& ReferenceCell(const std::optional<UnitCell> &cell);

    // 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; }

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

    void AddImage(const IntegrationOutcome& outcome, bool cc_mask = false);

    // Per-crystal CC1/2-delta rejection (CrystFEL deltaCChalf): returns a per-image flag
    // marking images whose removal would raise CC1/2 by a low-side outlier amount
    // (deltaCChalf < mean - nsigma*stddev). Skip the flagged images when merging.
    [[nodiscard]] std::vector<char> DeltaCChalfReject(const std::vector<IntegrationOutcome> &outcomes,
                                                      double nsigma) const;

    MergeStatistics MergeStats(const std::vector<MergedReflection> &merged,
                               const std::vector<IntegrationOutcome> &reflections,
                               const std::vector<MergedReflection> &reference = {});

    std::vector<MergedReflection> ExportReflections();
};

std::vector<MergedReflection> MergeAll(const DiffractionExperiment &x,
                                       const std::vector<IntegrationOutcome> &reflections,
                                       bool mask = false);