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 "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.
Posterior integrate_posterior(double I, double sigma, double sigma_wilson, bool centric, int npts,
                              std::vector<double> &logw) {
    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;

    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 * std::log(j))
                                          : (-j / sigma_wilson);
        logw[i] = log_prior - diff * diff * inv_2s2;
        max_logw = std::max(max_logw, logw[i]);
    }

    double sum_w = 0, sum_wI = 0, sum_wF = 0;
    for (int i = 0; i < npts; ++i) {
        const double j = (i + 0.5) * dj;
        const double w = std::exp(logw[i] - max_logw);
        if (!std::isfinite(w))
            continue;
        sum_w += w;
        sum_wI += w * j;
        sum_wF += w * std::sqrt(j);
    }
    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, int32_t space_group_number,
                       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);
    };

    const gemmi::SpaceGroup *sg = gemmi::find_spacegroup_by_number(space_group_number);
    if (sg == nullptr || merged.empty()) {
        for (auto &r : merged) naive_all(r);
        return;
    }
    const gemmi::GroupOps gops = sg->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.
    ResolutionShells shells(d_min * 0.999f, d_max * 1.001f, opts.num_shells);
    std::vector<double> shell_sum(opts.num_shells, 0.0);
    std::vector<int> shell_count(opts.num_shells, 0);
    double global_sum = 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}}));
    };
    for (const auto &r : merged) {
        if (!std::isfinite(r.I) || !std::isfinite(r.sigma) || r.sigma <= 0.0f)
            continue;
        const double i_over_eps = r.I / epsilon(r);
        global_sum += i_over_eps;
        ++global_count;
        if (const auto s = shells.GetShell(r.d)) {
            shell_sum[*s] += i_over_eps;
            ++shell_count[*s];
        }
    }
    const double global_mean = global_count > 0 ? std::max(global_sum / global_count, 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] = std::max(shell_sum[s] / shell_count[s], 1e-10);

    // 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.
    auto fw_one = [&](const MergedReflection &r, float I, float sigma, float &F, float &sigF,
                      std::vector<double> &logw) {
        if (!std::isfinite(I) || !std::isfinite(sigma) || sigma <= 0.0f) { naive_one(I, sigma, F, sigF); return; }
        // Strong reflections: the FW correction is negligible, <|F|> = sqrt(I).
        if (I > opts.strong_cutoff * sigma) { naive_one(I, sigma, F, sigF); return; }
        const auto s = shells.GetShell(r.d);
        const double sigma_wilson = epsilon(r) * (s ? shell_mean[*s] : global_mean);
        const bool centric = gops.is_reflection_centric({{r.h, r.k, r.l}});
        const Posterior post = integrate_posterior(I, sigma, sigma_wilson, centric,
                                                   opts.integration_points, logw);
        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;
    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];
            fw_one(r, r.I, r.sigma, r.F, r.sigmaF, logw);
            fw_one(r, r.I_plus, r.sigma_plus, r.F_plus, r.sigmaF_plus, logw);
            fw_one(r, r.I_minus, r.sigma_minus, r.F_minus, r.sigmaF_minus, logw);
        }
    };
    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] { do_chunk(lo, hi); }));
    }
    for (auto &f : futures) f.get();
}