diff --git a/image_analysis/structure_refinement/ModelScaling.cpp b/image_analysis/structure_refinement/ModelScaling.cpp index 0b7688974..6917699ae 100644 --- a/image_analysis/structure_refinement/ModelScaling.cpp +++ b/image_analysis/structure_refinement/ModelScaling.cpp @@ -523,15 +523,17 @@ ModelScaleReport ModelScaleGPUPool::Fit(gemmi::Scaling &scaling, double b std::vector stol2(n); std::vector fobs(n), sigma(n); std::vector 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(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> constraints; for (const gemmi::Vec6 &c : scaling.constraint_matrix) constraints.push_back(c); diff --git a/image_analysis/structure_refinement/Refinement.cpp b/image_analysis/structure_refinement/Refinement.cpp index c6d9a3d0c..1aaef71e6 100644 --- a/image_analysis/structure_refinement/Refinement.cpp +++ b/image_analysis/structure_refinement/Refinement.cpp @@ -391,16 +391,25 @@ struct Context { } std::optional s3("refit.scale"); + // the per-reflection loops on the threads, each element on its own + const int n_refl = static_cast(refl.Size()); + const size_t refl_threads = ThreadsForWork(refl.Size(), nthreads); gemmi::AsuData> fc, fm; gemmi::AsuData> 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 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(scaling.get_overall_scale_factor(refl.hkl[m])); - rows.fmask[m] = static_cast(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(scaling.get_overall_scale_factor(refl.hkl[m])); + rows.fmask[m] = static_cast(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(scaling.get_solvent_scale(cell.calculate_stol_sq(h))) * fmask[m]; @@ -2511,18 +2523,12 @@ RefinementResult RefineOnce(const gemmi::Structure &placed, const std::vector> undo(trial_slots.size()); - std::vector topo_before; - for (const TrialSlot &slot : trial_slots) - topo_before.push_back(*slot.topo); + std::vector topo_before(trial_slots.size()); const SolventPool pool_before = pool; - auto edit_atom = [&](int index, const std::function &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; + std::vector> edits; + auto edit_atom = [&](int index, Edit f) { edits.emplace_back(index, std::move(f)); }; std::vector 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(e.occupancy); a.b_iso = static_cast(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(t.special.size()); + t.special_index[index] = static_cast(t.special.size()); t.special.push_back(site); } }); @@ -2550,9 +2557,9 @@ RefinementResult RefineOnce(const gemmi::Structure &placed, const std::vectorsecond, [&](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 edit_stage("build.edit"); + std::vector 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]) diff --git a/tests/RefinementTest.cpp b/tests/RefinementTest.cpp index 6b7932280..9412d212d 100644 --- a/tests/RefinementTest.cpp +++ b/tests/RefinementTest.cpp @@ -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 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 cancel{true}; settings.cancel = &cancel;