Jungfraujoch/image_analysis/scale_merge/FrenchWilson.cpp

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

#include "FrenchWilson.h"

#include <algorithm>
#include <cmath>
#include <future>
#include <limits>
#include <vector>

#include "../../common/ResolutionShells.h"
#include "../../common/ThreadAffinity.h"
#include "gemmi/math.hpp"
#include "gemmi/symmetry.hpp"

namespace {

struct Posterior {
    double mean_I;   // <J>   (posterior mean true intensity)
    double mean_F;   // <|F|> (posterior mean amplitude)
};

// Posterior moments of the true intensity J >= 0 given a measurement I +/- sigma and the Wilson
// prior with mean sigma_wilson. Integrated numerically over J in [0, I + 8 sigma] with a log-shift
// so the exponentials never overflow/underflow. acentric: p(J) ~ exp(-J/S); centric:
// p(J) ~ exp(-J/2S)/sqrt(J).
// `logw` is caller-owned scratch of npts doubles (one per worker), so the integration allocates nothing.
// The quadrature grid is j = (i + 1/2) * dj, so sqrt(j) and log(j) separate into a per-reflection
// factor and a term that depends only on i: sqrt(j) = sqrt(dj) * sqrt(i + 1/2) and
// log(j) = log(dj) + log(i + 1/2). The i-dependent halves are the same for every reflection, so
// they are tabulated once instead of being recomputed npts times per intensity.
struct QuadratureGrid {
    std::vector<double> sqrt_i, log_i;
    explicit QuadratureGrid(int npts) : sqrt_i(npts), log_i(npts) {
        for (int i = 0; i < npts; ++i) {
            const double x = i + 0.5;
            sqrt_i[i] = std::sqrt(x);
            log_i[i] = std::log(x);
        }
    }
};

// Following French & Wilson (1978) Acta Cryst. A34, 517-525
Posterior integrate_posterior(double I, double sigma, double sigma_wilson, bool centric, int npts,
                              std::vector<double> &logw, const QuadratureGrid &grid) {
    const double inv_2s2 = 1.0 / (2.0 * sigma * sigma);
    // The posterior is the Gaussian likelihood tilted by the exponential prior, so it peaks at
    // I - sigma^2/S and decays over whichever of sigma and S is TIGHTER. Ranging to I + 8 sigma
    // regardless is wrong once sigma greatly exceeds S: with npts fixed the whole prior then falls
    // inside the first grid cell, the quadrature degenerates to that one point and returns
    // F = sqrt(dj/2) with sigmaF -> 0 - i.e. a reflection we know nothing about comes back looking
    // like the best measured one in the file.
    const double prior_scale = centric ? 2.0 * sigma_wilson : sigma_wilson;
    const double peak = std::max(I - sigma * sigma / prior_scale, 0.0);
    const double width = peak > 0.0 ? sigma : std::min(sigma, prior_scale);
    const double j_max = peak + 10.0 * width;
    const double dj = j_max / npts;

    const double log_dj = centric ? std::log(dj) : 0.0;
    double max_logw = -std::numeric_limits<double>::infinity();
    for (int i = 0; i < npts; ++i) {
        const double j = (i + 0.5) * dj;
        const double diff = I - j;
        const double log_prior = centric
                                     ? (-j / (2.0 * sigma_wilson) - 0.5 * (log_dj + grid.log_i[i]))
                                     : (-j / sigma_wilson);
        logw[i] = log_prior - diff * diff * inv_2s2;
        max_logw = std::max(max_logw, logw[i]);
    }

    // exp(-37) is 8e-17: a point that far below the peak cannot change a normalised sum of
    // doubles, and the posterior is sharply peaked, so most of the grid is skipped outright.
    constexpr double LOG_NEGLIGIBLE = -37.0;
    const double sqrt_dj = std::sqrt(dj);
    double sum_w = 0, sum_wI = 0, sum_wF = 0;
    for (int i = 0; i < npts; ++i) {
        const double shifted = logw[i] - max_logw;
        if (shifted < LOG_NEGLIGIBLE)
            continue;
        const double j = (i + 0.5) * dj;
        const double w = std::exp(shifted);
        if (!std::isfinite(w))
            continue;
        sum_w += w;
        sum_wI += w * j;
        sum_wF += w * (sqrt_dj * grid.sqrt_i[i]);
    }
    if (sum_w <= 0.0) {
        const double j = std::max(I, 0.0);
        return {j, std::sqrt(j)};
    }
    return {sum_wI / sum_w, sum_wF / sum_w};
}

}  // namespace

