Jungfraujoch/image_analysis/structure_refinement/RefinementXray.cpp

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

#include "RefinementXray.h"

#include <algorithm>
#include <cmath>
#include <map>
#include <mutex>
#include <numeric>
#include <set>

#include <fftw3.h>

#include "gemmi/dencalc.hpp"
#include "gemmi/it92.hpp"
#include "gemmi/math.hpp"
#include "gemmi/twin.hpp"

#include "../../common/Logger.h"
#include "../../common/ParallelFor.h"
#include "ModelFFT.h"
#include "ModelGrid.h"
#include "RefinementMath.h"

namespace {

using Table = gemmi::IT92<float>;

// The free reflections a likelihood bin is fitted on, at least.
constexpr size_t FREE_PER_BIN = 100;
constexpr size_t MAX_BINS = 30;

// The density calculator the refinement grids atoms with: DensityCalculator's settings for the grid of
// RigidBodyZoneGrid(), and the given blur.
gemmi::DensityCalculator<Table, float> Calculator(const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg,
                                                  double d_min, double blur) {
    gemmi::DensityCalculator<Table, float> dc;
    dc.d_min = d_min;
    dc.rate = 1.5;   // RigidBodyRefine.cpp's GRID_RATE, which RigidBodyZoneGrid sizes the grid with
    dc.grid.unit_cell = cell;
    dc.grid.spacegroup = &sg;
    dc.blur = blur;
    dc.cutoff = REFINEMENT_DENSITY_CUTOFF;
    return dc;
}

// The likelihood's variance of reflection i: epsilon beta, plus the measurement error of each of the structure
// factor's components that the observed amplitude measures - sigma^2 twice for an acentric reflection, once for a
// centric one (with_sigma = false leaves it out).
double Variance(const RefinementReflections &refl, size_t i, double beta, bool with_sigma) {
    const double s2 = with_sigma ? double(refl.sigma[i]) * refl.sigma[i] : 0.0;
    return refl.epsilon[i] * beta + (refl.centric[i] ? s2 : 2.0 * s2);
}

// The first and second derivatives of one row's -ln L (refine_math::AmplitudeNegLogLikelihood) in D and in the
// variance, at D and var.
struct RowDerivatives {
    double d, dd, v, vv;
};

RowDerivatives LikelihoodDerivatives(double fo, double fc, double d, double var, bool centric) {
    const double q = fo * fo + d * d * fc * fc;
    if (!centric) {
        // x = 2 D Fo Fc / var, m = I1(x) / I0(x) = d ln I0 / dx, m' = 1 - m / x - m^2
        const double a = 2.0 * fo * fc / var, x = a * d;
        double log_i0, m;
        refine_math::BesselI0I1(x, log_i0, m);
        const double m1 = x > 0 ? 1.0 - m / x - m * m : 0.5;
        return {2.0 * d * fc * fc / var - a * m,
                2.0 * fc * fc / var - a * a * m1,
                1.0 / var - q / (var * var) + m * x / var,
                (-1.0 + 2.0 * q / var - x * (m1 * x + 2.0 * m)) / (var * var)};
    }
    // x = D Fo Fc / var, t = tanh x = d ln cosh / dx
    const double a = fo * fc / var, x = a * d;
    const double t = std::tanh(x), t1 = 1.0 - t * t;
    return {d * fc * fc / var - a * t,
            fc * fc / var - a * a * t1,
            0.5 / var - q / (2.0 * var * var) + t * x / var,
            (-0.5 + q / var - x * (t1 * x + 2.0 * t)) / (var * var)};
}

// The minimum of a function over [lo, hi] from x, given its first and second derivative (derivatives(x, g, h)):
// Newton's method on the first derivative, kept inside the interval where that changes sign and bisecting it
// wherever a step would leave it or the curvature is not positive. A minimum on a bound is the bound.
template <typename F>
double NewtonMinimum(F derivatives, double x, double lo, double hi) {
    x = std::clamp(x, lo, hi);
    for (int iteration = 0; iteration < 200; iteration++) {
        double g = 0, h = 0;
        derivatives(x, g, h);
        if (g < 0)
            lo = x;
        else
            hi = x;
        double next = h > 0 ? x - g / h : lo - 1.0;
        if (!(next > lo && next < hi))
            next = 0.5 * (lo + hi);
        const bool converged = std::fabs(next - x) < 1e-10 || hi - lo < 1e-10;
        x = next;
        if (converged)
            break;
    }
    return x;
}

class RefinementEngineCPU : public RefinementEngine {
public:
    RefinementEngineCPU(const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg, double d_min,
                            const RefinementReflections &refl, size_t nthreads)
        : cell_(cell), sg_(sg), d_min_(d_min), refl_(refl), nthreads_(nthreads),
          grid_(RigidBodyZoneGrid(cell, sg, d_min)), composition_(grid_, d_min, refl.RowHkl()),
          orbit_leaders_(OrbitLeaders(grid_, nthreads)) {}

    void SetBlur(double blur) override { blur_ = blur; }
    void SetRowModel(const XrayRowModel &rows) override { rows_ = rows; }
    bool OnGPU() const override { return false; }

    void Fcalc(const gemmi::Model &model, std::vector<std::complex<float>> &f) override {
        std::vector<std::complex<double>> fd;
        Compose(WithAniso(model), fd);
        f.assign(fd.begin(), fd.end());
    }

    void Fmask(const gemmi::Model &model, std::vector<std::complex<float>> &f) override {
        gemmi::Grid<float> mask = grid_;
        PutMaskOnGrid(mask, model, orbit_leaders_, nthreads_);
        const gemmi::FPhiGrid<float> fm = MapToFPhi(mask);
        f.clear();
        for (const gemmi::Miller &h : refl_.RowHkl())
            f.push_back(fm.get_value_by_hkl(h));
    }

