From b4b68d368f4e26a02e4f68b1d3a10a4a6cf981bf Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Sun, 25 Oct 2026 01:22:54 +0200 Subject: [PATCH] Restrained refinement on the GPU: ML amplitude target, FFT gradient, dictionary restraints, R1 + R2 driver, blobs RefineModel (Refinement.h) refines a placed model against the merge, reading and writing nothing: - placement stage (R1): every chain of 50+ atoms as a rigid body at once (L-BFGS on six parameters per chain, their gradient from the per-atom one through Ceres jets), then one B shift per residue; - restrained stage (R2): positions and isotropic B of every atom of a residue the embedded dictionary describes, the rest held fixed but kept in Fcalc and the mask; riding hydrogens added by gemmi's topology, moved with their parents, re-placed every macro-cycle. Macro-cycles refit the mask, the scale and bulk solvent (FitModelScale / ModelScaleGPUPool) and the likelihood's D and beta (ML on the free set in resolution bins), rebuild the repulsion pairs (symmetry mates included), then L-BFGS on positions and, where chosen, on B. - weights: a ladder of full refinements at multiples of the gradient-balanced X-ray weight, chosen by R-free within bond/angle bounds (ties to the tighter geometry), then the B refinement tried at its own weights and kept only where it beats B held by more than the tie. The target and its gradient (RefinementXray, RefinementGPU): the Rice/centric amplitude likelihood (RefinementMath.h, shared host/device code), the gradient by Agarwal's FFT method - the adjoint of the reciprocal-space symmetry composition into a half-l (CPU) or half-u (cuFFT) spectrum, one inverse real transform, then each atom's density derivatives (xyz and B) summed over its box. On the GPU everything per evaluation stays on the device: ModelDensityGPU's gather, r2c, a fused compose + likelihood kernel, a CSR gather for the adjoint (no atomics), c2r, one warp per atom for the gradient; the restraints are one kernel per restraint plus a per-atom gather of gradient slots; cub reductions. Deterministic. Restraints (RefinementRestraints): bonds, angles, torsions (with esd, no phi/psi, no alternative sugar puckers), chiral volumes, least-squares planes from gemmi's Topo; repulsions beyond 1-4 (vdW radii from the dictionary's energy library, 2.6 A for H-bond pairs, 1.9 A metal-N/O/S); B similarity of bonded and angle-related atoms. DensityBlobs (D1): gemmi's flood-fill blobs of mFo-DFc above 3 sigma peaking above 4 sigma, placed beside the model and labelled WATER / ION / LIGAND / MODEL_ERROR by distance, contacts, volume and the anomalous map. tools/rugnux_refine_test refines a model against a finished run's p.mtz (developer utility, not installed). Tests: likelihood derivative, the embedded library, restraint values on ideal residues and every restraint kind's gradient against finite differences (CPU, and the GPU against the CPU), the X-ray gradient against finite differences in P1, P212121, C2 and P6122 (CPU and GPU), and a synthetic refinement that recovers a shaken model. Credits in docs/ACKNOWLEDGEMENT.md and at the algorithms. Measured on the open arm's battery outputs against REFMAC 5.8, 10 cycles, same p.mtz, free set and start model (R-free): 5reo 0.2228 vs 0.2236, 8dyz 0.1733 vs 0.1651, 9chw 0.1995 vs 0.1969 (CPU engine). Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_015CeEGgpKdfkCTJb3YNMJcu (cherry picked from commit 8efe6f2322b3bfdae580bb34f256e67730dfc3c2) --- docs/ACKNOWLEDGEMENT.md | 45 + .../structure_refinement/CMakeLists.txt | 8 + .../structure_refinement/DensityBlobs.cpp | 110 +++ .../structure_refinement/DensityBlobs.h | 41 + .../structure_refinement/Refinement.cpp | 781 ++++++++++++++++++ .../structure_refinement/Refinement.h | 103 +++ .../structure_refinement/RefinementGPU.cpp | 283 +++++++ .../structure_refinement/RefinementGPU.cu | 691 ++++++++++++++++ .../structure_refinement/RefinementGPU.h | 24 + .../RefinementGPUEngine.h | 78 ++ .../structure_refinement/RefinementMath.h | 268 ++++++ .../RefinementRestraints.cpp | 568 +++++++++++++ .../RefinementRestraints.h | 98 +++ .../structure_refinement/RefinementXray.cpp | 396 +++++++++ .../structure_refinement/RefinementXray.h | 117 +++ tests/CMakeLists.txt | 1 + tests/RefinementTest.cpp | 452 ++++++++++ tools/CMakeLists.txt | 5 + tools/rugnux_refine_test.cpp | 151 ++++ 19 files changed, 4220 insertions(+) create mode 100644 image_analysis/structure_refinement/DensityBlobs.cpp create mode 100644 image_analysis/structure_refinement/DensityBlobs.h create mode 100644 image_analysis/structure_refinement/Refinement.cpp create mode 100644 image_analysis/structure_refinement/Refinement.h create mode 100644 image_analysis/structure_refinement/RefinementGPU.cpp create mode 100644 image_analysis/structure_refinement/RefinementGPU.cu create mode 100644 image_analysis/structure_refinement/RefinementGPU.h create mode 100644 image_analysis/structure_refinement/RefinementGPUEngine.h create mode 100644 image_analysis/structure_refinement/RefinementMath.h create mode 100644 image_analysis/structure_refinement/RefinementRestraints.cpp create mode 100644 image_analysis/structure_refinement/RefinementRestraints.h create mode 100644 image_analysis/structure_refinement/RefinementXray.cpp create mode 100644 image_analysis/structure_refinement/RefinementXray.h create mode 100644 tests/RefinementTest.cpp create mode 100644 tools/rugnux_refine_test.cpp diff --git a/docs/ACKNOWLEDGEMENT.md b/docs/ACKNOWLEDGEMENT.md index 7238bb34f..51a688911 100644 --- a/docs/ACKNOWLEDGEMENT.md +++ b/docs/ACKNOWLEDGEMENT.md @@ -508,6 +508,51 @@ molecular replacement; the error itself is AlphaFold's: J. Jumper, R. Evans, A. accurate protein structure prediction with AlphaFold" (2021), Nature 596, 583-589 [doi:10.1038/s41586-021-03819-2](https://doi.org/10.1038/s41586-021-03819-2). +**Structure refinement** (`--refine`) — the placed model is refined against the maximum-likelihood +amplitude target, the Rice distribution of an observed amplitude given the model's (its centric +counterpart for a centric reflection): N. S. Pannu and R. J. Read, "Improved structure refinement +through maximum likelihood" (1996), Acta Cryst. A52, 659-668 +[doi:10.1107/S0108767396004370](https://doi.org/10.1107/S0108767396004370); G. N. Murshudov, +A. A. Vagin and E. J. Dodson, "Refinement of macromolecular structures by the maximum-likelihood method" +(1997), Acta Cryst. D53, 240-255 [doi:10.1107/S0907444996012255](https://doi.org/10.1107/S0907444996012255). +Its two parameters, D and the variance, are fitted by maximum likelihood on the free reflections in +resolution bins: V. Y. Lunin and T. P. Skovoroda, "R-free likelihood-based estimates of errors for phases +calculated from atomic models" (1995), Acta Cryst. A51, 880-887 +[doi:10.1107/S010876739500688X](https://doi.org/10.1107/S010876739500688X); V. Y. Lunin, P. V. Afonine and +A. G. Urzhumtsev, "Likelihood-based refinement. I. Irremovable model errors" (2002), Acta Cryst. A58, +270-282 [doi:10.1107/S0108767302001046](https://doi.org/10.1107/S0108767302001046). The gradient with +respect to every atom is taken by the FFT method - the derivatives of the target with respect to the +structure factors transformed into a map and summed against each atom's density derivatives: +R. C. Agarwal, "A new least-squares refinement technique based on the fast Fourier transform algorithm" +(1978), Acta Cryst. A34, 791-809 [doi:10.1107/S0567739478001618](https://doi.org/10.1107/S0567739478001618). +The positions and B factors are minimised by L-BFGS (Ceres): D. C. Liu and J. Nocedal, "On the limited +memory BFGS method for large scale optimization" (1989), Math. Program. 45, 503-528 +[doi:10.1007/BF01589116](https://doi.org/10.1007/BF01589116). The geometry restraints are those of the CCP4 +Monomer Library, a subset of which rugnux carries (see THIRD_PARTY_NOTICES.md): A. A. Vagin, +R. A. Steiner, A. A. Lebedev, L. Potterton, S. McNicholas, F. Long and G. N. Murshudov, "REFMAC5 +dictionary: organization of prior chemical knowledge and guidelines for its use" (2004), Acta Cryst. D60, +2184-2195 [doi:10.1107/S0907444904023510](https://doi.org/10.1107/S0907444904023510); applied to the model +through GEMMI's topology (links, modifications, planes, chiral volumes, torsions): M. Wojdyr, "GEMMI: A +library for structural biology" (2022), J. Open Source Softw. 7, 4200 +[doi:10.21105/joss.04200](https://doi.org/10.21105/joss.04200). The torsion derivatives follow +A. Blondel and M. Karplus, "New formulation for derivatives of torsion angles and improper torsion angles +in molecular mechanics: elimination of singularities" (1996), J. Comput. Chem. 17, 1132-1141 +[doi:10.1002/(SICI)1096-987X(19960715)17:9<1132::AID-JCC5>3.0.CO;2-T](https://doi.org/10.1002/(SICI)1096-987X(19960715)17:9%3C1132::AID-JCC5%3E3.0.CO;2-T). +The B factors of bonded atoms are restrained to each other, as REFMAC5 restrains them (Murshudov et al. +1997, above). The weight of the X-ray term against the geometry starts from the ratio of the two +gradients and is chosen among several full refinements by R-free with the geometry held within bounds, +as phenix.refine chooses it: P. V. Afonine, R. W. Grosse-Kunstleve, N. Echols, J. J. Headd, +N. W. Moriarty, M. Mustyakimov, T. C. Terwilliger, A. Urzhumtsev, P. H. Zwart and P. D. Adams, +"Towards automated crystallographic structure refinement with phenix.refine" (2012), Acta Cryst. D68, +352-367 [doi:10.1107/S0907444912001308](https://doi.org/10.1107/S0907444912001308). R-free itself: +A. T. Brünger, "Free R value: a novel statistical quantity for assessing the accuracy of crystal +structures" (1992), Nature 355, 472-475 [doi:10.1038/355472a0](https://doi.org/10.1038/355472a0). + +**Unmodelled density** — the blobs `--refine` 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 +(`find_blobs_by_flood_fill`): P. Emsley, B. Lohkamp, W. G. Scott and K. Cowtan, "Features and development +of Coot" (2010), Acta Cryst. D66, 486-501 [doi:10.1107/S0907444910007493](https://doi.org/10.1107/S0907444910007493). + ## Software and computing methods Decoding bitshuffle+LZ4 images on the GPU, rather than decompressing them on the host and uploading diff --git a/image_analysis/structure_refinement/CMakeLists.txt b/image_analysis/structure_refinement/CMakeLists.txt index 1f98395e0..44d65a19d 100644 --- a/image_analysis/structure_refinement/CMakeLists.txt +++ b/image_analysis/structure_refinement/CMakeLists.txt @@ -18,6 +18,14 @@ ADD_LIBRARY(JFJochStructureRefinement STATIC $<$:ModelScaleGPU.cu> ModelScaleGPU.h SigmaA.cpp SigmaA.h MonomerLibrary.cpp MonomerLibrary.h + RefinementMath.h + RefinementRestraints.cpp RefinementRestraints.h + RefinementXray.cpp RefinementXray.h + Refinement.cpp Refinement.h + DensityBlobs.cpp DensityBlobs.h + $<$:RefinementGPU.cpp> + $<$:RefinementGPU.cu> + RefinementGPU.h RefinementGPUEngine.h ${CMAKE_CURRENT_BINARY_DIR}/MonomerLibraryData.cpp) # The subset of the CCP4 Monomer Library that refinement restrains with (monomers/, chosen by diff --git a/image_analysis/structure_refinement/DensityBlobs.cpp b/image_analysis/structure_refinement/DensityBlobs.cpp new file mode 100644 index 000000000..59794e77b --- /dev/null +++ b/image_analysis/structure_refinement/DensityBlobs.cpp @@ -0,0 +1,110 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#include "DensityBlobs.h" + +#include +#include + +#include "gemmi/blob.hpp" +#include "gemmi/neighbor.hpp" + +namespace { + +constexpr double CUTOFF_SIGMA = 3.0; +constexpr double PEAK_SIGMA = 4.0; +// Closer than this to a model atom, the density is the model's own atom out of place. +constexpr double MODEL_ERROR_DISTANCE = 1.6; +// Waters and ions sit at least this far from the model's atoms. +constexpr double CONTACT_MIN = 2.2; +// An ion's peak, in sigma, where the anomalous map gives no help. +constexpr double ION_PEAK_SIGMA = 8.0; +constexpr double ION_ANOMALOUS_SIGMA = 4.0; +// The smallest blob called ligand-like, A^3: at high resolution a strong single atom's blob already spans four of +// the resolution's spheres. +constexpr double LIGAND_MIN_VOLUME = 8.0; + +bool Polar(const gemmi::Element &el) { + return el == gemmi::El::N || el == gemmi::El::O; +} + +} // namespace + +const char *DensityBlobKindName(DensityBlob::Kind kind) { + switch (kind) { + case DensityBlob::Kind::Water: return "WATER"; + case DensityBlob::Kind::Ion: return "ION"; + case DensityBlob::Kind::Ligand: return "LIGAND"; + case DensityBlob::Kind::ModelError: return "MODEL_ERROR"; + case DensityBlob::Kind::Unassigned: break; + } + return "UNASSIGNED"; +} + +std::vector FindDensityBlobs(const gemmi::Grid &difference, double rms, const gemmi::Structure &model, + double d_min, const gemmi::Grid *anomalous, double anomalous_rms) { + std::vector out; + if (!(rms > 0) || model.models.empty()) + return out; + gemmi::Grid map = difference; + for (float &v : map.data) + v = static_cast(v / rms); + gemmi::BlobCriteria criteria; + criteria.cutoff = CUTOFF_SIGMA; + criteria.min_volume = 0.0; + criteria.min_score = 0.0; + criteria.min_peak = PEAK_SIGMA; + const std::vector blobs = gemmi::find_blobs_by_flood_fill(map, criteria); + + gemmi::Model copy = model.models[0]; + gemmi::UnitCell cell = map.unit_cell; + cell.set_cell_images_from_spacegroup(map.spacegroup); + gemmi::NeighborSearch ns(copy, cell, 5.0); + ns.populate(false); + // A single atom's blob at this resolution: a sphere of half the resolution's radius. Four of them is ligand size. + const double atom_volume = 4.0 / 3.0 * 3.14159265358979 * std::pow(0.5 * d_min, 3); + + for (const gemmi::Blob &b : blobs) { + DensityBlob d; + d.peak_sigma = b.peak_value; + d.volume = b.volume; + d.score = b.score; + d.centroid = b.centroid; + // the copy of the blob beside the atom nearest to it + if (gemmi::NeighborSearch::Mark *m = ns.find_nearest_atom(b.centroid, 12.0)) { + const gemmi::CRA cra = m->to_cra(copy); + const gemmi::FTransform t = ns.get_image_transformation(m->image_idx); + const gemmi::Fractional fc = cell.fractionalize(b.centroid); + const gemmi::Fractional fi = t.apply(cell.fractionalize(cra.atom->pos)); + const gemmi::Fractional shifted(fc.x - std::round(fc.x - fi.x), fc.y - std::round(fc.y - fi.y), + fc.z - std::round(fc.z - fi.z)); + d.centroid = cell.orthogonalize(gemmi::FTransform(t.inverse()).apply(shifted)); + d.nearest_atom_distance = d.centroid.dist(cra.atom->pos); + d.nearest_atom = cra.atom->name + " " + cra.residue->name + " " + cra.chain->name + + cra.residue->seqid.str(); + } else + d.nearest_atom_distance = INFINITY; + ns.for_each(d.centroid, '\0', 3.4, [&](gemmi::NeighborSearch::Mark &m, double dist_sq) { + const double r = std::sqrt(dist_sq); + if (Polar(m.element) && r >= 2.4) + d.polar_contacts++; + if ((Polar(m.element) || m.element == gemmi::El::S) && r >= 1.8 && r <= 2.6) + d.metal_contacts++; + }); + if (anomalous != nullptr && anomalous_rms > 0) + d.anomalous_sigma = anomalous->interpolate_value(b.peak_pos) / anomalous_rms; + + if (d.nearest_atom_distance < MODEL_ERROR_DISTANCE) + d.kind = DensityBlob::Kind::ModelError; + else if (d.nearest_atom_distance >= 1.8 && d.metal_contacts >= 2 && + (d.anomalous_sigma >= ION_ANOMALOUS_SIGMA || d.peak_sigma >= ION_PEAK_SIGMA)) + d.kind = DensityBlob::Kind::Ion; + else if (d.volume > std::max(4 * atom_volume, LIGAND_MIN_VOLUME) && d.nearest_atom_distance >= CONTACT_MIN) + d.kind = DensityBlob::Kind::Ligand; + else if (d.nearest_atom_distance >= CONTACT_MIN && d.polar_contacts >= 1) + d.kind = DensityBlob::Kind::Water; + out.push_back(d); + } + std::sort(out.begin(), out.end(), [](const DensityBlob &a, const DensityBlob &b) { return a.score > b.score; }); + return out; +} diff --git a/image_analysis/structure_refinement/DensityBlobs.h b/image_analysis/structure_refinement/DensityBlobs.h new file mode 100644 index 000000000..657c22f51 --- /dev/null +++ b/image_analysis/structure_refinement/DensityBlobs.h @@ -0,0 +1,41 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#pragma once + +#include +#include + +#include "gemmi/grid.hpp" +#include "gemmi/model.hpp" + +// What a model leaves unexplained: the connected regions of positive mFo-DFc density above a cutoff ("blobs"), +// each placed beside the model, measured and given a likely reading - a water, an ion, something ligand-sized, +// or an error of the model itself (a misplaced or missing atom, an alternative conformation). +// +// The readings are rules on what is measured, meant to point a person at the density, not to identify a +// compound: a blob called LIGAND is density of ligand size and shape where the model has nothing, which a +// buffer molecule or an unmodelled side chain can be as well as the compound that was soaked. +struct DensityBlob { + enum class Kind { Water, Ion, Ligand, ModelError, Unassigned }; + Kind kind = Kind::Unassigned; + gemmi::Position centroid; // the copy nearest the model + double peak_sigma = 0; // the highest point, in the map's rms + double volume = 0; // A^3 above the cutoff + double score = 0; // the density above the cutoff integrated, sigma A^3 + double nearest_atom_distance = 0; + std::string nearest_atom; // e.g. "OG1 THR A42" + int polar_contacts = 0; // N or O of the model at 2.4-3.4 A + int metal_contacts = 0; // N, O or S of the model at 1.8-2.6 A + double anomalous_sigma = 0; // the anomalous difference map at the peak, 0 where there is none +}; + +const char *DensityBlobKindName(DensityBlob::Kind kind); + +// The blobs of `difference` (an mFo-DFc map of the whole cell, rms `rms`) above 3 sigma that peak above 4 sigma, +// strongest (by score) first. `model`: what the map was made from, in the map's cell. `anomalous`: the anomalous +// difference map on any grid of the same cell, or null. Following the flood fill of Coot's unmodelled blobs as +// gemmi implements it (find_blobs_by_flood_fill). +std::vector FindDensityBlobs(const gemmi::Grid &difference, double rms, const gemmi::Structure &model, + double d_min, const gemmi::Grid *anomalous = nullptr, + double anomalous_rms = 0); diff --git a/image_analysis/structure_refinement/Refinement.cpp b/image_analysis/structure_refinement/Refinement.cpp new file mode 100644 index 000000000..588bb140b --- /dev/null +++ b/image_analysis/structure_refinement/Refinement.cpp @@ -0,0 +1,781 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#include "Refinement.h" + +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include + +#include "gemmi/fourier.hpp" // get_size_for_hkl, get_f_phi_on_grid +#include "gemmi/modify.hpp" // remove_hydrogens +#include "gemmi/scaling.hpp" +#include "gemmi/stats.hpp" // calculate_data_statistics + +#include "../../common/JFJochException.h" +#include "../../common/JFJochMath.h" +#include "../../common/Logger.h" +#include "ModelFFT.h" +#include "ModelScaling.h" +#include "RefinementXray.h" +#include "SigmaA.h" +#ifdef JFJOCH_USE_CUDA +#include "../../common/CUDAWrapper.h" +#include "RefinementGPU.h" +#endif + +namespace { + +constexpr double B_FLOOR = 2.0; // A^2: no atom's B goes below it +// nor above twice the model's largest B, within these bounds (the GPU's gather is reserved for it) +constexpr double B_CEILING_MIN = 100.0, B_CEILING_MAX = 300.0; +// How far below the smallest B of the model a macro-cycle's B refinement may take an atom; the blur of the +// cycle is chosen for that. +constexpr double B_STEP_DOWN = 10.0; +constexpr double REPULSION_MARGIN = 0.5; // A: pairs this much beyond their contact distance are listed too +// The geometry the weight is chosen within (Afonine et al. 2012 select among weights by R-free with the +// geometry held to a target in the same way). +constexpr double MAX_RMSD_BONDS = 0.02, MAX_RMSD_ANGLES = 2.0; +// R-free differences smaller than this do not decide between two weights: the tighter geometry is taken. +constexpr double R_FREE_TIE = 0.001; + +std::string Joined(const std::vector &names) { + std::string out; + for (const std::string &n : names) + out += (out.empty() ? "" : " ") + n; + return out; +} + +std::vector AtomList(gemmi::Model &model) { + std::vector atoms; + for (gemmi::Chain &ch : model.chains) + for (gemmi::Residue &r : ch.residues) + for (gemmi::Atom &a : r.atoms) + atoms.push_back(&a); + 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. +gemmi::Structure PrepareModel(const gemmi::Structure &placed, const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg) { + gemmi::Structure st = placed; + st.models.resize(1); + st.cell = cell; + st.spacegroup_hm = sg.xhm(); + gemmi::remove_hydrogens(st.models[0]); + for (gemmi::Chain &ch : st.models[0].chains) + 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()) { + 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)); + } + } + gemmi::remove_empty_children(st); + st.setup_cell_images(); + return st; +} + +// What every stage works on: the model, the reflections, the engine and the row model of the cycle. +struct Context { + gemmi::Structure st; + std::vector atoms; + gemmi::UnitCell cell; + const gemmi::SpaceGroup *sg = nullptr; + double d_min = 0; + RefinementReflections refl; + std::unique_ptr engine; +#ifdef JFJOCH_USE_CUDA + std::unique_ptr scale_gpu; +#endif + size_t nthreads = 1; + XrayRowModel rows; + std::vector> fcalc, fmask; + double k_sol = 0, b_sol = 0; + double b_floor = B_FLOOR; // the cycle's lowest allowed B + double b_ceiling = B_CEILING_MAX; + // Riding hydrogens: per atom the parent it rides on (-1: none), and where it sits from the parent. + std::vector rider_parent; + std::vector rider_offset; + int evaluations = 0; + + gemmi::Model &Model() { return st.models[0]; } + + // The hydrogens moved with their parents, at the offsets they were last placed at, and given their B. + void MoveRiders() { + for (size_t i = 0; i < rider_parent.size(); i++) + if (rider_parent[i] >= 0) { + atoms[i]->pos = atoms[rider_parent[i]]->pos + rider_offset[i]; + atoms[i]->b_iso = atoms[rider_parent[i]]->b_iso; + } + } + // What the target says of a hydrogen is said of its parent. + void FoldRiders(std::vector> &g) const { + for (size_t i = 0; i < rider_parent.size(); i++) + if (rider_parent[i] >= 0) + for (int c = 0; c < 4; c++) + g[rider_parent[i]][c] += g[i][c]; + } + + std::vector> XyzB() const { + std::vector> out(atoms.size()); + for (size_t i = 0; i < atoms.size(); i++) + out[i] = {atoms[i]->pos.x, atoms[i]->pos.y, atoms[i]->pos.z, atoms[i]->b_iso}; + return out; + } + std::vector Positions() const { + std::vector out; + for (const gemmi::Atom *a : atoms) + out.push_back(a->pos); + return out; + } + void SetPositions(const std::vector &p) { + for (size_t i = 0; i < atoms.size(); i++) + atoms[i]->pos = p[i]; + } + std::vector B() const { + std::vector out; + for (const gemmi::Atom *a : atoms) + out.push_back(a->b_iso); + return out; + } + void SetB(const std::vector &b) { + for (size_t i = 0; i < atoms.size(); i++) + atoms[i]->b_iso = static_cast(b[i]); + } + + // A new macro-cycle: the blur for the B the cycle may reach, the bulk-solvent mask of the model as it is, + // the scale and the solvent fitted to the working set, and the likelihood's parameters from the free set. + // Returns the R-factors of the model as it stands. + RFactors Refit() { + double b_min = b_ceiling; + for (const gemmi::Atom *a : atoms) + b_min = std::min(b_min, a->b_iso); + b_floor = std::max(B_FLOOR, b_min - B_STEP_DOWN); + engine->SetBlur(RefinementBlur(cell, *sg, d_min, b_floor)); + engine->Fcalc(Model(), fcalc); + engine->Fmask(Model(), fmask); + + gemmi::AsuData> fc, fm; + gemmi::AsuData> fo; + fc.unit_cell_ = fm.unit_cell_ = fo.unit_cell_ = cell; + fc.spacegroup_ = fm.spacegroup_ = fo.spacegroup_ = sg; + for (size_t m = 0; m < refl.Size(); m++) { + fc.v.push_back({refl.hkl[m], fcalc[m]}); + fm.v.push_back({refl.hkl[m], fmask[m]}); + if (!refl.free[m]) + fo.v.push_back({refl.hkl[m], {refl.fo[m], 1.0f}}); + } + gemmi::Scaling scaling(cell, sg); + scaling.use_solvent = true; + scaling.prepare_points(fc, fo, &fm); +#ifdef JFJOCH_USE_CUDA + if (scale_gpu) + scale_gpu->Fit(scaling, SolventBMax(Model()), nthreads); + else +#endif + FitModelScale(scaling, {.b_hi = SolventBMax(Model())}, nthreads); + k_sol = scaling.k_sol; + b_sol = scaling.b_sol; + rows.k_total.resize(refl.Size()); + rows.fmask.resize(refl.Size()); + std::vector> fmodel(refl.Size()); + for (size_t m = 0; m < refl.Size(); m++) { + rows.k_total[m] = static_cast(scaling.get_overall_scale_factor(refl.hkl[m])); + rows.fmask[m] = static_cast(scaling.get_solvent_scale(cell.calculate_stol_sq(refl.hkl[m]))) * fmask[m]; + fmodel[m] = rows.k_total[m] * (fcalc[m] + rows.fmask[m]); + } + EstimateLikelihoodParameters(refl, fmodel, rows); + engine->SetRowModel(rows); + return ComputeRFactors(refl, fmodel); + } + + // F_model with the cycle's scale and solvent, for the atoms as they are now. + std::vector> FModel() { + engine->Fcalc(Model(), fcalc); + std::vector> fmodel(refl.Size()); + for (size_t m = 0; m < refl.Size(); m++) + fmodel[m] = rows.k_total[m] * (fcalc[m] + rows.fmask[m]); + return fmodel; + } +}; + +ceres::GradientProblemSolver::Options SolverOptions(int iterations) { + ceres::GradientProblemSolver::Options options; + options.line_search_direction_type = ceres::LBFGS; + options.max_num_iterations = iterations; + options.function_tolerance = 1e-7; + options.gradient_tolerance = 1e-10; + options.parameter_tolerance = 1e-9; + options.logging_type = ceres::SILENT; + return options; +} + +// The positions or the B factors of the refinable atoms, against w times the X-ray target plus a restraint set. +// Following Liu & Nocedal (1989) Math. Program. 45, 503-528 for the minimiser (Ceres' L-BFGS) +class AtomTarget : public ceres::FirstOrderFunction { +public: + enum class Mode { Xyz, B }; + AtomTarget(Context &ctx, Mode mode, const std::vector ¶m_atoms, double weight, int restraint_set) + : ctx_(ctx), mode_(mode), param_atoms_(param_atoms), weight_(weight), restraint_set_(restraint_set) {} + + int NumParameters() const override { + return static_cast(param_atoms_.size()) * (mode_ == Mode::Xyz ? 3 : 1); + } + + bool Evaluate(const double *p, double *cost, double *gradient) const override { + for (size_t k = 0; k < param_atoms_.size(); k++) { + gemmi::Atom &a = *ctx_.atoms[param_atoms_[k]]; + 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)); + } + ctx_.MoveRiders(); + ctx_.evaluations++; + std::vector> gx, gr; + const double xray = ctx_.engine->Evaluate(ctx_.Model(), gradient != nullptr ? &gx : nullptr); + if (gradient != nullptr) + ctx_.FoldRiders(gx); + double restraints = 0; + if (restraint_set_ >= 0) + restraints = ctx_.engine->EvaluateRestraints(restraint_set_, ctx_.XyzB(), gradient != nullptr ? &gr : nullptr); + *cost = weight_ * xray + restraints; + if (gradient != nullptr) + for (size_t k = 0; k < param_atoms_.size(); k++) { + const int i = param_atoms_[k]; + if (mode_ == Mode::Xyz) + for (int c = 0; c < 3; c++) + gradient[3 * k + c] = weight_ * gx[i][c] + (restraint_set_ >= 0 ? gr[i][c] : 0.0); + else { + // held at a bound, B does not move further out + const bool held = (p[k] <= ctx_.b_floor && gx[i][3] > 0) || (p[k] >= ctx_.b_ceiling && gx[i][3] < 0); + gradient[k] = held ? 0.0 : weight_ * gx[i][3] + (restraint_set_ >= 0 ? gr[i][3] : 0.0); + } + } + return std::isfinite(*cost); + } + +private: + Context &ctx_; + Mode mode_; + const std::vector ¶m_atoms_; + double weight_; + int restraint_set_; +}; + +void Minimise(Context &ctx, AtomTarget::Mode mode, const std::vector ¶m_atoms, double weight, + int restraint_set, int iterations) { + if (param_atoms.empty()) + return; + std::vector p; + for (int i : param_atoms) { + const gemmi::Atom &a = *ctx.atoms[i]; + if (mode == AtomTarget::Mode::Xyz) { + p.push_back(a.pos.x); + p.push_back(a.pos.y); + p.push_back(a.pos.z); + } else + p.push_back(a.b_iso); + } + ceres::GradientProblem problem(std::make_unique(ctx, mode, param_atoms, weight, restraint_set)); + ceres::GradientProblemSolver::Summary summary; + ceres::Solve(SolverOptions(iterations), problem, p.data(), &summary); + // the atoms where the solver ended, not where its last line search probed + for (size_t k = 0; k < param_atoms.size(); k++) { + gemmi::Atom &a = *ctx.atoms[param_atoms[k]]; + if (mode == AtomTarget::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)); + } + ctx.MoveRiders(); +} + +// Groups of atoms that move as rigid bodies (one per chain) against the X-ray target: per body an angle-axis +// rotation about its centroid, scaled by its rms radius so that a unit moves a typical atom by an angstrom as a +// unit of translation does, then a translation. +class RigidBodiesTarget : public ceres::FirstOrderFunction { +public: + RigidBodiesTarget(Context &ctx, const std::vector> &bodies) + : ctx_(ctx), bodies_(bodies), base_(ctx.Positions()) { + for (const std::vector &b : bodies) { + gemmi::Position c; + for (int i : b) + c += base_[i]; + c *= 1.0 / static_cast(b.size()); + double r2 = 0; + for (int i : b) + r2 += c.dist_sq(base_[i]); + centre_.push_back(c); + radius_.push_back(std::max(1.0, std::sqrt(r2 / static_cast(b.size())))); + } + } + int NumParameters() const override { return static_cast(6 * bodies_.size()); } + + bool Evaluate(const double *q, double *cost, double *gradient) const override { + Place(q); + ctx_.evaluations++; + std::vector> g; + *cost = ctx_.engine->Evaluate(ctx_.Model(), gradient != nullptr ? &g : nullptr); + if (gradient == nullptr) + return std::isfinite(*cost); + using Jet = ceres::Jet; + for (size_t b = 0; b < bodies_.size(); b++) { + const Jet aa[3] = {Jet(q[6 * b] / radius_[b], 0), Jet(q[6 * b + 1] / radius_[b], 1), + Jet(q[6 * b + 2] / radius_[b], 2)}; + double dq[6] = {0, 0, 0, 0, 0, 0}; + for (int i : bodies_[b]) { + const gemmi::Position d = base_[i] - centre_[b]; + const Jet p[3] = {Jet(d.x), Jet(d.y), Jet(d.z)}; + Jet rp[3]; + ceres::AngleAxisRotatePoint(aa, p, rp); + for (int c = 0; c < 3; c++) { + for (int k = 0; k < 3; k++) + dq[k] += g[i][c] * rp[c].v[k] / radius_[b]; + dq[3 + c] += g[i][c]; + } + } + std::copy(dq, dq + 6, gradient + 6 * b); + } + return std::isfinite(*cost); + } + + void Place(const double *q) const { + for (size_t b = 0; b < bodies_.size(); b++) { + const double aa[3] = {q[6 * b] / radius_[b], q[6 * b + 1] / radius_[b], q[6 * b + 2] / radius_[b]}; + for (int i : bodies_[b]) { + const gemmi::Position d = base_[i] - centre_[b]; + const double p[3] = {d.x, d.y, d.z}; + double rp[3]; + ceres::AngleAxisRotatePoint(aa, p, rp); + ctx_.atoms[i]->pos = gemmi::Position(rp[0] + centre_[b].x + q[6 * b + 3], rp[1] + centre_[b].y + q[6 * b + 4], + rp[2] + centre_[b].z + q[6 * b + 5]); + } + } + } + +private: + Context &ctx_; + const std::vector> &bodies_; + std::vector base_; + std::vector centre_; + std::vector radius_; +}; + +// One B shift per group of atoms (a residue), against the X-ray target. +class GroupBTarget : public ceres::FirstOrderFunction { +public: + GroupBTarget(Context &ctx, const std::vector> &groups) + : ctx_(ctx), groups_(groups), base_(ctx.B()) {} + int NumParameters() const override { return static_cast(groups_.size()); } + + bool Evaluate(const double *q, double *cost, double *gradient) const override { + Place(q); + ctx_.evaluations++; + std::vector> g; + *cost = ctx_.engine->Evaluate(ctx_.Model(), gradient != nullptr ? &g : nullptr); + if (gradient != nullptr) + for (size_t k = 0; k < groups_.size(); k++) { + gradient[k] = 0; + for (int i : groups_[k]) { + const double b = base_[i] + q[k]; + if (b > ctx_.b_floor && b < ctx_.b_ceiling) + gradient[k] += g[i][3]; + } + } + return std::isfinite(*cost); + } + + 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)); + } + +private: + Context &ctx_; + const std::vector> &groups_; + std::vector base_; +}; + +// The placement stage: rigid bodies, then group B. +void RefinePlacement(Context &ctx, RefinementResult &result) { + std::vector> bodies, groups; + size_t i = 0; + for (const gemmi::Chain &ch : ctx.Model().chains) { + std::vector body; + for (const gemmi::Residue &r : ch.residues) { + std::vector group; + for (size_t k = 0; k < r.atoms.size(); k++, i++) { + group.push_back(static_cast(i)); + if (!r.is_water()) + body.push_back(static_cast(i)); + } + groups.push_back(group); + } + // A chain of a few atoms (an ion, a ligand of its own chain) stays where it is. + if (body.size() >= 50) + bodies.push_back(body); + } + result.rigid_bodies = static_cast(bodies.size()); + if (!bodies.empty()) { + std::vector q(6 * bodies.size(), 0.0); + auto target = std::make_unique(ctx, bodies); + RigidBodiesTarget *rb = target.get(); + ceres::GradientProblem problem(std::move(target)); + ceres::GradientProblemSolver::Summary summary; + ceres::Solve(SolverOptions(30), problem, q.data(), &summary); + rb->Place(q.data()); + } + ctx.Refit(); + { + std::vector q(groups.size(), 0.0); + auto target = std::make_unique(ctx, groups); + GroupBTarget *gb = target.get(); + ceres::GradientProblem problem(std::move(target)); + ceres::GradientProblemSolver::Summary summary; + ceres::Solve(SolverOptions(30), problem, q.data(), &summary); + gb->Place(q.data()); + } +} + +// The X-ray weight that balances the two gradients: the norm of the geometry's gradient (taken on the model +// shaken by 0.1 A, where a refined model's geometry would have none) over the norm of the X-ray target's. +// Following Brunger (1992)-style gradient balancing as phenix.refine starts its weight search (Afonine et al. +// (2012) Acta Cryst. D68, 352-367) +double BalancedWeight(Context &ctx, const std::vector &refinable) { + std::vector> gx, gr; + ctx.engine->Evaluate(ctx.Model(), &gx); + std::vector> xyzb = ctx.XyzB(); + std::mt19937 rng(1); + std::normal_distribution shake(0.0, 0.1); + for (std::array &a : xyzb) + for (int c = 0; c < 3; c++) + a[c] += shake(rng); + ctx.engine->EvaluateRestraints(0, xyzb, &gr); + double nx = 0, ng = 0; + for (int i : refinable) + for (int c = 0; c < 3; c++) { + nx += gx[i][c] * gx[i][c]; + ng += gr[i][c] * gr[i][c]; + } + return nx > 0 ? std::sqrt(ng / nx) : 1.0; +} + +// The model's map coefficients and its difference map, as model validation makes them (sigma_A weighted). +void MakeMaps(Context &ctx, const std::vector &merged, double wavelength_A, RefinementResult &result) { + ctx.Refit(); + const std::vector> fmodel = ctx.FModel(); + const gemmi::GroupOps gops = ctx.sg->operations(); + std::vector sa; + for (size_t m = 0; m < ctx.refl.Size(); m++) + sa.push_back({ctx.refl.fo[m], std::abs(fmodel[m]), ctx.refl.inv_d2[m], static_cast(ctx.refl.epsilon[m]), + static_cast(ctx.refl.centric[m]), static_cast(ctx.refl.free[m])}); + const SigmaAResult sigma_a = EstimateSigmaA(sa, ctx.cell); + result.mean_fom = sigma_a.mean_fom; + + auto maps = std::make_shared(); + MtzTable &mtz = maps->coefficients; + mtz.space_group_name = ctx.sg->xhm(); + mtz.unit_cell = UnitCell{static_cast(ctx.cell.a), static_cast(ctx.cell.b), static_cast(ctx.cell.c), + static_cast(ctx.cell.alpha), static_cast(ctx.cell.beta), + static_cast(ctx.cell.gamma)}; + mtz.datasets = {{"HKL_base", "HKL_base", "HKL_base", 0.0f}, + {"refinement", "refinement", "refinement", static_cast(wavelength_A)}}; + const std::pair columns[] = {{"H", 'H'}, {"K", 'H'}, {"L", 'H'}, {"FP", 'F'}, {"FC", 'F'}, + {"PHIC", 'P'}, {"FWT", 'F'}, {"PHWT", 'P'}, {"DELFWT", 'F'}, + {"PHDELWT", 'P'}, {"FOM", 'W'}, {"FREE", 'I'}}; + for (const auto &[label, type] : columns) + mtz.columns.push_back({label, type, type == 'H' ? 0 : 1, {}}); + gemmi::AsuData> diff; + diff.unit_cell_ = ctx.cell; + diff.spacegroup_ = ctx.sg; + for (size_t m = 0; m < ctx.refl.Size(); m++) { + const double fo = ctx.refl.fo[m], fc = std::abs(fmodel[m]); + const double mw = sigma_a.weight[m].m, d = sigma_a.weight[m].d; + const double fwt = ctx.refl.centric[m] ? mw * fo : 2 * mw * fo - d * fc; + const double delfwt = mw * fo - d * fc; + const double phi = std::arg(fmodel[m]); + double phi_deg = phi * 180.0 / PI; + if (phi_deg < 0) + phi_deg += 360.0; + const float row[] = {static_cast(ctx.refl.hkl[m][0]), static_cast(ctx.refl.hkl[m][1]), + static_cast(ctx.refl.hkl[m][2]), static_cast(fo), static_cast(fc), + static_cast(phi_deg), static_cast(fwt), static_cast(phi_deg), + static_cast(delfwt), static_cast(phi_deg), static_cast(mw), + ctx.refl.free[m] ? 0.0f : 1.0f}; + for (size_t c = 0; c < mtz.columns.size(); c++) + mtz.columns[c].values.push_back(row[c]); + diff.v.push_back({ctx.refl.hkl[m], std::polar(static_cast(delfwt), static_cast(phi))}); + } + result.maps = maps; + const std::array size = gemmi::get_size_for_hkl(diff, {{0, 0, 0}}, 3.0); + result.difference_map = std::make_shared>(MapFromFPhi(gemmi::get_f_phi_on_grid(diff, size, true))); + result.difference_map_rms = gemmi::calculate_data_statistics(result.difference_map->data).rms; + + // The anomalous difference map: F(+) - F(-) on the model's phase turned back by 90 degrees, read on the + + // rows of an anomalous merge only (as model validation reads it). + // Following ANODE, Thorn & Sheldrick (2011) J. Appl. Cryst. 44, 1285-1287 + const gemmi::ReciprocalAsu asu(ctx.sg); + std::map danom; + for (const MergedReflection &r : merged) { + if (!std::isfinite(r.F_plus) || !std::isfinite(r.F_minus)) + continue; + const gemmi::Miller h{{r.h, r.k, r.l}}; + if (gops.is_reflection_centric(h)) + continue; + const auto [hasu, plus] = asu.to_asu_sign(h, gops); + if (plus) + danom[hasu] = r.F_plus - r.F_minus; + } + if (!danom.empty()) { + gemmi::AsuData> anom; + anom.unit_cell_ = ctx.cell; + anom.spacegroup_ = ctx.sg; + for (size_t m = 0; m < ctx.refl.Size(); m++) { + const auto it = danom.find(ctx.refl.hkl[m]); + if (it != danom.end()) + anom.v.push_back({ctx.refl.hkl[m], it->second * std::polar(1.0f, static_cast(std::arg(fmodel[m]) - PI / 2))}); + } + if (!anom.v.empty()) { + result.anomalous_map = std::make_shared>( + MapFromFPhi(gemmi::get_f_phi_on_grid(anom, gemmi::get_size_for_hkl(anom, {{0, 0, 0}}, 3.0), true))); + result.anomalous_map_rms = gemmi::calculate_data_statistics(result.anomalous_map->data).rms; + } + } +} + +// The engine for the model as it is now: on the GPU where it is wanted and takes the case within its budget, +// on the CPU otherwise. +void MakeEngine(Context &ctx, bool gpu, const RefinementSettings &settings, Logger &logger) { + ctx.engine.reset(); +#ifdef JFJOCH_USE_CUDA + ctx.scale_gpu.reset(); + if (gpu && get_gpu_count() > 0) { + try { + const int device = get_gpu(); + const size_t card = get_gpu_total_memory(device); + const size_t usable = card / 5 * 4; + const size_t budget = usable > settings.gpu_resident_bytes ? usable - settings.gpu_resident_bytes : 0; + size_t bytes = 0; + ctx.engine = MakeRefinementEngineGPU(device, ctx.cell, *ctx.sg, ctx.d_min, ctx.refl, ctx.Model(), + ctx.b_ceiling, budget, &bytes); + if (ctx.engine) + ctx.scale_gpu = std::make_unique(device, ctx.refl.Size(), 1); + else + logger.Info("Refinement: on the CPU - the GPU cannot take this case or {:.1f} GB is over its budget", + bytes / 1e9); + } catch (const JFJochException &e) { + logger.Warning("Refinement: the GPU failed ({}), on the CPU", e.what()); + ctx.engine.reset(); + ctx.scale_gpu.reset(); + } + } +#else + (void) gpu; + (void) logger; +#endif + if (!ctx.engine) + ctx.engine = MakeRefinementEngineCPU(ctx.cell, *ctx.sg, ctx.d_min, ctx.refl, settings.nthreads); +} + +} // namespace + +RefinementResult RefineModel(const gemmi::Structure &placed, const std::vector &merged, + const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg, double wavelength_A, + Logger &logger, const RefinementSettings &settings) { + const auto start = std::chrono::steady_clock::now(); + RefinementResult result; + Context ctx; + ctx.st = PrepareModel(placed, cell, sg); + ctx.cell = cell; + ctx.sg = &sg; + ctx.nthreads = settings.nthreads; + ctx.atoms = AtomList(ctx.Model()); + result.atoms = static_cast(ctx.atoms.size()); + double b_max = 0; + for (const gemmi::Atom *a : ctx.atoms) + b_max = std::max(b_max, a->b_iso); + 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"; + return result; + } + double max_inv_d2 = 0; + for (const MergedReflection &r : merged) + if (std::isfinite(r.F) && r.d > 0) + max_inv_d2 = std::max(max_inv_d2, 1.0 / (double(r.d) * r.d)); + if (max_inv_d2 <= 0) { + result.failure_reason = "no amplitudes"; + return result; + } + // just past the finest reflection, so that prepare_asu_data()'s strict limit keeps it + ctx.d_min = 0.999 / std::sqrt(max_inv_d2); + result.d_min = 1.0 / std::sqrt(max_inv_d2); + ctx.refl = PrepareRefinementReflections(merged, cell, sg, ctx.d_min); + if (ctx.refl.Size() < 100) { + result.failure_reason = "too few reflections"; + return result; + } + + MakeEngine(ctx, settings.gpu, settings, logger); + result.on_gpu = ctx.engine->OnGPU(); + + RFactors r = ctx.Refit(); + result.n_work = r.n_work; + result.n_free = r.n_free; + result.r_work_start = r.r_work; + result.r_free_start = r.r_free; + logger.Info("Refinement: {} atoms, {} reflections to {:.2f} A ({} free), on the {}; as given R-work {:.4f}, R-free {:.4f}", + ctx.atoms.size(), ctx.refl.Size(), result.d_min, r.n_free, result.on_gpu ? "GPU" : "CPU", r.r_work, + r.r_free); + + RefinePlacement(ctx, result); + r = ctx.Refit(); + result.r_work_placement = r.r_work; + result.r_free_placement = r.r_free; + logger.Info("Refinement, placement ({} rigid bodies, B per residue): R-work {:.4f}, R-free {:.4f}", + result.rigid_bodies, r.r_work, r.r_free); + + if (settings.restrained) { + RefinementTopology topo = BuildRefinementTopology(ctx.st, settings.riding_hydrogens); + // the hydrogens it added: the atoms and the engine made again for them + ctx.atoms = AtomList(ctx.Model()); + ctx.rider_parent = topo.rider_parent; + auto place_riders = [&]() { + PlaceRidingHydrogens(topo); + ctx.rider_offset.assign(ctx.atoms.size(), gemmi::Position()); + for (size_t i = 0; i < ctx.atoms.size(); i++) + if (ctx.rider_parent[i] >= 0) { + ctx.rider_offset[i] = ctx.atoms[i]->pos - ctx.atoms[ctx.rider_parent[i]]->pos; + ctx.atoms[i]->b_iso = ctx.atoms[ctx.rider_parent[i]]->b_iso; + } + }; + place_riders(); + if (static_cast(ctx.atoms.size()) != result.atoms) { + MakeEngine(ctx, result.on_gpu, settings, logger); + result.on_gpu = ctx.engine->OnGPU(); + } + ctx.Refit(); + result.restrained_residues = topo.restrained_residues; + result.fixed_residues = topo.fixed_residues; + result.fixed_names = topo.fixed_names; + std::vector refinable; + for (size_t i = 0; i < ctx.atoms.size(); i++) + if (topo.refinable[i]) + refinable.push_back(static_cast(i)); + result.refined_atoms = static_cast(refinable.size()); + logger.Info("Refinement: {} residues restrained, {} held fixed{}; {} bonds, {} angles, {} torsions, {} chiral " + "volumes, {} planes", + topo.restrained_residues, topo.fixed_residues, + topo.fixed_names.empty() ? std::string() : " (" + Joined(topo.fixed_names) + ")", + topo.bonds, topo.angles, topo.torsions, topo.chirals, topo.planes); + + auto restraints_now = [&]() { + RestraintSet set = topo.geometry; + set.Append(BuildRepulsions(topo, ctx.Model(), cell, sg, REPULSION_MARGIN)); + set.Finalize(ctx.atoms.size()); + return set; + }; + ctx.engine->SetRestraints(1, topo.adp); + { + const RestraintSet set = restraints_now(); + ctx.engine->SetRestraints(0, set); + RestraintSet repulsions = BuildRepulsions(topo, ctx.Model(), cell, sg, 0.0); + result.geometry_start = MeasureGeometry(topo.geometry, repulsions, ctx.XyzB()); + } + const double w0 = BalancedWeight(ctx, refinable); + logger.Info("Refinement: geometry as given rmsd bonds {:.4f} A, angles {:.2f} deg, {} clashes; balanced X-ray " + "weight {:.3g}", result.geometry_start.rmsd_bonds, result.geometry_start.rmsd_angles, + result.geometry_start.clashes, w0); + + const std::vector pos0 = ctx.Positions(); + const std::vector b0 = ctx.B(); + int best = -1; + std::vector best_pos; + std::vector best_b; + // One full refinement from the start: macro-cycles of the positions at X-ray weight w and, where adp_w is + // not 0, the B factors at X-ray weight adp_w; kept where it beats the best so far. The rungs come in order + // of trust - tighter geometry, then fewer parameters - and a later one displaces an earlier only by more + // than the tie in R-free. + auto trial = [&](double w, double adp_w) { + ctx.SetPositions(pos0); + ctx.SetB(b0); + for (int cycle = 0; cycle < settings.macro_cycles; cycle++) { + place_riders(); + ctx.Refit(); + ctx.engine->SetRestraints(0, restraints_now()); + Minimise(ctx, AtomTarget::Mode::Xyz, refinable, w, 0, settings.xyz_iterations); + if (adp_w > 0) + Minimise(ctx, AtomTarget::Mode::B, refinable, adp_w, 1, settings.b_iterations); + } + place_riders(); + RefinementWeightTrial t; + t.weight = w; + t.adp_weight = adp_w; + const RFactors rf = ctx.Refit(); + t.r_work = rf.r_work; + t.r_free = rf.r_free; + t.geometry = MeasureGeometry(topo.geometry, BuildRepulsions(topo, ctx.Model(), cell, sg, 0.0), ctx.XyzB()); + t.within_bounds = t.geometry.rmsd_bonds <= MAX_RMSD_BONDS && t.geometry.rmsd_angles <= MAX_RMSD_ANGLES; + logger.Info("Refinement, X-ray weight {:.3g}, B {}: 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, + t.within_bounds ? "" : " (geometry out of bounds)"); + result.trials.push_back(t); + const int k = static_cast(result.trials.size()) - 1; + const bool better = best < 0 ? true + : t.within_bounds && (!result.trials[best].within_bounds || + t.r_free < result.trials[best].r_free - R_FREE_TIE); + if (better) { + best = k; + best_pos = ctx.Positions(); + best_b = ctx.B(); + } + }; + // the positions' weight, with the B factors as they are + for (double ladder : settings.weight_ladder) + trial(w0 * ladder, 0.0); + // then the B factors refined too, at that weight, against their own weight + const double w = result.trials[best].weight; + for (double adp : settings.adp_ladder) + trial(w, w * adp); + ctx.SetPositions(best_pos); + ctx.SetB(best_b); + place_riders(); + const RefinementWeightTrial &chosen = result.trials[best]; + result.restrained = true; + result.weight = chosen.weight; + result.adp_weight = chosen.adp_weight; + result.r_work = chosen.r_work; + result.r_free = chosen.r_free; + result.geometry = chosen.geometry; + } + + MakeMaps(ctx, merged, wavelength_A, result); + result.blobs = FindDensityBlobs(*result.difference_map, result.difference_map_rms, ctx.st, result.d_min, + result.anomalous_map.get(), result.anomalous_map_rms); + result.k_sol = ctx.k_sol; + result.b_sol = ctx.b_sol; + result.model = std::make_shared(ctx.st); + result.ok = true; + result.seconds = std::chrono::duration(std::chrono::steady_clock::now() - start).count(); + logger.Info("Refinement: done in {:.1f} s, {} target evaluations{}", result.seconds, ctx.evaluations, + result.restrained ? fmt::format("; R-work {:.4f}, R-free {:.4f}, weight {:.3g}", result.r_work, + result.r_free, result.weight) + : std::string()); + return result; +} diff --git a/image_analysis/structure_refinement/Refinement.h b/image_analysis/structure_refinement/Refinement.h new file mode 100644 index 000000000..a1d1076b6 --- /dev/null +++ b/image_analysis/structure_refinement/Refinement.h @@ -0,0 +1,103 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#pragma once + +#include +#include +#include + +#include "gemmi/grid.hpp" +#include "gemmi/model.hpp" +#include "gemmi/symmetry.hpp" +#include "gemmi/unitcell.hpp" + +#include "../../common/Reflection.h" // MergedReflection +#include "DensityBlobs.h" +#include "ModelValidation.h" // ModelMaps +#include "RefinementRestraints.h" // GeometryStatistics + +class Logger; + +// Refinement of an atomic model already placed against the data (ValidateAgainstModel's placed model, or a +// molecular-replacement placement), in two stages: +// +// 1. Placement: each chain as a rigid body, then one B shift per residue, then the bulk solvent and the +// scale. Needs no dictionary. +// 2. Restrained: every atom's position and isotropic B, against the maximum-likelihood amplitude target with +// the geometry the dictionary describes (MonomerLibrary.h: bonds, angles, planes, chiral volumes, +// torsions, repulsions; B similarity of bonded atoms). Residues it does not describe are held fixed and +// stay in the structure factors and the bulk-solvent mask. The weight of the X-ray term against the +// geometry is chosen from a ladder of full refinements, by R-free among those whose geometry stays within +// bounds; then whether, and at what weight, to refine the B factors as well, the same way. +// +// The library reads and writes nothing: the model and the reflections come in, the refined model, the +// statistics and the map coefficients go out. Nothing it finds feeds back into processing decisions. +struct RefinementSettings { + bool restrained = true; // stage 2; false stops after the placement stage + int macro_cycles = 5; + int xyz_iterations = 25; // L-BFGS iterations per macro-cycle + int b_iterations = 15; + // The positions' X-ray weight, times the gradient-balanced one (BalancedWeight), tried in this order. + std::vector weight_ladder = {0.0078125, 0.015625, 0.03125, 0.0625}; + // The X-ray weights of the B refinement tried next, times the positions' chosen weight. + std::vector adp_ladder = {1.0, 4.0, 16.0}; + bool riding_hydrogens = true; // the dictionary's hydrogens added, riding on their parents + bool gpu = true; + size_t nthreads = 1; + size_t gpu_resident_bytes = 0; // held on the card by what runs beside +}; + +// One rung of the weight ladder. +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; + GeometryStatistics geometry; + bool within_bounds = false; // the geometry bounds the choice is made inside +}; + +struct RefinementResult { + bool ok = false; + std::string failure_reason; + bool on_gpu = false; + double seconds = 0; + double d_min = 0; + int n_work = 0, n_free = 0; + + // The model as given, scaled to these reflections (hydrogens removed, B made isotropic). + double r_work_start = 1, r_free_start = 1; + // After the placement stage. + int rigid_bodies = 0; + double r_work_placement = 1, r_free_placement = 1; + // After the restrained stage (NaN where it did not run). + bool restrained = false; + double r_work = NAN, r_free = NAN; + double weight = 0; + double adp_weight = 0; // 0: the B factors kept as placed + std::vector trials; + GeometryStatistics geometry_start, geometry; + + int atoms = 0, refined_atoms = 0; + int restrained_residues = 0, fixed_residues = 0; + std::vector fixed_names; + double k_sol = 0, b_sol = 0; + double mean_fom = 0; + + std::shared_ptr model; // the refined model, in the data's cell and frame + std::shared_ptr maps; // its 2mFo-DFc and mFo-DFc coefficients (as _maps.mtz holds them) + // The mFo-DFc map of the refined model and its rms, for finding what it leaves unexplained (DensityBlobs.h). + std::shared_ptr> difference_map; + double difference_map_rms = 0; + // The anomalous difference map on the refined model's phases, where the merge kept the Bijvoet pairs. + std::shared_ptr> anomalous_map; + double anomalous_map_rms = 0; + // What the refined model leaves unexplained in the difference map (FindDensityBlobs). + std::vector blobs; +}; + +// `placed`: the model in the data's cell and space group. `merged`: the merge, with French-Wilson amplitudes +// and free flags, in the same group. +RefinementResult RefineModel(const gemmi::Structure &placed, const std::vector &merged, + const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg, double wavelength_A, + Logger &logger, const RefinementSettings &settings); diff --git a/image_analysis/structure_refinement/RefinementGPU.cpp b/image_analysis/structure_refinement/RefinementGPU.cpp new file mode 100644 index 000000000..97e31b7b8 --- /dev/null +++ b/image_analysis/structure_refinement/RefinementGPU.cpp @@ -0,0 +1,283 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#include "RefinementGPU.h" + +#include +#include + +#include "gemmi/dencalc.hpp" +#include "gemmi/it92.hpp" +#include "gemmi/solmask.hpp" // SolventMasker, refmac_radius_for_bulk_solvent + +#include "RefinementGPUEngine.h" + +namespace { + +using Table = gemmi::IT92; + +// 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; + const auto &coef = Table::get(atom.element, atom.charge, atom.serial); + const CReal b = static_cast(b_iso + dc.blur); + const auto precal = coef.precalculate_density_iso(b, 0.0f); + const double radius = dc.estimate_radius(precal, b); + ModelDensityAtom a{}; + a.occ = atom.occ; + a.aniso = 0; + 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]; + a.b[k][0] = precal.b[k]; + } + return a; +} + +std::vector Atoms(const gemmi::Model &model) { + std::vector atoms; + for (const gemmi::Chain &ch : model.chains) + for (const gemmi::Residue &r : ch.residues) + for (const gemmi::Atom &a : r.atoms) + atoms.push_back(&a); + return atoms; +} + +class RefinementEngineGPU : public RefinementEngine { +public: + RefinementEngineGPU(int device, const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg, double d_min, + const RefinementReflections &refl, const gemmi::Model &model, double b_max) + : device_(device), cell_(cell), sg_(sg), d_min_(d_min) { + const gemmi::Grid grid = RigidBodyZoneGrid(cell, sg, d_min); + setup_.grid.nu = grid.nu; + setup_.grid.nv = grid.nv; + setup_.grid.nw = grid.nw; + for (int i = 0; i < 3; i++) + for (int j = 0; j < 3; j++) { + setup_.grid.orth[3 * i + j] = cell.orth.mat[i][j]; + setup_.grid.frac[3 * i + j] = cell.frac.mat[i][j]; + } + setup_.volume = cell.volume; + const gemmi::GroupOps gops = sg.operations(); + setup_.den = gemmi::Op::DEN; + setup_.centring = static_cast(gops.cen_ops.size()); + for (const gemmi::Op &op : gops.sym_ops) { + std::array o{}; + for (int i = 0; i < 3; i++) { + for (int j = 0; j < 3; j++) + o[3 * i + j] = op.rot[i][j]; + o[9 + i] = op.tran[i]; + } + setup_.sym_ops.push_back(o); + } + for (const gemmi::Op::Tran &cen : gops.cen_ops) + for (const gemmi::Op &op : gops.sym_ops) { + ModelMaskOp m{}; + for (int i = 0; i < 3; i++) { + for (int j = 0; j < 3; j++) + m.rot[3 * i + j] = static_cast(op.rot[i][j]) / gemmi::Op::DEN; + m.tran[i] = static_cast(op.tran[i] + cen[i]) / gemmi::Op::DEN; + } + setup_.mask_ops.push_back(m); + } + for (size_t m = 0; m < refl.Size(); m++) { + setup_.rows.push_back({refl.hkl[m][0], refl.hkl[m][1], refl.hkl[m][2]}); + setup_.inv_d2.push_back(refl.inv_d2[m]); + setup_.fo.push_back(refl.fo[m]); + setup_.work.push_back(!refl.free[m]); + setup_.centric.push_back(refl.centric[m]); + } + + // The adjoint of the composition on cuFFT's half-u spectrum: the term of row m and operator o, at + // k = hR, goes to k where k_u > 0, conjugated to -k where k_u < 0, and to both on the planes k_u = 0 + // and k_u = nu/2, which hold their own Friedel mates. + const int nu = grid.nu, nv = grid.nv, nw = grid.nw; + auto point = [&](int ku, int kv, int kw) { + return ku + (nu / 2 + 1) * (gemmi::modulo(kv, nv) + nv * gemmi::modulo(kw, nw)); + }; + std::vector> entries; + const int n_ops = static_cast(gops.sym_ops.size()); + for (size_t m = 0; m < refl.Size(); m++) + for (int o = 0; o < n_ops; o++) { + const gemmi::Miller k = gops.sym_ops[o].apply_to_hkl(refl.hkl[m]); + const int term = (static_cast(m) * n_ops + o) * 2; + if (k[0] > 0 && 2 * k[0] != nu) + entries.emplace_back(point(k[0], k[1], k[2]), term); + else if (k[0] < 0 && -2 * k[0] != nu) + entries.emplace_back(point(-k[0], -k[1], -k[2]), term + 1); + else { + const int ku = std::abs(k[0]); + entries.emplace_back(point(ku, k[1], k[2]), term); + entries.emplace_back(point(ku, -k[1], -k[2]), term + 1); + } + } + std::stable_sort(entries.begin(), entries.end(), + [](const auto &a, const auto &b) { return a.first < b.first; }); + for (size_t e = 0; e < entries.size(); e++) { + if (e == 0 || entries[e].first != entries[e - 1].first) { + setup_.scatter_point.push_back(entries[e].first); + setup_.scatter_start.push_back(static_cast(e)); + } + setup_.scatter_term.push_back(entries[e].second); + } + setup_.scatter_start.push_back(static_cast(entries.size())); + + // The gather's reserve for every atom at b_max with the largest blur a refinement sets (B down to 0). + dc_ = Calculator(RefinementBlur(cell, sg, d_min, 0.0)); + std::vector widest; + for (const gemmi::Atom *a : Atoms(model)) + widest.push_back(DensityAtom(*a, b_max, dc_)); + setup_.max_atoms = widest.size(); + setup_.max_pairs = ModelDensityGPU::PairBound(setup_.grid, widest); + supported_ = !setup_.rows.empty() && ModelDensityGPU::Supports(setup_.grid, widest) && + setup_.mask_ops.size() <= ModelMaskGPU::MAX_OPS; + if (supported_) { + fft_work_ = RefinementGPUEngine::FFTWorkBytes(device, setup_.grid); + // the restraints: about thirty records and slots per atom + bytes_ = RefinementGPUEngine::DeviceBytes(setup_, fft_work_, 30 * widest.size(), 60 * widest.size()); + } + } + + bool Supported() const { return supported_; } + size_t Bytes() const { return bytes_; } + void Reserve() { engine_ = std::make_unique(device_, setup_, fft_work_); } + + void SetBlur(double blur) override { + dc_ = Calculator(blur); + last_b_.clear(); + engine_->SetBlur(blur); + } + + void SetRowModel(const XrayRowModel &rows) override { + std::vector> fmask(rows.fmask.size()); + for (size_t m = 0; m < fmask.size(); m++) + fmask[m] = {rows.fmask[m].real(), rows.fmask[m].imag()}; + engine_->SetRowModel(rows.k_total, fmask, rows.d, rows.var); + } + + bool OnGPU() const override { return true; } + + void Fcalc(const gemmi::Model &model, std::vector> &f) override { + SetAtoms(model); + std::vector> out; + engine_->Fcalc(out); + f.resize(out.size()); + for (size_t m = 0; m < out.size(); m++) + f[m] = {out[m][0], out[m][1]}; + } + + void Fmask(const gemmi::Model &model, std::vector> &f) override { + const gemmi::SolventMasker masker(gemmi::AtomicRadiiSet::Refmac); + std::vector mask_atoms; + for (const gemmi::Atom *a : Atoms(model)) { + if ((masker.ignore_hydrogen && a->is_hydrogen()) || (masker.ignore_zero_occupancy_atoms && a->occ <= 0)) + continue; + const gemmi::Fractional f0 = cell_.fractionalize(a->pos); + double frac[3] = {f0.x, f0.y, f0.z}; + for (double &x : frac) { + x -= std::floor(x); + if (x >= 1.0) + x = 0.0; + } + mask_atoms.push_back({frac[0], frac[1], frac[2], + masker.constant_r + masker.rprobe + + gemmi::refmac_radius_for_bulk_solvent(a->element.elem)}); + } + std::vector> out; + engine_->Fmask(mask_atoms, out); + f.resize(out.size()); + for (size_t m = 0; m < out.size(); m++) + f[m] = {out[m][0], out[m][1]}; + } + + double Evaluate(const gemmi::Model &model, std::vector> *grad) override { + SetAtoms(model); + std::vector> g; + const double value = engine_->Evaluate(grad != nullptr ? &g : nullptr); + if (grad != nullptr) { + grad->resize(g.size()); + for (size_t i = 0; i < g.size(); i++) + (*grad)[i] = {g[i][0], g[i][1], g[i][2], g[i][3]}; + } + return value; + } + + void SetRestraints(int which, const RestraintSet &set) override { + engine_->SetRestraints(which, set.records, set.plane_atoms, set.slot_start, set.slot_index, set.slots); + } + + double EvaluateRestraints(int which, const std::vector> &xyzb, + std::vector> *grad) override { + std::vector> x(xyzb.size()), g; + for (size_t i = 0; i < xyzb.size(); i++) + x[i] = {static_cast(xyzb[i][0]), static_cast(xyzb[i][1]), static_cast(xyzb[i][2]), + static_cast(xyzb[i][3])}; + const double value = engine_->EvaluateRestraints(which, x, grad != nullptr ? &g : nullptr); + if (grad != nullptr) { + grad->resize(g.size()); + for (size_t i = 0; i < g.size(); i++) + (*grad)[i] = {g[i][0], g[i][1], g[i][2], g[i][3]}; + } + return value; + } + +private: + gemmi::DensityCalculator Calculator(double blur) const { + gemmi::DensityCalculator dc; + dc.d_min = d_min_; + dc.rate = 1.5; + dc.grid.unit_cell = cell_; + dc.grid.spacegroup = &sg_; + dc.blur = blur; + return dc; + } + + // The atoms' densities (made again only for the atoms whose B moved) and grid positions, to the device. + void SetAtoms(const gemmi::Model &model) { + const std::vector atoms = Atoms(model); + if (last_b_.size() != atoms.size()) { + last_b_.assign(atoms.size(), NAN); + atoms_.resize(atoms.size()); + } + std::vector> pos(atoms.size()); + const int n3[3] = {setup_.grid.nu, setup_.grid.nv, setup_.grid.nw}; + for (size_t i = 0; i < atoms.size(); i++) { + if (!(atoms[i]->b_iso == last_b_[i])) { + atoms_[i] = DensityAtom(*atoms[i], atoms[i]->b_iso, dc_); + last_b_[i] = atoms[i]->b_iso; + } + const gemmi::Fractional f = cell_.fractionalize(atoms[i]->pos); + pos[i] = {static_cast(f.x * n3[0]), static_cast(f.y * n3[1]), static_cast(f.z * n3[2]), 0}; + } + engine_->SetAtoms(atoms_, pos); + } + + int device_; + gemmi::UnitCell cell_; + const gemmi::SpaceGroup &sg_; + double d_min_; + RefinementGPUSetup setup_; + gemmi::DensityCalculator dc_; + bool supported_ = false; + size_t fft_work_ = 0, bytes_ = 0; + std::unique_ptr engine_; + std::vector last_b_; + std::vector atoms_; +}; + +} // namespace + +std::unique_ptr MakeRefinementEngineGPU(int device, const gemmi::UnitCell &cell, + const gemmi::SpaceGroup &sg, double d_min, + const RefinementReflections &refl, + const gemmi::Model &model, double b_max, size_t budget, + size_t *device_bytes) { + auto engine = std::make_unique(device, cell, sg, d_min, refl, model, b_max); + if (device_bytes != nullptr) + *device_bytes = engine->Bytes(); + if (!engine->Supported() || engine->Bytes() > budget) + return nullptr; + engine->Reserve(); + return engine; +} diff --git a/image_analysis/structure_refinement/RefinementGPU.cu b/image_analysis/structure_refinement/RefinementGPU.cu new file mode 100644 index 000000000..4dde8f429 --- /dev/null +++ b/image_analysis/structure_refinement/RefinementGPU.cu @@ -0,0 +1,691 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#include "RefinementGPUEngine.h" + +#include +#include + +#include +#include + +#include "../../common/JFJochException.h" +#include "../indexing/CUDAMemHelpers.h" + +namespace { + +void cuda_err(cudaError_t val) { + if (val != cudaSuccess) + throw JFJochException(JFJochExceptionCategory::GPUCUDAError, cudaGetErrorString(val)); +} + +void cufft_err(cufftResult val) { + if (val != CUFFT_SUCCESS) + throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "cuFFT error " + std::to_string(static_cast(val))); +} + +constexpr int THREADS = 256; +constexpr int WARP = 32; +constexpr double FOUR_PI2 = 4 * refine_math::PI * refine_math::PI; + +int blocks(size_t n) { + return static_cast((n + THREADS - 1) / THREADS); +} + +__device__ __forceinline__ int imod(int a, int n) { + const int r = a % n; + return r < 0 ? r + n : r; +} + +// gemmi's unsafe_expapprox(), as ModelDensityGPU.cu has it. +__device__ __forceinline__ float expapprox(float x) { + const float val = 12102203.1615614f * x + 1065353216.f; + const int vali = static_cast(val); + const float a = __int_as_float(vali & 0x7F800000); + const float b = __int_as_float((vali & 0x7FFFFF) | 0x3F800000); + return a * (0.509871020f + b * (0.312146713f + b * (0.166617139f + b * (-2.190619930e-3f + b * 1.3555747234e-2f)))); +} + +// exp(+2 pi i h.t) of row h and operator o, exact for the fractions a translation is made of. +__device__ __forceinline__ void op_phase(const int *h, const int *tran, int den, double &c, double &s) { + sincospi(2.0 * (h[0] * tran[0] + h[1] * tran[1] + h[2] * tran[2]) / den, &s, &c); +} + +// F of the atoms at each row, composed from the half-u transform of one copy (ModelStructureFactorsGPU.cu's +// compose_kernel), and with `evaluate` the target's term of a working row and the coefficient +// conj(dT/dF) n_cen unblur the gradient's map is made of. +__global__ void compose_target_kernel(const float2 *f1, float norm, int nu, int nv, int nw, const int *ops, int n_ops, + int den, const int *row_hkl, const double *row_scale, int rows, bool evaluate, + const float *k_total, const float2 *fmask, const float *ml_d, const float *ml_var, + const float *fo, const char *work, const char *centric, float2 *f, double *nll, + float2 *coef) { + const int m = blockIdx.x * blockDim.x + threadIdx.x; + if (m >= rows) + return; + const int h[3] = {row_hkl[3 * m], row_hkl[3 * m + 1], row_hkl[3 * m + 2]}; + double sr = 0, si = 0; + for (int o = 0; o < n_ops; o++) { + const int *rot = ops + 12 * o, *tran = rot + 9; + int k[3]; + for (int i = 0; i < 3; i++) + k[i] = (rot[i] * h[0] + rot[3 + i] * h[1] + rot[6 + i] * h[2]) / den; + const bool conj = k[0] < 0; + if (conj) + for (int i = 0; i < 3; i++) + k[i] = -k[i]; + const float2 out = f1[k[0] + static_cast(nu / 2 + 1) * (imod(k[1], nv) + static_cast(nv) * imod(k[2], nw))]; + const double vr = out.x * norm, vi = (conj ? out.y : -out.y) * norm; + double s, c; + op_phase(h, tran, den, c, s); + sr += c * vr - s * vi; + si += c * vi + s * vr; + } + const double scale = row_scale[m]; + sr *= scale; + si *= scale; + f[m] = make_float2(static_cast(sr), static_cast(si)); + if (!evaluate) + return; + double term = 0, cr = 0, ci = 0; + if (work[m]) { + const double k = k_total[m]; + const double mr = k * (sr + fmask[m].x), mi = k * (si + fmask[m].y); + const double fc = sqrt(mr * mr + mi * mi); + double dfc = 0; + term = refine_math::AmplitudeNegLogLikelihood(fo[m], fc, ml_d[m], ml_var[m], centric[m], dfc); + if (fc > 0) { + // dT/dF = k dT/d|F_model| F_model / |F_model|; the coefficient is its conjugate times the scale + const double g = k * dfc / fc * scale; + cr = g * mr; + ci = -g * mi; + } + } + nll[m] = term; + coef[m] = make_float2(static_cast(cr), static_cast(ci)); +} + +// The gradient's coefficients on the half-u spectrum: each point the sum, in order, of the terms +// (coef[row] exp(2 pi i h.t_op)) / 2 that land on it, conjugated where they land as a Friedel mate. +__global__ void scatter_kernel(const float2 *coef, const int *row_hkl, const int *ops, int n_ops, int den, + const int *point, const int *start, const int *term, int points, float2 *spectrum) { + const int s = blockIdx.x * blockDim.x + threadIdx.x; + if (s >= points) + return; + double sr = 0, si = 0; + for (int e = start[s]; e < start[s + 1]; e++) { + const int t = term[e]; + const bool conj = t & 1; + const int m = (t >> 1) / n_ops, o = (t >> 1) % n_ops; + double c, sn; + op_phase(row_hkl + 3 * m, ops + 12 * o + 9, den, c, sn); + const double ar = 0.5 * (coef[m].x * c - coef[m].y * sn), ai = 0.5 * (coef[m].x * sn + coef[m].y * c); + sr += ar; + si += conj ? -ai : ai; + } + spectrum[point[s]] = make_float2(static_cast(sr), static_cast(si)); +} + +// F_mask at each row, read at h directly (ModelStructureFactorsGPU.cu's fmask_kernel). +__global__ void fmask_kernel(const float2 *fm, float norm, const int *row_hkl, int rows, int nu, int nv, int nw, + float2 *fmask) { + const int m = blockIdx.x * blockDim.x + threadIdx.x; + if (m >= rows) + return; + int h = row_hkl[3 * m], k = row_hkl[3 * m + 1], l = row_hkl[3 * m + 2]; + const bool conj = h < 0; + if (conj) { + h = -h; + k = -k; + l = -l; + } + const float2 out = fm[h + static_cast(nu / 2 + 1) * (imod(k, nv) + static_cast(nv) * imod(l, nw))]; + const float re = out.x * norm, im = -out.y * norm; + fmask[m] = make_float2(re, conj ? -im : im); +} + +// 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). +__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) { + const int a = (blockIdx.x * blockDim.x + threadIdx.x) / WARP; + const int lane = threadIdx.x % WARP; + if (a >= n) + return; + const ModelDensityAtom at = atoms[a]; + const float4 p = pos[a]; + const float nctr[3] = {p.x, p.y, p.z}; + int lo[3], span[3]; + for (int k = 0; k < 3; k++) { + lo[k] = __float2int_rn(nctr[k]) - box[3 * a + k]; + span[k] = 2 * box[3 * a + k] + 1; + } + const int total = span[0] * span[1] * span[2]; + float gx = 0, gy = 0, gz = 0, gb = 0; + for (int t = lane; t < total; t += WARP) { + const int u = lo[0] + t % span[0], v = lo[1] + t / span[0] % span[1], w = lo[2] + t / (span[0] * span[1]); + const float fx = nctr[0] - static_cast(u); + const float fy = nctr[1] - static_cast(v); + const float fz = nctr[2] - static_cast(w); + const float z = g.m22 * fz; + const float y = g.m11 * fy + g.m12 * fz; + const float x = g.m00 * fx + (g.m01 * fy + g.m02 * fz); + const float r2 = (y * y + z * z) + x * x; + if (r2 > at.max_r2) + continue; + const float value = map[imod(u, g.nu) + static_cast(g.nu) * (imod(v, g.nv) + static_cast(g.nv) * imod(w, 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); + } + 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); + } + 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)); + } +} + +using V3f = refine_math::V3; + +__device__ V3f xyz(const float4 *p, int i) { + const float4 q = p[i]; + return {q.x, q.y, q.z}; +} + +// Every restraint's E, and its gradient into its slots, one thread per restraint. +__global__ void restraint_kernel(const RestraintRecord *records, int n, const int *plane_atoms, const float4 *xyzb, + float4 *contrib, double *energy) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= n) + return; + const RestraintRecord &r = records[i]; + V3f g[4] = {}; + int na = 0; + float e = 0; + switch (static_cast(r.kind)) { + case RestraintKind::Bond: + e = refine_math::BondTerm(xyz(xyzb, r.atom[0]), xyz(xyzb, r.atom[1]), r.ideal, r.sigma, g); + na = 2; + break; + case RestraintKind::Angle: + e = refine_math::AngleTerm(xyz(xyzb, r.atom[0]), xyz(xyzb, r.atom[1]), xyz(xyzb, r.atom[2]), r.ideal, + r.sigma, g); + na = 3; + break; + case RestraintKind::Torsion: + e = refine_math::TorsionTerm(xyz(xyzb, r.atom[0]), xyz(xyzb, r.atom[1]), xyz(xyzb, r.atom[2]), + xyz(xyzb, r.atom[3]), r.ideal, r.sigma, r.period, g); + na = 4; + break; + case RestraintKind::Chiral: + e = refine_math::ChiralTerm(xyz(xyzb, r.atom[0]), xyz(xyzb, r.atom[1]), xyz(xyzb, r.atom[2]), + xyz(xyzb, r.atom[3]), r.ideal, r.sigma, g); + na = 4; + break; + case RestraintKind::Repulsion: { + const V3f q = xyz(xyzb, r.atom[1]); + const V3f image{r.rot[0] * q.x + r.rot[1] * q.y + r.rot[2] * q.z + r.tran[0], + r.rot[3] * q.x + r.rot[4] * q.y + r.rot[5] * q.z + r.tran[1], + r.rot[6] * q.x + r.rot[7] * q.y + r.rot[8] * q.z + r.tran[2]}; + e = refine_math::RepulsionTerm(xyz(xyzb, r.atom[0]), image, r.ideal, r.sigma, g); + const V3f gi = g[1]; + g[1] = {r.rot[0] * gi.x + r.rot[3] * gi.y + r.rot[6] * gi.z, + r.rot[1] * gi.x + r.rot[4] * gi.y + r.rot[7] * gi.z, + r.rot[2] * gi.x + r.rot[5] * gi.y + r.rot[8] * gi.z}; + na = 2; + break; + } + case RestraintKind::Plane: { + V3f px[RefinementGPUEngine::MAX_PLANE_ATOMS], pg[RefinementGPUEngine::MAX_PLANE_ATOMS]; + const int m = r.atom[1]; + for (int k = 0; k < m; k++) + px[k] = xyz(xyzb, plane_atoms[r.atom[0] + k]); + e = refine_math::PlaneTerm(px, m, r.sigma, pg); + for (int k = 0; k < m; k++) + contrib[r.slot + k] = make_float4(pg[k].x, pg[k].y, pg[k].z, 0.0f); + break; + } + case RestraintKind::BSimilarity: { + const float rr = (xyzb[r.atom[0]].w - xyzb[r.atom[1]].w) / r.sigma; + e = rr * rr; + contrib[r.slot] = make_float4(0, 0, 0, 2 * rr / r.sigma); + contrib[r.slot + 1] = make_float4(0, 0, 0, -2 * rr / r.sigma); + break; + } + } + for (int k = 0; k < na; k++) + contrib[r.slot + k] = make_float4(g[k].x, g[k].y, g[k].z, 0.0f); + energy[i] = e; +} + +// Each atom's restraint gradient: its slots summed in order. +__global__ void slot_gather_kernel(const float4 *contrib, const int *start, const int *index, int n, float4 *grad) { + const int a = blockIdx.x * blockDim.x + threadIdx.x; + if (a >= n) + return; + float4 s = make_float4(0, 0, 0, 0); + for (int k = start[a]; k < start[a + 1]; k++) { + const float4 c = contrib[index[k]]; + s.x += c.x; + s.y += c.y; + s.z += c.z; + s.w += c.w; + } + grad[a] = s; +} + +CudaFFTShape ForwardShape(const ModelDensityGrid &g) { + return {.n = {g.nw, g.nv, g.nu}, .idist = g.nu * g.nv * g.nw, .odist = (g.nu / 2 + 1) * g.nv * g.nw, + .type = CUFFT_R2C}; +} + +CudaFFTShape InverseShape(const ModelDensityGrid &g) { + return {.n = {g.nw, g.nv, g.nu}, .idist = (g.nu / 2 + 1) * g.nv * g.nw, .odist = g.nu * g.nv * g.nw, + .type = CUFFT_C2R}; +} + +class DeviceGuard { +public: + explicit DeviceGuard(int device) { + cuda_err(cudaGetDevice(&previous_)); + cuda_err(cudaSetDevice(device)); + } + ~DeviceGuard() { cudaSetDevice(previous_); } + DeviceGuard(const DeviceGuard &) = delete; + DeviceGuard &operator=(const DeviceGuard &) = delete; + +private: + int previous_ = 0; +}; + +size_t GridPoints(const ModelDensityGrid &g) { + return static_cast(g.nu) * g.nv * g.nw; +} + +size_t ComplexPoints(const ModelDensityGrid &g) { + return static_cast(g.nu / 2 + 1) * g.nv * g.nw; +} + +template +void Upload(CudaDevicePtr &dst, const std::vector &src, cudaStream_t stream) { + if (!src.empty()) + cuda_err(cudaMemcpyAsync(dst, src.data(), src.size() * sizeof(T), cudaMemcpyHostToDevice, stream)); +} + +// A device buffer of at least n elements, kept between calls and grown when it is too small. +template +void Ensure(CudaDevicePtr &p, size_t &capacity, size_t n) { + if (n <= capacity && p) + return; + capacity = std::max(n, 1); + p = CudaDevicePtr(capacity, CudaAlloc::Synchronous); +} + +// cub's temporary storage to reduce n doubles. +size_t ReduceBytes(size_t n) { + size_t bytes = 0; + cuda_err(cub::DeviceReduce::Sum(nullptr, bytes, static_cast(nullptr), static_cast(nullptr), + static_cast(std::max(n, 1)))); + return bytes; +} + +} // namespace + +struct RefinementGPUEngineImpl { + struct Restraints { + int records = 0, atoms = 0; + size_t cap_records = 0, cap_planes = 0, cap_start = 0, cap_index = 0, cap_slots = 0, cap_reduce = 0, + cap_atoms = 0, cap_xyzb = 0, reduce_bytes = 0; + CudaDevicePtr rec; + CudaDevicePtr plane_atoms, slot_start, slot_index; + CudaDevicePtr contrib, xyzb, grad; + CudaDevicePtr energy, sum; + CudaDevicePtr reduce_temp; + }; + + RefinementGPUSetup setup; + size_t work_bytes; + int device = 0; + CudaStream stream; + int n_atoms = 0; + + CudaDevicePtr pos; + CudaDevicePtr box; + CudaDevicePtr mask_atoms; + CudaDevicePtr grid; + CudaDevicePtr spectrum; + CudaDevicePtr fft_work; + CudaDevicePtr ops, row_hkl; + CudaDevicePtr row_scale; + CudaDevicePtr k_total, ml_d, ml_var, fo; + CudaDevicePtr fmask_scaled, f, fm, coef; + CudaDevicePtr work, centric; + CudaDevicePtr nll, nll_sum; + CudaDevicePtr reduce_temp; + size_t reduce_bytes = 0; + CudaDevicePtr scatter_point, scatter_start, scatter_term; + CudaDevicePtr grad; + CudaDevicePtr atoms_copy; // the density gather keeps its own; this one is the gradient's + std::unique_ptr density; + std::unique_ptr mask; + CudaFFTPlan forward, inverse; + ModelDensityGPU::Geometry geom{}; + Restraints restraints[2]; + + RefinementGPUEngineImpl(int dev, const RefinementGPUSetup &s, size_t fft_work_bytes) : setup(s), work_bytes(fft_work_bytes), device(dev) { + const auto sync = CudaAlloc::Synchronous; + const size_t na = std::max(setup.max_atoms, 1), nr = std::max(setup.rows.size(), 1); + pos = CudaDevicePtr(na, sync); + box = CudaDevicePtr(3 * na, sync); + mask_atoms = CudaDevicePtr(na, sync); + atoms_copy = CudaDevicePtr(na, sync); + grad = CudaDevicePtr(na, sync); + grid = CudaDevicePtr(GridPoints(setup.grid), sync); + spectrum = CudaDevicePtr(ComplexPoints(setup.grid), sync); + fft_work = CudaDevicePtr(std::max(work_bytes, 1), sync); + ops = CudaDevicePtr(std::max(12 * setup.sym_ops.size(), 1), sync); + row_hkl = CudaDevicePtr(3 * nr, sync); + row_scale = CudaDevicePtr(nr, sync); + k_total = CudaDevicePtr(nr, sync); + ml_d = CudaDevicePtr(nr, sync); + ml_var = CudaDevicePtr(nr, sync); + fo = CudaDevicePtr(nr, sync); + fmask_scaled = CudaDevicePtr(nr, sync); + f = CudaDevicePtr(nr, sync); + fm = CudaDevicePtr(nr, sync); + coef = CudaDevicePtr(nr, sync); + work = CudaDevicePtr(nr, sync); + centric = CudaDevicePtr(nr, sync); + nll = CudaDevicePtr(nr, sync); + nll_sum = CudaDevicePtr(1, sync); + reduce_bytes = ReduceBytes(nr); + reduce_temp = CudaDevicePtr(std::max(reduce_bytes, 1), sync); + scatter_point = CudaDevicePtr(std::max(setup.scatter_point.size(), 1), sync); + scatter_start = CudaDevicePtr(std::max(setup.scatter_start.size(), 1), sync); + scatter_term = CudaDevicePtr(std::max(setup.scatter_term.size(), 1), sync); + density = std::make_unique(stream, setup.max_atoms, setup.max_pairs, + ModelDensityGPU::Bricks(setup.grid.nu, setup.grid.nv, setup.grid.nw)); + mask = std::make_unique(stream, GridPoints(setup.grid)); + ModelMaskGrid mask_grid; + mask_grid.nu = setup.grid.nu; + mask_grid.nv = setup.grid.nv; + mask_grid.nw = setup.grid.nw; + std::copy(setup.grid.orth, setup.grid.orth + 9, mask_grid.orth); + mask_grid.volume = setup.volume; + mask->SetGrid(mask_grid, setup.mask_ops); + + std::vector flat_ops; + for (const auto &o : setup.sym_ops) + flat_ops.insert(flat_ops.end(), o.begin(), o.end()); + std::vector hkl; + for (const auto &h : setup.rows) + hkl.insert(hkl.end(), h.begin(), h.end()); + Upload(ops, flat_ops, stream); + Upload(row_hkl, hkl, stream); + Upload(fo, setup.fo, stream); + Upload(work, setup.work, stream); + Upload(centric, setup.centric, stream); + Upload(scatter_point, setup.scatter_point, stream); + Upload(scatter_start, setup.scatter_start, stream); + Upload(scatter_term, setup.scatter_term, stream); + + const double inv_n[3] = {1.0 / setup.grid.nu, 1.0 / setup.grid.nv, 1.0 / setup.grid.nw}; + geom.nu = setup.grid.nu; + geom.nv = setup.grid.nv; + geom.nw = setup.grid.nw; + geom.m00 = static_cast(setup.grid.orth[0] * inv_n[0]); + geom.m01 = static_cast(setup.grid.orth[1] * inv_n[1]); + geom.m02 = static_cast(setup.grid.orth[2] * inv_n[2]); + geom.m11 = static_cast(setup.grid.orth[4] * inv_n[1]); + geom.m12 = static_cast(setup.grid.orth[5] * inv_n[2]); + geom.m22 = static_cast(setup.grid.orth[8] * inv_n[2]); + + forward = CudaFFTPlan(ForwardShape(setup.grid), stream, fft_work); + inverse = CudaFFTPlan(InverseShape(setup.grid), stream, fft_work); + if (forward.WorkBytes() > work_bytes || inverse.WorkBytes() > work_bytes) + throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "refinement: cuFFT work area over the reserve"); + cuda_err(cudaStreamSynchronize(stream)); + } + + ~RefinementGPUEngineImpl() { cudaStreamSynchronize(stream); } + + float Norm() const { return static_cast(setup.volume / static_cast(GridPoints(setup.grid))); } + + // The atoms' density on the grid and its transform into the spectrum. + void Transform() { + density->Compute(pos, grid); + cufft_err(forward.R2C(grid, reinterpret_cast(spectrum.get()))); + } + + void Compose(bool evaluate) { + const int rows = static_cast(setup.rows.size()); + if (rows == 0) + return; + compose_target_kernel<<>>( + spectrum, Norm(), setup.grid.nu, setup.grid.nv, setup.grid.nw, ops, static_cast(setup.sym_ops.size()), + setup.den, row_hkl, row_scale, rows, evaluate, k_total, fmask_scaled, ml_d, ml_var, fo, work, centric, f, + nll, coef); + cuda_err(cudaGetLastError()); + } +}; + +size_t RefinementGPUEngine::FFTWorkBytes(int device, const ModelDensityGrid &grid) { + DeviceGuard guard(device); + return std::max(CudaFFTPlan::WorkBytes(ForwardShape(grid)), CudaFFTPlan::WorkBytes(InverseShape(grid))); +} + +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)); + 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() + s.scatter_term.size()) * sizeof(int); + bytes += ModelDensityGPU::DeviceBytes(s.max_atoms, s.max_pairs, ModelDensityGPU::Bricks(s.grid.nu, s.grid.nv, s.grid.nw)); + bytes += ModelMaskGPU::DeviceBytes(GridPoints(s.grid)); + bytes += max_records * (sizeof(RestraintRecord) + sizeof(double)) + max_slots * (sizeof(float4) + sizeof(int)) + + na * (2 * sizeof(float4) + sizeof(int)); + return bytes; +} + +RefinementGPUEngine::RefinementGPUEngine(int device, const RefinementGPUSetup &setup, size_t fft_work_bytes) { + DeviceGuard guard(device); + impl_ = std::make_unique(device, setup, fft_work_bytes); +} + +RefinementGPUEngine::~RefinementGPUEngine() { + if (impl_) { + DeviceGuard guard(impl_->device); + impl_.reset(); + } +} + +void RefinementGPUEngine::SetBlur(double blur) { + RefinementGPUEngineImpl &e = *impl_; + DeviceGuard guard(e.device); + std::vector scale(e.setup.rows.size()); + for (size_t m = 0; m < scale.size(); m++) + scale[m] = e.setup.centring * std::exp(blur * 0.25 * e.setup.inv_d2[m]); + Upload(e.row_scale, scale, e.stream); + cuda_err(cudaStreamSynchronize(e.stream)); +} + +void RefinementGPUEngine::SetRowModel(const std::vector &k_total, const std::vector> &fmask, + const std::vector &d, const std::vector &var) { + RefinementGPUEngineImpl &e = *impl_; + DeviceGuard guard(e.device); + Upload(e.k_total, k_total, e.stream); + if (!fmask.empty()) + cuda_err(cudaMemcpyAsync(e.fmask_scaled, fmask.data(), fmask.size() * sizeof(float2), cudaMemcpyHostToDevice, + e.stream)); + Upload(e.ml_d, d, e.stream); + Upload(e.ml_var, var, e.stream); + cuda_err(cudaStreamSynchronize(e.stream)); +} + +void RefinementGPUEngine::SetAtoms(const std::vector &atoms, const std::vector> &pos) { + RefinementGPUEngineImpl &e = *impl_; + DeviceGuard guard(e.device); + if (atoms.size() > e.setup.max_atoms || pos.size() != atoms.size()) + throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "refinement: atoms over the reserve"); + e.density->SetAtoms(e.setup.grid, atoms); + e.n_atoms = static_cast(atoms.size()); + const std::vector b = ModelDensityGPU::Boxes(e.setup.grid, atoms); + if (!atoms.empty()) { + cuda_err(cudaMemcpyAsync(e.pos, pos.data(), pos.size() * sizeof(float4), cudaMemcpyHostToDevice, e.stream)); + cuda_err(cudaMemcpyAsync(e.atoms_copy, atoms.data(), atoms.size() * sizeof(ModelDensityAtom), + cudaMemcpyHostToDevice, e.stream)); + Upload(e.box, b, e.stream); + } + cuda_err(cudaStreamSynchronize(e.stream)); +} + +void RefinementGPUEngine::Fcalc(std::vector> &f) { + RefinementGPUEngineImpl &e = *impl_; + DeviceGuard guard(e.device); + e.Transform(); + e.Compose(false); + f.resize(e.setup.rows.size()); + if (!f.empty()) + cuda_err(cudaMemcpyAsync(f.data(), e.f, f.size() * sizeof(float2), cudaMemcpyDeviceToHost, e.stream)); + cuda_err(cudaStreamSynchronize(e.stream)); +} + +void RefinementGPUEngine::Fmask(const std::vector &mask_atoms, std::vector> &f) { + RefinementGPUEngineImpl &e = *impl_; + DeviceGuard guard(e.device); + if (mask_atoms.size() > e.setup.max_atoms) + throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "refinement: mask atoms over the reserve"); + if (!mask_atoms.empty()) + cuda_err(cudaMemcpyAsync(e.mask_atoms, mask_atoms.data(), mask_atoms.size() * sizeof(ModelMaskAtom), + cudaMemcpyHostToDevice, e.stream)); + e.mask->Compute(e.mask_atoms, static_cast(mask_atoms.size()), e.grid); + cufft_err(e.forward.R2C(e.grid, reinterpret_cast(e.spectrum.get()))); + const int rows = static_cast(e.setup.rows.size()); + f.resize(rows); + if (rows > 0) { + fmask_kernel<<>>(e.spectrum, e.Norm(), e.row_hkl, rows, e.setup.grid.nu, + e.setup.grid.nv, e.setup.grid.nw, e.fm); + cuda_err(cudaGetLastError()); + cuda_err(cudaMemcpyAsync(f.data(), e.fm, rows * sizeof(float2), cudaMemcpyDeviceToHost, e.stream)); + } + cuda_err(cudaStreamSynchronize(e.stream)); +} + +double RefinementGPUEngine::Evaluate(std::vector> *grad) { + RefinementGPUEngineImpl &e = *impl_; + DeviceGuard guard(e.device); + const int rows = static_cast(e.setup.rows.size()); + e.Transform(); + e.Compose(true); + double sum = 0; + if (rows > 0) { + size_t temp = e.reduce_bytes; + cuda_err(cub::DeviceReduce::Sum(e.reduce_temp.get(), temp, e.nll.get(), e.nll_sum.get(), rows, e.stream)); + cuda_err(cudaMemcpyAsync(&sum, e.nll_sum, sizeof(double), cudaMemcpyDeviceToHost, e.stream)); + } + if (grad != nullptr) { + const ModelDensityGrid &g = e.setup.grid; + cuda_err(cudaMemsetAsync(e.spectrum, 0, ComplexPoints(g) * sizeof(float2), e.stream)); + const int points = static_cast(e.setup.scatter_point.size()); + if (points > 0) { + scatter_kernel<<>>(e.coef, e.row_hkl, e.ops, + static_cast(e.setup.sym_ops.size()), e.setup.den, + e.scatter_point, e.scatter_start, e.scatter_term, + points, e.spectrum); + cuda_err(cudaGetLastError()); + } + cufft_err(e.inverse.C2R(reinterpret_cast(e.spectrum.get()), e.grid)); + if (e.n_atoms > 0) { + const size_t threads = static_cast(e.n_atoms) * WARP; + atom_gradient_kernel<<>>(e.atoms_copy, e.pos, e.box, e.n_atoms, e.geom, + e.grid, e.Norm(), e.grad); + cuda_err(cudaGetLastError()); + } + grad->resize(e.n_atoms); + if (e.n_atoms > 0) + cuda_err(cudaMemcpyAsync(grad->data(), e.grad, e.n_atoms * sizeof(float4), cudaMemcpyDeviceToHost, e.stream)); + } + cuda_err(cudaStreamSynchronize(e.stream)); + return sum; +} + +void RefinementGPUEngine::SetRestraints(int which, const std::vector &records, + const std::vector &plane_atoms, const std::vector &slot_start, + const std::vector &slot_index, int slots) { + RefinementGPUEngineImpl &e = *impl_; + DeviceGuard guard(e.device); + RefinementGPUEngineImpl::Restraints &r = e.restraints[which]; + for (const RestraintRecord &rec : records) + if (rec.kind == static_cast(RestraintKind::Plane) && rec.atom[1] > MAX_PLANE_ATOMS) + throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "refinement: a plane of too many atoms"); + r.records = static_cast(records.size()); + r.atoms = static_cast(slot_start.empty() ? 0 : slot_start.size() - 1); + Ensure(r.rec, r.cap_records, records.size()); + Ensure(r.energy, r.cap_reduce, records.size()); + Ensure(r.plane_atoms, r.cap_planes, plane_atoms.size()); + Ensure(r.slot_start, r.cap_start, slot_start.size()); + Ensure(r.slot_index, r.cap_index, slot_index.size()); + Ensure(r.contrib, r.cap_slots, static_cast(slots)); + Ensure(r.xyzb, r.cap_xyzb, r.atoms); + Ensure(r.grad, r.cap_atoms, r.atoms); + if (!r.sum) + r.sum = CudaDevicePtr(1, CudaAlloc::Synchronous); + const size_t bytes = ReduceBytes(r.cap_reduce); + if (!r.reduce_temp || bytes > r.reduce_bytes) { + r.reduce_bytes = bytes; + r.reduce_temp = CudaDevicePtr(std::max(bytes, 1), CudaAlloc::Synchronous); + } + Upload(r.rec, records, e.stream); + Upload(r.plane_atoms, plane_atoms, e.stream); + Upload(r.slot_start, slot_start, e.stream); + Upload(r.slot_index, slot_index, e.stream); + cuda_err(cudaStreamSynchronize(e.stream)); +} + +double RefinementGPUEngine::EvaluateRestraints(int which, const std::vector> &xyzb, + std::vector> *grad) { + RefinementGPUEngineImpl &e = *impl_; + DeviceGuard guard(e.device); + RefinementGPUEngineImpl::Restraints &r = e.restraints[which]; + if (static_cast(xyzb.size()) != r.atoms) + throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "refinement: restraint atoms do not match"); + if (r.atoms > 0) + cuda_err(cudaMemcpyAsync(r.xyzb, xyzb.data(), xyzb.size() * sizeof(float4), cudaMemcpyHostToDevice, e.stream)); + double sum = 0; + if (r.records > 0) { + restraint_kernel<<>>(r.rec, r.records, r.plane_atoms, r.xyzb, r.contrib, + r.energy); + cuda_err(cudaGetLastError()); + size_t temp = r.reduce_bytes; + cuda_err(cub::DeviceReduce::Sum(r.reduce_temp.get(), temp, r.energy.get(), r.sum.get(), r.records, e.stream)); + cuda_err(cudaMemcpyAsync(&sum, r.sum, sizeof(double), cudaMemcpyDeviceToHost, e.stream)); + } + if (grad != nullptr) { + grad->resize(r.atoms); + if (r.atoms > 0) { + if (r.records > 0) { + slot_gather_kernel<<>>(r.contrib, r.slot_start, r.slot_index, + r.atoms, r.grad); + cuda_err(cudaGetLastError()); + } else + cuda_err(cudaMemsetAsync(r.grad, 0, r.atoms * sizeof(float4), e.stream)); + cuda_err(cudaMemcpyAsync(grad->data(), r.grad, r.atoms * sizeof(float4), cudaMemcpyDeviceToHost, e.stream)); + } + } + cuda_err(cudaStreamSynchronize(e.stream)); + return sum; +} diff --git a/image_analysis/structure_refinement/RefinementGPU.h b/image_analysis/structure_refinement/RefinementGPU.h new file mode 100644 index 000000000..965fab1ad --- /dev/null +++ b/image_analysis/structure_refinement/RefinementGPU.h @@ -0,0 +1,24 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#pragma once + +// Restrained refinement's targets on a GPU (CUDA builds only): the RefinementEngine of RefinementXray.h with +// everything per evaluation on the device - the atoms' density (ModelDensityGPU's gather), its transform, the +// symmetry composed in reciprocal space, the likelihood target, the derivative map by the inverse transform +// of the target's derivatives and every atom's gradient read off it, and the geometry restraints. The host +// passes positions in and gets one value and one gradient per atom back. + +#include + +#include "RefinementXray.h" + +// The engine for refining `model` against `refl` on `device`, or null where the GPU cannot take this case (an +// atom's box wider than the cell, too many symmetry operators) or its memory is over `budget` bytes. +// `b_max`: the largest B an atom may reach, which the gather's reserve is sized for. Throws JFJochException on a +// CUDA failure. +std::unique_ptr MakeRefinementEngineGPU(int device, const gemmi::UnitCell &cell, + const gemmi::SpaceGroup &sg, double d_min, + const RefinementReflections &refl, + const gemmi::Model &model, double b_max, size_t budget, + size_t *device_bytes = nullptr); diff --git a/image_analysis/structure_refinement/RefinementGPUEngine.h b/image_analysis/structure_refinement/RefinementGPUEngine.h new file mode 100644 index 000000000..a8b34807a --- /dev/null +++ b/image_analysis/structure_refinement/RefinementGPUEngine.h @@ -0,0 +1,78 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#pragma once + +// The device half of restrained refinement on a GPU (RefinementGPU.h, CUDA builds only). The host half works +// out with gemmi everything that depends on the cell, the group and the reflections; this half grids the atoms, +// transforms, composes the symmetry, evaluates the likelihood target, takes its derivative back to every atom +// by the FFT method, and evaluates the geometry restraints. Every kernel is deterministic - gathers and fixed +// reduction trees, no floating-point atomics - so the same input gives the same bits on every run. + +#include +#include +#include +#include + +#include "ModelDensityGPU.h" +#include "ModelMaskGPU.h" +#include "RefinementMath.h" + +// Everything an engine needs that does not depend on the atoms. +struct RefinementGPUSetup { + ModelDensityGrid grid; // the structure-factor grid, u fastest + double volume = 0; + std::vector> sym_ops; // gemmi's Op: rot row-major and tran, in units of 1 / den + int den = 24; + double centring = 1; // the number of centring vectors + std::vector mask_ops; // every operator with every centring vector + std::vector> rows; // the reflections, in the order their numbers come back + std::vector inv_d2; + std::vector fo; + std::vector work, centric; + // The adjoint of the composition, as a gather: the half-u spectrum points the gradient's coefficients + // land on (ascending), and for each the terms (row * ops + op) * 2 + conj adding to it, in order. + std::vector scatter_point, scatter_start, scatter_term; + size_t max_atoms = 0, max_pairs = 0; +}; + +struct RefinementGPUEngineImpl; + +class RefinementGPUEngine { +public: + static size_t DeviceBytes(const RefinementGPUSetup &setup, size_t fft_work_bytes, size_t max_records, + size_t max_slots); + static size_t FFTWorkBytes(int device, const ModelDensityGrid &grid); + + RefinementGPUEngine(int device, const RefinementGPUSetup &setup, size_t fft_work_bytes); + ~RefinementGPUEngine(); + + // n_cen exp(blur s^2 / 4) per reflection: prepare_asu_data()'s unblur. + void SetBlur(double blur); + // Per macro-cycle: F_model = k_total (F + fmask), and the likelihood's D and variance. + void SetRowModel(const std::vector &k_total, const std::vector> &fmask, + const std::vector &d, const std::vector &var); + + // The atoms, as ModelStructureFactorsGPU takes them: densities, grid positions (fractional times the grid + // size, not wrapped) and the boxes of ModelDensityGPU::Boxes(). + void SetAtoms(const std::vector &atoms, const std::vector> &pos); + // F of the atoms at the rows. + void Fcalc(std::vector> &f); + // The bulk-solvent mask's F at the rows (the mask atoms fractional, wrapped). + void Fmask(const std::vector &mask_atoms, std::vector> &f); + // -ln L over the working rows for the atoms last set and, with `grad`, dT/d(x, y, z, B) per atom. + double Evaluate(std::vector> *grad); + + // Restraint set `which` (0 or 1), as RestraintSet holds it. + void SetRestraints(int which, const std::vector &records, const std::vector &plane_atoms, + const std::vector &slot_start, const std::vector &slot_index, int slots); + // The sum of E of set `which` at `xyzb` (x, y, z, B per atom) and dE/d(x, y, z, B) per atom. + double EvaluateRestraints(int which, const std::vector> &xyzb, + std::vector> *grad); + + // The most atoms of a plane restraint the kernel takes. + static constexpr int MAX_PLANE_ATOMS = 32; + +private: + std::unique_ptr impl_; +}; diff --git a/image_analysis/structure_refinement/RefinementMath.h b/image_analysis/structure_refinement/RefinementMath.h new file mode 100644 index 000000000..996a96ac6 --- /dev/null +++ b/image_analysis/structure_refinement/RefinementMath.h @@ -0,0 +1,268 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#pragma once + +// The per-term arithmetic of restrained refinement (Refinement.h), written once for the CPU and the GPU: +// the maximum-likelihood amplitude target of one reflection and the geometry restraints, each with its +// derivatives. Plain functions on plain numbers, so the CUDA engine (RefinementGPU.cu) and the CPU path +// (RefinementXray.cpp, RefinementRestraints.cpp) evaluate the same expressions. + +#include + +#if defined(__CUDACC__) +#define REFINE_HD __host__ __device__ +#else +#define REFINE_HD +#endif + +// What a restraint is, for the evaluators on either side (RefinementRestraints.h builds them). +enum class RestraintKind : int { Bond = 0, Angle, Torsion, Chiral, Plane, Repulsion, BSimilarity }; + +// One restraint. E = ((value - ideal) / sigma)^2, summed over a set. +struct RestraintRecord { + int kind = 0; + int atom[4] = {0, 0, 0, 0}; // Plane: atom[0] the first of its atoms in plane_atoms, atom[1] their number + float ideal = 0; // A, degrees, A^3 (chiral) or A (repulsion: the contact distance) + float sigma = 1; + int period = 0; // Torsion + int slot = 0; // the first of its gradient slots, one per atom + // Repulsion: where the second atom is, as the image x' = rot x + tran of its position (Cartesian); + // the identity elsewhere. + float rot[9] = {1, 0, 0, 0, 1, 0, 0, 0, 1}; + float tran[3] = {0, 0, 0}; +}; + +namespace refine_math { + +constexpr double PI = 3.14159265358979323846; +constexpr double DEG = PI / 180.0; + +// ln I0(x) and I1(x)/I0(x) for x >= 0, from the polynomial approximations of Abramowitz & Stegun (1964), +// Handbook of Mathematical Functions, 9.8.1-9.8.4 (relative error below 2e-7); for large x through the +// exponentially scaled form, so nothing overflows. +template +REFINE_HD inline void BesselI0I1(T x, T &log_i0, T &i1_over_i0) { + if (x < T(3.75)) { + const T t = (x / T(3.75)) * (x / T(3.75)); + const T i0 = T(1.0) + t * (T(3.5156229) + t * (T(3.0899424) + t * (T(1.2067492) + + t * (T(0.2659732) + t * (T(0.0360768) + t * T(0.0045813)))))); + const T i1 = x * (T(0.5) + t * (T(0.87890594) + t * (T(0.51498869) + t * (T(0.15084934) + + t * (T(0.02658733) + t * (T(0.00301532) + t * T(0.00032411))))))); + log_i0 = log(i0); + i1_over_i0 = i1 / i0; + } else { + const T t = T(3.75) / x; + // sqrt(x) exp(-x) I0(x) and sqrt(x) exp(-x) I1(x) + const T s0 = T(0.39894228) + t * (T(0.01328592) + t * (T(0.00225319) + t * (T(-0.00157565) + + t * (T(0.00916281) + t * (T(-0.02057706) + t * (T(0.02635537) + t * (T(-0.01647633) + + t * T(0.00392377)))))))); + const T s1 = T(0.39894228) + t * (T(-0.03988024) + t * (T(-0.00362018) + t * (T(0.00163801) + + t * (T(-0.01031555) + t * (T(0.02282967) + t * (T(-0.02895312) + t * (T(0.01787654) + + t * T(-0.00420059)))))))); + log_i0 = x + log(s0) - T(0.5) * log(x); + i1_over_i0 = s1 / s0; + } +} + +// The likelihood of one observed amplitude given the model amplitude Fc - the Rice distribution for an +// acentric reflection, its centric counterpart otherwise - with D the fraction of the model that is +// right and var the variance the rest of the structure and the measurement leave (epsilon beta plus the +// observation's own sigma^2). Returns -ln L without the terms that depend on neither Fc nor the +// parameters, and sets d(-ln L)/dFc. +// Following Pannu & Read (1996) Acta Cryst. A52, 659-668 and Murshudov, Vagin & Dodson (1997) Acta Cryst. +// D53, 240-255 +template +REFINE_HD inline T AmplitudeNegLogLikelihood(T fo, T fc, T d, T var, bool centric, T &dfc) { + if (!centric) { + const T x = T(2.0) * d * fo * fc / var; + T log_i0, ratio; + BesselI0I1(x, log_i0, ratio); + dfc = (T(2.0) * d * d * fc - ratio * T(2.0) * d * fo) / var; + return log(var) + (fo * fo + d * d * fc * fc) / var - log_i0; + } + const T x = d * fo * fc / var; + // ln cosh x = x + ln(1 + exp(-2x)) - ln 2 for x >= 0 + const T log_cosh = x + log(T(1.0) + exp(T(-2.0) * x)) - T(0.6931471805599453); + dfc = (d * d * fc - tanh(x) * d * fo) / var; + return T(0.5) * log(var) + (fo * fo + d * d * fc * fc) / (T(2.0) * var) - log_cosh; +} + +// --- geometry restraints: E = ((value - ideal) / sigma)^2 and dE/dx for each atom --- + +template +struct V3 { + T x, y, z; +}; + +template REFINE_HD inline V3 operator-(V3 a, V3 b) { return {a.x - b.x, a.y - b.y, a.z - b.z}; } +template REFINE_HD inline V3 operator+(V3 a, V3 b) { return {a.x + b.x, a.y + b.y, a.z + b.z}; } +template REFINE_HD inline V3 operator*(T s, V3 a) { return {s * a.x, s * a.y, s * a.z}; } +template REFINE_HD inline T Dot(V3 a, V3 b) { return a.x * b.x + a.y * b.y + a.z * b.z; } +template REFINE_HD inline V3 Cross(V3 a, V3 b) { + return {a.y * b.z - a.z * b.y, a.z * b.x - a.x * b.z, a.x * b.y - a.y * b.x}; +} +template REFINE_HD inline T Norm(V3 a) { return sqrt(Dot(a, a)); } + +// A bond length. +template +REFINE_HD inline T BondTerm(V3 a, V3 b, T ideal, T sigma, V3 g[2]) { + const V3 d = a - b; + const T len = Norm(d); + const T r = (len - ideal) / sigma; + const T f = len > T(0) ? T(2.0) * r / (sigma * len) : T(0); + g[0] = f * d; + g[1] = (-f) * d; + return r * r; +} + +// A bond angle at b, ideal and sigma in degrees. +template +REFINE_HD inline T AngleTerm(V3 a, V3 b, V3 c, T ideal, T sigma, V3 g[3]) { + const V3 u = a - b, v = c - b; + const T lu = Norm(u), lv = Norm(v); + const V3 cr = Cross(u, v); + const T theta = atan2(Norm(cr), Dot(u, v)); + const T s = T(DEG); + const T r = (theta - ideal * s) / (sigma * s); + // dtheta/du = (cos(theta) u/|u| - v/|v|) / (|u| sin(theta)), and the same with u and v swapped + const T sin_t = fmax(sin(theta), T(1e-6)), cos_t = cos(theta); + const T f = T(2.0) * r / (sigma * s); + const V3 du = (T(1.0) / (lu * sin_t)) * ((cos_t / lu) * u - (T(1.0) / lv) * v); + const V3 dv = (T(1.0) / (lv * sin_t)) * ((cos_t / lv) * v - (T(1.0) / lu) * u); + g[0] = f * du; + g[2] = f * dv; + g[1] = (-f) * (du + dv); + return r * r; +} + +// The dihedral angle p0-p1-p2-p3 in radians, and its derivatives with respect to the four positions. +// Following Blondel & Karplus (1996) J. Comput. Chem. 17, 1132-1141 +template +REFINE_HD inline T Dihedral(V3 p0, V3 p1, V3 p2, V3 p3, V3 d[4]) { + const V3 b1 = p1 - p0, b2 = p2 - p1, b3 = p3 - p2; + const V3 m = Cross(b1, b2), n = Cross(b2, b3); + const T lb2 = Norm(b2); + const T mm = fmax(Dot(m, m), T(1e-12)), nn = fmax(Dot(n, n), T(1e-12)); + const T phi = atan2(lb2 * Dot(b1, n), Dot(m, n)); + d[0] = (-lb2 / mm) * m; + d[3] = (lb2 / nn) * n; + const T bb = lb2 * lb2; + const T p = Dot(b1, b2) / bb, q = Dot(b3, b2) / bb; + d[1] = (-(T(1.0) + p)) * d[0] + q * d[3]; + d[2] = (-T(1.0)) * (d[0] + d[1] + d[3]); + return phi; +} + +// A torsion restrained to ideal (degrees) modulo 360/period; period 0 is read as 1. +template +REFINE_HD inline T TorsionTerm(V3 p0, V3 p1, V3 p2, V3 p3, T ideal, T sigma, int period, V3 g[4]) { + V3 d[4]; + const T phi = Dihedral(p0, p1, p2, p3, d) / T(DEG); + const T step = T(360.0) / T(period > 0 ? period : 1); + T delta = phi - ideal; + delta -= step * floor(delta / step + T(0.5)); + const T r = delta / sigma; + const T f = T(2.0) * r / (sigma * T(DEG)); + for (int i = 0; i < 4; i++) + g[i] = f * d[i]; + return r * r; +} + +// A chiral volume (c the centre): (a1 - c) . ((a2 - c) x (a3 - c)). +template +REFINE_HD inline T ChiralTerm(V3 c, V3 a1, V3 a2, V3 a3, T ideal, T sigma, V3 g[4]) { + const V3 x1 = a1 - c, x2 = a2 - c, x3 = a3 - c; + const T vol = Dot(x1, Cross(x2, x3)); + const T r = (vol - ideal) / sigma; + const T f = T(2.0) * r / sigma; + g[1] = f * Cross(x2, x3); + g[2] = f * Cross(x3, x1); + g[3] = f * Cross(x1, x2); + g[0] = (-T(1.0)) * (g[1] + g[2] + g[3]); + return r * r; +} + +// A repulsion below the contact distance d0 (none above it). +template +REFINE_HD inline T RepulsionTerm(V3 a, V3 b, T d0, T sigma, V3 g[2]) { + const V3 d = a - b; + const T len = Norm(d); + if (len >= d0 || len <= T(0)) { + g[0] = g[1] = V3{0, 0, 0}; + return T(0); + } + const T r = (d0 - len) / sigma; + const T f = T(-2.0) * r / (sigma * len); + g[0] = f * d; + g[1] = (-f) * d; + return r * r; +} + +// The unit eigenvector of the smallest eigenvalue of a symmetric 3x3 matrix (a00 a11 a22 a01 a02 a12), by +// cyclic Jacobi rotations. +template +REFINE_HD inline V3 SmallestEigenvector(const T m_in[6]) { + T a[3][3] = {{m_in[0], m_in[3], m_in[4]}, {m_in[3], m_in[1], m_in[5]}, {m_in[4], m_in[5], m_in[2]}}; + T v[3][3] = {{1, 0, 0}, {0, 1, 0}, {0, 0, 1}}; + for (int sweep = 0; sweep < 12; sweep++) + for (int p = 0; p < 2; p++) + for (int q = p + 1; q < 3; q++) { + if (fabs(a[p][q]) < T(1e-30)) + continue; + const T theta = (a[q][q] - a[p][p]) / (T(2.0) * a[p][q]); + const T t = (theta >= 0 ? T(1.0) : T(-1.0)) / (fabs(theta) + sqrt(theta * theta + T(1.0))); + const T c = T(1.0) / sqrt(t * t + T(1.0)), s = t * c; + for (int k = 0; k < 3; k++) { + const T akp = a[k][p], akq = a[k][q]; + a[k][p] = c * akp - s * akq; + a[k][q] = s * akp + c * akq; + } + for (int k = 0; k < 3; k++) { + const T apk = a[p][k], aqk = a[q][k]; + a[p][k] = c * apk - s * aqk; + a[q][k] = s * apk + c * aqk; + } + for (int k = 0; k < 3; k++) { + const T vkp = v[k][p], vkq = v[k][q]; + v[k][p] = c * vkp - s * vkq; + v[k][q] = s * vkp + c * vkq; + } + } + int i = 0; + if (a[1][1] < a[i][i]) + i = 1; + if (a[2][2] < a[i][i]) + i = 2; + return {v[0][i], v[1][i], v[2][i]}; +} + +// A plane: the distances of n atoms from their least-squares plane. The plane is optimal, so the +// derivative of each atom's distance carries no term from the plane moving: dE/dx_i = 2 (n.(x_i - c)) n / sigma^2. +template +REFINE_HD inline T PlaneTerm(const V3 *x, int n, T sigma, V3 *g) { + V3 c{0, 0, 0}; + for (int i = 0; i < n; i++) + c = c + x[i]; + c = (T(1.0) / T(n)) * c; + T m[6] = {0, 0, 0, 0, 0, 0}; + for (int i = 0; i < n; i++) { + const V3 d = x[i] - c; + m[0] += d.x * d.x; + m[1] += d.y * d.y; + m[2] += d.z * d.z; + m[3] += d.x * d.y; + m[4] += d.x * d.z; + m[5] += d.y * d.z; + } + const V3 nv = SmallestEigenvector(m); + T e = 0; + for (int i = 0; i < n; i++) { + const T r = Dot(nv, x[i] - c) / sigma; + e += r * r; + g[i] = (T(2.0) * r / sigma) * nv; + } + return e; +} + +} // namespace refine_math diff --git a/image_analysis/structure_refinement/RefinementRestraints.cpp b/image_analysis/structure_refinement/RefinementRestraints.cpp new file mode 100644 index 000000000..3b79485ef --- /dev/null +++ b/image_analysis/structure_refinement/RefinementRestraints.cpp @@ -0,0 +1,568 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#include "RefinementRestraints.h" + +#include +#include +#include +#include + +#include "gemmi/neighbor.hpp" +#include "gemmi/polyheur.hpp" // setup_entities +#include "gemmi/riding_h.hpp" // place_hydrogens_on_all_atoms +#include "gemmi/topo.hpp" + +#include "MonomerLibrary.h" +#include "RefinementMath.h" + +using refine_math::V3; + +namespace { + +constexpr float CHIRAL_SIGMA = 0.2f; // A^3 +constexpr float REPULSION_SIGMA = 0.2f; // A +constexpr float HBOND_CONTACT = 2.6f; // A +constexpr float METAL_CONTACT = 1.9f; // A, to N, O or S +// B similarity: main chain and the rest, for bonded pairs and pairs bonded to a common atom. +constexpr float B_SIGMA_MAIN_BOND = 1.5f, B_SIGMA_MAIN_ANGLE = 2.0f; +constexpr float B_SIGMA_SIDE_BOND = 3.0f, B_SIGMA_SIDE_ANGLE = 4.5f; + +bool IsMainChain(const gemmi::Atom &a) { + static const std::set names = {"N", "CA", "C", "O", "OXT", "P", "OP1", "OP2", "O5'", "C5'", + "C4'", "C3'", "O3'"}; + return names.count(a.name) > 0; +} + +bool IsMetal(const gemmi::Element &el) { + return el.is_metal(); +} + +// The radius van der Waals contacts are made of, without hydrogens: the dictionary's for the atom's +// energy type, or an element's typical value where the residue is not described. +float VdwRadius(const gemmi::MonLib &monlib, const std::string &resname, const gemmi::Atom &a) { + auto cc = monlib.monomers.find(resname); + if (cc != monlib.monomers.end()) { + auto it = cc->second.find_atom(a.name); + if (it != cc->second.atoms.end()) { + auto e = monlib.ener_lib.atoms.find(it->chem_type); + if (e != monlib.ener_lib.atoms.end() && std::isfinite(e->second.vdw_radius)) + return static_cast(e->second.vdw_radius); + } + } + switch (a.element.elem) { + case gemmi::El::C: return 1.70f; + case gemmi::El::N: return 1.55f; + case gemmi::El::O: return 1.52f; + case gemmi::El::S: return 1.80f; + case gemmi::El::P: return 1.80f; + default: return 1.70f; + } +} + +bool CanHydrogenBond(const gemmi::MonLib &monlib, const std::string &resname, const gemmi::Atom &a) { + if (a.element != gemmi::El::N && a.element != gemmi::El::O) + return false; + auto cc = monlib.monomers.find(resname); + if (cc != monlib.monomers.end()) { + auto it = cc->second.find_atom(a.name); + if (it != cc->second.atoms.end()) { + auto e = monlib.ener_lib.atoms.find(it->chem_type); + if (e != monlib.ener_lib.atoms.end()) + return e->second.hb_type == 'D' || e->second.hb_type == 'A' || e->second.hb_type == 'B' || + e->second.hb_type == 'H'; + } + } + return true; // N or O of an undescribed residue +} + +V3 At(const std::vector> &xyzb, int i) { + return {xyzb[i][0], xyzb[i][1], xyzb[i][2]}; +} + +V3 Image(const RestraintRecord &r, V3 p) { + return {r.rot[0] * p.x + r.rot[1] * p.y + r.rot[2] * p.z + r.tran[0], + r.rot[3] * p.x + r.rot[4] * p.y + r.rot[5] * p.z + r.tran[1], + r.rot[6] * p.x + r.rot[7] * p.y + r.rot[8] * p.z + r.tran[2]}; +} + +// The gradient at an image back on the atom it is an image of: rot^T g. +V3 FromImage(const RestraintRecord &r, V3 g) { + return {r.rot[0] * g.x + r.rot[3] * g.y + r.rot[6] * g.z, + r.rot[1] * g.x + r.rot[4] * g.y + r.rot[7] * g.z, + r.rot[2] * g.x + r.rot[5] * g.y + r.rot[8] * g.z}; +} + +} // namespace + +void RestraintSet::Add(RestraintRecord r, int n_atoms_of_r) { + r.slot = slots; + slots += n_atoms_of_r; + records.push_back(r); +} + +void RestraintSet::Append(const RestraintSet &other) { + const int plane_offset = static_cast(plane_atoms.size()); + for (RestraintRecord r : other.records) { + r.slot += slots; + if (r.kind == static_cast(RestraintKind::Plane)) + r.atom[0] += plane_offset; + records.push_back(r); + } + plane_atoms.insert(plane_atoms.end(), other.plane_atoms.begin(), other.plane_atoms.end()); + slots += other.slots; +} + +void RestraintSet::Finalize(size_t n_atoms) { + slot_atom.assign(slots, -1); + for (const RestraintRecord &r : records) { + if (r.kind == static_cast(RestraintKind::Plane)) { + for (int k = 0; k < r.atom[1]; k++) + slot_atom[r.slot + k] = plane_atoms[r.atom[0] + k]; + } else { + const int n = r.kind == static_cast(RestraintKind::Bond) || + r.kind == static_cast(RestraintKind::Repulsion) || + r.kind == static_cast(RestraintKind::BSimilarity) ? 2 + : r.kind == static_cast(RestraintKind::Angle) ? 3 : 4; + for (int k = 0; k < n; k++) + slot_atom[r.slot + k] = r.atom[k]; + } + } + slot_start.assign(n_atoms + 1, 0); + for (int a : slot_atom) + if (a >= 0) + slot_start[a + 1]++; + for (size_t a = 0; a < n_atoms; a++) + slot_start[a + 1] += slot_start[a]; + slot_index.assign(slot_start[n_atoms], 0); + std::vector fill(slot_start.begin(), slot_start.end() - 1); + for (int s = 0; s < slots; s++) + if (slot_atom[s] >= 0) + slot_index[fill[slot_atom[s]]++] = s; +} + +double RestraintValue(const RestraintSet &set, const RestraintRecord &r, const std::vector> &xyzb) { + V3 g[4]; + switch (static_cast(r.kind)) { + case RestraintKind::Bond: + return refine_math::Norm(At(xyzb, r.atom[0]) - At(xyzb, r.atom[1])); + case RestraintKind::Angle: { + const V3 u = At(xyzb, r.atom[0]) - At(xyzb, r.atom[1]), v = At(xyzb, r.atom[2]) - At(xyzb, r.atom[1]); + return std::atan2(refine_math::Norm(refine_math::Cross(u, v)), refine_math::Dot(u, v)) / refine_math::DEG; + } + case RestraintKind::Torsion: + return refine_math::Dihedral(At(xyzb, r.atom[0]), At(xyzb, r.atom[1]), At(xyzb, r.atom[2]), + At(xyzb, r.atom[3]), g) / refine_math::DEG; + case RestraintKind::Chiral: { + const V3 c = At(xyzb, r.atom[0]); + return refine_math::Dot(At(xyzb, r.atom[1]) - c, + refine_math::Cross(At(xyzb, r.atom[2]) - c, At(xyzb, r.atom[3]) - c)); + } + case RestraintKind::Repulsion: + return refine_math::Norm(At(xyzb, r.atom[0]) - Image(r, At(xyzb, r.atom[1]))); + case RestraintKind::BSimilarity: + return xyzb[r.atom[0]][3] - xyzb[r.atom[1]][3]; + case RestraintKind::Plane: + break; + } + (void) set; + return NAN; +} + +double EvaluateRestraints(const RestraintSet &set, const std::vector> &xyzb, + std::vector> *grad) { + double sum = 0; + std::vector> px, pg; + for (const RestraintRecord &r : set.records) { + V3 g[4] = {}; + int n = 0; + double e = 0; + const double ideal = r.ideal, sigma = r.sigma; + switch (static_cast(r.kind)) { + case RestraintKind::Bond: + e = refine_math::BondTerm(At(xyzb, r.atom[0]), At(xyzb, r.atom[1]), ideal, sigma, g); + n = 2; + break; + case RestraintKind::Angle: + e = refine_math::AngleTerm(At(xyzb, r.atom[0]), At(xyzb, r.atom[1]), At(xyzb, r.atom[2]), ideal, sigma, g); + n = 3; + break; + case RestraintKind::Torsion: + e = refine_math::TorsionTerm(At(xyzb, r.atom[0]), At(xyzb, r.atom[1]), At(xyzb, r.atom[2]), + At(xyzb, r.atom[3]), ideal, sigma, r.period, g); + n = 4; + break; + case RestraintKind::Chiral: + e = refine_math::ChiralTerm(At(xyzb, r.atom[0]), At(xyzb, r.atom[1]), At(xyzb, r.atom[2]), + At(xyzb, r.atom[3]), ideal, sigma, g); + n = 4; + break; + case RestraintKind::Repulsion: + e = refine_math::RepulsionTerm(At(xyzb, r.atom[0]), Image(r, At(xyzb, r.atom[1])), ideal, sigma, g); + g[1] = FromImage(r, g[1]); + n = 2; + break; + case RestraintKind::Plane: { + const int m = r.atom[1]; + px.resize(m); + pg.resize(m); + for (int k = 0; k < m; k++) + px[k] = At(xyzb, set.plane_atoms[r.atom[0] + k]); + e = refine_math::PlaneTerm(px.data(), m, sigma, pg.data()); + if (grad != nullptr) + for (int k = 0; k < m; k++) { + std::array &d = (*grad)[set.plane_atoms[r.atom[0] + k]]; + d[0] += pg[k].x; + d[1] += pg[k].y; + d[2] += pg[k].z; + } + break; + } + case RestraintKind::BSimilarity: { + const double rr = (xyzb[r.atom[0]][3] - xyzb[r.atom[1]][3]) / sigma; + e = rr * rr; + if (grad != nullptr) { + (*grad)[r.atom[0]][3] += 2 * rr / sigma; + (*grad)[r.atom[1]][3] -= 2 * rr / sigma; + } + break; + } + } + sum += e; + if (grad != nullptr) + for (int k = 0; k < n; k++) { + std::array &d = (*grad)[r.atom[k]]; + d[0] += g[k].x; + d[1] += g[k].y; + d[2] += g[k].z; + } + } + return sum; +} + +RefinementTopology BuildRefinementTopology(gemmi::Structure &st, bool hydrogens) { + RefinementTopology out; + gemmi::Model &model = st.models.at(0); + out.monlib = std::make_shared(EmbeddedMonomerLibrary()); + gemmi::MonLib &monlib = *out.monlib; + gemmi::setup_entities(st); + gemmi::Logger quiet; // missing atoms and unknown residues are what the fixed residues report + quiet.threshold = 0; + out.topo = gemmi::prepare_topology(st, monlib, 0, + hydrogens ? gemmi::HydrogenChange::ReAddButWater : gemmi::HydrogenChange::NoChange, + /*reorder=*/false, quiet, /*ignore_unknown_links=*/true, /*use_cispeps=*/true); + gemmi::Topo *topo = out.topo.get(); + + // the atoms as they are now, hydrogens added + std::unordered_map index; + std::vector residue_of; + for (const gemmi::Chain &ch : model.chains) + for (const gemmi::Residue &r : ch.residues) + for (const gemmi::Atom &a : r.atoms) { + index[&a] = static_cast(residue_of.size()); + residue_of.push_back(&r); + } + const size_t n = residue_of.size(); + out.refinable.assign(n, 0); + out.vdw_radius.assign(n, 1.7f); + out.hbond.assign(n, 0); + out.metal.assign(n, 0); + out.altloc.assign(n, 0); + out.element.assign(n, gemmi::El::X); + out.rider_parent.assign(n, -1); + for (const gemmi::Topo::Bond &b : topo->bonds) + for (int k = 0; k < 2; k++) + if (b.atoms[k]->is_hydrogen() && !b.atoms[1 - k]->is_hydrogen()) + out.rider_parent[index.at(b.atoms[k])] = index.at(b.atoms[1 - k]); + // A hydrogen gemmi could not place, or with no parent, scatters nothing. + for (gemmi::Chain &ch : model.chains) + for (gemmi::Residue &r : ch.residues) + for (gemmi::Atom &a : r.atoms) + if (a.is_hydrogen() && (out.rider_parent[index.at(&a)] < 0 || !std::isfinite(a.pos.x))) { + a.occ = 0; + out.rider_parent[index.at(&a)] = -1; + if (!std::isfinite(a.pos.x)) + a.pos = r.atoms.front().pos; + } + + // A residue is refined where the dictionary describes it and every one of its atoms. + std::set fixed_names; + for (const gemmi::Topo::ChainInfo &ci : topo->chain_infos) + for (const gemmi::Topo::ResInfo &ri : ci.res_infos) { + const gemmi::Residue &res = *ri.res; + bool known = ri.orig_chemcomp != nullptr && InEmbeddedMonomerLibrary(res.name); + for (const gemmi::Atom &a : res.atoms) + known = known && ri.get_final_chemcomp(a.altloc).has_atom(a.name); + if (known) + out.restrained_residues++; + else { + out.fixed_residues++; + fixed_names.insert(res.name); + } + for (const gemmi::Atom &a : res.atoms) { + const auto it = index.find(&a); + if (it == index.end()) + continue; + const int i = it->second; + out.refinable[i] = known && a.occ > 0 && !a.is_hydrogen(); + out.vdw_radius[i] = VdwRadius(monlib, res.name, a); + out.hbond[i] = res.is_water() || CanHydrogenBond(monlib, res.name, a); + out.metal[i] = res.atoms.size() == 1 && IsMetal(a.element); + out.altloc[i] = a.altloc; + out.element[i] = a.element.elem; + } + } + out.fixed_names.assign(fixed_names.begin(), fixed_names.end()); + + auto idx = [&](const gemmi::Atom *a) { return index.at(a); }; + auto has_h = [](std::initializer_list atoms) { + for (const gemmi::Atom *a : atoms) + if (a->is_hydrogen()) + return true; + return false; + }; + auto any_refinable = [&](std::initializer_list atoms) { + for (int a : atoms) + if (out.refinable[a]) + return true; + return false; + }; + std::vector> excluded(n); + std::vector> bonded(n); + for (size_t i = 0; i < n; i++) + excluded[i].insert(static_cast(i)); + + for (const gemmi::Topo::Bond &b : topo->bonds) { + if (b.asu == gemmi::Asu::Different || has_h({b.atoms[0], b.atoms[1]})) + continue; + const int i = idx(b.atoms[0]), j = idx(b.atoms[1]); + excluded[i].insert(j); + excluded[j].insert(i); + bonded[i].push_back(j); + bonded[j].push_back(i); + if (!any_refinable({i, j}) || !(b.restr->esd > 0)) + continue; + RestraintRecord r; + r.kind = static_cast(RestraintKind::Bond); + r.atom[0] = i; + r.atom[1] = j; + r.ideal = static_cast(b.restr->value); + r.sigma = static_cast(b.restr->esd); + out.geometry.Add(r, 2); + out.bonds++; + RestraintRecord rb; + rb.kind = static_cast(RestraintKind::BSimilarity); + rb.atom[0] = i; + rb.atom[1] = j; + const bool main = IsMainChain(*b.atoms[0]) && IsMainChain(*b.atoms[1]); + rb.sigma = main ? B_SIGMA_MAIN_BOND : B_SIGMA_SIDE_BOND; + out.adp.Add(rb, 2); + } + for (const gemmi::Topo::Angle &an : topo->angles) { + if (has_h({an.atoms[0], an.atoms[1], an.atoms[2]})) + continue; + const int i = idx(an.atoms[0]), j = idx(an.atoms[1]), k = idx(an.atoms[2]); + excluded[i].insert(k); + excluded[k].insert(i); + if (!any_refinable({i, j, k}) || !(an.restr->esd > 0)) + continue; + RestraintRecord r; + r.kind = static_cast(RestraintKind::Angle); + r.atom[0] = i; + r.atom[1] = j; + r.atom[2] = k; + r.ideal = static_cast(an.restr->value); + r.sigma = static_cast(an.restr->esd); + out.geometry.Add(r, 3); + out.angles++; + RestraintRecord rb; + rb.kind = static_cast(RestraintKind::BSimilarity); + rb.atom[0] = i; + rb.atom[1] = k; + const bool main = IsMainChain(*an.atoms[0]) && IsMainChain(*an.atoms[2]); + rb.sigma = main ? B_SIGMA_MAIN_ANGLE : B_SIGMA_SIDE_ANGLE; + out.adp.Add(rb, 2); + } + for (const gemmi::Topo::Torsion &t : topo->torsions) { + const std::string &label = t.restr->label; + if (has_h({t.atoms[0], t.atoms[1], t.atoms[2], t.atoms[3]})) + continue; + if (!(t.restr->esd > 0) || label.rfind("phi", 0) == 0 || label.rfind("psi", 0) == 0 || + label.rfind("C2e", 0) == 0 || label.rfind("C3e", 0) == 0) + continue; + const int a = idx(t.atoms[0]), b = idx(t.atoms[1]), c = idx(t.atoms[2]), d = idx(t.atoms[3]); + if (!any_refinable({a, b, c, d})) + continue; + RestraintRecord r; + r.kind = static_cast(RestraintKind::Torsion); + r.atom[0] = a; + r.atom[1] = b; + r.atom[2] = c; + r.atom[3] = d; + r.ideal = static_cast(t.restr->value); + r.sigma = static_cast(t.restr->esd); + r.period = t.restr->period; + out.geometry.Add(r, 4); + out.torsions++; + } + for (const gemmi::Topo::Chirality &ch : topo->chirs) { + if (ch.restr->sign == gemmi::ChiralityType::Both || has_h({ch.atoms[0], ch.atoms[1], ch.atoms[2], ch.atoms[3]})) + continue; + const int c = idx(ch.atoms[0]), a1 = idx(ch.atoms[1]), a2 = idx(ch.atoms[2]), a3 = idx(ch.atoms[3]); + if (!any_refinable({c, a1, a2, a3})) + continue; + const double volume = topo->ideal_chiral_abs_volume(ch); + RestraintRecord r; + r.kind = static_cast(RestraintKind::Chiral); + r.atom[0] = c; + r.atom[1] = a1; + r.atom[2] = a2; + r.atom[3] = a3; + r.ideal = static_cast(ch.restr->sign == gemmi::ChiralityType::Negative ? -volume : volume); + r.sigma = CHIRAL_SIGMA; + out.geometry.Add(r, 4); + out.chirals++; + } + for (const gemmi::Topo::Plane &p_all : topo->planes) { + gemmi::Topo::Plane p = p_all; + std::erase_if(p.atoms, [](const gemmi::Atom *a) { return a->is_hydrogen(); }); + if (p.atoms.size() < 4 || !(p.restr->esd > 0)) + continue; + bool any = false; + for (const gemmi::Atom *a : p.atoms) + any = any || out.refinable[idx(a)]; + if (!any) + continue; + RestraintRecord r; + r.kind = static_cast(RestraintKind::Plane); + r.atom[0] = static_cast(out.geometry.plane_atoms.size()); + r.atom[1] = static_cast(p.atoms.size()); + r.sigma = static_cast(p.restr->esd); + for (const gemmi::Atom *a : p.atoms) + out.geometry.plane_atoms.push_back(idx(a)); + out.geometry.Add(r, r.atom[1]); + out.planes++; + } + // and those three bonds away, whose distance the torsions and planes govern + std::vector> within3 = excluded; + for (size_t i = 0; i < n; i++) + for (int j : excluded[i]) + for (int k : bonded[j]) + within3[i].insert(k); + out.excluded.resize(n); + for (size_t i = 0; i < n; i++) + out.excluded[i].assign(within3[i].begin(), within3[i].end()); + out.geometry.Finalize(n); + out.adp.Finalize(n); + return out; +} + +RestraintSet BuildRepulsions(const RefinementTopology &topo, const gemmi::Model &model, const gemmi::UnitCell &cell_in, + const gemmi::SpaceGroup &sg, double margin) { + RestraintSet out; + gemmi::UnitCell cell = cell_in; + cell.set_cell_images_from_spacegroup(&sg); + gemmi::Model copy = model; + float max_radius = 0; + for (float r : topo.vdw_radius) + max_radius = std::max(max_radius, r); + const double reach = 2 * max_radius + margin; + gemmi::NeighborSearch ns(copy, cell, std::max(reach, 5.0)); + ns.populate(false); + + std::vector> first(copy.chains.size()); + std::vector pos; + for (size_t c = 0; c < copy.chains.size(); c++) + for (const gemmi::Residue &r : copy.chains[c].residues) { + first[c].push_back(static_cast(pos.size())); + for (const gemmi::Atom &a : r.atoms) + pos.push_back(a.pos); + } + + for (size_t i = 0; i < pos.size(); i++) { + const gemmi::Fractional fi = cell.fractionalize(pos[i]); + ns.for_each(pos[i], '\0', reach, [&](gemmi::NeighborSearch::Mark &m, double) { + const int j = first[m.chain_idx][m.residue_idx] + m.atom_idx; + if (topo.rider_parent[i] >= 0 || topo.rider_parent[j] >= 0 || topo.element[i] == gemmi::El::H || + topo.element[j] == gemmi::El::H) + return; + if (!topo.refinable[i] && !topo.refinable[j]) + return; + if (m.image_idx == 0 && (static_cast(i) >= j || + std::binary_search(topo.excluded[i].begin(), topo.excluded[i].end(), j))) + return; + if (m.image_idx != 0 && static_cast(i) > j) + return; + if (topo.altloc[i] && topo.altloc[j] && topo.altloc[i] != topo.altloc[j]) + return; + // The image of j nearest to i, as a Cartesian x -> R x + t. + const gemmi::FTransform ft = ns.get_image_transformation(m.image_idx); + const gemmi::Fractional fj = ft.apply(cell.fractionalize(pos[j])); + const gemmi::Vec3 shift(std::round(fi.x - fj.x), std::round(fi.y - fj.y), std::round(fi.z - fj.z)); + if (m.image_idx == 0 && (shift.x != 0 || shift.y != 0 || shift.z != 0)) + return; // the same copy's own lattice translate: the identity term covers the contact + const gemmi::Mat33 rot = cell.orth.mat.multiply(ft.mat).multiply(cell.frac.mat); + const gemmi::Vec3 tran = cell.orth.mat.multiply(ft.vec + shift); + const gemmi::Position image(rot.multiply(pos[j]) + tran); + const double d = pos[i].dist(image); + + float d0 = topo.vdw_radius[i] + topo.vdw_radius[j]; + if (topo.metal[i] || topo.metal[j]) { + const gemmi::El other = topo.element[topo.metal[i] ? j : static_cast(i)]; + if (other == gemmi::El::N || other == gemmi::El::O || other == gemmi::El::S) + d0 = METAL_CONTACT; + } else if (topo.hbond[i] && topo.hbond[j]) + d0 = HBOND_CONTACT; + if (!(d < d0 + margin)) + return; + RestraintRecord r; + r.kind = static_cast(RestraintKind::Repulsion); + r.atom[0] = static_cast(i); + r.atom[1] = j; + r.ideal = d0; + r.sigma = REPULSION_SIGMA; + for (int a = 0; a < 3; a++) { + for (int b = 0; b < 3; b++) + r.rot[3 * a + b] = static_cast(rot[a][b]); + r.tran[a] = static_cast(tran.at(a)); + } + out.Add(r, 2); + }); + } + out.Finalize(pos.size()); + return out; +} + +void PlaceRidingHydrogens(RefinementTopology &topo) { + gemmi::place_hydrogens_on_all_atoms(*topo.topo); +} + +GeometryStatistics MeasureGeometry(const RestraintSet &geometry, const RestraintSet &repulsions, + const std::vector> &xyzb) { + GeometryStatistics s; + double sb = 0, sa = 0, sc = 0; + int nc = 0; + for (const RestraintRecord &r : geometry.records) { + const RestraintKind kind = static_cast(r.kind); + if (kind == RestraintKind::Bond) { + const double d = RestraintValue(geometry, r, xyzb) - r.ideal; + sb += d * d; + s.bonds++; + } else if (kind == RestraintKind::Angle) { + const double d = RestraintValue(geometry, r, xyzb) - r.ideal; + sa += d * d; + s.angles++; + } else if (kind == RestraintKind::Chiral) { + const double d = RestraintValue(geometry, r, xyzb) - r.ideal; + sc += d * d; + nc++; + } + } + for (const RestraintRecord &r : repulsions.records) + if (RestraintValue(repulsions, r, xyzb) < r.ideal - 0.4) + s.clashes++; + s.rmsd_bonds = s.bonds > 0 ? std::sqrt(sb / s.bonds) : 0; + s.rmsd_angles = s.angles > 0 ? std::sqrt(sa / s.angles) : 0; + s.rmsd_chirals = nc > 0 ? std::sqrt(sc / nc) : 0; + return s; +} diff --git a/image_analysis/structure_refinement/RefinementRestraints.h b/image_analysis/structure_refinement/RefinementRestraints.h new file mode 100644 index 000000000..0fcb13761 --- /dev/null +++ b/image_analysis/structure_refinement/RefinementRestraints.h @@ -0,0 +1,98 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#pragma once + +#include +#include +#include + +#include + +#include "gemmi/model.hpp" +#include "gemmi/monlib.hpp" +#include "gemmi/topo.hpp" +#include "gemmi/unitcell.hpp" + +#include "RefinementMath.h" // RestraintRecord + +// The geometry side of restrained refinement (Refinement.h): what the dictionary says each refinable atom's +// surroundings should look like, as flat lists of restraints the CPU and the GPU evaluate alike +// (RefinementMath.h). Atoms are numbered in model order, as ModelPositions() lists them. + +// A list of restraints and how their gradients reach the atoms: each restraint writes one gradient slot per +// atom it involves, and atom a sums the slots slot_index[slot_start[a] ... slot_start[a + 1]) in that order - +// a gather, so that a GPU, summing in the same order, needs no atomics and gives the same bits every run. +struct RestraintSet { + std::vector records; + std::vector plane_atoms; + int slots = 0; + std::vector slot_start; // n_atoms + 1 + std::vector slot_index; + + void Add(RestraintRecord r, int n_atoms_of_r); + void Append(const RestraintSet &other); // other's records after these; Finalize() again afterwards + void Finalize(size_t n_atoms); + // Which atom each slot belongs to (Finalize() works it out); -1 for none. + std::vector slot_atom; +}; + +// The value of one restraint's own quantity (a distance, an angle, a volume ...) at `xyzb`, for statistics. +double RestraintValue(const RestraintSet &set, const RestraintRecord &r, const std::vector> &xyzb); + +// The sum of E over the set at `xyzb` (x, y, z, B per atom) and, with `grad`, dE/d(x, y, z, B) per atom (added +// to it). On the CPU, in double. +double EvaluateRestraints(const RestraintSet &set, const std::vector> &xyzb, + std::vector> *grad); + +// What refinement knows of a model's chemistry. +struct RefinementTopology { + RestraintSet geometry; // bonds, angles, torsions, chiral volumes, planes + RestraintSet adp; // B of bonded and angle-related atoms alike + std::vector refinable; // per atom: its residue is described by the dictionary + // Per atom: the atoms the repulsion leaves out - itself and those up to three bonds away (sorted). + std::vector> excluded; + std::vector vdw_radius; // A + std::vector hbond; // donor, acceptor or both (water included) + std::vector metal; // a metal ion + std::vector altloc; + std::vector element; + // Per atom: the heavy atom a hydrogen rides on, -1 for every other atom. + std::vector rider_parent; + // The topology itself and the dictionary it points into, for placing the hydrogens again. + std::shared_ptr monlib; + std::shared_ptr topo; + int restrained_residues = 0; + int fixed_residues = 0; + std::vector fixed_names; // the residue names held fixed, each once + int bonds = 0, angles = 0, torsions = 0, chirals = 0, planes = 0; +}; + +// The topology of the first model of `st` (hydrogens already removed) from the embedded monomer library +// (MonomerLibrary.h), through gemmi's prepare_topology(): links between residues (peptide, cis-peptide, +// nucleic-acid backbone, disulfides, the model's LINK records the library describes) and their +// modifications included. A residue the library does not describe, or with an atom it does not know, is +// held fixed: it stays in the structure factors and the bulk-solvent mask, and it is restrained against. +// Torsions: those with an esd, except the main-chain phi/psi and the sugar puckers given as alternatives. +// With `hydrogens`, the dictionary's hydrogens are added to the residues it describes (waters excepted) and +// placed; they ride on their parent atoms and take part in no restraint. +RefinementTopology BuildRefinementTopology(gemmi::Structure &st, bool hydrogens = false); + +// The riding hydrogens placed again on their parents as they are now (gemmi's place_hydrogens_on_all_atoms). +void PlaceRidingHydrogens(RefinementTopology &topo); + +// The repulsions of the model as it is now: every pair of atoms, symmetry mates included, closer than its +// contact distance plus `margin` and not excluded, where at least one of the two is refinable and the two are +// not alternative conformations of each other. Contact distances: van der Waals radii summed, 2.6 A where +// both can hydrogen-bond, 1.9 A between a metal and N, O or S. +RestraintSet BuildRepulsions(const RefinementTopology &topo, const gemmi::Model &model, const gemmi::UnitCell &cell, + const gemmi::SpaceGroup &sg, double margin); + +// The usual summary of a model's geometry. +struct GeometryStatistics { + double rmsd_bonds = 0, rmsd_angles = 0, rmsd_chirals = 0; + int bonds = 0, angles = 0; + int clashes = 0; // repulsions overlapping by more than 0.4 A +}; +GeometryStatistics MeasureGeometry(const RestraintSet &geometry, const RestraintSet &repulsions, + const std::vector> &xyzb); diff --git a/image_analysis/structure_refinement/RefinementXray.cpp b/image_analysis/structure_refinement/RefinementXray.cpp new file mode 100644 index 000000000..5a3da3f4d --- /dev/null +++ b/image_analysis/structure_refinement/RefinementXray.cpp @@ -0,0 +1,396 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#include "RefinementXray.h" + +#include +#include +#include + +#include + +#include "gemmi/dencalc.hpp" +#include "gemmi/it92.hpp" +#include "gemmi/math.hpp" + +#include "../../common/ParallelFor.h" +#include "ModelFFT.h" +#include "ModelGrid.h" +#include "RefinementMath.h" + +namespace { + +using Table = gemmi::IT92; + +// The free reflections a likelihood bin is fitted on, at least. +constexpr size_t FREE_PER_BIN = 100; +constexpr size_t MAX_BINS = 30; + +// The density calculator the refinement grids atoms with: DensityCalculator's settings for the grid of +// RigidBodyZoneGrid(), and the given blur. +gemmi::DensityCalculator Calculator(const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg, + double d_min, double blur) { + gemmi::DensityCalculator dc; + dc.d_min = d_min; + dc.rate = 1.5; // RigidBodyRefine.cpp's GRID_RATE, which RigidBodyZoneGrid sizes the grid with + dc.grid.unit_cell = cell; + dc.grid.spacegroup = &sg; + dc.blur = blur; + return dc; +} + +// -ln L summed over rows `idx` at D and beta. +double BinNegLogLikelihood(const RefinementReflections &refl, const std::vector &fc, + const std::vector &idx, double d, double beta) { + double sum = 0, dfc = 0; + for (size_t i : idx) { + const double var = refl.epsilon[i] * beta + double(refl.sigma[i]) * refl.sigma[i]; + sum += refine_math::AmplitudeNegLogLikelihood(refl.fo[i], fc[i], d, var, refl.centric[i], dfc); + } + return sum; +} + +// The minimum of f over [lo, hi] by golden-section search. +template +double GoldenMinimum(F f, double lo, double hi, int iterations) { + const double g = 0.5 * (std::sqrt(5.0) - 1.0); + double a = lo, b = hi; + double c = b - g * (b - a), d = a + g * (b - a); + double fc = f(c), fd = f(d); + for (int i = 0; i < iterations; i++) { + if (fc < fd) { + b = d; + d = c; + fd = fc; + c = b - g * (b - a); + fc = f(c); + } else { + a = c; + c = d; + fc = fd; + d = a + g * (b - a); + fd = f(d); + } + } + return 0.5 * (a + b); +} + +class RefinementEngineCPU : public RefinementEngine { +public: + RefinementEngineCPU(const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg, double d_min, + const RefinementReflections &refl, size_t nthreads) + : cell_(cell), sg_(sg), d_min_(d_min), refl_(refl), nthreads_(nthreads), + grid_(RigidBodyZoneGrid(cell, sg, d_min)), composition_(grid_, d_min, refl.hkl), + orbit_leaders_(OrbitLeaders(grid_, nthreads)) {} + + void SetBlur(double blur) override { blur_ = blur; } + void SetRowModel(const XrayRowModel &rows) override { rows_ = rows; } + bool OnGPU() const override { return false; } + + void Fcalc(const gemmi::Model &model, std::vector> &f) override { + std::vector> fd; + Compose(model, fd); + f.assign(fd.begin(), fd.end()); + } + + void Fmask(const gemmi::Model &model, std::vector> &f) override { + gemmi::Grid mask = grid_; + PutMaskOnGrid(mask, model, orbit_leaders_, nthreads_); + const gemmi::FPhiGrid fm = MapToFPhi(mask); + f.clear(); + for (const gemmi::Miller &h : refl_.hkl) + f.push_back(fm.get_value_by_hkl(h)); + } + + double Evaluate(const gemmi::Model &model, std::vector> *grad) override { + std::vector> f; + Compose(model, f); + const size_t n = f.size(); + // The target and, per reflection, dT/dF of the atoms' structure factor (dT/dRe F + i dT/dIm F). + std::vector> dt_df(n, 0.0); + double sum = 0; + for (size_t m = 0; m < n; m++) { + if (refl_.free[m]) + continue; + const double k = rows_.k_total[m]; + const std::complex fm = k * (f[m] + std::complex(rows_.fmask[m])); + const double fc = std::abs(fm); + double dfc = 0; + sum += refine_math::AmplitudeNegLogLikelihood(refl_.fo[m], fc, rows_.d[m], rows_.var[m], + refl_.centric[m], dfc); + if (fc > 0) + dt_df[m] = k * dfc * fm / fc; + } + if (grad == nullptr) + return sum; + + // The FFT method's map, dT/drho(x) = (V/N) Re sum_m sum_ops conj(dT/dF_m) scale_m exp(2 pi i h.t) + // exp(2 pi i (hR).x), as the inverse real transform of a half-l coefficient grid: a term A at index k + // goes in as A/2 at k, or as conj(A)/2 at -k where k is in the half not stored; on the planes l = 0 and + // l = nw/2, which hold both, as both. + const int nu = grid_.nu, nv = grid_.nv, nw = grid_.nw, nwh = nw / 2 + 1; + const size_t ncomplex = static_cast(nu) * nv * nwh; + std::complex *in = reinterpret_cast *>(fftwf_alloc_complex(ncomplex)); + float *out = fftwf_alloc_real(static_cast(nu) * nv * nw); + std::fill(in, in + ncomplex, std::complex(0, 0)); + const std::vector &terms = composition_.Terms(); + const size_t ops = composition_.Ops(); + auto add = [&](int ku, int kv, int kw, std::complex value) { + in[gemmi::modulo(ku, nu) + static_cast(nu) * (gemmi::modulo(kv, nv) + static_cast(nv) * kw)] += + std::complex(value); + }; + for (size_t m = 0; m < n; m++) { + if (dt_df[m] == 0.0) + continue; + const double scale = composition_.Centring() * std::exp(blur_ * 0.25 * composition_.InvD2()[m]); + const std::complex c = std::conj(dt_df[m]) * scale; + for (size_t o = 0; o < ops; o++) { + const SymmetryComposition::Term &t = terms[m * ops + o]; + const std::complex a = 0.5 * c * t.phase; + const gemmi::Miller &k = t.k; + if (k[2] > 0 && 2 * k[2] != nw) + add(k[0], k[1], k[2], a); + else if (k[2] < 0 && -2 * k[2] != nw) + add(-k[0], -k[1], -k[2], std::conj(a)); + else { + const int l = k[2] < 0 ? -k[2] : k[2]; + add(k[0], k[1], l, a); + add(-k[0], -k[1], l, std::conj(a)); + } + } + } + MapFromFPhiExecute(nu, nv, nw, in, out); + gemmi::Grid map = grid_; + const float norm = static_cast(cell_.volume / static_cast(map.point_count())); + for (size_t i = 0; i < map.data.size(); i++) + map.data[i] = out[i] * norm; + fftwf_free(in); + fftwf_free(out); + AtomGradientsFromMap(map, model, blur_, nthreads_, *grad); + return sum; + } + +private: + 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 + composition_.Compose(MapToFPhi(dc.grid), blur_, f, nullptr, nthreads_); + } + + gemmi::UnitCell cell_; + const gemmi::SpaceGroup &sg_; + double d_min_; + const RefinementReflections &refl_; + size_t nthreads_; + gemmi::Grid grid_; + SymmetryComposition composition_; + std::vector orbit_leaders_; + double blur_ = 0; + XrayRowModel rows_; +}; + +} // namespace + +RefinementReflections PrepareRefinementReflections(const std::vector &merged, + const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg, + double d_min) { + const gemmi::GroupOps gops = sg.operations(); + const gemmi::ReciprocalAsu asu(&sg); + // ASU, last one wins where a reflection is listed twice (both hands of an anomalous merge) + std::map by_hkl; + for (const MergedReflection &r : merged) { + if (!std::isfinite(r.F) || r.F < 0) + continue; + gemmi::Miller h{{r.h, r.k, r.l}}; + if (!asu.is_in(h)) + h = asu.to_asu(h, gops).first; + by_hkl[h] = &r; + } + std::vector hkl; + for (const auto &[h, r] : by_hkl) + hkl.push_back(h); + const SymmetryComposition composition(RigidBodyZoneGrid(cell, sg, d_min), d_min, hkl); + RefinementReflections out; + for (size_t i = 0; i < hkl.size(); i++) { + if (composition.Row()[i] < 0) + continue; + const MergedReflection &r = *by_hkl[hkl[i]]; + out.hkl.push_back(hkl[i]); + out.fo.push_back(r.F); + out.sigma.push_back(std::isfinite(r.sigmaF) ? r.sigmaF : 0.0f); + out.free.push_back(r.rfree_flag); + out.centric.push_back(gops.is_reflection_centric(hkl[i])); + out.epsilon.push_back(static_cast(gops.epsilon_factor(hkl[i]))); + out.inv_d2.push_back(cell.calculate_1_d2(hkl[i])); + } + return out; +} + +void EstimateLikelihoodParameters(const RefinementReflections &refl, const std::vector> &fmodel, + XrayRowModel &rows) { + const size_t n = refl.Size(); + std::vector fc(n); + for (size_t i = 0; i < n; i++) + fc[i] = std::abs(fmodel[i]); + std::vector free; + for (size_t i = 0; i < n; i++) + if (refl.free[i] && fc[i] > 0) + free.push_back(i); + // Too few free reflections to bin: the working set, which says the model is better than it is, but + // still gives the likelihood a scale. + if (free.size() < FREE_PER_BIN) + for (size_t i = 0; i < n; i++) + if (!refl.free[i] && fc[i] > 0) + free.push_back(i); + std::sort(free.begin(), free.end(), [&](size_t a, size_t b) { return refl.inv_d2[a] < refl.inv_d2[b]; }); + const size_t bins = std::clamp(free.size() / FREE_PER_BIN, 1, MAX_BINS); + + std::vector centre(bins), bin_d(bins), bin_beta(bins); + for (size_t b = 0; b < bins; b++) { + const std::vector idx(free.begin() + b * free.size() / bins, free.begin() + (b + 1) * free.size() / bins); + double s2 = 0, sfo = 0, sfc2 = 0, sfofc = 0; + for (size_t i : idx) { + s2 += refl.inv_d2[i]; + sfo += double(refl.fo[i]) * refl.fo[i] / refl.epsilon[i]; + sfc2 += fc[i] * fc[i] / refl.epsilon[i]; + sfofc += refl.fo[i] * fc[i] / refl.epsilon[i]; + } + centre[b] = s2 / static_cast(idx.size()); + double d = sfc2 > 0 ? std::max(sfofc / sfc2, 0.05) : 1.0; + double beta = std::max((sfo - d * d * sfc2) / static_cast(idx.size()), 1e-3 * sfo / idx.size()); + // Alternating one-dimensional minimisations of -ln L over ln D and ln beta. + for (int round = 0; round < 4; round++) { + const double ld = GoldenMinimum([&](double x) { return BinNegLogLikelihood(refl, fc, idx, std::exp(x), beta); }, + std::log(d) - 2.0, std::log(std::min(d * 7.4, 3.0)), 40); + d = std::exp(ld); + const double lb = GoldenMinimum([&](double x) { return BinNegLogLikelihood(refl, fc, idx, d, std::exp(x)); }, + std::log(beta) - 3.0, std::log(beta) + 3.0, 40); + beta = std::exp(lb); + } + bin_d[b] = d; + bin_beta[b] = beta; + } + rows.d.resize(n); + rows.var.resize(n); + for (size_t i = 0; i < n; i++) { + // linear in 1/d^2 between the bins' centres, flat beyond the outermost + const double s = refl.inv_d2[i]; + double d = bin_d[0], beta = bin_beta[0]; + if (s >= centre[bins - 1]) { + d = bin_d[bins - 1]; + beta = bin_beta[bins - 1]; + } else if (s > centre[0]) { + const size_t b = std::upper_bound(centre.begin(), centre.end(), s) - centre.begin(); + const double t = (s - centre[b - 1]) / (centre[b] - centre[b - 1]); + d = bin_d[b - 1] + t * (bin_d[b] - bin_d[b - 1]); + beta = bin_beta[b - 1] + t * (bin_beta[b] - bin_beta[b - 1]); + } + rows.d[i] = static_cast(d); + rows.var[i] = static_cast(refl.epsilon[i] * beta + double(refl.sigma[i]) * refl.sigma[i]); + } +} + +RFactors ComputeRFactors(const RefinementReflections &refl, const std::vector> &fmodel) { + double nw = 0, dw = 0, nf = 0, df = 0; + RFactors r; + for (size_t i = 0; i < refl.Size(); i++) { + const double diff = std::fabs(refl.fo[i] - std::abs(fmodel[i])); + if (refl.free[i]) { + nf += diff; + df += refl.fo[i]; + r.n_free++; + } else { + nw += diff; + dw += refl.fo[i]; + r.n_work++; + } + } + r.r_work = dw > 0 ? nw / dw : 1; + r.r_free = df > 0 ? nf / df : 1; + return r; +} + +double RefinementBlur(const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg, double d_min, double b_floor) { + // set_refmac_compatible_blur(): u_to_b / 1.1 * spacing^2 - b_min, not below 0, the spacing the one asked for + const gemmi::DensityCalculator dc = Calculator(cell, sg, d_min, 0); + const double spacing = dc.requested_grid_spacing(); + return std::max(0.0, gemmi::u_to_b() / 1.1 * spacing * spacing - b_floor); +} + +std::unique_ptr MakeRefinementEngineCPU(const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg, + double d_min, const RefinementReflections &refl, + size_t nthreads) { + return std::make_unique(cell, sg, d_min, refl, nthreads); +} + +void AtomGradientsFromMap(const gemmi::Grid &map, const gemmi::Model &model, double blur, size_t nthreads, + std::vector> &grad) { + const gemmi::DensityCalculator dc = Calculator(map.unit_cell, *map.spacegroup, 1.0, blur); + std::vector atoms; + for (const gemmi::Chain &ch : model.chains) + for (const gemmi::Residue &r : ch.residues) + for (const gemmi::Atom &a : r.atoms) + atoms.push_back(&a); + grad.assign(atoms.size(), {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; + 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); + constexpr double FOUR_PI2 = 4 * refine_math::PI * refine_math::PI; + ParallelChunks(static_cast(atoms.size()), nthreads, [&](int lo, int hi) { + for (int i = lo; i < hi; i++) { + const gemmi::Atom &atom = *atoms[i]; + 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); + const float max_r2 = static_cast(radius * radius); + 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 gx = 0, gy = 0, gz = 0, gb = 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; + const float r2 = yz + x * x; + if (r2 > max_r2) + continue; + const double value = map.data[row + gemmi::modulo(u, nu)]; + double dr2 = 0, db = 0; + for (size_t q = 0; q < std::size(precal.a); q++) { + const double e = precal.a[q] * gemmi::unsafe_expapprox(std::max(precal.b[q] * r2, -88.f)); + dr2 += e * precal.b[q]; + db += e * precal.b[q] * (1.5 + precal.b[q] * r2); + } + gx += value * dr2 * x; + gy += value * dr2 * y; + gz += value * dr2 * z; + gb += value * db; + } + } + } + // rho = occ sum_q a_q exp(b_q r^2): d rho / d(centre) = occ 2 (sum a_q b_q e_q) (centre - point), and + // with b_q = -4 pi^2 / (B_q + B + blur), d rho / dB = occ sum a_q e_q b_q (1.5 + b_q r^2) / (4 pi^2). + grad[i] = {2.0 * atom.occ * gx, 2.0 * atom.occ * gy, 2.0 * atom.occ * gz, atom.occ * gb / FOUR_PI2}; + } + }); +} diff --git a/image_analysis/structure_refinement/RefinementXray.h b/image_analysis/structure_refinement/RefinementXray.h new file mode 100644 index 000000000..9bbe976cd --- /dev/null +++ b/image_analysis/structure_refinement/RefinementXray.h @@ -0,0 +1,117 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#pragma once + +#include +#include +#include +#include +#include + +#include "gemmi/grid.hpp" +#include "gemmi/model.hpp" +#include "gemmi/symmetry.hpp" +#include "gemmi/unitcell.hpp" + +#include "../../common/Reflection.h" // MergedReflection +#include "RefinementRestraints.h" // RestraintSet +#include "RigidBodyRefine.h" // SymmetryComposition + +// The X-ray side of restrained refinement (Refinement.h): the observed amplitudes, the model's structure +// factors and the maximum-likelihood target that compares them, with its gradient with respect to every +// atom's position and B. + +// The reflections a refinement fits and scores: the observed amplitudes, in the group's reciprocal ASU, that +// the refinement grid composes (SymmetryComposition) - every observed one in practice. +struct RefinementReflections { + std::vector hkl; + std::vector fo, sigma; + std::vector free, centric; + std::vector epsilon; + std::vector inv_d2; + size_t Size() const { return hkl.size(); } +}; + +// The French-Wilson amplitudes of `merged` (NaN dropped) into the ASU of `sg`, those the grid of +// RigidBodyZoneGrid(cell, sg, d_min) composes. +RefinementReflections PrepareRefinementReflections(const std::vector &merged, + const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg, + double d_min); + +// What turns the atoms' structure factors F into the model the target compares with the data, per reflection: +// F_model = k_total (F + fmask), +// with k_total the overall scale and anisotropic B and fmask the bulk solvent already scaled by k_sol +// exp(-b_sol s^2/4); and the two parameters of the likelihood, D and the variance. Fixed for one macro-cycle. +struct XrayRowModel { + std::vector k_total; + std::vector> fmask; + std::vector d, var; +}; + +// The ML parameters per reflection from the free reflections: D and beta fitted by maximum likelihood in +// resolution bins of the free set and interpolated between the bins' centres; the variance of a reflection is +// epsilon beta + sigma^2. `fmodel` is k_total (F + fmask), without D. +// Following Lunin & Skovoroda (1995) Acta Cryst. A51, 880-887, and Afonine et al. (2012) Acta Cryst. D68, 352-367 +void EstimateLikelihoodParameters(const RefinementReflections &refl, const std::vector> &fmodel, + XrayRowModel &rows); + +// R-work and R-free of k_total (F + fmask) against the observations. +struct RFactors { + double r_work = 1, r_free = 1; + int n_work = 0, n_free = 0; +}; +RFactors ComputeRFactors(const RefinementReflections &refl, const std::vector> &fmodel); + +// 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. +class RefinementEngine { +public: + virtual ~RefinementEngine() = default; + // The blur the atoms are gridded with (DensityCalculator's), and the matching unblur of their transform. + // Every atom's B must stay above what the blur was chosen for (RefinementBlur) while it is in use. + virtual void SetBlur(double blur) = 0; + virtual void SetRowModel(const XrayRowModel &rows) = 0; + // F of the atoms at the reflections. + virtual void Fcalc(const gemmi::Model &model, std::vector> &f) = 0; + // The bulk-solvent mask's structure factors at the reflections (unscaled). + 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. + virtual double Evaluate(const gemmi::Model &model, std::vector> *grad) = 0; + virtual bool OnGPU() const = 0; + + // 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; } + virtual double EvaluateRestraints(int which, const std::vector> &xyzb, + std::vector> *grad) { + if (grad != nullptr) + grad->assign(xyzb.size(), {0, 0, 0, 0}); + return ::EvaluateRestraints(restraints_[which], xyzb, grad); + } + +private: + RestraintSet restraints_[2]; +}; + +// The blur for a model whose smallest B may go as low as b_floor: DensityCalculator's +// set_refmac_compatible_blur() for that B, on the refinement grid of d_min. +double RefinementBlur(const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg, double d_min, double b_floor); + +// The CPU engine: the density of one copy of the atoms on the grid of RigidBodyZoneGrid(cell, sg, d_min) +// (PutModelDensityOnGrid), its transform and the crystal's symmetry composed in reciprocal space +// (SymmetryComposition). The gradient by the FFT method: the derivative of the target with respect to each +// structure factor, taken back to real space by one inverse transform, is a map; each atom's derivatives are +// that map summed against the derivatives of its own density. +// Following Agarwal (1978) Acta Cryst. A34, 791-809 +std::unique_ptr MakeRefinementEngineCPU(const gemmi::UnitCell &cell, + const gemmi::SpaceGroup &sg, double d_min, + const RefinementReflections &refl, size_t nthreads); + +// 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. +void AtomGradientsFromMap(const gemmi::Grid &map, const gemmi::Model &model, double blur, size_t nthreads, + std::vector> &grad); diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index cecfe9c5a..36e005beb 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -184,6 +184,7 @@ ADD_EXECUTABLE(jfjoch_test MolecularReplacementTest.cpp ModelMaskGPUTest.cpp ModelDensityGPUTest.cpp + RefinementTest.cpp SearchSpaceGroupTest.cpp SearchSpaceGroupTwinTest.cpp SyntheticMergedReflections.h diff --git a/tests/RefinementTest.cpp b/tests/RefinementTest.cpp new file mode 100644 index 000000000..ace1d7659 --- /dev/null +++ b/tests/RefinementTest.cpp @@ -0,0 +1,452 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#include + +#include +#include +#include +#include + +#include + +#include "../common/CUDAWrapper.h" +#include "../common/Logger.h" +#include "../image_analysis/structure_refinement/MonomerLibrary.h" +#include "../image_analysis/structure_refinement/Refinement.h" +#include "../image_analysis/structure_refinement/RefinementMath.h" +#include "../image_analysis/structure_refinement/RefinementRestraints.h" +#include "../image_analysis/structure_refinement/RefinementXray.h" +#ifdef JFJOCH_USE_CUDA +#include "../image_analysis/structure_refinement/RefinementGPU.h" +#endif + +namespace { + +// Copies of residue `code` at the embedded dictionary's ideal coordinates (hydrogens left out), each turned and +// moved to a place of its own in the cell. +gemmi::Structure residue_crystal(const std::vector &codes, const gemmi::UnitCell &cell, const char *group, + unsigned seed) { + std::mt19937 rng(seed); + std::uniform_real_distribution frac(0.0, 1.0), angle(-3.1, 3.1); + gemmi::Structure st; + st.cell = cell; + st.spacegroup_hm = group; + st.setup_cell_images(); + gemmi::Model model(1); + gemmi::Chain chain("A"); + int seq = 1; + for (const std::string &code : codes) { + const gemmi::ChemComp &cc = EmbeddedMonomerLibrary().monomers.at(code); + gemmi::Residue res; + res.name = code; + res.seqid = gemmi::SeqId(seq++, ' '); + res.het_flag = 'H'; + res.entity_type = gemmi::EntityType::NonPolymer; + const gemmi::Position centre = cell.orthogonalize(gemmi::Fractional(frac(rng), frac(rng), frac(rng))); + const double ax = angle(rng), ay = angle(rng), az = angle(rng); + const gemmi::Mat33 rot = gemmi::Mat33(1, 0, 0, 0, std::cos(ax), -std::sin(ax), 0, std::sin(ax), std::cos(ax)) + .multiply(gemmi::Mat33(std::cos(ay), 0, std::sin(ay), 0, 1, 0, -std::sin(ay), 0, std::cos(ay))) + .multiply(gemmi::Mat33(std::cos(az), -std::sin(az), 0, std::sin(az), std::cos(az), 0, 0, 0, 1)); + for (const gemmi::ChemComp::Atom &ca : cc.atoms) { + if (ca.is_hydrogen()) + continue; + gemmi::Atom a; + a.name = ca.id; + a.element = ca.el; + a.pos = gemmi::Position(rot.multiply(ca.xyz)) + centre; + a.b_iso = 20.0f; + a.occ = 1.0f; + res.atoms.push_back(a); + } + chain.residues.push_back(res); + } + model.chains.push_back(chain); + st.models.push_back(model); + return st; +} + +// Atoms of a few elements scattered over the cell. +gemmi::Structure random_crystal(const gemmi::UnitCell &cell, const char *group, int n, unsigned seed) { + std::mt19937 rng(seed); + std::uniform_real_distribution frac(0.0, 1.0), b(10.0, 40.0); + const gemmi::El elements[] = {gemmi::El::C, gemmi::El::N, gemmi::El::O, gemmi::El::S}; + gemmi::Structure st; + st.cell = cell; + st.spacegroup_hm = group; + st.setup_cell_images(); + gemmi::Model model(1); + gemmi::Chain chain("A"); + gemmi::Residue res; + res.name = "UNK"; + res.seqid = gemmi::SeqId(1, ' '); + for (int i = 0; i < n; i++) { + gemmi::Atom a; + a.name = "X"; + a.element = elements[i % 4]; + a.pos = cell.orthogonalize(gemmi::Fractional(frac(rng), frac(rng), frac(rng))); + a.b_iso = static_cast(b(rng)); + a.occ = 1.0f; + res.atoms.push_back(a); + } + chain.residues.push_back(res); + model.chains.push_back(chain); + st.models.push_back(model); + return st; +} + +// Every reflection of the group's ASU to d_min, amplitude 1, one in twenty free. +std::vector asu_reflections(const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg, double d_min) { + std::vector out; + const gemmi::ReciprocalAsu asu(&sg); + const gemmi::GroupOps gops = sg.operations(); + const gemmi::Miller lim = cell.get_hkl_limits(d_min); + int i = 0; + for (int h = -lim[0]; h <= lim[0]; h++) + for (int k = -lim[1]; k <= lim[1]; k++) + for (int l = -lim[2]; l <= lim[2]; l++) { + const gemmi::Miller m{{h, k, l}}; + if ((h == 0 && k == 0 && l == 0) || !asu.is_in(m) || gops.is_systematically_absent(m)) + continue; + const double d = cell.calculate_d(m); + if (d < d_min) + continue; + MergedReflection r; + r.h = h; + r.k = k; + r.l = l; + r.d = static_cast(d); + r.F = 1.0f; + r.sigmaF = 0.05f; + r.rfree_flag = (i++ % 20) == 0; + out.push_back(r); + } + return out; +} + +// The observations set to |F| of `truth`, and a row model with no scale, no solvent and the likelihood's +// parameters estimated against `model`. +void set_observations(RefinementEngine &engine, RefinementReflections &refl, const gemmi::Model &truth, + const gemmi::Model &model) { + std::vector> f; + engine.Fcalc(truth, f); + for (size_t m = 0; m < refl.Size(); m++) + refl.fo[m] = std::abs(f[m]); + engine.Fcalc(model, f); + XrayRowModel rows; + rows.k_total.assign(refl.Size(), 1.0f); + rows.fmask.assign(refl.Size(), {0, 0}); + EstimateLikelihoodParameters(refl, f, rows); + engine.SetRowModel(rows); +} + +gemmi::Model shaken(const gemmi::Model &model, double sigma, double b_sigma, unsigned seed) { + std::mt19937 rng(seed); + std::normal_distribution d(0, sigma), db(0, b_sigma); + gemmi::Model out = model; + for (gemmi::Chain &ch : out.chains) + for (gemmi::Residue &r : ch.residues) + for (gemmi::Atom &a : r.atoms) { + a.pos += gemmi::Position(d(rng), d(rng), d(rng)); + a.b_iso = static_cast(std::max(5.0, a.b_iso + db(rng))); + } + return out; +} + +std::vector atoms_of(gemmi::Model &model) { + std::vector out; + for (gemmi::Chain &ch : model.chains) + for (gemmi::Residue &r : ch.residues) + for (gemmi::Atom &a : r.atoms) + out.push_back(&a); + return out; +} + +// The engine's gradient of a few atoms against central differences of its own target. +void check_xray_gradient(RefinementEngine &engine, gemmi::Model model, double tolerance) { + std::vector> grad; + engine.Evaluate(model, &grad); + std::vector atoms = atoms_of(model); + const double hx = 2e-3, hb = 2e-2; + double max_g = 0; + for (const auto &g : grad) + for (int c = 0; c < 4; c++) + max_g = std::max(max_g, std::fabs(g[c])); + for (size_t i = 0; i < atoms.size(); i += atoms.size() / 7) { + for (int c = 0; c < 4; c++) { + const gemmi::Position p0 = atoms[i]->pos; + const float b0 = atoms[i]->b_iso; + const double h = c < 3 ? hx : hb; + auto set = [&](double s) { + atoms[i]->pos = p0; + atoms[i]->b_iso = b0; + if (c < 3) + atoms[i]->pos.at(c) += s; + else + atoms[i]->b_iso = static_cast(b0 + s); + }; + set(h); + const double fp = engine.Evaluate(model, nullptr); + set(-h); + const double fm = engine.Evaluate(model, nullptr); + set(0); + const double fd = (fp - fm) / (2 * h); + CAPTURE(i, c, fd, grad[i][c], max_g); + CHECK(std::fabs(fd - grad[i][c]) < tolerance * std::max(std::fabs(fd), 0.05 * max_g)); + } + } +} + +} // namespace + +TEST_CASE("RefinementLikelihood_derivative", "[Refinement]") { + for (bool centric : {false, true}) + for (double fc : {0.3, 2.0, 9.0}) { + double dfc = 0, dummy = 0; + const double f0 = refine_math::AmplitudeNegLogLikelihood(4.0, fc, 0.8, 2.5, centric, dfc); + const double h = 1e-5; + const double fp = refine_math::AmplitudeNegLogLikelihood(4.0, fc + h, 0.8, 2.5, centric, dummy); + const double fm = refine_math::AmplitudeNegLogLikelihood(4.0, fc - h, 0.8, 2.5, centric, dummy); + CAPTURE(centric, fc, f0); + CHECK(dfc == Catch::Approx((fp - fm) / (2 * h)).epsilon(1e-5)); + } + // the Bessel approximation on both sides of its switch + double l0, r0, l1, r1; + refine_math::BesselI0I1(3.75 * (1 - 1e-12), l0, r0); + refine_math::BesselI0I1(3.75 * (1 + 1e-12), l1, r1); + CHECK(l0 == Catch::Approx(l1).epsilon(1e-5)); + CHECK(r0 == Catch::Approx(r1).epsilon(1e-5)); +} + +TEST_CASE("RefinementMonomerLibrary_embedded", "[Refinement]") { + const gemmi::MonLib &monlib = EmbeddedMonomerLibrary(); + for (const char *code : {"ALA", "PRO", "MSE", "DA", "U", "ATP", "GOL", "SO4", "HOH", "ZN"}) + CHECK(InEmbeddedMonomerLibrary(code)); + CHECK_FALSE(InEmbeddedMonomerLibrary("XYZ")); + for (const char *link : {"TRANS", "CIS", "PTRANS", "p", "SS"}) + CHECK(monlib.links.count(link) == 1); + CHECK(monlib.modifications.count("NH3") == 1); + CHECK(monlib.ener_lib.atoms.count("CH1") == 1); +} + +TEST_CASE("RefinementRestraints_ideal_residues", "[Refinement]") { + const gemmi::UnitCell cell(60, 65, 70, 90, 90, 90); + gemmi::Structure st = residue_crystal({"PHE", "DA", "ATP", "HOH", "ZN"}, cell, "P 21 21 21", 3); + const RefinementTopology topo = BuildRefinementTopology(st); + CHECK(topo.restrained_residues == 5); + CHECK(topo.fixed_residues == 0); + CHECK(topo.bonds > 60); + CHECK(topo.planes >= 3); + CHECK(topo.chirals >= 4); + std::vector> xyzb; + for (gemmi::Atom *a : atoms_of(st.models[0])) + xyzb.push_back({a->pos.x, a->pos.y, a->pos.z, a->b_iso}); + const GeometryStatistics g = MeasureGeometry(topo.geometry, {}, xyzb); + CHECK(g.rmsd_bonds < 0.02); + CHECK(g.rmsd_angles < 2.5); + CHECK(g.rmsd_chirals < 0.2); +} + +TEST_CASE("RefinementRestraints_gradient", "[Refinement]") { + const gemmi::UnitCell cell(14, 15, 16, 90, 90, 90); + gemmi::Structure st = residue_crystal({"TYR", "DG"}, cell, "P 21 21 21", 5); + st.models[0] = shaken(st.models[0], 0.15, 3.0, 2); + const RefinementTopology topo = BuildRefinementTopology(st); + // the geometry, the B similarity and the repulsions (symmetry mates included, at a contact distance that + // makes every pair within reach count) + RestraintSet set = topo.geometry; + set.Append(topo.adp); + RefinementTopology wide = topo; + for (float &r : wide.vdw_radius) + r = 2.5f; + std::fill(wide.hbond.begin(), wide.hbond.end(), 0); + const RestraintSet rep = BuildRepulsions(wide, st.models[0], cell, *st.find_spacegroup(), 0.0); + REQUIRE(rep.records.size() > 10); + bool image = false; + for (const RestraintRecord &r : rep.records) + image = image || r.rot[0] != 1.0f || r.tran[0] != 0.0f || r.tran[1] != 0.0f || r.tran[2] != 0.0f; + CAPTURE(rep.records.size()); + CHECK(image); + set.Append(rep); + set.Finalize(atoms_of(st.models[0]).size()); + + std::vector> xyzb; + for (gemmi::Atom *a : atoms_of(st.models[0])) + xyzb.push_back({a->pos.x, a->pos.y, a->pos.z, a->b_iso}); + // each kind of restraint on its own, against central differences + for (int kind = 0; kind <= static_cast(RestraintKind::BSimilarity); kind++) { + RestraintSet one; + for (const RestraintRecord &r : set.records) + if (r.kind == kind) { + RestraintRecord copy = r; + if (kind == static_cast(RestraintKind::Plane)) { + copy.atom[0] = static_cast(one.plane_atoms.size()); + for (int k = 0; k < r.atom[1]; k++) + one.plane_atoms.push_back(set.plane_atoms[r.atom[0] + k]); + } + one.Add(copy, kind == static_cast(RestraintKind::Plane) ? r.atom[1] : 4); + } + CAPTURE(kind, one.records.size()); + std::vector> grad(xyzb.size(), {0, 0, 0, 0}); + EvaluateRestraints(one, xyzb, &grad); + const double h = 1e-6; + for (size_t i = 0; i < xyzb.size(); i++) + for (int c = 0; c < 4; c++) { + auto x = xyzb; + x[i][c] += h; + const double fp = EvaluateRestraints(one, x, nullptr); + x[i][c] -= 2 * h; + const double fm = EvaluateRestraints(one, x, nullptr); + CAPTURE(i, c); + CHECK(grad[i][c] == Catch::Approx((fp - fm) / (2 * h)).epsilon(1e-4).margin(1e-3)); + } + } + +#ifdef JFJOCH_USE_CUDA + if (get_gpu_count() > 0) { + // the GPU's float evaluation of the same set + std::vector merged = asu_reflections(cell, *st.find_spacegroup(), 3.0); + RefinementReflections refl = PrepareRefinementReflections(merged, cell, *st.find_spacegroup(), 3.0); + auto gpu = MakeRefinementEngineGPU(0, cell, *st.find_spacegroup(), 3.0, refl, st.models[0], 100.0, SIZE_MAX); + REQUIRE(gpu); + gpu->SetRestraints(0, set); + std::vector> ggpu; + const double egpu = gpu->EvaluateRestraints(0, xyzb, &ggpu); + std::vector> gcpu(xyzb.size(), {0, 0, 0, 0}); + const double ecpu = EvaluateRestraints(set, xyzb, &gcpu); + CHECK(egpu == Catch::Approx(ecpu).epsilon(1e-4)); + for (size_t i = 0; i < xyzb.size(); i++) + for (int c = 0; c < 4; c++) + CHECK(ggpu[i][c] == Catch::Approx(gcpu[i][c]).epsilon(2e-3).margin(2e-2)); + } +#endif +} + +TEST_CASE("RefinementXray_gradient_CPU", "[Refinement]") { + for (const char *group : {"P 1", "P 21 21 21", "C 1 2 1", "P 61 2 2"}) { + CAPTURE(group); + const gemmi::SpaceGroup &sg = *gemmi::find_spacegroup_by_name(group); + const gemmi::UnitCell cell = sg.crystal_system() == gemmi::CrystalSystem::Hexagonal + ? gemmi::UnitCell(30, 30, 45, 90, 90, 120) + : gemmi::UnitCell(28, 31, 34, 90, sg.number == 5 ? 104 : 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.2), cell, sg, 2.2); + REQUIRE(refl.Size() > 500); + auto cpu = MakeRefinementEngineCPU(cell, sg, 2.2, refl, 4); + cpu->SetBlur(RefinementBlur(cell, sg, 2.2, 5.0)); + set_observations(*cpu, refl, truth.models[0], model); + check_xray_gradient(*cpu, model, 0.03); + } +} + +#ifdef JFJOCH_USE_CUDA +TEST_CASE("RefinementXray_gradient_GPU", "[Refinement]") { + if (get_gpu_count() == 0) + return; + for (const char *group : {"P 1", "P 21 21 21", "C 1 2 1", "P 61 2 2"}) { + CAPTURE(group); + const gemmi::SpaceGroup &sg = *gemmi::find_spacegroup_by_name(group); + const gemmi::UnitCell cell = sg.crystal_system() == gemmi::CrystalSystem::Hexagonal + ? gemmi::UnitCell(30, 30, 45, 90, 90, 120) + : gemmi::UnitCell(28, 31, 34, 90, sg.number == 5 ? 104 : 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.2), cell, sg, 2.2); + auto cpu = MakeRefinementEngineCPU(cell, sg, 2.2, refl, 4); + cpu->SetBlur(RefinementBlur(cell, sg, 2.2, 5.0)); + set_observations(*cpu, refl, truth.models[0], model); + // the GPU engine sees the same observations: it is made after they are set + auto gpu = MakeRefinementEngineGPU(0, cell, sg, 2.2, refl, model, 300.0, SIZE_MAX); + REQUIRE(gpu); + gpu->SetBlur(RefinementBlur(cell, sg, 2.2, 5.0)); + std::vector> f; + 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; + const double tc = cpu->Evaluate(model, &gc); + const double tg = gpu->Evaluate(model, &gg); + CHECK(tg == Catch::Approx(tc).epsilon(1e-4)); + double max_g = 0; + for (const auto &g : gc) + for (int c = 0; c < 4; c++) + max_g = std::max(max_g, std::fabs(g[c])); + for (size_t i = 0; i < gc.size(); i++) + for (int c = 0; c < 4; c++) { + CAPTURE(i, c); + CHECK(std::fabs(gg[i][c] - gc[i][c]) < 0.01 * max_g); + } + check_xray_gradient(*gpu, model, 0.05); + } +} +#endif + +TEST_CASE("Refinement_restrained_recovers_model", "[Refinement]") { + const gemmi::UnitCell cell(36, 40, 44, 90, 90, 90); + const gemmi::SpaceGroup &sg = *gemmi::find_spacegroup_by_name("P 21 21 21"); + gemmi::Structure truth = residue_crystal({"PHE", "TRP", "DA", "DG", "TYR", "HOH", "HOH", "HOH"}, cell, + "P 21 21 21", 11); + gemmi::Structure start = truth; + start.models[0] = shaken(truth.models[0], 0.25, 5.0, 12); + std::vector merged = asu_reflections(cell, sg, 1.8); + { + RefinementReflections refl = PrepareRefinementReflections(merged, cell, sg, 1.8 * 0.999); + auto cpu = MakeRefinementEngineCPU(cell, sg, 1.8 * 0.999, refl, 4); + cpu->SetBlur(RefinementBlur(cell, sg, 1.8, 5.0)); + std::vector> f; + cpu->Fcalc(truth.models[0], f); + std::map amplitude; + for (size_t m = 0; m < refl.Size(); m++) + amplitude[refl.hkl[m]] = std::abs(f[m]); + std::mt19937 rng(3); + std::normal_distribution noise(0.0, 0.03); + for (MergedReflection &r : merged) { + const auto it = amplitude.find(gemmi::Miller{{r.h, r.k, r.l}}); + r.F = it == amplitude.end() ? NAN : static_cast(it->second * (1 + noise(rng))); + r.sigmaF = 0.03f * r.F + 0.01f; + } + } + Logger logger("RefinementTest"); + for (bool gpu : {false, true}) { + if (gpu && get_gpu_count() == 0) + continue; + CAPTURE(gpu); + RefinementSettings settings; + settings.gpu = gpu; + settings.nthreads = 4; + settings.macro_cycles = 3; + settings.weight_ladder = {0.0625, 0.25}; + settings.adp_ladder = {4.0}; + settings.riding_hydrogens = false; // the data were made without them + const RefinementResult r = RefineModel(start, merged, cell, sg, 1.0, logger, settings); + REQUIRE(r.ok); + CHECK(r.on_gpu == gpu); + CHECK(r.restrained_residues == 8); + CHECK(r.r_free < 0.6 * r.r_free_start); + CHECK(r.geometry.rmsd_bonds < r.geometry_start.rmsd_bonds); + double rms = 0; + std::vector refined = atoms_of(r.model->models[0]); + std::erase_if(refined, [](const gemmi::Atom *a) { return a->is_hydrogen(); }); // the riding hydrogens + const std::vector ref = atoms_of(truth.models[0]); + REQUIRE(refined.size() == ref.size()); + for (size_t i = 0; i < ref.size(); i++) + rms += refined[i]->pos.dist_sq(ref[i]->pos); + rms = std::sqrt(rms / ref.size()); + CHECK(rms < 0.15); + } +} diff --git a/tools/CMakeLists.txt b/tools/CMakeLists.txt index b6754ec31..1ee3ad665 100644 --- a/tools/CMakeLists.txt +++ b/tools/CMakeLists.txt @@ -12,6 +12,11 @@ TARGET_LINK_LIBRARIES(jfjoch_extract_hkl JFJochReader) # rugnux, the single offline analysis CLI, lives in rugnux/ next to its library. +# Developer utility, not installed: refines a model against a finished run's merged MTZ (Refinement.h), for +# measuring the refinement on battery outputs without processing anything again. +ADD_EXECUTABLE(rugnux_refine_test rugnux_refine_test.cpp) +TARGET_LINK_LIBRARIES(rugnux_refine_test JFJochStructureRefinement JFJochCommon gemmi) + # In-place re-compress /entry/data/data of a _data file from bitshuffle/LZ4 to bitshuffle/zstd. ADD_EXECUTABLE(jfjoch_recompress jfjoch_recompress.cpp) TARGET_LINK_LIBRARIES(jfjoch_recompress JFJochHDF5Wrappers) diff --git a/tools/rugnux_refine_test.cpp b/tools/rugnux_refine_test.cpp new file mode 100644 index 000000000..4350d9b7a --- /dev/null +++ b/tools/rugnux_refine_test.cpp @@ -0,0 +1,151 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +// Refinement (image_analysis/structure_refinement/Refinement.h) of a model against a rugnux merged MTZ, without +// processing anything: for measuring the refinement on finished runs (a battery run's p.mtz and p_model.cif) +// against another program on the same files. Prints the outcome, and one JSON line for scripts. + +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include + +#include "../common/CUDAWrapper.h" +#include "../common/Logger.h" +#include "../image_analysis/structure_refinement/Refinement.h" + +namespace { + +void usage() { + std::cout << "Usage: rugnux_refine_test [options] \n" + " -N CPU threads (default 8)\n" + " --no-gpu the CPU engine\n" + " --placement-only rigid bodies and group B only\n" + " --cycles macro-cycles per weight (default as in RefinementSettings)\n" + " --xyz-iterations , --b-iterations L-BFGS iterations per macro-cycle\n" + " --ladder the weight ladder, times the gradient-balanced weight\n" + " -o write the refined model\n"; +} + +std::vector ParseList(const std::string &s) { + std::vector out; + std::stringstream ss(s); + std::string item; + while (std::getline(ss, item, ',')) + out.push_back(std::stod(item)); + return out; +} + +} // namespace + +int main(int argc, char **argv) { + RefinementSettings settings; + settings.nthreads = 8; + std::string out_path; + std::vector files; + for (int i = 1; i < argc; i++) { + const std::string a = argv[i]; + if (a == "-N" && i + 1 < argc) + settings.nthreads = std::stoul(argv[++i]); + else if (a == "--no-gpu") { + settings.gpu = false; + hide_gpus(); + } else if (a == "--placement-only") + settings.restrained = false; + else if (a == "--cycles" && i + 1 < argc) + settings.macro_cycles = std::stoi(argv[++i]); + else if (a == "--xyz-iterations" && i + 1 < argc) + settings.xyz_iterations = std::stoi(argv[++i]); + else if (a == "--adp-ladder" && i + 1 < argc) + settings.adp_ladder = ParseList(argv[++i]); + else if (a == "--b-iterations" && i + 1 < argc) + settings.b_iterations = std::stoi(argv[++i]); + else if (a == "--ladder" && i + 1 < argc) + settings.weight_ladder = ParseList(argv[++i]); + else if (a == "-o" && i + 1 < argc) + out_path = argv[++i]; + else if (!a.empty() && a[0] == '-') { + usage(); + return EXIT_FAILURE; + } else + files.push_back(a); + } + if (files.size() != 2) { + usage(); + return EXIT_FAILURE; + } + + gemmi::Mtz mtz = gemmi::read_mtz_file(files[0]); + const gemmi::SpaceGroup *sg = mtz.spacegroup; + if (sg == nullptr) { + std::cerr << "No space group in " << files[0] << std::endl; + return EXIT_FAILURE; + } + const int ih = mtz.column_with_label("H")->idx, ik = mtz.column_with_label("K")->idx, + il = mtz.column_with_label("L")->idx, iF = mtz.column_with_label("F")->idx, + isF = mtz.column_with_label("SIGF")->idx, ifree = mtz.column_with_label("FreeR_flag")->idx; + const gemmi::Mtz::Column *fp = mtz.column_with_label("F(+)"), *fm = mtz.column_with_label("F(-)"); + std::vector merged; + for (size_t r = 0; r < static_cast(mtz.nreflections); r++) { + const float *row = &mtz.data[r * mtz.columns.size()]; + MergedReflection m; + m.h = static_cast(row[ih]); + m.k = static_cast(row[ik]); + m.l = static_cast(row[il]); + m.F = row[iF]; + m.sigmaF = row[isF]; + m.rfree_flag = row[ifree] == 0; + m.d = static_cast(mtz.cell.calculate_d({{m.h, m.k, m.l}})); + if (fp != nullptr && fm != nullptr) { + m.F_plus = row[fp->idx]; + m.F_minus = row[fm->idx]; + } + merged.push_back(m); + } + gemmi::Structure st = gemmi::read_structure_gz(files[1]); + + Logger logger("rugnux_refine_test"); + const RefinementResult r = RefineModel(st, merged, mtz.cell, *sg, 1.0, logger, settings); + if (!r.ok) { + std::cerr << "Refinement failed: " << r.failure_reason << std::endl; + return EXIT_FAILURE; + } + if (!out_path.empty()) + { + std::ofstream f(out_path); + gemmi::cif::write_cif_to_stream(f, gemmi::make_mmcif_document(*r.model)); + } + + int kinds[5] = {0, 0, 0, 0, 0}; + for (const DensityBlob &b : r.blobs) + kinds[static_cast(b.kind)]++; + nlohmann::json blobs = nlohmann::json::array(); + for (const DensityBlob &b : r.blobs) + blobs.push_back({{"kind", DensityBlobKindName(b.kind)}, {"x", b.centroid.x}, {"y", b.centroid.y}, + {"z", b.centroid.z}, {"peak", b.peak_sigma}, {"volume", b.volume}, {"score", b.score}, + {"nearest", b.nearest_atom}, {"distance", b.nearest_atom_distance}, + {"anomalous", b.anomalous_sigma}}); + nlohmann::json j = {{"d_min", r.d_min}, {"gpu", r.on_gpu}, {"seconds", r.seconds}, + {"r_work_start", r.r_work_start}, {"r_free_start", r.r_free_start}, + {"r_work_placement", r.r_work_placement}, {"r_free_placement", r.r_free_placement}, + {"r_work", r.r_work}, {"r_free", r.r_free}, {"weight", r.weight}, {"adp_weight", r.adp_weight}, + {"rmsd_bonds", r.geometry.rmsd_bonds}, {"rmsd_angles", r.geometry.rmsd_angles}, + {"rmsd_bonds_start", r.geometry_start.rmsd_bonds}, + {"rmsd_angles_start", r.geometry_start.rmsd_angles}, {"clashes", r.geometry.clashes}, + {"atoms", r.atoms}, {"refined_atoms", r.refined_atoms}, {"fixed_residues", r.fixed_residues}, + {"blobs", blobs}}; + nlohmann::json trials = nlohmann::json::array(); + for (const RefinementWeightTrial &t : r.trials) + trials.push_back({{"weight", t.weight}, {"adp_weight", t.adp_weight}, {"r_work", t.r_work}, {"r_free", t.r_free}, + {"rmsd_bonds", t.geometry.rmsd_bonds}, {"rmsd_angles", t.geometry.rmsd_angles}}); + j["trials"] = trials; + std::cout << "JSON " << j.dump() << std::endl; + return EXIT_SUCCESS; +}