mx/Jungfraujoch
3
2 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 87 Wiki Activity
Files
9da38518442c9af8f568c92e1d50b1827c8399d3
Jungfraujoch/tests/RefinementTest.cpp
T
leonarski_f 9da3851844 Refinement: at most n weight trials at a time, and a cancel flag the caller can set 
- RefinementSettings::max_trials_at_once (rugnux_refine_test --trials-at-once): 0 as many as the ladders and
  the cards take, 1 one after another - the check that side by side changes nothing. Verified: refined
  model and report identical, side by side vs one after another, GPU (5reo, 9yl4, 6o2h, 9rcs) and CPU
  engine (6o2h, 9rcs, 5reo), and run to run.
- RefinementSettings::cancel, an atomic flag for the broker's live deadline: checked at every iteration of
  the L-BFGS (positions, B, rigid bodies, group B: the minimiser stops where it stands), at every
  macro-cycle and before every weight trial. The trials that ran are chosen among as usual; cut before the
  first, the model is the placement's. RefinementResult::cancelled says so.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_015CeEGgpKdfkCTJb3YNMJcu
(cherry picked from commit f708d3d227bb42215ffe3d9a15c766dbd73b783c)
2026-10-26 01:38:24 +01:00

466 lines
20 KiB
C++
Raw 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 <cmath>
#include <map>
#include <random>
#include <vector>
#include <gemmi/symmetry.hpp>
#include "../common/CUDAWrapper.h"
#include "../common/Logger.h"
#include "../image_analysis/structure_refinement/MonomerLibrary.h"
#include "../image_analysis/structure_refinement/Refinement.h"
#include "../image_analysis/structure_refinement/RefinementMath.h"
#include "../image_analysis/structure_refinement/RefinementRestraints.h"
#include "../image_analysis/structure_refinement/RefinementXray.h"
#ifdef JFJOCH_USE_CUDA
#include "../image_analysis/structure_refinement/RefinementGPU.h"
#endif
namespace {
// Copies of residue `code` at the embedded dictionary's ideal coordinates (hydrogens left out), each turned and
// moved to a place of its own in the cell.
gemmi::Structure residue_crystal(const std::vector<std::string> &codes, const gemmi::UnitCell &cell, const char *group,
unsigned seed) {
std::mt19937 rng(seed);
std::uniform_real_distribution<double> frac(0.0, 1.0), angle(-3.1, 3.1);
gemmi::Structure st;
st.cell = cell;
st.spacegroup_hm = group;
st.setup_cell_images();
gemmi::Model model(1);
gemmi::Chain chain("A");
int seq = 1;
for (const std::string &code : codes) {
const gemmi::ChemComp &cc = EmbeddedMonomerLibrary().monomers.at(code);
gemmi::Residue res;
res.name = code;
res.seqid = gemmi::SeqId(seq++, ' ');
res.het_flag = 'H';
res.entity_type = gemmi::EntityType::NonPolymer;
const gemmi::Position centre = cell.orthogonalize(gemmi::Fractional(frac(rng), frac(rng), frac(rng)));
const double ax = angle(rng), ay = angle(rng), az = angle(rng);
const gemmi::Mat33 rot = gemmi::Mat33(1, 0, 0, 0, std::cos(ax), -std::sin(ax), 0, std::sin(ax), std::cos(ax))
.multiply(gemmi::Mat33(std::cos(ay), 0, std::sin(ay), 0, 1, 0, -std::sin(ay), 0, std::cos(ay)))
.multiply(gemmi::Mat33(std::cos(az), -std::sin(az), 0, std::sin(az), std::cos(az), 0, 0, 0, 1));
for (const gemmi::ChemComp::Atom &ca : cc.atoms) {
if (ca.is_hydrogen())
continue;
gemmi::Atom a;
a.name = ca.id;
a.element = ca.el;
a.pos = gemmi::Position(rot.multiply(ca.xyz)) + centre;
a.b_iso = 20.0f;
a.occ = 1.0f;
res.atoms.push_back(a);
}
chain.residues.push_back(res);
}
model.chains.push_back(chain);
st.models.push_back(model);
return st;
}
// Atoms of a few elements scattered over the cell.
gemmi::Structure random_crystal(const gemmi::UnitCell &cell, const char *group, int n, unsigned seed) {
std::mt19937 rng(seed);
std::uniform_real_distribution<double> frac(0.0, 1.0), b(10.0, 40.0);
const gemmi::El elements[] = {gemmi::El::C, gemmi::El::N, gemmi::El::O, gemmi::El::S};
gemmi::Structure st;
st.cell = cell;
st.spacegroup_hm = group;
st.setup_cell_images();
gemmi::Model model(1);
gemmi::Chain chain("A");
gemmi::Residue res;
res.name = "UNK";
res.seqid = gemmi::SeqId(1, ' ');
for (int i = 0; i < n; i++) {
gemmi::Atom a;
a.name = "X";
a.element = elements[i % 4];
a.pos = cell.orthogonalize(gemmi::Fractional(frac(rng), frac(rng), frac(rng)));
a.b_iso = static_cast<float>(b(rng));
a.occ = 1.0f;
res.atoms.push_back(a);
}
chain.residues.push_back(res);
model.chains.push_back(chain);
st.models.push_back(model);
return st;
}
// Every reflection of the group's ASU to d_min, amplitude 1, one in twenty free.
std::vector<MergedReflection> asu_reflections(const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg, double d_min) {
std::vector<MergedReflection> out;
const gemmi::ReciprocalAsu asu(&sg);
const gemmi::GroupOps gops = sg.operations();
const gemmi::Miller lim = cell.get_hkl_limits(d_min);
int i = 0;
for (int h = -lim[0]; h <= lim[0]; h++)
for (int k = -lim[1]; k <= lim[1]; k++)
for (int l = -lim[2]; l <= lim[2]; l++) {
const gemmi::Miller m{{h, k, l}};
if ((h == 0 && k == 0 && l == 0) || !asu.is_in(m) || gops.is_systematically_absent(m))
continue;
const double d = cell.calculate_d(m);
if (d < d_min)
continue;
MergedReflection r;
r.h = h;
r.k = k;
r.l = l;
r.d = static_cast<float>(d);
r.F = 1.0f;
r.sigmaF = 0.05f;
r.rfree_flag = (i++ % 20) == 0;
out.push_back(r);
}
return out;
}
// The observations set to |F| of `truth`, and a row model with no scale, no solvent and the likelihood's
// parameters estimated against `model`.
void set_observations(RefinementEngine &engine, RefinementReflections &refl, const gemmi::Model &truth,
const gemmi::Model &model) {
std::vector<std::complex<float>> f;
engine.Fcalc(truth, f);
for (size_t m = 0; m < refl.Size(); m++)
refl.fo[m] = std::abs(f[m]);
engine.Fcalc(model, f);
XrayRowModel rows;
rows.k_total.assign(refl.Size(), 1.0f);
rows.fmask.assign(refl.Size(), {0, 0});
EstimateLikelihoodParameters(refl, f, rows);
engine.SetRowModel(rows);
}
gemmi::Model shaken(const gemmi::Model &model, double sigma, double b_sigma, unsigned seed) {
std::mt19937 rng(seed);
std::normal_distribution<double> d(0, sigma), db(0, b_sigma);
gemmi::Model out = model;
for (gemmi::Chain &ch : out.chains)
for (gemmi::Residue &r : ch.residues)
for (gemmi::Atom &a : r.atoms) {
a.pos += gemmi::Position(d(rng), d(rng), d(rng));
a.b_iso = static_cast<float>(std::max(5.0, a.b_iso + db(rng)));
}
return out;
}
std::vector<gemmi::Atom *> atoms_of(gemmi::Model &model) {
std::vector<gemmi::Atom *> out;
for (gemmi::Chain &ch : model.chains)
for (gemmi::Residue &r : ch.residues)
for (gemmi::Atom &a : r.atoms)
out.push_back(&a);
return out;
}
// The engine's gradient of a few atoms against central differences of its own target.
void check_xray_gradient(RefinementEngine &engine, gemmi::Model model, double tolerance) {
std::vector<std::array<double, 4>> grad;
engine.Evaluate(model, &grad);
std::vector<gemmi::Atom *> atoms = atoms_of(model);
const double hx = 2e-3, hb = 2e-2;
double max_g = 0;
for (const auto &g : grad)
for (int c = 0; c < 4; c++)
max_g = std::max(max_g, std::fabs(g[c]));
for (size_t i = 0; i < atoms.size(); i += atoms.size() / 7) {
for (int c = 0; c < 4; c++) {
const gemmi::Position p0 = atoms[i]->pos;
const float b0 = atoms[i]->b_iso;
const double h = c < 3 ? hx : hb;
auto set = [&](double s) {
atoms[i]->pos = p0;
atoms[i]->b_iso = b0;
if (c < 3)
atoms[i]->pos.at(c) += s;
else
atoms[i]->b_iso = static_cast<float>(b0 + s);
};
set(h);
const double fp = engine.Evaluate(model, nullptr);
set(-h);
const double fm = engine.Evaluate(model, nullptr);
set(0);
const double fd = (fp - fm) / (2 * h);
CAPTURE(i, c, fd, grad[i][c], max_g);
CHECK(std::fabs(fd - grad[i][c]) < tolerance * std::max(std::fabs(fd), 0.05 * max_g));
}
}
}
} // namespace
TEST_CASE("RefinementLikelihood_derivative", "[Refinement]") {
for (bool centric : {false, true})
for (double fc : {0.3, 2.0, 9.0}) {
double dfc = 0, dummy = 0;
const double f0 = refine_math::AmplitudeNegLogLikelihood<double>(4.0, fc, 0.8, 2.5, centric, dfc);
const double h = 1e-5;
const double fp = refine_math::AmplitudeNegLogLikelihood<double>(4.0, fc + h, 0.8, 2.5, centric, dummy);
const double fm = refine_math::AmplitudeNegLogLikelihood<double>(4.0, fc - h, 0.8, 2.5, centric, dummy);
CAPTURE(centric, fc, f0);
CHECK(dfc == Catch::Approx((fp - fm) / (2 * h)).epsilon(1e-5));
}
// the Bessel approximation on both sides of its switch
double l0, r0, l1, r1;
refine_math::BesselI0I1(3.75 * (1 - 1e-12), l0, r0);
refine_math::BesselI0I1(3.75 * (1 + 1e-12), l1, r1);
CHECK(l0 == Catch::Approx(l1).epsilon(1e-5));
CHECK(r0 == Catch::Approx(r1).epsilon(1e-5));
}
TEST_CASE("RefinementMonomerLibrary_embedded", "[Refinement]") {
const gemmi::MonLib &monlib = EmbeddedMonomerLibrary();
for (const char *code : {"ALA", "PRO", "MSE", "DA", "U", "ATP", "GOL", "SO4", "HOH", "ZN"})
CHECK(InEmbeddedMonomerLibrary(code));
CHECK_FALSE(InEmbeddedMonomerLibrary("XYZ"));
for (const char *link : {"TRANS", "CIS", "PTRANS", "p", "SS"})
CHECK(monlib.links.count(link) == 1);
CHECK(monlib.modifications.count("NH3") == 1);
CHECK(monlib.ener_lib.atoms.count("CH1") == 1);
}
TEST_CASE("RefinementRestraints_ideal_residues", "[Refinement]") {
const gemmi::UnitCell cell(60, 65, 70, 90, 90, 90);
gemmi::Structure st = residue_crystal({"PHE", "DA", "ATP", "HOH", "ZN"}, cell, "P 21 21 21", 3);
const RefinementTopology topo = BuildRefinementTopology(st);
CHECK(topo.restrained_residues == 5);
CHECK(topo.fixed_residues == 0);
CHECK(topo.bonds > 60);
CHECK(topo.planes >= 3);
CHECK(topo.chirals >= 4);
std::vector<std::array<double, 4>> xyzb;
for (gemmi::Atom *a : atoms_of(st.models[0]))
xyzb.push_back({a->pos.x, a->pos.y, a->pos.z, a->b_iso});
const GeometryStatistics g = MeasureGeometry(topo.geometry, {}, xyzb);
CHECK(g.rmsd_bonds < 0.02);
CHECK(g.rmsd_angles < 2.5);
CHECK(g.rmsd_chirals < 0.2);
}
TEST_CASE("RefinementRestraints_gradient", "[Refinement]") {
const gemmi::UnitCell cell(14, 15, 16, 90, 90, 90);
gemmi::Structure st = residue_crystal({"TYR", "DG"}, cell, "P 21 21 21", 5);
st.models[0] = shaken(st.models[0], 0.15, 3.0, 2);
const RefinementTopology topo = BuildRefinementTopology(st);
// the geometry, the B similarity and the repulsions (symmetry mates included, at a contact distance that
// makes every pair within reach count)
RestraintSet set = topo.geometry;
set.Append(topo.adp);
RefinementTopology wide = topo;
for (float &r : wide.vdw_radius)
r = 2.5f;
std::fill(wide.hbond.begin(), wide.hbond.end(), 0);
const RestraintSet rep = BuildRepulsions(wide, st.models[0], cell, *st.find_spacegroup(), 0.0);
REQUIRE(rep.records.size() > 10);
bool image = false;
for (const RestraintRecord &r : rep.records)
image = image || r.rot[0] != 1.0f || r.tran[0] != 0.0f || r.tran[1] != 0.0f || r.tran[2] != 0.0f;
CAPTURE(rep.records.size());
CHECK(image);
set.Append(rep);
set.Finalize(atoms_of(st.models[0]).size());
std::vector<std::array<double, 4>> xyzb;
for (gemmi::Atom *a : atoms_of(st.models[0]))
xyzb.push_back({a->pos.x, a->pos.y, a->pos.z, a->b_iso});
// each kind of restraint on its own, against central differences
for (int kind = 0; kind <= static_cast<int>(RestraintKind::BSimilarity); kind++) {
RestraintSet one;
for (const RestraintRecord &r : set.records)
if (r.kind == kind) {
RestraintRecord copy = r;
if (kind == static_cast<int>(RestraintKind::Plane)) {
copy.atom[0] = static_cast<int>(one.plane_atoms.size());
for (int k = 0; k < r.atom[1]; k++)
one.plane_atoms.push_back(set.plane_atoms[r.atom[0] + k]);
}
one.Add(copy, kind == static_cast<int>(RestraintKind::Plane) ? r.atom[1] : 4);
}
CAPTURE(kind, one.records.size());
std::vector<std::array<double, 4>> grad(xyzb.size(), {0, 0, 0, 0});
EvaluateRestraints(one, xyzb, &grad);
const double h = 1e-6;
for (size_t i = 0; i < xyzb.size(); i++)
for (int c = 0; c < 4; c++) {
auto x = xyzb;
x[i][c] += h;
const double fp = EvaluateRestraints(one, x, nullptr);
x[i][c] -= 2 * h;
const double fm = EvaluateRestraints(one, x, nullptr);
CAPTURE(i, c);
CHECK(grad[i][c] == Catch::Approx((fp - fm) / (2 * h)).epsilon(1e-4).margin(1e-3));
}
}
#ifdef JFJOCH_USE_CUDA
if (get_gpu_count() > 0) {
// the GPU's float evaluation of the same set
// (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
}
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);
}
}
Reference in new issue View Git Blame Copy Permalink