    double Evaluate(const gemmi::Model &model_in, std::vector<std::array<double, 4>> *grad) override {
        const gemmi::Model &model = WithAniso(model_in);
        std::vector<std::complex<double>> f;
        Compose(model, f);
        const size_t n = f.size();
        // The target and, per reflection, dT/dF of the atoms' structure factor (dT/dRe F + i dT/dIm F).
        std::vector<std::complex<double>> dt_df(n, 0.0);
        double sum = 0;
        const size_t laws = refl_.Laws();
        double own_fraction = 1.0;
        for (float a : rows_.twin_fractions)
            own_fraction -= a;
        for (size_t m = 0; m < refl_.Size(); m++) {
            if (refl_.free[m])
                continue;
            const double k = rows_.k_total[m];
            if (!refl_.Twinned()) {
                const std::complex<double> fm = k * (f[m] + std::complex<double>(rows_.fmask[m]));
                const double fc = std::abs(fm);
                double dfc = 0;
                sum += refine_math::AmplitudeNegLogLikelihood<double>(refl_.fo[m], fc, rows_.d[m], rows_.var[m],
                                                                     refl_.centric[m], dfc);
                if (fc > 0)
                    dt_df[m] = k * dfc * fm / fc;
                continue;
            }
            // twinned: |F_model| = k sqrt(a_0 |U(h)|^2 + sum_d a_d |U(T_d h)|^2), U = F + fmask, so that
            // dT/dU(T_d h) = dT/d|F_model| k^2 a_d U(T_d h) / |F_model|
            const std::complex<double> u1 = f[m] + std::complex<double>(rows_.fmask[m]);
            double i_model = own_fraction * std::norm(u1);
            for (size_t d = 0; d < laws; d++) {
                const int j = refl_.mates[d][m];
                i_model += rows_.twin_fractions[d] * std::norm(f[j] + std::complex<double>(rows_.fmask[j]));
            }
            const double fc = k * std::sqrt(std::max(i_model, 0.0));
            double dfc = 0;
            sum += refine_math::AmplitudeNegLogLikelihood<double>(refl_.fo[m], fc, rows_.d[m], rows_.var[m],
                                                                 refl_.centric[m], dfc);
            if (fc > 0) {
                const double g = dfc * k * k / fc;
                dt_df[m] += g * own_fraction * u1;
                for (size_t d = 0; d < laws; d++) {
                    const int j = refl_.mates[d][m];
                    dt_df[j] += g * double(rows_.twin_fractions[d]) * (f[j] + std::complex<double>(rows_.fmask[j]));
                }
            }
        }
        if (grad == nullptr)
            return sum;

        // The FFT method's map, dT/drho(x) = (V/N) Re sum_m sum_ops conj(dT/dF_m) scale_m exp(2 pi i h.t)
        // exp(2 pi i (hR).x), as the inverse real transform of a half-l coefficient grid: a term A at index k
        // goes in as A/2 at k, or as conj(A)/2 at -k where k is in the half not stored; on the planes l = 0 and
        // l = nw/2, which hold both, as both.
        const int nu = grid_.nu, nv = grid_.nv, nw = grid_.nw, nwh = nw / 2 + 1;
        const size_t ncomplex = static_cast<size_t>(nu) * nv * nwh;
        std::complex<float> *in = reinterpret_cast<std::complex<float> *>(fftwf_alloc_complex(ncomplex));
        float *out = fftwf_alloc_real(static_cast<size_t>(nu) * nv * nw);
        std::fill(in, in + ncomplex, std::complex<float>(0, 0));
        const std::vector<SymmetryComposition::Term> &terms = composition_.Terms();
        const size_t ops = composition_.Ops();
        auto add = [&](int ku, int kv, int kw, std::complex<double> value) {
            in[gemmi::modulo(ku, nu) + static_cast<size_t>(nu) * (gemmi::modulo(kv, nv) + static_cast<size_t>(nv) * kw)] +=
                std::complex<float>(value);
        };
        for (size_t m = 0; m < n; m++) {
            if (dt_df[m] == 0.0)
                continue;
            const double scale = composition_.Centring() * std::exp(blur_ * 0.25 * composition_.InvD2()[m]);
            const std::complex<double> c = std::conj(dt_df[m]) * scale;
            for (size_t o = 0; o < ops; o++) {
                const SymmetryComposition::Term &t = terms[m * ops + o];
                const std::complex<double> a = 0.5 * c * t.phase;
                const gemmi::Miller &k = t.k;
                if (k[2] > 0 && 2 * k[2] != nw)
                    add(k[0], k[1], k[2], a);
                else if (k[2] < 0 && -2 * k[2] != nw)
                    add(-k[0], -k[1], -k[2], std::conj(a));
                else {
                    const int l = k[2] < 0 ? -k[2] : k[2];
                    add(k[0], k[1], l, a);
                    add(-k[0], -k[1], l, std::conj(a));
                }
            }
        }
        MapFromFPhiExecute(nu, nv, nw, in, out);
        gemmi::Grid<float> map = grid_;
        const float norm = static_cast<float>(cell_.volume / static_cast<double>(map.point_count()));
        for (size_t i = 0; i < map.data.size(); i++)
            map.data[i] = out[i] * norm;
        fftwf_free(in);
        fftwf_free(out);
        const bool aniso = AnyAniso();
        AtomGradientsFromMap(map, model, blur_, nthreads_, *grad, aniso ? &grad_u_ : nullptr);
        if (!aniso)
            grad_u_.clear();
        return sum;
    }

private:
    // The model as the engine grids it: where some atom is anisotropic, a copy with every atom's U set (ApplyAniso).
    const gemmi::Model &WithAniso(const gemmi::Model &model) {
        if (!AnyAniso())
            return model;
        aniso_model_ = model;
        ApplyAniso(aniso_model_, *aniso_);
        return aniso_model_;
    }