void ApplyFrenchWilson(std::vector<MergedReflection> &merged, const gemmi::SpaceGroup &space_group,
                       const FrenchWilsonOptions &opts) {
    // Naive amplitude sqrt(max(I,0)) for a missing / strong / untrusted intensity; NaN in -> NaN out
    // (a missing Bijvoet hand stays missing). Fills one (F, sigmaF) pair.
    auto naive_one = [](float I, float sigma, float &F, float &sigF) {
        if (!std::isfinite(I)) { F = NAN; sigF = NAN; return; }
        const double ip = std::max(I, 0.0f);
        F = static_cast<float>(std::sqrt(ip));
        sigF = (ip > 0.0 && std::isfinite(sigma)) ? static_cast<float>(sigma / (2.0 * std::sqrt(ip))) : NAN;
    };
    // The mean intensity and each measured hand share the reflection's Wilson prior, so fill all three.
    auto naive_all = [&](MergedReflection &r) {
        naive_one(r.I, r.sigma, r.F, r.sigmaF);
        naive_one(r.I_plus, r.sigma_plus, r.F_plus, r.sigmaF_plus);
        naive_one(r.I_minus, r.sigma_minus, r.F_minus, r.sigmaF_minus);
    };

    if (merged.empty())
        return;
    const gemmi::GroupOps gops = space_group.operations();

    float d_min = std::numeric_limits<float>::max(), d_max = 0.0f;
    for (const auto &r : merged)
        if (std::isfinite(r.d) && r.d > 0.0f) {
            d_min = std::min(d_min, r.d);
            d_max = std::max(d_max, r.d);
        }
    if (!(d_min < d_max && d_min > 0.0f)) {
        for (auto &r : merged) naive_all(r);
        return;
    }

    // Wilson mean intensity <I/epsilon> per resolution shell, and with an anisotropy tensor the
    // direction it is expected along. The prior of reflection h is then epsilon * K * a(h), with
    // a(h) = exp(-1/2 h^T Q h) - the form the tensor was fitted in - and the shell constant
    // K = sum(I/epsilon) / sum(a) over the shell, so the priors of a shell still average to its
    // measured mean. A shell's weak direction gets a prior matched to its own intensities instead of
    // one set mostly by the strong direction, which turned the weak direction's noise into amplitude.
    // Isotropic data give a near-zero tensor, a(h) ~ 1 and the isotropic prior back, so the tensor is
    // used whenever one was fitted, with no detection gate in front of it: the prior needs the best
    // estimate of <I> along h, whether the anisotropy comes from the crystal or not.
    // Following the anisotropic prior of CCP4 ctruncate (Ballard & Stein; Winn et al. (2011) Acta Cryst. D67, 235-242)
    ResolutionShells shells(d_min * 0.999f, d_max * 1.001f, opts.num_shells);
    std::vector<double> shell_sum(opts.num_shells, 0.0), shell_norm(opts.num_shells, 0.0);
    std::vector<int> shell_count(opts.num_shells, 0);
    double global_sum = 0.0, global_norm = 0.0;
    int global_count = 0;
    auto epsilon = [&](const MergedReflection &r) {
        return std::max(1, gops.epsilon_factor_without_centering({{r.h, r.k, r.l}}));
    };
    auto anisotropy = [&](const MergedReflection &r) {
        const gemmi::Vec3 h(r.h, r.k, r.l);
        return std::exp(-0.5 * h.dot(opts.anisotropy_hkl.multiply(h)));
    };
    // An intensity that gets no amplitude (below reject_below, see fw_one) stays out of the prior too,
    // as ctruncate leaves its outliers out of the norm: kept in, the systematically negative
    // intensities of a background over-subtracted on a powder ring pull the shell mean down and with
    // it every weak amplitude of the shell.
    for (const auto &r : merged) {
        if (!std::isfinite(r.I) || !std::isfinite(r.sigma) || r.sigma <= 0.0f
            || r.I < opts.reject_below * r.sigma)
            continue;
        const double i_over_eps = r.I / epsilon(r);
        const double a = anisotropy(r);
        global_sum += i_over_eps;
        global_norm += a;
        ++global_count;
        if (const auto s = shells.GetShell(r.d)) {
            shell_sum[*s] += i_over_eps;
            shell_norm[*s] += a;
            ++shell_count[*s];
        }
    }
    const double global_mean = global_count > 0 ? std::max(global_sum / global_norm, 1e-10) : 1.0;
    std::vector<double> shell_mean(opts.num_shells, global_mean);
    for (int s = 0; s < opts.num_shells; ++s)
        if (shell_count[s] >= opts.min_reflections_per_shell)
            shell_mean[s] = shell_sum[s] / shell_norm[s];
    // A shell whose mean intensity is not positive has no measurable signal, but a prior of ~0 would
    // still take every amplitude in it to ~0 - more confidently than any measurement says. It takes
    // the nearest lower-resolution shell's mean instead (shell 0 is the lowest resolution), which is
    // where ctruncate's smooth Wilson curve also stays positive.
    for (int s = 0; s < opts.num_shells; ++s)
        if (!(shell_mean[s] > 0.0))
            shell_mean[s] = s > 0 ? shell_mean[s - 1] : global_mean;

    // French-Wilson |F| for one intensity of reflection r (its mean, or one Bijvoet hand); the shell
    // Wilson prior, epsilon and centric flag are the reflection's, shared by all three.
    // The Wilson prior and the centric flag belong to the reflection, not to the intensity, so they
    // are looked up once and shared by its mean and both Bijvoet hands - three symmetry lookups
    // became one.
    auto fw_one = [&](const MergedReflection &r, float I, float sigma, float &F, float &sigF,
                      std::vector<double> &logw, const QuadratureGrid &grid, double sigma_wilson,
                      bool centric) {
        if (!std::isfinite(I) || !std::isfinite(sigma) || sigma <= 0.0f) { naive_one(I, sigma, F, sigF); return; }
        // Far below zero the measurement contradicts any non-negative true intensity, and the
        // posterior would turn it into a small, confident amplitude; ctruncate gives it none
        // ("unphysical") and neither does this. The intensity column keeps it. As CCP4 ctruncate.
        if (I < opts.reject_below * sigma) { F = NAN; sigF = NAN; return; }
        // Strong reflections: the FW correction is negligible, <|F|> = sqrt(I). Not at 4 sigma: the
        // posterior still pulls a 4-6 sigma amplitude down by 2-3% on a weak shell's prior, and a
        // switch there left those amplitudes that much above ctruncate's; at 20 sigma the pull is
        // below 0.3%.
        if (I > opts.strong_cutoff * sigma) { naive_one(I, sigma, F, sigF); return; }
        const Posterior post = integrate_posterior(I, sigma, sigma_wilson, centric,
                                                   opts.integration_points, logw, grid);
        F = static_cast<float>(post.mean_F);
        sigF = static_cast<float>(std::sqrt(std::max(0.0, post.mean_I - post.mean_F * post.mean_F)));
    };
    // Each reflection's amplitudes depend only on itself and the shell priors above, so the loop is
    // data-parallel over contiguous chunks and gives the same result whatever the worker count.
    const int n = static_cast<int>(merged.size());
    const int nt = std::clamp(opts.num_threads, 1, n);
    const int chunk = (n + nt - 1) / nt;
    const QuadratureGrid grid(opts.integration_points);
    auto do_chunk = [&](int lo, int hi) {
        std::vector<double> logw(opts.integration_points);
        for (int i = lo; i < hi; ++i) {
            MergedReflection &r = merged[i];
            const auto s = shells.GetShell(r.d);
            const double sigma_wilson = epsilon(r) * (s ? shell_mean[*s] : global_mean) * anisotropy(r);
            const bool centric = gops.is_reflection_centric({{r.h, r.k, r.l}});
            fw_one(r, r.I, r.sigma, r.F, r.sigmaF, logw, grid, sigma_wilson, centric);
            fw_one(r, r.I_plus, r.sigma_plus, r.F_plus, r.sigmaF_plus, logw, grid, sigma_wilson, centric);
            fw_one(r, r.I_minus, r.sigma_minus, r.F_minus, r.sigmaF_minus, logw, grid, sigma_wilson, centric);
        }
    };
    if (nt == 1) {
        do_chunk(0, n);
        return;
    }
    std::vector<std::future<void>> futures;
    futures.reserve(nt);
    for (int t = 0; t < nt; ++t) {
        const int lo = t * chunk, hi = std::min(n, lo + chunk);
        if (lo >= hi) break;
        futures.emplace_back(std::async(std::launch::async, [&do_chunk, lo, hi] {
            RestoreThreadAffinity();   // not the node of a merging engine that called this (ScopedGpuPin)
            do_chunk(lo, hi);
        }));
    }
    for (auto &f : futures) f.get();
}