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

Refinement: the per-evaluation host chain on the threads and in place

Browse Source 
Byte-identical: every pass is element by element in the serial loop's order, or a copy that is no longer made.

- The engine's gradients are read in place from its pinned buffers and widened to double on the threads
  (RefinementGPUEngine::HostGradient/HostRestraintGradient); the positions go up from a pinned vector of the
  engine's host side; the model's atom pointers are kept between calls (the model is the context's own).
- Context: the riders moved and folded on the threads - a parent takes its riders in index order through a
  CSR list (SetRiderParents), as the serial walk added them.
- AtomTarget: the parameters set and the gradient assembled on the threads; the per-atom gradient vectors kept
  between evaluations (7 MB each on a large model: made anew, their pages fault in every time).
- RefinementMinimiser: the sweeps over the parameters (x + alpha d, s, y) on the threads it is given, the dot
  products the serial chains they were; Direction()'s working vector and a let-go pair's vectors reused.

A/B against the frozen combined2 binary, one lock hold each, outputs identical but provenance lines:
cytc refinement 1.600 -> 1.548 s (live), 1.724 -> 1.668 (CLI); thau 1.906 -> 1.885, 2.031 -> 2.027;
4nwv (with the commit before) 39.9 -> 35.5 s, the ladder 14.6 -> 12.7, each step 8.3 -> 7.2.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_015CeEGgpKdfkCTJb3YNMJcu
This commit is contained in:
leonarski_fandClaude Fable 5.1 committed 2026-10-26 22:48:28 +01:00
1 parent b16873036c
commit d122118d73
5 files changed
+174 -99
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

image_analysis/structure_refinement/Refinement.cpp
+67 -34
View File
@@ -28,6 +28,7 @@
#include "../../common/JFJochException.h"
#include "../../common/JFJochMath.h"
#include "../../common/Logger.h"
#include "../../common/ParallelFor.h"
#include "ModelFFT.h"
#include "ModelScaling.h"
#include "RefinementMath.h"
@@ -137,6 +138,7 @@ struct Context {
double b_ceiling = B_CEILING_MAX;
// Riding hydrogens: per atom the parent it rides on (-1: none), and where it sits from the parent.
std::vector<int> rider_parent;
std::vector<int> rider_start, rider_index; // SetRiderParents()
std::vector<gemmi::Position> rider_offset;
int evaluations = 0;
const std::atomic<bool> *cancel = nullptr; // RefinementSettings::cancel
@@ -145,20 +147,42 @@ struct Context {
gemmi::Model &Model() { return st.models[0]; }
// The hydrogens moved with their parents, at the offsets they were last placed at, and given their B.
void MoveRiders() {
for (size_t i = 0; i < rider_parent.size(); i++)
if (rider_parent[i] >= 0) {
atoms[i]->pos = atoms[rider_parent[i]]->pos + rider_offset[i];
atoms[i]->b_iso = atoms[rider_parent[i]]->b_iso;
}
}
// What the target says of a hydrogen is said of its parent.
void FoldRiders(std::vector<std::array<double, 4>> &g) const {
// The riders' parents, and per parent its riders in index order (rider_start, rider_index: CSR over the atoms).
void SetRiderParents(const std::vector<int> &parents) {
rider_parent = parents;
rider_start.assign(atoms.size() + 1, 0);
for (int p : rider_parent)
if (p >= 0)
rider_start[p + 1]++;
for (size_t i = 0; i < atoms.size(); i++)
rider_start[i + 1] += rider_start[i];
rider_index.resize(rider_start.back());
std::vector<int> fill(rider_start.begin(), rider_start.end() - 1);
for (size_t i = 0; i < rider_parent.size(); i++)
if (rider_parent[i] >= 0)
for (int c = 0; c < 4; c++)
g[rider_parent[i]][c] += g[i][c];
rider_index[fill[rider_parent[i]]++] = static_cast<int>(i);
}
// The hydrogens moved with their parents, at the offsets they were last placed at, and given their B.
void MoveRiders() {
const int n = static_cast<int>(rider_parent.size());
ParallelChunks(n, ThreadsForWork(n, nthreads, 16384), [&](int lo, int hi) {
for (int i = lo; i < hi; i++)
if (rider_parent[i] >= 0) {
atoms[i]->pos = atoms[rider_parent[i]]->pos + rider_offset[i];
atoms[i]->b_iso = atoms[rider_parent[i]]->b_iso;
}
});
}
// What the target says of a hydrogen is said of its parent: each parent takes its riders in index order, as
// the serial walk over the atoms added them (a rider has no riders of its own).
void FoldRiders(std::vector<std::array<double, 4>> &g) const {
const int n = static_cast<int>(rider_parent.size());
ParallelChunks(n, ThreadsForWork(n, nthreads, 16384), [&](int lo, int hi) {
for (int i = lo; i < hi; i++)
for (int r = rider_start[i]; r < rider_start[i + 1]; r++)
for (int c = 0; c < 4; c++)
g[i][c] += g[rider_index[r]][c];
});
}
std::vector<std::array<double, 4>> XyzB() const {
@@ -509,13 +533,16 @@ public:
}
// The atoms set to the parameters, their riders with them.
void Set(const double *p) const {
for (size_t k = 0; k < param_atoms_.size(); k++) {
gemmi::Atom &a = *ctx_.atoms[param_atoms_[k]];
if (mode_ == Mode::Xyz)
a.pos = gemmi::Position(p[3 * k], p[3 * k + 1], p[3 * k + 2]);
else
a.b_iso = static_cast<float>(std::clamp(p[k], ctx_.b_floor, ctx_.b_ceiling));
}
const int n = static_cast<int>(param_atoms_.size());
ParallelChunks(n, ThreadsForWork(n, ctx_.nthreads, 16384), [&](int lo, int hi) {
for (int k = lo; k < hi; k++) {
gemmi::Atom &a = *ctx_.atoms[param_atoms_[k]];
if (mode_ == Mode::Xyz)
a.pos = gemmi::Position(p[3 * k], p[3 * k + 1], p[3 * k + 2]);
else
a.b_iso = static_cast<float>(std::clamp(p[k], ctx_.b_floor, ctx_.b_ceiling));
}
});
ctx_.MoveRiders();
}
// the most a parameter moves in the first trial of a line search
@@ -525,7 +552,7 @@ public:
const RefinementStage stage(mode_ == Mode::Xyz ? "eval.xyz" : "eval.b");
Set(p);
ctx_.evaluations++;
std::vector<std::array<double, 4>> gx, gr;
std::vector<std::array<double, 4>> &gx = gx_, &gr = gr_;
double xray = 0, restraints = 0;
{
const RefinementStage s1("eval.engine");
@@ -536,18 +563,21 @@ public:
if (gradient == nullptr)
return std::isfinite(*cost);
ctx_.FoldRiders(gx);
for (size_t k = 0; k < param_atoms_.size(); k++) {
const int i = param_atoms_[k];
if (mode_ == Mode::Xyz)
for (int c = 0; c < 3; c++)
gradient[3 * k + c] = weight_ * gx[i][c] + gr[i][c];
else {
// held at a bound, B does not move further out: by the whole target, restraints included
const double g = weight_ * gx[i][3] + gr[i][3];
const bool held = (p[k] <= ctx_.b_floor && g > 0) || (p[k] >= ctx_.b_ceiling && g < 0);
gradient[k] = held ? 0.0 : g;
const int n = static_cast<int>(param_atoms_.size());
ParallelChunks(n, ThreadsForWork(n, ctx_.nthreads, 16384), [&](int lo, int hi) {
for (int k = lo; k < hi; k++) {
const int i = param_atoms_[k];
if (mode_ == Mode::Xyz)
for (int c = 0; c < 3; c++)
gradient[3 * k + c] = weight_ * gx[i][c] + gr[i][c];
else {
// held at a bound, B does not move further out: by the whole target, restraints included
const double g = weight_ * gx[i][3] + gr[i][3];
const bool held = (p[k] <= ctx_.b_floor && g > 0) || (p[k] >= ctx_.b_ceiling && g < 0);
gradient[k] = held ? 0.0 : g;
}
}
}
});
return std::isfinite(*cost);
}
@@ -557,6 +587,9 @@ private:
const std::vector<int> &param_atoms_;
double weight_;
int restraint_set_;
// the gradients per atom, kept from evaluation to evaluation (7 MB each on a large model: made anew, their
// pages fault in every time)
mutable std::vector<std::array<double, 4>> gx_, gr_;
};
// The minimiser's diagonal for the B factors of `param_atoms` (RefinementMinimiser::SetDiagonal): weight times the
@@ -998,7 +1031,7 @@ std::unique_ptr<OwnSlot> MakeOwnSlot(const gemmi::Structure &before_topology, co
ctx.mask = std::make_unique<BulkSolventMask>(ctx.cell, *ctx.sg, ctx.d_min, nthreads);
slot->topo = BuildRefinementTopology(ctx.st, settings.riding_hydrogens);
ctx.atoms = AtomList(ctx.Model());
ctx.rider_parent = slot->topo.rider_parent;
ctx.SetRiderParents(slot->topo.rider_parent);
#ifdef JFJOCH_USE_CUDA
if (main.engine->OnGPU()) {
try {
@@ -1077,7 +1110,7 @@ TrialOutcome RunTrial(TrialSlot slot, const std::vector<int> &refinable, const s
ctx.SetB(b0);
// the positions and the B factors in turn, each with a memory of its own kept from macro-cycle to macro-cycle
const AtomTarget xyz(ctx, AtomTarget::Mode::Xyz, refinable, w, 0), b(ctx, AtomTarget::Mode::B, refinable, adp_w, 1);
RefinementMinimiser xyz_minimiser, b_minimiser;
RefinementMinimiser xyz_minimiser(20, ctx.nthreads), b_minimiser(20, ctx.nthreads);
double xray_scale = 0;
for (int cycle = 0; cycle < cycles && !ctx.Cancelled(); cycle++) {
PlaceRiders(ctx, topo);
@@ -1241,7 +1274,7 @@ RefinementResult RefineModel(const gemmi::Structure &placed, const std::vector<M
topo = BuildRefinementTopology(ctx.st, settings.riding_hydrogens);
topo_stage.reset();
ctx.atoms = AtomList(ctx.Model());
ctx.rider_parent = topo.rider_parent;
ctx.SetRiderParents(topo.rider_parent);
PlaceRiders(ctx, topo);
}