mx/Jungfraujoch
3
2 Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases 90 Wiki Activity
Files
rc183
Jungfraujoch/tests/RefinementTest.cpp
T
leonarski_fandClaude Opus 5.5 08594fe6c9 rugnux: fill 2mFo-DFc at the unmeasured reflections with D Fc, keep an unfilled copy 
FWT/PHWT in <prefix>_maps.mtz and <prefix>_refined_maps.mtz now carries, beside
the observed reflections, every reflection of the data's resolution range that the
data do not have, at D Fc - its expected value given the model - with D that of
the observed reflection nearest in resolution. This is the convention of
phenix.refine (fill_missing_f_obs) and BUSTER: a missing reflection left at zero
biases the map towards it, which at atomic resolution shows as broader, lower
atoms where completeness is short. The filled rows have FP, FOM, FREE and the
difference coefficients missing. FWT_no_fill/PHWT_no_fill is the same map on the
observed reflections only (phenix's 2FOFCWT_no_fill), and DELFWT/PHDELWT is never
filled. For the refined model, F_model at the filled reflections is computed on
the CPU with the bulk solvent and scale of the last fit.

The map-coefficient tests check the filled FWT and the missing FWT_no_fill/DELFWT
at unmeasured reflections. The R-factors, the blobs and the maps the refinement reads are unchanged. On an
86%-complete 0.99 A set the median FWHM of well-ordered carbons drops from 1.061
to 1.034 A and their peak rises by 4%; on a 99%-complete set nothing changes.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01UqXYdL4qBGYRiSXSGXfCLu
2026-11-05 16:48:36 +01:00

1778 lines
83 KiB
C++
Raw Permalink Blame History

// 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 <algorithm>
#include <atomic>
#include <chrono>
#include <cmath>
#include <map>
#include <random>
#include <set>
#include <vector>
#include <gemmi/mmread_gz.hpp>
#include <gemmi/neighbor.hpp>
#include <gemmi/modify.hpp>
#include <gemmi/symmetry.hpp>
#include "../common/CUDAWrapper.h"
#include "../common/Logger.h"
#include "../common/UnitCell.h"
#include "../image_analysis/scale_merge/RfreeFlags.h"
#include "../image_analysis/structure_refinement/ModelGrid.h"
#include "../image_analysis/structure_refinement/MonomerLibrary.h"
#include "../image_analysis/structure_refinement/Refinement.h"
#include "../image_analysis/structure_refinement/RefinementADP.h"
#include "../image_analysis/structure_refinement/RefinementMath.h"
#include "../image_analysis/structure_refinement/RefinementMinimiser.h"
#include "../image_analysis/structure_refinement/RefinementNcs.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);
}
// A zinc bound by four cysteines (the dictionary's own ZN-CYS link) and a magnesium by six waters (links made from
// the library's metal-site statistics), both from LINK records: the bonds come out at the expected distances,
// and a metal set off its site returns there under them alone.
TEST_CASE("RefinementRestraints_metal_links", "[Refinement]") {
const gemmi::UnitCell cell(40, 40, 40, 90, 90, 90);
gemmi::Structure st = residue_crystal({"CYS", "CYS", "CYS", "CYS", "ZN", "HOH", "HOH", "HOH", "HOH", "HOH", "HOH",
"MG"}, cell, "P 1", 11);
double mg_o = 0, mg_sigma = 0;
REQUIRE(EmbeddedMetalDistance(gemmi::El::Mg, gemmi::El::O, 6, mg_o, mg_sigma));
CHECK(mg_o == Catch::Approx(2.07).margin(0.02)); // Mg-O of an octahedral site
constexpr double ZN_S = 2.34; // the ZN-CYS link
std::vector<gemmi::Residue> &res = st.models[0].chains[0].residues;
// the donors on a tetrahedron and an octahedron around their sites, each residue moved as a whole
const gemmi::Position zn_site(10, 10, 10), mg_site(25, 25, 25);
const double t = 1.0 / std::sqrt(3.0);
const gemmi::Position tetra[4] = {{t, t, t}, {t, -t, -t}, {-t, t, -t}, {-t, -t, t}};
const gemmi::Position octa[6] = {{1, 0, 0}, {-1, 0, 0}, {0, 1, 0}, {0, -1, 0}, {0, 0, 1}, {0, 0, -1}};
auto donor_to = [](gemmi::Residue &r, const std::string &atom, const gemmi::Position &target) {
const gemmi::Position shift = target - r.find_atom(atom, '*')->pos;
for (gemmi::Atom &a : r.atoms)
a.pos += shift;
};
for (int k = 0; k < 4; k++)
donor_to(res[k], "SG", zn_site + tetra[k] * ZN_S);
for (int k = 0; k < 6; k++)
donor_to(res[5 + k], "O", mg_site + octa[k] * mg_o);
// the metals off their sites by 0.3 A
res[4].atoms[0].pos = zn_site + gemmi::Position(0.2, -0.15, 0.17);
res[11].atoms[0].pos = mg_site + gemmi::Position(-0.1, 0.25, 0.12);
auto link = [&](int metal, int donor, const char *atom) {
gemmi::Connection c;
c.type = gemmi::Connection::MetalC;
c.asu = gemmi::Asu::Same;
c.partner1 = gemmi::AtomAddress("A", res[donor].seqid, res[donor].name, atom);
c.partner2 = gemmi::AtomAddress("A", res[metal].seqid, res[metal].name, res[metal].atoms[0].name);
st.connections.push_back(c);
};
for (int k = 0; k < 4; k++)
link(4, k, "SG");
for (int k = 0; k < 6; k++)
link(11, 5 + k, "O");
const RefinementTopology topo = BuildRefinementTopology(st);
CHECK(topo.metal_links == 6); // the Mg-O links; Zn-S is the dictionary's
std::vector<gemmi::Atom *> atoms = atoms_of(st.models[0]);
const int zn = static_cast<int>(std::find(atoms.begin(), atoms.end(), &res[4].atoms[0]) - atoms.begin());
const int mg = static_cast<int>(std::find(atoms.begin(), atoms.end(), &res[11].atoms[0]) - atoms.begin());
// the metals' bonds kept apart from the covalent ones, none left among them
for (const RestraintRecord &r : topo.geometry.records)
CHECK_FALSE((r.kind == static_cast<int>(RestraintKind::Bond) &&
(r.atom[0] == zn || r.atom[1] == zn || r.atom[0] == mg || r.atom[1] == mg)));
const RestraintSet &metal_bonds = topo.metal_bonds;
REQUIRE(metal_bonds.records.size() == 10);
for (const RestraintRecord &r : metal_bonds.records) {
const bool is_zn = r.atom[0] == zn || r.atom[1] == zn;
CHECK((is_zn || r.atom[0] == mg || r.atom[1] == mg));
CHECK(r.ideal == Catch::Approx(is_zn ? ZN_S : mg_o).margin(1e-4));
}
// the two metals moved down the gradient of their bonds, everything else held
std::vector<std::array<double, 4>> xyzb;
for (gemmi::Atom *a : atoms)
xyzb.push_back({a->pos.x, a->pos.y, a->pos.z, a->b_iso});
double e = EvaluateRestraints(metal_bonds, xyzb, nullptr), step = 1e-4;
for (int it = 0; it < 2000 && step > 1e-12; it++) {
std::vector<std::array<double, 4>> grad(xyzb.size(), {0, 0, 0, 0}), trial = xyzb;
EvaluateRestraints(metal_bonds, xyzb, &grad);
for (int a : {zn, mg})
for (int k = 0; k < 3; k++)
trial[a][k] -= step * grad[a][k];
const double e_trial = EvaluateRestraints(metal_bonds, trial, nullptr);
if (e_trial < e) {
xyzb = trial;
e = e_trial;
step *= 1.5;
} else
step *= 0.5;
}
for (const RestraintRecord &r : metal_bonds.records)
CHECK(RestraintValue(metal_bonds, r, xyzb) == Catch::Approx(r.ideal).margin(0.005));
}
// A magnesium with six waters at its inner-sphere distance and a seventh at 2.9 A, all seven from LINK records: the
// six become bonds, the seventh - a record at an outer-sphere distance - stays a contact, and a second topology of
// the same model (a second round of refinement) finds the same, the records keeping the names they were read with.
TEST_CASE("RefinementRestraints_metal_links_outer_sphere", "[Refinement]") {
const gemmi::UnitCell cell(40, 40, 40, 90, 90, 90);
gemmi::Structure st = residue_crystal({"HOH", "HOH", "HOH", "HOH", "HOH", "HOH", "HOH", "MG"}, cell, "P 1", 7);
std::vector<gemmi::Residue> &res = st.models[0].chains[0].residues;
const gemmi::Position site(20, 20, 20);
const gemmi::Position octa[6] = {{1, 0, 0}, {-1, 0, 0}, {0, 1, 0}, {0, -1, 0}, {0, 0, 1}, {0, 0, -1}};
for (int k = 0; k < 6; k++)
res[k].atoms[0].pos = site + octa[k] * 2.07;
res[6].atoms[0].pos = site + gemmi::Position(1, 1, 1) * (2.9 / std::sqrt(3.0));
res[7].atoms[0].pos = site;
for (int k = 0; k < 7; k++) {
gemmi::Connection c;
c.type = gemmi::Connection::MetalC;
c.asu = gemmi::Asu::Same;
c.partner1 = gemmi::AtomAddress("A", res[k].seqid, "HOH", "O");
c.partner2 = gemmi::AtomAddress("A", res[7].seqid, "MG", res[7].atoms[0].name);
st.connections.push_back(c);
}
std::vector<gemmi::Atom *> atoms = atoms_of(st.models[0]);
auto index = [&](const gemmi::Atom &a) {
return static_cast<int>(std::find(atoms.begin(), atoms.end(), &a) - atoms.begin());
};
const int mg = index(res[7].atoms[0]), outer = index(res[6].atoms[0]);
for (int round = 0; round < 2; round++) {
const RefinementTopology topo = BuildRefinementTopology(st);
CHECK(topo.metal_links == 6);
CHECK(topo.metal_links_refused == 1);
REQUIRE(topo.metal_bonds.records.size() == 6);
for (const RestraintRecord &r : topo.metal_bonds.records)
CHECK((r.atom[0] != outer && r.atom[1] != outer));
for (const gemmi::Connection &c : st.connections)
CHECK(c.link_id.empty());
// the outer water is a contact of the metal's (listed with a margin that reaches it)
const RestraintSet contacts = BuildRepulsions(topo, st.models[0], st.cell, *st.find_spacegroup(), 1.5);
int found = 0;
for (const RestraintRecord &r : contacts.records)
found += (r.atom[0] == mg && r.atom[1] == outer) || (r.atom[0] == outer && r.atom[1] == mg);
CHECK(found == 1);
}
}
// Deposited models cut down to a ligand and the residues within 4.2 A of it (test_data/ligands, from the PDB entries
// their names start with; the tRNA as its modified nucleotides and their neighbours in the chain): every ligand is
// restrained, its own restraints included, the deposited geometry agrees with the dictionary, and the links the
// ligands need are made - the RNA backbone to and from each modified nucleotide, the iron of a heme to its histidine
// (and a heme C's to its methionine) from the LINK records, and a heme C's two thioethers to its cysteines.
TEST_CASE("RefinementRestraints_deposited_ligands", "[Refinement]") {
const std::vector<std::string> modified_nucleotides = {"PSU", "5MC", "1MA", "OMG", "OMC", "H2U", "M2G", "7MG",
"2MG", "5MU"};
const std::vector<std::pair<std::string, std::vector<std::string>>> cases = {
{"1ehz_psu", modified_nucleotides}, {"8bkh_hem", {"HEM"}}, {"6k9i_hec", {"HEC"}}, {"2w3v_ndp", {"NDP"}},
{"5vml_nap", {"NAP"}}, {"6f3p_nad", {"NAD"}}, {"8con_nai", {"NAI"}}, {"3tk0_fad", {"FAD"}},
{"7tcd_fmn", {"FMN"}}, {"6w4h_sam", {"SAM"}}, {"5epe_sah", {"SAH"}}};
for (const auto &[name, ligands] : cases) {
DYNAMIC_SECTION(name) {
gemmi::Structure st = gemmi::read_structure_gz("../../tests/test_data/ligands/" + name + ".pdb.gz");
const RefinementTopology topo = BuildRefinementTopology(st);
std::vector<gemmi::Atom *> atoms = atoms_of(st.models[0]);
std::vector<const gemmi::Residue *> residue_of;
for (const gemmi::Chain &ch : st.models[0].chains)
for (const gemmi::Residue &r : ch.residues)
for (size_t k = 0; k < r.atoms.size(); k++)
residue_of.push_back(&r);
auto index_of = [&](const gemmi::Residue &r, const char *atom) {
const gemmi::Atom *a = r.find_atom(atom, '*');
return a == nullptr ? -1 : static_cast<int>(std::find(atoms.begin(), atoms.end(), a) - atoms.begin());
};
auto bonded = [&](const RestraintSet &set, int i, int j) {
for (const RestraintRecord &r : set.records)
if (r.kind == static_cast<int>(RestraintKind::Bond) &&
((r.atom[0] == i && r.atom[1] == j) || (r.atom[0] == j && r.atom[1] == i)))
return true;
return false;
};
std::map<std::string, int> restraints, found;
for (const RestraintRecord &r : topo.geometry.records)
if (r.kind != static_cast<int>(RestraintKind::Plane))
restraints[residue_of[r.atom[0]]->name]++;
for (size_t i = 0; i < atoms.size(); i++)
if (std::find(ligands.begin(), ligands.end(), residue_of[i]->name) != ligands.end()) {
found[residue_of[i]->name]++;
CHECK(topo.refinable[i]);
}
for (const std::string &code : ligands) {
INFO(code);
CHECK(found[code] > 0);
CHECK(restraints[code] > 0);
CHECK(std::find(topo.fixed_names.begin(), topo.fixed_names.end(), code) == topo.fixed_names.end());
}
std::vector<std::array<double, 4>> xyzb;
for (gemmi::Atom *a : atoms)
xyzb.push_back({a->pos.x, a->pos.y, a->pos.z, a->b_iso});
const GeometryStatistics g = MeasureGeometry(topo.geometry, {}, xyzb, &topo.metal_bonds);
CHECK(g.rmsd_bonds < 0.03);
CHECK(g.rmsd_angles < 3.0);
CHECK(g.rmsd_metal < 0.1);
if (name == "1ehz_psu") {
// each modified nucleotide bonded to the residue before it in the chain, where that one is there
int links = 0;
for (const gemmi::Residue &r : st.models[0].chains[0].residues) {
const gemmi::Residue *prev = &r == &st.models[0].chains[0].residues.front() ? nullptr : &r - 1;
if (prev == nullptr || prev->seqid.num != r.seqid.num - 1 ||
std::find(ligands.begin(), ligands.end(), r.name) == ligands.end())
continue;
INFO(r.name << " " << *r.seqid.num);
CHECK(bonded(topo.geometry, index_of(*prev, "O3'"), index_of(r, "P")));
links++;
}
CHECK(links >= 10);
}
if (name == "8bkh_hem" || name == "6k9i_hec") {
const gemmi::Residue *heme = nullptr, *his = nullptr;
for (const gemmi::Residue &r : st.models[0].chains[0].residues) {
if (r.name == ligands[0])
heme = &r;
if (r.name == "HIS" && r.seqid.num == (name == "8bkh_hem" ? 93 : 18))
his = &r;
}
REQUIRE(heme != nullptr);
REQUIRE(his != nullptr);
const int fe = index_of(*heme, "FE");
CHECK(topo.metal[fe]);
CHECK(topo.metal_links == (name == "8bkh_hem" ? 1 : 2)); // the histidine; and the methionine
CHECK(bonded(topo.metal_bonds, fe, index_of(*his, "NE2")));
CHECK(bonded(topo.metal_bonds, fe, index_of(*heme, "NA"))); // the porphyrin's own, among the metal's
if (name == "6k9i_hec")
for (auto [cys, carbon] : {std::pair{14, "CAB"}, std::pair{17, "CAC"}})
for (const gemmi::Residue &r : st.models[0].chains[0].residues)
if (r.name == "CYS" && r.seqid.num == cys)
CHECK(bonded(topo.geometry, index_of(r, "SG"), index_of(*heme, carbon)));
}
}
}
}
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
// (the engine is made for a cell wide enough for its gather; the restraints do not read the cell)
const gemmi::UnitCell wide_cell(40, 40, 40, 90, 90, 90);
std::vector<MergedReflection> merged = asu_reflections(wide_cell, *st.find_spacegroup(), 3.0);
RefinementReflections refl = PrepareRefinementReflections(merged, wide_cell, *st.find_spacegroup(), 3.0);
auto gpu = MakeRefinementEngineGPU(0, wide_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
}
// Contacts with other copies of the molecule: in P1 two waters 2.4 A apart across the a face (17.6 A apart within
// the cell) repel each other once, and a water 1.4 A from a 3-fold axis, 2.4 A from its images x -> T x and
// x -> T^-1 x (one contact per asymmetric unit), repels its image once.
TEST_CASE("RefinementRestraints_lattice_and_symmetry_contacts", "[Refinement]") {
auto oxygens = [](gemmi::Structure &st, const std::vector<gemmi::Fractional> &where) {
std::vector<gemmi::Atom *> atoms = atoms_of(st.models[0]);
REQUIRE(atoms.size() == where.size());
for (size_t i = 0; i < where.size(); i++)
atoms[i]->pos = st.cell.orthogonalize(where[i]);
};
{
const gemmi::UnitCell cell(20, 21, 22, 90, 90, 90);
gemmi::Structure st = residue_crystal({"HOH", "HOH"}, cell, "P 1", 1);
oxygens(st, {gemmi::Fractional(0.01, 0.5, 0.5), gemmi::Fractional(0.89, 0.5, 0.5)});
const RefinementTopology topo = BuildRefinementTopology(st);
const RestraintSet rep = BuildRepulsions(topo, st.models[0], cell, *st.find_spacegroup(), 0.0);
REQUIRE(rep.records.size() == 1);
CHECK(rep.records[0].tran[0] == Catch::Approx(-20.0f));
CHECK(BuildBContacts(topo, st.models[0], cell, *st.find_spacegroup()).records.size() == 1);
}
{
const gemmi::UnitCell cell(20, 20, 22, 90, 90, 120);
gemmi::Structure st = residue_crystal({"HOH"}, cell, "P 3", 1);
oxygens(st, {cell.fractionalize(gemmi::Position(1.4, 0.0, 11.0))});
const RefinementTopology topo = BuildRefinementTopology(st);
const RestraintSet rep = BuildRepulsions(topo, st.models[0], cell, *st.find_spacegroup(), 0.0);
CHECK(rep.records.size() == 1);
}
}
// A complete model whose every atom is off by a Gaussian error of sigma_x per axis: D(s) = <cos(2 pi s.dx)> =
// exp(-2 pi^2 sigma_x^2 s^2) (Luzzati). With a negligible measurement error D comes out at that; with an error of
// half the rms amplitude of the outermost shell it falls below it there, where the error is a large part of the
// spread and is counted twice, and stays where the amplitudes are strong.
TEST_CASE("RefinementLikelihood_D_coordinate_error", "[Refinement]") {
const char *group = "P 1";
const gemmi::SpaceGroup &sg = *gemmi::find_spacegroup_by_name(group);
const gemmi::UnitCell cell(40, 45, 50, 90, 90, 90);
const double d_min = 1.5, sigma_x = 0.25;
gemmi::Structure truth = random_crystal(cell, group, 1500, 3);
gemmi::Model model = truth.models[0];
std::mt19937 rng(4);
std::normal_distribution<double> dx(0, sigma_x);
for (gemmi::Atom &a : model.chains[0].residues[0].atoms)
a.pos += gemmi::Position(dx(rng), dx(rng), dx(rng));
RefinementReflections refl = PrepareRefinementReflections(asu_reflections(cell, sg, d_min), cell, sg, d_min);
auto cpu = MakeRefinementEngineCPU(cell, sg, d_min, refl, 4);
cpu->SetBlur(RefinementBlur(cell, sg, d_min, 5.0));
std::vector<std::complex<float>> f;
cpu->Fcalc(truth.models[0], f);
for (size_t m = 0; m < refl.Size(); m++)
refl.fo[m] = std::abs(f[m]);
cpu->Fcalc(model, f);
// mean D and mean Luzzati D in five shells of 1/d^2
constexpr int SHELLS = 5;
const double s2_max = 1.0 / (d_min * d_min);
auto shells = [&](const XrayRowModel &rows, std::vector<double> &d, std::vector<double> &luzzati) {
std::vector<int> n(SHELLS, 0);
d.assign(SHELLS, 0);
luzzati.assign(SHELLS, 0);
for (size_t m = 0; m < refl.Size(); m++) {
const int k = std::min(SHELLS - 1, static_cast<int>(refl.inv_d2[m] / s2_max * SHELLS));
d[k] += rows.d[m];
luzzati[k] += std::exp(-2 * M_PI * M_PI * sigma_x * sigma_x * refl.inv_d2[m]);
n[k]++;
}
for (int k = 0; k < SHELLS; k++) {
d[k] /= n[k];
luzzati[k] /= n[k];
}
};
XrayRowModel rows;
rows.k_total.assign(refl.Size(), 1.0f);
rows.fmask.assign(refl.Size(), {0, 0});
for (float &s : refl.sigma)
s = 0.01f;
EstimateLikelihoodParameters(refl, f, rows);
std::vector<double> d_clean, luzzati;
shells(rows, d_clean, luzzati);
for (int k = 0; k < SHELLS; k++) {
CAPTURE(k, d_clean[k], luzzati[k]);
CHECK(std::fabs(d_clean[k] - luzzati[k]) < 0.05);
}
CHECK(luzzati[SHELLS - 1] < 0.65);
// a measurement error of half the rms amplitude of the outermost shell
double fo2 = 0;
int n = 0;
for (size_t m = 0; m < refl.Size(); m++)
if (refl.inv_d2[m] > s2_max * (SHELLS - 1) / SHELLS) {
fo2 += double(refl.fo[m]) * refl.fo[m];
n++;
}
for (float &s : refl.sigma)
s = static_cast<float>(0.5 * std::sqrt(fo2 / n));
EstimateLikelihoodParameters(refl, f, rows);
std::vector<double> d_noisy;
shells(rows, d_noisy, luzzati);
CHECK(d_noisy[0] > d_clean[0] - 0.03);
CHECK(d_noisy[SHELLS - 1] < d_clean[SHELLS - 1] - 0.1);
}
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("RefinementMinimiser_quadratic", "[Refinement]") {
// f = sum a_k (x_k - 1)^2 with curvatures spread over two decades
const size_t n = 50;
std::vector<double> a(n);
for (size_t k = 0; k < n; k++)
a[k] = std::pow(100.0, static_cast<double>(k) / (n - 1));
const RefinementMinimiser::Target target = [&](const std::vector<double> &x, double &f, std::vector<double> &g) {
f = 0;
for (size_t k = 0; k < n; k++) {
f += a[k] * (x[k] - 1) * (x[k] - 1);
g[k] = 2 * a[k] * (x[k] - 1);
}
return true;
};
auto distance = [&](const std::vector<double> &x) {
double d = 0;
for (double v : x)
d = std::max(d, std::fabs(v - 1));
return d;
};
// the exact diagonal: one step
RefinementMinimiser exact;
std::vector<double> h(n);
for (size_t k = 0; k < n; k++)
h[k] = 2 * a[k];
exact.SetDiagonal(h);
std::vector<double> x(n, 0.0);
exact.Minimise(target, x, 1, 0.0, 10.0);
CHECK(distance(x) < 1e-9);
// without it, the memory carried from one call to the next does what one long call does
RefinementMinimiser split, whole;
std::vector<double> xs(n, 0.0), xw(n, 0.0);
for (int call = 0; call < 6; call++)
split.Minimise(target, xs, 10, 0.0, 10.0);
whole.Minimise(target, xw, 60, 0.0, 10.0);
CHECK(distance(xw) < 1e-3);
for (size_t k = 0; k < n; k++)
CHECK(xs[k] == Catch::Approx(xw[k]).margin(1e-12));
}
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.weight_ladder = {0.0625};
settings.adp_ladder = {4.0};
settings.max_steps = 2;
settings.final_ladder = {1.0};
settings.riding_hydrogens = false; // the data were made without them
settings.solvent = true;
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()); // the solvent build's placeholders not among them
for (const gemmi::Chain &ch : r.model->models[0].chains)
CHECK(ch.name != "~");
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);
// the map coefficients as amplitude and phase: FWT, FWT_no_fill and DELFWT non-negative, their phases in
// [0, 360); FWT filled where there is no FP, the other two missing there
REQUIRE(r.maps);
const auto &columns = r.maps->coefficients.columns;
REQUIRE(columns[6].label == "FWT");
REQUIRE(columns[8].label == "FWT_no_fill");
REQUIRE(columns[10].label == "DELFWT");
for (size_t c : {6, 8, 10})
for (size_t i = 0; i < columns[c].values.size(); i++) {
if (c != 6 && std::isnan(columns[3].values[i])) {
CHECK(std::isnan(columns[c].values[i]));
continue;
}
CHECK(columns[c].values[i] >= 0.0f);
CHECK(columns[c + 1].values[i] >= 0.0f);
CHECK(columns[c + 1].values[i] < 360.0f);
}
// one trial at a time instead of three side by side (sharing a card's tables, mask and scale fit): the same
// model, bit for bit
RefinementSettings one = settings;
one.max_trials_at_once = 1;
const RefinementResult s = RefineModel(start, merged, cell, sg, 1.0, logger, one);
REQUIRE(s.ok);
CHECK(s.r_free == r.r_free);
// the same trials kept, those run ahead where the slots take them left out
CHECK(s.trials.size() == r.trials.size());
CHECK(s.adp_trials.size() == r.adp_trials.size());
CHECK(s.adp_model == r.adp_model);
const std::vector<gemmi::Atom *> a_side = atoms_of(r.model->models[0]), a_one = atoms_of(s.model->models[0]);
REQUIRE(a_one.size() == a_side.size());
size_t same = 0;
for (size_t i = 0; i < a_one.size(); i++)
same += a_one[i]->pos.x == a_side[i]->pos.x && a_one[i]->pos.y == a_side[i]->pos.y &&
a_one[i]->pos.z == a_side[i]->pos.z && a_one[i]->b_iso == a_side[i]->b_iso;
CHECK(same == a_one.size());
// and the maps and what is read off them after it
CHECK(s.mean_fom == r.mean_fom);
CHECK(s.difference_map_rms == r.difference_map_rms);
REQUIRE(s.blobs.size() == r.blobs.size());
for (size_t i = 0; i < s.blobs.size(); i++) {
CHECK(s.blobs[i].centroid.dist(r.blobs[i].centroid) == 0.0);
CHECK(s.blobs[i].score == r.blobs[i].score);
}
REQUIRE(s.solvent);
REQUIRE(r.solvent);
CHECK(s.solvent->size() == r.solvent->size());
// cut short before the first weight trial: the model as placed comes back, no trial chosen
const std::atomic<bool> cancel{true};
settings.cancel = &cancel;
const RefinementResult c = RefineModel(start, merged, cell, sg, 1.0, logger, settings);
REQUIRE(c.ok);
CHECK(c.cancelled);
CHECK_FALSE(c.restrained);
CHECK(c.trials.empty());
CHECK(c.model != nullptr);
CHECK_FALSE(r.cancelled);
}
}
// Two copies of a tRNA: with a slot for every ladder trial, the NCS rungs run ahead at both B weights and the
// intermediate TLS cuts with the others, those not needed left out; one trial at a time runs only what is needed.
// The same trials kept and the same model, bit for bit.
TEST_CASE("Refinement_trials_run_ahead", "[Refinement]") {
gemmi::Structure truth = gemmi::read_structure_gz("../../tests/test_data/ligands/1ehz_psu.pdb.gz");
truth.models.resize(1);
gemmi::remove_hydrogens(truth.models[0]);
truth.cell = gemmi::UnitCell(120, 120, 120, 90, 90, 90);
truth.spacegroup_hm = "P 1";
truth.setup_cell_images();
gemmi::Chain copy = truth.models[0].chains.at(0);
copy.name = "B";
const double a = 0.7;
const gemmi::Mat33 rot(std::cos(a), -std::sin(a), 0, std::sin(a), std::cos(a), 0, 0, 0, 1);
for (gemmi::Residue &r : copy.residues)
for (gemmi::Atom &at : r.atoms)
at.pos = gemmi::Position(rot.multiply(at.pos)) + gemmi::Position(40, 10, -5);
truth.models[0].chains.push_back(copy);
truth.connections.clear(); // the LINK records name chain A only
gemmi::Structure start = truth;
start.models[0] = shaken(truth.models[0], 0.3, 5.0, 21);
const gemmi::SpaceGroup &sg = *gemmi::find_spacegroup_by_name("P 1");
const double d_min = 5.0;
std::vector<MergedReflection> merged = asu_reflections(truth.cell, sg, d_min);
{
RefinementReflections refl = PrepareRefinementReflections(merged, truth.cell, sg, d_min * 0.999);
auto cpu = MakeRefinementEngineCPU(truth.cell, sg, d_min * 0.999, refl, 8);
cpu->SetBlur(RefinementBlur(truth.cell, sg, d_min, 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]);
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);
r.sigmaF = 0.03f * r.F + 0.01f;
}
}
Logger logger("RefinementTest");
RefinementSettings settings;
settings.gpu = false;
settings.nthreads = 8;
settings.weight_ladder = {1.0 / 128, 1.0 / 16};
settings.macro_cycles = 1;
settings.max_steps = 1;
settings.final_ladder = {1.0};
settings.riding_hydrogens = false;
settings.solvent = false;
const RefinementResult ahead = RefineModel(start, merged, truth.cell, sg, 1.0, logger, settings);
RefinementSettings one = settings;
one.max_trials_at_once = 1;
const RefinementResult needed = RefineModel(start, merged, truth.cell, sg, 1.0, logger, one);
REQUIRE(ahead.ok);
REQUIRE(needed.ok);
CHECK(ahead.ncs_copies == 2);
CHECK(ahead.trials.size() == needed.trials.size());
CHECK(ahead.adp_trials.size() == needed.adp_trials.size());
CHECK(ahead.adp_model == needed.adp_model);
CHECK(ahead.r_free == needed.r_free);
std::vector<gemmi::Atom *> x = atoms_of(ahead.model->models[0]), y = atoms_of(needed.model->models[0]);
REQUIRE(x.size() == y.size());
size_t same = 0;
for (size_t i = 0; i < x.size(); i++)
same += x[i]->pos.x == y[i]->pos.x && x[i]->pos.y == y[i]->pos.y && x[i]->pos.z == y[i]->pos.z &&
x[i]->b_iso == y[i]->b_iso;
CHECK(same == x.size());
}
// A K+ on the 4-fold axis of P4 at occupancy 0.25 (its four coinciding copies one ion), 0.2 A off the axis in the
// model as given, with a water 2.8 A from it: the topology puts the ion on the axis, its contact with the water is
// listed once (not once per copy) and none with itself; refinement keeps it on the axis at the occupancy given,
// while the rest of the model moves.
TEST_CASE("Refinement_special_position", "[Refinement]") {
const gemmi::UnitCell cell(30, 30, 40, 90, 90, 90);
const gemmi::SpaceGroup &sg = *gemmi::find_spacegroup_by_name("P 4");
gemmi::Structure truth = residue_crystal({"PHE", "TRP", "TYR", "HOH", "HOH", "K"}, cell, "P 4", 5);
std::vector<gemmi::Residue> &res = truth.models[0].chains[0].residues;
const gemmi::Position axis(0, 0, 20);
res[5].atoms[0].pos = axis;
res[5].atoms[0].occ = 0.25f;
res[4].atoms[0].pos = axis + gemmi::Position(2.8, 0, 0.3);
gemmi::Structure start = truth;
start.models[0] = shaken(truth.models[0], 0.25, 5.0, 6);
gemmi::Atom &k_start = start.models[0].chains[0].residues[5].atoms[0];
k_start.pos = axis + gemmi::Position(0.2, 0, 0.1);
{
gemmi::Structure st = start;
const std::vector<SpecialPosition> sp = FindSpecialPositions(st);
REQUIRE(sp.size() == 1);
CHECK(sp[0].multiplicity == 4);
const RefinementTopology topo = BuildRefinementTopology(st);
const std::vector<gemmi::Atom *> atoms = atoms_of(st.models[0]);
const int k = static_cast<int>(atoms.size()) - 1, water = k - 1;
REQUIRE(topo.special.size() == 1);
CHECK(topo.special[0].atom == k);
CHECK(atoms[k]->pos.x == Catch::Approx(0).margin(1e-9));
CHECK(atoms[k]->pos.y == Catch::Approx(0).margin(1e-9));
CHECK(atoms[k]->pos.z == Catch::Approx(20.1));
const RestraintSet rep = BuildRepulsions(topo, st.models[0], cell, sg, 1.5);
int with_water = 0, with_itself = 0;
for (const RestraintRecord &r : rep.records) {
with_water += (r.atom[0] == k && r.atom[1] == water) || (r.atom[0] == water && r.atom[1] == k);
with_itself += r.atom[0] == k && r.atom[1] == k;
}
CHECK(with_water == 1);
CHECK(with_itself == 0);
// the water's neighbours: the ion once, not four times
gemmi::UnitCell images = cell;
images.set_cell_images_from_spacegroup(&sg);
gemmi::NeighborSearch ns(st.models[0], images, 5.0);
ns.populate(false);
int ion = 0;
for (const ModelNeighbour &n : AtomsNear(ns, st.models[0], atoms[water]->pos, 4.0))
ion += n.residue == 5;
CHECK(ion == 1);
}
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]);
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);
r.sigmaF = 0.03f * r.F + 0.01f;
}
}
Logger logger("RefinementTest");
RefinementSettings settings;
settings.gpu = false;
settings.nthreads = 4;
settings.weight_ladder = {0.0625};
settings.adp_ladder = {4.0};
settings.max_steps = 2;
settings.final_ladder = {1.0};
settings.riding_hydrogens = false;
settings.solvent = false;
const RefinementResult r = RefineModel(start, merged, cell, sg, 1.0, logger, settings);
REQUIRE(r.ok);
REQUIRE(r.restrained);
const gemmi::Atom &k = r.model->models[0].chains[0].residues[5].atoms[0];
CHECK(std::hypot(k.pos.x, k.pos.y) < 1e-6);
CHECK(k.pos.z == Catch::Approx(20).margin(0.15));
CHECK(k.occ == 0.25f);
CHECK(r.r_free < 0.6 * r.r_free_start);
}
// Single-atom ions scatter as the ions they are (10 electrons for Mg2+, 18 for Cl-), every other atom as the
// neutral atom whatever charge the file gives it; the CPU and the GPU engines score such a model alike.
TEST_CASE("Refinement_ion_form_factors", "[Refinement]") {
CHECK(IonFormFactor<float>(gemmi::El::Mg, 2).calculate_sf(0.0f) == Catch::Approx(10.0).margin(0.02));
CHECK(IonFormFactor<float>(gemmi::El::Mg, 0).calculate_sf(0.0f) == Catch::Approx(12.0).margin(0.02));
CHECK(IonFormFactor<double>(gemmi::El::Cl, -1).calculate_sf(0.0) == Catch::Approx(18.0).margin(0.02));
CHECK(IonFormFactor<float>(gemmi::El::N, 1).calculate_sf(0.0f) == Catch::Approx(7.0).margin(0.02)); // none tabulated
const gemmi::UnitCell cell(30, 32, 34, 90, 90, 90);
const gemmi::SpaceGroup &sg = *gemmi::find_spacegroup_by_name("P 21 21 21");
gemmi::Structure st = residue_crystal({"LYS", "TRP", "MG", "CL", "HOH"}, cell, "P 21 21 21", 3);
for (gemmi::Atom &a : st.models[0].chains[0].residues[0].atoms)
if (a.name == "NZ")
a.charge = 1; // as a file may give it
std::vector<MergedReflection> merged = asu_reflections(cell, sg, 2.0);
std::mt19937 rng(4);
std::uniform_real_distribution<float> f(10.0f, 100.0f);
for (MergedReflection &r : merged) {
r.F = f(rng);
r.sigmaF = 1.0f;
}
Logger logger("RefinementTest");
RefinementSettings settings;
settings.score_only = true;
settings.nthreads = 4;
settings.gpu = false;
const RefinementResult cpu = RefineModel(st, merged, cell, sg, 1.0, logger, settings);
REQUIRE(cpu.ok);
for (const gemmi::Residue &r : cpu.model->models[0].chains[0].residues)
for (const gemmi::Atom &a : r.atoms)
CHECK(static_cast<int>(a.charge) == (r.name == "MG" ? 2 : r.name == "CL" ? -1 : 0));
if (get_gpu_count() > 0) {
settings.gpu = true;
const RefinementResult gpu = RefineModel(st, merged, cell, sg, 1.0, logger, settings);
REQUIRE(gpu.ok);
CHECK(gpu.on_gpu);
CHECK(gpu.r_work_start == Catch::Approx(cpu.r_work_start).margin(1e-4));
}
}
// Two copies of one RNA chain in the asymmetric unit, the second turned and moved, a fragment of it and the chain in
// another conformation: the two found as copies of one molecule, the fragment and the other conformation not, held
// to each other by local distance restraints whose gradient is right on the CPU and the GPU, and let go where the
// copies really differ. A single copy has none.
TEST_CASE("RefinementNcs_two_copies", "[Refinement]") {
gemmi::Structure st = gemmi::read_structure_gz("../../tests/test_data/ligands/1ehz_psu.pdb.gz");
st.models.resize(1);
gemmi::remove_hydrogens(st.models[0]);
st.cell = gemmi::UnitCell(120, 120, 120, 90, 90, 90);
st.spacegroup_hm = "P 1";
st.setup_cell_images();
{
gemmi::Structure single = st;
const RefinementTopology topo = BuildRefinementTopology(single);
CHECK(FindNcsCopies(single.models[0], topo).Empty());
}
gemmi::Chain copy = st.models[0].chains.at(0);
copy.name = "B";
const double a = 0.7;
const gemmi::Mat33 rot(std::cos(a), -std::sin(a), 0, std::sin(a), std::cos(a), 0, 0, 0, 1);
for (gemmi::Residue &r : copy.residues)
for (gemmi::Atom &at : r.atoms)
at.pos = gemmi::Position(rot.multiply(at.pos)) + gemmi::Position(40, 10, -5);
st.models[0].chains.push_back(copy);
// and a fragment of the molecule, its first 18 residues: not a copy of the whole
gemmi::Chain fragment = st.models[0].chains.at(0);
fragment.name = "C";
fragment.residues.resize(18);
for (gemmi::Residue &r : fragment.residues)
for (gemmi::Atom &at : r.atoms)
at.pos += gemmi::Position(-30, 40, 20);
st.models[0].chains.push_back(fragment);
// and the whole molecule in another conformation, each residue moved 3 A its own way: the same sequence, not a copy
gemmi::Chain other = st.models[0].chains.at(0);
other.name = "D";
std::mt19937 shift_rng(11);
std::uniform_real_distribution<double> shift(-1.0, 1.0);
for (gemmi::Residue &r : other.residues) {
const gemmi::Vec3 v = gemmi::Vec3(shift(shift_rng), shift(shift_rng), shift(shift_rng)).normalized() * 3.0;
for (gemmi::Atom &at : r.atoms)
at.pos += gemmi::Position(v) + gemmi::Position(30, -40, 25);
}
st.models[0].chains.push_back(other);
st.connections.clear(); // the LINK records name chain A only
const RefinementTopology topo = BuildRefinementTopology(st);
const NcsCopies ncs = FindNcsCopies(st.models[0], topo);
REQUIRE_FALSE(ncs.Empty());
CHECK(ncs.groups == 1);
CHECK(ncs.copies == 2);
CHECK(ncs.description == "A=B");
CHECK(ncs.anchors.size() == 1);
CHECK(ncs.anchors[0].size() > 25);
CAPTURE(ncs.distances.records.size());
CHECK(ncs.distances.records.size() > 1000);
std::vector<gemmi::Atom *> atoms = atoms_of(st.models[0]);
std::vector<std::array<double, 4>> xyzb;
for (gemmi::Atom *at : atoms)
xyzb.push_back({at->pos.x, at->pos.y, at->pos.z, at->b_iso});
RestraintSet set = NcsRestraintsAt(ncs, 0.1);
set.Finalize(xyzb.size());
// the same molecule twice: nothing to restrain
CHECK(EvaluateRestraints(set, xyzb, nullptr) == Catch::Approx(0.0).margin(1e-6));
CHECK(NcsRmsd(ncs, xyzb) < 1e-3);
CHECK(NcsReleasedResidues(ncs, xyzb) == 0);
// the second copy shaken: the gradient against central differences
std::mt19937 rng(7);
std::normal_distribution<double> d(0, 0.1);
const size_t half = gemmi::count_atom_sites(st.models[0].chains[0]);
for (size_t i = half; i < 2 * half; i++)
for (int c = 0; c < 3; c++)
xyzb[i][c] += d(rng);
std::vector<std::array<double, 4>> grad(xyzb.size(), {0, 0, 0, 0});
const double e = EvaluateRestraints(set, xyzb, &grad);
CHECK(e > 0);
const double h = 1e-6;
for (size_t i = 0; i < xyzb.size(); i += 17)
for (int c = 0; c < 3; c++) {
auto x = xyzb;
x[i][c] += h;
const double fp = EvaluateRestraints(set, x, nullptr);
x[i][c] -= 2 * h;
const double fm = EvaluateRestraints(set, 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) {
const gemmi::SpaceGroup &p1 = *gemmi::find_spacegroup_by_name("P 1");
const gemmi::UnitCell wide_cell(40, 40, 40, 90, 90, 90);
std::vector<MergedReflection> merged = asu_reflections(wide_cell, p1, 3.0);
RefinementReflections refl = PrepareRefinementReflections(merged, wide_cell, p1, 3.0);
auto gpu = MakeRefinementEngineGPU(0, wide_cell, p1, 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);
CHECK(egpu == Catch::Approx(e).epsilon(1e-3));
for (size_t i = 0; i < xyzb.size(); i++)
for (int c = 0; c < 3; c++)
CHECK(ggpu[i][c] == Catch::Approx(grad[i][c]).epsilon(2e-3).margin(2e-2));
}
#endif
// one residue of the second copy (its 16th) distorted, each atom moved by 1.5 A in a direction of its own: its
// restraints let go, the others hold
size_t first = gemmi::count_atom_sites(st.models[0].chains[0]);
for (int k = 0; k < 15; k++)
first += st.models[0].chains[1].residues[k].atoms.size();
std::vector<std::array<double, 4>> moved_xyzb = xyzb;
std::uniform_real_distribution<double> u(-1.0, 1.0);
for (size_t i = first; i < first + st.models[0].chains[1].residues[15].atoms.size(); i++) {
const gemmi::Vec3 v = gemmi::Vec3(u(rng), u(rng), u(rng)).normalized() * 1.5;
moved_xyzb[i][0] += v.x;
moved_xyzb[i][1] += v.y;
moved_xyzb[i][2] += v.z;
}
CHECK(NcsReleasedResidues(ncs, xyzb) == 0);
CHECK(NcsReleasedResidues(ncs, moved_xyzb) >= 1);
CHECK(NcsReleasedResidues(ncs, moved_xyzb) <= 3);
}
namespace {
// Every other atom anisotropic, its dev random but traceless, up to a third of its B.
std::vector<AtomAniso> random_aniso(const gemmi::Model &model, unsigned seed) {
std::mt19937 rng(seed);
std::uniform_real_distribution<double> d(-1.0, 1.0);
std::vector<AtomAniso> out;
int i = 0;
for (const gemmi::Chain &ch : model.chains)
for (const gemmi::Residue &r : ch.residues)
for (const gemmi::Atom &a : r.atoms) {
AtomAniso an;
an.on = (i++ % 2) == 0;
if (an.on) {
const double u = a.b_iso / gemmi::u_to_b() / 3.0;
double v[6];
for (double &x : v)
x = u * d(rng);
const double t = (v[0] + v[1] + v[2]) / 3.0;
an.dev = {static_cast<float>(v[0] - t), static_cast<float>(v[1] - t), static_cast<float>(v[2] - t),
static_cast<float>(v[3] / 2), static_cast<float>(v[4] / 2), static_cast<float>(v[5] / 2)};
}
out.push_back(an);
}
return out;
}
// The engine's U gradient of a few anisotropic atoms against central differences of its own target.
void check_u_gradient(RefinementEngine &engine, const gemmi::Model &model, std::vector<AtomAniso> &aniso,
double tolerance) {
std::vector<std::array<double, 4>> grad;
engine.SetAniso(&aniso);
engine.Evaluate(model, &grad);
const std::vector<std::array<double, 6>> gu = engine.UGradient();
REQUIRE(gu.size() == aniso.size());
double max_g = 0;
for (size_t i = 0; i < gu.size(); i++)
for (int c = 0; c < 6; c++)
max_g = std::max(max_g, std::fabs(gu[i][c]));
const double h = 2e-3; // A^2
for (size_t i = 0; i < aniso.size(); i += 2 * (aniso.size() / 14)) {
REQUIRE(aniso[i].on);
for (int c = 0; c < 6; c++) {
const float v0 = aniso[i].dev[c];
aniso[i].dev[c] = static_cast<float>(v0 + h);
const double fp = engine.Evaluate(model, nullptr);
aniso[i].dev[c] = static_cast<float>(v0 - h);
const double fm = engine.Evaluate(model, nullptr);
aniso[i].dev[c] = v0;
const double fd = (fp - fm) / (2 * h);
CAPTURE(i, c, fd, gu[i][c], max_g);
CHECK(std::fabs(fd - gu[i][c]) < tolerance * std::max(std::fabs(fd), 0.05 * max_g));
}
}
}
} // namespace
TEST_CASE("RefinementXray_aniso_gradient_CPU", "[Refinement]") {
for (const char *group : {"P 1", "P 21 21 21"}) {
CAPTURE(group);
const gemmi::SpaceGroup &sg = *gemmi::find_spacegroup_by_name(group);
const gemmi::UnitCell cell(28, 31, 34, 90, 90, 90);
gemmi::Structure truth = random_crystal(cell, group, 240, 7);
gemmi::Model model = shaken(truth.models[0], 0.3, 4.0, 8);
RefinementReflections refl = PrepareRefinementReflections(asu_reflections(cell, sg, 2.0), cell, sg, 2.0);
auto cpu = MakeRefinementEngineCPU(cell, sg, 2.0, refl, 4);
cpu->SetBlur(RefinementBlur(cell, sg, 2.0, 5.0));
set_observations(*cpu, refl, truth.models[0], model);
std::vector<AtomAniso> aniso = random_aniso(model, 9);
check_u_gradient(*cpu, model, aniso, 0.06);
// the positions and the B along the isotropic direction, the shapes held
check_xray_gradient(*cpu, model, 0.03);
}
}
#ifdef JFJOCH_USE_CUDA
TEST_CASE("RefinementXray_aniso_gradient_GPU", "[Refinement]") {
if (get_gpu_count() == 0)
return;
const char *group = "P 21 21 21";
const gemmi::SpaceGroup &sg = *gemmi::find_spacegroup_by_name(group);
const gemmi::UnitCell cell(28, 31, 34, 90, 90, 90);
gemmi::Structure truth = random_crystal(cell, group, 240, 7);
gemmi::Model model = shaken(truth.models[0], 0.3, 4.0, 8);
RefinementReflections refl = PrepareRefinementReflections(asu_reflections(cell, sg, 2.0), cell, sg, 2.0);
auto cpu = MakeRefinementEngineCPU(cell, sg, 2.0, refl, 4);
cpu->SetBlur(RefinementBlur(cell, sg, 2.0, 5.0));
set_observations(*cpu, refl, truth.models[0], model);
auto gpu = MakeRefinementEngineGPU(0, cell, sg, 2.0, refl, model, 300.0, SIZE_MAX);
REQUIRE(gpu);
gpu->SetBlur(RefinementBlur(cell, sg, 2.0, 5.0));
std::vector<std::complex<float>> f;
std::vector<AtomAniso> aniso = random_aniso(model, 9);
cpu->SetAniso(&aniso);
gpu->SetAniso(&aniso);
cpu->Fcalc(model, f);
XrayRowModel rows;
rows.k_total.assign(refl.Size(), 1.0f);
rows.fmask.assign(refl.Size(), {0, 0});
EstimateLikelihoodParameters(refl, f, rows);
cpu->SetRowModel(rows);
gpu->SetRowModel(rows);
std::vector<std::complex<float>> fg;
gpu->Fcalc(model, fg);
double num = 0, den = 0;
for (size_t m = 0; m < f.size(); m++) {
num += std::abs(f[m] - fg[m]);
den += std::abs(f[m]);
}
CHECK(num / den < 1e-3);
std::vector<std::array<double, 4>> gc, gg;
cpu->Evaluate(model, &gc);
const std::vector<std::array<double, 6>> uc = cpu->UGradient();
gpu->Evaluate(model, &gg);
const std::vector<std::array<double, 6>> ug = gpu->UGradient();
REQUIRE(ug.size() == uc.size());
double max_g = 0, max_u = 0;
for (size_t i = 0; i < gc.size(); i++) {
for (int c = 0; c < 4; c++)
max_g = std::max(max_g, std::fabs(gc[i][c]));
for (int c = 0; c < 6; c++)
max_u = std::max(max_u, std::fabs(uc[i][c]));
}
for (size_t i = 0; i < gc.size(); i++) {
CAPTURE(i);
for (int c = 0; c < 4; c++)
CHECK(std::fabs(gg[i][c] - gc[i][c]) < 0.01 * max_g);
for (int c = 0; c < 6; c++)
CHECK(std::fabs(ug[i][c] - uc[i][c]) < 0.01 * max_u);
}
check_u_gradient(*gpu, model, aniso, 0.08);
}
// The GPU engine's evaluation with the gradient, timed on a protein-sized case: about 5500 atoms in the asymmetric
// unit of a P 21 21 21 cell at 1.6 A, the density of a protein's atoms with their hydrogens. Run by hand, under a
// profiler for the time of each kernel.
TEST_CASE("RefinementGPU_Benchmark", "[.][refine_bench]") {
if (get_gpu_count() == 0)
return;
const char *group = "P 21 21 21";
const gemmi::SpaceGroup &sg = *gemmi::find_spacegroup_by_name(group);
const gemmi::UnitCell cell(60, 64, 69, 90, 90, 90);
const double d_min = 1.6;
gemmi::Structure truth = random_crystal(cell, group, 5500, 7);
gemmi::Model model = shaken(truth.models[0], 0.3, 4.0, 8);
RefinementReflections refl = PrepareRefinementReflections(asu_reflections(cell, sg, d_min), cell, sg, d_min);
auto gpu = MakeRefinementEngineGPU(0, cell, sg, d_min, refl, model, 300.0, SIZE_MAX);
REQUIRE(gpu);
gpu->SetBlur(RefinementBlur(cell, sg, d_min, 5.0));
set_observations(*gpu, refl, truth.models[0], model);
std::vector<std::array<double, 4>> g;
gpu->Evaluate(model, &g);
const int n = 200;
const auto start = std::chrono::steady_clock::now();
double t = 0;
for (int i = 0; i < n; i++)
t += gpu->Evaluate(model, &g);
const double ms = std::chrono::duration<double, std::milli>(std::chrono::steady_clock::now() - start).count();
printf("RefinementGPU_Benchmark: %zu atoms, %zu reflections, %.3f ms per evaluation (target %.6g)\n",
atoms_of(model).size(), refl.Size(), ms / n, t / n);
}
#endif
TEST_CASE("RefinementADP_tls_gradient", "[Refinement]") {
std::mt19937 rng(4);
std::uniform_real_distribution<double> d(-1.0, 1.0);
gemmi::Model model(1);
gemmi::Chain chain("A");
gemmi::Residue res;
for (int i = 0; i < 30; i++) {
gemmi::Atom a;
a.pos = gemmi::Position(10 + 8 * d(rng), -3 + 8 * d(rng), 5 + 8 * d(rng));
res.atoms.push_back(a);
}
chain.residues.push_back(res);
model.chains.push_back(chain);
std::vector<gemmi::Atom *> atoms = atoms_of(model);
TlsBody body;
body.origin = gemmi::Position(9, -2, 4);
std::array<double, TlsBody::PARAMETERS> p{};
for (double &x : p)
x = 0.01 * d(rng);
body.Set(p);
for (int i = 0; i < 30; i++)
body.atoms.push_back(i);
// U is linear in the parameters: U(p) is sum p_k dU/dp_k
for (gemmi::Atom *a : atoms) {
const gemmi::SMat33<double> u = body.U(a->pos);
const auto dk = body.Derivatives(a->pos);
const double uu[6] = {u.u11, u.u22, u.u33, u.u12, u.u13, u.u23};
for (int c = 0; c < 6; c++) {
double sum = 0;
for (int k = 0; k < TlsBody::PARAMETERS; k++)
sum += p[k] * dk[k][c];
CHECK(sum == Catch::Approx(uu[c]).margin(1e-12));
}
}
// the chain rule in matrix form against the sum over dU/dp
std::vector<std::array<double, 6>> gu(atoms.size());
for (auto &g : gu)
for (double &x : g)
x = d(rng);
const std::array<double, TlsBody::PARAMETERS> g = body.Gradient(atoms, gu);
for (int k = 0; k < TlsBody::PARAMETERS; k++) {
double sum = 0;
for (size_t i = 0; i < atoms.size(); i++) {
const auto dk = body.Derivatives(atoms[i]->pos);
for (int c = 0; c < 6; c++)
sum += gu[i][c] * dk[k][c];
}
CAPTURE(k);
CHECK(g[k] == Catch::Approx(sum).margin(1e-9));
}
// and through the metadata in degrees and back
gemmi::Model m2 = model;
const std::vector<gemmi::TlsGroup> groups = TlsBodiesToGroups(m2, {body});
REQUIRE(groups.size() == 1);
const std::vector<TlsBody> back = TlsBodiesFromModel(m2, groups);
REQUIRE(back.size() == 1);
CHECK(back[0].atoms.size() == 30);
const gemmi::SMat33<double> u0 = body.U(atoms[3]->pos), u1 = back[0].U(atoms[3]->pos);
CHECK(u1.u12 == Catch::Approx(u0.u12).margin(1e-12));
CHECK(u1.u33 == Catch::Approx(u0.u33).margin(1e-12));
}
TEST_CASE("RefinementADP_u_restraints_gradient", "[Refinement]") {
std::mt19937 rng(5);
std::uniform_real_distribution<double> d(-1.0, 1.0);
const int n = 6;
std::vector<gemmi::Position> pos;
std::vector<std::array<double, 6>> u(n);
for (int i = 0; i < n; i++) {
pos.emplace_back(1.5 * i + 0.2 * d(rng), 0.3 * d(rng), 0.3 * d(rng));
u[i] = {0.2 + 0.05 * d(rng), 0.2 + 0.05 * d(rng), 0.2 + 0.05 * d(rng), 0.03 * d(rng), 0.03 * d(rng), 0.03 * d(rng)};
}
const URestraints r = MakeURestraints({{0, 1}, {1, 2}, {0, 2}, {3, 4}}, {1, 1, 1, 1, 1, 1});
CHECK(r.pairs.size() == 4);
REQUIRE(r.isotropic == std::vector<int>{5});
std::vector<std::array<double, 6>> g(n, {0, 0, 0, 0, 0, 0});
EvaluateURestraints(r, pos, u, &g);
const double h = 1e-6;
for (int i = 0; i < n; i++)
for (int c = 0; c < 6; c++) {
std::vector<std::array<double, 6>> up = u, um = u;
up[i][c] += h;
um[i][c] -= h;
const double fd = (EvaluateURestraints(r, pos, up, nullptr) - EvaluateURestraints(r, pos, um, nullptr)) / (2 * h);
CAPTURE(i, c);
CHECK(g[i][c] == Catch::Approx(fd).epsilon(1e-5).margin(1e-3));
}
// a rigid pair: the same U along the bond, however different across it, costs no rigid-bond term
const URestraints one = MakeURestraints({{0, 1}}, {1, 1});
const std::vector<gemmi::Position> line = {{0, 0, 0}, {0, 0, 1.5}};
const std::vector<std::array<double, 6>> across = {{0.1, 0.3, 0.2, 0, 0, 0}, {0.3, 0.1, 0.2, 0.05, 0, 0}};
const double sim = (0.2 * 0.2 + 0.2 * 0.2 + 2 * 0.05 * 0.05) / (URestraints::SIGMA_SIMILAR * URestraints::SIGMA_SIMILAR);
CHECK(EvaluateURestraints(one, line, across, nullptr) == Catch::Approx(sim));
}
TEST_CASE("RefinementADP_tls_segments", "[Refinement]") {
// a chain of 60 residues along x whose B follows one quadratic on the first 25 residues and another after
std::vector<std::vector<int>> residues;
std::vector<gemmi::Position> pos;
std::vector<double> b;
for (int r = 0; r < 60; r++) {
std::vector<int> ra;
for (int a = 0; a < 5; a++) {
const gemmi::Position p(3.8 * r, 2.0 * std::sin(a + r), 2.0 * std::cos(a * r));
ra.push_back(static_cast<int>(pos.size()));
pos.push_back(p);
const double y = p.y, z = p.z;
b.push_back(r < 25 ? 20 + 0.5 * (y * y + z * z) : 60 + 3.0 * (y * y) - 0.02 * p.x);
}
residues.push_back(ra);
}
CHECK(TlsSegments(residues, pos, b, 1, 10) == std::vector<int>{0});
CHECK(TlsSegments(residues, pos, b, 2, 10) == std::vector<int>{0, 25});
// too short for three segments of 25
CHECK(TlsSegments(residues, pos, b, 3, 25).size() == 2);
}
namespace {
// A twinned test case in P 4, whose lattice adds the twin law: the reflections of the group's ASU to 2.2 A with
// every fourth of the zone l = 1 left out, so that some twin mates are rows of their own, and the free set drawn
// on the lattice's holohedry (AssignRfreeFlags) so that no twin pair is split.
struct TwinCase {
gemmi::UnitCell cell{30, 30, 40, 90, 90, 90};
const gemmi::SpaceGroup &sg = *gemmi::find_spacegroup_by_name("P 4");
gemmi::Op law;
std::vector<MergedReflection> merged;
RefinementReflections refl;
TwinCase() {
const std::vector<gemmi::Op> laws = RefinementTwinLaws(cell, sg);
REQUIRE(laws.size() == 1);
law = laws[0];
merged = asu_reflections(cell, sg, 2.2);
int i = 0;
std::erase_if(merged, [&](const MergedReflection &r) { return r.l == 1 && i++ % 4 == 0; });
AssignRfreeFlags(merged, sg, 0.05, 0,
UnitCell{static_cast<float>(cell.a), static_cast<float>(cell.b), static_cast<float>(cell.c),
90, 90, 90});
refl = PrepareRefinementReflections(merged, cell, sg, 2.2);
CHECK(SetTwinLaw(refl, cell, sg, 2.2, law) == 0);
}
};
} // namespace
TEST_CASE("RefinementTwin_laws_and_mates", "[Refinement]") {
TwinCase t;
CHECK(RefinementTwinLaws(t.cell, *gemmi::find_spacegroup_by_name("P 4 2 2")).empty());
REQUIRE(t.refl.Twinned());
CHECK(!t.refl.mate_hkl.empty());
const gemmi::GroupOps gops = t.sg.operations();
const gemmi::ReciprocalAsu asu(&t.sg);
const std::vector<gemmi::Miller> rows = t.refl.RowHkl();
REQUIRE(t.refl.Laws() == 1);
for (size_t m = 0; m < t.refl.Size(); m++) {
const int j = t.refl.mates[0][m];
CAPTURE(m, j);
REQUIRE(j >= 0);
if (j != static_cast<int>(m)) {
CHECK(rows[j] == asu.to_asu(t.law.apply_to_hkl(rows[m]), gops).first);
CHECK(t.refl.mates_from[0][j] == static_cast<int>(m));
if (j < static_cast<int>(t.refl.Size()))
CHECK(t.refl.mates[0][j] == static_cast<int>(m)); // an involution
}
}
// a free set drawn on the merging group alone splits twin pairs
RefinementReflections split = t.refl;
for (size_t m = 0; m < split.Size(); m++)
split.free[m] = m % 20 == 0;
CHECK(SetTwinLaw(split, t.cell, t.sg, 2.2, t.law) > 0);
}
// The free set closed under the lattice's domains: one drawn on the holohedry is left as it is; one drawn on
// another cell's (P 2 2 2 on a tetragonal metric, its flags drawn on an orthorhombic one: twin pairs split) ends with
// every pair in one set and about the free fraction it had - where a majority of each pair would empty it, and a
// draw that repeats the flag of one member would inflate it.
TEST_CASE("RefinementTwin_free_set_closed", "[Refinement]") {
const gemmi::UnitCell cell{30, 40, 30, 90, 90, 90};
const gemmi::SpaceGroup &sg = *gemmi::find_spacegroup_by_name("P 2 2 2");
std::vector<MergedReflection> merged = asu_reflections(cell, sg, 1.5);
AssignRfreeFlags(merged, sg, 0.05, 0, UnitCell{30, 40, 30, 90, 90, 90});
CHECK(CloseFreeSet(merged, cell, sg) == 0);
AssignRfreeFlags(merged, sg, 0.05, 0, UnitCell{30, 40, 36, 90, 90, 90});
const auto free_count = [&] {
return std::count_if(merged.begin(), merged.end(), [](const MergedReflection &r) { return r.rfree_flag; });
};
const auto before = free_count();
CHECK(CloseFreeSet(merged, cell, sg) > 0);
CHECK(free_count() == Catch::Approx(before).epsilon(0.1));
std::map<gemmi::Miller, bool> flag;
const gemmi::GroupOps gops = sg.operations();
const gemmi::ReciprocalAsu asu(&sg);
for (const MergedReflection &r : merged)
flag[asu.to_asu(gemmi::Miller{{r.h, r.k, r.l}}, gops).first] = r.rfree_flag;
int split = 0;
for (const MergedReflection &r : merged)
for (const gemmi::Op &law : RefinementTwinDomains(cell, sg)) {
const auto it = flag.find(asu.to_asu(law.apply_to_hkl({{r.h, r.k, r.l}}), gops).first);
split += it != flag.end() && it->second != r.rfree_flag;
}
CHECK(split == 0);
CHECK(CloseFreeSet(merged, cell, sg) == 0);
}
TEST_CASE("RefinementTwin_rule_with_data_evidence", "[Refinement]") {
// the null alone: 3 standard errors
CHECK(AdoptTwinLaw(0.1, 3.0, false));
CHECK_FALSE(AdoptTwinLaw(0.1, 2.9, false));
// the data read twinned as well: 2 standard errors
CHECK(AdoptTwinLaw(0.1, 2.0, true));
CHECK(AdoptTwinLaw(0.1, 2.5, true));
CHECK_FALSE(AdoptTwinLaw(0.1, 1.9, true));
// no fraction, no twin, whatever the evidence
CHECK_FALSE(AdoptTwinLaw(0.0, 10.0, true));
}
TEST_CASE("RefinementTwin_gradient", "[Refinement]") {
TwinCase t;
gemmi::Structure truth = random_crystal(t.cell, "P 4", 240, 17);
const gemmi::Model model = shaken(truth.models[0], 0.3, 4.0, 18);
auto cpu = MakeRefinementEngineCPU(t.cell, t.sg, 2.2, t.refl, 4);
cpu->SetBlur(RefinementBlur(t.cell, t.sg, 2.2, 5.0));
XrayRowModel rows;
rows.k_total.assign(t.refl.Size(), 1.0f);
rows.fmask.assign(t.refl.Rows(), {0, 0});
rows.twin_fractions = {0.3f};
std::vector<std::complex<float>> f;
cpu->Fcalc(truth.models[0], f);
REQUIRE(f.size() == t.refl.Rows());
const std::vector<std::complex<float>> twinned = TwinnedModel(t.refl, rows, f);
for (size_t m = 0; m < t.refl.Size(); m++)
t.refl.fo[m] = std::abs(twinned[m]);
// the fraction comes back from the data
CHECK(FitTwinFraction(t.refl, rows, f) == Catch::Approx(0.3).margin(0.01));
cpu->Fcalc(model, f);
EstimateLikelihoodParameters(t.refl, TwinnedModel(t.refl, rows, f), rows);
cpu->SetRowModel(rows);
// the CPU engine kept a reference to t.refl and reads the observations set above
check_xray_gradient(*cpu, model, 0.03);
#ifdef JFJOCH_USE_CUDA
if (get_gpu_count() > 0) {
auto gpu = MakeRefinementEngineGPU(0, t.cell, t.sg, 2.2, t.refl, model, 300.0, SIZE_MAX);
REQUIRE(gpu);
gpu->SetBlur(RefinementBlur(t.cell, t.sg, 2.2, 5.0));
gpu->SetRowModel(rows);
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);
}
}
#endif
}
TEST_CASE("RefinementTwin_domains_gradient", "[Refinement]") {
// P 2 on a tetragonal metric (a = c, beta = 90): three domain operators beside the identity, a two-fold whose
// square is in the group and two four-fold cosets
const gemmi::UnitCell cell{30, 40, 30, 90, 90, 90};
const gemmi::SpaceGroup &sg = *gemmi::find_spacegroup_by_name("P 1 2 1");
const std::vector<gemmi::Op> laws = RefinementTwinDomains(cell, sg);
REQUIRE(laws.size() == 3);
std::vector<MergedReflection> merged = asu_reflections(cell, sg, 2.4);
int i = 0;
std::erase_if(merged, [&](const MergedReflection &r) { return r.k == 1 && i++ % 4 == 0; });
AssignRfreeFlags(merged, sg, 0.05, 0, UnitCell{30, 40, 30, 90, 90, 90});
RefinementReflections refl = PrepareRefinementReflections(merged, cell, sg, 2.4);
CHECK(SetTwinLaws(refl, cell, sg, 2.4, laws) == 0);
REQUIRE(refl.Laws() == 3);
CHECK(!refl.mate_hkl.empty());
gemmi::Structure truth = random_crystal(cell, "P 1 2 1", 240, 27);
const gemmi::Model model = shaken(truth.models[0], 0.3, 4.0, 28);
auto cpu = MakeRefinementEngineCPU(cell, sg, 2.4, refl, 4);
cpu->SetBlur(RefinementBlur(cell, sg, 2.4, 5.0));
XrayRowModel rows;
rows.k_total.assign(refl.Size(), 1.0f);
rows.fmask.assign(refl.Rows(), {0, 0});
rows.twin_fractions = {0.2f, 0.1f, 0.15f};
std::vector<std::complex<float>> f;
cpu->Fcalc(truth.models[0], f);
const std::vector<std::complex<float>> twinned = TwinnedModel(refl, rows, f);
for (size_t m = 0; m < refl.Size(); m++)
refl.fo[m] = std::abs(twinned[m]);
const std::vector<float> fitted = FitTwinFractions(refl, rows, f);
REQUIRE(fitted.size() == 3);
for (size_t d = 0; d < 3; d++)
CHECK(fitted[d] == Catch::Approx(rows.twin_fractions[d]).margin(0.01));
cpu->Fcalc(model, f);
EstimateLikelihoodParameters(refl, TwinnedModel(refl, rows, f), rows);
cpu->SetRowModel(rows);
check_xray_gradient(*cpu, model, 0.03);
#ifdef JFJOCH_USE_CUDA
if (get_gpu_count() > 0) {
auto gpu = MakeRefinementEngineGPU(0, cell, sg, 2.4, refl, model, 300.0, SIZE_MAX);
REQUIRE(gpu);
gpu->SetBlur(RefinementBlur(cell, sg, 2.4, 5.0));
gpu->SetRowModel(rows);
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 a = 0; a < gc.size(); a++)
for (int c = 0; c < 4; c++) {
CAPTURE(a, c);
CHECK(std::fabs(gg[a][c] - gc[a][c]) < 0.01 * max_g);
}
}
#endif
}
TEST_CASE("Refinement_twin_decided", "[Refinement]") {
TwinCase t;
gemmi::Structure truth = residue_crystal({"PHE", "TRP", "DA", "DG", "TYR", "HOH", "HOH", "HOH"}, t.cell, "P 4", 21);
gemmi::Structure start = truth;
start.models[0] = shaken(truth.models[0], 0.2, 4.0, 22);
std::vector<std::complex<float>> f;
{
RefinementReflections refl = t.refl;
auto cpu = MakeRefinementEngineCPU(t.cell, t.sg, 2.2, refl, 4);
cpu->SetBlur(RefinementBlur(t.cell, t.sg, 2.2, 5.0));
cpu->Fcalc(truth.models[0], f);
}
Logger logger("RefinementTest");
RefinementSettings settings;
settings.gpu = false;
settings.nthreads = 4;
settings.weight_ladder = {0.0625};
settings.adp_ladder = {0.0};
settings.max_steps = 1;
settings.final_ladder = {1.0};
settings.riding_hydrogens = false;
settings.solvent = false;
for (double fraction : {0.0, 0.3}) {
CAPTURE(fraction);
// the observations as twinned at `fraction`, with 3% noise
XrayRowModel rows;
rows.k_total.assign(t.refl.Size(), 1.0f);
rows.fmask.assign(t.refl.Rows(), {0, 0});
rows.twin_fractions = {static_cast<float>(fraction)};
const std::vector<std::complex<float>> twinned = TwinnedModel(t.refl, rows, f);
std::map<gemmi::Miller, double> amplitude;
for (size_t m = 0; m < t.refl.Size(); m++)
amplitude[t.refl.hkl[m]] = std::abs(twinned[m]);
std::vector<MergedReflection> merged = t.merged;
std::mt19937 rng(5);
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;
}
const RefinementResult r = RefineModel(start, merged, t.cell, t.sg, 1.0, logger, settings);
REQUIRE(r.ok);
CHECK(r.twin_law == t.law.as_hkl().triplet('h'));
if (fraction > 0) {
CHECK(r.twin_decision == "TWINNED");
CHECK(r.twinned);
CHECK(r.twin_fraction == Catch::Approx(fraction).margin(0.03));
CHECK(r.r_free < r.untwinned_r_free - 0.02);
CHECK(r.twin_r_free < r.twin_r_free_null - 0.02);
CHECK(r.twin_null_z >= 3.0);
} else {
CHECK((r.twin_decision == "ZERO_FRACTION" || r.twin_null_z < 3.0));
CHECK(r.twin_decision != "TWINNED");
CHECK_FALSE(r.twinned);
// the untwinned refinement, as it is without the twin hypothesis
RefinementSettings plain = settings;
plain.twin = false;
const RefinementResult u = RefineModel(start, merged, t.cell, t.sg, 1.0, logger, plain);
CHECK(u.r_free == r.r_free);
CHECK(u.r_work == r.r_work);
}
// scored only, the model as given is tested against the twin law by the same rule
RefinementSettings score = settings;
score.score_only = true;
const RefinementResult s = RefineModel(truth, merged, t.cell, t.sg, 1.0, logger, score);
REQUIRE(s.ok);
CHECK((s.twin_decision == "TWINNED") == (fraction > 0));
if (fraction > 0)
CHECK(s.r_free_start < s.untwinned_r_free - 0.02);
}
}
Reference in new issue View Git Blame Copy Permalink