Jungfraujoch/image_analysis/structure_refinement/Refinement.h

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

#pragma once

#include <array>
#include <atomic>
#include <cmath>
#include <complex>
#include <memory>
#include <optional>
#include <set>
#include <string>
#include <vector>

#include "gemmi/grid.hpp"
#include "gemmi/model.hpp"
#include "gemmi/symmetry.hpp"
#include "gemmi/unitcell.hpp"

#include "../../common/Reflection.h"   // MergedReflection
#include "DensityBlobs.h"
#include "ModelValidation.h"           // ModelMaps
#include "RefinementRestraints.h"      // GeometryStatistics
#include "SolventBuild.h"
#include "SolventSites.h"

class Logger;

// Refinement of an atomic model already placed against the data (ValidateAgainstModel's placed model, or a
// molecular-replacement placement), in two stages:
//
// 1. Placement: each chain as a rigid body, then one B shift per residue, each kept only where it lowers
//    R-free. Needs no dictionary.
// 2. Restrained: every atom's position and isotropic B, against the maximum-likelihood amplitude target with
//    the geometry the dictionary describes (MonomerLibrary.h: bonds, angles, planes, chiral volumes,
//    torsions, repulsions; B similarity of bonded atoms and of atoms in contact, in proportion to B). Residues
//    it does not describe are held fixed and stay in the structure factors and the bulk-solvent mask. The
//    weights of the X-ray term against the geometry and against the B similarity, and whether to refine the B
//    factors, are chosen by R-free among trials whose geometry stays within bounds: a ladder from the placed
//    model, then a trajectory of steps that each try the weights as they are and changed, from where the last
//    step ended, until no step gains; then a short final ladder.
//
// The library reads and writes nothing: the model and the reflections come in, the refined model, the
// statistics and the map coefficients go out. Nothing it finds feeds back into processing decisions.
struct RefinementSettings {
    bool restrained = true;          // stage 2; false stops after the placement stage
    bool score_only = false;         // only score the model as given (RefinementResult::score_start), no stage
    bool group_b = true;             // the placement stage's B shift per residue
    int macro_cycles = 2;            // per weight trial, from the model the trial starts from
    int xyz_iterations = 30;         // L-BFGS iterations per macro-cycle
    int b_iterations = 20;
    // The first trials, from the placed model: the positions' X-ray weight, times the gradient-balanced one
    // (BalancedWeight), each with the B factors' X-ray weight their own balanced one times each of adp_ladder (0: B
    // held).
    std::vector<double> weight_ladder = {1.0 / 1024, 1.0 / 128, 1.0 / 16};
    std::vector<double> adp_ladder = {0.0, 1.0};
    // Then at most this many steps of the trajectory, each a set of trials from where the last one ended,
    int max_steps = 8;
    // and the final trials from where it ended, at its weights times each of these;
    std::vector<double> final_ladder = {1.0, 2.0};
    double max_weight = 1.0 / 8;     // the positions' X-ray weight at most, times the balanced one
    bool riding_hydrogens = true;    // the dictionary's hydrogens added, riding on their parents
    // Where the model has copies of a molecule (RefinementNcs.h), the trials try them restrained to each other, at
    // NCS restraint sigmas from ncs_sigma, beside the trials without; R-free decides.
    bool ncs = true;
    double ncs_sigma = 0.1;          // A
    // Waters added at the difference map's peaks and removed where the maps no longer hold them, and waters and ions
    // relabelled, within the weight trajectory (SolventBuild.h), to SOLVENT_BUILD_D_MIN, and the refined model's
    // solvent audited (SolventSites.h); off (the default, rugnux --refine-solvent turns it on) the solvent stays
    // as given and is not audited.
    bool solvent = false;
    // After the weights are chosen, the displacement model decided on the free set: the isotropic B refined on, TLS
    // groups (the model's own, one per chain, the chains cut into TLSMD segments) with a residual B per atom, and an
    // anisotropic U per atom with rigid-bond and similarity restraints (false: the isotropic B, or the model's U held).
    bool adp_hypotheses = true;
    // A refinement against the data as twinned by each twin law the lattice admits (RefinementTwinLaws), beside the
    // untwinned one, and the model of whichever the evidence favours (RefineModel).
    bool twin = true;
    // The data's own intensity statistics read twinned - the caller's verdict, from the merge (the L-test): the law
    // is then taken from 2 standard errors over its null instead of 3 (AdoptTwinLaw).
    bool twin_data_indicated = false;
    // The model as given a candidate of the final choice (the restrained stage, 5.): only where its R-free is a free
    // statistic - false where the model has most likely been refined against the free reflections too (a deposited
    // model of this crystal scored on a free set drawn here).
    bool as_given_competes = true;
    bool gpu = true;
    size_t nthreads = 1;
    size_t gpu_resident_bytes = 0;   // held on the card by what runs beside
    // The most weight trials run side by side, 0 for as many as the ladders and the cards take; the result does not
    // depend on it (1 runs them one after another, as a check of that).
    size_t max_trials_at_once = 0;
    // Set by the caller (from any thread) to cut a running refinement short: checked between the minimiser's
    // iterations and before every weight trial. The trials that ran are chosen among as usual, those that did not
    // are left out (RefinementResult::cancelled); cut before the first trial, the model is the placement's.
    const std::atomic<bool> *cancel = nullptr;
};