    void Compose(const gemmi::Model &model, std::vector<std::complex<double>> &f) const {
        gemmi::DensityCalculator<Table, float> dc = Calculator(cell_, sg_, d_min_, blur_);
        PutModelDensityOnGrid(dc, model, {}, nthreads_, true, true);   // one copy: the symmetry is composed below
        composition_.Compose(MapToFPhi(dc.grid), blur_, f, nullptr, nthreads_);
    }

    gemmi::UnitCell cell_;
    const gemmi::SpaceGroup &sg_;
    double d_min_;
    const RefinementReflections &refl_;
    size_t nthreads_;
    gemmi::Grid<float> grid_;
    SymmetryComposition composition_;
    std::vector<size_t> orbit_leaders_;
    double blur_ = 0;
    XrayRowModel rows_;
    gemmi::Model aniso_model_{1};
};

}  // namespace

RefinementReflections PrepareRefinementReflections(const std::vector<MergedReflection> &merged,
                                                   const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg,
                                                   double d_min) {
    const gemmi::GroupOps gops = sg.operations();
    const gemmi::ReciprocalAsu asu(&sg);
    // ASU, last one wins where a reflection is listed twice (both hands of an anomalous merge)
    std::map<gemmi::Miller, const MergedReflection *> by_hkl;
    for (const MergedReflection &r : merged) {
        if (!std::isfinite(r.F) || r.F < 0)
            continue;
        gemmi::Miller h{{r.h, r.k, r.l}};
        if (!asu.is_in(h))
            h = asu.to_asu(h, gops).first;
        by_hkl[h] = &r;
    }
    std::vector<gemmi::Miller> hkl;
    for (const auto &[h, r] : by_hkl)
        hkl.push_back(h);
    const SymmetryComposition composition(RigidBodyZoneGrid(cell, sg, d_min), d_min, hkl);
    RefinementReflections out;
    for (size_t i = 0; i < hkl.size(); i++) {
        if (composition.Row()[i] < 0)
            continue;
        const MergedReflection &r = *by_hkl[hkl[i]];
        out.hkl.push_back(hkl[i]);
        out.fo.push_back(r.F);
        out.sigma.push_back(std::isfinite(r.sigmaF) ? r.sigmaF : 0.0f);
        out.free.push_back(r.rfree_flag);
        out.centric.push_back(gops.is_reflection_centric(hkl[i]));
        out.epsilon.push_back(static_cast<float>(gops.epsilon_factor(hkl[i])));
        out.inv_d2.push_back(cell.calculate_1_d2(hkl[i]));
    }
    return out;
}

std::vector<gemmi::Miller> RefinementReflections::RowHkl() const {
    std::vector<gemmi::Miller> out = hkl;
    out.insert(out.end(), mate_hkl.begin(), mate_hkl.end());
    return out;
}

std::vector<gemmi::Op> RefinementTwinLaws(const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg) {
    const gemmi::GroupOps gops = sg.operations();
    std::vector<gemmi::Op> laws;
    std::vector<gemmi::Op> found;
    try {
        found = gemmi::find_twin_laws(cell, &sg, 3.0, false);
    } catch (const std::exception &) {
        return laws;   // a setting whose centring the lattice search does not reproduce: no law is looked for
    }
    for (const gemmi::Op &op : found) {
        const gemmi::Op law{op.rot, {0, 0, 0}, 'x'};
        const gemmi::Op square = law * law;
        if (gops.find_by_rotation(square.rot) != nullptr || gops.find_by_rotation(square.negated_rot()) != nullptr)
            laws.push_back(law);
    }
    return laws;
}

std::vector<gemmi::Op> RefinementTwinDomains(const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg) {
    std::vector<gemmi::Op> laws;
    try {
        for (const gemmi::Op &op : gemmi::find_twin_laws(cell, &sg, 3.0, false))
            laws.push_back(gemmi::Op{op.rot, {0, 0, 0}, 'x'});
    } catch (const std::exception &) {
    }
    return laws;
}

int SetTwinLaws(RefinementReflections &refl, const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg, double d_min,
                const std::vector<gemmi::Op> &laws) {
    const gemmi::GroupOps gops = sg.operations();
    const gemmi::ReciprocalAsu asu(&sg);
    const size_t n = refl.Size();
    std::map<gemmi::Miller, int> row;
    for (size_t m = 0; m < n; m++)
        row[refl.hkl[m]] = static_cast<int>(m);
    // each observed reflection's mates in the ASU, and those not observed - where the grid composes them - as rows
    // of their own
    std::vector<std::vector<gemmi::Miller>> mate_of(laws.size(), std::vector<gemmi::Miller>(n));
    std::set<gemmi::Miller> unobserved;
    for (size_t d = 0; d < laws.size(); d++)
        for (size_t m = 0; m < n; m++) {
            mate_of[d][m] = asu.to_asu(laws[d].apply_to_hkl(refl.hkl[m]), gops).first;
            if (row.find(mate_of[d][m]) == row.end())
                unobserved.insert(mate_of[d][m]);
        }
    const std::vector<gemmi::Miller> candidates(unobserved.begin(), unobserved.end());
    const SymmetryComposition composition(RigidBodyZoneGrid(cell, sg, d_min), d_min, candidates);
    refl.twin_laws = laws;
    refl.mate_hkl.clear();
    for (size_t i = 0; i < candidates.size(); i++)
        if (composition.Row()[i] >= 0) {
            row[candidates[i]] = static_cast<int>(n + refl.mate_hkl.size());
            refl.mate_hkl.push_back(candidates[i]);
        }
    const size_t rows = refl.Rows();
    refl.mates.assign(laws.size(), std::vector<int>(rows));
    refl.mates_from.assign(laws.size(), std::vector<int>(rows, -1));
    int split = 0;
    for (size_t d = 0; d < laws.size(); d++) {
        std::vector<int> &mate = refl.mates[d];
        std::iota(mate.begin(), mate.end(), 0);
        for (size_t m = 0; m < n; m++) {
            const auto it = row.find(mate_of[d][m]);
            if (it != row.end() && it->second != static_cast<int>(m)) {
                mate[m] = it->second;
                refl.mates_from[d][it->second] = static_cast<int>(m);
            }
        }
        for (size_t m = 0; m < n; m++) {
            const int j = mate[m];
            if (j != static_cast<int>(m) && j < static_cast<int>(n) && refl.free[m] != refl.free[j] &&
                (refl.mates[d][j] != static_cast<int>(m) || j > static_cast<int>(m)))
                split++;
        }
    }
    return split;
}

