Jungfraujoch/image_analysis/structure_refinement/MolecularReplacement.h

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

#pragma once

// Molecular replacement of ONE copy of a predicted model (rugnux --af-model, experimental): a global
// search for the model's orientation and position in the cell, so that a model with no placement of
// its own - an AlphaFold prediction - can be scored against the data, and the space groups the data
// cannot tell apart can be scored against it.
//
//  - The model: residues with pLDDT < 70 removed, B-factors from pLDDT, hydrogens dropped.
//  - Rotation: a uniform sampling of the point group's rotation ASU, each orientation scored by the
//    correlation of the observed E^2 with the model's own symmetry-summed |F|^2 - read off the
//    model's molecular transform, computed once by FFT and interpolated per orientation.
//  - Translation: the full-symmetry translation function on one FFT per orientation and space group.
//  - Rescoring: the best placements through the rigid body, bulk solvent, R-free, a packing check
//    against the symmetry mates and a null of the same model in random orientations.
//
// Everything up to the rescoring runs on the CPU, in both builds; the rigid body runs on the GPU
// where there is one.

#include <complex>
#include <memory>
#include <string>
#include <vector>

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

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

class Logger;

// The predicted model as the search uses it: trimmed, B-factors from pLDDT, centroid at the origin.
struct SearchModel {
    gemmi::Structure st;
    int residues_read = 0;
    int residues_kept = 0;
    double mass_da = 0;        // of the residues kept, 110 Da each
    double radius_A = 0;       // the farthest atom from the centroid
    std::string failure_reason;
};
SearchModel PrepareSearchModel(const std::string &path);

// Rotations covering SO(3) with neighbours about `step_deg` apart: the identity, then shells of
// rotation angle kappa = step, 2 step, ... 180 deg, each with its axes on a golden-spiral lattice of
// the sphere as dense as 2 sin(kappa/2) - the distance a change of axis moves a rotation - needs.
std::vector<gemmi::Mat33> UniformRotations(double step_deg);

// The rotations of the space group's point group, in Cartesian coordinates of `cell`.
std::vector<gemmi::Mat33> CartesianPointGroup(const gemmi::SpaceGroup &sg, const gemmi::UnitCell &cell);

// The rotations of `all` in the rotation ASU of `point_group`: the crystal puts a copy of the model
// in every orientation S R, so only one rotation of each class {S R} is worth scoring - here the one
// closest to the identity (largest trace), the classes' Dirichlet region about it.
std::vector<gemmi::Mat33> RotationAsu(const std::vector<gemmi::Mat33> &all,
                                      const std::vector<gemmi::Mat33> &point_group);

// The angle in degrees between a and b, the nearest over the classes {S b}.
double RotationDistanceDeg(const gemmi::Mat33 &a, const gemmi::Mat33 &b,
                           const std::vector<gemmi::Mat33> &point_group);

// The structure factors of one copy of the model, at its centroid, at any reciprocal vector.
class MolecularTransform {
public:
    MolecularTransform(const SearchModel &model, double d_min, size_t nthreads);
    // F at the Cartesian reciprocal vector s (1/A), interpolated trilinearly.
    std::complex<double> At(const gemmi::Vec3 &s) const;
    // The grid At() interpolates (half-l, XYZ), its box edge in A and the blur it is unblurred by.
    const gemmi::FPhiGrid<float> &Grid() const { return grid_; }
    double Box() const { return box_; }
    double Blur() const { return blur_; }

private:
    gemmi::FPhiGrid<float> grid_;   // half-l, on a P1 box `box_` A wide
    double box_ = 0;
    double blur_ = 0;
};

// The observed reflections the search scores against: E^2 in resolution bins, with each unique
// reflection's multiplicity on the whole sphere.
struct SearchReflections {
    std::vector<gemmi::Miller> hkl;
    std::vector<double> e2;
    std::vector<double> multiplicity;
    std::vector<int> bin;
    int bins = 0;
};
SearchReflections MakeSearchReflections(const std::vector<MergedReflection> &merged, const gemmi::UnitCell &cell,
                                        const gemmi::SpaceGroup &sg, double d_min, double d_max);

// Copy j of the model contributes |F_model(s)|^2 at s = h R_j, the reflection's index under the j-th
// operator of the point group as a Cartesian reciprocal vector: refl.hkl.size() x |point group|, index-major.
std::vector<gemmi::Vec3> CopyVectors(const SearchReflections &refl, const gemmi::UnitCell &cell,
                                     const gemmi::SpaceGroup &sg);

// The rotation function at each of `rotations`: the correlation of E^2 - 1 with the bin-normalised
// sum over the point group's copies of |F_model|^2.
std::vector<double> RotationFunction(const MolecularTransform &ft, const SearchReflections &refl,
                                     const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg,
                                     const std::vector<gemmi::Mat33> &rotations, size_t nthreads);

