From 6472ce3ff1bba9c4ab80d918344a5828bcdc1cb9 Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Tue, 27 Oct 2026 17:34:27 +0100 Subject: [PATCH] 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 Claude-Session: https://claude.ai/code/session_01HVeDQqN9qRD57vnShzTC37 --- docs/ACKNOWLEDGEMENT.md | 23 + docs/CPU_DATA_ANALYSIS_REFINEMENT.md | 46 +- docs/RUGNUX_REPORT.md | 5 +- .../structure_refinement/CMakeLists.txt | 1 + .../structure_refinement/ModelGrid.cpp | 2 +- .../structure_refinement/ModelGrid.h | 12 + .../ModelStructureFactorsGPU.cpp | 3 +- .../structure_refinement/Refinement.cpp | 703 +++++++++++++++++- .../structure_refinement/Refinement.h | 16 + .../structure_refinement/RefinementADP.cpp | 319 ++++++++ .../structure_refinement/RefinementADP.h | 78 ++ .../structure_refinement/RefinementGPU.cpp | 57 +- .../structure_refinement/RefinementGPU.cu | 82 +- .../RefinementGPUEngine.h | 3 + .../structure_refinement/RefinementXray.cpp | 126 +++- .../structure_refinement/RefinementXray.h | 40 +- rugnux/ResultReport.cpp | 4 + tests/RefinementTest.cpp | 253 +++++++ writer/WriteModel.cpp | 59 +- 19 files changed, 1764 insertions(+), 68 deletions(-) create mode 100644 image_analysis/structure_refinement/RefinementADP.cpp create mode 100644 image_analysis/structure_refinement/RefinementADP.h diff --git a/docs/ACKNOWLEDGEMENT.md b/docs/ACKNOWLEDGEMENT.md index ac8c70426..78164d014 100644 --- a/docs/ACKNOWLEDGEMENT.md +++ b/docs/ACKNOWLEDGEMENT.md @@ -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 diff --git a/docs/CPU_DATA_ANALYSIS_REFINEMENT.md b/docs/CPU_DATA_ANALYSIS_REFINEMENT.md index 0c9d74b4f..3dbb0b848 100644 --- a/docs/CPU_DATA_ANALYSIS_REFINEMENT.md +++ b/docs/CPU_DATA_ANALYSIS_REFINEMENT.md @@ -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). diff --git a/docs/RUGNUX_REPORT.md b/docs/RUGNUX_REPORT.md index 8865f17d4..720ba1798 100644 --- a/docs/RUGNUX_REPORT.md +++ b/docs/RUGNUX_REPORT.md @@ -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 `_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, 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 | diff --git a/image_analysis/structure_refinement/CMakeLists.txt b/image_analysis/structure_refinement/CMakeLists.txt index d7a5b82ef..e5cff4feb 100644 --- a/image_analysis/structure_refinement/CMakeLists.txt +++ b/image_analysis/structure_refinement/CMakeLists.txt @@ -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 diff --git a/image_analysis/structure_refinement/ModelGrid.cpp b/image_analysis/structure_refinement/ModelGrid.cpp index 4e3177182..365c308dc 100644 --- a/image_analysis/structure_refinement/ModelGrid.cpp +++ b/image_analysis/structure_refinement/ModelGrid.cpp @@ -451,7 +451,7 @@ void PutModelDensityOnGrid(gemmi::DensityCalculator &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); diff --git a/image_analysis/structure_refinement/ModelGrid.h b/image_analysis/structure_refinement/ModelGrid.h index 1eea50952..4552ed6e2 100644 --- a/image_analysis/structure_refinement/ModelGrid.h +++ b/image_analysis/structure_refinement/ModelGrid.h @@ -3,6 +3,8 @@ #pragma once +#include +#include #include #include @@ -22,6 +24,16 @@ // (empty in P1). Both functions below take them, for a grid of the same size. std::vector OrbitLeaders(const gemmi::Grid &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 +Real AnisoRadiusB(const gemmi::SMat33 &aniso_b) { + const gemmi::SMat33 m{aniso_b.u11, aniso_b.u22, aniso_b.u33, aniso_b.u12, aniso_b.u13, aniso_b.u23}; + const std::array ev = m.calculate_eigenvalues(); + return static_cast(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. diff --git a/image_analysis/structure_refinement/ModelStructureFactorsGPU.cpp b/image_analysis/structure_refinement/ModelStructureFactorsGPU.cpp index 7f2f9373a..5f08709e0 100644 --- a/image_analysis/structure_refinement/ModelStructureFactorsGPU.cpp +++ b/image_analysis/structure_refinement/ModelStructureFactorsGPU.cpp @@ -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(radius); a.max_r2 = static_cast(radius * radius); diff --git a/image_analysis/structure_refinement/Refinement.cpp b/image_analysis/structure_refinement/Refinement.cpp index ed119cd3f..cf1fdb2a5 100644 --- a/image_analysis/structure_refinement/Refinement.cpp +++ b/image_analysis/structure_refinement/Refinement.cpp @@ -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 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(gemmi::u_to_b() * a.aniso.trace() / 3.0); - a.aniso = {0, 0, 0, 0, 0, 0}; - } a.b_iso = static_cast(std::clamp(a.b_iso, B_FLOOR, B_CEILING_MAX)); } } @@ -143,6 +143,110 @@ struct Context { std::vector rider_offset; int evaluations = 0; const std::atomic *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 aniso; + std::vector 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 &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 &d = aniso[i].dev; + const gemmi::SMat33 dev{d[0], d[1], d[2], d[3], d[4], d[5]}; + const std::array 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 tls; + std::vector 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 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 b = u.scaled(gemmi::u_to_b()).added_kI(residual); + std::array 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(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((b.u11 - t) / gemmi::u_to_b()), static_cast((b.u22 - t) / gemmi::u_to_b()), + static_cast((b.u33 - t) / gemmi::u_to_b()), static_cast(b.u12 / gemmi::u_to_b()), + static_cast(b.u13 / gemmi::u_to_b()), static_cast(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> UOfAtoms() const { + std::vector> 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 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 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(B_CEILING_MAX / top)); + at.b_iso = static_cast(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(at.aniso.u11 - t), static_cast(at.aniso.u22 - t), + static_cast(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()); + } 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(b_min, a->b_iso); + for (size_t i = 0; i < atoms.size(); i++) + b_min = std::min(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(std::clamp(atoms[i]->b_iso + b_iso, B_FLOOR, b_ceiling)); - b_min = std::min(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(std::clamp(atoms[i]->b_iso + b_iso, B_FLOOR + lo, + std::max(B_FLOOR + lo, b_ceiling - hi))); + b_min = std::min(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(std::clamp(p[k], ctx_.b_floor, ctx_.b_ceiling)); + a.b_iso = static_cast(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> 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> FoldedUGradient(const Context &ctx, const std::vector> &gx) { + std::vector> 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 ¶m_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(body.atoms.size(), 1)))); + } + } + + std::vector Get() const { + std::vector p; + if (mode_ == Mode::U) { + const std::vector> 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 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(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 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> gx; + const double xray = ctx_.engine->Evaluate(ctx_.Model(), gradient != nullptr ? &gx : nullptr); + double restraints = 0; + std::vector> 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> 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 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 b{p[0], p[1], p[2], p[3], p[4], p[5]}; + const std::array 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 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(t); + AtomAniso &a = ctx_.aniso[i]; + a.on = true; + a.dev = {static_cast((b.u11 - t) / gemmi::u_to_b()), static_cast((b.u22 - t) / gemmi::u_to_b()), + static_cast((b.u33 - t) / gemmi::u_to_b()), static_cast(b.u12 / gemmi::u_to_b()), + static_cast(b.u13 / gemmi::u_to_b()), static_cast(b.u23 / gemmi::u_to_b())}; + } + + Context &ctx_; + Mode mode_; + const std::vector ¶m_atoms_; + double weight_; + const URestraints *restraints_; + std::vector 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 BDiagonal(const std::vector &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 +void Minimise(Context &ctx, const Target &target, int iterations, RefinementMinimiser &minimiser, std::vector diagonal = {}) { std::vector 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(std::clamp(base_[i] + q[k], ctx_.b_floor, ctx_.b_ceiling)); + ctx_.atoms[i]->b_iso = static_cast(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 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 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 OwnSlotDevices(const Context &main, int wanted, const Refinemen return std::vector(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 bodies; // Tls, with the part of each atom's B they hold at the start (Context::tls_b) + std::vector tls_b; + std::vector> 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 pos; std::vector b; + std::vector aniso; + std::vector tls; + std::vector 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 &refinable, const std::vector &pos0, - const std::vector &b0, const Context::Fit *fit, double w, double adp_w, int cycles, - const RefinementSettings &settings) { + const std::vector &b0, const std::vector &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 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 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 &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 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 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 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 xray = ctx.XrayBCurvatureShape(); @@ -1165,15 +1479,26 @@ TrialOutcome RunTrial(TrialSlot slot, const std::vector &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 &refinable, const s std::vector RunTrials(const std::vector &trial_slots, const std::vector> &trials, const std::vector &refinable, const std::vector &pos0, - const std::vector &b0, const Context::Fit *fit, int cycles, - const RefinementSettings &settings) { + const std::vector &b0, const std::vector &aniso0, + const Context::Fit *fit, int cycles, const RefinementSettings &settings, + const std::vector &hyps = {}) { std::vector out(trials.size()); std::vector 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 RunTrials(const std::vector &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 RunTrials(const std::vector &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> 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 &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(b_min, ctx.atoms[i]->b_iso); + } + body.origin /= static_cast(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 AdpHypotheses(const Context &ctx, const RefinementTopology &topo, + const std::vector &input) { + constexpr size_t MIN_TLS_ATOMS = 50; + constexpr int MIN_TLS_RESIDUES = 10; + std::vector 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>> chains; + int i = 0; + for (const gemmi::Chain &ch : ctx.st.models[0].chains) { + std::vector> residues; + size_t atoms = 0; + for (const gemmi::Residue &r : ch.residues) { + std::vector 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 pos = ctx.Positions(); + const std::vector b = ctx.B(); + size_t last = 0; + for (int k : {1, 2, 4, 8, 16}) { + std::vector> groups; + for (const std::vector> &residues : chains) { + std::vector starts = TlsSegments(residues, pos, b, k, MIN_TLS_RESIDUES); + starts.push_back(static_cast(residues.size())); + for (size_t s = 0; s + 1 < starts.size(); s++) { + std::vector 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(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 &u_atoms) { + std::vector> gx; + ctx.engine->Evaluate(ctx.Model(), &gx); + const std::vector> gu = FoldedUGradient(ctx, gx); + std::vector> u = ctx.UOfAtoms(); + std::mt19937 rng(3); + std::normal_distribution 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> 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 &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>> 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 b0 = ctx.B(); + const std::vector aniso0 = ctx.aniso; + struct Reading { + const char *name; + std::vector b; + std::vector aniso; + }; + std::vector readings; + for (int residual = 0; residual < 2; residual++) { + Reading read{residual ? "TLS, B residual" : "TLS, B total", b0, + aniso0.empty() ? std::vector(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(u.u11 - t), static_cast(u.u22 - t), + static_cast(u.u33 - t), static_cast(u.u12), static_cast(u.u13), + static_cast(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 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(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(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 &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(ctx.atoms.size()); double b_max = 0; for (const gemmi::Atom *a : ctx.atoms) - b_max = std::max(b_max, a->b_iso); + b_max = std::max(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 tls_bodies = + ctx.st.meta.get_tls_groups() != nullptr ? TlsBodiesFromModel(ctx.Model(), *ctx.st.meta.get_tls_groups()) + : std::vector(); + 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::vectorOnGPU(); - RFactors r = ctx.Refit(); + std::vector 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 *groups = ctx.st.meta.get_tls_groups()) + result.tls_groups = static_cast(groups->size()); result.model = std::make_shared(ctx.st); result.ok = true; result.seconds = std::chrono::duration(std::chrono::steady_clock::now() - start).count(); @@ -1464,7 +2018,8 @@ RefinementResult RefineOnce(const gemmi::Structure &placed, const std::vector outcomes = RunTrials(trial_slots, ladder, refinable, pos0, b0, nullptr, + const std::vector aniso0 = ctx.aniso; + std::vector 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 ran = RunTrials(trial_slots, missing, refinable, state.pos, state.b, &state.fit, - settings.macro_cycles, settings); + std::vector 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= 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 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 keep = ctx.aniso; + std::vector all = keep.empty() ? std::vector(ctx.atoms.size()) : keep; + std::vector refines_u(ctx.atoms.size(), 0); + std::vector 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> trials = {{w, adp}}; + std::vector 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(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(added.size()); diff --git a/image_analysis/structure_refinement/Refinement.h b/image_analysis/structure_refinement/Refinement.h index cefde1799..a0dd1b16f 100644 --- a/image_analysis/structure_refinement/Refinement.h +++ b/image_analysis/structure_refinement/Refinement.h @@ -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 trials; + // The displacement model kept (RefinementSettings::adp_hypotheses): "isotropic", "TLS, 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 adp_trials; + int aniso_atoms = 0, tls_groups = 0; GeometryStatistics geometry_start, geometry; int atoms = 0, refined_atoms = 0; diff --git a/image_analysis/structure_refinement/RefinementADP.cpp b/image_analysis/structure_refinement/RefinementADP.cpp new file mode 100644 index 000000000..810400540 --- /dev/null +++ b/image_analysis/structure_refinement/RefinementADP.cpp @@ -0,0 +1,319 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#include "RefinementADP.h" + +#include +#include +#include + +#include + +namespace { + +constexpr double DEG = 3.14159265358979323846 / 180.0; + +} // namespace + +gemmi::SMat33 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 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 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 &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, TlsBody::PARAMETERS> TlsBody::Derivatives(const gemmi::Position &pos) const { + std::array, PARAMETERS> d{}; + TlsBody unit; + unit.origin = origin; + for (int k = 0; k < PARAMETERS; k++) { + std::array p{}; + p[k] = 1; + unit.Set(p); + const gemmi::SMat33 u = unit.U(pos); + d[k] = {u.u11, u.u22, u.u33, u.u12, u.u13, u.u23}; + } + return d; +} + +std::array TlsBody::Gradient(const std::vector &atom_list, + const std::vector> &grad_u) const { + gemmi::Mat33 mt(0), ml(0), ms(0); + for (int i : atoms) { + const std::array &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 TlsBodiesFromModel(const gemmi::Model &model, const std::vector &groups) { + std::vector 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(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(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 TlsBodiesToGroups(const gemmi::Model &model, const std::vector &bodies) { + // per atom its body, per residue the body of its first atom in one + std::vector 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(b); + std::vector 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(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> &b_similarity_pairs, const std::vector &refines_u) { + URestraints r; + std::vector 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(i)); + return r; +} + +double EvaluateURestraints(const URestraints &r, const std::vector &pos, + const std::vector> &u, std::vector> *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 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 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 TlsSegments(const std::vector> &residue_atoms, const std::vector &pos, + const std::vector &b, int k, int min_residues) { + using Mat = Eigen::Matrix; + using Vec = Eigen::Matrix; + const int n = static_cast(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 block_first; + for (int r = 0; r < n; r += per_block) + block_first.push_back(r); + const int nb = static_cast(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 &ra : residue_atoms) + for (int a : ra) { + c += pos[a]; + count++; + } + if (count == 0) + return {0}; + c /= count; + std::vector ata(nb + 1, Mat::Zero()); + std::vector atb(nb + 1, Vec::Zero()); + std::vector 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::infinity(); + // best[s][j]: blocks [0, j) in s + 1 segments; from[s][j]: where the last of them starts + std::vector> best(k, std::vector(nb + 1, INF)); + std::vector> from(k, std::vector(nb + 1, 0)); + std::vector> seg(nb + 1, std::vector(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 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; +} diff --git a/image_analysis/structure_refinement/RefinementADP.h b/image_analysis/structure_refinement/RefinementADP.h new file mode 100644 index 000000000..0c817ba64 --- /dev/null +++ b/image_analysis/structure_refinement/RefinementADP.h @@ -0,0 +1,78 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#pragma once + +#include +#include +#include + +#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 T{0, 0, 0, 0, 0, 0}, L{0, 0, 0, 0, 0, 0}; + gemmi::Mat33 S{0}; + std::vector atoms; // in model order + + gemmi::SMat33 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 Get() const; + void Set(const std::array &p); + // dU/dp at `pos`: per parameter the six U11 U22 U33 U12 U13 U23 (U is linear in the parameters). + std::array, 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 Gradient(const std::vector &atoms, + const std::vector> &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 TlsBodiesFromModel(const gemmi::Model &model, const std::vector &groups); + +// The bodies as the metadata holds them (degrees), each selection one run of consecutive residues of a chain. +std::vector TlsBodiesToGroups(const gemmi::Model &model, const std::vector &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> pairs; + std::vector 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> &b_similarity_pairs, const std::vector &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 &pos, + const std::vector> &u, std::vector> *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 TlsSegments(const std::vector> &residue_atoms, const std::vector &pos, + const std::vector &b, int k, int min_residues); diff --git a/image_analysis/structure_refinement/RefinementGPU.cpp b/image_analysis/structure_refinement/RefinementGPU.cpp index dc886f577..06c28db73 100644 --- a/image_analysis/structure_refinement/RefinementGPU.cpp +++ b/image_analysis/structure_refinement/RefinementGPU.cpp @@ -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; +// 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 &u, + const gemmi::DensityCalculator &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(radius); + a.max_r2 = static_cast(radius * radius); + for (int k = 0; k < 5; k++) { + a.a[k] = precal.a[k]; + const gemmi::SMat33 &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 &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 Calculator(double blur) const { gemmi::DensityCalculator dc; dc.d_min = d_min_; @@ -302,16 +343,23 @@ private: const int n = static_cast(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 *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 engine_; std::vector atom_ptrs_; // the model's atoms at the last SetAtoms() std::vector last_b_; + std::vector last_aniso_; PinnedVector atoms_; // pinned: a B rung sends them every evaluation (37 MB on a large model) PinnedVector> pos_; std::vector changed_; diff --git a/image_analysis/structure_refinement/RefinementGPU.cu b/image_analysis/structure_refinement/RefinementGPU.cu index 90d9672ec..f8bffe29e 100644 --- a/image_analysis/structure_refinement/RefinementGPU.cu +++ b/image_analysis/structure_refinement/RefinementGPU.cu @@ -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(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(g.nu) * ((mv < g.nv ? mv : mv - g.nv) + static_cast(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(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(FOUR_PI2)); + else + grad[a] = make_float4(s * gx, s * gy, s * gz, o * (gu[0] + gu[1] + gu[2]) / static_cast(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 host_sums; CudaHostPtr host_grad, host_xyzb, host_restraint_grad; + CudaHostPtr host_grad_u; // six per atom CudaDevicePtr mask_atoms; CudaDevicePtr grid; CudaDevicePtr spectrum; @@ -503,6 +539,8 @@ struct RefinementGPUEngineImpl { CudaDevicePtr phase; // exp(2 pi i h.t) over the integer h.t from phase_min int phase_min = 0; CudaDevicePtr grad; + CudaDevicePtr grad_u; + int n_aniso = 0; // the anisotropic atoms of the last SetDensities() std::unique_ptr density; std::unique_ptr mask; CudaFFTPlan forward, inverse, mask_forward; @@ -516,10 +554,12 @@ struct RefinementGPUEngineImpl { pos = CudaDevicePtr(na, sync); host_sums = CudaHostPtr(2); host_grad = CudaHostPtr(na); + host_grad_u = CudaHostPtr(6 * na); host_xyzb = CudaHostPtr(na); host_restraint_grad = CudaHostPtr(na); mask_atoms = CudaDevicePtr(na, sync); grad = CudaDevicePtr(na, sync); + grad_u = CudaDevicePtr(6 * na, sync); grid = CudaDevicePtr(GridPoints(setup.grid), sync); spectrum = CudaDevicePtr(ComplexPoints(setup.grid), sync); fft_work = CudaDevicePtr(std::max(work_bytes, 1), sync); @@ -662,10 +702,14 @@ struct RefinementGPUEngineImpl { const size_t threads = static_cast(n_atoms) * WARP; atom_gradient_kernel<<>>(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(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 *pos, size_t n) { @@ -850,6 +898,10 @@ const std::array *RefinementGPUEngine::HostGradient() const { return reinterpret_cast *>(impl_->host_grad.get()); } +const float *RefinementGPUEngine::HostUGradient() const { + return impl_->host_grad_u.get(); +} + const std::array *RefinementGPUEngine::HostRestraintGradient() const { return reinterpret_cast *>(impl_->host_restraint_grad.get()); } diff --git a/image_analysis/structure_refinement/RefinementGPUEngine.h b/image_analysis/structure_refinement/RefinementGPUEngine.h index 4b8c86d7c..052d719ae 100644 --- a/image_analysis/structure_refinement/RefinementGPUEngine.h +++ b/image_analysis/structure_refinement/RefinementGPUEngine.h @@ -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 *HostGradient() const; const std::array *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 &records, const std::vector &plane_atoms, diff --git a/image_analysis/structure_refinement/RefinementXray.cpp b/image_analysis/structure_refinement/RefinementXray.cpp index faba0a11d..66c47c257 100644 --- a/image_analysis/structure_refinement/RefinementXray.cpp +++ b/image_analysis/structure_refinement/RefinementXray.cpp @@ -116,7 +116,7 @@ public: void Fcalc(const gemmi::Model &model, std::vector> &f) override { std::vector> 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> *grad) override { + double Evaluate(const gemmi::Model &model_in, std::vector> *grad) override { + const gemmi::Model &model = WithAniso(model_in); std::vector> 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> &f) const { gemmi::DensityCalculator 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 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 AnisoU(const AtomAniso &a, double b) { + const double u = b / gemmi::u_to_b(); + gemmi::SMat33 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 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(m.u11), static_cast(m.u22), static_cast(m.u33), + static_cast(m.u12), static_cast(m.u13), static_cast(m.u23)}; +} + +void ApplyAniso(gemmi::Model &model, const std::vector &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> ModelXyzB(const gemmi::Model &model) { std::vector> out; for (const gemmi::Chain &ch : model.chains) @@ -458,8 +496,82 @@ std::unique_ptr MakeRefinementEngineCPU(const gemmi::UnitCell return std::make_unique(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 &map, const gemmi::DensityCalculator &dc, + const gemmi::Atom &atom, double blur, std::array &grad, std::array *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(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(om.a11), m01 = static_cast(om.a12), m02 = static_cast(om.a13), + m11 = static_cast(om.a22), m12 = static_cast(om.a23), m22 = static_cast(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(n3[k]); + c[k] = gemmi::iround(n3[k]); + } + d[0] = std::min(static_cast(std::ceil(radius / map.spacing[0])), nu - 1); + d[1] = std::min(static_cast(std::ceil(radius / map.spacing[1])), nv - 1); + d[2] = std::min(static_cast(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(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(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(nu) * (gemmi::modulo(v, nv) + static_cast(nv) * w_); + for (int u = c[0] - d[0]; u <= c[0] + d[0]; u++) { + const float x = m00 * (nctr[0] - static_cast(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 &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 &map, const gemmi::Model &model, double blur, size_t nthreads, - std::vector> &grad) { + std::vector> &grad, std::vector> *grad_u) { const gemmi::DensityCalculator dc = Calculator(map.unit_cell, *map.spacegroup, 1.0, blur); std::vector atoms; for (const gemmi::Chain &ch : model.chains) @@ -467,6 +579,8 @@ void AtomGradientsFromMap(const gemmi::Grid &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 &map, const gemmi::Model &mod ParallelChunks(static_cast(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(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); diff --git a/image_analysis/structure_refinement/RefinementXray.h b/image_analysis/structure_refinement/RefinementXray.h index d7062afbf..9c51ba9ed 100644 --- a/image_analysis/structure_refinement/RefinementXray.h +++ b/image_analysis/structure_refinement/RefinementXray.h @@ -101,6 +101,20 @@ RFactors ComputeRFactors(const RefinementReflections &refl, const std::vector> 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 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 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 &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> &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> *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 *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> &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 *aniso_ = nullptr; + std::vector> grad_u_; + private: RestraintSet restraints_[2]; }; @@ -179,6 +212,9 @@ std::unique_ptr 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 &map, const gemmi::Model &model, double blur, size_t nthreads, - std::vector> &grad); + std::vector> &grad, + std::vector> *grad_u = nullptr); diff --git a/rugnux/ResultReport.cpp b/rugnux/ResultReport.cpp index ec448e6b6..0096ead32 100644 --- a/rugnux/ResultReport.cpp +++ b/rugnux/ResultReport.cpp @@ -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)); diff --git a/tests/RefinementTest.cpp b/tests/RefinementTest.cpp index d9fd7a09f..134fdc5b3 100644 --- a/tests/RefinementTest.cpp +++ b/tests/RefinementTest.cpp @@ -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 random_aniso(const gemmi::Model &model, unsigned seed) { + std::mt19937 rng(seed); + std::uniform_real_distribution d(-1.0, 1.0); + std::vector 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(v[0] - t), static_cast(v[1] - t), static_cast(v[2] - t), + static_cast(v[3] / 2), static_cast(v[4] / 2), static_cast(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 &aniso, + double tolerance) { + std::vector> grad; + engine.SetAniso(&aniso); + engine.Evaluate(model, &grad); + const std::vector> 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(v0 + h); + const double fp = engine.Evaluate(model, nullptr); + aniso[i].dev[c] = static_cast(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 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> f; + std::vector 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> 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> gc, gg; + cpu->Evaluate(model, &gc); + const std::vector> uc = cpu->UGradient(); + gpu->Evaluate(model, &gg); + const std::vector> 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 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 atoms = atoms_of(model); + TlsBody body; + body.origin = gemmi::Position(9, -2, 4); + std::array 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 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> gu(atoms.size()); + for (auto &g : gu) + for (double &x : g) + x = d(rng); + const std::array 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 groups = TlsBodiesToGroups(m2, {body}); + REQUIRE(groups.size() == 1); + const std::vector back = TlsBodiesFromModel(m2, groups); + REQUIRE(back.size() == 1); + CHECK(back[0].atoms.size() == 30); + const gemmi::SMat33 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 d(-1.0, 1.0); + const int n = 6; + std::vector pos; + std::vector> 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{5}); + std::vector> 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> 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 line = {{0, 0, 0}, {0, 0, 1.5}}; + const std::vector> 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> residues; + std::vector pos; + std::vector b; + for (int r = 0; r < 60; r++) { + std::vector 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(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{0}); + CHECK(TlsSegments(residues, pos, b, 2, 10) == std::vector{0, 25}); + // too short for three segments of 25 + CHECK(TlsSegments(residues, pos, b, 3, 25).size() == 2); +} diff --git a/writer/WriteModel.cpp b/writer/WriteModel.cpp index 745ebf001..932a70366 100644 --- a/writer/WriteModel.cpp +++ b/writer/WriteModel.cpp @@ -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 *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(os); } catch (const std::exception &e) { logger.Warning("Model validation: the placed model could not also be written as PDB ({}); "