mx/Jungfraujoch
3
2 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 87 Wiki Activity

Refinement: the weight trials side by side, each on a model, topology and engine of its own

Browse Source 
Every rung of the weight ladder starts from the same model and depends on no other, so the six position
trials, and then the three B trials at the chosen weight, run at the same time: on the main model and engine
and on up to five slots beside it, each a copy of the model before the topology was made (the topology is
built again on it, hydrogens included) with its own engine. The slots go on the GPUs in turn from the main
engine's, as many as each card's budget takes, so a machine with several cards spreads them; a slot whose
engine cannot be made is left out, and with none the trials run one after another on the main engine as
before. The threads are shared out among the slots. The log lines and the choice are made afterwards in
ladder order, so the log, the trial table and the refined model are the same, bit for bit (GPU and CPU
engines checked).

5reo 3.9 -> 2.5 s, 9yl4 13.4 -> 7.8 s, 6o2h 3.1 -> 1.9 s, 8dyz 4.6 -> 2.7 s, 9rcs 2.9 -> 2.2 s (one RTX 5080, -N 8).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_015CeEGgpKdfkCTJb3YNMJcu
(cherry picked from commit 0ea6a6db8aba640afc5249fe748922ddfaa2537a)
(cherry picked from commit 0c0176c8bf571ed26f8a68a9de40222487c6200e)
This commit is contained in:
leonarski_f committed 2026-10-26 01:38:23 +01:00
1 parent 5569f3d1f2
commit 8933a72a89
1 file changed
+245 -67
Download Patch File Download Diff File Expand all files Collapse all files
image_analysis/structure_refinement/Refinement.cpp
+245 -67
View File
@@ -9,6 +9,7 @@
#include <optional>
#include <cmath>
#include <random>
#include <thread>
#include <unordered_map>
#include <ceres/gradient_problem.h>
@@ -101,6 +102,8 @@ struct Context {
std::unique_ptr<ModelScaleGPUPool> scale_gpu;
#endif
size_t nthreads = 1;
int device = -1; // the engine's GPU
size_t engine_bytes = 0; // and what it holds there
XrayRowModel rows;
std::vector<std::complex<float>> fcalc, fmask;
double k_sol = 0, b_sol = 0;
@@ -611,6 +614,8 @@ void MakeEngine(Context &ctx, bool gpu, const RefinementSettings &settings, Logg
size_t bytes = 0;
ctx.engine = MakeRefinementEngineGPU(device, ctx.cell, *ctx.sg, ctx.d_min, ctx.refl, ctx.Model(),
ctx.b_ceiling, budget, &bytes, settings.nthreads);
ctx.device = device;
ctx.engine_bytes = bytes;
if (ctx.engine)
ctx.scale_gpu = std::make_unique<ModelScaleGPUPool>(device, ctx.refl.Size(), 1);
else
@@ -630,6 +635,187 @@ void MakeEngine(Context &ctx, bool gpu, const RefinementSettings &settings, Logg
ctx.engine = MakeRefinementEngineCPU(ctx.cell, *ctx.sg, ctx.d_min, ctx.refl, settings.nthreads);
}
// The riding hydrogens placed on their parents as they are, and the offsets and B they ride with.
void PlaceRiders(Context &ctx, RefinementTopology &topo) {
const RefinementStage stage("riders");
PlaceRidingHydrogens(topo);
ctx.rider_offset.assign(ctx.atoms.size(), gemmi::Position());
for (size_t i = 0; i < ctx.atoms.size(); i++)
if (ctx.rider_parent[i] >= 0) {
ctx.rider_offset[i] = ctx.atoms[i]->pos - ctx.atoms[ctx.rider_parent[i]]->pos;
ctx.atoms[i]->b_iso = ctx.atoms[ctx.rider_parent[i]]->b_iso;
}
}
// The geometry and the repulsions of the model as it is.
RestraintSet RestraintsNow(Context &ctx, const RefinementTopology &topo) {
const RefinementStage stage("repulsions");
RestraintSet set = topo.geometry;
set.Append(BuildRepulsions(topo, ctx.Model(), ctx.cell, *ctx.sg, REPULSION_MARGIN, ctx.nthreads));
set.Finalize(ctx.atoms.size());
return set;
}
// Where a weight trial runs: a model with its topology and its engine.
struct TrialSlot {
Context *ctx;
RefinementTopology *topo;
};
// A slot of its own beside the main one.
struct OwnSlot {
Context ctx;
RefinementTopology topo;
};
// A copy of the model as it was before the topology was made (the same topology comes out of it, hydrogens
// included), with an engine of the main one's kind - a GPU's on `device`. Null where the card cannot make one.
std::unique_ptr<OwnSlot> MakeOwnSlot(const gemmi::Structure &before_topology, const Context &main, int device,
size_t nthreads, const RefinementSettings &settings) {
auto slot = std::make_unique<OwnSlot>();
Context &ctx = slot->ctx;
ctx.st = before_topology;
ctx.cell = main.cell;
ctx.sg = main.sg;
ctx.d_min = main.d_min;
ctx.refl = main.refl;
ctx.nthreads = nthreads;
ctx.b_ceiling = main.b_ceiling;
slot->topo = BuildRefinementTopology(ctx.st, settings.riding_hydrogens);
ctx.atoms = AtomList(ctx.Model());
ctx.rider_parent = slot->topo.rider_parent;
#ifdef JFJOCH_USE_CUDA
if (main.engine->OnGPU()) {
try {
ctx.engine = MakeRefinementEngineGPU(device, ctx.cell, *ctx.sg, ctx.d_min, ctx.refl, ctx.Model(),
ctx.b_ceiling, SIZE_MAX, nullptr, nthreads);
if (!ctx.engine)
return nullptr;
ctx.scale_gpu = std::make_unique<ModelScaleGPUPool>(device, ctx.refl.Size(), 1);
} catch (const JFJochException &) {
return nullptr;
}
ctx.device = device;
}
#else
(void) device;
#endif
if (!ctx.engine)
ctx.engine = MakeRefinementEngineCPU(ctx.cell, *ctx.sg, ctx.d_min, ctx.refl, nthreads);
ctx.engine->SetRestraints(1, slot->topo.adp);
return slot;
}
// The devices of the slots beside the main one, `wanted` of them at most: the GPUs in turn from the main
// engine's, each taking engines while they fit in its budget. -1 for each one on the CPU.
std::vector<int> OwnSlotDevices(const Context &main, int wanted, const RefinementSettings &settings) {
#ifdef JFJOCH_USE_CUDA
if (main.engine->OnGPU()) {
const int n = get_gpu_count();
std::vector<size_t> used(n, 0);
used[main.device] = main.engine_bytes;
std::vector<int> devices;
for (int k = 1; k <= wanted; k++) {
const int d = (main.device + k) % n;
const size_t usable = get_gpu_total_memory(d) / 5 * 4;
const size_t budget = usable > settings.gpu_resident_bytes ? usable - settings.gpu_resident_bytes : 0;
if (used[d] + main.engine_bytes <= budget) {
used[d] += main.engine_bytes;
devices.push_back(d);
}
}
return devices;
}
#else
(void) settings;
#endif
return std::vector<int>(std::max(wanted, 0), -1);
}
// One weight trial: a full refinement from pos0 and b0, macro-cycles of the positions at X-ray weight w and,
// where adp_w is not 0, the B factors at X-ray weight adp_w.
struct TrialOutcome {
RefinementWeightTrial trial;
std::vector<gemmi::Position> pos;
std::vector<double> b;
std::string line; // what the log says of it
};
TrialOutcome RunTrial(TrialSlot slot, const std::vector<int> &refinable, const std::vector<gemmi::Position> &pos0,
const std::vector<double> &b0, double w, double adp_w, const RefinementSettings &settings) {
Context &ctx = *slot.ctx;
RefinementTopology &topo = *slot.topo;
const RefinementStage stage(adp_w > 0 ? "rung.b" : "rung.xyz");
ctx.SetPositions(pos0);
ctx.SetB(b0);
for (int cycle = 0; cycle < settings.macro_cycles; cycle++) {
PlaceRiders(ctx, topo);
ctx.Refit();
{
const RestraintSet set = RestraintsNow(ctx, topo);
const RefinementStage s1("set_restraints");
ctx.engine->SetRestraints(0, set);
}
Minimise(ctx, AtomTarget::Mode::Xyz, refinable, w, 0, settings.xyz_iterations);
if (adp_w > 0)
Minimise(ctx, AtomTarget::Mode::B, refinable, adp_w, 1, settings.b_iterations);
}
PlaceRiders(ctx, topo);
TrialOutcome out;
RefinementWeightTrial &t = out.trial;
t.weight = w;
t.adp_weight = adp_w;
const RFactors rf = ctx.Refit();
t.r_work = rf.r_work;
t.r_free = rf.r_free;
{
const RefinementStage s1("measure");
t.geometry = MeasureGeometry(topo.geometry,
BuildRepulsions(topo, ctx.Model(), ctx.cell, *ctx.sg, 0.0, ctx.nthreads), ctx.XyzB());
}
t.within_bounds = t.geometry.rmsd_bonds <= MAX_RMSD_BONDS && t.geometry.rmsd_angles <= MAX_RMSD_ANGLES;
out.line = fmt::format("Refinement, X-ray weight {:.3g}, B {}: R-work {:.4f}, R-free {:.4f}, rmsd bonds {:.4f} A, "
"angles {:.2f} deg{}",
w, adp_w > 0 ? fmt::format("refined at weight {:.3g}", adp_w) : std::string("held"),
rf.r_work, rf.r_free, t.geometry.rmsd_bonds, t.geometry.rmsd_angles,
t.within_bounds ? "" : " (geometry out of bounds)");
out.pos = ctx.Positions();
out.b = ctx.B();
return out;
}
// The trials (w, adp_w) side by side: trial i in slot i % slots, each slot's trials in turn on a thread of its
// own (on its GPU), returned in the order given. Every trial starts from pos0 and b0 and takes nothing from the
// one before it in its slot, so the outcome is the same for any number of slots.
std::vector<TrialOutcome> RunTrials(const std::vector<TrialSlot> &trial_slots,
const std::vector<std::pair<double, double>> &trials,
const std::vector<int> &refinable, const std::vector<gemmi::Position> &pos0,
const std::vector<double> &b0, const RefinementSettings &settings) {
std::vector<TrialOutcome> out(trials.size());
std::vector<std::exception_ptr> errors(trial_slots.size());
auto run_slot = [&](size_t s) {
try {
#ifdef JFJOCH_USE_CUDA
const ScopedGpuPin pin(trial_slots[s].ctx->device >= 0 ? trial_slots[s].ctx->device : get_gpu());
#endif
for (size_t i = s; i < trials.size(); i += trial_slots.size())
out[i] = RunTrial(trial_slots[s], refinable, pos0, b0, trials[i].first, trials[i].second, settings);
} catch (...) {
errors[s] = std::current_exception();
}
};
std::vector<std::thread> threads;
for (size_t s = 1; s < std::min(trial_slots.size(), trials.size()); s++)
threads.emplace_back(run_slot, s);
run_slot(0);
for (std::thread &t : threads)
t.join();
for (const std::exception_ptr &e : errors)
if (e)
std::rethrow_exception(e);
return out;
}
} // namespace
RefinementResult RefineModel(const gemmi::Structure &placed, const std::vector<MergedReflection> &merged,
@@ -689,23 +875,14 @@ RefinementResult RefineModel(const gemmi::Structure &placed, const std::vector<M
result.rigid_bodies, result.group_b ? "kept" : "not kept", r.r_work, r.r_free);
if (settings.restrained) {
const gemmi::Structure before_topology = ctx.st;
std::optional<RefinementStage> topo_stage("topology");
RefinementTopology topo = BuildRefinementTopology(ctx.st, settings.riding_hydrogens);
topo_stage.reset();
// the hydrogens it added: the atoms and the engine made again for them
ctx.atoms = AtomList(ctx.Model());
ctx.rider_parent = topo.rider_parent;
auto place_riders = [&]() {
const RefinementStage stage("riders");
PlaceRidingHydrogens(topo);
ctx.rider_offset.assign(ctx.atoms.size(), gemmi::Position());
for (size_t i = 0; i < ctx.atoms.size(); i++)
if (ctx.rider_parent[i] >= 0) {
ctx.rider_offset[i] = ctx.atoms[i]->pos - ctx.atoms[ctx.rider_parent[i]]->pos;
ctx.atoms[i]->b_iso = ctx.atoms[ctx.rider_parent[i]]->b_iso;
}
};
place_riders();
PlaceRiders(ctx, topo);
if (static_cast<int>(ctx.atoms.size()) != result.atoms) {
MakeEngine(ctx, result.on_gpu, settings, logger);
result.on_gpu = ctx.engine->OnGPU();
@@ -725,16 +902,9 @@ RefinementResult RefineModel(const gemmi::Structure &placed, const std::vector<M
topo.fixed_names.empty() ? std::string() : " (" + Joined(topo.fixed_names) + ")",
topo.bonds, topo.angles, topo.torsions, topo.chirals, topo.planes);
auto restraints_now = [&]() {
const RefinementStage stage("repulsions");
RestraintSet set = topo.geometry;
set.Append(BuildRepulsions(topo, ctx.Model(), cell, sg, REPULSION_MARGIN, ctx.nthreads));
set.Finalize(ctx.atoms.size());
return set;
};
ctx.engine->SetRestraints(1, topo.adp);
{
const RestraintSet set = restraints_now();
const RestraintSet set = RestraintsNow(ctx, topo);
ctx.engine->SetRestraints(0, set);
RestraintSet repulsions = BuildRepulsions(topo, ctx.Model(), cell, sg, 0.0, ctx.nthreads);
result.geometry_start = MeasureGeometry(topo.geometry, repulsions, ctx.XyzB());
@@ -750,67 +920,75 @@ RefinementResult RefineModel(const gemmi::Structure &placed, const std::vector<M
const std::vector<gemmi::Position> pos0 = ctx.Positions();
const std::vector<double> b0 = ctx.B();
// The trials run side by side, as many at a time as the longer ladder has rungs and the cards take engines:
// on the main model and engine, and on copies with engines of their own; each with a share of the threads.
std::optional<RefinementStage> slot_stage("slots");
const int wanted = static_cast<int>(std::max(settings.weight_ladder.size(), settings.adp_ladder.size()));
const std::vector<int> devices = OwnSlotDevices(ctx, wanted - 1, settings);
const size_t slot_threads = std::max<size_t>(1, settings.nthreads / (devices.size() + 1));
std::vector<std::unique_ptr<OwnSlot>> own(devices.size());
{
std::vector<std::thread> threads;
for (size_t k = 0; k < devices.size(); k++)
threads.emplace_back([&, k]() {
#ifdef JFJOCH_USE_CUDA
const ScopedGpuPin pin(devices[k] >= 0 ? devices[k] : get_gpu());
#endif
own[k] = MakeOwnSlot(before_topology, ctx, devices[k], slot_threads, settings);
});
for (std::thread &t : threads)
t.join();
}
const size_t main_threads = ctx.nthreads;
ctx.nthreads = slot_threads;
std::vector<TrialSlot> trial_slots{{&ctx, &topo}};
for (const std::unique_ptr<OwnSlot> &o : own)
if (o)
trial_slots.push_back({&o->ctx, &o->topo});
slot_stage.reset();
logger.Debug("Refinement: {} weight trials at a time", trial_slots.size());
int best = -1;
std::vector<gemmi::Position> best_pos;
std::vector<double> best_b;
// One full refinement from the start: macro-cycles of the positions at X-ray weight w and, where adp_w is
// not 0, the B factors at X-ray weight adp_w; kept where it beats the best so far. The rungs come in order
// of trust - tighter geometry, then fewer parameters - and a later one displaces an earlier only by more
// than the tie in R-free.
auto trial = [&](double w, double adp_w) {
const RefinementStage stage(adp_w > 0 ? "rung.b" : "rung.xyz");
ctx.SetPositions(pos0);
ctx.SetB(b0);
for (int cycle = 0; cycle < settings.macro_cycles; cycle++) {
place_riders();
ctx.Refit();
{
const RestraintSet set = restraints_now();
const RefinementStage s1("set_restraints");
ctx.engine->SetRestraints(0, set);
// A trial is kept where it beats the best so far. The rungs come in order of trust - tighter geometry,
// then fewer parameters - and a later one displaces an earlier only by more than the tie in R-free.
auto choose = [&](std::vector<TrialOutcome> &outcomes) {
for (TrialOutcome &o : outcomes) {
logger.Info("{}", o.line);
result.trials.push_back(o.trial);
const RefinementWeightTrial &t = result.trials.back();
const int k = static_cast<int>(result.trials.size()) - 1;
const bool better = best < 0 ? true
: t.within_bounds && (!result.trials[best].within_bounds ||
t.r_free < result.trials[best].r_free - R_FREE_TIE);
if (better) {
best = k;
best_pos = std::move(o.pos);
best_b = std::move(o.b);
}
Minimise(ctx, AtomTarget::Mode::Xyz, refinable, w, 0, settings.xyz_iterations);
if (adp_w > 0)
Minimise(ctx, AtomTarget::Mode::B, refinable, adp_w, 1, settings.b_iterations);
}
place_riders();
RefinementWeightTrial t;
t.weight = w;
t.adp_weight = adp_w;
const RFactors rf = ctx.Refit();
t.r_work = rf.r_work;
t.r_free = rf.r_free;
{
const RefinementStage s1("measure");
t.geometry = MeasureGeometry(topo.geometry, BuildRepulsions(topo, ctx.Model(), cell, sg, 0.0, ctx.nthreads), ctx.XyzB());
}
t.within_bounds = t.geometry.rmsd_bonds <= MAX_RMSD_BONDS && t.geometry.rmsd_angles <= MAX_RMSD_ANGLES;
logger.Info("Refinement, X-ray weight {:.3g}, B {}: R-work {:.4f}, R-free {:.4f}, rmsd bonds {:.4f} A, "
"angles {:.2f} deg{}",
w, adp_w > 0 ? fmt::format("refined at weight {:.3g}", adp_w) : std::string("held"), rf.r_work,
rf.r_free, t.geometry.rmsd_bonds, t.geometry.rmsd_angles,
t.within_bounds ? "" : " (geometry out of bounds)");
result.trials.push_back(t);
const int k = static_cast<int>(result.trials.size()) - 1;
const bool better = best < 0 ? true
: t.within_bounds && (!result.trials[best].within_bounds ||
t.r_free < result.trials[best].r_free - R_FREE_TIE);
if (better) {
best = k;
best_pos = ctx.Positions();
best_b = ctx.B();
}
};
// the positions' weight, with the B factors as they are
std::vector<std::pair<double, double>> xyz_trials;
for (double ladder : settings.weight_ladder)
trial(w0 * ladder, 0.0);
xyz_trials.emplace_back(w0 * ladder, 0.0);
std::vector<TrialOutcome> outcomes = RunTrials(trial_slots, xyz_trials, refinable, pos0, b0, settings);
choose(outcomes);
// then the B factors refined too, at that weight, against their own weight
const double w = result.trials[best].weight;
std::vector<std::pair<double, double>> b_trials;
for (double adp : settings.adp_ladder)
trial(w, w * adp);
b_trials.emplace_back(w, w * adp);
outcomes = RunTrials(trial_slots, b_trials, refinable, pos0, b0, settings);
choose(outcomes);
for (size_t s = 1; s < trial_slots.size(); s++)
ctx.evaluations += trial_slots[s].ctx->evaluations;
ctx.nthreads = main_threads;
ctx.SetPositions(best_pos);
ctx.SetB(best_b);
place_riders();
PlaceRiders(ctx, topo);
const RefinementWeightTrial &chosen = result.trials[best];
result.restrained = true;
result.weight = chosen.weight;