RESULT-CHANGING: Refinement: anisotropic atoms and TLS groups, the displacement model decided on the free set

S1 - the model's displacements kept: an atom's ANISOU is no longer made isotropic; where the model has TLS
groups whose atoms bring no U, their U is made from the groups and the B column read as the total or as the
residual to them, whichever (or the isotropic B) scores the lowest R-free as given. Refined models are written
with ANISOU and TLS groups (mmCIF _pdbx_refine_tls*, and PDB REMARK 3, which gemmi's writer leaves out).

S2/S3 - both engines grid an atom with its full U (U = dev + B/(8 pi^2) I, so the B machinery works on it with
its shape held) and give the target's derivatives in U by the FFT method, CPU and GPU (one warp per atom, no
atomics). After the weights are chosen, the refinement tries, as trials side by side: the B factors refined on;
TLS groups with a residual B (the model's own, one per chain, and chains cut into 2/4/8/16 TLSMD segments by
dynamic programming over a quadratic B(r) fit); and a U per atom with rigid-bond (RIGU form), similarity and
isotropy restraints at three weights. A trial displaces another only within the geometry bounds and with both
R-free (beyond the tie) and the free -ln L lower. The water-build rounds hold the chosen U shapes.

An anisotropic atom's density is now cut at the radius of its largest eigenvalue rather than its largest
diagonal element (also in the validation's density and the GPU gather), which left skewed atoms truncated.

Refine arm + open-arm subset (11 sets, vs 51cbe0ef7), refined R-free: 1exr 0.2064 -> 0.1377 (SHELXL U kept),
193l 0.1689 -> 0.1322 (aniso), 2r8s 0.2262 -> 0.2054 (TLS 47), 1hr2 0.2463 -> 0.2316 (TLS 30), 1k4c 0.2110 ->
0.1881 (TLS 42), 36gk 0.2445 -> 0.2221, 7l6j 0.1766 -> 0.1553, 7os3 0.2040 -> 0.2014, 8agq 0.1948 -> 0.1595,
8sa8 0.1830 -> 0.1548, 6f3p 0.1821 -> 0.1453. Every output U positive definite. The first refinement takes
about 2-3x as long (the hypothesis trials; measured on a shared GPU).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01HVeDQqN9qRD57vnShzTC37
This commit is contained in:
leonarski_fandClaude Opus 5.5 committed 2026-10-27 17:34:27 +01:00
1 parent 51cbe0ef7d
commit 6472ce3ff1
19 files changed
+1764 -68

No files matched your search

+23
View File
@@ -595,6 +595,29 @@ predicted distances around metal sites in proteins" (2006), Acta Cryst. D62, 678
[doi:10.1107/S0907444906014594](https://doi.org/10.1107/S0907444906014594), and, for Mg²⁺, Na⁺ and K⁺ at
nucleic acids, F. Leonarski, L. D'Ascenzo and P. Auffinger, "Mg2+ ions: do they bind to nucleobase
nitrogens?" (2017), Nucleic Acids Res. 45, 987-1004 [doi:10.1093/nar/gkw1175](https://doi.org/10.1093/nar/gkw1175).
The anisotropic displacements: a TLS group's U is that of a rigid body's translation, libration and their
correlation, U = T + A L A^T + A S + S^T A^T, of V. Schomaker and K. N. Trueblood, "On the rigid-body motion of
molecules in crystals" (1968), Acta Cryst. B24, 63-76
[doi:10.1107/S0567740868001718](https://doi.org/10.1107/S0567740868001718), refined against the likelihood with a
residual B per atom, and a model's B column read as the total or as the residual to its TLS groups, as in
M. D. Winn, M. N. Isupov and G. N. Murshudov, "Use of TLS parameters to model anisotropic displacements in
macromolecular refinement" (2001), Acta Cryst. D57, 122-133
[doi:10.1107/S0907444900014736](https://doi.org/10.1107/S0907444900014736). The chains are cut into TLS groups as
TLSMD cuts them - contiguous segments, by dynamic programming over the fit of an isotropic TLS model to the
refined B factors: J. Painter and E. A. Merritt, "Optimal description of a protein structure in terms of multiple
groups undergoing TLS motion" (2006), Acta Cryst. D62, 439-450
[doi:10.1107/S0907444906005270](https://doi.org/10.1107/S0907444906005270). Atoms refined with a U of their own are
restrained by the rigid-bond postulate of F. L. Hirshfeld, "Can X-ray data distinguish bonding effects from
vibrational smearing?" (1976), Acta Cryst. A32, 239-244
[doi:10.1107/S0567739476000533](https://doi.org/10.1107/S0567739476000533), in the form of SHELXL's RIGU - the
components of U along the bond and across it in the bond's plane: A. Thorn, B. Dittrich and G. M. Sheldrick,
"Enhanced rigid-bond restraints" (2012), Acta Cryst. A68, 448-451
[doi:10.1107/S0108767312014535](https://doi.org/10.1107/S0108767312014535) - with SHELXL's similarity (SIMU) and
approximate-isotropy (ISOR) restraints as G. M. Sheldrick, "Crystal structure refinement with SHELXL" (2015), Acta
Cryst. C71, 3-8 [doi:10.1107/S2053229614024218](https://doi.org/10.1107/S2053229614024218) documents them. Which
displacement model a refinement keeps - isotropic, TLS, or a U per atom - is decided on the free set, as a test of
what the data support rather than a rule by resolution: E. A. Merritt, "To B or not to B: a question of
resolution?" (2012), Acta Cryst. D68, 468-477 [doi:10.1107/S0907444911028320](https://doi.org/10.1107/S0907444911028320).
**Unmodelled density** — the blobs the refinement lists are the connected regions of the difference map above
a cutoff, found by flood fill as Coot's "unmodelled blobs" finds them, in GEMMI's implementation
+41 -5
View File
@@ -27,8 +27,12 @@ statistics and its map coefficients go out; the restraint dictionary is compiled
### 15.1 What is refined, and in what order
The model is taken as the refinement can use it: the first model of the file only, its hydrogens removed,
every anisotropic B replaced by its isotropic equivalent, and every B clamped to the range the refinement
allows (2 Ų up to twice the model's largest B, within 100-300 Ų). Occupancies are kept as they are, and
and every B clamped to the range the refinement allows (2 Ų up to twice the model's largest B, within
100-300 Ų). An atom's anisotropic U (ANISOU, `_atom_site_anisotrop`) is kept. Where the model has TLS
groups (`_pdbx_refine_tls*`, PDB `REMARK 3`) whose atoms bring no U of their own, their U is made from the
groups, and the B column is read either as the total B (as phenix.refine writes it) or as the residual to the
groups (as REFMAC writes it) - whichever of the two, and of the isotropic B as it stands, scores the lowest
R-free on the model as given (`REFINE_ADP_INPUT`). Occupancies are kept as they are, and
so are alternative conformations, each refined with its own restraints.
Two stages follow.
@@ -184,6 +188,24 @@ read, without the names of the links made here. The rmsd of these bonds - with t
a metal within a ligand, a heme iron's to its porphyrin - is reported apart from the covalent bonds'
(`REFINE_RMSD_METAL_BONDS`).
**Anisotropic atoms.** An atom's U is held as its equivalent isotropic B and a traceless anisotropy,
$U = U_\mathrm{dev} + B/(8\pi^2)\,I$, so that everything that refines, restrains and bounds B works on an anisotropic
atom as on an isotropic one, its shape held. Both engines grid it with its full U (an anisotropic Gaussian per
form-factor term, cut at the radius of its widest axis) and give the target's derivatives in its six U elements by the
same FFT method: for $\rho_q = a_q \exp(\mathbf{r}^T M_q \mathbf{r})$, $M_q = -4\pi^2(8\pi^2 U + B_q + B_\mathrm{blur})^{-1}$,
$\partial\rho_q/\partial U = \rho_q\,(M_q + 2 M_q\mathbf{r}\,\mathbf{r}^T M_q)$. Every U stays positive definite: its
smallest eigenvalue (as B) at least the cycle's floor, its largest at most the ceiling of §15.1.
A **TLS group** gives each of its atoms $U = T + A L A^T + A S + S^T A^T$ at $\mathbf{r}$ from its origin,
$A = [[0, z, -y], [-z, 0, x], [y, -x, 0]]$ (Schomaker and Trueblood), plus a residual isotropic B per atom; its 20
parameters (the trace of S does not enter U) are refined against the X-ray target alone, the residual B held, with
the derivatives by the chain rule from the atoms' U derivatives ($G$ for T, $A^T G A$ for L, $2 A^T G$ for S, $G$ an
atom's symmetric matrix of derivatives). An atom refined with a **U of its own** is restrained, for each pair of
bonded or angle-related atoms, by the rigid bond $|(U_i - U_j)\mathbf{n}|^2/\sigma^2$, $\sigma$ = 0.004 Ų along the
line $\mathbf{n}$ between them (the components SHELXL's RIGU restrains), and by the similarity
$|U_i - U_j|^2/\sigma^2$, $\sigma$ = 0.04 Ų (SIMU); an atom in no pair (a water, an ion) by
$|U - \mathrm{tr}(U)/3\,I|^2/\sigma^2$, $\sigma$ = 0.1 Ų (ISOR).
**B restraints.** The B factors of two atoms are restrained to each other in proportion to their size, by
the symmetric Kullback-Leibler divergence of two isotropic Gaussians of variances $B_i$ and $B_j$:
@@ -247,6 +269,17 @@ are:
whose free set the depositor's refinement has already fitted, R-work falls while R-free rises, and the
trajectory ends where it began.
3. **The final ladder**: the trajectory's weights times $\{0.5, 1, 2\}$, kept where one displaces it.
4. **The displacement model**, from where the weights ended, each trial a macro-cycle longer than a step: the B
factors refined on as they are (isotropic, or the model's U held); TLS groups with a residual B - the model's own,
one per chain of at least 50 refinable atoms, and the chains cut into 2, 4, 8 and 16 contiguous segments of at
least 10 residues as TLSMD cuts them (by dynamic programming over the least-squares fit of an isotropic TLS model,
a quadratic B(r), to the refined B factors), each cut where it gives more groups than the one before; and a U per
atom, at the U weight balanced as $w_{B,0}$ is (on the U shaken by 20% of its B) times $\{1/4, 1, 4\}$. A trial
displaces the one before it only where its geometry is within bounds and both its R-free (beyond the tie) and the
$-\ln L$ of the free reflections are lower: the free set tests the parameters a model adds, so the choice is made
on the evidence of each run, not at a resolution (Merritt 2012). The refined U are written as ANISOU and the TLS
groups as `_pdbx_refine_tls*` and `REMARK 3` (`REFINE_ADP_MODEL`, `REFINE_ANISO_ATOMS`, `REFINE_TLS_GROUPS`).
The rounds of the water build (§15.10) keep the model chosen: they hold each atom's U shape and refine its B.
$w$ never exceeds $w_0/8$. The trials of one set run side by side - on one GPU, on several, or on the CPU's
threads - each from the same start and taking nothing from the one before it, so the result does not depend
@@ -384,9 +417,10 @@ What the refinement does not do yet - each a known gap, not a design choice:
- **Custom ligands and other monomer libraries.** Only the 84 compounds of §15.4 are restrained. Any other
residue is held fixed (it stays in $F_\mathrm{calc}$, the mask and the repulsions); no restraints are
generated for a new ligand, and neither a user's restraint CIF nor an external monomer library can be given.
- **TLS** refinement.
- **Anisotropic ADPs**: an anisotropic B is replaced by its isotropic equivalent, and only isotropic B is
refined.
- **TLS and anisotropic U after the first refinement**: the water build's rounds hold the U shapes and the TLS
parameters the first refinement chose and refine only B; a TLS group's T and L are not held positive
semi-definite (only every atom's total U is kept positive definite); the TLS groups are contiguous segments of one
chain, never a domain made of several.
- **NCS restraints**: copies are refined independently.
- **Twin refinement**: the target is not a twinned one, also where the twin hypothesis' merge was written.
- **Ion identification**: waters are built and removed (§15.10), but no ion is placed or relabelled; the solvent
@@ -411,6 +445,8 @@ What the refinement does not do yet - each a known gap, not a design choice:
- P. J. Huber, "Robust estimation of a location parameter", *Ann. Math. Statist.* **35** (1964), 73-101, and P. Charbonnier, L. Blanc-Féraud, G. Aubert & M. Barlaud, "Deterministic edge-preserving regularization in computed imaging", *IEEE Trans. Image Process.* **6** (1997), 298-311 (the robust scale fit and its smooth, pseudo-Huber loss, §15.3).
- A. A. Vagin, R. A. Steiner, A. A. Lebedev, L. Potterton, S. McNicholas, F. Long & G. N. Murshudov, "REFMAC5 dictionary: organization of prior chemical knowledge and guidelines for its use", *Acta Cryst.* **D60** (2004), 2184-2195 (the restraint dictionary, §15.4).
- M. Wojdyr, "GEMMI: A library for structural biology", *J. Open Source Softw.* **7** (2022), 4200 (the topology, the riding hydrogens and the blob search, §15.4, §15.9).
- V. Schomaker & K. N. Trueblood, *Acta Cryst.* **B24** (1968), 63-76, and M. D. Winn, M. N. Isupov & G. N. Murshudov, *Acta Cryst.* **D57** (2001), 122-133 (TLS, and the B column total or residual to it, §15.1, §15.5); J. Painter & E. A. Merritt, *Acta Cryst.* **D62** (2006), 439-450 (TLSMD's segments, §15.7).
- F. L. Hirshfeld, *Acta Cryst.* **A32** (1976), 239-244, A. Thorn, B. Dittrich & G. M. Sheldrick, *Acta Cryst.* **A68** (2012), 448-451, and G. M. Sheldrick, *Acta Cryst.* **C71** (2015), 3-8 (the rigid-bond, similarity and isotropy restraints of anisotropic atoms, §15.5); E. A. Merritt, *Acta Cryst.* **D68** (2012), 468-477 (the displacement model decided by test, §15.7).
- A. Blondel & M. Karplus, "New formulation for derivatives of torsion angles and improper torsion angles in molecular mechanics: elimination of singularities", *J. Comput. Chem.* **17** (1996), 1132-1141 (the torsion derivatives, §15.5).
- M. M. Harding, "Geometry of metal-ligand interactions in proteins", *Acta Cryst.* **D57** (2001), 401-411, and "Small revisions to predicted distances around metal sites in proteins", *Acta Cryst.* **D62** (2006), 678-682; F. Leonarski, L. D'Ascenzo & P. Auffinger, "Mg2+ ions: do they bind to nucleobase nitrogens?", *Nucleic Acids Res.* **45** (2017), 987-1004 (the metal-site distances, §15.5).
- J. Nocedal, "Updating quasi-Newton matrices with limited storage", *Math. Comp.* **35** (1980), 773-782, and D. C. Liu & J. Nocedal, "On the limited memory BFGS method for large scale optimization", *Math. Program.* **45** (1989), 503-528 (the minimiser, §15.6).
+4 -1
View File
@@ -1188,7 +1188,8 @@ Where a model was placed and refined ([Refinement](RUGNUX_MODEL_AND_REFERENCE.md
| `REFINE` | `PERFORMED` \| `NOT_PERFORMED` (with `REFINE_REASON`) |
| `REFINE_DEVICE`, `REFINE_SECONDS` | *dev*. `GPU` or `CPU`, and the time the refinement took |
| `REFINE_D_MIN` | The finest reflection refined against, Å |
| `REFINE_R_WORK_START`, `REFINE_R_FREE_START` | The model as given, scored as the refinement scores (riding hydrogens, B isotropic, the fine solvent mask, the robust scale): REFMAC's 0-cycle R-factors to about 0.002 on the same model and data. Not `R_WORK`/`R_FREE`, which are the validation's own scoring of the placed model |
| `REFINE_ADP_INPUT` | How the model's displacements were read ([§15.1](CPU_DATA_ANALYSIS_REFINEMENT.md#151-what-is-refined-and-in-what-order)): `isotropic`, `ANISOU` (its atoms' own U), or `TLS, B total` / `TLS, B residual` (U from its TLS groups, its B column the total or the residual) |
| `REFINE_R_WORK_START`, `REFINE_R_FREE_START` | The model as given, scored as the refinement scores (riding hydrogens, its U as read, the fine solvent mask, the robust scale): REFMAC's 0-cycle R-factors to about 0.002 on the same model and data. Not `R_WORK`/`R_FREE`, which are the validation's own scoring of the placed model |
| `REFINE_RIGID_BODIES` | The chains moved as rigid bodies, where that step was kept (0: not kept) |
| `REFINE_R_WORK_PLACEMENT`, `REFINE_R_FREE_PLACEMENT` | After the placement stage (rigid bodies, then one B per residue, each kept only where it lowers R-free) |
| `REFINE_R_WORK`, `REFINE_R_FREE` | After the restrained stage: the refined model, the one written as `<prefix>_refined_model.cif` |
@@ -1197,6 +1198,8 @@ Where a model was placed and refined ([Refinement](RUGNUX_MODEL_AND_REFERENCE.md
| `REFINE_RMSD_BONDS_START`, `REFINE_RMSD_ANGLES_START` | *dev*. The same for the model as given |
| `REFINE_CLASHES` | Pairs of atoms in the refined model more than 0.4 Å closer than their contact distance (symmetry mates included) |
| `REFINE_XRAY_WEIGHT` | *dev*. The X-ray weight of the positions that was kept |
| `REFINE_ADP_MODEL` | The displacement model the free set chose ([§15.7](CPU_DATA_ANALYSIS_REFINEMENT.md#157-the-weights-decided-by-trial)): `isotropic`, `the model's U held`, `TLS, <n> groups (...)` or `anisotropic` |
| `REFINE_ANISO_ATOMS`, `REFINE_TLS_GROUPS` | The atoms of the refined model with an anisotropic U (ANISOU), and its TLS groups |
| `REFINE_ATOMS`, `REFINE_ATOMS_REFINED` | The model's atoms (hydrogens not counted) and those refined: the atoms of the residues the dictionary describes, with an occupancy above zero |
| `REFINE_RESIDUES_FIXED`, `REFINE_RESIDUES_FIXED_NAMES` | The residues the embedded dictionary does not describe, held fixed but kept in the structure factors, the solvent mask and the clash terms, and their names |
| `REFINE_BULK_SOLVENT_K_SOL`, `REFINE_BULK_SOLVENT_B_SOL` | *dev*. The flat bulk solvent of the refined model |
@@ -21,6 +21,7 @@ ADD_LIBRARY(JFJochStructureRefinement STATIC
RefinementMath.h RefinementMinimiser.h
RefinementRestraints.cpp RefinementRestraints.h
RefinementXray.cpp RefinementXray.h
RefinementADP.cpp RefinementADP.h
Refinement.cpp Refinement.h
DensityBlobs.cpp DensityBlobs.h
SolventSites.cpp SolventSites.h
@@ -451,7 +451,7 @@ void PutModelDensityOnGrid(gemmi::DensityCalculator<Table, float> &dc, const gem
boxes[i] = MakeDensityBox(dc.grid, fpos, dc.estimate_radius(d.iso, b));
} else {
const auto aniso_b = atom.aniso.scaled(CReal(gemmi::u_to_b())).added_kI(CReal(dc.blur));
const CReal b_max = std::max(std::max(aniso_b.u11, aniso_b.u22), aniso_b.u33);
const CReal b_max = AnisoRadiusB(aniso_b);
const double radius = dc.estimate_radius(coef.precalculate_density_iso(b_max, addend), b_max);
d.aniso = coef.precalculate_density_aniso_b(aniso_b, addend);
boxes[i] = MakeDensityBox(dc.grid, fpos, radius);
@@ -3,6 +3,8 @@
#pragma once
#include <algorithm>
#include <array>
#include <cstddef>
#include <vector>
@@ -22,6 +24,16 @@
// (empty in P1). Both functions below take them, for a grid of the same size.
std::vector<size_t> OrbitLeaders(const gemmi::Grid<float> &grid, size_t nthreads);
// The B an anisotropic atom's density is cut at the radius of: its largest eigenvalue (of 8 pi^2 U plus the blur), so
// that the density is cut where it falls to the cutoff along its widest axis. gemmi takes the largest diagonal element,
// which a U with large off-diagonal elements leaves short of that axis.
template <typename Real>
Real AnisoRadiusB(const gemmi::SMat33<Real> &aniso_b) {
const gemmi::SMat33<double> m{aniso_b.u11, aniso_b.u22, aniso_b.u33, aniso_b.u12, aniso_b.u13, aniso_b.u23};
const std::array<double, 3> ev = m.calculate_eigenvalues();
return static_cast<Real>(std::max({ev[0], ev[1], ev[2], double(aniso_b.u11), double(aniso_b.u22), double(aniso_b.u33)}));
}
// `dc` set up as for put_model_density_on_grid(): d_min, rate, blur, the grid's cell and group. On macOS the
// Apple GPU puts the density on the grid (ModelDensityMetal.h), to the same bits; `gpu` false keeps it on the
// CPU.
@@ -13,6 +13,7 @@
#include "gemmi/fourier.hpp" // get_size_for_hkl, get_f_phi_on_grid
#include "gemmi/solmask.hpp" // SolventMasker, refmac_radius_for_bulk_solvent
#include "ModelGrid.h" // AnisoRadiusB
#include "RigidBodyRefine.h" // RigidBodyZoneGrid
namespace {
@@ -86,7 +87,7 @@ void ModelDensityAtoms(const gemmi::Model &model, const gemmi::DensityCalculator
}
} else {
const auto aniso_b = atom.aniso.scaled(CReal(gemmi::u_to_b())).added_kI(CReal(dc.blur));
const CReal b_max = std::max(std::max(aniso_b.u11, aniso_b.u22), aniso_b.u33);
const CReal b_max = AnisoRadiusB(aniso_b);
const double radius = dc.estimate_radius(coef.precalculate_density_iso(b_max, addend), b_max);
a.radius = static_cast<float>(radius);
a.max_r2 = static_cast<float>(radius * radius);
@@ -31,7 +31,9 @@
#include "../../common/Logger.h"
#include "../../common/ParallelFor.h"
#include "ModelFFT.h"
#include "ModelGrid.h" // AnisoRadiusB
#include "ModelScaling.h"
#include "RefinementADP.h"
#include "RefinementMath.h"
#include "RefinementMinimiser.h"
#include "RefinementXray.h"
@@ -84,8 +86,8 @@ std::vector<gemmi::Atom *> AtomList(gemmi::Model &model) {
return atoms;
}
// The model as refinement takes it: the first model only, without hydrogens and atoms of unknown element,
// every B isotropic (an anisotropic one becomes its equivalent isotropic B) and inside the allowed range.
// The model as refinement takes it: the first model only, without hydrogens and atoms of unknown element, every B
// inside the allowed range - an anisotropic atom's equivalent isotropic B, its U kept (Context::TakeAniso() takes it).
gemmi::Structure PrepareModel(const gemmi::Structure &placed, const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg) {
gemmi::Structure st = placed;
st.models.resize(1);
@@ -96,10 +98,8 @@ gemmi::Structure PrepareModel(const gemmi::Structure &placed, const gemmi::UnitC
for (gemmi::Residue &r : ch.residues) {
std::erase_if(r.atoms, [](const gemmi::Atom &a) { return a.element == gemmi::El::X; });
for (gemmi::Atom &a : r.atoms) {
if (a.aniso.nonzero()) {
if (a.aniso.nonzero())
a.b_iso = static_cast<float>(gemmi::u_to_b() * a.aniso.trace() / 3.0);
a.aniso = {0, 0, 0, 0, 0, 0};
}
a.b_iso = static_cast<float>(std::clamp<double>(a.b_iso, B_FLOOR, B_CEILING_MAX));
}
}
@@ -143,6 +143,110 @@ struct Context {
std::vector<gemmi::Position> rider_offset;
int evaluations = 0;
const std::atomic<bool> *cancel = nullptr; // RefinementSettings::cancel
// The anisotropic atoms (empty: none), which the engine reads (RefinementEngine::SetAniso), and per atom how far
// its B stays inside the bounds so that its U does: its smallest eigenvalue above b_floor, its largest below
// b_ceiling (in B, A^2).
std::vector<AtomAniso> aniso;
std::vector<double> b_lo, b_hi;
double BLow(size_t i) const { return b_floor + (b_lo.empty() ? 0.0 : b_lo[i]); }
double BHigh(size_t i) const { return std::max(BLow(i), b_ceiling - (b_hi.empty() ? 0.0 : b_hi[i])); }
// The anisotropy of `a` (empty: none) for the atoms as they are, and the bounds of their B with it.
void SetAniso(const std::vector<AtomAniso> &a) {
aniso = a;
b_lo.clear();
b_hi.clear();
if (aniso.empty())
return;
b_lo.assign(atoms.size(), 0.0);
b_hi.assign(atoms.size(), 0.0);
for (size_t i = 0; i < aniso.size(); i++)
if (aniso[i].on) {
const std::array<float, 6> &d = aniso[i].dev;
const gemmi::SMat33<double> dev{d[0], d[1], d[2], d[3], d[4], d[5]};
const std::array<double, 3> ev = dev.calculate_eigenvalues();
b_lo[i] = std::max(0.0, -gemmi::u_to_b() * std::min({ev[0], ev[1], ev[2]}));
b_hi[i] = std::max(0.0, gemmi::u_to_b() * std::max({ev[0], ev[1], ev[2]}));
}
}
// TLS bodies (empty: none) and per atom the part of its B they hold as they were last applied (8 pi^2 tr(U_TLS) / 3,
// and what bringing U inside the bounds added): the rest of its B is its residual.
std::vector<TlsBody> tls;
std::vector<double> tls_b;
// The bodies' U at the atoms as they are now, each atom's residual B kept: its anisotropy the traceless part of
// U_TLS, its B the residual plus the trace.
void ApplyTls() {
if (tls.empty())
return;
if (aniso.size() != atoms.size())
aniso.assign(atoms.size(), AtomAniso());
tls_b.resize(atoms.size(), 0.0);
for (const TlsBody &body : tls)
for (int i : body.atoms) {
const gemmi::SMat33<double> u = body.U(atoms[i]->pos);
double t = u.trace() / 3.0;
const double residual = atoms[i]->b_iso - tls_b[i];
// the total U inside the bounds: its largest eigenvalue (in B) at most the ceiling, its smallest at least
// the floor - the whole U scaled down, or raised by a multiple of I
gemmi::SMat33<double> b = u.scaled(gemmi::u_to_b()).added_kI(residual);
std::array<double, 3> ev = b.calculate_eigenvalues();
const double top = std::max({ev[0], ev[1], ev[2]});
if (top > b_ceiling) {
b = b.scaled(b_ceiling / top);
ev = b.calculate_eigenvalues();
}
const double lowest = std::min({ev[0], ev[1], ev[2]});
if (!(lowest >= b_floor))
b = b.added_kI(b_floor - (std::isfinite(lowest) ? lowest : 0.0));
t = b.trace() / 3.0;
atoms[i]->b_iso = static_cast<float>(t);
// what the bodies hold of B, the bounds' share included: the residual stays as it was
tls_b[i] = atoms[i]->b_iso - residual;
aniso[i].on = true;
aniso[i].dev = {static_cast<float>((b.u11 - t) / gemmi::u_to_b()), static_cast<float>((b.u22 - t) / gemmi::u_to_b()),
static_cast<float>((b.u33 - t) / gemmi::u_to_b()), static_cast<float>(b.u12 / gemmi::u_to_b()),
static_cast<float>(b.u13 / gemmi::u_to_b()), static_cast<float>(b.u23 / gemmi::u_to_b())};
}
MoveRiders();
SetAniso(aniso);
}
// The U of every atom (AnisoU for the anisotropic, B / (8 pi^2) I for the others), in A^2.
std::vector<std::array<double, 6>> UOfAtoms() const {
std::vector<std::array<double, 6>> u(atoms.size());
for (size_t i = 0; i < atoms.size(); i++) {
const double b = atoms[i]->b_iso / gemmi::u_to_b();
if (!aniso.empty() && aniso[i].on) {
const gemmi::SMat33<float> m = AnisoU(aniso[i], atoms[i]->b_iso);
u[i] = {m.u11, m.u22, m.u33, m.u12, m.u13, m.u23};
} else
u[i] = {b, b, b, 0, 0, 0};
}
return u;
}
// The anisotropic U the atoms of the model bring (PrepareModel() keeps it) as their AtomAniso, taken off them.
void TakeAniso() {
std::vector<AtomAniso> a(atoms.size());
bool any = false;
for (size_t i = 0; i < atoms.size(); i++) {
gemmi::Atom &at = *atoms[i];
if (!at.aniso.nonzero())
continue;
// its widest axis no wider than the refinement allows
const double top = AnisoRadiusB(at.aniso.scaled(gemmi::u_to_b()));
if (top > B_CEILING_MAX) {
at.aniso = at.aniso.scaled(static_cast<float>(B_CEILING_MAX / top));
at.b_iso = static_cast<float>(gemmi::u_to_b() * at.aniso.trace() / 3.0);
}
const double t = at.aniso.trace() / 3.0;
a[i].on = true;
a[i].dev = {static_cast<float>(at.aniso.u11 - t), static_cast<float>(at.aniso.u22 - t),
static_cast<float>(at.aniso.u33 - t), at.aniso.u12, at.aniso.u13, at.aniso.u23};
at.aniso = {0, 0, 0, 0, 0, 0};
any = true;
}
SetAniso(any ? a : std::vector<AtomAniso>());
}
bool Cancelled() const { return cancel != nullptr && cancel->load(std::memory_order_relaxed); }
@@ -247,8 +351,8 @@ struct Context {
RFactors Refit() {
const RefinementStage stage("refit");
double b_min = b_ceiling;
for (const gemmi::Atom *a : atoms)
b_min = std::min<double>(b_min, a->b_iso);
for (size_t i = 0; i < atoms.size(); i++)
b_min = std::min<double>(b_min, atoms[i]->b_iso - (b_lo.empty() ? 0.0 : b_lo[i]));
b_floor = std::max(B_FLOOR, b_min - B_STEP_DOWN);
engine->SetBlur(RefinementBlur(cell, *sg, d_min, b_floor));
{
@@ -318,8 +422,10 @@ struct Context {
double b_min = b_ceiling;
for (size_t i = 0; i < atoms.size(); i++)
if (rider_parent.empty() || rider_parent[i] < 0) {
atoms[i]->b_iso = static_cast<float>(std::clamp(atoms[i]->b_iso + b_iso, B_FLOOR, b_ceiling));
b_min = std::min<double>(b_min, atoms[i]->b_iso);
const double lo = b_lo.empty() ? 0.0 : b_lo[i], hi = b_hi.empty() ? 0.0 : b_hi[i];
atoms[i]->b_iso = static_cast<float>(std::clamp(atoms[i]->b_iso + b_iso, B_FLOOR + lo,
std::max(B_FLOOR + lo, b_ceiling - hi)));
b_min = std::min<double>(b_min, atoms[i]->b_iso - lo);
}
MoveRiders();
b_floor = std::max(B_FLOOR, b_min - B_STEP_DOWN);
@@ -541,7 +647,7 @@ public:
if (mode_ == Mode::Xyz)
a.pos = gemmi::Position(p[3 * k], p[3 * k + 1], p[3 * k + 2]);
else
a.b_iso = static_cast<float>(std::clamp(p[k], ctx_.b_floor, ctx_.b_ceiling));
a.b_iso = static_cast<float>(std::clamp(p[k], ctx_.BLow(param_atoms_[k]), ctx_.BHigh(param_atoms_[k])));
}
});
ctx_.MoveRiders();
@@ -574,7 +680,7 @@ public:
else {
// held at a bound, B does not move further out: by the whole target, restraints included
const double g = weight_ * gx[i][3] + gr[i][3];
const bool held = (p[k] <= ctx_.b_floor && g > 0) || (p[k] >= ctx_.b_ceiling && g < 0);
const bool held = (p[k] <= ctx_.BLow(i) && g > 0) || (p[k] >= ctx_.BHigh(i) && g < 0);
gradient[k] = held ? 0.0 : g;
}
}
@@ -593,6 +699,138 @@ private:
mutable std::vector<std::array<double, 4>> gx_, gr_;
};
// The X-ray target's derivatives in U of every atom, with what the target says of a riding hydrogen's B (its parent's
// equivalent isotropic B, 8 pi^2 tr(U) / 3) said of its parent's U.
std::vector<std::array<double, 6>> FoldedUGradient(const Context &ctx, const std::vector<std::array<double, 4>> &gx) {
std::vector<std::array<double, 6>> g = ctx.engine->UGradient();
g.resize(ctx.atoms.size(), {0, 0, 0, 0, 0, 0});
for (size_t r = 0; r < ctx.rider_parent.size(); r++)
if (const int p = ctx.rider_parent[r]; p >= 0)
for (int c = 0; c < 3; c++)
g[p][c] += gx[r][3] * gemmi::u_to_b() / 3.0;
return g;
}
// The displacements of the anisotropic atoms against w times the X-ray target. Mode::U: a U per atom of
// `param_atoms`, as 8 pi^2 U (in the units of B), plus the U restraints; its smallest eigenvalue held at the cycle's
// floor and its diagonal under the ceiling. Mode::Tls: the parameters of each of ctx.tls, in the units of B at the
// body's radius (T, L r^2, S r, each times 8 pi^2), the atoms' residual B held; X-ray only.
class AdpTarget {
public:
enum class Mode { U, Tls };
AdpTarget(Context &ctx, Mode mode, const std::vector<int> &param_atoms, double weight, const URestraints *restraints)
: ctx_(ctx), mode_(mode), param_atoms_(param_atoms), weight_(weight), restraints_(restraints) {
for (const TlsBody &body : ctx.tls) {
double r2 = 0;
for (int i : body.atoms)
r2 += body.origin.dist_sq(ctx.atoms[i]->pos);
radius_.push_back(std::max(1.0, std::sqrt(r2 / std::max<size_t>(body.atoms.size(), 1))));
}
}
std::vector<double> Get() const {
std::vector<double> p;
if (mode_ == Mode::U) {
const std::vector<std::array<double, 6>> u = ctx_.UOfAtoms();
for (int i : param_atoms_)
for (int c = 0; c < 6; c++)
p.push_back(gemmi::u_to_b() * u[i][c]);
} else
for (size_t b = 0; b < ctx_.tls.size(); b++) {
const std::array<double, TlsBody::PARAMETERS> q = ctx_.tls[b].Get();
for (int k = 0; k < TlsBody::PARAMETERS; k++)
p.push_back(q[k] * Scale(b, k));
}
return p;
}
void Set(const double *p) const {
if (mode_ == Mode::U) {
const int n = static_cast<int>(param_atoms_.size());
ParallelChunks(n, ThreadsForWork(n, ctx_.nthreads, 16384), [&](int lo, int hi) {
for (int k = lo; k < hi; k++)
SetU(param_atoms_[k], p + 6 * k);
});
ctx_.MoveRiders();
ctx_.SetAniso(ctx_.aniso);
} else {
for (size_t b = 0; b < ctx_.tls.size(); b++) {
std::array<double, TlsBody::PARAMETERS> q{};
for (int k = 0; k < TlsBody::PARAMETERS; k++)
q[k] = p[b * TlsBody::PARAMETERS + k] / Scale(b, k);
ctx_.tls[b].Set(q);
}
ctx_.ApplyTls();
}
}
double MaxStep() const { return MAX_STEP_B; }
bool Evaluate(const double *p, double *cost, double *gradient) const {
const RefinementStage stage(mode_ == Mode::U ? "eval.u" : "eval.tls");
Set(p);
ctx_.evaluations++;
std::vector<std::array<double, 4>> gx;
const double xray = ctx_.engine->Evaluate(ctx_.Model(), gradient != nullptr ? &gx : nullptr);
double restraints = 0;
std::vector<std::array<double, 6>> gr;
if (mode_ == Mode::U && restraints_ != nullptr) {
if (gradient != nullptr)
gr.assign(ctx_.atoms.size(), {0, 0, 0, 0, 0, 0});
restraints = EvaluateURestraints(*restraints_, ctx_.Positions(), ctx_.UOfAtoms(), gradient != nullptr ? &gr : nullptr);
}
*cost = weight_ * xray + restraints;
if (gradient == nullptr)
return std::isfinite(*cost);
const std::vector<std::array<double, 6>> gu = FoldedUGradient(ctx_, gx);
if (mode_ == Mode::U) {
for (size_t k = 0; k < param_atoms_.size(); k++)
for (int c = 0; c < 6; c++)
gradient[6 * k + c] = (weight_ * gu[param_atoms_[k]][c] + (gr.empty() ? 0.0 : gr[param_atoms_[k]][c])) /
gemmi::u_to_b();
} else {
for (size_t b = 0; b < ctx_.tls.size(); b++) {
const std::array<double, TlsBody::PARAMETERS> d = ctx_.tls[b].Gradient(ctx_.atoms, gu);
for (int k = 0; k < TlsBody::PARAMETERS; k++)
gradient[b * TlsBody::PARAMETERS + k] = weight_ * d[k] / Scale(b, k);
}
}
return std::isfinite(*cost);
}
private:
double Scale(size_t body, int k) const {
const double r = radius_[body];
return gemmi::u_to_b() * (k < 6 ? 1.0 : k < 12 ? r * r : r);
}
// Atom i's U from 8 pi^2 U = p[0..5]: its smallest eigenvalue raised to the cycle's floor and its largest brought
// down to the ceiling where they are outside, then its B and its anisotropy.
void SetU(int i, const double *p) const {
gemmi::SMat33<double> b{p[0], p[1], p[2], p[3], p[4], p[5]};
const std::array<double, 3> ev = b.calculate_eigenvalues();
const double lowest = std::min({ev[0], ev[1], ev[2]});
if (!(lowest >= ctx_.b_floor))
b = b.added_kI(ctx_.b_floor - (std::isfinite(lowest) ? lowest : 0.0));
const std::array<double, 3> ev2 = b.calculate_eigenvalues();
const double top = std::max({ev2[0], ev2[1], ev2[2]});
if (top > ctx_.b_ceiling)
b = b.scaled(ctx_.b_ceiling / top);
const double t = b.trace() / 3.0;
ctx_.atoms[i]->b_iso = static_cast<float>(t);
AtomAniso &a = ctx_.aniso[i];
a.on = true;
a.dev = {static_cast<float>((b.u11 - t) / gemmi::u_to_b()), static_cast<float>((b.u22 - t) / gemmi::u_to_b()),
static_cast<float>((b.u33 - t) / gemmi::u_to_b()), static_cast<float>(b.u12 / gemmi::u_to_b()),
static_cast<float>(b.u13 / gemmi::u_to_b()), static_cast<float>(b.u23 / gemmi::u_to_b())};
}
Context &ctx_;
Mode mode_;
const std::vector<int> &param_atoms_;
double weight_;
const URestraints *restraints_;
std::vector<double> radius_;
};
// The minimiser's diagonal for the B factors of `param_atoms` (RefinementMinimiser::SetDiagonal): weight times the
// X-ray target's curvature in B plus the B restraints', each element at least a thousandth of the median. The
// positions take none: their restraints' curvature is that of stiff bonds and angles, which move atoms together, so
@@ -611,7 +849,8 @@ std::vector<double> BDiagonal(const std::vector<double> &xray, const std::vector
}
// `minimiser` carries its memory from the call before; `diagonal` is the curvature it scales the steps by (none: 1).
void Minimise(Context &ctx, const AtomTarget &target, int iterations, RefinementMinimiser &minimiser,
template <typename Target>
void Minimise(Context &ctx, const Target &target, int iterations, RefinementMinimiser &minimiser,
std::vector<double> diagonal = {}) {
std::vector<double> p = target.Get();
if (p.empty())
@@ -714,7 +953,7 @@ public:
gradient[k] = 0;
for (int i : groups_[k]) {
const double b = base_[i] + q[k];
if (b > ctx_.b_floor && b < ctx_.b_ceiling)
if (b > ctx_.BLow(i) && b < ctx_.BHigh(i))
gradient[k] += g[i][3];
}
}
@@ -724,7 +963,7 @@ public:
void Place(const double *q) const {
for (size_t k = 0; k < groups_.size(); k++)
for (int i : groups_[k])
ctx_.atoms[i]->b_iso = static_cast<float>(std::clamp(base_[i] + q[k], ctx_.b_floor, ctx_.b_ceiling));
ctx_.atoms[i]->b_iso = static_cast<float>(std::clamp(base_[i] + q[k], ctx_.BLow(i), ctx_.BHigh(i)));
}
private:
@@ -972,6 +1211,7 @@ void MakeEngine(Context &ctx, bool gpu, const RefinementSettings &settings, Logg
#endif
if (!ctx.engine)
ctx.engine = MakeRefinementEngineCPU(ctx.cell, *ctx.sg, ctx.d_min, ctx.refl, settings.nthreads);
ctx.engine->SetAniso(&ctx.aniso);
}
// The riding hydrogens placed on their parents as they are, and the offsets and B they ride with.
@@ -1036,6 +1276,7 @@ std::unique_ptr<OwnSlot> MakeOwnSlot(const gemmi::Structure &before_topology, co
slot->topo = BuildRefinementTopology(ctx.st, settings.riding_hydrogens);
ctx.atoms = AtomList(ctx.Model());
ctx.SetRiderParents(slot->topo.rider_parent);
ctx.TakeAniso(); // each trial sets the anisotropy it starts from
#ifdef JFJOCH_USE_CUDA
if (main.engine->OnGPU()) {
try {
@@ -1054,6 +1295,7 @@ std::unique_ptr<OwnSlot> MakeOwnSlot(const gemmi::Structure &before_topology, co
#endif
if (!ctx.engine)
ctx.engine = MakeRefinementEngineCPU(ctx.cell, *ctx.sg, ctx.d_min, ctx.refl, nthreads);
ctx.engine->SetAniso(&ctx.aniso);
ctx.engine->SetRestraints(1, slot->topo.adp);
return slot;
}
@@ -1093,28 +1335,71 @@ std::vector<int> OwnSlotDevices(const Context &main, int wanted, const Refinemen
return std::vector<int>(std::max(beside, 0), -1);
}
// A displacement model a trial refines (RefinementSettings::adp_hypotheses): the B factors as they are (isotropic, or
// the shapes of the anisotropic atoms held), TLS bodies with a residual B per atom, or a U per atom.
struct AdpHypothesis {
enum class Kind { Isotropic, Tls, Anisotropic };
Kind kind = Kind::Isotropic;
std::string name;
std::vector<TlsBody> bodies; // Tls, with the part of each atom's B they hold at the start (Context::tls_b)
std::vector<double> tls_b;
std::vector<std::pair<int, int>> pairs; // Anisotropic: the pairs of the U restraints
};
// One weight trial: `cycles` macro-cycles from pos0 and b0, of the positions at X-ray weight w and, where adp_w is
// not 0, the B factors at X-ray weight adp_w. `fit`: what was fitted at pos0 and b0 already, null to fit it here.
struct TrialOutcome {
RefinementWeightTrial trial;
std::vector<gemmi::Position> pos;
std::vector<double> b;
std::vector<AtomAniso> aniso;
std::vector<TlsBody> tls;
std::vector<double> tls_b;
Context::Fit fit; // the scale, the solvent and the likelihood's parameters fitted where it ended
std::string line; // what the log says of it
bool ran = false; // false: left out, the refinement cancelled before it
};
TrialOutcome RunTrial(TrialSlot slot, const std::vector<int> &refinable, const std::vector<gemmi::Position> &pos0,
const std::vector<double> &b0, const Context::Fit *fit, double w, double adp_w, int cycles,
const RefinementSettings &settings) {
const std::vector<double> &b0, const std::vector<AtomAniso> &aniso0, const Context::Fit *fit,
double w, double adp_w, int cycles, const RefinementSettings &settings,
const AdpHypothesis *hyp = nullptr) {
Context &ctx = *slot.ctx;
RefinementTopology &topo = *slot.topo;
const RefinementStage stage(adp_w > 0 ? "rung.b" : "rung.xyz");
const RefinementStage stage(hyp != nullptr ? "rung.adp" : adp_w > 0 ? "rung.b" : "rung.xyz");
ctx.SetPositions(pos0);
ctx.SetB(b0);
ctx.SetAniso(aniso0);
ctx.tls.clear();
ctx.tls_b.clear();
const AdpHypothesis::Kind kind = hyp != nullptr ? hyp->kind : AdpHypothesis::Kind::Isotropic;
// The displacement model of the hypothesis: its TLS bodies applied to their atoms, or a U of its own for every
// refinable atom (the isotropic ones from their B) with the U restraints.
std::vector<int> u_atoms;
URestraints u_restraints;
if (kind == AdpHypothesis::Kind::Tls) {
PlaceRiders(ctx, topo);
ctx.tls = hyp->bodies;
ctx.tls_b = hyp->tls_b;
ctx.ApplyTls();
} else if (kind == AdpHypothesis::Kind::Anisotropic) {
if (ctx.aniso.size() != ctx.atoms.size())
ctx.aniso.assign(ctx.atoms.size(), AtomAniso());
std::vector<char> refines_u(ctx.atoms.size(), 0);
for (int i : refinable)
if (!ctx.atoms[i]->is_hydrogen()) {
ctx.aniso[i].on = true;
refines_u[i] = 1;
u_atoms.push_back(i);
}
ctx.SetAniso(ctx.aniso);
u_restraints = MakeURestraints(hyp->pairs, refines_u);
}
// the positions and the B factors in turn, each with a memory of its own kept from macro-cycle to macro-cycle
const AtomTarget xyz(ctx, AtomTarget::Mode::Xyz, refinable, w, 0), b(ctx, AtomTarget::Mode::B, refinable, adp_w, 1);
RefinementMinimiser xyz_minimiser(20, ctx.nthreads), b_minimiser(20, ctx.nthreads);
const AdpTarget tls_target(ctx, AdpTarget::Mode::Tls, refinable, 1.0, nullptr);
const AdpTarget u_target(ctx, AdpTarget::Mode::U, u_atoms, adp_w, &u_restraints);
RefinementMinimiser xyz_minimiser(20, ctx.nthreads), b_minimiser(20, ctx.nthreads), tls_minimiser(20, ctx.nthreads);
double xray_scale = 0;
for (int cycle = 0; cycle < cycles && !ctx.Cancelled(); cycle++) {
PlaceRiders(ctx, topo);
@@ -1137,7 +1422,36 @@ TrialOutcome RunTrial(TrialSlot slot, const std::vector<int> &refinable, const s
const RefinementStage s2("minimise.xyz");
Minimise(ctx, xyz, settings.xyz_iterations, xyz_minimiser);
}
if (adp_w > 0) {
if (kind == AdpHypothesis::Kind::Tls) {
ctx.ApplyTls(); // the bodies' U at where the atoms moved
const RefinementStage s4("minimise.tls");
Minimise(ctx, tls_target, settings.b_iterations, tls_minimiser);
}
if (adp_w > 0 && kind == AdpHypothesis::Kind::Anisotropic) {
// each U's six elements take a third of the curvature of its B
std::vector<double> xray = ctx.XrayBCurvatureShape();
if (cycle == 0)
xray_scale = ctx.CalibrateBCurvature(refinable, xray);
// and the restraints' curvature in 8 pi^2 U: per pair 2 / sigma^2 of the similarity and about a third of
// the rigid bond's, per isotropic atom 2 / sigma^2
std::vector<double> restraint(xray.size(), 0.0);
const double u2 = gemmi::u_to_b() * gemmi::u_to_b();
const double per_pair = (2.0 / (URestraints::SIGMA_SIMILAR * URestraints::SIGMA_SIMILAR) +
2.0 / (3.0 * URestraints::SIGMA_RIGID * URestraints::SIGMA_RIGID)) / u2;
for (const auto &[i, j] : u_restraints.pairs) {
restraint[i] += per_pair;
restraint[j] += per_pair;
}
for (int i : u_restraints.isotropic)
restraint[i] += 2.0 / (URestraints::SIGMA_ISOTROPIC * URestraints::SIGMA_ISOTROPIC) / u2;
for (double &x : xray)
x *= xray_scale / 3.0;
std::vector<double> diagonal;
for (double h : BDiagonal(xray, restraint, u_atoms, adp_w))
diagonal.insert(diagonal.end(), 6, h);
const RefinementStage s3("minimise.u");
Minimise(ctx, u_target, settings.b_iterations, b_minimiser, diagonal);
} else if (adp_w > 0) {
const RestraintSet b_set = BRestraintsNow(ctx, topo);
ctx.engine->SetRestraints(1, b_set);
std::vector<double> xray = ctx.XrayBCurvatureShape();
@@ -1165,15 +1479,26 @@ TrialOutcome RunTrial(TrialSlot slot, const std::vector<int> &refinable, const s
&topo.metal_bonds);
}
t.within_bounds = t.geometry.rmsd_bonds <= MAX_RMSD_BONDS && t.geometry.rmsd_angles <= MAX_RMSD_ANGLES;
out.line = fmt::format("X-ray weight {:.3g}, B {}: R-work {:.4f}, R-free {:.4f}, rmsd bonds {:.4f} A, "
if (hyp != nullptr) {
t.adp = hyp->name;
t.nll_free = ctx.Score().nll_free;
}
out.line = fmt::format("{}X-ray weight {:.3g}, {} {}: R-work {:.4f}, R-free {:.4f}{}, rmsd bonds {:.4f} A, "
"angles {:.2f} deg{}{}",
w, adp_w > 0 ? fmt::format("refined at weight {:.3g}", adp_w) : std::string("held"),
rf.r_work, rf.r_free, t.geometry.rmsd_bonds, t.geometry.rmsd_angles,
hyp != nullptr ? hyp->name + ", " : std::string(), w,
kind == AdpHypothesis::Kind::Anisotropic ? "U" : kind == AdpHypothesis::Kind::Tls ? "residual B" : "B",
adp_w > 0 ? fmt::format("refined at weight {:.3g}", adp_w) : std::string("held"),
rf.r_work, rf.r_free,
hyp != nullptr ? fmt::format(", -LL free {:.1f}", t.nll_free) : std::string(),
t.geometry.rmsd_bonds, t.geometry.rmsd_angles,
t.geometry.metal_bonds > 0 ? fmt::format(", metal bonds {:.3f} A", t.geometry.rmsd_metal)
: std::string(),
t.within_bounds ? "" : " (geometry out of bounds)");
out.pos = ctx.Positions();
out.b = ctx.B();
out.aniso = ctx.aniso;
out.tls = ctx.tls;
out.tls_b = ctx.tls_b;
out.fit = ctx.Fitted();
out.ran = true;
return out;
@@ -1185,8 +1510,9 @@ TrialOutcome RunTrial(TrialSlot slot, const std::vector<int> &refinable, const s
std::vector<TrialOutcome> RunTrials(const std::vector<TrialSlot> &trial_slots,
const std::vector<std::pair<double, double>> &trials,
const std::vector<int> &refinable, const std::vector<gemmi::Position> &pos0,
const std::vector<double> &b0, const Context::Fit *fit, int cycles,
const RefinementSettings &settings) {
const std::vector<double> &b0, const std::vector<AtomAniso> &aniso0,
const Context::Fit *fit, int cycles, const RefinementSettings &settings,
const std::vector<const AdpHypothesis *> &hyps = {}) {
std::vector<TrialOutcome> out(trials.size());
std::vector<std::exception_ptr> errors(trial_slots.size());
// the threads shared among the slots that run (the outcome does not depend on how many each has)
@@ -1206,8 +1532,8 @@ std::vector<TrialOutcome> RunTrials(const std::vector<TrialSlot> &trial_slots,
#endif
for (size_t j = s; j < trials.size(); j += trial_slots.size())
if (const size_t i = order[j]; !trial_slots[s].ctx->Cancelled())
out[i] = RunTrial(trial_slots[s], refinable, pos0, b0, fit, trials[i].first, trials[i].second,
cycles, settings);
out[i] = RunTrial(trial_slots[s], refinable, pos0, b0, aniso0, fit, trials[i].first,
trials[i].second, cycles, settings, hyps.empty() ? nullptr : hyps[i]);
} catch (...) {
errors[s] = std::current_exception();
}
@@ -1228,6 +1554,204 @@ std::vector<TrialOutcome> RunTrials(const std::vector<TrialSlot> &trial_slots,
namespace {
// TLS bodies with their atoms (the refinable ones but hydrogens), about their centroid, T = t I with t half of what
// the lowest B among them holds, L = S = 0: so they leave every B as it is (tls_b), the residual taking the rest.
AdpHypothesis TlsHypothesis(const Context &ctx, std::string name, std::vector<std::vector<int>> groups) {
AdpHypothesis h;
h.kind = AdpHypothesis::Kind::Tls;
h.name = std::move(name);
h.tls_b.assign(ctx.atoms.size(), 0.0);
for (std::vector<int> &atoms : groups) {
if (atoms.empty())
continue;
TlsBody body;
body.id = std::to_string(h.bodies.size() + 1);
double b_min = ctx.b_ceiling;
for (int i : atoms) {
body.origin += ctx.atoms[i]->pos;
b_min = std::min<double>(b_min, ctx.atoms[i]->b_iso);
}
body.origin /= static_cast<double>(atoms.size());
const double t = 0.5 * b_min / gemmi::u_to_b();
body.T = {t, t, t, 0, 0, 0};
for (int i : atoms)
h.tls_b[i] = gemmi::u_to_b() * t;
body.atoms = std::move(atoms);
h.bodies.push_back(std::move(body));
}
return h;
}
// The displacement models the refinement chooses among once its weights are chosen, from the model as it stands: the
// TLS bodies of the model's own groups (`input`, where it has them); one body per chain of at least MIN_TLS_ATOMS
// refinable atoms, and the chains cut into 2, 4, 8 and 16 TLSMD segments of at least MIN_TLS_RESIDUES residues
// (TlsSegments) - each cut a hypothesis where it differs from the one before; and a U per atom.
std::vector<AdpHypothesis> AdpHypotheses(const Context &ctx, const RefinementTopology &topo,
const std::vector<TlsBody> &input) {
constexpr size_t MIN_TLS_ATOMS = 50;
constexpr int MIN_TLS_RESIDUES = 10;
std::vector<AdpHypothesis> out;
if (!input.empty()) {
AdpHypothesis h;
h.kind = AdpHypothesis::Kind::Tls;
h.name = fmt::format("TLS, the model's {} groups", input.size());
h.bodies = input;
h.tls_b.assign(ctx.atoms.size(), 0.0);
for (const TlsBody &body : h.bodies)
for (int i : body.atoms)
h.tls_b[i] = gemmi::u_to_b() * body.U(ctx.atoms[i]->pos).trace() / 3.0;
out.push_back(std::move(h));
}
// per chain its residues' refinable atoms but hydrogens
std::vector<std::vector<std::vector<int>>> chains;
int i = 0;
for (const gemmi::Chain &ch : ctx.st.models[0].chains) {
std::vector<std::vector<int>> residues;
size_t atoms = 0;
for (const gemmi::Residue &r : ch.residues) {
std::vector<int> ra;
for (size_t k = 0; k < r.atoms.size(); k++, i++)
if (!r.is_water() && topo.refinable[i] && !r.atoms[k].is_hydrogen())
ra.push_back(i);
if (!ra.empty()) {
atoms += ra.size();
residues.push_back(std::move(ra));
}
}
if (atoms >= MIN_TLS_ATOMS)
chains.push_back(std::move(residues));
}
if (!chains.empty()) {
const std::vector<gemmi::Position> pos = ctx.Positions();
const std::vector<double> b = ctx.B();
size_t last = 0;
for (int k : {1, 2, 4, 8, 16}) {
std::vector<std::vector<int>> groups;
for (const std::vector<std::vector<int>> &residues : chains) {
std::vector<int> starts = TlsSegments(residues, pos, b, k, MIN_TLS_RESIDUES);
starts.push_back(static_cast<int>(residues.size()));
for (size_t s = 0; s + 1 < starts.size(); s++) {
std::vector<int> g;
for (int r = starts[s]; r < starts[s + 1]; r++)
g.insert(g.end(), residues[r].begin(), residues[r].end());
groups.push_back(std::move(g));
}
}
if (groups.size() == last)
continue;
last = groups.size();
std::string name = k == 1 ? fmt::format("TLS, {} groups (one per chain)", groups.size())
: fmt::format("TLS, {} groups (chains in up to {} segments)", groups.size(), k);
out.push_back(TlsHypothesis(ctx, std::move(name), std::move(groups)));
}
}
AdpHypothesis aniso;
aniso.kind = AdpHypothesis::Kind::Anisotropic;
aniso.name = "anisotropic";
for (const RestraintRecord &rec : topo.adp.records)
if (rec.kind == static_cast<int>(RestraintKind::BSimilarity))
aniso.pairs.emplace_back(rec.atom[0], rec.atom[1]);
out.push_back(std::move(aniso));
return out;
}
// The U X-ray weight that balances the gradients, as BalancedADPWeight() does for B: the U restraints' gradient with
// every U given an anisotropy of 20% of its B at random, over the X-ray target's, for the atoms that refine a U.
double BalancedUWeight(Context &ctx, const URestraints &restraints, const std::vector<int> &u_atoms) {
std::vector<std::array<double, 4>> gx;
ctx.engine->Evaluate(ctx.Model(), &gx);
const std::vector<std::array<double, 6>> gu = FoldedUGradient(ctx, gx);
std::vector<std::array<double, 6>> u = ctx.UOfAtoms();
std::mt19937 rng(3);
std::normal_distribution<double> shake(0.0, 0.2);
for (int i : u_atoms) {
const double t = (u[i][0] + u[i][1] + u[i][2]) / 3.0;
for (int c = 0; c < 6; c++)
u[i][c] += t * shake(rng);
}
std::vector<std::array<double, 6>> gr(u.size(), {0, 0, 0, 0, 0, 0});
EvaluateURestraints(restraints, ctx.Positions(), u, &gr);
double nx = 0, ng = 0;
for (int i : u_atoms)
for (int c = 0; c < 6; c++) {
nx += gu[i][c] * gu[i][c];
ng += gr[i][c] * gr[i][c];
}
return nx > 0 ? std::sqrt(ng / nx) : 1.0;
}
// The displacements of the model as given, read where it has TLS groups whose atoms bring no U of their own: their U
// from the groups, its B column taken as the total B (the U's equivalent isotropic B, as phenix.refine writes it) or as
// the residual to the groups (as REFMAC writes it), and the isotropic B as it stands - the one of the three whose
// R-free is lowest, beyond the tie in that order. `tls`: the groups, their atoms those of ctx.atoms but the hydrogens
// (which ride on their parents' B). Returns the R-factors of the reading kept.
// Following Winn, Isupov & Murshudov (2001) Acta Cryst. D57, 122-133 (the B column total or residual to the TLS)
RFactors ReadInputDisplacements(Context &ctx, const std::vector<TlsBody> &tls, Logger &logger, RefinementResult &result) {
int own = 0;
for (const AtomAniso &a : ctx.aniso)
own += a.on;
result.adp_input = own > 0 ? "ANISOU" : "isotropic";
result.adp_input_aniso_atoms = own;
std::vector<std::pair<int, gemmi::SMat33<double>>> tls_u; // atom, U_TLS
for (const TlsBody &body : tls)
for (int i : body.atoms)
if (ctx.aniso.empty() || !ctx.aniso[i].on)
tls_u.emplace_back(i, body.U(ctx.atoms[i]->pos));
RFactors r = ctx.Refit();
if (tls_u.empty())
return r;
const std::vector<double> b0 = ctx.B();
const std::vector<AtomAniso> aniso0 = ctx.aniso;
struct Reading {
const char *name;
std::vector<double> b;
std::vector<AtomAniso> aniso;
};
std::vector<Reading> readings;
for (int residual = 0; residual < 2; residual++) {
Reading read{residual ? "TLS, B residual" : "TLS, B total", b0,
aniso0.empty() ? std::vector<AtomAniso>(ctx.atoms.size()) : aniso0};
for (const auto &[i, u] : tls_u) {
const double t = u.trace() / 3.0;
read.aniso[i].on = true;
read.aniso[i].dev = {static_cast<float>(u.u11 - t), static_cast<float>(u.u22 - t),
static_cast<float>(u.u33 - t), static_cast<float>(u.u12), static_cast<float>(u.u13),
static_cast<float>(u.u23)};
if (residual)
read.b[i] = std::min(b0[i] + gemmi::u_to_b() * t, ctx.b_ceiling);
}
readings.push_back(std::move(read));
}
readings.push_back({result.adp_input == "ANISOU" ? "ANISOU" : "isotropic", b0, aniso0});
int best = -1;
std::vector<RFactors> rf;
for (size_t k = 0; k < readings.size(); k++) {
ctx.SetB(readings[k].b);
ctx.SetAniso(readings[k].aniso);
for (size_t i = 0; i < ctx.atoms.size(); i++)
if (!ctx.rider_parent.empty() && ctx.rider_parent[i] >= 0)
ctx.atoms[i]->b_iso = ctx.atoms[ctx.rider_parent[i]]->b_iso;
ctx.ForgetStarts();
rf.push_back(ctx.Refit());
logger.Info("Refinement: the model's {} TLS groups ({} atoms) read as {}: R-work {:.4f}, R-free {:.4f}", tls.size(),
tls_u.size(), readings[k].name, rf[k].r_work, rf[k].r_free);
if (best < 0 || rf[k].r_free < rf[best].r_free - R_FREE_TIE)
best = static_cast<int>(k);
}
ctx.SetB(readings[best].b);
ctx.SetAniso(readings[best].aniso);
for (size_t i = 0; i < ctx.atoms.size(); i++)
if (!ctx.rider_parent.empty() && ctx.rider_parent[i] >= 0)
ctx.atoms[i]->b_iso = ctx.atoms[ctx.rider_parent[i]]->b_iso;
ctx.ForgetStarts();
r = ctx.Refit();
if (best < 2) {
result.adp_input = readings[best].name;
result.adp_input_aniso_atoms = own + static_cast<int>(tls_u.size());
}
return r;
}
// One refinement of a model whose atoms stay as they are (RefineModel adds and removes waters between calls).
RefinementResult RefineOnce(const gemmi::Structure &placed, const std::vector<MergedReflection> &merged,
const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg, double wavelength_A,
@@ -1245,7 +1769,14 @@ RefinementResult RefineOnce(const gemmi::Structure &placed, const std::vector<Me
result.atoms = static_cast<int>(ctx.atoms.size());
double b_max = 0;
for (const gemmi::Atom *a : ctx.atoms)
b_max = std::max<double>(b_max, a->b_iso);
b_max = std::max<double>(b_max, a->aniso.nonzero() ? AnisoRadiusB(a->aniso.scaled(gemmi::u_to_b())) : a->b_iso);
// The model's TLS groups, where it has them; read as the residual to them, its B may go up by their trace.
const std::vector<TlsBody> tls_bodies =
ctx.st.meta.get_tls_groups() != nullptr ? TlsBodiesFromModel(ctx.Model(), *ctx.st.meta.get_tls_groups())
: std::vector<TlsBody>();
for (const TlsBody &body : tls_bodies)
for (int i : body.atoms)
b_max = std::max(b_max, ctx.atoms[i]->b_iso + gemmi::u_to_b() * body.U(ctx.atoms[i]->pos).trace() / 3.0);
ctx.b_ceiling = std::clamp(2 * b_max, B_CEILING_MIN, B_CEILING_MAX);
if (ctx.atoms.empty()) {
result.failure_reason = "the model has no atoms";
@@ -1285,10 +1816,17 @@ RefinementResult RefineOnce(const gemmi::Structure &placed, const std::vector<Me
PlaceRiders(ctx, topo);
}
ctx.TakeAniso();
MakeEngine(ctx, settings.gpu, settings, logger);
result.on_gpu = ctx.engine->OnGPU();
RFactors r = ctx.Refit();
std::vector<TlsBody> input_tls;
if (ctx.st.meta.get_tls_groups() != nullptr) {
input_tls = TlsBodiesFromModel(ctx.Model(), *ctx.st.meta.get_tls_groups());
for (TlsBody &body : input_tls)
std::erase_if(body.atoms, [&](int i) { return ctx.atoms[i]->is_hydrogen(); });
}
RFactors r = ReadInputDisplacements(ctx, input_tls, logger, result);
result.n_work = r.n_work;
result.n_free = r.n_free;
result.r_work_start = r.r_work;
@@ -1317,6 +1855,22 @@ RefinementResult RefineOnce(const gemmi::Structure &placed, const std::vector<Me
result.b_ceiling = ctx.b_ceiling;
result.k_sol = ctx.k_sol;
result.b_sol = ctx.b_sol;
// the anisotropic atoms' U into the model written out (ANISOU), and the TLS groups refined (the model's own
// stay where the U of their atoms was held, and go where every atom refined a U of its own)
ApplyAniso(ctx.Model(), ctx.aniso);
for (const AtomAniso &a : ctx.aniso)
result.aniso_atoms += a.on;
if (!ctx.tls.empty()) {
if (ctx.st.meta.refinement.empty())
ctx.st.meta.refinement.emplace_back();
ctx.st.meta.refinement[0].tls_groups = TlsBodiesToGroups(ctx.Model(), ctx.tls);
for (size_t k = 1; k < ctx.st.meta.refinement.size(); k++)
ctx.st.meta.refinement[k].tls_groups.clear();
} else if (result.adp_model == "anisotropic")
for (gemmi::RefinementInfo &ref : ctx.st.meta.refinement)
ref.tls_groups.clear();
if (const std::vector<gemmi::TlsGroup> *groups = ctx.st.meta.get_tls_groups())
result.tls_groups = static_cast<int>(groups->size());
result.model = std::make_shared<gemmi::Structure>(ctx.st);
result.ok = true;
result.seconds = std::chrono::duration<double>(std::chrono::steady_clock::now() - start).count();
@@ -1464,7 +2018,8 @@ RefinementResult RefineOnce(const gemmi::Structure &placed, const std::vector<Me
for (double rung : settings.weight_ladder)
for (double adp : settings.adp_ladder)
ladder.emplace_back(w0 * rung, adp_w0 * adp);
std::vector<TrialOutcome> outcomes = RunTrials(trial_slots, ladder, refinable, pos0, b0, nullptr,
const std::vector<AtomAniso> aniso0 = ctx.aniso;
std::vector<TrialOutcome> outcomes = RunTrials(trial_slots, ladder, refinable, pos0, b0, aniso0, nullptr,
settings.macro_cycles, settings);
int k = choose(outcomes, "ladder");
TrialOutcome state;
@@ -1494,8 +2049,8 @@ RefinementResult RefineOnce(const gemmi::Structure &placed, const std::vector<Me
if (2 * w <= w_max)
trials.emplace_back(2 * w, adp);
}
outcomes = RunTrials(trial_slots, trials, refinable, state.pos, state.b, &state.fit, settings.macro_cycles,
settings);
outcomes = RunTrials(trial_slots, trials, refinable, state.pos, state.b, state.aniso, &state.fit,
settings.macro_cycles, settings);
if (state_is_best)
from_best = outcomes;
k = choose(outcomes, fmt::format("step {}", step + 1).c_str());
@@ -1531,8 +2086,8 @@ RefinementResult RefineOnce(const gemmi::Structure &placed, const std::vector<Me
if (!outcomes[i].ran)
missing.push_back(trials[i]);
}
std::vector<TrialOutcome> ran = RunTrials(trial_slots, missing, refinable, state.pos, state.b, &state.fit,
settings.macro_cycles, settings);
std::vector<TrialOutcome> ran = RunTrials(trial_slots, missing, refinable, state.pos, state.b, state.aniso,
&state.fit, settings.macro_cycles, settings);
for (size_t i = 0, m = 0; i < trials.size(); i++)
if (!outcomes[i].ran)
outcomes[i] = std::move(ran[m++]);
@@ -1545,12 +2100,85 @@ RefinementResult RefineOnce(const gemmi::Structure &placed, const std::vector<Me
if (k >= 0 && displaces(outcomes[k].trial, state.trial))
state = std::move(outcomes[k]);
}
// 4. The displacement model, decided late, on the free set: from where the weights ended, the B factors refined
// on as they are, beside TLS bodies (the model's own groups, one per chain, the chains in TLSMD segments) with a
// residual B, and a U per atom at three weights - each trial a macro-cycle longer than a step. One displaces
// the one before it only where its geometry is within bounds and both its R-free (beyond the tie) and the -ln L
// of the free reflections are lower: the free set is the test of the parameters it adds.
// Following Merritt (2012) Acta Cryst. D68, 468-477 (the displacement model the data support, decided by test)
if (settings.adp_hypotheses && state.ran && !ctx.Cancelled()) {
const RefinementStage adp_stage("adp_hypotheses");
ctx.SetPositions(state.pos);
ctx.SetB(state.b);
ctx.SetAniso(state.aniso);
ctx.Restore(state.fit);
PlaceRiders(ctx, topo);
std::vector<AdpHypothesis> hyps = AdpHypotheses(ctx, topo, input_tls);
AdpHypothesis iso;
iso.name = result.adp_input == "isotropic" ? "isotropic" : "the model's U held";
const double w = state.trial.weight, adp = state.trial.adp_weight > 0 ? state.trial.adp_weight : adp_w0;
// the U weight balanced on the model with every refinable atom given a U
double u_w0 = adp;
{
const std::vector<AtomAniso> keep = ctx.aniso;
std::vector<AtomAniso> all = keep.empty() ? std::vector<AtomAniso>(ctx.atoms.size()) : keep;
std::vector<char> refines_u(ctx.atoms.size(), 0);
std::vector<int> u_atoms;
for (int i : refinable)
if (!ctx.atoms[i]->is_hydrogen()) {
all[i].on = true;
refines_u[i] = 1;
u_atoms.push_back(i);
}
ctx.SetAniso(all);
u_w0 = BalancedUWeight(ctx, MakeURestraints(hyps.back().pairs, refines_u), u_atoms);
ctx.SetAniso(keep);
}
std::vector<std::pair<double, double>> trials = {{w, adp}};
std::vector<const AdpHypothesis *> which = {&iso};
for (const AdpHypothesis &h : hyps) {
if (h.kind == AdpHypothesis::Kind::Anisotropic) {
for (double f : {0.25, 1.0, 4.0}) {
trials.emplace_back(w, u_w0 * f);
which.push_back(&h);
}
} else {
trials.emplace_back(w, adp);
which.push_back(&h);
}
}
outcomes = RunTrials(trial_slots, trials, refinable, state.pos, state.b, state.aniso, &state.fit,
settings.macro_cycles + 1, settings, which);
int best = -1;
for (size_t k = 0; k < outcomes.size(); k++) {
const TrialOutcome &o = outcomes[k];
if (!o.ran)
continue;
logger.Info("Refinement, displacements: {}", o.line);
result.adp_trials.push_back(o.trial);
if (best < 0 || (o.trial.within_bounds && (!outcomes[best].trial.within_bounds ||
(o.trial.r_free < outcomes[best].trial.r_free - R_FREE_TIE &&
o.trial.nll_free < outcomes[best].trial.nll_free))))
best = static_cast<int>(k);
}
// The B factors refined on are a step further on the same weights: kept where they displace where the
// weights ended, as any step.
if (best >= 0) {
result.adp_model = outcomes[best].trial.adp;
logger.Info("Refinement: displacement model {}", result.adp_model);
}
if (best > 0 || (best == 0 && displaces(outcomes[0].trial, state.trial)))
state = std::move(outcomes[best]);
}
for (size_t s = 1; s < trial_slots.size(); s++)
ctx.evaluations += trial_slots[s].ctx->evaluations;
ctx.nthreads = main_threads;
if (state.ran) {
ctx.SetPositions(state.pos);
ctx.SetB(state.b);
ctx.SetAniso(state.aniso);
ctx.tls = state.tls;
ctx.tls_b = state.tls_b;
PlaceRiders(ctx, topo);
result.restrained = true;
result.weight = state.trial.weight;
@@ -1616,7 +2244,10 @@ RefinementResult RefineModel(const gemmi::Structure &placed, const std::vector<M
}
gemmi::Structure changed = *result.model;
ApplySolventEdits(changed, picked.edits);
RefinementResult next = RefineOnce(changed, merged, cell, sg, wavelength_A, logger, settings);
// the displacement model as the first refinement chose it: its U held, the B factors refined
RefinementSettings round_settings = settings;
round_settings.adp_hypotheses = false;
RefinementResult next = RefineOnce(changed, merged, cell, sg, wavelength_A, logger, round_settings);
if (!next.ok || !next.restrained || next.cancelled || !(next.r_free < result.r_free)) {
logger.Info("Solvent build, round {}: not kept (R-free {:.4f}, before {:.4f})", round, next.r_free,
result.r_free);
@@ -1643,6 +2274,10 @@ RefinementResult RefineModel(const gemmi::Structure &placed, const std::vector<M
next.group_b = result.group_b;
next.r_work_placement = result.r_work_placement;
next.r_free_placement = result.r_free_placement;
next.adp_input = result.adp_input;
next.adp_input_aniso_atoms = result.adp_input_aniso_atoms;
next.adp_model = result.adp_model;
next.adp_trials = result.adp_trials;
result = std::move(next);
}
b.added = static_cast<int>(added.size());
@@ -62,6 +62,10 @@ struct RefinementSettings {
// Waters added at the difference map's peaks and removed where the maps no longer hold them, between refinements
// (SolventBuild.h), to SOLVENT_BUILD_D_MIN.
bool solvent = true;
// 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;
bool gpu = true;
size_t nthreads = 1;
size_t gpu_resident_bytes = 0; // held on the card by what runs beside
@@ -111,8 +115,10 @@ 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 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 {
@@ -127,6 +133,11 @@ struct RefinementResult {
// 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
@@ -138,6 +149,11 @@ struct RefinementResult {
double weight = 0;
double adp_weight = 0; // 0: the B factors kept as placed
std::vector<RefinementWeightTrial> trials;
// 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;
@@ -0,0 +1,319 @@
// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only
#include "RefinementADP.h"
#include <algorithm>
#include <cmath>
#include <limits>
#include <Eigen/Dense>
namespace {
constexpr double DEG = 3.14159265358979323846 / 180.0;
} // namespace
gemmi::SMat33<double> TlsBody::U(const gemmi::Position &pos) const {
const double x = pos.x - origin.x, y = pos.y - origin.y, z = pos.z - origin.z;
const gemmi::Mat33 a(0, z, -y, -z, 0, x, y, -x, 0);
const gemmi::Mat33 ala = a.multiply(L.as_mat33()).multiply(a.transpose());
const gemmi::Mat33 as = a.multiply(S);
gemmi::SMat33<double> u = T;
u.u11 += ala[0][0] + 2 * as[0][0];
u.u22 += ala[1][1] + 2 * as[1][1];
u.u33 += ala[2][2] + 2 * as[2][2];
u.u12 += ala[0][1] + as[0][1] + as[1][0];
u.u13 += ala[0][2] + as[0][2] + as[2][0];
u.u23 += ala[1][2] + as[1][2] + as[2][1];
return u;
}
std::array<double, TlsBody::PARAMETERS> TlsBody::Get() const {
return {T.u11, T.u22, T.u33, T.u12, T.u13, T.u23, L.u11, L.u22, L.u33, L.u12, L.u13, L.u23,
S[0][0], S[1][1], S[0][1], S[0][2], S[1][0], S[1][2], S[2][0], S[2][1]};
}
void TlsBody::Set(const std::array<double, PARAMETERS> &p) {
T = {p[0], p[1], p[2], p[3], p[4], p[5]};
L = {p[6], p[7], p[8], p[9], p[10], p[11]};
S = gemmi::Mat33(p[12], p[14], p[15], p[16], p[13], p[17], p[18], p[19], -(p[12] + p[13]));
}
std::array<std::array<double, 6>, TlsBody::PARAMETERS> TlsBody::Derivatives(const gemmi::Position &pos) const {
std::array<std::array<double, 6>, PARAMETERS> d{};
TlsBody unit;
unit.origin = origin;
for (int k = 0; k < PARAMETERS; k++) {
std::array<double, PARAMETERS> p{};
p[k] = 1;
unit.Set(p);
const gemmi::SMat33<double> u = unit.U(pos);
d[k] = {u.u11, u.u22, u.u33, u.u12, u.u13, u.u23};
}
return d;
}
std::array<double, TlsBody::PARAMETERS> TlsBody::Gradient(const std::vector<gemmi::Atom *> &atom_list,
const std::vector<std::array<double, 6>> &grad_u) const {
gemmi::Mat33 mt(0), ml(0), ms(0);
for (int i : atoms) {
const std::array<double, 6> &g = grad_u[i];
const gemmi::Mat33 gm(g[0], g[3] / 2, g[4] / 2, g[3] / 2, g[1], g[5] / 2, g[4] / 2, g[5] / 2, g[2]);
const gemmi::Position &pos = atom_list[i]->pos;
const double x = pos.x - origin.x, y = pos.y - origin.y, z = pos.z - origin.z;
const gemmi::Mat33 a(0, z, -y, -z, 0, x, y, -x, 0);
const gemmi::Mat33 atg = a.transpose().multiply(gm);
const gemmi::Mat33 atga = atg.multiply(a);
for (int u = 0; u < 3; u++)
for (int v = 0; v < 3; v++) {
mt[u][v] += gm[u][v];
ml[u][v] += atga[u][v];
ms[u][v] += 2 * atg[u][v];
}
}
// symmetric T and L: an off-diagonal parameter takes both of its elements; S33 = -(S11 + S22)
return {mt[0][0], mt[1][1], mt[2][2], 2 * mt[0][1], 2 * mt[0][2], 2 * mt[1][2],
ml[0][0], ml[1][1], ml[2][2], 2 * ml[0][1], 2 * ml[0][2], 2 * ml[1][2],
ms[0][0] - ms[2][2], ms[1][1] - ms[2][2], ms[0][1], ms[0][2], ms[1][0], ms[1][2], ms[2][0], ms[2][1]};
}
std::vector<TlsBody> TlsBodiesFromModel(const gemmi::Model &model, const std::vector<gemmi::TlsGroup> &groups) {
std::vector<TlsBody> bodies;
for (const gemmi::TlsGroup &g : groups) {
TlsBody b;
b.id = g.id;
b.origin = g.origin;
b.T = g.T;
b.L = g.L.scaled(DEG * DEG);
for (int i = 0; i < 3; i++)
for (int j = 0; j < 3; j++)
b.S[i][j] = g.S[i][j] * DEG;
// the trace of S does not enter U; S33 is taken as -(S11 + S22)
const double t = (b.S[0][0] + b.S[1][1] + b.S[2][2]) / 3.0;
for (int i = 0; i < 3; i++)
b.S[i][i] -= t;
bool finite = std::isfinite(b.origin.x) && std::isfinite(b.origin.y) && std::isfinite(b.origin.z);
for (double v : b.Get())
finite = finite && std::isfinite(v);
if (finite)
bodies.push_back(b);
}
// the atoms of each group's residues, an atom to the first group that names it
int i = 0;
for (const gemmi::Chain &ch : model.chains)
for (const gemmi::Residue &r : ch.residues) {
int body = -1;
if (!r.is_water())
for (size_t k = 0; k < groups.size() && body < 0; k++) {
for (const gemmi::TlsGroup::Selection &sel : groups[k].selections)
if (sel.chain == ch.name && !(r.seqid < sel.res_begin) && !(sel.res_end < r.seqid))
body = static_cast<int>(k);
}
// the group index among those kept
int kept = -1;
if (body >= 0)
for (size_t k = 0; k < bodies.size(); k++)
if (bodies[k].id == groups[body].id)
kept = static_cast<int>(k);
for (size_t a = 0; a < r.atoms.size(); a++, i++)
if (kept >= 0)
bodies[kept].atoms.push_back(i);
}
std::erase_if(bodies, [](const TlsBody &b) { return b.atoms.empty(); });
return bodies;
}
std::vector<gemmi::TlsGroup> TlsBodiesToGroups(const gemmi::Model &model, const std::vector<TlsBody> &bodies) {
// per atom its body, per residue the body of its first atom in one
std::vector<int> atom_body;
for (const gemmi::Chain &ch : model.chains)
for (const gemmi::Residue &r : ch.residues)
atom_body.insert(atom_body.end(), r.atoms.size(), -1);
for (size_t b = 0; b < bodies.size(); b++)
for (int a : bodies[b].atoms)
atom_body[a] = static_cast<int>(b);
std::vector<gemmi::TlsGroup> groups(bodies.size());
for (size_t b = 0; b < bodies.size(); b++) {
gemmi::TlsGroup &g = groups[b];
g.id = std::to_string(b + 1);
g.num_id = static_cast<short>(b + 1);
g.origin = bodies[b].origin;
g.T = bodies[b].T;
g.L = bodies[b].L.scaled(1.0 / (DEG * DEG));
for (int i = 0; i < 3; i++)
for (int j = 0; j < 3; j++)
g.S[i][j] = bodies[b].S[i][j] / DEG;
}
size_t i = 0;
for (const gemmi::Chain &ch : model.chains) {
int open = -1; // the body whose run of residues is open in this chain
for (const gemmi::Residue &r : ch.residues) {
int body = -1;
for (size_t a = 0; a < r.atoms.size(); a++)
if (atom_body[i + a] >= 0 && body < 0)
body = atom_body[i + a];
i += r.atoms.size();
if (body >= 0 && body == open)
groups[body].selections.back().res_end = r.seqid;
else if (body >= 0) {
groups[body].selections.push_back({ch.name, r.seqid, r.seqid, ""});
}
open = body;
}
}
return groups;
}
URestraints MakeURestraints(const std::vector<std::pair<int, int>> &b_similarity_pairs, const std::vector<char> &refines_u) {
URestraints r;
std::vector<char> paired(refines_u.size(), 0);
for (const auto &[i, j] : b_similarity_pairs)
if (refines_u[i] && refines_u[j]) {
r.pairs.emplace_back(i, j);
paired[i] = paired[j] = 1;
}
for (size_t i = 0; i < refines_u.size(); i++)
if (refines_u[i] && !paired[i])
r.isotropic.push_back(static_cast<int>(i));
return r;
}
double EvaluateURestraints(const URestraints &r, const std::vector<gemmi::Position> &pos,
const std::vector<std::array<double, 6>> &u, std::vector<std::array<double, 6>> *grad) {
// a symmetric matrix as its six elements: the off-diagonal ones count twice in a Frobenius norm, and a derivative
// with respect to U_ab (a != b) is the sum of the matrix elements ab and ba
double e = 0;
const double wr = 1.0 / (URestraints::SIGMA_RIGID * URestraints::SIGMA_RIGID);
const double ws = 1.0 / (URestraints::SIGMA_SIMILAR * URestraints::SIGMA_SIMILAR);
const double wi = 1.0 / (URestraints::SIGMA_ISOTROPIC * URestraints::SIGMA_ISOTROPIC);
for (const auto &[i, j] : r.pairs) {
std::array<double, 6> d;
for (int c = 0; c < 6; c++)
d[c] = u[i][c] - u[j][c];
gemmi::Vec3 n = pos[j] - pos[i];
const double len = n.length();
std::array<double, 6> g{};
if (len > 0) {
n = n / len;
// v = (U_i - U_j) n
const double v[3] = {d[0] * n.x + d[3] * n.y + d[4] * n.z, d[3] * n.x + d[1] * n.y + d[5] * n.z,
d[4] * n.x + d[5] * n.y + d[2] * n.z};
const double nn[3] = {n.x, n.y, n.z};
e += wr * (v[0] * v[0] + v[1] * v[1] + v[2] * v[2]);
g[0] += 2 * wr * v[0] * nn[0];
g[1] += 2 * wr * v[1] * nn[1];
g[2] += 2 * wr * v[2] * nn[2];
g[3] += 2 * wr * (v[0] * nn[1] + v[1] * nn[0]);
g[4] += 2 * wr * (v[0] * nn[2] + v[2] * nn[0]);
g[5] += 2 * wr * (v[1] * nn[2] + v[2] * nn[1]);
}
e += ws * (d[0] * d[0] + d[1] * d[1] + d[2] * d[2] + 2 * (d[3] * d[3] + d[4] * d[4] + d[5] * d[5]));
for (int c = 0; c < 3; c++)
g[c] += 2 * ws * d[c];
for (int c = 3; c < 6; c++)
g[c] += 4 * ws * d[c];
if (grad != nullptr)
for (int c = 0; c < 6; c++) {
(*grad)[i][c] += g[c];
(*grad)[j][c] -= g[c];
}
}
for (int i : r.isotropic) {
const double t = (u[i][0] + u[i][1] + u[i][2]) / 3.0;
const double d[6] = {u[i][0] - t, u[i][1] - t, u[i][2] - t, u[i][3], u[i][4], u[i][5]};
e += wi * (d[0] * d[0] + d[1] * d[1] + d[2] * d[2] + 2 * (d[3] * d[3] + d[4] * d[4] + d[5] * d[5]));
if (grad != nullptr) {
// d/dU_aa of the traceless part's norm: 2 d_aa (the trace's own derivative sums to zero against d)
for (int c = 0; c < 3; c++)
(*grad)[i][c] += 2 * wi * d[c];
for (int c = 3; c < 6; c++)
(*grad)[i][c] += 4 * wi * d[c];
}
}
return e;
}
std::vector<int> TlsSegments(const std::vector<std::vector<int>> &residue_atoms, const std::vector<gemmi::Position> &pos,
const std::vector<double> &b, int k, int min_residues) {
using Mat = Eigen::Matrix<double, 10, 10>;
using Vec = Eigen::Matrix<double, 10, 1>;
const int n = static_cast<int>(residue_atoms.size());
k = std::min(k, n / std::max(min_residues, 1));
if (k <= 1)
return {0};
// The chain in blocks of consecutive residues (cuts between blocks only), so that a long chain stays a few hundred
// blocks: the dynamic programming is quadratic in them.
const int per_block = std::max(1, (n + 299) / 300);
const int min_blocks = std::max(1, (min_residues + per_block - 1) / per_block);
std::vector<int> block_first;
for (int r = 0; r < n; r += per_block)
block_first.push_back(r);
const int nb = static_cast<int>(block_first.size());
k = std::min(k, nb / min_blocks);
if (k <= 1)
return {0};
// the normal equations of each block, about the chain's centroid in units of 10 A (for their conditioning)
gemmi::Position c;
int count = 0;
for (const std::vector<int> &ra : residue_atoms)
for (int a : ra) {
c += pos[a];
count++;
}
if (count == 0)
return {0};
c /= count;
std::vector<Mat> ata(nb + 1, Mat::Zero());
std::vector<Vec> atb(nb + 1, Vec::Zero());
std::vector<double> bb(nb + 1, 0.0);
for (int q = 0; q < nb; q++) {
ata[q + 1] = ata[q];
atb[q + 1] = atb[q];
bb[q + 1] = bb[q];
const int last = q + 1 < nb ? block_first[q + 1] : n;
for (int r = block_first[q]; r < last; r++)
for (int a : residue_atoms[r]) {
const double x = (pos[a].x - c.x) / 10, y = (pos[a].y - c.y) / 10, z = (pos[a].z - c.z) / 10;
Vec f;
f << 1, x, y, z, x * x, y * y, z * z, x * y, x * z, y * z;
ata[q + 1] += f * f.transpose();
atb[q + 1] += b[a] * f;
bb[q + 1] += b[a] * b[a];
}
}
// the residual of the fit of blocks [i, j)
auto cost = [&](int i, int j) {
const Mat m = ata[j] - ata[i] + 1e-6 * Mat::Identity();
const Vec v = atb[j] - atb[i];
const Vec x = m.ldlt().solve(v);
return std::max(0.0, bb[j] - bb[i] - v.dot(x));
};
const double INF = std::numeric_limits<double>::infinity();
// best[s][j]: blocks [0, j) in s + 1 segments; from[s][j]: where the last of them starts
std::vector<std::vector<double>> best(k, std::vector<double>(nb + 1, INF));
std::vector<std::vector<int>> from(k, std::vector<int>(nb + 1, 0));
std::vector<std::vector<double>> seg(nb + 1, std::vector<double>(nb + 1, INF));
for (int i = 0; i < nb; i++)
for (int j = i + min_blocks; j <= nb; j++)
seg[i][j] = cost(i, j);
for (int j = min_blocks; j <= nb; j++)
best[0][j] = seg[0][j];
for (int s = 1; s < k; s++)
for (int j = (s + 1) * min_blocks; j <= nb; j++)
for (int i = s * min_blocks; i <= j - min_blocks; i++)
if (best[s - 1][i] + seg[i][j] < best[s][j]) {
best[s][j] = best[s - 1][i] + seg[i][j];
from[s][j] = i;
}
std::vector<int> starts;
int j = nb;
for (int s = k - 1; s > 0; s--) {
j = from[s][j];
starts.push_back(block_first[j]);
}
starts.push_back(0);
std::reverse(starts.begin(), starts.end());
return starts;
}
@@ -0,0 +1,78 @@
// 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/math.hpp"
#include "gemmi/metadata.hpp" // TlsGroup
#include "gemmi/model.hpp"
// The anisotropic side of restrained refinement (Refinement.h): TLS groups, and the restraints of atoms refined
// with a U of their own. Atoms are numbered in model order, as everywhere in the refinement.
// A group of atoms whose displacement is in part that of a rigid body: U_TLS(r) = T + A L A^T + A S + S^T A^T at
// r - origin = (x, y, z), A = [[0, z, -y], [-z, 0, x], [y, -x, 0]]. T in A^2, L in rad^2, S in rad A (the files'
// L and S are in degrees). The trace of S does not enter U: S33 is -(S11 + S22).
// Following Schomaker & Trueblood (1968) Acta Cryst. B24, 63-76, and Winn, Isupov & Murshudov (2001) Acta Cryst. D57, 122-133
struct TlsBody {
std::string id;
gemmi::Position origin;
gemmi::SMat33<double> T{0, 0, 0, 0, 0, 0}, L{0, 0, 0, 0, 0, 0};
gemmi::Mat33 S{0};
std::vector<int> atoms; // in model order
gemmi::SMat33<double> U(const gemmi::Position &pos) const;
// The 20 parameters (T11 T22 T33 T12 T13 T23, L likewise, S11 S22 S12 S13 S21 S23 S31 S32) and back.
static constexpr int PARAMETERS = 20;
std::array<double, PARAMETERS> Get() const;
void Set(const std::array<double, PARAMETERS> &p);
// dU/dp at `pos`: per parameter the six U11 U22 U33 U12 U13 U23 (U is linear in the parameters).
std::array<std::array<double, 6>, PARAMETERS> Derivatives(const gemmi::Position &pos) const;
// The derivatives of a target in the parameters, from its derivatives in U11 U22 U33 U12 U13 U23 of every atom
// (`grad_u`, in model order as `atoms`): sum over the body's atoms of grad_u . dU/dp, by the chain rule in matrix
// form - with G the symmetric matrix of an atom's derivatives (an off-diagonal element half the U_ab one),
// dT = tr(G dU) gives G for T, A^T G A for L and 2 A^T G for S.
std::array<double, PARAMETERS> Gradient(const std::vector<gemmi::Atom *> &atoms,
const std::vector<std::array<double, 6>> &grad_u) const;
};
// The TLS groups of a model's metadata (mmCIF _pdbx_refine_tls*, PDB REMARK 3) with their atoms: the atoms of every
// residue of a selection's chain from its first residue to its last, waters left out; a group without numbers or
// without atoms is left out.
std::vector<TlsBody> TlsBodiesFromModel(const gemmi::Model &model, const std::vector<gemmi::TlsGroup> &groups);
// The bodies as the metadata holds them (degrees), each selection one run of consecutive residues of a chain.
std::vector<gemmi::TlsGroup> TlsBodiesToGroups(const gemmi::Model &model, const std::vector<TlsBody> &bodies);
// The restraints of atoms refined with a U of their own (in A^2, U11 U22 U33 U12 U13 U23 per atom):
// - rigid bond, for each pair of bonded or angle-related atoms (the pairs of the B similarity): the difference of
// their U along the line between them, |(U_i - U_j) n|^2 / sigma^2, sigma 0.004 A^2 - the components U33, U13 and
// U23 of the difference in a frame whose z is n, as SHELXL's RIGU restrains them;
// - similarity, for the same pairs: |U_i - U_j|^2 (Frobenius) / sigma^2, sigma 0.04 A^2 (SHELXL's SIMU);
// - near isotropy, for an atom in no pair (a water, an ion): |U - tr(U)/3 I|^2 / sigma^2, sigma 0.1 A^2 (ISOR).
// Following Hirshfeld (1976) Acta Cryst. A32, 239-244, and Thorn, Dittrich & Sheldrick (2012) Acta Cryst. A68, 448-451
struct URestraints {
static constexpr double SIGMA_RIGID = 0.004, SIGMA_SIMILAR = 0.04, SIGMA_ISOTROPIC = 0.1;
std::vector<std::pair<int, int>> pairs;
std::vector<int> isotropic;
};
// The pairs of `b_similarity_pairs` whose atoms both refine a U, and the refined atoms of no pair.
URestraints MakeURestraints(const std::vector<std::pair<int, int>> &b_similarity_pairs, const std::vector<char> &refines_u);
// The sum of E at the positions `pos` and the U of every atom `u` (atoms that refine none are not read) and, with
// `grad` (sized as u), dE/dU per atom added to it.
double EvaluateURestraints(const URestraints &r, const std::vector<gemmi::Position> &pos,
const std::vector<std::array<double, 6>> &u, std::vector<std::array<double, 6>> *grad);
// Contiguous segments of a chain (TLSMD's idea): residues cut into `k` runs so that the B factors of each run's atoms
// are best fitted, in least squares, by an isotropic TLS model of its own - B(r) a general quadratic in r (tr(A L A^T)
// is tr(L)|r|^2 - r^T L r and tr(A S) linear in r, ten parameters as Painter & Merritt's isotropic TLS), by dynamic
// programming over where the runs end. `residue_atoms`: per residue (in chain order) its atoms, `pos` and `b` per
// atom. Returns per segment its first residue (the first is 0); fewer than k where the chain has too few residues for
// `min_residues` each.
// Following Painter & Merritt (2006) Acta Cryst. D62, 439-450
std::vector<int> TlsSegments(const std::vector<std::vector<int>> &residue_atoms, const std::vector<gemmi::Position> &pos,
const std::vector<double> &b, int k, int min_residues);
@@ -12,12 +12,36 @@
#include "gemmi/solmask.hpp" // SolventMasker, refmac_radius_for_bulk_solvent
#include "../../common/ParallelFor.h"
#include "ModelGrid.h" // AnisoRadiusB
#include "RefinementGPUEngine.h"
namespace {
using Table = gemmi::IT92<float>;
// One atom's density as the gather takes it, anisotropic with U (as PutModelDensityOnGrid() grids it), with `blur`.
ModelDensityAtom AnisoDensityAtom(const gemmi::Atom &atom, const gemmi::SMat33<float> &u,
const gemmi::DensityCalculator<Table, float> &dc) {
using CReal = Table::Coef::coef_type;
const auto &coef = Table::get(atom.element, atom.charge, atom.serial);
const auto aniso_b = u.scaled(CReal(gemmi::u_to_b())).added_kI(CReal(dc.blur));
const CReal b_max = AnisoRadiusB(aniso_b);
const double radius = dc.estimate_radius(coef.precalculate_density_iso(b_max, 0.0f), b_max);
const auto precal = coef.precalculate_density_aniso_b(aniso_b, 0.0f);
ModelDensityAtom a{};
a.occ = atom.occ;
a.aniso = 1;
a.radius = static_cast<float>(radius);
a.max_r2 = static_cast<float>(radius * radius);
for (int k = 0; k < 5; k++) {
a.a[k] = precal.a[k];
const gemmi::SMat33<float> &m = precal.b[k];
const float e[6] = {m.u11, m.u22, m.u33, m.u12, m.u13, m.u23};
std::copy(e, e + 6, a.b[k]);
}
return a;
}
// One atom's density as the gather takes it, isotropic, with `blur`.
ModelDensityAtom DensityAtom(const gemmi::Atom &atom, double b_iso, const gemmi::DensityCalculator<Table, float> &dc) {
using CReal = Table::Coef::coef_type;
@@ -231,8 +255,10 @@ public:
const RefinementStage stage("gpu.evaluate");
value = engine_->Evaluate(grad != nullptr);
}
if (grad != nullptr)
if (grad != nullptr) {
Widen(engine_->HostGradient(), atom_ptrs_.size(), *grad, nthreads_);
CopyUGradient();
}
return value;
}
@@ -249,8 +275,10 @@ public:
}
}, restraint_grad != nullptr, restraints);
}
if (grad != nullptr)
if (grad != nullptr) {
Widen(engine_->HostGradient(), atom_ptrs_.size(), *grad, nthreads_);
CopyUGradient();
}
if (restraint_grad != nullptr)
Widen(engine_->HostRestraintGradient(), atom_ptrs_.size(), *restraint_grad, nthreads_);
return value;
@@ -280,6 +308,19 @@ public:
}
private:
// The U derivatives of the last evaluation, where some atom is anisotropic.
void CopyUGradient() {
if (!AnyAniso()) {
grad_u_.clear();
return;
}
const float *g = engine_->HostUGradient();
grad_u_.resize(atom_ptrs_.size());
for (size_t i = 0; i < atom_ptrs_.size(); i++)
for (int c = 0; c < 6; c++)
grad_u_[i][c] = g[6 * i + c];
}
gemmi::DensityCalculator<Table, float> Calculator(double blur) const {
gemmi::DensityCalculator<Table, float> dc;
dc.d_min = d_min_;
@@ -302,16 +343,23 @@ private:
const int n = static_cast<int>(atoms.size());
if (last_b_.size() != atoms.size()) {
last_b_.assign(atoms.size(), NAN);
last_aniso_.assign(atoms.size(), AtomAniso());
atoms_.resize(atoms.size());
pos_.resize(atoms.size());
}
changed_.assign(atoms.size(), 0);
const int n3[3] = {setup_.grid.nu, setup_.grid.nv, setup_.grid.nw};
const std::vector<AtomAniso> *aniso = aniso_ != nullptr && aniso_->size() == atoms.size() ? aniso_ : nullptr;
ParallelChunks(n, ThreadsForWork(atoms.size(), nthreads_, 4096), [&](int lo, int hi) {
for (int i = lo; i < hi; i++) {
if (!(atoms[i]->b_iso == last_b_[i])) {
atoms_[i] = DensityAtom(*atoms[i], atoms[i]->b_iso, dc_);
const AtomAniso none;
const AtomAniso &an = aniso != nullptr ? (*aniso)[i] : none;
if (!(atoms[i]->b_iso == last_b_[i]) || an.on != last_aniso_[i].on ||
(an.on && an.dev != last_aniso_[i].dev)) {
atoms_[i] = an.on ? AnisoDensityAtom(*atoms[i], AnisoU(an, atoms[i]->b_iso), dc_)
: DensityAtom(*atoms[i], atoms[i]->b_iso, dc_);
last_b_[i] = atoms[i]->b_iso;
last_aniso_[i] = an;
changed_[i] = 1;
}
const gemmi::Fractional f = cell_.fractionalize(atoms[i]->pos);
@@ -337,6 +385,7 @@ private:
std::unique_ptr<RefinementGPUEngine> engine_;
std::vector<const gemmi::Atom *> atom_ptrs_; // the model's atoms at the last SetAtoms()
std::vector<float> last_b_;
std::vector<AtomAniso> last_aniso_;
PinnedVector<ModelDensityAtom> atoms_; // pinned: a B rung sends them every evaluation (37 MB on a large model)
PinnedVector<std::array<float, 4>> pos_;
std::vector<char> changed_;
@@ -157,8 +157,11 @@ __global__ void fmask_kernel(const float2 *fm, float norm, const int *row_hkl, i
// dT/d(x, y, z, B) of every atom: the derivative map summed against the derivatives of the atom's density over
// its box, one warp per atom, the lanes' sums added in a fixed tree. The box, the points within the radius and
// their offsets from the centre are the density gather's (ModelDensityGPU.cu).
// An anisotropic atom (rho = occ sum_q a_q exp(r^T M_q r), RefinementXray.cpp's AnisoGradient()) gives its B derivative
// along U + dB / (8 pi^2) I and, into grad_u, its six U derivatives; an isotropic one writes zeros there.
__global__ void atom_gradient_kernel(const ModelDensityAtom *atoms, const float4 *pos, const int *box, int n,
ModelDensityGPU::Geometry g, const float *map, float map_scale, float4 *grad) {
ModelDensityGPU::Geometry g, const float *map, float map_scale, float4 *grad,
float *grad_u) {
const int a = (blockIdx.x * blockDim.x + threadIdx.x) / WARP;
const int lane = threadIdx.x % WARP;
if (a >= n)
@@ -191,6 +194,8 @@ __global__ void atom_gradient_kernel(const ModelDensityAtom *atoms, const float4
}
};
float gx = 0, gy = 0, gz = 0, gb = 0;
float gu[6] = {0, 0, 0, 0, 0, 0};
const bool aniso = at.aniso != 0;
for (int t = lane; t < total; t += WARP, next()) {
const int u = lo[0] + iu, v = lo[1] + iv, w = lo[2] + iw;
const float fx = nctr[0] - static_cast<float>(u);
@@ -206,27 +211,57 @@ __global__ void atom_gradient_kernel(const ModelDensityAtom *atoms, const float4
const float value = map[(mu < g.nu ? mu : mu - g.nu) +
static_cast<size_t>(g.nu) * ((mv < g.nv ? mv : mv - g.nv) +
static_cast<size_t>(g.nv) * (mw < g.nw ? mw : mw - g.nw))];
float dr2 = 0, db = 0;
for (int q = 0; q < 5; q++) {
const float b = at.b[q][0];
const float e = at.a[q] * expapprox(fmaxf(b * r2, -88.f));
dr2 += e * b;
db += e * b * (1.5f + b * r2);
if (!aniso) {
float dr2 = 0, db = 0;
for (int q = 0; q < 5; q++) {
const float b = at.b[q][0];
const float e = at.a[q] * expapprox(fmaxf(b * r2, -88.f));
dr2 += e * b;
db += e * b * (1.5f + b * r2);
}
gx += value * dr2 * x;
gy += value * dr2 * y;
gz += value * dr2 * z;
gb += value * db;
} else {
for (int q = 0; q < 5; q++) {
const float *m = at.b[q];
const float rur = x * x * m[0] + y * y * m[1] + z * z * m[2] + 2 * (x * y * m[3] + x * z * m[4] + y * z * m[5]);
const float e = value * at.a[q] * expapprox(fmaxf(rur, -88.f));
const float mr0 = m[0] * x + m[3] * y + m[4] * z, mr1 = m[3] * x + m[1] * y + m[5] * z,
mr2 = m[4] * x + m[5] * y + m[2] * z;
gx += e * mr0;
gy += e * mr1;
gz += e * mr2;
gu[0] += e * (m[0] + 2 * mr0 * mr0);
gu[1] += e * (m[1] + 2 * mr1 * mr1);
gu[2] += e * (m[2] + 2 * mr2 * mr2);
gu[3] += e * (m[3] + 2 * mr0 * mr1);
gu[4] += e * (m[4] + 2 * mr0 * mr2);
gu[5] += e * (m[5] + 2 * mr1 * mr2);
}
}
gx += value * dr2 * x;
gy += value * dr2 * y;
gz += value * dr2 * z;
gb += value * db;
}
for (int off = WARP / 2; off > 0; off /= 2) {
gx += __shfl_down_sync(0xffffffff, gx, off);
gy += __shfl_down_sync(0xffffffff, gy, off);
gz += __shfl_down_sync(0xffffffff, gz, off);
gb += __shfl_down_sync(0xffffffff, gb, off);
for (int c = 0; c < 6; c++)
gu[c] += __shfl_down_sync(0xffffffff, gu[c], off);
}
if (lane == 0) {
const float s = 2.0f * at.occ * map_scale;
grad[a] = make_float4(s * gx, s * gy, s * gz, at.occ * map_scale * gb / static_cast<float>(FOUR_PI2));
const float o = at.occ * map_scale;
if (!aniso)
grad[a] = make_float4(s * gx, s * gy, s * gz, o * gb / static_cast<float>(FOUR_PI2));
else
grad[a] = make_float4(s * gx, s * gy, s * gz, o * (gu[0] + gu[1] + gu[2]) / static_cast<float>(2 * FOUR_PI2));
if (grad_u != nullptr) {
const float f[6] = {o * gu[0], o * gu[1], o * gu[2], 2 * o * gu[3], 2 * o * gu[4], 2 * o * gu[5]};
for (int c = 0; c < 6; c++)
grad_u[6 * a + c] = f[c];
}
}
}
@@ -483,6 +518,7 @@ struct RefinementGPUEngineImpl {
// sums, their gradients, and the restraints' positions going up.
CudaHostPtr<double> host_sums;
CudaHostPtr<float4> host_grad, host_xyzb, host_restraint_grad;
CudaHostPtr<float> host_grad_u; // six per atom
CudaDevicePtr<ModelMaskAtom> mask_atoms;
CudaDevicePtr<float> grid;
CudaDevicePtr<float2> spectrum;
@@ -503,6 +539,8 @@ struct RefinementGPUEngineImpl {
CudaDevicePtr<double2> phase; // exp(2 pi i h.t) over the integer h.t from phase_min
int phase_min = 0;
CudaDevicePtr<float4> grad;
CudaDevicePtr<float> grad_u;
int n_aniso = 0; // the anisotropic atoms of the last SetDensities()
std::unique_ptr<ModelDensityGPU> density;
std::unique_ptr<ModelMaskGPU> mask;
CudaFFTPlan forward, inverse, mask_forward;
@@ -516,10 +554,12 @@ struct RefinementGPUEngineImpl {
pos = CudaDevicePtr<float4>(na, sync);
host_sums = CudaHostPtr<double>(2);
host_grad = CudaHostPtr<float4>(na);
host_grad_u = CudaHostPtr<float>(6 * na);
host_xyzb = CudaHostPtr<float4>(na);
host_restraint_grad = CudaHostPtr<float4>(na);
mask_atoms = CudaDevicePtr<ModelMaskAtom>(na, sync);
grad = CudaDevicePtr<float4>(na, sync);
grad_u = CudaDevicePtr<float>(6 * na, sync);
grid = CudaDevicePtr<float>(GridPoints(setup.grid), sync);
spectrum = CudaDevicePtr<float2>(ComplexPoints(setup.grid), sync);
fft_work = CudaDevicePtr<char>(std::max<size_t>(work_bytes, 1), sync);
@@ -662,10 +702,14 @@ struct RefinementGPUEngineImpl {
const size_t threads = static_cast<size_t>(n_atoms) * WARP;
atom_gradient_kernel<<<blocks(threads), THREADS, 0, stream>>>(density->DeviceAtoms(), pos,
density->DeviceBoxes(), n_atoms, geom,
this->grid, Norm(), this->grad);
this->grid, Norm(), this->grad,
n_aniso > 0 ? this->grad_u.get() : nullptr);
cuda_err(cudaGetLastError());
cuda_err(cudaMemcpyAsync(host_grad.get(), this->grad, n_atoms * sizeof(float4), cudaMemcpyDeviceToHost,
stream));
if (n_aniso > 0)
cuda_err(cudaMemcpyAsync(host_grad_u.get(), this->grad_u, 6 * n_atoms * sizeof(float),
cudaMemcpyDeviceToHost, stream));
}
}
}
@@ -700,7 +744,7 @@ struct RefinementGPUEngineImpl {
// the host, so that a replay finds them so).
void LaunchEvaluate(bool grad) {
host_sums[0] = 0;
evaluate.Launch(stream, 2L * n_atoms + (grad ? 1 : 0), [&]() { QueueEvaluate(grad); });
evaluate.Launch(stream, 4L * n_atoms + (n_aniso > 0 ? 2 : 0) + (grad ? 1 : 0), [&]() { QueueEvaluate(grad); });
}
void LaunchRestraints(int which, bool grad) {
@@ -732,7 +776,8 @@ size_t RefinementGPUEngine::FFTWorkBytes(int device, const RefinementGPUSetup &s
size_t RefinementGPUEngine::DeviceBytes(const RefinementGPUSetup &s, size_t fft_work_bytes, size_t max_records,
size_t max_slots) {
const size_t na = s.max_atoms + 1, nr = s.rows.size() + 1;
size_t bytes = na * (2 * sizeof(float4) + 3 * sizeof(int) + sizeof(ModelMaskAtom) + sizeof(ModelDensityAtom));
size_t bytes = na * (2 * sizeof(float4) + 6 * sizeof(float) + 3 * sizeof(int) + sizeof(ModelMaskAtom) +
sizeof(ModelDensityAtom));
bytes += GridPoints(s.grid) * sizeof(float) + ComplexPoints(s.grid) * sizeof(float2) + fft_work_bytes;
bytes += nr * (3 * sizeof(int) + 2 * sizeof(double) + 4 * sizeof(float) + 4 * sizeof(float2) + 2);
bytes += (s.scatter_point.size() + s.scatter_start.size() + 2 * s.scatter_term.size()) * sizeof(int);
@@ -791,6 +836,9 @@ void RefinementGPUEngine::SetDensities(const ModelDensityAtom *atoms, size_t n)
throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "refinement: atoms over the reserve");
e.density->SetAtoms(e.setup.grid, atoms, n);
e.n_atoms = static_cast<int>(n);
e.n_aniso = 0;
for (size_t i = 0; i < n; i++)
e.n_aniso += atoms[i].aniso != 0;
}
void RefinementGPUEngine::SetPositions(const std::array<float, 4> *pos, size_t n) {
@@ -850,6 +898,10 @@ const std::array<float, 4> *RefinementGPUEngine::HostGradient() const {
return reinterpret_cast<const std::array<float, 4> *>(impl_->host_grad.get());
}
const float *RefinementGPUEngine::HostUGradient() const {
return impl_->host_grad_u.get();
}
const std::array<float, 4> *RefinementGPUEngine::HostRestraintGradient() const {
return reinterpret_cast<const std::array<float, 4> *>(impl_->host_restraint_grad.get());
}
@@ -72,6 +72,9 @@ public:
// The gradients of the last evaluation, in the engine's pinned buffers: valid until the next one.
const std::array<float, 4> *HostGradient() const;
const std::array<float, 4> *HostRestraintGradient() const;
// dT/d(U11 U22 U33 U12 U13 U23) of the anisotropic atoms (six per atom, zero for the isotropic ones) of the last
// evaluation with a gradient where some atom was anisotropic.
const float *HostUGradient() const;
// Restraint set `which` (0 or 1), as RestraintSet holds it.
void SetRestraints(int which, const std::vector<RestraintRecord> &records, const std::vector<int> &plane_atoms,
@@ -116,7 +116,7 @@ public:
void Fcalc(const gemmi::Model &model, std::vector<std::complex<float>> &f) override {
std::vector<std::complex<double>> fd;
Compose(model, fd);
Compose(WithAniso(model), fd);
f.assign(fd.begin(), fd.end());
}
@@ -129,7 +129,8 @@ public:
f.push_back(fm.get_value_by_hkl(h));
}
double Evaluate(const gemmi::Model &model, std::vector<std::array<double, 4>> *grad) override {
double Evaluate(const gemmi::Model &model_in, std::vector<std::array<double, 4>> *grad) override {
const gemmi::Model &model = WithAniso(model_in);
std::vector<std::complex<double>> f;
Compose(model, f);
const size_t n = f.size();
@@ -193,11 +194,23 @@ public:
map.data[i] = out[i] * norm;
fftwf_free(in);
fftwf_free(out);
AtomGradientsFromMap(map, model, blur_, nthreads_, *grad);
const bool aniso = AnyAniso();
AtomGradientsFromMap(map, model, blur_, nthreads_, *grad, aniso ? &grad_u_ : nullptr);
if (!aniso)
grad_u_.clear();
return sum;
}
private:
// The model as the engine grids it: where some atom is anisotropic, a copy with every atom's U set (ApplyAniso).
const gemmi::Model &WithAniso(const gemmi::Model &model) {
if (!AnyAniso())
return model;
aniso_model_ = model;
ApplyAniso(aniso_model_, *aniso_);
return aniso_model_;
}
void Compose(const gemmi::Model &model, std::vector<std::complex<double>> &f) const {
gemmi::DensityCalculator<Table, float> dc = Calculator(cell_, sg_, d_min_, blur_);
PutModelDensityOnGrid(dc, model, {}, nthreads_); // one copy: the symmetry is composed below
@@ -214,6 +227,7 @@ private:
std::vector<size_t> orbit_leaders_;
double blur_ = 0;
XrayRowModel rows_;
gemmi::Model aniso_model_{1};
};
} // namespace
@@ -412,6 +426,30 @@ RFactors ComputeRFactors(const RefinementReflections &refl, const std::vector<st
return r;
}
gemmi::SMat33<float> AnisoU(const AtomAniso &a, double b) {
const double u = b / gemmi::u_to_b();
gemmi::SMat33<double> m{a.dev[0] + u, a.dev[1] + u, a.dev[2] + u, a.dev[3], a.dev[4], a.dev[5]};
const std::array<double, 3> ev = m.calculate_eigenvalues();
const double lowest = std::min({ev[0], ev[1], ev[2]});
if (!(lowest >= ANISO_U_MIN))
m = m.added_kI(ANISO_U_MIN - (std::isfinite(lowest) ? lowest : 0.0));
return {static_cast<float>(m.u11), static_cast<float>(m.u22), static_cast<float>(m.u33),
static_cast<float>(m.u12), static_cast<float>(m.u13), static_cast<float>(m.u23)};
}
void ApplyAniso(gemmi::Model &model, const std::vector<AtomAniso> &aniso) {
size_t i = 0;
for (gemmi::Chain &ch : model.chains)
for (gemmi::Residue &r : ch.residues)
for (gemmi::Atom &a : r.atoms) {
if (i < aniso.size() && aniso[i].on)
a.aniso = AnisoU(aniso[i], a.b_iso);
else
a.aniso = {0, 0, 0, 0, 0, 0};
i++;
}
}
std::vector<std::array<double, 4>> ModelXyzB(const gemmi::Model &model) {
std::vector<std::array<double, 4>> out;
for (const gemmi::Chain &ch : model.chains)
@@ -458,8 +496,82 @@ std::unique_ptr<RefinementEngine> MakeRefinementEngineCPU(const gemmi::UnitCell
return std::make_unique<RefinementEngineCPU>(cell, sg, d_min, refl, nthreads);
}
namespace {
// AtomGradientsFromMap() for an anisotropic atom, gridded as PutModelDensityOnGrid() grids it: rho = occ sum_q a_q
// exp(r^T M_q r), M_q = -4 pi^2 (B_q I + 8 pi^2 U + blur I)^-1 and r = centre - point. Then d rho / d(centre) =
// occ sum_q 2 e_q M_q r, and, as d a_q / dU = -4 pi^2 a_q (B..)^-1 = a_q M_q and d(r^T M_q r) / dU = 2 (M_q r)(M_q r)^T,
// d rho / dU = occ sum_q e_q (M_q + 2 (M_q r)(M_q r)^T) - an off-diagonal U_ij twice its matrix element.
void AnisoGradient(const gemmi::Grid<float> &map, const gemmi::DensityCalculator<Table, float> &dc,
const gemmi::Atom &atom, double blur, std::array<double, 4> &grad, std::array<double, 6> *grad_u) {
using CReal = Table::Coef::coef_type;
const auto &coef = Table::get(atom.element, atom.charge, atom.serial);
const auto aniso_b = atom.aniso.scaled(CReal(gemmi::u_to_b())).added_kI(CReal(blur));
const CReal b_max = AnisoRadiusB(aniso_b);
const double radius = dc.estimate_radius(coef.precalculate_density_iso(b_max, 0.0f), b_max);
const auto precal = coef.precalculate_density_aniso_b(aniso_b, 0.0f);
const float max_r2 = static_cast<float>(radius * radius);
const int nu = map.nu, nv = map.nv, nw = map.nw;
const gemmi::UpperTriangularMat33 &om = map.orth_n;
const float m00 = static_cast<float>(om.a11), m01 = static_cast<float>(om.a12), m02 = static_cast<float>(om.a13),
m11 = static_cast<float>(om.a22), m12 = static_cast<float>(om.a23), m22 = static_cast<float>(om.a33);
const gemmi::Fractional fpos = map.unit_cell.fractionalize(atom.pos);
const double n3[3] = {fpos.x * nu, fpos.y * nv, fpos.z * nw};
float nctr[3];
int c[3], d[3];
for (int k = 0; k < 3; k++) {
nctr[k] = static_cast<float>(n3[k]);
c[k] = gemmi::iround(n3[k]);
}
d[0] = std::min(static_cast<int>(std::ceil(radius / map.spacing[0])), nu - 1);
d[1] = std::min(static_cast<int>(std::ceil(radius / map.spacing[1])), nv - 1);
d[2] = std::min(static_cast<int>(std::ceil(radius / map.spacing[2])), nw - 1);
double g[3] = {0, 0, 0}, gu[6] = {0, 0, 0, 0, 0, 0};
for (int w = c[2] - d[2]; w <= c[2] + d[2]; w++) {
const float fz = nctr[2] - static_cast<float>(w);
const float z = m22 * fz;
const int w_ = gemmi::modulo(w, nw);
for (int v = c[1] - d[1]; v <= c[1] + d[1]; v++) {
const float fy = nctr[1] - static_cast<float>(v);
const float y = m11 * fy + m12 * fz;
const float yz = y * y + z * z;
if (yz > max_r2)
continue;
const float x_yz = m01 * fy + m02 * fz;
const size_t row = static_cast<size_t>(nu) * (gemmi::modulo(v, nv) + static_cast<size_t>(nv) * w_);
for (int u = c[0] - d[0]; u <= c[0] + d[0]; u++) {
const float x = m00 * (nctr[0] - static_cast<float>(u)) + x_yz;
if (yz + x * x > max_r2)
continue;
const double value = map.data[row + gemmi::modulo(u, nu)];
for (size_t q = 0; q < std::size(precal.a); q++) {
const gemmi::SMat33<float> &m = precal.b[q];
const gemmi::Vec3f r(x, y, z);
const double e = value * precal.a[q] * gemmi::unsafe_expapprox(std::max(m.r_u_r(r), -88.f));
const double mr[3] = {m.u11 * x + m.u12 * y + m.u13 * z, m.u12 * x + m.u22 * y + m.u23 * z,
m.u13 * x + m.u23 * y + m.u33 * z};
for (int k = 0; k < 3; k++)
g[k] += e * mr[k];
gu[0] += e * (m.u11 + 2 * mr[0] * mr[0]);
gu[1] += e * (m.u22 + 2 * mr[1] * mr[1]);
gu[2] += e * (m.u33 + 2 * mr[2] * mr[2]);
gu[3] += e * (m.u12 + 2 * mr[0] * mr[1]);
gu[4] += e * (m.u13 + 2 * mr[0] * mr[2]);
gu[5] += e * (m.u23 + 2 * mr[1] * mr[2]);
}
}
}
}
const double occ = atom.occ;
grad = {2 * occ * g[0], 2 * occ * g[1], 2 * occ * g[2], occ * (gu[0] + gu[1] + gu[2]) / gemmi::u_to_b()};
if (grad_u != nullptr)
*grad_u = {occ * gu[0], occ * gu[1], occ * gu[2], 2 * occ * gu[3], 2 * occ * gu[4], 2 * occ * gu[5]};
}
} // namespace
void AtomGradientsFromMap(const gemmi::Grid<float> &map, const gemmi::Model &model, double blur, size_t nthreads,
std::vector<std::array<double, 4>> &grad) {
std::vector<std::array<double, 4>> &grad, std::vector<std::array<double, 6>> *grad_u) {
const gemmi::DensityCalculator<Table, float> dc = Calculator(map.unit_cell, *map.spacegroup, 1.0, blur);
std::vector<const gemmi::Atom *> atoms;
for (const gemmi::Chain &ch : model.chains)
@@ -467,6 +579,8 @@ void AtomGradientsFromMap(const gemmi::Grid<float> &map, const gemmi::Model &mod
for (const gemmi::Atom &a : r.atoms)
atoms.push_back(&a);
grad.assign(atoms.size(), {0, 0, 0, 0});
if (grad_u != nullptr)
grad_u->assign(atoms.size(), {0, 0, 0, 0, 0, 0});
const int nu = map.nu, nv = map.nv, nw = map.nw;
// As ModelGrid.cpp walks a density box: the orthogonalization over the grid size in single precision.
const gemmi::UpperTriangularMat33 &om = map.orth_n;
@@ -476,6 +590,10 @@ void AtomGradientsFromMap(const gemmi::Grid<float> &map, const gemmi::Model &mod
ParallelChunks(static_cast<int>(atoms.size()), nthreads, [&](int lo, int hi) {
for (int i = lo; i < hi; i++) {
const gemmi::Atom &atom = *atoms[i];
if (atom.aniso.nonzero()) {
AnisoGradient(map, dc, atom, blur, grad[i], grad_u != nullptr ? &(*grad_u)[i] : nullptr);
continue;
}
const float b = static_cast<float>(atom.b_iso + blur);
const auto precal = Table::get(atom.element, atom.charge, atom.serial).precalculate_density_iso(b, 0.0f);
const double radius = dc.estimate_radius(precal, b);
@@ -101,6 +101,20 @@ RFactors ComputeRFactors(const RefinementReflections &refl, const std::vector<st
// x, y, z and B of every atom of `model`, in model order.
std::vector<std::array<double, 4>> ModelXyzB(const gemmi::Model &model);
// An atom's displacement where it is anisotropic: U = dev + B / (8 pi^2) I, B the atom's b_iso (its equivalent
// isotropic B, as dev is traceless) and dev the anisotropy, U11 U22 U33 U12 U13 U23 (Cartesian, A^2). So everything
// that refines, restrains and bounds B works on an anisotropic atom as on an isotropic one, its shape held.
struct AtomAniso {
bool on = false;
std::array<float, 6> dev{};
};
// U of an anisotropic atom of b_iso `b`, its smallest eigenvalue at least ANISO_U_MIN (a U that is not positive
// definite has no density).
constexpr double ANISO_U_MIN = 0.5 / (8 * 3.14159265358979323846 * 3.14159265358979323846); // B 0.5 A^2
gemmi::SMat33<float> AnisoU(const AtomAniso &a, double b);
// The model with every anisotropic atom of `aniso` given its U (AnisoU), the others as they are.
void ApplyAniso(gemmi::Model &model, const std::vector<AtomAniso> &aniso);
// What evaluates the targets for a set of atom positions and B: the CPU below, or a GPU (RefinementGPU.h). The
// atoms are the model's, all of them, in model order; the grid and the reflections are fixed when the engine is
// made, the row model each macro-cycle, the blur where SetBlur is called.
@@ -117,9 +131,17 @@ public:
virtual void Fmask(const gemmi::Model &model, std::vector<std::complex<float>> &f) = 0;
// The sum of -ln L over the working reflections and, with `grad`, its derivatives with respect to every
// atom's x, y, z (Cartesian, A) and B (A^2), in model order.
// For an anisotropic atom (SetAniso) the B derivative is the one along dev + B / (8 pi^2) I, its shape held.
virtual double Evaluate(const gemmi::Model &model, std::vector<std::array<double, 4>> *grad) = 0;
virtual bool OnGPU() const = 0;
// Which atoms are anisotropic and their shapes, in model order (AtomAniso), read at every Fcalc() and
// Evaluate() until it is set again; null or empty: every atom isotropic (the model's own aniso is not read).
void SetAniso(const std::vector<AtomAniso> *aniso) { aniso_ = aniso; }
// dT/d(U11 U22 U33 U12 U13 U23) of every anisotropic atom at the last Evaluate() with a gradient, zero for the
// isotropic ones; empty where none is anisotropic.
const std::vector<std::array<double, 6>> &UGradient() const { return grad_u_; }
// Restraint set `which` (0: the geometry and the repulsions, 1: the B similarity) and its sum of E at xyzb
// (x, y, z, B per atom) with, given `grad`, dE/d(x, y, z, B) per atom. On the CPU unless the engine has its own.
virtual void SetRestraints(int which, const RestraintSet &set) { restraints_[which] = set; }
@@ -138,6 +160,17 @@ public:
return xray;
}
protected:
bool AnyAniso() const {
if (aniso_ != nullptr)
for (const AtomAniso &a : *aniso_)
if (a.on)
return true;
return false;
}
const std::vector<AtomAniso> *aniso_ = nullptr;
std::vector<std::array<double, 6>> grad_u_;
private:
RestraintSet restraints_[2];
};
@@ -179,6 +212,9 @@ std::unique_ptr<RefinementEngine> MakeRefinementEngineCPU(const gemmi::UnitCell
// The derivative map of the X-ray target on a grid (the FFT method's map), summed against each atom's density
// derivatives: dT/d(x, y, z, B) per atom. `map` is u fastest on the grid of RigidBodyZoneGrid(), the atoms
// gridded with `blur`. Exposed for the tests.
// gridded with `blur`. Exposed for the tests. An atom with a nonzero aniso is gridded anisotropically, as
// PutModelDensityOnGrid() grids it: its B derivative is that along U + dB / (8 pi^2) I, and with `grad_u` (resized
// to the atoms, zero for the isotropic ones) its derivatives in U11 U22 U33 U12 U13 U23.
void AtomGradientsFromMap(const gemmi::Grid<float> &map, const gemmi::Model &model, double blur, size_t nthreads,
std::vector<std::array<double, 4>> &grad);
std::vector<std::array<double, 4>> &grad,
std::vector<std::array<double, 6>> *grad_u = nullptr);
+4
View File
@@ -530,6 +530,7 @@ namespace {
Add(s, KeyText("REFINE_DEVICE", rf.on_gpu ? "GPU" : "CPU", true));
Add(s, KeyReal("REFINE_SECONDS", rf.seconds, "{:.1f}", true));
Add(s, KeyReal("REFINE_D_MIN", rf.d_min, "{:.2f}"));
Add(s, KeyText("REFINE_ADP_INPUT", rf.adp_input));
Add(s, KeyReal("REFINE_R_WORK_START", rf.r_work_start, "{:.4f}"));
Add(s, KeyReal("REFINE_R_FREE_START", rf.r_free_start, "{:.4f}"));
Add(s, KeyInt("REFINE_RIGID_BODIES", rf.rigid_bodies));
@@ -546,6 +547,9 @@ namespace {
Add(s, KeyReal("REFINE_RMSD_ANGLES_START", rf.geometry_start.rmsd_angles, "{:.2f}", true));
Add(s, KeyInt("REFINE_CLASHES", rf.geometry.clashes));
Add(s, KeyReal("REFINE_XRAY_WEIGHT", rf.weight, "{:.4g}", true));
Add(s, KeyText("REFINE_ADP_MODEL", rf.adp_model));
Add(s, KeyInt("REFINE_ANISO_ATOMS", rf.aniso_atoms));
Add(s, KeyInt("REFINE_TLS_GROUPS", rf.tls_groups));
Add(s, KeyInt("REFINE_ATOMS", rf.atoms));
Add(s, KeyInt("REFINE_ATOMS_REFINED", rf.refined_atoms));
Add(s, KeyInt("REFINE_RESIDUES_FIXED", rf.fixed_residues));
+253
View File
@@ -17,6 +17,7 @@
#include "../common/Logger.h"
#include "../image_analysis/structure_refinement/MonomerLibrary.h"
#include "../image_analysis/structure_refinement/Refinement.h"
#include "../image_analysis/structure_refinement/RefinementADP.h"
#include "../image_analysis/structure_refinement/RefinementMath.h"
#include "../image_analysis/structure_refinement/RefinementMinimiser.h"
#include "../image_analysis/structure_refinement/RefinementRestraints.h"
@@ -836,3 +837,255 @@ TEST_CASE("Refinement_restrained_recovers_model", "[Refinement]") {
CHECK_FALSE(r.cancelled);
}
}
namespace {
// Every other atom anisotropic, its dev random but traceless, up to a third of its B.
std::vector<AtomAniso> random_aniso(const gemmi::Model &model, unsigned seed) {
std::mt19937 rng(seed);
std::uniform_real_distribution<double> d(-1.0, 1.0);
std::vector<AtomAniso> out;
int i = 0;
for (const gemmi::Chain &ch : model.chains)
for (const gemmi::Residue &r : ch.residues)
for (const gemmi::Atom &a : r.atoms) {
AtomAniso an;
an.on = (i++ % 2) == 0;
if (an.on) {
const double u = a.b_iso / gemmi::u_to_b() / 3.0;
double v[6];
for (double &x : v)
x = u * d(rng);
const double t = (v[0] + v[1] + v[2]) / 3.0;
an.dev = {static_cast<float>(v[0] - t), static_cast<float>(v[1] - t), static_cast<float>(v[2] - t),
static_cast<float>(v[3] / 2), static_cast<float>(v[4] / 2), static_cast<float>(v[5] / 2)};
}
out.push_back(an);
}
return out;
}
// The engine's U gradient of a few anisotropic atoms against central differences of its own target.
void check_u_gradient(RefinementEngine &engine, const gemmi::Model &model, std::vector<AtomAniso> &aniso,
double tolerance) {
std::vector<std::array<double, 4>> grad;
engine.SetAniso(&aniso);
engine.Evaluate(model, &grad);
const std::vector<std::array<double, 6>> gu = engine.UGradient();
REQUIRE(gu.size() == aniso.size());
double max_g = 0;
for (size_t i = 0; i < gu.size(); i++)
for (int c = 0; c < 6; c++)
max_g = std::max(max_g, std::fabs(gu[i][c]));
const double h = 2e-3; // A^2
for (size_t i = 0; i < aniso.size(); i += 2 * (aniso.size() / 14)) {
REQUIRE(aniso[i].on);
for (int c = 0; c < 6; c++) {
const float v0 = aniso[i].dev[c];
aniso[i].dev[c] = static_cast<float>(v0 + h);
const double fp = engine.Evaluate(model, nullptr);
aniso[i].dev[c] = static_cast<float>(v0 - h);
const double fm = engine.Evaluate(model, nullptr);
aniso[i].dev[c] = v0;
const double fd = (fp - fm) / (2 * h);
CAPTURE(i, c, fd, gu[i][c], max_g);
CHECK(std::fabs(fd - gu[i][c]) < tolerance * std::max(std::fabs(fd), 0.05 * max_g));
}
}
}
} // namespace
TEST_CASE("RefinementXray_aniso_gradient_CPU", "[Refinement]") {
for (const char *group : {"P 1", "P 21 21 21"}) {
CAPTURE(group);
const gemmi::SpaceGroup &sg = *gemmi::find_spacegroup_by_name(group);
const gemmi::UnitCell cell(28, 31, 34, 90, 90, 90);
gemmi::Structure truth = random_crystal(cell, group, 240, 7);
gemmi::Model model = shaken(truth.models[0], 0.3, 4.0, 8);
RefinementReflections refl = PrepareRefinementReflections(asu_reflections(cell, sg, 2.0), cell, sg, 2.0);
auto cpu = MakeRefinementEngineCPU(cell, sg, 2.0, refl, 4);
cpu->SetBlur(RefinementBlur(cell, sg, 2.0, 5.0));
set_observations(*cpu, refl, truth.models[0], model);
std::vector<AtomAniso> aniso = random_aniso(model, 9);
check_u_gradient(*cpu, model, aniso, 0.06);
// the positions and the B along the isotropic direction, the shapes held
check_xray_gradient(*cpu, model, 0.03);
}
}
#ifdef JFJOCH_USE_CUDA
TEST_CASE("RefinementXray_aniso_gradient_GPU", "[Refinement]") {
if (get_gpu_count() == 0)
return;
const char *group = "P 21 21 21";
const gemmi::SpaceGroup &sg = *gemmi::find_spacegroup_by_name(group);
const gemmi::UnitCell cell(28, 31, 34, 90, 90, 90);
gemmi::Structure truth = random_crystal(cell, group, 240, 7);
gemmi::Model model = shaken(truth.models[0], 0.3, 4.0, 8);
RefinementReflections refl = PrepareRefinementReflections(asu_reflections(cell, sg, 2.0), cell, sg, 2.0);
auto cpu = MakeRefinementEngineCPU(cell, sg, 2.0, refl, 4);
cpu->SetBlur(RefinementBlur(cell, sg, 2.0, 5.0));
set_observations(*cpu, refl, truth.models[0], model);
auto gpu = MakeRefinementEngineGPU(0, cell, sg, 2.0, refl, model, 300.0, SIZE_MAX);
REQUIRE(gpu);
gpu->SetBlur(RefinementBlur(cell, sg, 2.0, 5.0));
std::vector<std::complex<float>> f;
std::vector<AtomAniso> aniso = random_aniso(model, 9);
cpu->SetAniso(&aniso);
gpu->SetAniso(&aniso);
cpu->Fcalc(model, f);
XrayRowModel rows;
rows.k_total.assign(refl.Size(), 1.0f);
rows.fmask.assign(refl.Size(), {0, 0});
EstimateLikelihoodParameters(refl, f, rows);
cpu->SetRowModel(rows);
gpu->SetRowModel(rows);
std::vector<std::complex<float>> fg;
gpu->Fcalc(model, fg);
double num = 0, den = 0;
for (size_t m = 0; m < f.size(); m++) {
num += std::abs(f[m] - fg[m]);
den += std::abs(f[m]);
}
CHECK(num / den < 1e-3);
std::vector<std::array<double, 4>> gc, gg;
cpu->Evaluate(model, &gc);
const std::vector<std::array<double, 6>> uc = cpu->UGradient();
gpu->Evaluate(model, &gg);
const std::vector<std::array<double, 6>> ug = gpu->UGradient();
REQUIRE(ug.size() == uc.size());
double max_g = 0, max_u = 0;
for (size_t i = 0; i < gc.size(); i++) {
for (int c = 0; c < 4; c++)
max_g = std::max(max_g, std::fabs(gc[i][c]));
for (int c = 0; c < 6; c++)
max_u = std::max(max_u, std::fabs(uc[i][c]));
}
for (size_t i = 0; i < gc.size(); i++) {
CAPTURE(i);
for (int c = 0; c < 4; c++)
CHECK(std::fabs(gg[i][c] - gc[i][c]) < 0.01 * max_g);
for (int c = 0; c < 6; c++)
CHECK(std::fabs(ug[i][c] - uc[i][c]) < 0.01 * max_u);
}
check_u_gradient(*gpu, model, aniso, 0.08);
}
#endif
TEST_CASE("RefinementADP_tls_gradient", "[Refinement]") {
std::mt19937 rng(4);
std::uniform_real_distribution<double> d(-1.0, 1.0);
gemmi::Model model(1);
gemmi::Chain chain("A");
gemmi::Residue res;
for (int i = 0; i < 30; i++) {
gemmi::Atom a;
a.pos = gemmi::Position(10 + 8 * d(rng), -3 + 8 * d(rng), 5 + 8 * d(rng));
res.atoms.push_back(a);
}
chain.residues.push_back(res);
model.chains.push_back(chain);
std::vector<gemmi::Atom *> atoms = atoms_of(model);
TlsBody body;
body.origin = gemmi::Position(9, -2, 4);
std::array<double, TlsBody::PARAMETERS> p{};
for (double &x : p)
x = 0.01 * d(rng);
body.Set(p);
for (int i = 0; i < 30; i++)
body.atoms.push_back(i);
// U is linear in the parameters: U(p) is sum p_k dU/dp_k
for (gemmi::Atom *a : atoms) {
const gemmi::SMat33<double> u = body.U(a->pos);
const auto dk = body.Derivatives(a->pos);
const double uu[6] = {u.u11, u.u22, u.u33, u.u12, u.u13, u.u23};
for (int c = 0; c < 6; c++) {
double sum = 0;
for (int k = 0; k < TlsBody::PARAMETERS; k++)
sum += p[k] * dk[k][c];
CHECK(sum == Catch::Approx(uu[c]).margin(1e-12));
}
}
// the chain rule in matrix form against the sum over dU/dp
std::vector<std::array<double, 6>> gu(atoms.size());
for (auto &g : gu)
for (double &x : g)
x = d(rng);
const std::array<double, TlsBody::PARAMETERS> g = body.Gradient(atoms, gu);
for (int k = 0; k < TlsBody::PARAMETERS; k++) {
double sum = 0;
for (size_t i = 0; i < atoms.size(); i++) {
const auto dk = body.Derivatives(atoms[i]->pos);
for (int c = 0; c < 6; c++)
sum += gu[i][c] * dk[k][c];
}
CAPTURE(k);
CHECK(g[k] == Catch::Approx(sum).margin(1e-9));
}
// and through the metadata in degrees and back
gemmi::Model m2 = model;
const std::vector<gemmi::TlsGroup> groups = TlsBodiesToGroups(m2, {body});
REQUIRE(groups.size() == 1);
const std::vector<TlsBody> back = TlsBodiesFromModel(m2, groups);
REQUIRE(back.size() == 1);
CHECK(back[0].atoms.size() == 30);
const gemmi::SMat33<double> u0 = body.U(atoms[3]->pos), u1 = back[0].U(atoms[3]->pos);
CHECK(u1.u12 == Catch::Approx(u0.u12).margin(1e-12));
CHECK(u1.u33 == Catch::Approx(u0.u33).margin(1e-12));
}
TEST_CASE("RefinementADP_u_restraints_gradient", "[Refinement]") {
std::mt19937 rng(5);
std::uniform_real_distribution<double> d(-1.0, 1.0);
const int n = 6;
std::vector<gemmi::Position> pos;
std::vector<std::array<double, 6>> u(n);
for (int i = 0; i < n; i++) {
pos.emplace_back(1.5 * i + 0.2 * d(rng), 0.3 * d(rng), 0.3 * d(rng));
u[i] = {0.2 + 0.05 * d(rng), 0.2 + 0.05 * d(rng), 0.2 + 0.05 * d(rng), 0.03 * d(rng), 0.03 * d(rng), 0.03 * d(rng)};
}
const URestraints r = MakeURestraints({{0, 1}, {1, 2}, {0, 2}, {3, 4}}, {1, 1, 1, 1, 1, 1});
CHECK(r.pairs.size() == 4);
REQUIRE(r.isotropic == std::vector<int>{5});
std::vector<std::array<double, 6>> g(n, {0, 0, 0, 0, 0, 0});
EvaluateURestraints(r, pos, u, &g);
const double h = 1e-6;
for (int i = 0; i < n; i++)
for (int c = 0; c < 6; c++) {
std::vector<std::array<double, 6>> up = u, um = u;
up[i][c] += h;
um[i][c] -= h;
const double fd = (EvaluateURestraints(r, pos, up, nullptr) - EvaluateURestraints(r, pos, um, nullptr)) / (2 * h);
CAPTURE(i, c);
CHECK(g[i][c] == Catch::Approx(fd).epsilon(1e-5).margin(1e-3));
}
// a rigid pair: the same U along the bond, however different across it, costs no rigid-bond term
const URestraints one = MakeURestraints({{0, 1}}, {1, 1});
const std::vector<gemmi::Position> line = {{0, 0, 0}, {0, 0, 1.5}};
const std::vector<std::array<double, 6>> across = {{0.1, 0.3, 0.2, 0, 0, 0}, {0.3, 0.1, 0.2, 0.05, 0, 0}};
const double sim = (0.2 * 0.2 + 0.2 * 0.2 + 2 * 0.05 * 0.05) / (URestraints::SIGMA_SIMILAR * URestraints::SIGMA_SIMILAR);
CHECK(EvaluateURestraints(one, line, across, nullptr) == Catch::Approx(sim));
}
TEST_CASE("RefinementADP_tls_segments", "[Refinement]") {
// a chain of 60 residues along x whose B follows one quadratic on the first 25 residues and another after
std::vector<std::vector<int>> residues;
std::vector<gemmi::Position> pos;
std::vector<double> b;
for (int r = 0; r < 60; r++) {
std::vector<int> ra;
for (int a = 0; a < 5; a++) {
const gemmi::Position p(3.8 * r, 2.0 * std::sin(a + r), 2.0 * std::cos(a * r));
ra.push_back(static_cast<int>(pos.size()));
pos.push_back(p);
const double y = p.y, z = p.z;
b.push_back(r < 25 ? 20 + 0.5 * (y * y + z * z) : 60 + 3.0 * (y * y) - 0.02 * p.x);
}
residues.push_back(ra);
}
CHECK(TlsSegments(residues, pos, b, 1, 10) == std::vector<int>{0});
CHECK(TlsSegments(residues, pos, b, 2, 10) == std::vector<int>{0, 25});
// too short for three segments of 25
CHECK(TlsSegments(residues, pos, b, 3, 25).size() == 2);
}
+58 -1
View File
@@ -53,6 +53,61 @@ gemmi::Structure LabelledModel(const gemmi::Structure &placed, const UnitCell &c
return st;
}
// The TLS groups of the model as REMARK 3 TLS records (the layout REFMAC writes and gemmi reads back), in place of
// any REMARK 3 the model was read with: gemmi's PDB writer writes none of them. Nothing where it has no groups.
std::string WithTlsRemarks(const gemmi::Structure &st, const std::string &pdb) {
const std::vector<gemmi::TlsGroup> *groups = st.meta.get_tls_groups();
if (groups == nullptr || groups->empty())
return pdb;
std::ostringstream r;
auto line = [&r](const std::string &text) { r << fmt::format("{:<80}\n", "REMARK 3" + text); };
line("");
line(" REFINEMENT.");
line(" PROGRAM : rugnux");
line("");
line(" DATA USED IN REFINEMENT.");
line("");
line(" TLS DETAILS");
line(fmt::format(" NUMBER OF TLS GROUPS : {:4}", groups->size()));
for (const gemmi::TlsGroup &g : *groups) {
line("");
line(fmt::format(" TLS GROUP : {:5}", g.id));
line(fmt::format(" NUMBER OF COMPONENTS GROUP : {:4}", g.selections.size()));
line(" COMPONENTS C SSSEQI TO C SSSEQI");
for (const gemmi::TlsGroup::Selection &sel : g.selections)
line(fmt::format(" RESIDUE RANGE :{:>5}{:>6} {:>5}{:>6}", sel.chain, sel.res_begin.str(), sel.chain,
sel.res_end.str()));
line(fmt::format(" ORIGIN FOR THE GROUP (A):{:9.4f}{:9.4f}{:9.4f}", g.origin.x, g.origin.y, g.origin.z));
line(" T TENSOR");
line(fmt::format(" T11: {:8.4f} T22: {:8.4f}", g.T.u11, g.T.u22));
line(fmt::format(" T33: {:8.4f} T12: {:8.4f}", g.T.u33, g.T.u12));
line(fmt::format(" T13: {:8.4f} T23: {:8.4f}", g.T.u13, g.T.u23));
line(" L TENSOR");
line(fmt::format(" L11: {:8.4f} L22: {:8.4f}", g.L.u11, g.L.u22));
line(fmt::format(" L33: {:8.4f} L12: {:8.4f}", g.L.u33, g.L.u12));
line(fmt::format(" L13: {:8.4f} L23: {:8.4f}", g.L.u13, g.L.u23));
line(" S TENSOR");
for (int i = 0; i < 3; i++)
line(fmt::format(" S{}1: {:8.4f} S{}2: {:8.4f} S{}3: {:8.4f}", i + 1, g.S[i][0], i + 1, g.S[i][1], i + 1,
g.S[i][2]));
}
// before the crystal's record; the model's own REMARK 3 left out
std::istringstream in(pdb);
std::ostringstream out;
std::string text;
bool written = false;
while (std::getline(in, text)) {
if (text.rfind("REMARK 3", 0) == 0)
continue;
if (!written && (text.rfind("CRYST1", 0) == 0 || text.rfind("ATOM", 0) == 0 || text.rfind("HETATM", 0) == 0)) {
out << r.str();
written = true;
}
out << text << '\n';
}
return out.str();
}
} // namespace
void WritePlacedModel(const gemmi::Structure &placed,
@@ -90,8 +145,10 @@ void WritePlacedModel(const gemmi::Structure &placed,
}
bool pdb_written = false;
try {
std::ostringstream pdb;
gemmi::write_pdb(st, pdb);
std::ofstream os(pdb_path);
gemmi::write_pdb(st, os);
os << WithTlsRemarks(st, pdb.str());
pdb_written = static_cast<bool>(os);
} catch (const std::exception &e) {
logger.Warning("Model validation: the placed model could not also be written as PDB ({}); "