Jungfraujoch/image_analysis/structure_refinement/SolventSites.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 <string>
#include <vector>

#include "gemmi/grid.hpp"
#include "gemmi/model.hpp"
#include "gemmi/monlib.hpp"

// What the solvent of a model - every water and every single-atom ion - is surrounded by, and which species
// that surrounding supports: water, Mg2+, Na+, K+, Ca2+ or Cl- (NH4+ scored too, but only reported). An audit
// of the model as it is: nothing here changes the model.
//
// Each site's neighbours (symmetry mates included) are measured in the coordination bands of Cat_Wiz and typed
// by their chemistry (AtomChemistry); bond-valence sums are taken for each cation; the maps are read at the
// site. Each species then gets a score, a plain sum of terms (SolventSites.cpp says which), and the best one
// is reported, AMBIGUOUS where the runner-up is close.

// What an atom can do in a hydrogen bond or a salt bridge.
struct AtomChemistry {
    bool donor = false;
    bool acceptor = false;
    bool positive = false;       // Lys NZ, Arg NE/NH1/NH2, a polypeptide's N-terminal N
    bool negative = false;       // Asp/Glu carboxylate O, a C-terminal carboxylate, a nucleotide's phosphate O
    bool flip_ambiguous = false; // His ring N, Asn/Gln amide ends: which end is which is the model's guess
};

// The chemistry of `atom` of `residue`: donor or acceptor from the dictionary's energy type (hb_type D, A, B),
// charge from the residue's chemistry; a water is both; an N or O of a residue the dictionary does not
// describe is both. `n_terminus`: the residue begins a polypeptide.
AtomChemistry ChemistryOf(const gemmi::MonLib &monlib, const gemmi::Residue &residue, const gemmi::Atom &atom,
                          bool n_terminus);

enum class SolventSpecies { Water, Mg, Na, K, Ca, Cl, NH4 };
constexpr int SOLVENT_SPECIES = 7;
// The residue name of a species: HOH, MG, NA, K, CA, CL, NH4.
const char *SolventSpeciesName(SolventSpecies s);

// The cations bond valences are summed for, in this order.
constexpr int BVS_CATIONS = 4;   // Mg, Na, K, Ca

struct SolventSite {
    std::string label;           // the residue name as modelled: HOH, MG, NA, ...
    std::string chain, seqid;
    char altloc = '\0';
    gemmi::Position pos;
    double occupancy = 1, b_iso = 0;
    double sigma_xyz = 0;        // the coordinate error assumed for the site, A

    // N and O neighbours, by distance: the bands of Cat_Wiz, and the clash shells beyond
    int n_below_19 = 0, n_19_23 = 0, n_23_26 = 0, n_26_32 = 0, n_32_34 = 0, n_34_36 = 0;
    // the N and O neighbours closer than 3.4 A, by what they can do in a hydrogen bond
    int donors = 0, acceptors = 0, donor_acceptors = 0;
    // the shortest distance between two N or O neighbours closer than 2.6 A to the site; NaN with fewer than two
    double min_ligand_distance = 0;
    // the nearest positive and negative group (a charged atom of the model or an ion), within 6 A; INFINITY none
    double nearest_positive = 0, nearest_negative = 0;
    // the positive groups closer than 3.5 A, the negative ones closer than 3.2 A
    int positive_contacts = 0, negative_contacts = 0;
    // the other ions of the model (not waters) within 4.2 and within 6 A, and the nearest one (INFINITY none)
    int ions_within_42 = 0, ions_within_6 = 0;
    double nearest_ion = 0;
    // bond-valence sum and normalised vector sum for each of the BVS cations (Mg, Na, K, Ca), from the N and
    // O neighbours closer than 3.5 A
    std::array<double, BVS_CATIONS> bvs{}, nvecsum{};
    // the maps at the site, in their rms; NaN where there is no such map
    double map_sigma = 0, difference_sigma = 0, anomalous_sigma = 0;
    // An ion is read where its own density puts it, not where the restraints of its label hold it: the highest point
    // of 2mFo-DFc within SOLVENT_PEAK_REACH of the modelled position (its shell, bands and bond valences are measured
    // from there), and how far that is from it (A). A water is read where it is modelled (shift 0).
    gemmi::Position measured_at;
    double peak_shift = 0;
    // The mean B of the N and O closer than 3.5 A (the environment's B, as CheckMyMetal compares a metal's with its
    // ligands'); NaN with none.
    double b_env = 0;
    // The electrons the 2mFo-DFc map shows at the site: its height there over the height of one electron of the
    // neighbours' (each neighbour's height over what its element at its B gives, the median of them), at the
    // environment's B and as a water's shape - 10 for a water as dense as its neighbours; NaN without the map or
    // neighbours. `electron_spread` is the spread of the neighbours' calibration (rms of the log), floored.
    double electrons = 0, electron_spread = 0;
    // The anomalous map's height expected at the site for each species (f'' at the wavelength, on the scale the
    // model's S and P show), in the map's rms; NaN where there is no anomalous map or nothing to put it on a scale.
    std::array<double, SOLVENT_SPECIES> anomalous_expected{};
    // Purine N7 among the first-shell ligands (closer than 2.6 A), and those of them a divalent ion could not hold:
    // the site more than 1.0 A out of the base's plane, or not 3.5-3.9 A from its O6 or N6 (NAR 2017).
    int n7_ligands = 0, n7_off_plane = 0;

    std::array<double, SOLVENT_SPECIES> score{};   // log-score per species, SolventSpecies order
    SolventSpecies best = SolventSpecies::Water;   // the best-scoring of water, Mg, Na, K, Ca and Cl
    bool ambiguous = false;                        // the runner-up within SOLVENT_AMBIGUOUS_MARGIN of it
    std::string neighbours;                        // the N and O neighbours closer than 3.4 A, nearest first
};

constexpr double SOLVENT_AMBIGUOUS_MARGIN = 2.0;

// The maps a site is read in: 2mFo-DFc, mFo-DFc and the anomalous difference map, each with its rms; any of
// them may be missing (null). `wavelength_A`: the data's, for the anomalous scattering expected (0: unknown).
struct SolventMaps {
    const gemmi::Grid<float> *map = nullptr;
    double map_rms = 0;
    const gemmi::Grid<float> *difference = nullptr;
    double difference_rms = 0;
    const gemmi::Grid<float> *anomalous = nullptr;
    double anomalous_rms = 0;
    double wavelength_A = 0;
};

// How far from an ion's modelled position its density's highest point is looked for (A).
constexpr double SOLVENT_PEAK_REACH = 0.6;

// Every water and single-atom ion of the first model of `st` (in its cell and space group; hydrogens are
// ignored), in model order. `d_min`: the resolution the model was refined at, for the coordinate error
// (0 for a model with no data).
// The sites are measured on `nthreads` threads; the result does not depend on their number.
std::vector<SolventSite> AnalyseSolventSites(const gemmi::Structure &st, double d_min, const SolventMaps &maps,
                                             size_t nthreads = 1);

// Whether the site is labelled a water; whether its label is one of the species decided here (water, Mg, Na, K,
// Ca, Cl); whether the best-supported species is its label.
bool SolventLabelIsWater(const SolventSite &s);
bool SolventLabelScored(const SolventSite &s);
bool SolventLabelAgrees(const SolventSite &s);