std::vector<std::complex<float>> TwinnedModel(const RefinementReflections &refl, const XrayRowModel &rows,
                                              const std::vector<std::complex<float>> &f) {
    const size_t n = refl.Size();
    std::vector<std::complex<float>> out(n);
    if (!refl.Twinned()) {
        for (size_t m = 0; m < n; m++)
            out[m] = rows.k_total[m] * (f[m] + rows.fmask[m]);
        return out;
    }
    double own_fraction = 1.0;
    for (float a : rows.twin_fractions)
        own_fraction -= a;
    for (size_t m = 0; m < n; m++) {
        const std::complex<double> u1 = std::complex<double>(f[m]) + std::complex<double>(rows.fmask[m]);
        const double i1 = std::norm(u1);
        double i_model = own_fraction * i1;
        for (size_t d = 0; d < refl.Laws(); d++) {
            const int j = refl.mates[d][m];
            i_model += rows.twin_fractions[d] *
                       std::norm(std::complex<double>(f[j]) + std::complex<double>(rows.fmask[j]));
        }
        const double amplitude = rows.k_total[m] * std::sqrt(std::max(i_model, 0.0));
        out[m] = i1 > 0 ? std::complex<float>(u1 * (amplitude / std::sqrt(i1))) : std::complex<float>(amplitude, 0);
    }
    return out;
}

std::vector<float> FitTwinFractions(const RefinementReflections &refl, const XrayRowModel &rows,
                                    const std::vector<std::complex<float>> &f) {
    const size_t laws = refl.Laws();
    std::vector<float> out(laws, 0.0f);
    if (laws == 0)
        return out;
    // the working reflections with a mate that is another reflection: Fo^2 and the model's intensity at h and at
    // each T_d h
    const size_t nx = laws + 1;
    struct Term {
        double inv_d2, io;
        std::vector<double> x;   // the model's intensity at h, then at T_d h
    };
    std::vector<Term> terms;
    for (size_t m = 0; m < refl.Size(); m++) {
        if (refl.free[m])
            continue;
        bool related = false;
        for (size_t d = 0; d < laws; d++)
            related |= refl.mates[d][m] != static_cast<int>(m);
        if (!related)
            continue;
        const double k2 = double(rows.k_total[m]) * rows.k_total[m];
        Term t{refl.inv_d2[m], double(refl.fo[m]) * refl.fo[m], std::vector<double>(nx)};
        t.x[0] = k2 * std::norm(std::complex<double>(f[m]) + std::complex<double>(rows.fmask[m]));
        for (size_t d = 0; d < laws; d++) {
            const int j = refl.mates[d][m];
            t.x[d + 1] = k2 * std::norm(std::complex<double>(f[j]) + std::complex<double>(rows.fmask[j]));
        }
        terms.push_back(std::move(t));
    }
    if (terms.size() < nx + 2)
        return out;
    // normalised per resolution shell of TWIN_FRACTION_SHELL reflections: Fo^2 by its mean, the model's by the mean
    // of its own intensity
    std::stable_sort(terms.begin(), terms.end(), [](const Term &a, const Term &b) { return a.inv_d2 < b.inv_d2; });
    const size_t shells = std::max<size_t>(1, terms.size() / TWIN_FRACTION_SHELL);
    // the normal equations of y = b0 + sum_p b_p x_p, unknowns 0 (the intercept) to nx
    const size_t nu = nx + 1;
    std::vector<double> nm(nu * nu, 0.0), r(nu, 0.0);
    std::vector<double> x(nu);
    for (size_t k = 0; k < shells; k++) {
        const size_t lo = k * terms.size() / shells, hi = (k + 1) * terms.size() / shells;
        double so = 0, sc = 0;
        for (size_t i = lo; i < hi; i++) {
            so += terms[i].io;
            sc += terms[i].x[0];
        }
        if (so <= 0 || sc <= 0)
            continue;
        for (size_t i = lo; i < hi; i++) {
            x[0] = 1.0;
            for (size_t p = 0; p < nx; p++)
                x[p + 1] = terms[i].x[p] / sc;
            const double y = terms[i].io / so;
            for (size_t p = 0; p < nu; p++) {
                r[p] += x[p] * y;
                for (size_t q = 0; q < nu; q++)
                    nm[p * nu + q] += x[p] * x[q];
            }
        }
    }
    // solved over the unknowns still in, by Gaussian elimination; a law whose coefficient is negative is taken out
    // (the most negative first) and the rest solved again
    std::vector<char> in(nu, 1);
    std::vector<double> b(nu, 0.0);
    for (size_t round = 0; round < nu; round++) {
        std::vector<size_t> idx;
        for (size_t p = 0; p < nu; p++)
            if (in[p])
                idx.push_back(p);
        const size_t k = idx.size();
        std::vector<double> a(k * (k + 1));
        for (size_t p = 0; p < k; p++) {
            for (size_t q = 0; q < k; q++)
                a[p * (k + 1) + q] = nm[idx[p] * nu + idx[q]];
            a[p * (k + 1) + k] = r[idx[p]];
        }
        bool singular = false;
        for (size_t c = 0; c < k && !singular; c++) {
            size_t piv = c;
            for (size_t p = c + 1; p < k; p++)
                if (std::fabs(a[p * (k + 1) + c]) > std::fabs(a[piv * (k + 1) + c]))
                    piv = p;
            if (!(std::fabs(a[piv * (k + 1) + c]) > 0)) {
                singular = true;
                break;
            }
            for (size_t q = 0; q <= k; q++)
                std::swap(a[c * (k + 1) + q], a[piv * (k + 1) + q]);
            for (size_t p = 0; p < k; p++)
                if (p != c) {
                    const double factor = a[p * (k + 1) + c] / a[c * (k + 1) + c];
                    for (size_t q = c; q <= k; q++)
                        a[p * (k + 1) + q] -= factor * a[c * (k + 1) + q];
                }
        }
        if (singular)
            return out;
        std::fill(b.begin(), b.end(), 0.0);
        for (size_t p = 0; p < k; p++)
            b[idx[p]] = a[p * (k + 1) + k] / a[p * (k + 1) + p];
        size_t worst = 0;
        for (size_t p = 2; p < nu; p++)
            if (in[p] && b[p] < 0 && (worst == 0 || b[p] < b[worst]))
                worst = p;
        if (worst == 0)
            break;
        in[worst] = 0;
    }
    double total = 0;
    for (size_t p = 1; p < nu; p++)
        total += b[p];
    if (!(total > 0))
        return out;
    for (size_t d = 0; d < laws; d++)
        out[d] = static_cast<float>(std::max(0.0, b[d + 2] / total));
    if (laws == 1)
        out[0] = std::min(out[0], 0.5f);
    return out;
}

