diff --git a/image_analysis/structure_refinement/Refinement.cpp b/image_analysis/structure_refinement/Refinement.cpp index c0e3e5770..c53ab5ccc 100644 --- a/image_analysis/structure_refinement/Refinement.cpp +++ b/image_analysis/structure_refinement/Refinement.cpp @@ -180,6 +180,26 @@ struct Context { atoms[i]->b_iso = static_cast(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 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 OwnSlotDevices(const Context &main, int wanted, const Refinemen return std::vector(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 pos; std::vector 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 &refinable, const std::vector &pos0, - const std::vector &b0, double w, double adp_w, const RefinementSettings &settings) { + const std::vector &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 &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 &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 &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 &refinable, const s std::vector RunTrials(const std::vector &trial_slots, const std::vector> &trials, const std::vector &refinable, const std::vector &pos0, - const std::vector &b0, const RefinementSettings &settings) { + const std::vector &b0, const Context::Fit *fit, int cycles, + const RefinementSettings &settings) { std::vector out(trials.size()); std::vector 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(1, settings.nthreads / running); auto run_slot = [&](size_t s) { try { #ifdef JFJOCH_USE_CUDA @@ -1038,7 +1069,8 @@ std::vector RunTrials(const std::vector &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 pos0 = ctx.Positions(); const std::vector 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 slot_stage("slots"); - 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(), + settings.final_ladder.size(), size_t{3}})); 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); @@ -1238,58 +1271,85 @@ RefinementResult RefineModel(const gemmi::Structure &placed, const std::vector best_pos; - std::vector 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 &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 &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(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(k); } + return best; }; - // the positions' weight, with the B factors as they are - std::vector> xyz_trials; - for (double ladder : settings.weight_ladder) - xyz_trials.emplace_back(w0 * ladder, 0.0); - 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 - if (best >= 0 && !ctx.Cancelled()) { - const double w = result.trials[best].weight; - std::vector> 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> 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 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> 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> 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 } diff --git a/image_analysis/structure_refinement/Refinement.h b/image_analysis/structure_refinement/Refinement.h index a2c10b2d2..e083b6000 100644 --- a/image_analysis/structure_refinement/Refinement.h +++ b/image_analysis/structure_refinement/Refinement.h @@ -31,9 +31,10 @@ class Logger; // 2. Restrained: every atom's position and isotropic B, against the maximum-likelihood amplitude target with // the geometry the dictionary describes (MonomerLibrary.h: bonds, angles, planes, chiral volumes, // torsions, repulsions; B similarity of bonded atoms). Residues it does not describe are held fixed and -// stay in the structure factors and the bulk-solvent mask. The weight of the X-ray term against the -// geometry is chosen from a ladder of full refinements, by R-free among those whose geometry stays within -// bounds; then whether, and at what weight, to refine the B factors as well, the same way. +// stay in the structure factors and the bulk-solvent mask. The weights of the X-ray term against the +// geometry, and whether to refine the B factors, are chosen by R-free among trials whose geometry stays +// within bounds: a ladder from the placed model, then a trajectory of steps that each try the weights as +// they are, halved and doubled from where the last step ended until no step gains, then a short final ladder. // // The library reads and writes nothing: the model and the reflections come in, the refined model, the // statistics and the map coefficients go out. Nothing it finds feeds back into processing decisions. @@ -41,13 +42,18 @@ struct RefinementSettings { bool restrained = true; // stage 2; false stops after the placement stage bool score_only = false; // only score the model as given (RefinementResult::score_start), no stage bool group_b = true; // the placement stage's B shift per residue - int macro_cycles = 10; - int xyz_iterations = 10; // L-BFGS iterations per macro-cycle - int b_iterations = 5; - // The positions' X-ray weight, times the gradient-balanced one (BalancedWeight), tried in this order. - std::vector weight_ladder = {1.0 / 1024, 1.0 / 512, 1.0 / 256, 1.0 / 128, 1.0 / 64, 1.0 / 32}; - // The X-ray weights of the B refinement tried next, times the positions' chosen weight. - std::vector adp_ladder = {1.0, 4.0, 16.0}; + int macro_cycles = 2; // per weight trial, from the model the trial starts from + int xyz_iterations = 30; // L-BFGS iterations per macro-cycle + int b_iterations = 20; + // The first trials, from the placed model: the positions' X-ray weight, times the gradient-balanced one + // (BalancedWeight), with the B factors' X-ray weight that weight times each of adp_ladder (0: B held). + std::vector weight_ladder = {1.0 / 1024, 1.0 / 128, 1.0 / 16}; + std::vector adp_ladder = {0.0, 4.0}; + // Then at most this many steps of the trajectory, each a set of trials from where the last one ended, + int max_steps = 8; + // and the final trials from where it ended, at its weights times each of these; + std::vector final_ladder = {0.5, 1.0, 2.0}; + double max_weight = 1.0 / 8; // the positions' X-ray weight at most, times the balanced one bool riding_hydrogens = true; // the dictionary's hydrogens added, riding on their parents bool gpu = true; size_t nthreads = 1; diff --git a/tests/RefinementTest.cpp b/tests/RefinementTest.cpp index 7081ba73d..826087805 100644 --- a/tests/RefinementTest.cpp +++ b/tests/RefinementTest.cpp @@ -631,9 +631,10 @@ TEST_CASE("Refinement_restrained_recovers_model", "[Refinement]") { RefinementSettings settings; settings.gpu = gpu; settings.nthreads = 4; - settings.macro_cycles = 3; - settings.weight_ladder = {0.0625, 0.25}; + settings.weight_ladder = {0.0625}; settings.adp_ladder = {4.0}; + settings.max_steps = 2; + settings.final_ladder = {1.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); diff --git a/tools/rugnux_refine_test.cpp b/tools/rugnux_refine_test.cpp index e9d008682..f45272611 100644 --- a/tools/rugnux_refine_test.cpp +++ b/tools/rugnux_refine_test.cpp @@ -31,10 +31,13 @@ void usage() { " --score-only score the model as given, refine nothing\n" " --no-group-b no B shift per residue in the placement stage\n" " --no-hydrogens no riding hydrogens\n" - " --adp-ladder the B refinement's weights, times the positions' chosen weight\n" - " --cycles macro-cycles per weight (default as in RefinementSettings)\n" + " --ladder the first weight trials, times the gradient-balanced weight\n" + " --adp-ladder each with the B factors' weight that weight times these (0: B held)\n" + " --steps the trajectory's steps at most\n" + " --final-ladder the final trials, times the trajectory's weights\n" + " --max-weight the positions' weight at most, times the gradient-balanced weight\n" + " --cycles macro-cycles per weight trial (default as in RefinementSettings)\n" " --xyz-iterations , --b-iterations L-BFGS iterations per macro-cycle\n" - " --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" @@ -83,6 +86,12 @@ int main(int argc, char **argv) { settings.b_iterations = std::stoi(argv[++i]); else if (a == "--ladder" && i + 1 < argc) settings.weight_ladder = ParseList(argv[++i]); + else if (a == "--steps" && i + 1 < argc) + settings.max_steps = std::stoi(argv[++i]); + else if (a == "--final-ladder" && i + 1 < argc) + settings.final_ladder = ParseList(argv[++i]); + else if (a == "--max-weight" && i + 1 < argc) + settings.max_weight = std::stod(argv[++i]); else if (a == "-v") verbose = true; else if (a == "--trials-at-once" && i + 1 < argc)