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;

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);

// 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);

// The translation function of the model in orientation `rotation` in `sg`: 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 SearchReflections &refl,
                                    const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg,
                                    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 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;
    // 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);