double FitTwinFraction(const RefinementReflections &refl, const XrayRowModel &rows,
                       const std::vector<std::complex<float>> &f) {
    double sum = 0;
    for (float a : FitTwinFractions(refl, rows, f))
        sum += a;
    return sum;
}

BulkSolventMask::BulkSolventMask(const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg, double d_min,
                                 size_t nthreads)
    : grid_(MaskGrid(cell, sg, d_min)), nthreads_(nthreads), spacing_(MaskSpacing(cell, d_min)) {}

double BulkSolventMask::MaskSpacing(const gemmi::UnitCell &cell, double d_min) {
    const double spacing = std::max(std::min(d_min / 3.0, MASK_SPACING_MAX), std::cbrt(cell.volume / MASK_POINTS_MAX));
    return std::min({spacing, MASK_SPACING_SHRINK, d_min / 3.0});
}

gemmi::Grid<float> BulkSolventMask::MaskGrid(const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg, double d_min) {
    return SolventMaskGrid(cell, sg, d_min, MaskSpacing(cell, d_min));
}

void BulkSolventMask::Fmask(const gemmi::Model &model, const std::vector<gemmi::Miller> &hkl,
                            std::vector<std::complex<float>> &f) const {
    if (orbit_leaders_.empty())   // at the first mask: a GPU engine's refinement never draws one here
        orbit_leaders_ = OrbitLeaders(grid_, nthreads_);
    gemmi::Grid<float> mask = grid_;
    mask.data.resize(mask.point_count());
    PutMaskOnGrid(mask, model, orbit_leaders_, nthreads_);
    const gemmi::FPhiGrid<float> fm = MapToFPhi(mask);
    f.clear();
    for (const gemmi::Miller &h : hkl)
        f.push_back(fm.get_value_by_hkl(h));
}

