Jungfraujoch/image_analysis/scale_merge/ResolutionCutoff.cpp

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

#include "ResolutionCutoff.h"

#include <algorithm>
#include <array>
#include <cmath>
#include <limits>

#include "../../common/CorrelationCoefficient.h"

namespace {
    // Fine CC1/2 bins for the fit (finer than the 10 reported shells, per the design). A bin needs
    // this many merged reflections for its CC1/2 to be trusted.
    constexpr int   N_FIT_BINS = 25;
    constexpr int   MIN_BIN_COUNT = 10;
    constexpr int   MIN_FIT_BINS = 5;          // need at least this many usable bins for a fit
    constexpr int   MAX_CC_GROUPS = 7;         // a bin's CC1/2 is the median over this many groups (odd)
    constexpr int   EXTEND_BINS_PAST_FALLOFF = 2;   // bins kept beyond the first sub-target bin
    constexpr double SHELLS_FOR_EXTENSION = 10.0;   // "+1 shell" = one 10-shell width in s (report-independent)

    double Logistic(double s, double k, double s0) {
        return 1.0 / (1.0 + std::exp(k * (s - s0)));
    }

    // Weighted (equal-weight) sum of squared residuals of the logistic against the binned CC1/2.
    double FitSSE(const std::vector<double> &s, const std::vector<double> &cc, double k, double s0) {
        double sse = 0.0;
        for (size_t i = 0; i < s.size(); ++i) {
            const double r = cc[i] - Logistic(s[i], k, s0);
            sse += r * r;
        }
        return sse;
    }
}

namespace {
    double FrameScaleOr1(const std::vector<double> &frame_scale, size_t f) {
        return std::isfinite(frame_scale[f]) && frame_scale[f] > 0.0 ? frame_scale[f] : 1.0;
    }
}

double TypicalFrameScale(const std::vector<double> &frame_scale, const std::vector<int64_t> &n_obs) {
    double s3 = 0.0, s2 = 0.0;
    for (size_t f = 0; f < frame_scale.size(); ++f) {
        const double g = FrameScaleOr1(frame_scale, f);
        s3 += n_obs[f] * g * g * g;
        s2 += n_obs[f] * g * g;
    }
    return s2 > 0.0 ? s3 / s2 : 1.0;
}

std::vector<double> CCHalfFrameFactors(const std::vector<double> &frame_scale,
                                       const std::vector<int64_t> &n_obs) {
    const double g_ref = TypicalFrameScale(frame_scale, n_obs);
    std::vector<double> factor(frame_scale.size());
    for (size_t f = 0; f < frame_scale.size(); ++f) {
        const double r = g_ref / FrameScaleOr1(frame_scale, f);
        factor[f] = std::max(1.0, r * r);
    }
    return factor;
}

