diff --git a/image_analysis/structure_refinement/Refinement.cpp b/image_analysis/structure_refinement/Refinement.cpp index 981e2b5e5..9772c78e0 100644 --- a/image_analysis/structure_refinement/Refinement.cpp +++ b/image_analysis/structure_refinement/Refinement.cpp @@ -117,6 +117,9 @@ struct Context { std::vector rider_parent; std::vector rider_offset; int evaluations = 0; + const std::atomic *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(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 ¶m } ceres::GradientProblem problem(std::make_unique(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 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(ctx.cell, *ctx.sg, ctx.d_min, nthreads); slot->topo = BuildRefinementTopology(ctx.st, settings.riding_hydrogens); @@ -817,6 +841,7 @@ struct TrialOutcome { std::vector pos; std::vector 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 &refinable, const std::vector &pos0, @@ -826,7 +851,7 @@ TrialOutcome RunTrial(TrialSlot slot, const std::vector &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 &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 RunTrials(const std::vector &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(ctx.atoms.size()); double b_max = 0; @@ -1047,7 +1075,9 @@ RefinementResult RefineModel(const gemmi::Structure &placed, const std::vector slot_stage("slots"); - const int wanted = static_cast(std::max(settings.weight_ladder.size(), settings.adp_ladder.size())); + int wanted = static_cast(std::max(settings.weight_ladder.size(), settings.adp_ladder.size())); + if (settings.max_trials_at_once > 0) + wanted = std::min(wanted, static_cast(settings.max_trials_at_once)); const std::vector devices = OwnSlotDevices(ctx, wanted - 1, settings); const size_t slot_threads = std::max(1, settings.nthreads / (devices.size() + 1)); std::vector> own(devices.size()); @@ -1079,6 +1109,8 @@ RefinementResult RefineModel(const gemmi::Structure &placed, const std::vector &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 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> 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> 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, diff --git a/image_analysis/structure_refinement/Refinement.h b/image_analysis/structure_refinement/Refinement.h index 45fca1e7b..214bad54b 100644 --- a/image_analysis/structure_refinement/Refinement.h +++ b/image_analysis/structure_refinement/Refinement.h @@ -4,6 +4,7 @@ #pragma once #include +#include #include #include #include @@ -51,6 +52,13 @@ struct RefinementSettings { bool gpu = true; size_t nthreads = 1; size_t gpu_resident_bytes = 0; // held on the card by what runs beside + // The most weight trials run side by side, 0 for as many as the ladders and the cards take; the result does not + // depend on it (1 runs them one after another, as a check of that). + size_t max_trials_at_once = 0; + // Set by the caller (from any thread) to cut a running refinement short: checked between the minimiser's + // iterations and before every weight trial. The trials that ran are chosen among as usual, those that did not + // are left out (RefinementResult::cancelled); cut before the first trial, the model is the placement's. + const std::atomic *cancel = nullptr; }; // R by resolution: one of RefinementScore's shells, equal in width in 1/d^2 between the coarsest and the @@ -112,6 +120,7 @@ struct RefinementResult { double r_work_placement = 1, r_free_placement = 1; // After the restrained stage (NaN where it did not run). bool restrained = false; + bool cancelled = false; // cut short by RefinementSettings::cancel double r_work = NAN, r_free = NAN; double weight = 0; double adp_weight = 0; // 0: the B factors kept as placed diff --git a/tests/RefinementTest.cpp b/tests/RefinementTest.cpp index 926aa3254..78c741632 100644 --- a/tests/RefinementTest.cpp +++ b/tests/RefinementTest.cpp @@ -450,5 +450,16 @@ TEST_CASE("Refinement_restrained_recovers_model", "[Refinement]") { 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 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); } } diff --git a/tools/rugnux_refine_test.cpp b/tools/rugnux_refine_test.cpp index 9c367854e..e9d008682 100644 --- a/tools/rugnux_refine_test.cpp +++ b/tools/rugnux_refine_test.cpp @@ -37,6 +37,7 @@ void usage() { " --ladder the weight ladder, times the gradient-balanced weight\n" " -o write the refined model\n" " --dump with --score-only: h k l |Fcalc| phi |Fmask| phi k_total D var per reflection\n" + " --trials-at-once at most n weight trials side by side (1: one after another)\n" " -v debug lines (TIMING per refinement stage)\n"; } @@ -84,6 +85,8 @@ int main(int argc, char **argv) { settings.weight_ladder = ParseList(argv[++i]); else if (a == "-v") verbose = true; + else if (a == "--trials-at-once" && i + 1 < argc) + settings.max_trials_at_once = std::stoul(argv[++i]); else if (a == "--dump" && i + 1 < argc) dump_path = argv[++i]; else if (a == "-o" && i + 1 < argc)