Jungfraujoch/image_analysis/structure_refinement/ModelValidation.h

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

#pragma once

#include <cmath>
#include <functional>
#include <future>
#include <memory>
#include <optional>
#include <string>
#include <vector>

#include "gemmi/symmetry.hpp"

#include "../../common/Reflection.h"   // MergedReflection
#include "../../common/UnitCell.h"

class Logger;
namespace gemmi { struct Structure; }

// Result of validating merged intensities against an atomic model.
struct ModelValidationResult {
    bool ok = false;
    std::string model_path;        // the file that was asked for, so the report can name it
    // Why it did not run, when ok is false: the file could not be read, had no atoms, no cell or no
    // usable space group. A caller that asked for a model must be able to say so rather than finish
    // quietly with no R-free and no maps, which reads exactly like a run that was never asked.
    std::string failure_reason;
    double r_work = 0.0;
    double r_free = 0.0;
    double k_sol = 0.0;
    double b_sol = 0.0;
    double k_overall = 0.0;
    int n_work = 0;
    int n_free = 0;

    // R over ALL the reflections - r_model as the maps are scaled, r_model_shell_scaled with one free
    // scale per resolution shell instead - and how far those per-shell scales had to depart from a
    // flat one. The pair exists to be read BETWEEN runs, where r_work / r_free cannot be.
    //
    // The scale the maps carry can only bend as k*exp(-B s^2), deliberately (see ModelValidation.cpp),
    // so where two reductions of one crystal differ in the radial profile of their amplitudes, the
    // part of the difference that shape cannot follow is reported as R - by more, measured, than a
    // real change in the data moves it. A per-shell scale removes the radial profile and only the
    // radial profile, leaving the agreement inside each shell.
    //
    // All the reflections and not the free set, because nothing here is refined: the coordinates, the
    // B factors and the occupancies are the model's own, so work and free estimate the same quantity
    // and the 5% split buys no cross-validation while costing a factor of sqrt(20) in precision.
    double r_model = NAN;
    double r_model_shell_scaled = NAN;
    int n_model = 0;
    // The RMS of ln(k_shell) about its own mean. Near zero says k*exp(-B s^2) already described this
    // dataset's radial profile, and then the two R-factors agree; a large value says it did not, and
    // is worth seeing on its own - it is the gap between the two R-factors, named.
    double radial_misfit = NAN;
    int shell_scale_bins = 0;
    double mean_atom_density_sigma = 0.0;  // mean 2mFo-DFc value at atom centres, in sigma
    std::string maps_prefix;               // where the .ccp4 / _maps.mtz were written

    // The rigid-body placement of the model against these data, where it was committed. The step is
    // taken only when it lowers R-free, so an unmoved model means the data said no, not that nothing
    // was tried; r_free_before_rigid_body is what it would have been without it.
    bool rigid_body_applied = false;
    double rigid_body_angle_deg = 0.0;
    double rigid_body_shift_A = 0.0;
    double r_free_before_rigid_body = 0.0;

    // The sigma_A weighting the map coefficients carry.
    double mean_fom = 1.0;
    int sigma_a_shells = 0;