ResolutionCutoffResult ComputeCCHalfLogisticCutoff(const std::vector<MergedReflection> &merged,
                                                   double cc_target, Logger &logger) {
    ResolutionCutoffResult result;
    if (!(cc_target > 0.0 && cc_target < 1.0)) {
        result.note = "invalid CC target";
        return result;
    }

    // Ice sits out of this fit, and only this fit. A powder ring is reproducible: past the crystal's
    // own limit the two half-sets still agree ABOUT THE ICE, and a Pearson CC1/2 cannot tell that
    // agreement from diffraction - shells have been measured at CC1/2 0.765 where <I/sigma> is -0.1
    // and R_meas is 470%, and it is those shells that hold the cut open. Measured against archived
    // references, a run with ice rings is written 12% finer than one without reaches 4%. The ring
    // also drags the curve the other way where the data are good, dipping one shell of forty to 0.51
    // because two ice lines cross it. Both signs are the same cause and both are removed here. The
    // flag is the one scaling, the error model and the space-group search already keep out of their
    // own fits; what is new is only that the resolution DECISION reads the same curve they do. The
    // reflections are still merged, still written, still counted in the shell table.
    //
    // s = 1/d^2 range over the merged reflections that carry a half-set pair.
    auto usable = [](const MergedReflection &m) {
        return !m.on_ice_ring && m.d > 0.0f && std::isfinite(m.I_half[0]) && std::isfinite(m.I_half[1]);
    };
    double s_lo = std::numeric_limits<double>::max(), s_hi = 0.0;
    double d_data_min = std::numeric_limits<double>::max();
    for (const auto &m : merged) {
        if (!usable(m))
            continue;
        const double s = 1.0 / (static_cast<double>(m.d) * m.d);
        s_lo = std::min(s_lo, s);
        s_hi = std::max(s_hi, s);
        d_data_min = std::min(d_data_min, static_cast<double>(m.d));
    }
    if (!(s_lo < s_hi)) {
        result.note = "no half-set data for CC1/2 fit";
        return result;
    }

    // Bin CC1/2 against s (equal width in s, matching the reporting shells which are equal in 1/d^2).
    const double bin_w = (s_hi - s_lo) / N_FIT_BINS;
    auto bin_of = [&](const MergedReflection &m) {
        const double s = 1.0 / (static_cast<double>(m.d) * m.d);
        return std::clamp(static_cast<int>((s - s_lo) / bin_w), 0, N_FIT_BINS - 1);
    };
    std::vector<int> bin_n(N_FIT_BINS, 0);
    for (const auto &m : merged)
        if (usable(m)) ++bin_n[bin_of(m)];

    // A bin's CC1/2 is the median of the Pearson CC1/2 of a few groups its reflections are dealt into by
    // their indices, not one Pearson over the whole bin. A Pearson CC has no bound on
    // what one pair can do to it: in a weak shell a single reflection with a huge intensity in both
    // halves - an artefact the halves reproduce, or one observation with a sigma fifty times the shell's
    // - carries nearly all of the covariance and reads the shell at CC1/2 0.98 where the rest of it is
    // noise, and the cut runs on through it. Here such a reflection takes one group with it, and the
    // median does not follow. On a shell without one, every group estimates the same CC1/2 and the
    // median is the Pearson value. Groups of fewer than MIN_BIN_COUNT reflections are not formed: a bin
    // too small for three groups is one group, the plain Pearson CC1/2.
    std::vector<int> bin_groups(N_FIT_BINS);
    for (int b = 0; b < N_FIT_BINS; ++b) {
        const int g = std::min(MAX_CC_GROUPS, std::max(1, bin_n[b] / MIN_BIN_COUNT));
        bin_groups[b] = g % 2 == 1 ? g : g - 1;
    }
    std::vector<std::array<CorrelationCoefficient, MAX_CC_GROUPS>> group_cc(N_FIT_BINS);
    for (const auto &m : merged) {
        if (!usable(m)) continue;
        const int b = bin_of(m);
        const uint32_t hash = static_cast<uint32_t>(m.h) * 73856093u ^ static_cast<uint32_t>(m.k) * 19349663u
                              ^ static_cast<uint32_t>(m.l) * 83492791u;
        group_cc[b][hash % bin_groups[b]].Add(m.I_half[0], m.I_half[1], m.cc_weight);
    }

    // Usable bins (enough counts), in ascending-s order, with their bin-centre s.
    std::vector<double> s_bin, cc_bin;
    for (int b = 0; b < N_FIT_BINS; ++b) {
        if (bin_n[b] < MIN_BIN_COUNT) continue;
        std::vector<double> cc;
        for (int g = 0; g < bin_groups[b]; ++g)
            if (std::isfinite(group_cc[b][g].GetCC()))
                cc.push_back(group_cc[b][g].GetCC());
        if (cc.empty()) continue;
        std::nth_element(cc.begin(), cc.begin() + cc.size() / 2, cc.end());
        s_bin.push_back(s_lo + (b + 0.5) * bin_w);
        cc_bin.push_back(cc[cc.size() / 2]);
    }
    if (static_cast<int>(s_bin.size()) < MIN_FIT_BINS) {
        result.note = "too few usable CC1/2 bins";
        return result;
    }

    // Restrict to the contiguous fall-off from low res: keep bins up to a couple past the first one
    // that drops below cc_target, so a high-res noise blip cannot pull the fit back up. If the lowest
    // bin is already below cc_target there is no low-res plateau to anchor on - bail out.
    if (cc_bin.front() < cc_target) {
        result.note = "no low-resolution CC1/2 plateau";
        return result;
    }
    size_t first_below = cc_bin.size();
    size_t keep = cc_bin.size();
    for (size_t i = 0; i < cc_bin.size(); ++i) {
        if (cc_bin[i] < cc_target) {
            first_below = i;
            keep = std::min(cc_bin.size(), i + 1 + EXTEND_BINS_PAST_FALLOFF);
            break;
        }
    }
    s_bin.resize(keep);
    cc_bin.resize(keep);
    if (static_cast<int>(s_bin.size()) < MIN_FIT_BINS) {
        result.note = "too few CC1/2 bins in the fall-off region";
        return result;
    }

    // Fit the logistic by a grid search over (k>0, s0) then a local coordinate-descent refine
    // (dependency-free; the fall-off is smooth and the grid lands close). s0 spans the s range; k
    // spans transitions from very gradual to very sharp relative to that range.
    const double s_range = s_hi - s_lo;
    double best_k = 0.0, best_s0 = 0.0, best_sse = std::numeric_limits<double>::max();
    constexpr int N_S0 = 60, N_K = 40;
    const double k_min = 2.0 / s_range, k_max = 200.0 / s_range;
    for (int ik = 0; ik < N_K; ++ik) {
        const double k = k_min * std::pow(k_max / k_min, static_cast<double>(ik) / (N_K - 1));
        for (int is = 0; is < N_S0; ++is) {
            const double s0 = s_lo + s_range * static_cast<double>(is) / (N_S0 - 1);
            const double sse = FitSSE(s_bin, cc_bin, k, s0);
            if (sse < best_sse) { best_sse = sse; best_k = k; best_s0 = s0; }
        }
    }

    double k = best_k, s0 = best_s0;
    double step_s0 = s_range / N_S0, step_k = best_k * 0.5;
    for (int iter = 0; iter < 200; ++iter) {
        bool improved = false;
        for (const double ds : {step_s0, -step_s0}) {
            const double sse = FitSSE(s_bin, cc_bin, k, s0 + ds);
            if (sse < best_sse) { best_sse = sse; s0 += ds; improved = true; }
        }
        for (const double dk : {step_k, -step_k}) {
            const double kt = k + dk;
            if (kt <= 0.0) continue;
            const double sse = FitSSE(s_bin, cc_bin, kt, s0);
            if (sse < best_sse) { best_sse = sse; k = kt; improved = true; }
        }
        if (!improved) { step_s0 *= 0.5; step_k *= 0.5; }
        if (step_s0 < 1e-6 * s_range && step_k < 1e-6 * best_k) break;
    }

    // s where the fitted CC1/2 crosses cc_target, then "one shell too far". The extension is one
    // reported-shell width, measured over the range that is actually kept and reported (low-res
    // plateau -> the fall-off crossing), NOT the full measured range: when the detector reaches far
    // past where the crystal diffracts (a high-res-configured detector on a low-res crystal), the
    // full range is dominated by high-res noise, so s_range/10 would be a huge over-extension.
    double s_cross = s0 + std::log(1.0 / cc_target - 1.0) / k;

    // The crossing has to be one the bins show, not one the logistic extrapolates to. On a real
    // fall-off it sits between the last bin at the target and the first bin below it, with the two
    // extension bins to spare. A crossing past the last fitted bin means the logistic never followed
    // these data down - it settled on a shallow slope through a fall-off region that is not one -
    // and every resolution read off it is one the bins refute. The bins' own crossing is the honest
    // answer then: where the measured CC1/2 passes the target, between the two bins that straddle it.
    bool logistic = true;
    if (s_cross > s_bin.back()) {
        if (first_below >= cc_bin.size()) {
            result.note = "CC1/2 does not fall off within the measured range";
            return result;
        }
        logistic = false;
        const double f = (cc_bin[first_below - 1] - cc_target)
                         / (cc_bin[first_below - 1] - cc_bin[first_below]);
        s_cross = s_bin[first_below - 1] + f * (s_bin[first_below] - s_bin[first_below - 1]);
    }

    const double delta_s = (s_cross - s_lo) / SHELLS_FOR_EXTENSION;
    const double s_final = s_cross + delta_s;

    // The crossing itself is what the data are judged to reach, and it is reported even when the
    // extension then pushes the cut past the measured edge (or no cut is applied at all): the two are
    // separate questions. Only a crossing inside the measured range is a measurement.
    if (s_cross > s_lo && s_cross < s_hi)
        result.d_fit = 1.0 / std::sqrt(s_cross);

    // No cut if the fall-off is beyond the measured edge (CC1/2 still healthy at the highest s).
    if (s_final >= s_hi) {
        result.note = "CC1/2 does not fall off within the measured range";
        return result;
    }
    // Low-resolution floor: never cut into good low-res data. A fit that puts the cutoff within two
    // shells of the lowest-res data is not a real fall-off - keep the full range and warn.
    if (s_final <= s_lo + 2.0 * delta_s) {
        logger.Warning("Resolution cutoff fit landed at low resolution (degenerate CC1/2 fall-off); "
                       "keeping the full resolution range");
        result.note = "degenerate low-resolution fit";
        return result;
    }

    double d_cut = 1.0 / std::sqrt(s_final);
    d_cut = std::max(d_cut, d_data_min);   // cannot cut beyond the highest-resolution reflection

    // A cut that is not meaningfully coarser than the data edge is a no-op.
    if (d_cut <= d_data_min * 1.001) {
        result.note = "CC1/2 healthy to the detector edge";
        return result;
    }

    result.d_cut = d_cut;
    result.note = logistic ? "CC1/2 logistic fit"
                           : "the logistic fit extrapolates past the bins it was made over, so the "
                             "measured CC1/2";
    return result;
}