void EstimateLikelihoodParameters(const RefinementReflections &refl, const std::vector<std::complex<float>> &fmodel,
                                  XrayRowModel &rows, size_t nthreads) {
    const size_t n = refl.Size();
    std::vector<double> fc(n);
    for (size_t i = 0; i < n; i++)
        fc[i] = std::abs(fmodel[i]);
    std::vector<size_t> free, work;
    for (size_t i = 0; i < n; i++)
        if (fc[i] > 0)
            (refl.free[i] ? free : work).push_back(i);
    // Too few free reflections to bin: the working set, which says the model is better than it is, but
    // still gives the likelihood a scale.
    if (free.size() < FREE_PER_BIN)
        free.insert(free.end(), work.begin(), work.end());
    std::sort(free.begin(), free.end(), [&](size_t a, size_t b) { return refl.inv_d2[a] < refl.inv_d2[b]; });
    const size_t bins = std::clamp<size_t>(free.size() / FREE_PER_BIN, 1, MAX_BINS);
    // the working reflections of each bin: those from its first free reflection's 1/d^2 to the next bin's
    std::vector<double> lower(bins);
    for (size_t b = 0; b < bins; b++)
        lower[b] = refl.inv_d2[free[b * free.size() / bins]];
    std::vector<std::vector<size_t>> work_in_bin(bins);
    for (size_t i : work) {
        const size_t b = std::upper_bound(lower.begin(), lower.end(), refl.inv_d2[i]) - lower.begin();
        work_in_bin[b > 0 ? b - 1 : 0].push_back(i);
    }

    std::vector<double> centre(bins), bin_d(bins), bin_beta(bins);
    const bool warm = rows.bin_d.size() == bins;
    // each bin is fitted on its own
    ParallelFor(static_cast<int>(bins), nthreads, [&](int b) {
        const std::vector<size_t> idx(free.begin() + b * free.size() / bins, free.begin() + (b + 1) * free.size() / bins);
        double s2 = 0, sfo = 0, sfc2 = 0, sfofc = 0;
        for (size_t i : idx) {
            s2 += refl.inv_d2[i];
            sfo += double(refl.fo[i]) * refl.fo[i] / refl.epsilon[i];
            sfc2 += fc[i] * fc[i] / refl.epsilon[i];
            sfofc += refl.fo[i] * fc[i] / refl.epsilon[i];
        }
        centre[b] = s2 / static_cast<double>(idx.size());
        double d = sfc2 > 0 ? std::max(sfofc / sfc2, 0.05) : 1.0;
        double beta = std::max((sfo - d * d * sfc2) / static_cast<double>(idx.size()), 1e-3 * sfo / idx.size());
        // beta: the spread of the free amplitudes about D Fc, measurement error included, by maximum likelihood
        // (alternating one-dimensional minimisations over ln D and ln beta)
        for (int round = 0; round < 4; round++) {
            const double ld = NewtonMinimum([&](double x, double &g, double &h) {
                const double dx = std::exp(x);
                for (size_t i : idx) {
                    const RowDerivatives r = LikelihoodDerivatives(refl.fo[i], fc[i], dx, refl.epsilon[i] * beta,
                                                                   refl.centric[i]);
                    g += dx * r.d;
                    h += dx * dx * r.dd + dx * r.d;
                }
            }, std::log(d), std::log(d) - 2.0, std::log(std::min(d * 7.4, 3.0)));
            d = std::exp(ld);
            const double lb = NewtonMinimum([&](double x, double &g, double &h) {
                for (size_t i : idx) {
                    const double var = refl.epsilon[i] * std::exp(x);
                    const RowDerivatives r = LikelihoodDerivatives(refl.fo[i], fc[i], d, var, refl.centric[i]);
                    g += var * r.v;
                    h += var * var * r.vv + var * r.v;
                }
            }, std::log(beta), std::log(beta) - 3.0, std::log(beta) + 3.0);
            beta = std::exp(lb);
        }
        // D: by maximum likelihood on the working reflections, with the measurement error added to that beta, from
        // where the last fit found it. Where the error is a large part of the spread, the variance counts it twice
        // and D falls towards 0 - the high-resolution D and figure of merit REFMAC 5.8 reports.
        // Following Murshudov et al. (2011) Acta Cryst. D67, 355-367 (D from the working set, the variance from the free)
        if (!work_in_bin[b].empty()) {
            std::vector<double> var;
            for (size_t i : work_in_bin[b])
                var.push_back(Variance(refl, i, beta, true));
            d = NewtonMinimum([&](double x, double &g, double &h) {
                for (size_t k = 0; k < var.size(); k++) {
                    const size_t i = work_in_bin[b][k];
                    const RowDerivatives r = LikelihoodDerivatives(refl.fo[i], fc[i], x, var[k], refl.centric[i]);
                    g += r.d;
                    h += r.dd;
                }
            }, warm ? rows.bin_d[b] : d, 0.0, 3.0);
        }
        bin_d[b] = d;
        bin_beta[b] = beta;
    });
    rows.bin_d = bin_d;
    rows.d.resize(n);
    rows.var.resize(n);
    for (size_t i = 0; i < n; i++) {
        // linear in 1/d^2 between the bins' centres, flat beyond the outermost
        const double s = refl.inv_d2[i];
        double d = bin_d[0], beta = bin_beta[0];
        if (s >= centre[bins - 1]) {
            d = bin_d[bins - 1];
            beta = bin_beta[bins - 1];
        } else if (s > centre[0]) {
            const size_t b = std::upper_bound(centre.begin(), centre.end(), s) - centre.begin();
            const double t = (s - centre[b - 1]) / (centre[b] - centre[b - 1]);
            d = bin_d[b - 1] + t * (bin_d[b] - bin_d[b - 1]);
            beta = bin_beta[b - 1] + t * (bin_beta[b] - bin_beta[b - 1]);
        }
        rows.d[i] = static_cast<float>(d);
        rows.var[i] = static_cast<float>(Variance(refl, i, beta, true));
    }
}

RFactors ComputeRFactors(const RefinementReflections &refl, const std::vector<std::complex<float>> &fmodel) {
    double nw = 0, dw = 0, nf = 0, df = 0;
    RFactors r;
    for (size_t i = 0; i < refl.Size(); i++) {
        const double diff = std::fabs(refl.fo[i] - std::abs(fmodel[i]));
        if (refl.free[i]) {
            nf += diff;
            df += refl.fo[i];
            r.n_free++;
        } else {
            nw += diff;
            dw += refl.fo[i];
            r.n_work++;
        }
    }
    r.r_work = dw > 0 ? nw / dw : 1;
    r.r_free = df > 0 ? nf / df : 1;
    return r;
}

gemmi::SMat33<float> AnisoU(const AtomAniso &a, double b) {
    const double u = b / gemmi::u_to_b();
    gemmi::SMat33<double> m{a.dev[0] + u, a.dev[1] + u, a.dev[2] + u, a.dev[3], a.dev[4], a.dev[5]};
    const std::array<double, 3> ev = m.calculate_eigenvalues();
    const double lowest = std::min({ev[0], ev[1], ev[2]});
    if (!(lowest >= ANISO_U_MIN))
        m = m.added_kI(ANISO_U_MIN - (std::isfinite(lowest) ? lowest : 0.0));
    return {static_cast<float>(m.u11), static_cast<float>(m.u22), static_cast<float>(m.u33),
            static_cast<float>(m.u12), static_cast<float>(m.u13), static_cast<float>(m.u23)};
}

void ApplyAniso(gemmi::Model &model, const std::vector<AtomAniso> &aniso) {
    size_t i = 0;
    for (gemmi::Chain &ch : model.chains)
        for (gemmi::Residue &r : ch.residues)
            for (gemmi::Atom &a : r.atoms) {
                if (i < aniso.size() && aniso[i].on)
                    a.aniso = AnisoU(aniso[i], a.b_iso);
                else
                    a.aniso = {0, 0, 0, 0, 0, 0};
                i++;
            }
}

std::vector<std::array<double, 4>> ModelXyzB(const gemmi::Model &model) {
    std::vector<std::array<double, 4>> out;
    for (const gemmi::Chain &ch : model.chains)
        for (const gemmi::Residue &r : ch.residues)
            for (const gemmi::Atom &a : r.atoms)
                out.push_back({a.pos.x, a.pos.y, a.pos.z, a.b_iso});
    return out;
}

namespace {
std::mutex stage_mutex;
std::map<std::string, std::pair<double, int>> stage_times;
}  // namespace

RefinementStage::RefinementStage(const char *name)
    : name_(name), start_(std::chrono::steady_clock::now()), nvtx_(NvtxRangeStart(name)) {}