    // CC(model, data) by resolution shell: the Pearson correlation between the merged intensities and
    // the placed, scaled model's |F_model|^2, over every reflection of the shell. Intensities and not
    // amplitudes, so that the row can be read straight across from the merge table's CC1/2 and CCref,
    // which are both correlations of intensities; it is also how Karplus & Diederichs (2012) define
    // CC_work, as the correlation of F^2_calc with F^2_obs. The observed value is the merged intensity
    // itself rather than the French-Wilson |F|^2 the R-factors use: the French-Wilson amplitude is a
    // posterior mean under a Wilson prior, which pulls a weak reflection towards its shell mean, and
    // that is a correlation with the prior in exactly the outer shells this statistic is read in.
    //
    // The model is placed and scaled with a handful of parameters, so there is no free/work
    // distinction to make here and the correlation is formed on all the reflections: the free set
    // alone in an outer shell is a few hundred reflections and the correlation on it correspondingly
    // noisy.
    //
    // The test it supports is ONE-SIDED. A correlation significantly above its null in a shell is
    // evidence that the shell carries signal, so it argues that a resolution limit could be pushed
    // OUTWARDS. A correlation near zero proves nothing at all - the model may be the thing at fault -
    // and must never pull a limit in. "Significantly" is the hard part: the data and a scaled model
    // agree in everything smooth in reciprocal space - the falloff across a shell, the anisotropic B
    // that was fitted to these data, the bulk solvent - whatever the model is, so a model of an
    // unrelated protein correlates with them far beyond what independent reflections allow. The
    // null is made from the same data and model for that reason (see sigma below).
    struct ModelDataCCShell {
        float d_min = 0.0f;   // the shell's high-resolution bound, as the merge table labels its rows
        double cc = NAN;
        int n = 0;            // reflections the correlation was formed from
        // How far the correlation sits above its null, in Fisher-z: the null is the correlation of the
        // data at h with the model at h + d for the 26 nearest lattice offsets d, which keeps
        // everything the two share smoothly in reciprocal space and loses what the model knows of h
        // itself. sigma = (atanh(cc) - mean) / sd over those 26, the sd never narrower than Fisher's
        // 1/sqrt(n-3). A guide to the strength of the claim rather than an exact p.
        double sigma = NAN;
        // The correlation those offset pairs reach on average: what an unrelated model scaled to these
        // data would show in this shell.
        double cc_null = NAN;
        // The shell's mean merged intensity over the mean |F_model|^2 of the SCALED model - the model
        // having already been given an overall scale, an anisotropic B and a bulk solvent fitted over
        // the whole range. That family is smooth and nearly exponential in s^2, so it can follow a
        // wrong RADIAL profile only so far; what is left here is the part of the run's resolution
        // dependence the model does not agree with, and it is the only number in the report that can
        // see it. Every internal statistic - R_meas, CC1/2, chi^2, the correction surfaces' own
        // cross-validation - is blind to a factor that multiplies a whole ASU group, because each
        // group's reference is its own mean. 1.0 across the table means the two agree; a monotone
        // run away from 1.0 means one of the two radial profiles is wrong. It cannot say which: a
        // real crystal can out-diffract or under-diffract its deposited model.
        double i_over_model = NAN;
    };
    std::vector<ModelDataCCShell> cc_model_shells;
    double cc_model_overall = NAN;
    int cc_model_n = 0;

    // The anomalous difference map read at the model's own atoms: the strongest sites, highest
    // first. Empty when the merge kept no Bijvoet split, and so had nothing to make the map from.
    struct AnomalousSite {
        std::string label;   // atom, residue and chain, e.g. "SG CYS A7"
        double sigma = 0.0;  // height of the anomalous map at that atom, in map sigma
    };
    std::vector<AnomalousSite> anomalous_sites;
    int anomalous_pairs = 0;               // Bijvoet pairs the anomalous map was computed from
    // The same map's mean height at the model's anomalous scatterers - every atom from phosphorus up
    // (S of Met/Cys, metals, Cl, I, ...) - one number for how much anomalous signal the data carry.
    int anomalous_scatterers = 0;
    double anomalous_scatterer_mean_sigma = NAN;

    // Set when the anomalous map is inverted at the model's atoms - deeper troughs than peaks - which
    // says the data and the model are in opposite hands. Named so the report can quote the evidence.
    bool anomalous_hands_disagree = false;
    std::string anomalous_deepest_site;
    double anomalous_deepest_sigma = 0.0;

    // Does the model describe these data at all? A model is a hypothesis, and an R-factor on its own
    // cannot test it: what a model that explains nothing gives against these data depends on its atom
    // count and its B-factors as much as on the data, so no fixed threshold on R separates the two.
    // The same model is therefore refitted from `null_replicates` random orientations about its own
    // centroid, by the identical path, and the real fit is asked how far above that distribution it
    // sits. R-work rather than R-free: nothing is refined here, so R-work carries no optimism, and it
    // is decided on an order of magnitude more reflections.
    //
    // The null is run on every validation that judges a model (ModelValidationSchedule::fit_test):
    // model_fits also says whether the rigid body placed the model at all. fit_tested is false only
    // where the null could not be built, or where the verdict is another validation's.
    bool fit_tested = false;
    bool model_fits = false;
    double null_r_work_mean = 0.0;
    double null_r_work_sd = 0.0;
    double r_work_sigma = 0.0;     // (null mean - r_work) / null sd; positive = better than random
    int null_replicates = 0;

    // The change of basis the model was put through to reach the data's description of the same
    // lattice - an axis permutation, a sign flip, I-centred against C-centred. Identity where the two
    // already agreed, which is the ordinary case. setting_as_read is the setting the model arrived
    // in, kept only when it changed, so the report can say what was done rather than only that
    // something was.
    gemmi::Op change_of_basis_op = gemmi::Op::identity();
    std::string setting_as_read;

