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

Refinement: time per stage - an NVTX range per call and one TIMING line per stage with its sum and calls

Browse Source 
RefinementStage sums the time of each stage over a refinement (Refit and its parts, the evaluations and
their X-ray, restraint, host and upload shares, the repulsion build, the riding hydrogens, the rungs) and
logs it at debug level at the end; rugnux_refine_test -v shows it. No change to any result.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_015CeEGgpKdfkCTJb3YNMJcu
(cherry picked from commit 9491f54e0beccf36c65ebcec9afff90d08485f86)
(cherry picked from commit 0ce206a88df960754476ce2f75ffaed90fbf3123)
This commit is contained in:
leonarski_f committed 2026-10-26 01:38:22 +01:00
1 parent 07ae590e54
commit 24b1711d64
5 files changed
+123 -13
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

image_analysis/structure_refinement/Refinement.cpp
+56 -10
View File
@@ -6,6 +6,7 @@
#include <algorithm>
#include <chrono>
#include <map>
#include <optional>
#include <cmath>
#include <random>
#include <unordered_map>
@@ -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<double>(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<RefinementStage> s3("refit.scale");
gemmi::AsuData<std::complex<float>> fc, fm;
gemmi::AsuData<gemmi::ValueSigma<float>> 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<float>(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<std::array<double, 4>> 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<int> &param
int restraint_set, int iterations) {
if (param_atoms.empty())
return;
const RefinementStage stage(mode == AtomTarget::Mode::Xyz ? "minimise.xyz" : "minimise.b");
std::vector<double> 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<int>(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<std::array<double, 4>> g;
@@ -382,6 +404,7 @@ public:
int NumParameters() const override { return static_cast<int>(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<std::array<double, 4>> 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<std::vector<int>> 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<int> &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<MergedReflection> &merged, double wavelength_A, RefinementResult &result) {
const RefinementStage stage("maps");
ctx.Refit();
const std::vector<std::complex<float>> fmodel = ctx.FModel();
const gemmi::GroupOps gops = ctx.sg->operations();
@@ -574,6 +599,7 @@ void MakeMaps(Context &ctx, const std::vector<MergedReflection> &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<M
result.rigid_bodies, result.group_b ? "kept" : "not kept", r.r_work, r.r_free);
if (settings.restrained) {
std::optional<RefinementStage> 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<M
topo.bonds, topo.angles, topo.torsions, topo.chirals, topo.planes);
auto restraints_now = [&]() {
const RefinementStage stage("repulsions");
RestraintSet set = topo.geometry;
set.Append(BuildRepulsions(topo, ctx.Model(), cell, sg, REPULSION_MARGIN));
set.Finalize(ctx.atoms.size());
@@ -710,7 +740,11 @@ RefinementResult RefineModel(const gemmi::Structure &placed, const std::vector<M
RestraintSet repulsions = BuildRepulsions(topo, ctx.Model(), cell, sg, 0.0);
result.geometry_start = MeasureGeometry(topo.geometry, repulsions, ctx.XyzB());
}
const double w0 = BalancedWeight(ctx, refinable);
double w0 = 0;
{
const RefinementStage stage("balanced_weight");
w0 = BalancedWeight(ctx, refinable);
}
logger.Info("Refinement: geometry as given rmsd bonds {:.4f} A, angles {:.2f} deg, {} clashes; balanced X-ray "
"weight {:.3g}", result.geometry_start.rmsd_bonds, result.geometry_start.rmsd_angles,
result.geometry_start.clashes, w0);
@@ -725,12 +759,17 @@ RefinementResult RefineModel(const gemmi::Structure &placed, const std::vector<M
// of trust - tighter geometry, then fewer parameters - and a later one displaces an earlier only by more
// than the tie in R-free.
auto trial = [&](double w, double adp_w) {
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++) {
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<M
const RFactors rf = ctx.Refit();
t.r_work = rf.r_work;
t.r_free = rf.r_free;
t.geometry = MeasureGeometry(topo.geometry, BuildRepulsions(topo, ctx.Model(), cell, sg, 0.0), ctx.XyzB());
{
const RefinementStage s1("measure");
t.geometry = MeasureGeometry(topo.geometry, BuildRepulsions(topo, ctx.Model(), cell, sg, 0.0), ctx.XyzB());
}
t.within_bounds = t.geometry.rmsd_bonds <= MAX_RMSD_BONDS && t.geometry.rmsd_angles <= MAX_RMSD_ANGLES;
logger.Info("Refinement, X-ray weight {:.3g}, B {}: R-work {:.4f}, R-free {:.4f}, rmsd bonds {:.4f} A, "
"angles {:.2f} deg{}",
@@ -780,8 +822,11 @@ RefinementResult RefineModel(const gemmi::Structure &placed, const std::vector<M
}
MakeMaps(ctx, merged, wavelength_A, result);
result.blobs = FindDensityBlobs(*result.difference_map, result.difference_map_rms, ctx.st, result.d_min,
result.anomalous_map.get(), result.anomalous_map_rms);
{
const RefinementStage stage("blobs");
result.blobs = FindDensityBlobs(*result.difference_map, result.difference_map_rms, ctx.st, result.d_min,
result.anomalous_map.get(), result.anomalous_map_rms);
}
result.k_sol = ctx.k_sol;
result.b_sol = ctx.b_sol;
result.model = std::make_shared<gemmi::Structure>(ctx.st);
@@ -791,5 +836,6 @@ RefinementResult RefineModel(const gemmi::Structure &placed, const std::vector<M
result.restrained ? fmt::format("; R-work {:.4f}, R-free {:.4f}, weight {:.3g}", result.r_work,
result.r_free, result.weight)
: std::string());
RefinementStage::Log(logger);
return result;
}