The likelihood's variance per bin (beta) is fitted by maximum likelihood on the free reflections as the whole spread of Fo about D Fc, and D is then fitted on the working reflections with the variance eps beta + 2 sigma^2 (acentric; + sigma^2 centric), which is also the target's variance. Murshudov et al. (2011) Acta D67 355: D from the working set, the variance from the free set. The measurement error is thereby counted twice, so D and the figure of merit fall where sigma is a large part of the spread - at high resolution. Measured on REFMAC 5.8.0431 0-cycle output (black box) of the 14 parity sets: its FOM inverted per reflection gives var = eps a + 2 sigma^2 (acentric, fitted coefficient 2.00) and eps a + sigma^2 (centric); its D (FC_ALL/FC_ALL_LS = the log's SigmaA_Fc1) and a are reproduced by this estimate (on REFMAC's own Fc, one set to rms 0.009 in D over all 20 bins). The previous per-bin free-set ML put D at 0.39-0.48 where REFMAC has 0.06-0.13 in the outermost bins. rugnux score vs REFMAC 0-cycle, 14 sets: mean |FOM difference| 0.056 -> 0.012 (max 0.112 -> 0.042), per-bin D rms 0.107 -> 0.068; R-work/R-free unchanged (D does not enter them). The score's FOM is now the mean over the working reflections, as REFMAC reports it, and the refined model's maps are weighted by the refinement's own m and D rather than a separate sigma_A estimate. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_015CeEGgpKdfkCTJb3YNMJcu (cherry picked from commit d00d635e2874bff6cd216b7f24367362d0cc2e56) (cherry picked from commit b93e0ea2fc1ffcc6bb40dae9b77816a321f4d6ec)
623 lines
28 KiB
C++
623 lines
28 KiB
C++
// 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 <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);
|
|
}
|
|
|
|
// 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));
|
|
}
|
|
|
|
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
|
|
}
|
|
|
|
// 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("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);
|
|
|
|
// 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);
|
|
}
|
|
}
|