    // The alternative-indexing operator picked by R-free, identity where none was needed or where the
    // choice did not beat its own null. AdoptModelFrame below applies it to the reflections that are
    // written out, so the file, the R-factors and the maps all describe one indexing.
    gemmi::Op indexing_op = gemmi::Op::identity();
    // The margin the choice was made on - the runner-up's R-free minus the winner's - beside what the
    // same margin comes out at for a model in a random orientation. A winner that leads by no more
    // than chance has decided nothing, and the data keep the indexing they were merged in.
    bool indexing_probed = false;
    // The probe's winner beat its own null on a model that fits: the indexing ambiguity is settled by
    // the model, whether that kept the data's indexing or moved it to indexing_op.
    bool indexing_decided = false;
    double indexing_margin = 0.0;
    double indexing_margin_null_mean = 0.0;
    double indexing_margin_null_sd = 0.0;
    double indexing_margin_sigma = 0.0;

    // The enantiomorph the model asserts, read off the two group numbers alone and so known before
    // anything has been fitted.
    bool model_enantiomorph_candidate = false;
    // Set when that assertion is taken up for the written reflections, which needs the model to have
    // fitted and the anomalous map not to have vetoed it. This is a change of LABEL only: the two
    // groups have the same rotation operations, so they index and split into Bijvoet hands
    // identically, and the reflections themselves must not move. See the note in ValidateAgainstModel.
    bool adopted_model_enantiomorph = false;
    int model_space_group_number = 0;

    // Against the data as twinned by the twin law ValidateAgainstModel was given (none given: not
    // made): R-work and R-free with |F_twin|^2 = (1-a)|F_model(h)|^2 + a|F_model(Th)|^2 and one overall
    // scale, at the fraction a (0 to 0.5) that gives the lowest R-work.
    bool twin_target = false;
    double twin_fraction = NAN;
    double r_work_twin = NAN;
    double r_free_twin = NAN;

    // The model as it was scored: re-fractionalized into the data cell and moved by the rigid-body
    // step, everything else exactly as it was read. Null where the validation did not run. The frame
    // it is written in is only settled by AdoptModelFrame below, so it is the caller that puts it on
    // disk, with WritePlacedModel (writer/WriteModel.h).
    std::shared_ptr<gemmi::Structure> placed_model;
};

// Given merged intensities and an atomic model (PDB or mmCIF, gzipped or not - the format is taken
// from the file's content), scale the model structure factors (with an
// optimized flat bulk-solvent contribution and an overall anisotropic B) to the WORKING observed
// amplitudes, then report R-work / R-free and write 2Fo-Fc and Fo-Fc electron-density maps
// (CCP4) plus an MTZ of map coefficients next to output_prefix. Nothing the reported R-free is
// computed against was fitted on the free reflections: the scale, like the rigid-body placement
// below, sees the working set only and is then applied to all of them.
//
// The structure itself is not refined. The model is re-fractionalized into the data unit cell and
// then placed as ONE RIGID BODY - an angle-axis rotation about its centroid plus a translation,
// refined over a coarse-to-fine resolution ladder and kept only where it lowers R-free - so a
// deposited model from a non-isomorphous crystal still sits where the density is. The map
// coefficients are sigma_A weighted (2mFo-DFc and mFo-DFc). Returns ok=false (and logs) on any
// failure.
// data_space_group is the space group the data was merged in. If it is the enantiomorph of
// the model's space group (e.g. data P4(1)2(1)2, model P4(3)2(1)2 - indistinguishable from merged
// intensities), the model's group is adopted as a label, without touching the reflections - but only
// where the model was shown to fit these data in the first place. A model is a hypothesis: it is
// always scored, and its R-factors and maps are always reported, but it decides nothing about the
// written reflections until it has beaten a null made of itself in random orientations.
//
// probe_indexing_ambiguity controls how a merohedral (alternative-indexing) ambiguity is resolved.
// Pass false when a reference (MTZ) already fixed the indexing - the data are then kept in that
// authoritative indexing. Pass true (model only, no reference) to resolve it here as a fallback, by
// fitting each candidate reindexing to the model and keeping the lowest R-free - where its lead over
// the runner-up beats the lead the same model in a random orientation takes, since a random model
// also picks a winner. Holohedral crystals (no twin laws) are unaffected either way.
//
// nthreads is what the null's replicates and the candidate changes of basis run on - each is
// independent of the others, so they run at once. The answer does not depend on it.
//
// twin_law, where given, is a twin law of the data's group (a Miller-index operator in the data's
// basis): the model is then also scored against the data as twinned by it (twin_target and below).
//
// model, where given, is the model itself - one already placed in the data's cell, by molecular
// replacement - and model_path only names it; no other setting is then probed for it.
//
// report_shell_d_min is the merge statistics' own shell bounds, coarse to fine, and it is what
// cc_model_shells is binned on. Sharing the grid is the point: a reader has to be able to put a
// CC(model, data) row beside that shell's CC1/2 and know the two describe the same reflections.
// Empty (the default) means no shells were given and none are reported.
// What a validation knows of its outcome before it has one: the setting the model was put into and
// the indexing the probe prefers - which is what it decides where the model turns out to fit and the
// probe's lead turns out to be real. Enough for a caller to start what follows on that forecast and
// keep it only where the outcome agrees.
struct ModelFrameForecast {
    gemmi::Op change_of_basis_op = gemmi::Op::identity();
    gemmi::Op indexing_op = gemmi::Op::identity();
    bool enantiomorph_candidate = false;
    // Whether a second validation of this size may run beside this one (ModelValidationGPUPlan), and
    // what this one holds on its card meanwhile - which the second is given as gpu_beside_bytes.
    bool may_run_beside = true;
    size_t gpu_bytes = 0;
};

