Jungfraujoch/image_analysis/structure_refinement/ModelScaling.h

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

#pragma once

#include <condition_variable>
#include <cstddef>
#include <memory>
#include <mutex>
#include <vector>

#include "gemmi/model.hpp"
#include "gemmi/scaling.hpp"

// Fit the overall scale, the anisotropic B and the flat bulk solvent of a model to observed
// amplitudes, writing the answer into `scaling` (k_overall, b_star, k_sol, b_sol) so that
// gemmi's own scale_data() can then apply it.
//
// The caller decides which reflections are fitted: FitModelScale reads scaling.points, which
// prepare_points() filled. Pass the WORKING set to keep the free reflections out of the scale
// model - that is a property of the call site, not of this function.
//
// Replaces gemmi's fit_isotropic_b_approximately() + fit_parameters(), which is an unbounded
// Levenberg-Marquardt: nothing there stops k_sol and b_sol reaching values a flat solvent model
// is meaningless at (b_sol of 1707 A^2 was measured on this corpus). Here k_sol and b_sol come
// from a grid search over the physical ranges, so an unphysical answer cannot be represented.
// The box the solvent parameters are searched in. The defaults are the physically reasonable
// range reported by Fokine & Urzhumtsev and adopted as phenix's default grid: k_sol in (0.1, 0.8)
// and b_sol in (10, 80) A^2, of which phenix searches b_sol 10-80 in steps of 5. A model with high
// atomic B gets a higher b_sol ceiling (SolventBMax).
struct ModelScaleBox {
    double k_lo = 0.10, k_hi = 0.60;
    double b_lo = 10.0, b_hi = 80.0;
    // The overall scale and anisotropic B of the points of the fine pass by robust least squares
    // (ModelScaling.cpp, RobustReweight) rather than plain least squares on |F|: what refinement scores a model with.
    bool robust = false;
};

// The b_sol ceiling for `model`: 80 A^2, or twice the model's mean atomic B where that is higher.
// b_sol is how far the solvent's edge is smeared, and the edge can be no sharper than the atoms it
// lies against: a model at B of 300 A^2 beside a solvent held at 80 gets a solvent term that reaches
// far too deep in resolution, which only a too-small k_sol can then offset. Fokine & Urzhumtsev's
// 10-80 A^2 is for models at the B of ordinary entries (their mean b_sol 46 A^2), where the ceiling
// stays at 80.
double SolventBMax(const gemmi::Model &model);

struct ModelScaleReport {
    int n_points = 0;         // reflections the fit actually used
    int n_grid = 0;           // grid points evaluated
    double r_work_fit = 1.0;  // R on the fitted reflections, at the chosen solution
};

// nthreads: the grid points are fitted in parallel on up to this many threads; the answer does not
// depend on it.
ModelScaleReport FitModelScale(gemmi::Scaling<float> &scaling, ModelScaleBox box = {}, size_t nthreads = 1);

#ifdef JFJOCH_USE_CUDA
// FitModelScale on a GPU (ModelScaleGPU::FitSolvent): the same grid over the same box, the same fits and
// the same winner, with the sums over the reflections taken in double on the device where the host takes
// them in long double - so the answer is FitModelScale's to rounding, not to the bit. A few engines,
// reserved once, each a stream and room for one fit at a time; a fit takes a free one and waits for one
// where all are taken. Every kernel is deterministic, so which engine a fit gets changes no number.
class ModelScaleGPUPool {
public:
    // What `engines` engines for fits of up to `max_points` reflections take on the card.
    static size_t DeviceBytes(size_t max_points, size_t engines, bool robust = false);
    // On `device`. Throws JFJochException on a CUDA failure.
    ModelScaleGPUPool(int device, size_t max_points, size_t engines);
    ~ModelScaleGPUPool();

    // FitModelScale(scaling, {.b_hi = b_hi, .robust = robust}, nthreads). Where FitModelScale would walk the
    // grid point to point (five or fewer reflections for the isotropic start), it is FitModelScale on the host.
    // Safe to call from several threads. Throws JFJochException on a CUDA failure. The robust fit's reweighting
    // runs on the device; it is FitModelScale's robust fit to rounding (float residuals and weights, the
    // anisotropic factor of the steps in float), deterministic run to run.
    ModelScaleReport Fit(gemmi::Scaling<float> &scaling, double b_hi, size_t nthreads, bool robust = false);
    // What the pool holds on the card.
    size_t Bytes() const { return DeviceBytes(max_points_, engines_.size()); }

private:
    struct Engine;
    int device_;
    size_t max_points_;
    std::vector<std::unique_ptr<Engine>> engines_;
    std::vector<Engine *> idle_;
    std::mutex m_;
    std::condition_variable cv_;
};
#endif