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

Refinement: at most n weight trials at a time, and a cancel flag the caller can set

Browse Source 
- 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)
This commit is contained in:
leonarski_f committed 2026-10-26 01:38:24 +01:00
1 parent 8718065a7c
commit 9da3851844
4 files changed
+85 -22
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

image_analysis/structure_refinement/Refinement.cpp
+62 -22
View File
@@ -117,6 +117,9 @@ struct Context {
std::vector<int> rider_parent;
std::vector<gemmi::Position> rider_offset;
int evaluations = 0;
const std::atomic<bool> *cancel = nullptr; // RefinementSettings::cancel
bool Cancelled() const { return cancel != nullptr && cancel->load(std::memory_order_relaxed); }
gemmi::Model &Model() { return st.models[0]; }
@@ -315,6 +318,23 @@ ceres::GradientProblemSolver::Options SolverOptions(int iterations) {
return options;
}
// SolverOptions(iterations) that stop at the next iteration, where the minimiser stands, once ctx is cancelled.
class CancelCheck : public ceres::IterationCallback {
public:
explicit CancelCheck(const Context &ctx) : ctx_(ctx) {}
ceres::CallbackReturnType operator()(const ceres::IterationSummary &) override {
return ctx_.Cancelled() ? ceres::SOLVER_TERMINATE_SUCCESSFULLY : ceres::SOLVER_CONTINUE;
}
ceres::GradientProblemSolver::Options Options(int iterations) const {
ceres::GradientProblemSolver::Options options = SolverOptions(iterations);
options.callbacks.push_back(const_cast<CancelCheck *>(this));
return options;
}
private:
const Context &ctx_;
};
// The positions or the B factors of the refinable atoms, against w times the X-ray target plus a restraint set.
// Following Liu & Nocedal (1989) Math. Program. 45, 503-528 for the minimiser (Ceres' L-BFGS)
class AtomTarget : public ceres::FirstOrderFunction {
@@ -391,7 +411,8 @@ void Minimise(Context &ctx, AtomTarget::Mode mode, const std::vector<int> &param
}
ceres::GradientProblem problem(std::make_unique<AtomTarget>(ctx, mode, param_atoms, weight, restraint_set));
ceres::GradientProblemSolver::Summary summary;
ceres::Solve(SolverOptions(iterations), problem, p.data(), &summary);
const CancelCheck cancel_check(ctx);
ceres::Solve(cancel_check.Options(iterations), problem, p.data(), &summary);
// the atoms where the solver ended, not where its last line search probed
for (size_t k = 0; k < param_atoms.size(); k++) {
gemmi::Atom &a = *ctx.atoms[param_atoms[k]];
@@ -542,7 +563,8 @@ void RefinePlacement(Context &ctx, bool group_b, RefinementResult &result) {
RigidBodiesTarget *rb = target.get();
ceres::GradientProblem problem(std::move(target));
ceres::GradientProblemSolver::Summary summary;
ceres::Solve(SolverOptions(30), problem, q.data(), &summary);
const CancelCheck cancel_check(ctx);
ceres::Solve(cancel_check.Options(30), problem, q.data(), &summary);
rb->Place(q.data());
const RFactors after = ctx.Refit();
if (after.r_free < before.r_free - R_FREE_TIE) {
@@ -560,7 +582,8 @@ void RefinePlacement(Context &ctx, bool group_b, RefinementResult &result) {
GroupBTarget *gb = target.get();
ceres::GradientProblem problem(std::move(target));
ceres::GradientProblemSolver::Summary summary;
ceres::Solve(SolverOptions(30), problem, q.data(), &summary);
const CancelCheck cancel_check(ctx);
ceres::Solve(cancel_check.Options(30), problem, q.data(), &summary);
gb->Place(q.data());
result.group_b = ctx.Refit().r_free < before.r_free - R_FREE_TIE;
if (!result.group_b)
@@ -757,6 +780,7 @@ std::unique_ptr<OwnSlot> MakeOwnSlot(const gemmi::Structure &before_topology, co
ctx.d_min = main.d_min;
ctx.refl = main.refl;
ctx.nthreads = nthreads;
ctx.cancel = main.cancel;
ctx.b_ceiling = main.b_ceiling;
ctx.mask = std::make_unique<BulkSolventMask>(ctx.cell, *ctx.sg, ctx.d_min, nthreads);
slot->topo = BuildRefinementTopology(ctx.st, settings.riding_hydrogens);
@@ -817,6 +841,7 @@ struct TrialOutcome {
std::vector<gemmi::Position> pos;
std::vector<double> b;
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,
@@ -826,7 +851,7 @@ TrialOutcome RunTrial(TrialSlot slot, const std::vector<int> &refinable, const s
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++) {
for (int cycle = 0; cycle < settings.macro_cycles && !ctx.Cancelled(); cycle++) {
PlaceRiders(ctx, topo);
ctx.Refit();
{
@@ -859,6 +884,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.ran = true;
return out;
}
@@ -877,7 +903,8 @@ std::vector<TrialOutcome> RunTrials(const std::vector<TrialSlot> &trial_slots,
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);
if (!trial_slots[s].ctx->Cancelled())
out[i] = RunTrial(trial_slots[s], refinable, pos0, b0, trials[i].first, trials[i].second, settings);
} catch (...) {
errors[s] = std::current_exception();
}
@@ -906,6 +933,7 @@ RefinementResult RefineModel(const gemmi::Structure &placed, const std::vector<M
ctx.cell = cell;
ctx.sg = &sg;
ctx.nthreads = settings.nthreads;
ctx.cancel = settings.cancel;
ctx.atoms = AtomList(ctx.Model());
result.atoms = static_cast<int>(ctx.atoms.size());
double b_max = 0;
@@ -1047,7 +1075,9 @@ RefinementResult RefineModel(const gemmi::Structure &placed, const std::vector<M
// 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()));
int wanted = static_cast<int>(std::max(settings.weight_ladder.size(), settings.adp_ladder.size()));
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);
const size_t slot_threads = std::max<size_t>(1, settings.nthreads / (devices.size() + 1));
std::vector<std::unique_ptr<OwnSlot>> own(devices.size());
@@ -1079,6 +1109,8 @@ RefinementResult RefineModel(const gemmi::Structure &placed, const std::vector<M
// 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) {
if (!o.ran)
continue;
logger.Info("{}", o.line);
result.trials.push_back(o.trial);
const RefinementWeightTrial &t = result.trials.back();
@@ -1100,26 +1132,34 @@ RefinementResult RefineModel(const gemmi::Structure &placed, const std::vector<M
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)
b_trials.emplace_back(w, w * adp);
outcomes = RunTrials(trial_slots, b_trials, refinable, pos0, b0, settings);
choose(outcomes);
if (best >= 0 && !ctx.Cancelled()) {
const double w = result.trials[best].weight;
std::vector<std::pair<double, double>> b_trials;
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);
}
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);
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;
if (best >= 0) {
ctx.SetPositions(best_pos);
ctx.SetB(best_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;
} else
ctx.Refit(); // no trial ran: the model as placed, its scale and solvent for the maps
}
result.cancelled = ctx.Cancelled();
if (result.cancelled)
logger.Warning("Refinement: cancelled, {} weight trials ran", result.trials.size());
finish();
logger.Info("Refinement: done in {:.1f} s, {} target evaluations{}", result.seconds, ctx.evaluations,