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

Refinement: one trajectory of short weight trials instead of a ladder of full refinements from the start

Browse Source 
Every rung of the old ladders (6 positional, then 3 B) was a full 10-cycle refinement from the placed model, so
~90% of the evaluations were thrown away and the kept trajectory had 100 iterations. The run is now:

1. a ladder from the placed model, 2 macro-cycles per trial: three positional weights (w0/1024, /128, /16), each
   with the B factors held and refined (adp = 4 w), chosen by the old rule (tighter first, a later trial wins only
   within the geometry bounds and by more than the R-free tie);
2. a trajectory of steps from where the last one ended, each two trials side by side - the weights as they are and
   the X-ray term raised, in turn the positions' (x2, at most w0/8) and the B factors' (x4) - until a step lowers
   neither R-work nor R-free by more than the tie (at most 8 steps);
3. a final ladder from there at x1/2, x1, x2.

A macro-cycle is 30 positional and 20 B iterations (the minimiser's memory kept within a trial). A trial starts
from the scale, solvent and likelihood fit its predecessor ended with (the same computation, so result-neutral),
and the slots that run share the threads.

On 5 battery MD starts (start 1 of 5f6m 1.09 A, 8dqb 2.04, 7arr 0.92, 11if 1.37, 7d98 3.32; R-free by rugnux's
scorer, REFMAC 0-cycle within 0.002), previous commit -> this, against REFMAC run to convergence (40-100 cycles):
5f6m 0.215 -> 0.167 (REFMAC 0.168), 8dqb 0.341 -> 0.290 (0.284), 7arr 0.356 -> 0.289 (0.288), 11if 0.295 ->
0.236 (0.244), 7d98 0.393 -> 0.337 (0.360). Time 5.4 / 8.6 / 7.3 / 3.9 / 10.0 s on the RTX 5080 (base 6-7 s).
Halving a weight during the trajectory never won in the logs, so steps only try raising it; the final ladder
can lower it.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_015CeEGgpKdfkCTJb3YNMJcu
(cherry picked from commit 2bfb84bc4f3c72e8da37e54b728839b65a636519)
This commit is contained in:
leonarski_f committed 2026-10-26 01:38:24 +01:00
1 parent 4424d8884d
commit 6369a263b1
4 files changed
+141 -65
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

image_analysis/structure_refinement/Refinement.cpp
+110 -50
View File
@@ -180,6 +180,26 @@ struct Context {
atoms[i]->b_iso = static_cast<float>(b[i]);
}
// What Refit() fitted, to start another context from without fitting it again.
struct Fit {
double b_floor = B_FLOOR;
double k_sol = 0, b_sol = 0, k_overall = 0;
gemmi::SMat33<double> b_star{0, 0, 0, 0, 0, 0};
XrayRowModel rows;
};
Fit Fitted() const { return {b_floor, k_sol, b_sol, k_overall, b_star, rows}; }
// Refit() as it was where `fit` was taken, for the atoms as they were there.
void Restore(const Fit &fit) {
b_floor = fit.b_floor;
engine->SetBlur(RefinementBlur(cell, *sg, d_min, b_floor));
k_sol = fit.k_sol;
b_sol = fit.b_sol;
k_overall = fit.k_overall;
b_star = fit.b_star;
rows = fit.rows;
engine->SetRowModel(rows);
}
// A new macro-cycle: the blur for the B the cycle may reach, the bulk-solvent mask of the model as it is,
// the scale and the solvent fitted to the working set, and the likelihood's parameters from the free set.
// Returns the R-factors of the model as it stands.
@@ -950,18 +970,20 @@ std::vector<int> OwnSlotDevices(const Context &main, int wanted, const Refinemen
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.
// One weight trial: `cycles` macro-cycles from pos0 and b0, of the positions at X-ray weight w and, where adp_w is
// not 0, the B factors at X-ray weight adp_w. `fit`: what was fitted at pos0 and b0 already, null to fit it here.
struct TrialOutcome {
RefinementWeightTrial trial;
std::vector<gemmi::Position> pos;
std::vector<double> b;
Context::Fit fit; // the scale, the solvent and the likelihood's parameters fitted where it ended
std::string line; // what the log says of it
bool ran = false; // false: left out, the refinement cancelled before 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) {
const std::vector<double> &b0, const Context::Fit *fit, double w, double adp_w, int cycles,
const RefinementSettings &settings) {
Context &ctx = *slot.ctx;
RefinementTopology &topo = *slot.topo;
const RefinementStage stage(adp_w > 0 ? "rung.b" : "rung.xyz");
@@ -971,9 +993,12 @@ TrialOutcome RunTrial(TrialSlot slot, const std::vector<int> &refinable, const s
RefinementMinimiser xyz_minimiser, b_minimiser;
BCurvature curvature;
double xray_scale = 0;
for (int cycle = 0; cycle < settings.macro_cycles && !ctx.Cancelled(); cycle++) {
for (int cycle = 0; cycle < cycles && !ctx.Cancelled(); cycle++) {
PlaceRiders(ctx, topo);
ctx.Refit();
if (cycle == 0 && fit != nullptr)
ctx.Restore(*fit);
else
ctx.Refit();
{
const RestraintSet set = RestraintsNow(ctx, topo);
const RefinementStage s1("set_restraints");
@@ -1009,7 +1034,7 @@ TrialOutcome RunTrial(TrialSlot slot, const std::vector<int> &refinable, const s
&topo.metal_bonds);
}
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, "
out.line = fmt::format("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,
@@ -1018,6 +1043,7 @@ TrialOutcome RunTrial(TrialSlot slot, const std::vector<int> &refinable, const s
t.within_bounds ? "" : " (geometry out of bounds)");
out.pos = ctx.Positions();
out.b = ctx.B();
out.fit = ctx.Fitted();
out.ran = true;
return out;
}
@@ -1028,9 +1054,14 @@ TrialOutcome RunTrial(TrialSlot slot, const std::vector<int> &refinable, const s
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) {
const std::vector<double> &b0, const Context::Fit *fit, int cycles,
const RefinementSettings &settings) {
std::vector<TrialOutcome> out(trials.size());
std::vector<std::exception_ptr> errors(trial_slots.size());
// the threads shared among the slots that run (the outcome does not depend on how many each has)
const size_t running = std::min(trial_slots.size(), trials.size());
for (size_t s = 0; s < running; s++)
trial_slots[s].ctx->nthreads = std::max<size_t>(1, settings.nthreads / running);
auto run_slot = [&](size_t s) {
try {
#ifdef JFJOCH_USE_CUDA
@@ -1038,7 +1069,8 @@ std::vector<TrialOutcome> RunTrials(const std::vector<TrialSlot> &trial_slots,
#endif
for (size_t i = s; i < trials.size(); i += trial_slots.size())
if (!trial_slots[s].ctx->Cancelled())
out[i] = RunTrial(trial_slots[s], refinable, pos0, b0, trials[i].first, trials[i].second, settings);
out[i] = RunTrial(trial_slots[s], refinable, pos0, b0, fit, trials[i].first, trials[i].second,
cycles, settings);
} catch (...) {
errors[s] = std::current_exception();
}
@@ -1208,10 +1240,11 @@ 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:
// The trials run side by side, as many at a time as the longest set of them has 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");
int wanted = static_cast<int>(std::max(settings.weight_ladder.size(), settings.adp_ladder.size()));
int wanted = static_cast<int>(std::max({settings.weight_ladder.size() * settings.adp_ladder.size(),
settings.final_ladder.size(), size_t{3}}));
if (settings.max_trials_at_once > 0)
wanted = std::min(wanted, static_cast<int>(settings.max_trials_at_once));
const std::vector<int> devices = OwnSlotDevices(ctx, wanted - 1, settings);
@@ -1238,58 +1271,85 @@ RefinementResult RefineModel(const gemmi::Structure &placed, const std::vector<M
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;
// 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) {
// Of a set of trials, the one kept: they come in order of trust and a later one displaces the best so far
// only where its geometry is within bounds and its R-free lower by more than the tie. -1: none ran.
auto choose = [&](std::vector<TrialOutcome> &outcomes, const char *stage) {
int best = -1;
for (size_t k = 0; k < outcomes.size(); k++) {
const TrialOutcome &o = outcomes[k];
if (!o.ran)
continue;
logger.Info("{}", o.line);
logger.Info("Refinement, {}: {}", stage, 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);
}
const RefinementWeightTrial &t = o.trial, *b = best < 0 ? nullptr : &outcomes[best].trial;
if (b == nullptr || (t.within_bounds && (!b->within_bounds || t.r_free < b->r_free - R_FREE_TIE)))
best = static_cast<int>(k);
}
return best;
};
// the positions' weight, with the B factors as they are
std::vector<std::pair<double, double>> xyz_trials;
for (double ladder : settings.weight_ladder)
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
if (best >= 0 && !ctx.Cancelled()) {
const double w = result.trials[best].weight;
std::vector<std::pair<double, double>> b_trials;
const double w_max = w0 * settings.max_weight;
// 1. The weights from the placed model: the ladder's rungs, tighter geometry first, each with the B factors
// held and refined.
std::vector<std::pair<double, double>> ladder;
for (double rung : settings.weight_ladder)
for (double adp : settings.adp_ladder)
b_trials.emplace_back(w, w * adp);
outcomes = RunTrials(trial_slots, b_trials, refinable, pos0, b0, settings);
choose(outcomes);
ladder.emplace_back(w0 * rung, w0 * rung * adp);
std::vector<TrialOutcome> outcomes = RunTrials(trial_slots, ladder, refinable, pos0, b0, nullptr,
settings.macro_cycles, settings);
int k = choose(outcomes, "ladder");
TrialOutcome state;
if (k >= 0)
state = std::move(outcomes[k]);
// 2. The trajectory: steps of a few macro-cycles from where the last one ended, each trying the weights as they
// are and the X-ray term's weight raised - in turn the positions' (doubled) and the B factors' (times
// four) - until a step lowers neither R-work nor R-free by more than the tie.
for (int step = 0; step < settings.max_steps && state.ran && !ctx.Cancelled(); step++) {
const double w = state.trial.weight, adp = state.trial.adp_weight;
std::vector<std::pair<double, double>> trials = {{w, adp}};
if (step % 2 == 1 && adp > 0)
trials.emplace_back(w, 4 * adp);
else if (2 * w <= w_max)
trials.emplace_back(2 * w, 2 * adp);
outcomes = RunTrials(trial_slots, trials, refinable, state.pos, state.b, &state.fit, settings.macro_cycles,
settings);
k = choose(outcomes, fmt::format("step {}", step + 1).c_str());
if (k < 0)
break;
const bool converged = state.trial.r_free - outcomes[k].trial.r_free < R_FREE_TIE &&
state.trial.r_work - outcomes[k].trial.r_work < R_FREE_TIE;
state = std::move(outcomes[k]);
if (converged)
break;
}
// 3. The final weights: a short ladder from where the trajectory ended, tighter geometry first.
if (state.ran && !ctx.Cancelled()) {
const double w = state.trial.weight, adp = state.trial.adp_weight;
std::vector<std::pair<double, double>> trials;
for (double f : settings.final_ladder)
if (w * f <= w_max)
trials.emplace_back(w * f, adp * f);
outcomes = RunTrials(trial_slots, trials, refinable, state.pos, state.b, &state.fit, settings.macro_cycles,
settings);
k = choose(outcomes, "final");
if (k >= 0)
state = std::move(outcomes[k]);
}
for (size_t s = 1; s < trial_slots.size(); s++)
ctx.evaluations += trial_slots[s].ctx->evaluations;
ctx.nthreads = main_threads;
if (best >= 0) {
ctx.SetPositions(best_pos);
ctx.SetB(best_b);
if (state.ran) {
ctx.SetPositions(state.pos);
ctx.SetB(state.b);
PlaceRiders(ctx, topo);
const RefinementWeightTrial &chosen = result.trials[best];
result.restrained = true;
result.weight = chosen.weight;
result.adp_weight = chosen.adp_weight;
result.r_work = chosen.r_work;
result.r_free = chosen.r_free;
result.geometry = chosen.geometry;
result.weight = state.trial.weight;
result.adp_weight = state.trial.adp_weight;
result.r_work = state.trial.r_work;
result.r_free = state.trial.r_free;
result.geometry = state.trial.geometry;
} else
ctx.Refit(); // no trial ran: the model as placed, its scale and solvent for the maps
}