std::optional<double> ApplyResolutionCutoff(std::vector<MergedReflection> &merged,
                                            std::optional<double> manual_limit,
                                            ResolutionCutoffMethod method,
                                            double cc_target,
                                            bool for_search,
                                            Logger &logger,
                                            std::optional<double> *fit_limit_out) {
    std::optional<double> effective_d_min = manual_limit;
    if (!effective_d_min && !for_search && method == ResolutionCutoffMethod::CCHalfLogistic) {
        const auto rc = ComputeCCHalfLogisticCutoff(merged, cc_target, logger);
        if (fit_limit_out) *fit_limit_out = rc.d_fit;
        if (rc.d_cut) {
            effective_d_min = rc.d_cut;
            if (rc.d_fit)
                logger.Info("Auto resolution cutoff: {:.2f} A ({} crosses {:.2f} at {:.2f} A, "
                            "written one shell past it; override with --scaling-high-resolution)",
                            *rc.d_cut, rc.note, cc_target, *rc.d_fit);
            else
                logger.Info("Auto resolution cutoff: {:.2f} A ({}; override with --scaling-high-resolution)",
                            *rc.d_cut, rc.note);
        } else {
            logger.Info("Auto resolution cutoff: none ({}); keeping the full resolution range", rc.note);
        }
    }
    if (effective_d_min)
        merged.erase(std::remove_if(merged.begin(), merged.end(),
            [&](const MergedReflection &m) { return std::isfinite(m.d) && m.d < *effective_d_min; }),
            merged.end());
    return effective_d_min;
}