RefinementStage::~RefinementStage() {
    NvtxRangeEnd(nvtx_);
    const double seconds = std::chrono::duration<double>(std::chrono::steady_clock::now() - start_).count();
    std::lock_guard lock(stage_mutex);
    std::pair<double, int> &t = stage_times[name_];
    t.first += seconds;
    t.second++;
}

void RefinementStage::Log(Logger &logger) {
    std::lock_guard lock(stage_mutex);
    for (const auto &[name, t] : stage_times)
        logger.Debug("TIMING refine.{} {:.3f} {}", name, t.first, t.second);
    stage_times.clear();
}

double RefinementBlur(const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg, double d_min, double b_floor) {
    // set_refmac_compatible_blur(): u_to_b / 1.1 * spacing^2 - b_min, not below 0, the spacing the one asked for
    const gemmi::DensityCalculator<Table, float> dc = Calculator(cell, sg, d_min, 0);
    const double spacing = dc.requested_grid_spacing();
    return std::max(0.0, gemmi::u_to_b() / 1.1 * spacing * spacing - b_floor);
}

std::unique_ptr<RefinementEngine> MakeRefinementEngineCPU(const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg,
                                                          double d_min, const RefinementReflections &refl,
                                                          size_t nthreads) {
    return std::make_unique<RefinementEngineCPU>(cell, sg, d_min, refl, nthreads);
}

namespace {

// AtomGradientsFromMap() for an anisotropic atom, gridded as PutModelDensityOnGrid() grids it: rho = occ sum_q a_q
// exp(r^T M_q r), M_q = -4 pi^2 (B_q I + 8 pi^2 U + blur I)^-1 and r = centre - point. Then d rho / d(centre) =
// occ sum_q 2 e_q M_q r, and, as d a_q / dU = -4 pi^2 a_q (B..)^-1 = a_q M_q and d(r^T M_q r) / dU = 2 (M_q r)(M_q r)^T,
// d rho / dU = occ sum_q e_q (M_q + 2 (M_q r)(M_q r)^T) - an off-diagonal U_ij twice its matrix element.
void AnisoGradient(const gemmi::Grid<float> &map, const gemmi::DensityCalculator<Table, float> &dc,
                   const gemmi::Atom &atom, double blur, std::array<double, 4> &grad, std::array<double, 6> *grad_u) {
    using CReal = Table::Coef::coef_type;
    const auto &coef = IonFormFactor<float>(atom.element.elem, atom.charge);
    const auto aniso_b = atom.aniso.scaled(CReal(gemmi::u_to_b())).added_kI(CReal(blur));
    const CReal b_max = AnisoRadiusB(aniso_b);
    const double radius = dc.estimate_radius(coef.precalculate_density_iso(b_max, 0.0f), b_max);
    const auto precal = coef.precalculate_density_aniso_b(aniso_b, 0.0f);
    const float max_r2 = static_cast<float>(radius * radius);
    const int nu = map.nu, nv = map.nv, nw = map.nw;
    const gemmi::UpperTriangularMat33 &om = map.orth_n;
    const float m00 = static_cast<float>(om.a11), m01 = static_cast<float>(om.a12), m02 = static_cast<float>(om.a13),
                m11 = static_cast<float>(om.a22), m12 = static_cast<float>(om.a23), m22 = static_cast<float>(om.a33);
    const gemmi::Fractional fpos = map.unit_cell.fractionalize(atom.pos);
    const double n3[3] = {fpos.x * nu, fpos.y * nv, fpos.z * nw};
    float nctr[3];
    int c[3], d[3];
    for (int k = 0; k < 3; k++) {
        nctr[k] = static_cast<float>(n3[k]);
        c[k] = gemmi::iround(n3[k]);
    }
    d[0] = std::min(static_cast<int>(std::ceil(radius / map.spacing[0])), nu - 1);
    d[1] = std::min(static_cast<int>(std::ceil(radius / map.spacing[1])), nv - 1);
    d[2] = std::min(static_cast<int>(std::ceil(radius / map.spacing[2])), nw - 1);
    double g[3] = {0, 0, 0}, gu[6] = {0, 0, 0, 0, 0, 0};
    for (int w = c[2] - d[2]; w <= c[2] + d[2]; w++) {
        const float fz = nctr[2] - static_cast<float>(w);
        const float z = m22 * fz;
        const int w_ = gemmi::modulo(w, nw);
        for (int v = c[1] - d[1]; v <= c[1] + d[1]; v++) {
            const float fy = nctr[1] - static_cast<float>(v);
            const float y = m11 * fy + m12 * fz;
            const float yz = y * y + z * z;
            if (yz > max_r2)
                continue;
            const float x_yz = m01 * fy + m02 * fz;
            const size_t row = static_cast<size_t>(nu) * (gemmi::modulo(v, nv) + static_cast<size_t>(nv) * w_);
            for (int u = c[0] - d[0]; u <= c[0] + d[0]; u++) {
                const float x = m00 * (nctr[0] - static_cast<float>(u)) + x_yz;
                if (yz + x * x > max_r2)
                    continue;
                const double value = map.data[row + gemmi::modulo(u, nu)];
                for (size_t q = 0; q < std::size(precal.a); q++) {
                    const gemmi::SMat33<float> &m = precal.b[q];
                    const gemmi::Vec3f r(x, y, z);
                    const double e = value * precal.a[q] * gemmi::unsafe_expapprox(std::max(m.r_u_r(r), -88.f));
                    const double mr[3] = {m.u11 * x + m.u12 * y + m.u13 * z, m.u12 * x + m.u22 * y + m.u23 * z,
                                          m.u13 * x + m.u23 * y + m.u33 * z};
                    for (int k = 0; k < 3; k++)
                        g[k] += e * mr[k];
                    gu[0] += e * (m.u11 + 2 * mr[0] * mr[0]);
                    gu[1] += e * (m.u22 + 2 * mr[1] * mr[1]);
                    gu[2] += e * (m.u33 + 2 * mr[2] * mr[2]);
                    gu[3] += e * (m.u12 + 2 * mr[0] * mr[1]);
                    gu[4] += e * (m.u13 + 2 * mr[0] * mr[2]);
                    gu[5] += e * (m.u23 + 2 * mr[1] * mr[2]);
                }
            }
        }
    }
    const double occ = atom.occ;
    grad = {2 * occ * g[0], 2 * occ * g[1], 2 * occ * g[2], occ * (gu[0] + gu[1] + gu[2]) / gemmi::u_to_b()};
    if (grad_u != nullptr)
        *grad_u = {occ * gu[0], occ * gu[1], occ * gu[2], 2 * occ * gu[3], 2 * occ * gu[4], 2 * occ * gu[5]};
}

}  // namespace