struct ModelValidationSchedule {
    // Called once, before the null, with the forecast above.
    std::function<void(const ModelFrameForecast &)> on_forecast;
    // Where valid, waited on before the first file is written: false returns without writing any.
    std::shared_future<bool> write_gate;
    // Asked once the verdict, the setting and the indexing are decided, before any map is made: true
    // where another validation will write this one's maps and map coefficients, which are then not made.
    std::function<bool(const ModelValidationResult &)> maps_superseded;
    // The null that says whether the model fits. False where the verdict is another validation's
    // (KeepModelVerdict) or was reached elsewhere (a molecular-replacement placement, scored for its maps).
    bool fit_test = true;
    // GPU memory held on the card for as long as the validation runs by what came before it (a merge
    // engine still at work beside it): counted against the card when what goes on the GPU is decided.
    size_t gpu_resident_bytes = 0;
    // GPU memory another validation running beside this one holds there. It limits how many engines
    // this one takes and never what runs on the GPU, so no number depends on whether the other runs.
    size_t gpu_beside_bytes = 0;
};

// What a model validation needs of its card, part by part: the bytes of one engine of each, how many
// engines of it the validation can use, and what else is on the card. 0 bytes: that part cannot run on
// the GPU (the cell is too small for its gridding), or there is none (the null's structure factors
// where the null is scored to the data's own resolution and shares the engine to d_min).
struct ModelValidationGPUNeed {
    size_t card = 0;                     // the card's total memory
    size_t resident = 0;                 // ModelValidationSchedule::gpu_resident_bytes
    size_t beside = 0;                   // ModelValidationSchedule::gpu_beside_bytes
    size_t structure_factors = 0;        // to the data's resolution, with the maps
    size_t null_structure_factors = 0;   // to the null's coarser limit
    size_t rigid_body = 0;
    size_t rigid_body_engines = 0;
    size_t scale_fit = 0;
    size_t scale_fit_engines = 0;
};

struct ModelValidationGPUPlan {
    size_t budget = 0;                   // 80% of the card, less what is resident there
    bool structure_factors = false;
    bool null_structure_factors = false;
    size_t rigid_body_engines = 0;       // 0: the rigid body on the CPU
    size_t scale_fit_engines = 0;        // 0: the scale fits on the CPU
    size_t bytes = 0;                    // what the plan holds on the card
    bool may_run_beside = false;
};

// What model validation and molecular replacement may hold on a card: 80% of its TOTAL memory - never what
// is free at the moment, so the same input on the same card always runs the same way - less what is
// resident there already (a merge engine still at work beside them).
size_t ModelGPUBudget(size_t card, size_t resident);