// R by resolution: one of RefinementScore's shells, equal in width in 1/d^2 between the coarsest and the
// finest reflection.
struct RefinementShell {
    double d_max = 0, d_min = 0;
    int n_work = 0, n_free = 0;
    double r_work = NAN, r_free = NAN;
};

// What scoring a model against the data gives, every term of it: the model's structure factors are
// F_model = k_overall exp(-s^T B s / 4) (Fcalc + k_sol exp(-b_sol s^2 / 4) Fmask), the scale and the solvent
// fitted to the working reflections; R-work and R-free; and the likelihood target, -ln L summed over the
// working and over the free reflections, at D and the variance estimated from the free set.
struct RefinementScore {
    double k_overall = 0;
    double b_iso = 0;                // the overall B: a third of the trace of B (Cartesian, A^2)
    std::array<double, 6> b_aniso{}; // B less b_iso: B11 B22 B33 B12 B13 B23 (Cartesian, A^2)
    double k_sol = 0, b_sol = 0;
    double mask_spacing = 0;         // A, the grid the solvent mask is drawn on
    double r_work = NAN, r_free = NAN;
    int n_work = 0, n_free = 0;
    double nll_work = NAN, nll_free = NAN;
    double fom = NAN;                // the likelihood's mean figure of merit, working reflections
    double twin_fraction = 0;        // against a twin law: the fraction fitted with the scale
    int hydrogens = 0;               // riding hydrogens in the model scored
    std::vector<RefinementShell> shells;
    // Per reflection, with settings.score_only only: the atoms' Fcalc, the solvent's k_sol exp(-b_sol s^2/4) Fmask
    // and k_overall exp(-s^T B s / 4), in the order of `hkl`, and the likelihood's parameters.
    std::vector<gemmi::Miller> hkl;
    std::vector<std::complex<float>> fcalc, fmask;
    std::vector<float> k_total;
    std::vector<float> d, var;       // the likelihood's D and variance (epsilon beta + sigma^2)
};

// One rung of the weight ladder.
struct RefinementWeightTrial {
    double weight = 0;               // the X-ray term's weight
    double adp_weight = 0;           // the B refinement's X-ray weight; 0: the B factors were held
    double ncs_sigma = 0;            // the NCS restraints' sigma (A); 0: none
    double ncs_rmsd = NAN;           // the copies' C-alpha (P) rmsd after superposition, where there are copies
    double r_work = 1, r_free = 1;
    double nll_free = NAN;           // -ln L of the free reflections
    GeometryStatistics geometry;
    bool within_bounds = false;      // the geometry bounds the choice is made inside
    std::string adp;                 // the displacement model: "isotropic", "TLS ..." or "anisotropic"
};

