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 ({}); "