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 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_015CeEGgpKdfkCTJb3YNMJcu (cherry picked from commit 8efe6f2322b3bfdae580bb34f256e67730dfc3c2)
This commit is contained in:
leonarski_f committed
1 parent
e58453a693
commit
b4b68d368f
19 files changed
+4220
No files matched your search
@@ -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
|
||||
|
||||
@@ -18,6 +18,14 @@ ADD_LIBRARY(JFJochStructureRefinement STATIC
|
||||
$<$<BOOL:${JFJOCH_CUDA_AVAILABLE}>: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
|
||||
$<$<BOOL:${JFJOCH_CUDA_AVAILABLE}>:RefinementGPU.cpp>
|
||||
$<$<BOOL:${JFJOCH_CUDA_AVAILABLE}>: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
|
||||
|
||||
@@ -0,0 +1,110 @@
|
||||
// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
|
||||
// SPDX-License-Identifier: GPL-3.0-only
|
||||
|
||||
#include "DensityBlobs.h"
|
||||
|
||||
#include <algorithm>
|
||||
#include <cmath>
|
||||
|
||||
#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<DensityBlob> FindDensityBlobs(const gemmi::Grid<float> &difference, double rms, const gemmi::Structure &model,
|
||||
double d_min, const gemmi::Grid<float> *anomalous, double anomalous_rms) {
|
||||
std::vector<DensityBlob> out;
|
||||
if (!(rms > 0) || model.models.empty())
|
||||
return out;
|
||||
gemmi::Grid<float> map = difference;
|
||||
for (float &v : map.data)
|
||||
v = static_cast<float>(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<gemmi::Blob> 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;
|
||||
}
|
||||
@@ -0,0 +1,41 @@
|
||||
// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
|
||||
// SPDX-License-Identifier: GPL-3.0-only
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
#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<DensityBlob> FindDensityBlobs(const gemmi::Grid<float> &difference, double rms, const gemmi::Structure &model,
|
||||
double d_min, const gemmi::Grid<float> *anomalous = nullptr,
|
||||
double anomalous_rms = 0);
|
||||
@@ -0,0 +1,781 @@
|
||||
// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
|
||||
// SPDX-License-Identifier: GPL-3.0-only
|
||||
|
||||
#include "Refinement.h"
|
||||
|
||||
#include <algorithm>
|
||||
#include <chrono>
|
||||
#include <map>
|
||||
#include <cmath>
|
||||
#include <random>
|
||||
#include <unordered_map>
|
||||
|
||||
#include <ceres/gradient_problem.h>
|
||||
#include <ceres/gradient_problem_solver.h>
|
||||
#include <ceres/jet.h>
|
||||
#include <ceres/rotation.h>
|
||||
|
||||
#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<std::string> &names) {
|
||||
std::string out;
|
||||
for (const std::string &n : names)
|
||||
out += (out.empty() ? "" : " ") + n;
|
||||
return out;
|
||||
}
|
||||
|
||||
std::vector<gemmi::Atom *> AtomList(gemmi::Model &model) {
|
||||
std::vector<gemmi::Atom *> 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<float>(gemmi::u_to_b() * a.aniso.trace() / 3.0);
|
||||
a.aniso = {0, 0, 0, 0, 0, 0};
|
||||
}
|
||||
a.b_iso = static_cast<float>(std::clamp<double>(a.b_iso, B_FLOOR, B_CEILING_MAX));
|
||||
}
|
||||
}
|
||||
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<gemmi::Atom *> atoms;
|
||||
gemmi::UnitCell cell;
|
||||
const gemmi::SpaceGroup *sg = nullptr;
|
||||
double d_min = 0;
|
||||
RefinementReflections refl;
|
||||
std::unique_ptr<RefinementEngine> engine;
|
||||
#ifdef JFJOCH_USE_CUDA
|
||||
std::unique_ptr<ModelScaleGPUPool> scale_gpu;
|
||||
#endif
|
||||
size_t nthreads = 1;
|
||||
XrayRowModel rows;
|
||||
std::vector<std::complex<float>> 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<int> rider_parent;
|
||||
std::vector<gemmi::Position> 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<std::array<double, 4>> &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<std::array<double, 4>> XyzB() const {
|
||||
std::vector<std::array<double, 4>> 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<gemmi::Position> Positions() const {
|
||||
std::vector<gemmi::Position> out;
|
||||
for (const gemmi::Atom *a : atoms)
|
||||
out.push_back(a->pos);
|
||||
return out;
|
||||
}
|
||||
void SetPositions(const std::vector<gemmi::Position> &p) {
|
||||
for (size_t i = 0; i < atoms.size(); i++)
|
||||
atoms[i]->pos = p[i];
|
||||
}
|
||||
std::vector<double> B() const {
|
||||
std::vector<double> out;
|
||||
for (const gemmi::Atom *a : atoms)
|
||||
out.push_back(a->b_iso);
|
||||
return out;
|
||||
}
|
||||
void SetB(const std::vector<double> &b) {
|
||||
for (size_t i = 0; i < atoms.size(); i++)
|
||||
atoms[i]->b_iso = static_cast<float>(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<double>(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<std::complex<float>> fc, fm;
|
||||
gemmi::AsuData<gemmi::ValueSigma<float>> 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<float> 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<std::complex<float>> fmodel(refl.Size());
|
||||
for (size_t m = 0; m < refl.Size(); m++) {
|
||||
rows.k_total[m] = static_cast<float>(scaling.get_overall_scale_factor(refl.hkl[m]));
|
||||
rows.fmask[m] = static_cast<float>(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<std::complex<float>> FModel() {
|
||||
engine->Fcalc(Model(), fcalc);
|
||||
std::vector<std::complex<float>> 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<int> ¶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<int>(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<float>(std::clamp(p[k], ctx_.b_floor, ctx_.b_ceiling));
|
||||
}
|
||||
ctx_.MoveRiders();
|
||||
ctx_.evaluations++;
|
||||
std::vector<std::array<double, 4>> 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<int> ¶m_atoms_;
|
||||
double weight_;
|
||||
int restraint_set_;
|
||||
};
|
||||
|
||||
void Minimise(Context &ctx, AtomTarget::Mode mode, const std::vector<int> ¶m_atoms, double weight,
|
||||
int restraint_set, int iterations) {
|
||||
if (param_atoms.empty())
|
||||
return;
|
||||
std::vector<double> 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<AtomTarget>(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<float>(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<std::vector<int>> &bodies)
|
||||
: ctx_(ctx), bodies_(bodies), base_(ctx.Positions()) {
|
||||
for (const std::vector<int> &b : bodies) {
|
||||
gemmi::Position c;
|
||||
for (int i : b)
|
||||
c += base_[i];
|
||||
c *= 1.0 / static_cast<double>(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<double>(b.size()))));
|
||||
}
|
||||
}
|
||||
int NumParameters() const override { return static_cast<int>(6 * bodies_.size()); }
|
||||
|
||||
bool Evaluate(const double *q, double *cost, double *gradient) const override {
|
||||
Place(q);
|
||||
ctx_.evaluations++;
|
||||
std::vector<std::array<double, 4>> g;
|
||||
*cost = ctx_.engine->Evaluate(ctx_.Model(), gradient != nullptr ? &g : nullptr);
|
||||
if (gradient == nullptr)
|
||||
return std::isfinite(*cost);
|
||||
using Jet = ceres::Jet<double, 3>;
|
||||
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<std::vector<int>> &bodies_;
|
||||
std::vector<gemmi::Position> base_;
|
||||
std::vector<gemmi::Position> centre_;
|
||||
std::vector<double> 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<std::vector<int>> &groups)
|
||||
: ctx_(ctx), groups_(groups), base_(ctx.B()) {}
|
||||
int NumParameters() const override { return static_cast<int>(groups_.size()); }
|
||||
|
||||
bool Evaluate(const double *q, double *cost, double *gradient) const override {
|
||||
Place(q);
|
||||
ctx_.evaluations++;
|
||||
std::vector<std::array<double, 4>> 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<float>(std::clamp(base_[i] + q[k], ctx_.b_floor, ctx_.b_ceiling));
|
||||
}
|
||||
|
||||
private:
|
||||
Context &ctx_;
|
||||
const std::vector<std::vector<int>> &groups_;
|
||||
std::vector<double> base_;
|
||||
};
|
||||
|
||||
// The placement stage: rigid bodies, then group B.
|
||||
void RefinePlacement(Context &ctx, RefinementResult &result) {
|
||||
std::vector<std::vector<int>> bodies, groups;
|
||||
size_t i = 0;
|
||||
for (const gemmi::Chain &ch : ctx.Model().chains) {
|
||||
std::vector<int> body;
|
||||
for (const gemmi::Residue &r : ch.residues) {
|
||||
std::vector<int> group;
|
||||
for (size_t k = 0; k < r.atoms.size(); k++, i++) {
|
||||
group.push_back(static_cast<int>(i));
|
||||
if (!r.is_water())
|
||||
body.push_back(static_cast<int>(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<int>(bodies.size());
|
||||
if (!bodies.empty()) {
|
||||
std::vector<double> q(6 * bodies.size(), 0.0);
|
||||
auto target = std::make_unique<RigidBodiesTarget>(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<double> q(groups.size(), 0.0);
|
||||
auto target = std::make_unique<GroupBTarget>(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<int> &refinable) {
|
||||
std::vector<std::array<double, 4>> gx, gr;
|
||||
ctx.engine->Evaluate(ctx.Model(), &gx);
|
||||
std::vector<std::array<double, 4>> xyzb = ctx.XyzB();
|
||||
std::mt19937 rng(1);
|
||||
std::normal_distribution<double> shake(0.0, 0.1);
|
||||
for (std::array<double, 4> &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<MergedReflection> &merged, double wavelength_A, RefinementResult &result) {
|
||||
ctx.Refit();
|
||||
const std::vector<std::complex<float>> fmodel = ctx.FModel();
|
||||
const gemmi::GroupOps gops = ctx.sg->operations();
|
||||
std::vector<SigmaAReflection> 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<int>(ctx.refl.epsilon[m]),
|
||||
static_cast<bool>(ctx.refl.centric[m]), static_cast<bool>(ctx.refl.free[m])});
|
||||
const SigmaAResult sigma_a = EstimateSigmaA(sa, ctx.cell);
|
||||
result.mean_fom = sigma_a.mean_fom;
|
||||
|
||||
auto maps = std::make_shared<ModelMaps>();
|
||||
MtzTable &mtz = maps->coefficients;
|
||||
mtz.space_group_name = ctx.sg->xhm();
|
||||
mtz.unit_cell = UnitCell{static_cast<float>(ctx.cell.a), static_cast<float>(ctx.cell.b), static_cast<float>(ctx.cell.c),
|
||||
static_cast<float>(ctx.cell.alpha), static_cast<float>(ctx.cell.beta),
|
||||
static_cast<float>(ctx.cell.gamma)};
|
||||
mtz.datasets = {{"HKL_base", "HKL_base", "HKL_base", 0.0f},
|
||||
{"refinement", "refinement", "refinement", static_cast<float>(wavelength_A)}};
|
||||
const std::pair<const char *, char> 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<std::complex<float>> 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<float>(ctx.refl.hkl[m][0]), static_cast<float>(ctx.refl.hkl[m][1]),
|
||||
static_cast<float>(ctx.refl.hkl[m][2]), static_cast<float>(fo), static_cast<float>(fc),
|
||||
static_cast<float>(phi_deg), static_cast<float>(fwt), static_cast<float>(phi_deg),
|
||||
static_cast<float>(delfwt), static_cast<float>(phi_deg), static_cast<float>(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<float>(delfwt), static_cast<float>(phi))});
|
||||
}
|
||||
result.maps = maps;
|
||||
const std::array<int, 3> size = gemmi::get_size_for_hkl(diff, {{0, 0, 0}}, 3.0);
|
||||
result.difference_map = std::make_shared<gemmi::Grid<float>>(MapFromFPhi(gemmi::get_f_phi_on_grid<float>(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<gemmi::Miller, float> 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<std::complex<float>> 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<float>(std::arg(fmodel[m]) - PI / 2))});
|
||||
}
|
||||
if (!anom.v.empty()) {
|
||||
result.anomalous_map = std::make_shared<gemmi::Grid<float>>(
|
||||
MapFromFPhi(gemmi::get_f_phi_on_grid<float>(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<ModelScaleGPUPool>(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<MergedReflection> &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<int>(ctx.atoms.size());
|
||||
double b_max = 0;
|
||||
for (const gemmi::Atom *a : ctx.atoms)
|
||||
b_max = std::max<double>(b_max, a->b_iso);
|
||||
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<int>(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<int> refinable;
|
||||
for (size_t i = 0; i < ctx.atoms.size(); i++)
|
||||
if (topo.refinable[i])
|
||||
refinable.push_back(static_cast<int>(i));
|
||||
result.refined_atoms = static_cast<int>(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<gemmi::Position> pos0 = ctx.Positions();
|
||||
const std::vector<double> b0 = ctx.B();
|
||||
int best = -1;
|
||||
std::vector<gemmi::Position> best_pos;
|
||||
std::vector<double> 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<int>(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<gemmi::Structure>(ctx.st);
|
||||
result.ok = true;
|
||||
result.seconds = std::chrono::duration<double>(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;
|
||||
}
|
||||
@@ -0,0 +1,103 @@
|
||||
// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
|
||||
// SPDX-License-Identifier: GPL-3.0-only
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <memory>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
#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<double> 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<double> 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<RefinementWeightTrial> trials;
|
||||
GeometryStatistics geometry_start, geometry;
|
||||
|
||||
int atoms = 0, refined_atoms = 0;
|
||||
int restrained_residues = 0, fixed_residues = 0;
|
||||
std::vector<std::string> fixed_names;
|
||||
double k_sol = 0, b_sol = 0;
|
||||
double mean_fom = 0;
|
||||
|
||||
std::shared_ptr<gemmi::Structure> model; // the refined model, in the data's cell and frame
|
||||
std::shared_ptr<ModelMaps> 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<gemmi::Grid<float>> 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<gemmi::Grid<float>> anomalous_map;
|
||||
double anomalous_map_rms = 0;
|
||||
// What the refined model leaves unexplained in the difference map (FindDensityBlobs).
|
||||
std::vector<DensityBlob> 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<MergedReflection> &merged,
|
||||
const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg, double wavelength_A,
|
||||
Logger &logger, const RefinementSettings &settings);
|
||||
@@ -0,0 +1,283 @@
|
||||
// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
|
||||
// SPDX-License-Identifier: GPL-3.0-only
|
||||
|
||||
#include "RefinementGPU.h"
|
||||
|
||||
#include <algorithm>
|
||||
#include <cmath>
|
||||
|
||||
#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<float>;
|
||||
|
||||
// One atom's density as the gather takes it, isotropic, with `blur`.
|
||||
ModelDensityAtom DensityAtom(const gemmi::Atom &atom, double b_iso, const gemmi::DensityCalculator<Table, float> &dc) {
|
||||
using CReal = Table::Coef::coef_type;
|
||||
const auto &coef = Table::get(atom.element, atom.charge, atom.serial);
|
||||
const CReal b = static_cast<CReal>(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<float>(radius);
|
||||
a.max_r2 = static_cast<float>(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<const gemmi::Atom *> Atoms(const gemmi::Model &model) {
|
||||
std::vector<const gemmi::Atom *> 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<float> 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<double>(gops.cen_ops.size());
|
||||
for (const gemmi::Op &op : gops.sym_ops) {
|
||||
std::array<int, 12> 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<double>(op.rot[i][j]) / gemmi::Op::DEN;
|
||||
m.tran[i] = static_cast<double>(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<std::pair<int, int>> entries;
|
||||
const int n_ops = static_cast<int>(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<int>(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<int>(e));
|
||||
}
|
||||
setup_.scatter_term.push_back(entries[e].second);
|
||||
}
|
||||
setup_.scatter_start.push_back(static_cast<int>(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<ModelDensityAtom> 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<RefinementGPUEngine>(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<std::array<float, 2>> 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<std::complex<float>> &f) override {
|
||||
SetAtoms(model);
|
||||
std::vector<std::array<float, 2>> 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<std::complex<float>> &f) override {
|
||||
const gemmi::SolventMasker masker(gemmi::AtomicRadiiSet::Refmac);
|
||||
std::vector<ModelMaskAtom> 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<std::array<float, 2>> 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<std::array<double, 4>> *grad) override {
|
||||
SetAtoms(model);
|
||||
std::vector<std::array<float, 4>> 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<std::array<double, 4>> &xyzb,
|
||||
std::vector<std::array<double, 4>> *grad) override {
|
||||
std::vector<std::array<float, 4>> x(xyzb.size()), g;
|
||||
for (size_t i = 0; i < xyzb.size(); i++)
|
||||
x[i] = {static_cast<float>(xyzb[i][0]), static_cast<float>(xyzb[i][1]), static_cast<float>(xyzb[i][2]),
|
||||
static_cast<float>(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<Table, float> Calculator(double blur) const {
|
||||
gemmi::DensityCalculator<Table, float> 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<const gemmi::Atom *> atoms = Atoms(model);
|
||||
if (last_b_.size() != atoms.size()) {
|
||||
last_b_.assign(atoms.size(), NAN);
|
||||
atoms_.resize(atoms.size());
|
||||
}
|
||||
std::vector<std::array<float, 4>> 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<float>(f.x * n3[0]), static_cast<float>(f.y * n3[1]), static_cast<float>(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<Table, float> dc_;
|
||||
bool supported_ = false;
|
||||
size_t fft_work_ = 0, bytes_ = 0;
|
||||
std::unique_ptr<RefinementGPUEngine> engine_;
|
||||
std::vector<float> last_b_;
|
||||
std::vector<ModelDensityAtom> atoms_;
|
||||
};
|
||||
|
||||
} // namespace
|
||||
|
||||
std::unique_ptr<RefinementEngine> 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<RefinementEngineGPU>(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;
|
||||
}
|
||||
@@ -0,0 +1,691 @@
|
||||
// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
|
||||
// SPDX-License-Identifier: GPL-3.0-only
|
||||
|
||||
#include "RefinementGPUEngine.h"
|
||||
|
||||
#include <algorithm>
|
||||
#include <string>
|
||||
|
||||
#include <cub/cub.cuh>
|
||||
#include <cufft.h>
|
||||
|
||||
#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<int>(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<int>((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<int>(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<size_t>(nu / 2 + 1) * (imod(k[1], nv) + static_cast<size_t>(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<float>(sr), static_cast<float>(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<double>(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<float>(cr), static_cast<float>(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<float>(sr), static_cast<float>(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<size_t>(nu / 2 + 1) * (imod(k, nv) + static_cast<size_t>(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<float>(u);
|
||||
const float fy = nctr[1] - static_cast<float>(v);
|
||||
const float fz = nctr[2] - static_cast<float>(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<size_t>(g.nu) * (imod(v, g.nv) + static_cast<size_t>(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<float>(FOUR_PI2));
|
||||
}
|
||||
}
|
||||
|
||||
using V3f = refine_math::V3<float>;
|
||||
|
||||
__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<RestraintKind>(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<size_t>(g.nu) * g.nv * g.nw;
|
||||
}
|
||||
|
||||
size_t ComplexPoints(const ModelDensityGrid &g) {
|
||||
return static_cast<size_t>(g.nu / 2 + 1) * g.nv * g.nw;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
void Upload(CudaDevicePtr<T> &dst, const std::vector<T> &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 <typename T>
|
||||
void Ensure(CudaDevicePtr<T> &p, size_t &capacity, size_t n) {
|
||||
if (n <= capacity && p)
|
||||
return;
|
||||
capacity = std::max<size_t>(n, 1);
|
||||
p = CudaDevicePtr<T>(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<double *>(nullptr), static_cast<double *>(nullptr),
|
||||
static_cast<int>(std::max<size_t>(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<RestraintRecord> rec;
|
||||
CudaDevicePtr<int> plane_atoms, slot_start, slot_index;
|
||||
CudaDevicePtr<float4> contrib, xyzb, grad;
|
||||
CudaDevicePtr<double> energy, sum;
|
||||
CudaDevicePtr<char> reduce_temp;
|
||||
};
|
||||
|
||||
RefinementGPUSetup setup;
|
||||
size_t work_bytes;
|
||||
int device = 0;
|
||||
CudaStream stream;
|
||||
int n_atoms = 0;
|
||||
|
||||
CudaDevicePtr<float4> pos;
|
||||
CudaDevicePtr<int> box;
|
||||
CudaDevicePtr<ModelMaskAtom> mask_atoms;
|
||||
CudaDevicePtr<float> grid;
|
||||
CudaDevicePtr<float2> spectrum;
|
||||
CudaDevicePtr<char> fft_work;
|
||||
CudaDevicePtr<int> ops, row_hkl;
|
||||
CudaDevicePtr<double> row_scale;
|
||||
CudaDevicePtr<float> k_total, ml_d, ml_var, fo;
|
||||
CudaDevicePtr<float2> fmask_scaled, f, fm, coef;
|
||||
CudaDevicePtr<char> work, centric;
|
||||
CudaDevicePtr<double> nll, nll_sum;
|
||||
CudaDevicePtr<char> reduce_temp;
|
||||
size_t reduce_bytes = 0;
|
||||
CudaDevicePtr<int> scatter_point, scatter_start, scatter_term;
|
||||
CudaDevicePtr<float4> grad;
|
||||
CudaDevicePtr<ModelDensityAtom> atoms_copy; // the density gather keeps its own; this one is the gradient's
|
||||
std::unique_ptr<ModelDensityGPU> density;
|
||||
std::unique_ptr<ModelMaskGPU> 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<size_t>(setup.max_atoms, 1), nr = std::max<size_t>(setup.rows.size(), 1);
|
||||
pos = CudaDevicePtr<float4>(na, sync);
|
||||
box = CudaDevicePtr<int>(3 * na, sync);
|
||||
mask_atoms = CudaDevicePtr<ModelMaskAtom>(na, sync);
|
||||
atoms_copy = CudaDevicePtr<ModelDensityAtom>(na, sync);
|
||||
grad = CudaDevicePtr<float4>(na, sync);
|
||||
grid = CudaDevicePtr<float>(GridPoints(setup.grid), sync);
|
||||
spectrum = CudaDevicePtr<float2>(ComplexPoints(setup.grid), sync);
|
||||
fft_work = CudaDevicePtr<char>(std::max<size_t>(work_bytes, 1), sync);
|
||||
ops = CudaDevicePtr<int>(std::max<size_t>(12 * setup.sym_ops.size(), 1), sync);
|
||||
row_hkl = CudaDevicePtr<int>(3 * nr, sync);
|
||||
row_scale = CudaDevicePtr<double>(nr, sync);
|
||||
k_total = CudaDevicePtr<float>(nr, sync);
|
||||
ml_d = CudaDevicePtr<float>(nr, sync);
|
||||
ml_var = CudaDevicePtr<float>(nr, sync);
|
||||
fo = CudaDevicePtr<float>(nr, sync);
|
||||
fmask_scaled = CudaDevicePtr<float2>(nr, sync);
|
||||
f = CudaDevicePtr<float2>(nr, sync);
|
||||
fm = CudaDevicePtr<float2>(nr, sync);
|
||||
coef = CudaDevicePtr<float2>(nr, sync);
|
||||
work = CudaDevicePtr<char>(nr, sync);
|
||||
centric = CudaDevicePtr<char>(nr, sync);
|
||||
nll = CudaDevicePtr<double>(nr, sync);
|
||||
nll_sum = CudaDevicePtr<double>(1, sync);
|
||||
reduce_bytes = ReduceBytes(nr);
|
||||
reduce_temp = CudaDevicePtr<char>(std::max<size_t>(reduce_bytes, 1), sync);
|
||||
scatter_point = CudaDevicePtr<int>(std::max<size_t>(setup.scatter_point.size(), 1), sync);
|
||||
scatter_start = CudaDevicePtr<int>(std::max<size_t>(setup.scatter_start.size(), 1), sync);
|
||||
scatter_term = CudaDevicePtr<int>(std::max<size_t>(setup.scatter_term.size(), 1), sync);
|
||||
density = std::make_unique<ModelDensityGPU>(stream, setup.max_atoms, setup.max_pairs,
|
||||
ModelDensityGPU::Bricks(setup.grid.nu, setup.grid.nv, setup.grid.nw));
|
||||
mask = std::make_unique<ModelMaskGPU>(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<int> flat_ops;
|
||||
for (const auto &o : setup.sym_ops)
|
||||
flat_ops.insert(flat_ops.end(), o.begin(), o.end());
|
||||
std::vector<int> 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<float>(setup.grid.orth[0] * inv_n[0]);
|
||||
geom.m01 = static_cast<float>(setup.grid.orth[1] * inv_n[1]);
|
||||
geom.m02 = static_cast<float>(setup.grid.orth[2] * inv_n[2]);
|
||||
geom.m11 = static_cast<float>(setup.grid.orth[4] * inv_n[1]);
|
||||
geom.m12 = static_cast<float>(setup.grid.orth[5] * inv_n[2]);
|
||||
geom.m22 = static_cast<float>(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<float>(setup.volume / static_cast<double>(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<cufftComplex *>(spectrum.get())));
|
||||
}
|
||||
|
||||
void Compose(bool evaluate) {
|
||||
const int rows = static_cast<int>(setup.rows.size());
|
||||
if (rows == 0)
|
||||
return;
|
||||
compose_target_kernel<<<blocks(rows), THREADS, 0, stream>>>(
|
||||
spectrum, Norm(), setup.grid.nu, setup.grid.nv, setup.grid.nw, ops, static_cast<int>(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<RefinementGPUEngineImpl>(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<double> 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<float> &k_total, const std::vector<std::array<float, 2>> &fmask,
|
||||
const std::vector<float> &d, const std::vector<float> &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<ModelDensityAtom> &atoms, const std::vector<std::array<float, 4>> &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<int>(atoms.size());
|
||||
const std::vector<int> 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<std::array<float, 2>> &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<ModelMaskAtom> &mask_atoms, std::vector<std::array<float, 2>> &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<int>(mask_atoms.size()), e.grid);
|
||||
cufft_err(e.forward.R2C(e.grid, reinterpret_cast<cufftComplex *>(e.spectrum.get())));
|
||||
const int rows = static_cast<int>(e.setup.rows.size());
|
||||
f.resize(rows);
|
||||
if (rows > 0) {
|
||||
fmask_kernel<<<blocks(rows), THREADS, 0, e.stream>>>(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<std::array<float, 4>> *grad) {
|
||||
RefinementGPUEngineImpl &e = *impl_;
|
||||
DeviceGuard guard(e.device);
|
||||
const int rows = static_cast<int>(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<int>(e.setup.scatter_point.size());
|
||||
if (points > 0) {
|
||||
scatter_kernel<<<blocks(points), THREADS, 0, e.stream>>>(e.coef, e.row_hkl, e.ops,
|
||||
static_cast<int>(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<cufftComplex *>(e.spectrum.get()), e.grid));
|
||||
if (e.n_atoms > 0) {
|
||||
const size_t threads = static_cast<size_t>(e.n_atoms) * WARP;
|
||||
atom_gradient_kernel<<<blocks(threads), THREADS, 0, e.stream>>>(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<RestraintRecord> &records,
|
||||
const std::vector<int> &plane_atoms, const std::vector<int> &slot_start,
|
||||
const std::vector<int> &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<int>(RestraintKind::Plane) && rec.atom[1] > MAX_PLANE_ATOMS)
|
||||
throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "refinement: a plane of too many atoms");
|
||||
r.records = static_cast<int>(records.size());
|
||||
r.atoms = static_cast<int>(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<size_t>(slots));
|
||||
Ensure(r.xyzb, r.cap_xyzb, r.atoms);
|
||||
Ensure(r.grad, r.cap_atoms, r.atoms);
|
||||
if (!r.sum)
|
||||
r.sum = CudaDevicePtr<double>(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<char>(std::max<size_t>(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<std::array<float, 4>> &xyzb,
|
||||
std::vector<std::array<float, 4>> *grad) {
|
||||
RefinementGPUEngineImpl &e = *impl_;
|
||||
DeviceGuard guard(e.device);
|
||||
RefinementGPUEngineImpl::Restraints &r = e.restraints[which];
|
||||
if (static_cast<int>(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<<<blocks(r.records), THREADS, 0, e.stream>>>(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<<<blocks(r.atoms), THREADS, 0, e.stream>>>(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;
|
||||
}
|
||||
@@ -0,0 +1,24 @@
|
||||
// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
|
||||
// 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 <memory>
|
||||
|
||||
#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<RefinementEngine> 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);
|
||||
@@ -0,0 +1,78 @@
|
||||
// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
|
||||
// SPDX-License-Identifier: GPL-3.0-only
|
||||
|
||||
#pragma once
|
||||
|
||||
// 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 <array>
|
||||
#include <cstddef>
|
||||
#include <memory>
|
||||
#include <vector>
|
||||
|
||||
#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<std::array<int, 12>> 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<ModelMaskOp> mask_ops; // every operator with every centring vector
|
||||
std::vector<std::array<int, 3>> rows; // the reflections, in the order their numbers come back
|
||||
std::vector<double> inv_d2;
|
||||
std::vector<float> fo;
|
||||
std::vector<char> 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<int> 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<float> &k_total, const std::vector<std::array<float, 2>> &fmask,
|
||||
const std::vector<float> &d, const std::vector<float> &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<ModelDensityAtom> &atoms, const std::vector<std::array<float, 4>> &pos);
|
||||
// F of the atoms at the rows.
|
||||
void Fcalc(std::vector<std::array<float, 2>> &f);
|
||||
// The bulk-solvent mask's F at the rows (the mask atoms fractional, wrapped).
|
||||
void Fmask(const std::vector<ModelMaskAtom> &mask_atoms, std::vector<std::array<float, 2>> &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<std::array<float, 4>> *grad);
|
||||
|
||||
// Restraint set `which` (0 or 1), as RestraintSet holds it.
|
||||
void SetRestraints(int which, const std::vector<RestraintRecord> &records, const std::vector<int> &plane_atoms,
|
||||
const std::vector<int> &slot_start, const std::vector<int> &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<std::array<float, 4>> &xyzb,
|
||||
std::vector<std::array<float, 4>> *grad);
|
||||
|
||||
// The most atoms of a plane restraint the kernel takes.
|
||||
static constexpr int MAX_PLANE_ATOMS = 32;
|
||||
|
||||
private:
|
||||
std::unique_ptr<RefinementGPUEngineImpl> impl_;
|
||||
};
|
||||
@@ -0,0 +1,268 @@
|
||||
// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
|
||||
// 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 <cmath>
|
||||
|
||||
#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 <typename T>
|
||||
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 <typename T>
|
||||
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 <typename T>
|
||||
struct V3 {
|
||||
T x, y, z;
|
||||
};
|
||||
|
||||
template <typename T> REFINE_HD inline V3<T> operator-(V3<T> a, V3<T> b) { return {a.x - b.x, a.y - b.y, a.z - b.z}; }
|
||||
template <typename T> REFINE_HD inline V3<T> operator+(V3<T> a, V3<T> b) { return {a.x + b.x, a.y + b.y, a.z + b.z}; }
|
||||
template <typename T> REFINE_HD inline V3<T> operator*(T s, V3<T> a) { return {s * a.x, s * a.y, s * a.z}; }
|
||||
template <typename T> REFINE_HD inline T Dot(V3<T> a, V3<T> b) { return a.x * b.x + a.y * b.y + a.z * b.z; }
|
||||
template <typename T> REFINE_HD inline V3<T> Cross(V3<T> a, V3<T> 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 <typename T> REFINE_HD inline T Norm(V3<T> a) { return sqrt(Dot(a, a)); }
|
||||
|
||||
// A bond length.
|
||||
template <typename T>
|
||||
REFINE_HD inline T BondTerm(V3<T> a, V3<T> b, T ideal, T sigma, V3<T> g[2]) {
|
||||
const V3<T> 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 <typename T>
|
||||
REFINE_HD inline T AngleTerm(V3<T> a, V3<T> b, V3<T> c, T ideal, T sigma, V3<T> g[3]) {
|
||||
const V3<T> u = a - b, v = c - b;
|
||||
const T lu = Norm(u), lv = Norm(v);
|
||||
const V3<T> 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<T> du = (T(1.0) / (lu * sin_t)) * ((cos_t / lu) * u - (T(1.0) / lv) * v);
|
||||
const V3<T> 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 <typename T>
|
||||
REFINE_HD inline T Dihedral(V3<T> p0, V3<T> p1, V3<T> p2, V3<T> p3, V3<T> d[4]) {
|
||||
const V3<T> b1 = p1 - p0, b2 = p2 - p1, b3 = p3 - p2;
|
||||
const V3<T> 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 <typename T>
|
||||
REFINE_HD inline T TorsionTerm(V3<T> p0, V3<T> p1, V3<T> p2, V3<T> p3, T ideal, T sigma, int period, V3<T> g[4]) {
|
||||
V3<T> 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 <typename T>
|
||||
REFINE_HD inline T ChiralTerm(V3<T> c, V3<T> a1, V3<T> a2, V3<T> a3, T ideal, T sigma, V3<T> g[4]) {
|
||||
const V3<T> 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 <typename T>
|
||||
REFINE_HD inline T RepulsionTerm(V3<T> a, V3<T> b, T d0, T sigma, V3<T> g[2]) {
|
||||
const V3<T> d = a - b;
|
||||
const T len = Norm(d);
|
||||
if (len >= d0 || len <= T(0)) {
|
||||
g[0] = g[1] = V3<T>{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 <typename T>
|
||||
REFINE_HD inline V3<T> 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 <typename T>
|
||||
REFINE_HD inline T PlaneTerm(const V3<T> *x, int n, T sigma, V3<T> *g) {
|
||||
V3<T> 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<T> 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<T> 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
|
||||
@@ -0,0 +1,568 @@
|
||||
// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
|
||||
// SPDX-License-Identifier: GPL-3.0-only
|
||||
|
||||
#include "RefinementRestraints.h"
|
||||
|
||||
#include <algorithm>
|
||||
#include <cmath>
|
||||
#include <set>
|
||||
#include <unordered_map>
|
||||
|
||||
#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<std::string> 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<float>(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<double> At(const std::vector<std::array<double, 4>> &xyzb, int i) {
|
||||
return {xyzb[i][0], xyzb[i][1], xyzb[i][2]};
|
||||
}
|
||||
|
||||
V3<double> Image(const RestraintRecord &r, V3<double> 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<double> FromImage(const RestraintRecord &r, V3<double> 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<int>(plane_atoms.size());
|
||||
for (RestraintRecord r : other.records) {
|
||||
r.slot += slots;
|
||||
if (r.kind == static_cast<int>(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<int>(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<int>(RestraintKind::Bond) ||
|
||||
r.kind == static_cast<int>(RestraintKind::Repulsion) ||
|
||||
r.kind == static_cast<int>(RestraintKind::BSimilarity) ? 2
|
||||
: r.kind == static_cast<int>(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<int> 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<std::array<double, 4>> &xyzb) {
|
||||
V3<double> g[4];
|
||||
switch (static_cast<RestraintKind>(r.kind)) {
|
||||
case RestraintKind::Bond:
|
||||
return refine_math::Norm(At(xyzb, r.atom[0]) - At(xyzb, r.atom[1]));
|
||||
case RestraintKind::Angle: {
|
||||
const V3<double> 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<double> 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<std::array<double, 4>> &xyzb,
|
||||
std::vector<std::array<double, 4>> *grad) {
|
||||
double sum = 0;
|
||||
std::vector<V3<double>> px, pg;
|
||||
for (const RestraintRecord &r : set.records) {
|
||||
V3<double> g[4] = {};
|
||||
int n = 0;
|
||||
double e = 0;
|
||||
const double ideal = r.ideal, sigma = r.sigma;
|
||||
switch (static_cast<RestraintKind>(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<double, 4> &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<double, 4> &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<gemmi::MonLib>(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<const gemmi::Atom *, int> index;
|
||||
std::vector<const gemmi::Residue *> 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<int>(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<std::string> 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<const gemmi::Atom *> atoms) {
|
||||
for (const gemmi::Atom *a : atoms)
|
||||
if (a->is_hydrogen())
|
||||
return true;
|
||||
return false;
|
||||
};
|
||||
auto any_refinable = [&](std::initializer_list<int> atoms) {
|
||||
for (int a : atoms)
|
||||
if (out.refinable[a])
|
||||
return true;
|
||||
return false;
|
||||
};
|
||||
std::vector<std::set<int>> excluded(n);
|
||||
std::vector<std::vector<int>> bonded(n);
|
||||
for (size_t i = 0; i < n; i++)
|
||||
excluded[i].insert(static_cast<int>(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<int>(RestraintKind::Bond);
|
||||
r.atom[0] = i;
|
||||
r.atom[1] = j;
|
||||
r.ideal = static_cast<float>(b.restr->value);
|
||||
r.sigma = static_cast<float>(b.restr->esd);
|
||||
out.geometry.Add(r, 2);
|
||||
out.bonds++;
|
||||
RestraintRecord rb;
|
||||
rb.kind = static_cast<int>(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<int>(RestraintKind::Angle);
|
||||
r.atom[0] = i;
|
||||
r.atom[1] = j;
|
||||
r.atom[2] = k;
|
||||
r.ideal = static_cast<float>(an.restr->value);
|
||||
r.sigma = static_cast<float>(an.restr->esd);
|
||||
out.geometry.Add(r, 3);
|
||||
out.angles++;
|
||||
RestraintRecord rb;
|
||||
rb.kind = static_cast<int>(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<int>(RestraintKind::Torsion);
|
||||
r.atom[0] = a;
|
||||
r.atom[1] = b;
|
||||
r.atom[2] = c;
|
||||
r.atom[3] = d;
|
||||
r.ideal = static_cast<float>(t.restr->value);
|
||||
r.sigma = static_cast<float>(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<int>(RestraintKind::Chiral);
|
||||
r.atom[0] = c;
|
||||
r.atom[1] = a1;
|
||||
r.atom[2] = a2;
|
||||
r.atom[3] = a3;
|
||||
r.ideal = static_cast<float>(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<int>(RestraintKind::Plane);
|
||||
r.atom[0] = static_cast<int>(out.geometry.plane_atoms.size());
|
||||
r.atom[1] = static_cast<int>(p.atoms.size());
|
||||
r.sigma = static_cast<float>(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<std::set<int>> 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<std::vector<int>> first(copy.chains.size());
|
||||
std::vector<gemmi::Position> 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<int>(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<int>(i) >= j ||
|
||||
std::binary_search(topo.excluded[i].begin(), topo.excluded[i].end(), j)))
|
||||
return;
|
||||
if (m.image_idx != 0 && static_cast<int>(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<int>(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<int>(RestraintKind::Repulsion);
|
||||
r.atom[0] = static_cast<int>(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<float>(rot[a][b]);
|
||||
r.tran[a] = static_cast<float>(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<std::array<double, 4>> &xyzb) {
|
||||
GeometryStatistics s;
|
||||
double sb = 0, sa = 0, sc = 0;
|
||||
int nc = 0;
|
||||
for (const RestraintRecord &r : geometry.records) {
|
||||
const RestraintKind kind = static_cast<RestraintKind>(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;
|
||||
}
|
||||
@@ -0,0 +1,98 @@
|
||||
// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
|
||||
// SPDX-License-Identifier: GPL-3.0-only
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <array>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
#include <memory>
|
||||
|
||||
#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<RestraintRecord> records;
|
||||
std::vector<int> plane_atoms;
|
||||
int slots = 0;
|
||||
std::vector<int> slot_start; // n_atoms + 1
|
||||
std::vector<int> 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<int> 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<std::array<double, 4>> &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<std::array<double, 4>> &xyzb,
|
||||
std::vector<std::array<double, 4>> *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<char> 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<std::vector<int>> excluded;
|
||||
std::vector<float> vdw_radius; // A
|
||||
std::vector<char> hbond; // donor, acceptor or both (water included)
|
||||
std::vector<char> metal; // a metal ion
|
||||
std::vector<char> altloc;
|
||||
std::vector<gemmi::El> element;
|
||||
// Per atom: the heavy atom a hydrogen rides on, -1 for every other atom.
|
||||
std::vector<int> rider_parent;
|
||||
// The topology itself and the dictionary it points into, for placing the hydrogens again.
|
||||
std::shared_ptr<gemmi::MonLib> monlib;
|
||||
std::shared_ptr<gemmi::Topo> topo;
|
||||
int restrained_residues = 0;
|
||||
int fixed_residues = 0;
|
||||
std::vector<std::string> 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<std::array<double, 4>> &xyzb);
|
||||
@@ -0,0 +1,396 @@
|
||||
// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
|
||||
// SPDX-License-Identifier: GPL-3.0-only
|
||||
|
||||
#include "RefinementXray.h"
|
||||
|
||||
#include <algorithm>
|
||||
#include <cmath>
|
||||
#include <map>
|
||||
|
||||
#include <fftw3.h>
|
||||
|
||||
#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<float>;
|
||||
|
||||
// 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<Table, float> Calculator(const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg,
|
||||
double d_min, double blur) {
|
||||
gemmi::DensityCalculator<Table, float> 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<double> &fc,
|
||||
const std::vector<size_t> &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<double>(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 <typename F>
|
||||
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<std::complex<float>> &f) override {
|
||||
std::vector<std::complex<double>> fd;
|
||||
Compose(model, fd);
|
||||
f.assign(fd.begin(), fd.end());
|
||||
}
|
||||
|
||||
void Fmask(const gemmi::Model &model, std::vector<std::complex<float>> &f) override {
|
||||
gemmi::Grid<float> mask = grid_;
|
||||
PutMaskOnGrid(mask, model, orbit_leaders_, nthreads_);
|
||||
const gemmi::FPhiGrid<float> 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<std::array<double, 4>> *grad) override {
|
||||
std::vector<std::complex<double>> 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<std::complex<double>> 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<double> fm = k * (f[m] + std::complex<double>(rows_.fmask[m]));
|
||||
const double fc = std::abs(fm);
|
||||
double dfc = 0;
|
||||
sum += refine_math::AmplitudeNegLogLikelihood<double>(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<size_t>(nu) * nv * nwh;
|
||||
std::complex<float> *in = reinterpret_cast<std::complex<float> *>(fftwf_alloc_complex(ncomplex));
|
||||
float *out = fftwf_alloc_real(static_cast<size_t>(nu) * nv * nw);
|
||||
std::fill(in, in + ncomplex, std::complex<float>(0, 0));
|
||||
const std::vector<SymmetryComposition::Term> &terms = composition_.Terms();
|
||||
const size_t ops = composition_.Ops();
|
||||
auto add = [&](int ku, int kv, int kw, std::complex<double> value) {
|
||||
in[gemmi::modulo(ku, nu) + static_cast<size_t>(nu) * (gemmi::modulo(kv, nv) + static_cast<size_t>(nv) * kw)] +=
|
||||
std::complex<float>(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<double> 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<double> 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<float> map = grid_;
|
||||
const float norm = static_cast<float>(cell_.volume / static_cast<double>(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<std::complex<double>> &f) const {
|
||||
gemmi::DensityCalculator<Table, float> dc = Calculator(cell_, sg_, d_min_, blur_);
|
||||
PutModelDensityOnGrid(dc, model, {}, nthreads_); // one copy: the symmetry is composed below
|
||||
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<float> grid_;
|
||||
SymmetryComposition composition_;
|
||||
std::vector<size_t> orbit_leaders_;
|
||||
double blur_ = 0;
|
||||
XrayRowModel rows_;
|
||||
};
|
||||
|
||||
} // namespace
|
||||
|
||||
RefinementReflections PrepareRefinementReflections(const std::vector<MergedReflection> &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<gemmi::Miller, const MergedReflection *> 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<gemmi::Miller> 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<float>(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<std::complex<float>> &fmodel,
|
||||
XrayRowModel &rows) {
|
||||
const size_t n = refl.Size();
|
||||
std::vector<double> fc(n);
|
||||
for (size_t i = 0; i < n; i++)
|
||||
fc[i] = std::abs(fmodel[i]);
|
||||
std::vector<size_t> 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<size_t>(free.size() / FREE_PER_BIN, 1, MAX_BINS);
|
||||
|
||||
std::vector<double> centre(bins), bin_d(bins), bin_beta(bins);
|
||||
for (size_t b = 0; b < bins; b++) {
|
||||
const std::vector<size_t> 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<double>(idx.size());
|
||||
double d = sfc2 > 0 ? std::max(sfofc / sfc2, 0.05) : 1.0;
|
||||
double beta = std::max((sfo - d * d * sfc2) / static_cast<double>(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<float>(d);
|
||||
rows.var[i] = static_cast<float>(refl.epsilon[i] * beta + double(refl.sigma[i]) * refl.sigma[i]);
|
||||
}
|
||||
}
|
||||
|
||||
RFactors ComputeRFactors(const RefinementReflections &refl, const std::vector<std::complex<float>> &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<Table, float> 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<RefinementEngine> MakeRefinementEngineCPU(const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg,
|
||||
double d_min, const RefinementReflections &refl,
|
||||
size_t nthreads) {
|
||||
return std::make_unique<RefinementEngineCPU>(cell, sg, d_min, refl, nthreads);
|
||||
}
|
||||
|
||||
void AtomGradientsFromMap(const gemmi::Grid<float> &map, const gemmi::Model &model, double blur, size_t nthreads,
|
||||
std::vector<std::array<double, 4>> &grad) {
|
||||
const gemmi::DensityCalculator<Table, float> dc = Calculator(map.unit_cell, *map.spacegroup, 1.0, blur);
|
||||
std::vector<const gemmi::Atom *> atoms;
|
||||
for (const gemmi::Chain &ch : model.chains)
|
||||
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<float>(om.a11), m01 = static_cast<float>(om.a12), m02 = static_cast<float>(om.a13),
|
||||
m11 = static_cast<float>(om.a22), m12 = static_cast<float>(om.a23), m22 = static_cast<float>(om.a33);
|
||||
constexpr double FOUR_PI2 = 4 * refine_math::PI * refine_math::PI;
|
||||
ParallelChunks(static_cast<int>(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<float>(atom.b_iso + blur);
|
||||
const auto precal = Table::get(atom.element, atom.charge, atom.serial).precalculate_density_iso(b, 0.0f);
|
||||
const double radius = dc.estimate_radius(precal, b);
|
||||
const float max_r2 = static_cast<float>(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<float>(n3[k]);
|
||||
c[k] = gemmi::iround(n3[k]);
|
||||
}
|
||||
d[0] = std::min(static_cast<int>(std::ceil(radius / map.spacing[0])), nu - 1);
|
||||
d[1] = std::min(static_cast<int>(std::ceil(radius / map.spacing[1])), nv - 1);
|
||||
d[2] = std::min(static_cast<int>(std::ceil(radius / map.spacing[2])), nw - 1);
|
||||
double 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<float>(w);
|
||||
const float z = m22 * fz;
|
||||
const int w_ = gemmi::modulo(w, nw);
|
||||
for (int v = c[1] - d[1]; v <= c[1] + d[1]; v++) {
|
||||
const float fy = nctr[1] - static_cast<float>(v);
|
||||
const float y = m11 * fy + m12 * fz;
|
||||
const float yz = y * y + z * z;
|
||||
if (yz > max_r2)
|
||||
continue;
|
||||
const float x_yz = m01 * fy + m02 * fz;
|
||||
const size_t row = static_cast<size_t>(nu) * (gemmi::modulo(v, nv) + static_cast<size_t>(nv) * w_);
|
||||
for (int u = c[0] - d[0]; u <= c[0] + d[0]; u++) {
|
||||
const float x = m00 * (nctr[0] - static_cast<float>(u)) + x_yz;
|
||||
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};
|
||||
}
|
||||
});
|
||||
}
|
||||
@@ -0,0 +1,117 @@
|
||||
// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
|
||||
// SPDX-License-Identifier: GPL-3.0-only
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <array>
|
||||
#include <complex>
|
||||
#include <memory>
|
||||
#include <optional>
|
||||
#include <vector>
|
||||
|
||||
#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<gemmi::Miller> hkl;
|
||||
std::vector<float> fo, sigma;
|
||||
std::vector<char> free, centric;
|
||||
std::vector<float> epsilon;
|
||||
std::vector<double> 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<MergedReflection> &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<float> k_total;
|
||||
std::vector<std::complex<float>> fmask;
|
||||
std::vector<float> 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<std::complex<float>> &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<std::complex<float>> &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<std::complex<float>> &f) = 0;
|
||||
// The bulk-solvent mask's structure factors at the reflections (unscaled).
|
||||
virtual void Fmask(const gemmi::Model &model, std::vector<std::complex<float>> &f) = 0;
|
||||
// The sum of -ln L over the working reflections and, with `grad`, its derivatives with respect to every
|
||||
// atom's x, y, z (Cartesian, A) and B (A^2), in model order.
|
||||
virtual double Evaluate(const gemmi::Model &model, std::vector<std::array<double, 4>> *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<std::array<double, 4>> &xyzb,
|
||||
std::vector<std::array<double, 4>> *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<RefinementEngine> 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<float> &map, const gemmi::Model &model, double blur, size_t nthreads,
|
||||
std::vector<std::array<double, 4>> &grad);
|
||||
@@ -184,6 +184,7 @@ ADD_EXECUTABLE(jfjoch_test
|
||||
MolecularReplacementTest.cpp
|
||||
ModelMaskGPUTest.cpp
|
||||
ModelDensityGPUTest.cpp
|
||||
RefinementTest.cpp
|
||||
SearchSpaceGroupTest.cpp
|
||||
SearchSpaceGroupTwinTest.cpp
|
||||
SyntheticMergedReflections.h
|
||||
|
||||
@@ -0,0 +1,452 @@
|
||||
// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
|
||||
// SPDX-License-Identifier: GPL-3.0-only
|
||||
|
||||
#include <catch2/catch_all.hpp>
|
||||
|
||||
#include <cmath>
|
||||
#include <map>
|
||||
#include <random>
|
||||
#include <vector>
|
||||
|
||||
#include <gemmi/symmetry.hpp>
|
||||
|
||||
#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<std::string> &codes, const gemmi::UnitCell &cell, const char *group,
|
||||
unsigned seed) {
|
||||
std::mt19937 rng(seed);
|
||||
std::uniform_real_distribution<double> 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<double> 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<float>(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<MergedReflection> asu_reflections(const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg, double d_min) {
|
||||
std::vector<MergedReflection> 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<float>(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<std::complex<float>> 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<double> 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<float>(std::max(5.0, a.b_iso + db(rng)));
|
||||
}
|
||||
return out;
|
||||
}
|
||||
|
||||
std::vector<gemmi::Atom *> atoms_of(gemmi::Model &model) {
|
||||
std::vector<gemmi::Atom *> 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<std::array<double, 4>> grad;
|
||||
engine.Evaluate(model, &grad);
|
||||
std::vector<gemmi::Atom *> 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<float>(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<double>(4.0, fc, 0.8, 2.5, centric, dfc);
|
||||
const double h = 1e-5;
|
||||
const double fp = refine_math::AmplitudeNegLogLikelihood<double>(4.0, fc + h, 0.8, 2.5, centric, dummy);
|
||||
const double fm = refine_math::AmplitudeNegLogLikelihood<double>(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<std::array<double, 4>> 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<std::array<double, 4>> 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<int>(RestraintKind::BSimilarity); kind++) {
|
||||
RestraintSet one;
|
||||
for (const RestraintRecord &r : set.records)
|
||||
if (r.kind == kind) {
|
||||
RestraintRecord copy = r;
|
||||
if (kind == static_cast<int>(RestraintKind::Plane)) {
|
||||
copy.atom[0] = static_cast<int>(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<int>(RestraintKind::Plane) ? r.atom[1] : 4);
|
||||
}
|
||||
CAPTURE(kind, one.records.size());
|
||||
std::vector<std::array<double, 4>> 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<MergedReflection> 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<std::array<double, 4>> ggpu;
|
||||
const double egpu = gpu->EvaluateRestraints(0, xyzb, &ggpu);
|
||||
std::vector<std::array<double, 4>> 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<std::complex<float>> 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<std::complex<float>> fg;
|
||||
gpu->Fcalc(model, fg);
|
||||
double num = 0, den = 0;
|
||||
for (size_t m = 0; m < f.size(); m++) {
|
||||
num += std::abs(f[m] - fg[m]);
|
||||
den += std::abs(f[m]);
|
||||
}
|
||||
CHECK(num / den < 1e-3);
|
||||
|
||||
std::vector<std::array<double, 4>> gc, gg;
|
||||
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<MergedReflection> 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<std::complex<float>> f;
|
||||
cpu->Fcalc(truth.models[0], f);
|
||||
std::map<gemmi::Miller, double> 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<double> 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<float>(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<gemmi::Atom *> 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<gemmi::Atom *> 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);
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
|
||||
@@ -0,0 +1,151 @@
|
||||
// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
|
||||
// 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 <cmath>
|
||||
#include <cstdlib>
|
||||
#include <iostream>
|
||||
#include <sstream>
|
||||
|
||||
#include <gemmi/mmread_gz.hpp>
|
||||
#include <gemmi/mtz.hpp>
|
||||
#include <gemmi/to_cif.hpp>
|
||||
#include <gemmi/to_mmcif.hpp>
|
||||
#include <fstream>
|
||||
#include <nlohmann/json.hpp>
|
||||
|
||||
#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] <merged.mtz> <model.cif|.pdb>\n"
|
||||
" -N <threads> CPU threads (default 8)\n"
|
||||
" --no-gpu the CPU engine\n"
|
||||
" --placement-only rigid bodies and group B only\n"
|
||||
" --cycles <n> macro-cycles per weight (default as in RefinementSettings)\n"
|
||||
" --xyz-iterations <n>, --b-iterations <n> L-BFGS iterations per macro-cycle\n"
|
||||
" --ladder <a,b,...> the weight ladder, times the gradient-balanced weight\n"
|
||||
" -o <file.cif> write the refined model\n";
|
||||
}
|
||||
|
||||
std::vector<double> ParseList(const std::string &s) {
|
||||
std::vector<double> 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<std::string> 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<MergedReflection> merged;
|
||||
for (size_t r = 0; r < static_cast<size_t>(mtz.nreflections); r++) {
|
||||
const float *row = &mtz.data[r * mtz.columns.size()];
|
||||
MergedReflection m;
|
||||
m.h = static_cast<int>(row[ih]);
|
||||
m.k = static_cast<int>(row[ik]);
|
||||
m.l = static_cast<int>(row[il]);
|
||||
m.F = row[iF];
|
||||
m.sigmaF = row[isF];
|
||||
m.rfree_flag = row[ifree] == 0;
|
||||
m.d = static_cast<float>(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<int>(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;
|
||||
}
|
||||
Reference in new issue
Block a user