// The one rule for what a validation puts on its GPU, decided before anything is reserved, from the
// card's TOTAL memory and never from what is free at the moment, so the same input on the same card
// always runs the same way. Everything the validation holds is alive at once - the engines are reserved
// up front and kept to the end - so it is their sum that is held against the budget, 80% of the card
// less what is resident. Each part takes one engine where it still fits, in the order of what it saves:
// the structure factors to d_min, the rigid body, the null's structure factors, the scale fits; a part
// that does not fit runs on the CPU. More rigid-body engines and then more scale-fit engines are added
// only within half of the budget, less what a validation beside holds: they change only how many fits
// run at once, and buy little - on a large cell a fit already fills the card (measured on a 16 GB card:
// 13.6 s with one rigid-body engine, 13.1 s with two, 12.6 s with five) - so the card keeps headroom
// rather than filling with them. A second validation may run beside where twice the whole of this one's GPU parts, at every
// engine it can use, fits the budgets of the cards together.
ModelValidationGPUPlan PlanModelValidationGPU(const ModelValidationGPUNeed &need, size_t cards = 1);

ModelValidationResult ValidateAgainstModel(const std::vector<MergedReflection> &merged,
                                           const UnitCell &cell,
                                           const std::string &model_path,
                                           const std::string &output_prefix,
                                           Logger &logger,
                                           const gemmi::SpaceGroup *data_space_group = nullptr,
                                           bool probe_indexing_ambiguity = true,
                                           size_t nthreads = 1,
                                           double wavelength_A = 0.0,
                                           const std::vector<float> &report_shell_d_min = {},
                                           const ModelValidationSchedule &schedule = {},
                                           const gemmi::Op *twin_law = nullptr,
                                           const gemmi::Structure *model = nullptr);

// Reindex `merged` into the frame ValidateAgainstModel reported, so the reflection files that are
// written describe the same indexing as the R-factors and the maps. Returns the space group they are
// now in - the model's, where its enantiomorph was adopted, and `data_space_group` otherwise
// (a change of indexing is metric- and group-preserving). A no-op where the model needed neither,
// and a no-op throughout where the model did not fit.
const gemmi::SpaceGroup *AdoptModelFrame(const ModelValidationResult &validation,
                                         std::vector<MergedReflection> &merged,
                                         const gemmi::SpaceGroup &data_space_group,
                                         bool merge_friedel,
                                         Logger &logger);

// The decisions `decided` took - whether the model fits, against what null, the indexing it chose and
// the enantiomorph it asserted - copied onto `remade`, a validation made afterwards on the same data
// relabelled into the model's own setting. That one faces a model already in the data's setting and
// asserts nothing, so it decides nothing; its R-factors, maps and placed model are the ones kept,
// because they are on the axes the files are written on.
void KeepModelVerdict(ModelValidationResult &remade, const ModelValidationResult &decided);

// The anomalous difference map of ValidateAgainstModel, and nothing else, for a model that is already
// placed and scored in the data's cell (`placed` carries the cell and the model's group; k_sol and b_sol
// are the bulk solvent it was scored with): F_calc and F_mask to the data's resolution, the map
// F(+) - F(-) on the phases of F_calc + solvent, and its reading at the model's atoms - anomalous_pairs,
// the scatterers' mean, the sites and anomalous_hands_disagree, by the same code and the same rule as the
// validation's. No placement, no scale fit, no null, no file. For a caller that has already placed the
// model and judged its fit (the --mr hand check). On the GPU where the structure factors fit the budget
// (ModelGPUBudget) beside `gpu_resident_bytes`.
ModelValidationResult ReadModelAnomalousMap(const std::vector<MergedReflection> &merged,
                                            const gemmi::Structure &placed, const gemmi::SpaceGroup &data_space_group,
                                            double k_sol, double b_sol, Logger &logger, size_t nthreads,
                                            size_t gpu_resident_bytes = 0);

// Reference intensities computed from an atomic model: |F_model|^2, with a flat bulk-solvent
// contribution at the standard constants, for every reflection to `d_min`. This is the same
// information a reference MTZ carries, and it is what lets a model break the indexing ambiguity of a
// serial-stills run: there the ambiguity has to be broken per image, at integration time, long before
// there is a merge to fit a model to. `cell` and `space_group` override the model's own, where
// the run already knows them. Empty on any failure (which is logged).
std::vector<MergedReflection> ModelReferenceIntensities(const std::string &model_path,
                                                        const std::optional<UnitCell> &cell,
                                                        const gemmi::SpaceGroup *space_group,
                                                        double d_min,
                                                        Logger &logger);