void AtomGradientsFromMap(const gemmi::Grid<float> &map, const gemmi::Model &model, double blur, size_t nthreads,
                          std::vector<std::array<double, 4>> &grad, std::vector<std::array<double, 6>> *grad_u) {
    const gemmi::DensityCalculator<Table, float> dc = Calculator(map.unit_cell, *map.spacegroup, 1.0, blur);
    std::vector<const gemmi::Atom *> atoms;
    for (const gemmi::Chain &ch : model.chains)
        for (const gemmi::Residue &r : ch.residues)
            for (const gemmi::Atom &a : r.atoms)
                atoms.push_back(&a);
    grad.assign(atoms.size(), {0, 0, 0, 0});
    if (grad_u != nullptr)
        grad_u->assign(atoms.size(), {0, 0, 0, 0, 0, 0});
    const int nu = map.nu, nv = map.nv, nw = map.nw;
    // As ModelGrid.cpp walks a density box: the orthogonalization over the grid size in single precision.
    const gemmi::UpperTriangularMat33 &om = map.orth_n;
    const float m00 = static_cast<float>(om.a11), m01 = static_cast<float>(om.a12), m02 = static_cast<float>(om.a13),
                m11 = static_cast<float>(om.a22), m12 = static_cast<float>(om.a23), m22 = static_cast<float>(om.a33);
    constexpr double FOUR_PI2 = 4 * refine_math::PI * refine_math::PI;
    ParallelChunks(static_cast<int>(atoms.size()), nthreads, [&](int lo, int hi) {
        for (int i = lo; i < hi; i++) {
            const gemmi::Atom &atom = *atoms[i];
            if (atom.aniso.nonzero()) {
                AnisoGradient(map, dc, atom, blur, grad[i], grad_u != nullptr ? &(*grad_u)[i] : nullptr);
                continue;
            }
            const float b = static_cast<float>(atom.b_iso + blur);
            const auto precal = IonFormFactor<float>(atom.element.elem, atom.charge).precalculate_density_iso(b, 0.0f);
            const double radius = dc.estimate_radius(precal, b);
            const float max_r2 = static_cast<float>(radius * radius);
            const gemmi::Fractional fpos = map.unit_cell.fractionalize(atom.pos);
            const double n3[3] = {fpos.x * nu, fpos.y * nv, fpos.z * nw};
            float nctr[3];
            int c[3], d[3];
            for (int k = 0; k < 3; k++) {
                nctr[k] = static_cast<float>(n3[k]);
                c[k] = gemmi::iround(n3[k]);
            }
            d[0] = std::min(static_cast<int>(std::ceil(radius / map.spacing[0])), nu - 1);
            d[1] = std::min(static_cast<int>(std::ceil(radius / map.spacing[1])), nv - 1);
            d[2] = std::min(static_cast<int>(std::ceil(radius / map.spacing[2])), nw - 1);
            double gx = 0, gy = 0, gz = 0, gb = 0;
            for (int w = c[2] - d[2]; w <= c[2] + d[2]; w++) {
                const float fz = nctr[2] - static_cast<float>(w);
                const float z = m22 * fz;
                const int w_ = gemmi::modulo(w, nw);
                for (int v = c[1] - d[1]; v <= c[1] + d[1]; v++) {
                    const float fy = nctr[1] - static_cast<float>(v);
                    const float y = m11 * fy + m12 * fz;
                    const float yz = y * y + z * z;
                    if (yz > max_r2)
                        continue;
                    const float x_yz = m01 * fy + m02 * fz;
                    const size_t row = static_cast<size_t>(nu) * (gemmi::modulo(v, nv) + static_cast<size_t>(nv) * w_);
                    for (int u = c[0] - d[0]; u <= c[0] + d[0]; u++) {
                        const float x = m00 * (nctr[0] - static_cast<float>(u)) + x_yz;
                        const float r2 = yz + x * x;
                        if (r2 > max_r2)
                            continue;
                        const double value = map.data[row + gemmi::modulo(u, nu)];
                        double dr2 = 0, db = 0;
                        for (size_t q = 0; q < std::size(precal.a); q++) {
                            const double e = precal.a[q] * gemmi::unsafe_expapprox(std::max(precal.b[q] * r2, -88.f));
                            dr2 += e * precal.b[q];
                            db += e * precal.b[q] * (1.5 + precal.b[q] * r2);
                        }
                        gx += value * dr2 * x;
                        gy += value * dr2 * y;
                        gz += value * dr2 * z;
                        gb += value * db;
                    }
                }
            }
            // rho = occ sum_q a_q exp(b_q r^2): d rho / d(centre) = occ 2 (sum a_q b_q e_q) (centre - point), and
            // with b_q = -4 pi^2 / (B_q + B + blur), d rho / dB = occ sum a_q e_q b_q (1.5 + b_q r^2) / (4 pi^2).
            grad[i] = {2.0 * atom.occ * gx, 2.0 * atom.occ * gy, 2.0 * atom.occ * gz, atom.occ * gb / FOUR_PI2};
        }
    });
}