struct RefinementResult {
    bool ok = false;
    std::string failure_reason;
    bool on_gpu = false;
    double seconds = 0;
    double d_min = 0;
    int n_work = 0, n_free = 0;

    // The model as given, scaled to these reflections (B made isotropic; the dictionary's hydrogens riding
    // where the restrained stage runs or the model is only scored), and every term of that score.
    double r_work_start = 1, r_free_start = 1;
    RefinementScore score_start;
    // How the model's displacements were read: "isotropic", "ANISOU" (the atoms' own U), or "TLS, B total" / "TLS, B
    // residual" (U from the model's TLS groups, its B column taken as the total B or as the residual to them) - the
    // reading of the TLS groups the model's own R-free chose among these and the isotropic B.
    std::string adp_input = "isotropic";
    int adp_input_aniso_atoms = 0;   // the atoms the model gave a U (its own or by its TLS groups)
    // After the placement stage.
    int rigid_bodies = 0;            // the chains moved as rigid bodies, where that was kept
    bool group_b = false;            // the B shift per residue was kept
    double r_work_placement = 1, r_free_placement = 1;
    // After the restrained stage (NaN where it did not run).
    bool restrained = false;
    bool as_given = false;           // the model as given kept: no trial displaced it (the restrained stage ran)
    bool cancelled = false;          // cut short by RefinementSettings::cancel
    double r_work = NAN, r_free = NAN;
    double weight = 0;
    double adp_weight = 0;           // 0: the B factors kept as placed
    std::vector<RefinementWeightTrial> trials;
    // Non-crystallographic symmetry: the molecules present more than once, their chains, the restraints between the
    // copies, and the sigma chosen (0: the copies refined independently); the copies' C-alpha (P) rmsd after
    // superposition, averaged over the pairs restrained, as given and refined; the residues whose restraints let go.
    int ncs_groups = 0, ncs_copies = 0, ncs_restraints = 0;
    std::string ncs_chains;
    double ncs_sigma = 0;
    double ncs_rmsd_start = NAN, ncs_rmsd = NAN;
    int ncs_released = 0;
    // The displacement model kept (RefinementSettings::adp_hypotheses): "isotropic", "TLS, <n> groups (...)" or
    // "anisotropic", and the trials it was chosen among; the atoms with a U of their own and the TLS groups written.
    std::string adp_model = "isotropic";
    std::vector<RefinementWeightTrial> adp_trials;
    int aniso_atoms = 0, tls_groups = 0;
    GeometryStatistics geometry_start, geometry;

    int atoms = 0, refined_atoms = 0;
    int restrained_residues = 0, fixed_residues = 0;
    std::vector<std::string> fixed_names;
    double k_sol = 0, b_sol = 0;
    double mean_fom = 0;

    std::shared_ptr<gemmi::Structure> model;   // the refined model, in the data's cell and frame
    std::shared_ptr<ModelMaps> maps;           // its 2mFo-DFc and mFo-DFc coefficients (as _maps.mtz holds them)
    // The 2mFo-DFc map of the refined model and its rms, for reading the solvent sites in (SolventSites.h); released
    // once they are read.
    std::shared_ptr<gemmi::Grid<float>> map;
    double map_rms = 0;
    // The mFo-DFc map of the refined model and its rms, for finding what it leaves unexplained (DensityBlobs.h).
    std::shared_ptr<gemmi::Grid<float>> difference_map;
    double difference_map_rms = 0;
    // The anomalous difference map on the refined model's phases, where the merge kept the Bijvoet pairs.
    std::shared_ptr<gemmi::Grid<float>> anomalous_map;
    double anomalous_map_rms = 0;
    // What the refined model leaves unexplained in the difference map (FindDensityBlobs).
    std::vector<DensityBlob> blobs;
    // The refined model's waters and ions and the species their surroundings support (AnalyseSolventSites): an
    // audit, nothing in the model is changed by it.
    std::optional<std::vector<SolventSite>> solvent;   // the solvent audit (settings.solvent)
    double b_ceiling = 0;                      // the highest B the refinement allowed
    SolventBuildResult solvent_build;          // the solvent build's rounds (settings.solvent)