// What the translation function of one space group needs that does not depend on the orientation:
// for every reflection not absent in the group (a "term") and every operator j, s = h R_j as a
// Cartesian vector and exp(2 pi i h.t_j); the term's weight - its multiplicity times E^2 - 1 - and bin;
// and for every term and every ordered pair of operators (j, k), j != k, where the coefficient goes:
// the index h R_j - h R_k on the function's grid, or u = -1 where its l is negative (only l >= 0 is kept).
struct TranslationSetup {
    int nu = 0, nv = 0, nw = 0;      // the function's grid
    int ops = 0;
    std::vector<gemmi::Vec3> s;                     // terms x ops
    std::vector<std::complex<double>> phase;        // terms x ops
    std::vector<double> weight;                     // terms
    std::vector<int> bin;                           // terms
    int bins = 0;
    std::vector<int> pair_j, pair_k;                // the ordered pairs
    std::vector<int> target_u, target_v, target_l;  // terms x pairs
    size_t Terms() const { return weight.size(); }
};
TranslationSetup MakeTranslationSetup(const SearchReflections &refl, const gemmi::UnitCell &cell,
                                      const gemmi::SpaceGroup &sg);

// The translation function of the model in orientation `rotation`: its highest point as a Z-score over
// every point of the function, and where it is (the fractional position of the model's centroid).
struct TranslationPeak {
    double z = 0;
    gemmi::Fractional t;
};
TranslationPeak TranslationFunction(const MolecularTransform &ft, const TranslationSetup &setup,
                                    const gemmi::Mat33 &rotation);

// The model placed: rotated about its centroid by `rotation`, the centroid at `t`, in `cell` and `sg`.
gemmi::Structure PlaceSearchModel(const SearchModel &model, const gemmi::Mat33 &rotation, const gemmi::Fractional &t,
                                  const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg);

// The fraction of the model's C-alpha atoms within 3 A of a C-alpha of a symmetry mate.
double ClashFraction(const gemmi::Structure &placed);

// The space groups with the rotations and the centring of `sg`, in its setting: those merged
// intensities cannot tell apart, apart from systematic absences. `sg` itself first.
std::vector<gemmi::SpaceGroup> SameRotationGroups(const gemmi::SpaceGroup &sg);

// One merge the model is searched against: the answer's, or the twin hypothesis' in a subgroup.
struct MRMerge {
    std::string name;                                // "MERGE" or "TWIN_SUBGROUP_MERGE"
    const std::vector<MergedReflection> *merged = nullptr;
    UnitCell cell;
    gemmi::SpaceGroup space_group;                   // the group it was merged in
};

// One space group's best single-copy placement.
struct MRCandidate {
    int merge = 0;                                   // index into the merges searched
    gemmi::SpaceGroup space_group;
    double vm = 0;                                   // Matthews coefficient of one copy
    int copies_likely = 1;
    double rf_z = 0;                                 // of the orientation, over the rotation function
    double tf_z = 0;
    double clash_fraction = 0;
    double r_work = NAN, r_free = NAN;               // after the rigid body, to d_min_score
    double k_sol = NAN, b_sol = NAN;                 // the bulk solvent R was computed with
    double null_sigma = NAN;                         // R-work above its random-orientation null, in sd
    std::string verdict;                             // PLACED, or why not
    std::shared_ptr<gemmi::Structure> placed;
};

struct MolecularReplacementResult {
    bool ok = false;
    std::string failure_reason;
    std::string model_path;
    int residues_read = 0, residues_kept = 0;
    double mass_da = 0;
    double d_min_search = 0, d_max_search = 0;
    double d_min_score = 0;                          // the limit R-work / R-free are computed to
    std::vector<MRCandidate> candidates;
    // The decision: the candidate the model chooses, or -1 where it chooses none, and in words why.
    int chosen = -1;
    std::string decision;
    bool gpu_search = false;   // the rotation and translation functions ran on the GPU
    // An ENANTIOMORPH decision is taken into the answer's label only as model validation takes a model's
    // enantiomorph (ValidateAgainstModel): the placed model has to fit against its null and the anomalous
    // map at its atoms must not be inverted. Filled by the caller, which runs that validation.
    struct HandCheck {
        bool performed = false;
        bool adopted = false;
        bool model_fits = false;
        double fit_sigma = NAN;
        bool anomalous_inverted = false;
        int anomalous_pairs = 0;
        double anomalous_scatterer_mean_sigma = NAN;
        std::string deepest_site;
        double deepest_sigma = NAN;
        double seconds = 0;
    } hand_check;
    // Seconds per step.
    double t_model = 0, t_transform = 0, t_rotation = 0, t_translation = 0, t_rescore = 0, t_null = 0;
};

// Search `merges` (the answer's first) for one copy of the predicted model at `model_path`, in every
// space group with the rotations of each merge's group, and decide. A candidate is PLACED where its
// translation function stands out (TF Z >= 8), it packs, one copy fits the cell, and its R-work beats
// its own null. The model chooses among the answer merge's groups where one placed candidate leads
// every other by an R-free margin; it chooses the twin hypothesis' merge only where none of the
// answer's groups can hold the model at all (Matthews, packing) and a group of the subgroup can.
MolecularReplacementResult RunMolecularReplacement(const std::string &model_path, const std::vector<MRMerge> &merges,
                                                   Logger &logger, size_t nthreads);