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

Refinement: the refit's per-reflection loops and the solvent build's per-slot edits on the threads

Browse Source 
- Refit: the AsuData of F_calc and F_mask, the row model's k_total and scaled F_mask, and the scale
  pool's packing of the points are filled in parallel, each element on its own (refit.scale on 4YBB
  84 -> 70 s summed over the slots).
- The solvent build's edits are collected once and applied to every slot's model and topology on a
  thread per slot, with the slot's topology copy for the undo: they were made slot after slot, so
  their cost grew with the number of slots. Timed as refine.build.edit.
- Test: Refinement_restrained_recovers_model refines again with one trial at a time and requires the
  same model bit for bit as three side by side.

4YBB refined model CIF byte-identical to the 1-slot reference run.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01HVeDQqN9qRD57vnShzTC37
This commit is contained in:
leonarski_fandClaude Opus 5.5 committed 2026-10-27 23:39:14 +01:00
1 parent a003887337
commit 18ea16d5d3
3 files changed
+73 -34
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

image_analysis/structure_refinement/ModelScaling.cpp
+11 -9
View File
@@ -523,15 +523,17 @@ ModelScaleReport ModelScaleGPUPool::Fit(gemmi::Scaling<float> &scaling, double b
std::vector<double> stol2(n);
std::vector<float> fobs(n), sigma(n);
std::vector<float2> fc(n), fm(n);
for (size_t i = 0; i < n; i++) {
const auto &p = scaling.points[i];
hkl[i] = {p.hkl[0], p.hkl[1], p.hkl[2]};
stol2[i] = p.stol2;
fobs[i] = p.fobs;
sigma[i] = p.sigma;
fc[i] = {p.fcmol.real(), p.fcmol.imag()};
fm[i] = {p.fmask.real(), p.fmask.imag()};
}
ParallelChunks(static_cast<int>(n), ThreadsForWork(n, nthreads), [&](int lo, int hi) {
for (int i = lo; i < hi; i++) {
const auto &p = scaling.points[i];
hkl[i] = {p.hkl[0], p.hkl[1], p.hkl[2]};
stol2[i] = p.stol2;
fobs[i] = p.fobs;
sigma[i] = p.sigma;
fc[i] = {p.fcmol.real(), p.fcmol.imag()};
fm[i] = {p.fmask.real(), p.fmask.imag()};
}
});
std::vector<std::array<double, 6>> constraints;
for (const gemmi::Vec6 &c : scaling.constraint_matrix)
constraints.push_back(c);
image_analysis/structure_refinement/Refinement.cpp
+47 -25
View File
@@ -391,16 +391,25 @@ struct Context {
}
std::optional<RefinementStage> s3("refit.scale");
// the per-reflection loops on the threads, each element on its own
const int n_refl = static_cast<int>(refl.Size());
const size_t refl_threads = ThreadsForWork(refl.Size(), nthreads);
gemmi::AsuData<std::complex<float>> fc, fm;
gemmi::AsuData<gemmi::ValueSigma<float>> fo;
fc.unit_cell_ = fm.unit_cell_ = fo.unit_cell_ = cell;
fc.spacegroup_ = fm.spacegroup_ = fo.spacegroup_ = sg;
for (size_t m = 0; m < refl.Size(); m++) {
fc.v.push_back({refl.hkl[m], fcalc[m]});
fm.v.push_back({refl.hkl[m], fmask[m]});
fc.v.resize(refl.Size());
fm.v.resize(refl.Size());
ParallelChunks(n_refl, refl_threads, [&](int lo, int hi) {
for (int m = lo; m < hi; m++) {
fc.v[m] = {refl.hkl[m], fcalc[m]};
fm.v[m] = {refl.hkl[m], fmask[m]};
}
});
fo.v.reserve(refl.Size());
for (size_t m = 0; m < refl.Size(); m++)
if (!refl.free[m])
fo.v.push_back({refl.hkl[m], {refl.fo[m], 1.0f}});
}
gemmi::Scaling<float> scaling(cell, sg);
scaling.use_solvent = true;
// the scale of the score: by robust least squares (ModelScaleBox::robust)
@@ -420,10 +429,13 @@ struct Context {
b_star = scaling.b_star;
rows.k_total.resize(refl.Size());
rows.fmask.resize(refl.Rows());
for (size_t m = 0; m < refl.Size(); m++) {
rows.k_total[m] = static_cast<float>(scaling.get_overall_scale_factor(refl.hkl[m]));
rows.fmask[m] = static_cast<float>(scaling.get_solvent_scale(cell.calculate_stol_sq(refl.hkl[m]))) * fmask[m];
}
ParallelChunks(n_refl, refl_threads, [&](int lo, int hi) {
for (int m = lo; m < hi; m++) {
rows.k_total[m] = static_cast<float>(scaling.get_overall_scale_factor(refl.hkl[m]));
rows.fmask[m] = static_cast<float>(scaling.get_solvent_scale(cell.calculate_stol_sq(refl.hkl[m]))) *
fmask[m];
}
});
for (size_t m = refl.Size(); m < refl.Rows(); m++) {
const gemmi::Miller &h = refl.mate_hkl[m - refl.Size()];
rows.fmask[m] = static_cast<float>(scaling.get_solvent_scale(cell.calculate_stol_sq(h))) * fmask[m];
@@ -2511,18 +2523,12 @@ RefinementResult RefineOnce(const gemmi::Structure &placed, const std::vector<Me
gemmi::EntityType entity;
};
std::vector<std::vector<Undo>> undo(trial_slots.size());
std::vector<RefinementTopology> topo_before;
for (const TrialSlot &slot : trial_slots)
topo_before.push_back(*slot.topo);
std::vector<RefinementTopology> topo_before(trial_slots.size());
const SolventPool pool_before = pool;
auto edit_atom = [&](int index, const std::function<void(gemmi::Atom &, gemmi::Residue &, size_t)> &f) {
for (size_t sl = 0; sl < trial_slots.size(); sl++) {
gemmi::Residue &r = *residue_of[sl][index];
gemmi::Atom &a = *trial_slots[sl].ctx->atoms[index];
undo[sl].push_back({index, a, r.name, r.entity_type});
f(a, r, sl);
}
};
// the edits, in order, made below in every slot - each slot on a thread of its own
using Edit = std::function<void(gemmi::Atom &, gemmi::Residue &, size_t)>;
std::vector<std::pair<int, Edit>> edits;
auto edit_atom = [&](int index, Edit f) { edits.emplace_back(index, std::move(f)); };
std::vector<SolventEdit> kept;
for (SolventEdit &e : picked.edits) {
if (e.kind == SolventEdit::Kind::Added) {
@@ -2535,14 +2541,15 @@ RefinementResult RefineOnce(const gemmi::Structure &placed, const std::vector<Me
// on a special position (PickSolvent put it exactly there), held on it
SpecialPosition site = SiteSymmetry(view.cell, e.pos);
site.atom = pool.first + k;
edit_atom(pool.first + k, [&](gemmi::Atom &a, gemmi::Residue &, size_t sl) {
const int index = pool.first + k;
edit_atom(index, [&trial_slots, e, index, site](gemmi::Atom &a, gemmi::Residue &, size_t sl) {
a.pos = e.pos;
a.occ = static_cast<float>(e.occupancy);
a.b_iso = static_cast<float>(e.b_iso);
RefinementTopology &t = *trial_slots[sl].topo;
t.refinable[pool.first + k] = 1;
t.refinable[index] = 1;
if (!site.ops.empty()) {
t.special_index[pool.first + k] = static_cast<int>(t.special.size());
t.special_index[index] = static_cast<int>(t.special.size());
t.special.push_back(site);
}
});
@@ -2550,9 +2557,9 @@ RefinementResult RefineOnce(const gemmi::Structure &placed, const std::vector<Me
const auto it = by_name.find(e.site + (e.altloc ? std::string("|") + e.altloc : std::string()));
if (it == by_name.end())
continue;
edit_atom(it->second, [&](gemmi::Atom &a, gemmi::Residue &, size_t sl) {
edit_atom(it->second, [&trial_slots, index = it->second](gemmi::Atom &a, gemmi::Residue &, size_t sl) {
a.occ = 0;
trial_slots[sl].topo->refinable[it->second] = 0;
trial_slots[sl].topo->refinable[index] = 0;
});
} else if (e.kind == SolventEdit::Kind::Relabelled) {
const auto it = by_name.find(e.site);
@@ -2570,7 +2577,7 @@ RefinementResult RefineOnce(const gemmi::Structure &placed, const std::vector<Me
for (const InnerSphereDonor &d : InnerSphereDonors(view.models[0], r, as_metal.atoms[0]))
donors.push_back(view.models[0].chains[d.chain].residues[d.residue].atoms[d.atom].serial);
}
edit_atom(i, [&](gemmi::Atom &a, gemmi::Residue &r, size_t sl) {
edit_atom(i, [&trial_slots, e, i, donors](gemmi::Atom &a, gemmi::Residue &r, size_t sl) {
const bool water = e.species == "HOH";
r.name = e.species;
r.entity_type = water ? gemmi::EntityType::Water : gemmi::EntityType::NonPolymer;
@@ -2583,6 +2590,21 @@ RefinementResult RefineOnce(const gemmi::Structure &placed, const std::vector<Me
continue;
kept.push_back(e);
}
std::optional<RefinementStage> edit_stage("build.edit");
std::vector<std::thread> editors;
for (size_t sl = 0; sl < trial_slots.size(); sl++)
editors.emplace_back([&, sl]() {
topo_before[sl] = *trial_slots[sl].topo;
for (const auto &[index, f] : edits) {
gemmi::Residue &r = *residue_of[sl][index];
gemmi::Atom &a = *trial_slots[sl].ctx->atoms[index];
undo[sl].push_back({index, a, r.name, r.entity_type});
f(a, r, sl);
}
});
for (std::thread &t : editors)
t.join();
edit_stage.reset();
refinable.clear();
for (size_t i = 0; i < ctx.atoms.size(); i++)
if (topo.refinable[i])
tests/RefinementTest.cpp
+15
View File
@@ -834,6 +834,21 @@ TEST_CASE("Refinement_restrained_recovers_model", "[Refinement]") {
rms = std::sqrt(rms / ref.size());
CHECK(rms < 0.15);
// one trial at a time instead of three side by side (sharing a card's tables, mask and scale fit): the same
// model, bit for bit
RefinementSettings one = settings;
one.max_trials_at_once = 1;
const RefinementResult s = RefineModel(start, merged, cell, sg, 1.0, logger, one);
REQUIRE(s.ok);
CHECK(s.r_free == r.r_free);
const std::vector<gemmi::Atom *> a_side = atoms_of(r.model->models[0]), a_one = atoms_of(s.model->models[0]);
REQUIRE(a_one.size() == a_side.size());
size_t same = 0;
for (size_t i = 0; i < a_one.size(); i++)
same += a_one[i]->pos.x == a_side[i]->pos.x && a_one[i]->pos.y == a_side[i]->pos.y &&
a_one[i]->pos.z == a_side[i]->pos.z && a_one[i]->b_iso == a_side[i]->b_iso;
CHECK(same == a_one.size());
// cut short before the first weight trial: the model as placed comes back, no trial chosen
const std::atomic<bool> cancel{true};
settings.cancel = &cancel;