    // Twin refinement (settings.twin): what was decided (twin_decision: NO_TWIN_LAW, FREE_SET_SPLIT - the free
    // set does not keep the law's twin pairs together -, ZERO_FRACTION - the model asks for none -, UNTWINNED,
    // TWINNED), the law tested and the fraction it refined to, and the evidence: the R-free of the untwinned and
    // the twinned refinement, and the null of the twinned one - its final model at its fraction with every mate
    // replaced by an unrelated reflection of the same resolution - the R-free that averaging alone buys,
    // and the law's lead over it on the free set in standard errors.
    // `twinned`: the model and every number above it are the twinned refinement's.
    std::string twin_decision;
    bool twinned = false;
    std::string twin_law;                      // hkl triplet, in the data's basis
    double twin_fraction = NAN;
    int twin_laws = 0;                         // the laws the lattice admits
    int twin_split_pairs = 0;                  // the tested law's twin pairs split between work and free
    double twin_r_work = NAN, twin_r_free = NAN;
    double untwinned_r_work = NAN, untwinned_r_free = NAN;
    double twin_r_free_null = NAN, twin_fraction_null = NAN;
    double twin_null_z = NAN;                  // the law's lead over the null on the free set, in standard errors
    // Where tested to UNTWINNED or TWINNED, the rule that decided: NULL (the null alone, 3 standard errors), or
    // NULL+LTEST (taken at 2-3 standard errors because the data read twinned, settings.twin_data_indicated).
    std::string twin_rule;
    // The domains beyond two (DecideTwin's second rung): every other domain of the lattice added to the adopted law,
    // its decision (MULTI_TWINNED where the model is that refinement's, MULTI_UNTWINNED, or empty: not tested), the
    // laws and their fractions (law=fraction, ';'-separated), its R-free, and the null of the domains added - their
    // mates unrelated, the adopted law's kept - with the domains' lead over it in standard errors.
    std::string twin_multi_decision;
    std::string twin_multi_domains;
    double twin_multi_r_work = NAN, twin_multi_r_free = NAN, twin_multi_r_free_null = NAN, twin_multi_null_z = NAN;
    int twin_free_closed = 0;                  // reflections moved between work and free to close the free set
                                               // under the lattice's domains (CloseFreeSet)
    std::vector<double> twin_domain_fractions; // per law of the refinement (twin_law's order), as refined
};

// The free set of `merged` closed under the lattice's domain operators (RefinementTwinDomains): the reflections that
// any domain relates - an orbit of the lattice's holohedry, as far as the data have it - all free or all working, so
// that no twin law splits a twin group between the working and the free set. An orbit the flags already keep in one
// set stays there; a split one goes free with the probability of its own share of free flags, drawn by a hash of the
// orbit, so the free fraction stays as it was (a majority would empty the free set where the flags were drawn on
// smaller orbits). Returns how many reflections changed set - none for a free set drawn on the lattice's holohedry
// (AssignRfreeFlags), some for one a file brings or one keyed on another cell.
// As REFMAC5 keeps twin-related reflections in one set (Roversi et al. (2012) Acta Cryst. D68, 418-424)
int CloseFreeSet(std::vector<MergedReflection> &merged, const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg);

// Whether the twinned refinement is taken (RefineModel): a fraction above 0, and the law ahead of its null on the
// free set by 3 standard errors - or by 2 where the data's own intensity statistics read twinned as well.
bool AdoptTwinLaw(double twin_fraction, double null_z, bool data_indicated);

// `placed`: the model in the data's cell and space group. With settings.solvent (and restrained, not score_only),
// the weight trajectory is followed by up to SOLVENT_BUILD_ROUNDS rounds of the solvent build, each the trajectory
// again from the model with its solvent changed, kept where it ends at a lower R-free (Refinement.cpp). `merged`: the
// merge, with French-Wilson amplitudes and free flags, in the same group.
RefinementResult RefineModel(const gemmi::Structure &placed, const std::vector<MergedReflection> &merged,
                             const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg, double wavelength_A,
                             Logger &logger, const RefinementSettings &settings);