diff --git a/image_analysis/structure_refinement/Refinement.cpp b/image_analysis/structure_refinement/Refinement.cpp index 734d27540..ea8771af8 100644 --- a/image_analysis/structure_refinement/Refinement.cpp +++ b/image_analysis/structure_refinement/Refinement.cpp @@ -6,6 +6,7 @@ #include #include #include +#include #include #include #include @@ -159,13 +160,21 @@ struct Context { // 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. RFactors Refit() { + const RefinementStage stage("refit"); double b_min = b_ceiling; for (const gemmi::Atom *a : atoms) b_min = std::min(b_min, a->b_iso); b_floor = std::max(B_FLOOR, b_min - B_STEP_DOWN); engine->SetBlur(RefinementBlur(cell, *sg, d_min, b_floor)); - engine->Fcalc(Model(), fcalc); - engine->Fmask(Model(), fmask); + { + const RefinementStage s1("refit.fcalc"); + engine->Fcalc(Model(), fcalc); + } + { + const RefinementStage s2("refit.fmask"); + engine->Fmask(Model(), fmask); + } + std::optional s3("refit.scale"); gemmi::AsuData> fc, fm; gemmi::AsuData> fo; @@ -186,6 +195,7 @@ struct Context { else #endif FitModelScale(scaling, {.b_hi = SolventBMax(Model())}, nthreads); + s3.reset(); k_sol = scaling.k_sol; b_sol = scaling.b_sol; rows.k_total.resize(refl.Size()); @@ -196,7 +206,10 @@ struct Context { rows.fmask[m] = static_cast(scaling.get_solvent_scale(cell.calculate_stol_sq(refl.hkl[m]))) * fmask[m]; fmodel[m] = rows.k_total[m] * (fcalc[m] + rows.fmask[m]); } - EstimateLikelihoodParameters(refl, fmodel, rows); + { + const RefinementStage s4("refit.ml"); + EstimateLikelihoodParameters(refl, fmodel, rows); + } engine->SetRowModel(rows); return ComputeRFactors(refl, fmodel); } @@ -235,6 +248,7 @@ public: } bool Evaluate(const double *p, double *cost, double *gradient) const override { + const RefinementStage stage(mode_ == Mode::Xyz ? "eval.xyz" : "eval.b"); for (size_t k = 0; k < param_atoms_.size(); k++) { gemmi::Atom &a = *ctx_.atoms[param_atoms_[k]]; if (mode_ == Mode::Xyz) @@ -245,12 +259,18 @@ public: ctx_.MoveRiders(); ctx_.evaluations++; std::vector> gx, gr; - const double xray = ctx_.engine->Evaluate(ctx_.Model(), gradient != nullptr ? &gx : nullptr); + double xray = 0; + { + const RefinementStage s1("eval.xray"); + xray = ctx_.engine->Evaluate(ctx_.Model(), gradient != nullptr ? &gx : nullptr); + } if (gradient != nullptr) ctx_.FoldRiders(gx); double restraints = 0; - if (restraint_set_ >= 0) + if (restraint_set_ >= 0) { + const RefinementStage s2("eval.restraints"); restraints = ctx_.engine->EvaluateRestraints(restraint_set_, ctx_.XyzB(), gradient != nullptr ? &gr : nullptr); + } *cost = weight_ * xray + restraints; if (gradient != nullptr) for (size_t k = 0; k < param_atoms_.size(); k++) { @@ -279,6 +299,7 @@ void Minimise(Context &ctx, AtomTarget::Mode mode, const std::vector ¶m int restraint_set, int iterations) { if (param_atoms.empty()) return; + const RefinementStage stage(mode == AtomTarget::Mode::Xyz ? "minimise.xyz" : "minimise.b"); std::vector p; for (int i : param_atoms) { const gemmi::Atom &a = *ctx.atoms[i]; @@ -325,6 +346,7 @@ public: int NumParameters() const override { return static_cast(6 * bodies_.size()); } bool Evaluate(const double *q, double *cost, double *gradient) const override { + const RefinementStage stage("eval.rigid"); Place(q); ctx_.evaluations++; std::vector> g; @@ -382,6 +404,7 @@ public: int NumParameters() const override { return static_cast(groups_.size()); } bool Evaluate(const double *q, double *cost, double *gradient) const override { + const RefinementStage stage("eval.group_b"); Place(q); ctx_.evaluations++; std::vector> g; @@ -412,6 +435,7 @@ private: // The placement stage: rigid bodies, then group B. void RefinePlacement(Context &ctx, bool group_b, RefinementResult &result) { + const RefinementStage stage("placement"); std::vector> bodies, groups; size_t i = 0; for (const gemmi::Chain &ch : ctx.Model().chains) { @@ -490,6 +514,7 @@ double BalancedWeight(Context &ctx, const std::vector &refinable) { // The model's map coefficients and its difference map, as model validation makes them (sigma_A weighted). void MakeMaps(Context &ctx, const std::vector &merged, double wavelength_A, RefinementResult &result) { + const RefinementStage stage("maps"); ctx.Refit(); const std::vector> fmodel = ctx.FModel(); const gemmi::GroupOps gops = ctx.sg->operations(); @@ -574,6 +599,7 @@ void MakeMaps(Context &ctx, const std::vector &merged, double // The engine for the model as it is now: on the GPU where it is wanted and takes the case within its budget, // on the CPU otherwise. void MakeEngine(Context &ctx, bool gpu, const RefinementSettings &settings, Logger &logger) { + const RefinementStage stage("engine"); ctx.engine.reset(); #ifdef JFJOCH_USE_CUDA ctx.scale_gpu.reset(); @@ -664,11 +690,14 @@ RefinementResult RefineModel(const gemmi::Structure &placed, const std::vector 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++) @@ -698,6 +727,7 @@ RefinementResult RefineModel(const gemmi::Structure &placed, const std::vector 0 ? "rung.b" : "rung.xyz"); ctx.SetPositions(pos0); ctx.SetB(b0); for (int cycle = 0; cycle < settings.macro_cycles; cycle++) { place_riders(); ctx.Refit(); - ctx.engine->SetRestraints(0, restraints_now()); + { + const RestraintSet set = restraints_now(); + 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); @@ -742,7 +781,10 @@ RefinementResult RefineModel(const gemmi::Structure &placed, const std::vector(ctx.st); @@ -791,5 +836,6 @@ RefinementResult RefineModel(const gemmi::Structure &placed, const std::vector #include +#include #include "gemmi/dencalc.hpp" #include "gemmi/it92.hpp" @@ -194,7 +195,11 @@ public: double Evaluate(const gemmi::Model &model, std::vector> *grad) override { SetAtoms(model); std::vector> g; - const double value = engine_->Evaluate(grad != nullptr ? &g : nullptr); + double value = 0; + { + const RefinementStage stage("gpu.evaluate"); + value = engine_->Evaluate(grad != nullptr ? &g : nullptr); + } if (grad != nullptr) { grad->resize(g.size()); for (size_t i = 0; i < g.size(); i++) @@ -213,7 +218,11 @@ public: for (size_t i = 0; i < xyzb.size(); i++) x[i] = {static_cast(xyzb[i][0]), static_cast(xyzb[i][1]), static_cast(xyzb[i][2]), static_cast(xyzb[i][3])}; - const double value = engine_->EvaluateRestraints(which, x, grad != nullptr ? &g : nullptr); + double value = 0; + { + const RefinementStage stage("gpu.restraints"); + value = engine_->EvaluateRestraints(which, x, grad != nullptr ? &g : nullptr); + } if (grad != nullptr) { grad->resize(g.size()); for (size_t i = 0; i < g.size(); i++) @@ -235,6 +244,7 @@ private: // The atoms' densities (made again only for the atoms whose B moved) and grid positions, to the device. void SetAtoms(const gemmi::Model &model) { + std::optional stage("gpu.atoms_host"); const std::vector atoms = Atoms(model); if (last_b_.size() != atoms.size()) { last_b_.assign(atoms.size(), NAN); @@ -250,6 +260,8 @@ private: const gemmi::Fractional f = cell_.fractionalize(atoms[i]->pos); pos[i] = {static_cast(f.x * n3[0]), static_cast(f.y * n3[1]), static_cast(f.z * n3[2]), 0}; } + stage.reset(); + const RefinementStage upload("gpu.upload"); engine_->SetAtoms(atoms_, pos); } diff --git a/image_analysis/structure_refinement/RefinementXray.cpp b/image_analysis/structure_refinement/RefinementXray.cpp index 5a3da3f4d..ccbe8e3ab 100644 --- a/image_analysis/structure_refinement/RefinementXray.cpp +++ b/image_analysis/structure_refinement/RefinementXray.cpp @@ -6,6 +6,7 @@ #include #include #include +#include #include @@ -13,6 +14,7 @@ #include "gemmi/it92.hpp" #include "gemmi/math.hpp" +#include "../../common/Logger.h" #include "../../common/ParallelFor.h" #include "ModelFFT.h" #include "ModelGrid.h" @@ -310,6 +312,30 @@ RFactors ComputeRFactors(const RefinementReflections &refl, const std::vector> stage_times; +} // namespace + +RefinementStage::RefinementStage(const char *name) + : name_(name), start_(std::chrono::steady_clock::now()), nvtx_(NvtxRangeStart(name)) {} + +RefinementStage::~RefinementStage() { + NvtxRangeEnd(nvtx_); + const double seconds = std::chrono::duration(std::chrono::steady_clock::now() - start_).count(); + std::lock_guard lock(stage_mutex); + std::pair &t = stage_times[name_]; + t.first += seconds; + t.second++; +} + +void RefinementStage::Log(Logger &logger) { + std::lock_guard lock(stage_mutex); + for (const auto &[name, t] : stage_times) + logger.Debug("TIMING refine.{} {:.3f} {}", name, t.first, t.second); + stage_times.clear(); +} + double RefinementBlur(const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg, double d_min, double b_floor) { // set_refmac_compatible_blur(): u_to_b / 1.1 * spacing^2 - b_min, not below 0, the spacing the one asked for const gemmi::DensityCalculator dc = Calculator(cell, sg, d_min, 0); diff --git a/image_analysis/structure_refinement/RefinementXray.h b/image_analysis/structure_refinement/RefinementXray.h index 9bbe976cd..a0a459f21 100644 --- a/image_analysis/structure_refinement/RefinementXray.h +++ b/image_analysis/structure_refinement/RefinementXray.h @@ -4,6 +4,8 @@ #pragma once #include +#include +#include #include #include #include @@ -18,6 +20,8 @@ #include "RefinementRestraints.h" // RestraintSet #include "RigidBodyRefine.h" // SymmetryComposition +class Logger; + // The X-ray side of restrained refinement (Refinement.h): the observed amplitudes, the model's structure // factors and the maximum-likelihood target that compares them, with its gradient with respect to every // atom's position and B. @@ -96,6 +100,23 @@ private: RestraintSet restraints_[2]; }; +// The time a refinement spends in one stage, summed over every call (and over threads) until Log(): the scope of +// an instance is one call, and an NVTX range of the stage's name. +class RefinementStage { +public: + explicit RefinementStage(const char *name); + ~RefinementStage(); + RefinementStage(const RefinementStage &) = delete; + RefinementStage &operator=(const RefinementStage &) = delete; + // "TIMING refine. " per stage at debug level, then the sums start again. + static void Log(Logger &logger); + +private: + const char *name_; + std::chrono::steady_clock::time_point start_; + uint64_t nvtx_; +}; + // The blur for a model whose smallest B may go as low as b_floor: DensityCalculator's // set_refmac_compatible_blur() for that B, on the refinement grid of d_min. double RefinementBlur(const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg, double d_min, double b_floor); diff --git a/tools/rugnux_refine_test.cpp b/tools/rugnux_refine_test.cpp index 9c1147cd1..037613063 100644 --- a/tools/rugnux_refine_test.cpp +++ b/tools/rugnux_refine_test.cpp @@ -34,7 +34,8 @@ void usage() { " --cycles macro-cycles per weight (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"; + " -o write the refined model\n" + " -v debug lines (TIMING per refinement stage)\n"; } std::vector ParseList(const std::string &s) { @@ -52,6 +53,7 @@ int main(int argc, char **argv) { RefinementSettings settings; settings.nthreads = 8; std::string out_path; + bool verbose = false; std::vector files; for (int i = 1; i < argc; i++) { const std::string a = argv[i]; @@ -76,6 +78,8 @@ 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 == "-v") + verbose = true; else if (a == "-o" && i + 1 < argc) out_path = argv[++i]; else if (!a.empty() && a[0] == '-') { @@ -119,6 +123,7 @@ int main(int argc, char **argv) { gemmi::Structure st = gemmi::read_structure_gz(files[1]); Logger logger("rugnux_refine_test"); + logger.Verbose(verbose); const RefinementResult r = RefineModel(st, merged, mtz.cell, *sg, 1.0, logger, settings); if (!r.ok) { std::cerr << "Refinement failed: " << r.failure_reason << std::endl;