From 5d931222d700cd198e19081888c48212f9a0491f Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Sat, 19 Sep 2026 12:58:03 +0200 Subject: [PATCH] rugnux --model: evaluate the rigid-body Jacobian's columns in parallel The rigid-body target is an Fcalc + solvent mask + scale re-fit per evaluation, and its forward-difference Jacobian made six of them one after the other after the central one - 7 of the ~8 evaluations per LM iteration. The six shifted evaluations now run in parallel, column j on its own Evaluator over its own copy of the model (an evaluation moves every atom, so two cannot share one), set to the same zone and given the central evaluation's bulk-solvent pair - the pair the serial loop's shifted evaluations used, since the zone's solvent is fitted on the zone's first evaluation and every Evaluate starts with the central one. Each column's arithmetic is the serial one and the Jacobian is assembled in column order, so the refinement is bit-identical; the evaluation count is kept as the serial loop kept it (up to and including a failing column). The per-zone solvent fit (FitModelScale) gets the thread count as well. Inside a null replicate (a pool worker) the columns run inline, as before. To check: RIGID_BODY_* and MODEL_* keys and md5 of maps/.mtz/_model.cif identical with and without this commit on the audit set; rigid-body seconds on a large model. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_013nW6FNRP1bBJJ8pfHiByAT --- rugnux/ModelValidation.cpp | 2 +- rugnux/RigidBodyRefine.cpp | 63 ++++++++++++++++++++++++++++------- rugnux/RigidBodyRefine.h | 4 ++- tests/ModelValidationTest.cpp | 40 ++++++++++++++++++++++ 4 files changed, 95 insertions(+), 14 deletions(-) diff --git a/rugnux/ModelValidation.cpp b/rugnux/ModelValidation.cpp index 06b3e4790..4d09805ab 100644 --- a/rugnux/ModelValidation.cpp +++ b/rugnux/ModelValidation.cpp @@ -610,7 +610,7 @@ ModelValidationResult ValidateAgainstModel(const std::vector & // exactly where it was read. const std::vector before = ModelPositions(ms.st.models[0]); const RigidBodyRefineResult rb = - RefineRigidBody(ms.st.models[0], ucell, *sg, out.fit.fobs_work, d_min, logger); + RefineRigidBody(ms.st.models[0], ucell, *sg, out.fit.fobs_work, d_min, logger, nthreads); if (!rb.converged) { // RefineRigidBody leaves the model wherever the solver left it, usable answer or not, so // the restore cannot be conditional on the same flag the re-fit is. diff --git a/rugnux/RigidBodyRefine.cpp b/rugnux/RigidBodyRefine.cpp index 91d25843c..8b36f44d5 100644 --- a/rugnux/RigidBodyRefine.cpp +++ b/rugnux/RigidBodyRefine.cpp @@ -4,6 +4,7 @@ #include "RigidBodyRefine.h" #include +#include #include #include #include @@ -21,6 +22,7 @@ #include "ModelScaling.h" // FitModelScale #include "../common/JFJochMath.h" // PI #include "../common/Logger.h" +#include "../common/ParallelFor.h" namespace { @@ -72,8 +74,8 @@ struct Placement { class Evaluator { public: Evaluator(gemmi::Model &model, const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg, - const std::vector &base, const Placement &placement) - : model_(model), cell_(cell), sg_(sg), base_(base), placement_(placement) {} + const std::vector &base, const Placement &placement, size_t nthreads) + : model_(model), cell_(cell), sg_(sg), base_(base), placement_(placement), nthreads_(nthreads) {} // The zone's observations, and the scale the residuals are expressed in. void SetZone(const gemmi::AsuData> &fobs, double d_min) { @@ -86,6 +88,14 @@ public: solvent_fitted_ = false; } + // Hold the zone's bulk solvent at a pair another evaluator already fitted, as if this one had + // fitted it itself (the Jacobian's column evaluators take the central evaluation's). + void UseSolvent(double k, double b) { + k_sol = k; + b_sol = b; + solvent_fitted_ = true; + } + size_t NumObservations() const { return fobs_.v.size(); } double JacobianStep() const { return JACOBIAN_STEP_FRACTION * d_min_; } int evaluations = 0; @@ -135,7 +145,7 @@ public: if (scaling.points.empty()) return false; if (!solvent_fitted_) { - FitModelScale(scaling); + FitModelScale(scaling, {}, nthreads_); k_sol = scaling.k_sol; b_sol = scaling.b_sol; solvent_fitted_ = true; @@ -179,6 +189,7 @@ private: double d_min_ = 0; double f_mean_ = 1; bool solvent_fitted_ = false; + size_t nthreads_ = 1; }; // Ceres' own numeric differentiation steps by |x| * relative_step_size, which is zero at the start of @@ -186,9 +197,17 @@ private: // forward differences at a step chosen in the parameters' units. Analytic dF/dp would need // derivatives GEMMI's structure-factor path does not have, and at six parameters it is not worth it: // a Jacobian costs seven evaluations, and the evaluations at 6-3.5 A are cheap. +// +// The six shifted evaluations are independent of each other, so they run in parallel, column j on +// `columns[j]`: an evaluator over its own copy of the model (an evaluation moves every atom to the +// placement it is asked about, so two cannot share one), set to the same zone and handed the solvent +// pair the central evaluation fitted - which is the pair the serial loop's shifted evaluations used, +// since the zone's solvent is fitted by the first evaluation of the zone and every Evaluate starts +// with the central one. Each column's arithmetic is the serial one, so the Jacobian is too. class RigidBodyCost : public ceres::CostFunction { public: - explicit RigidBodyCost(Evaluator &ev) : ev_(ev) { + RigidBodyCost(Evaluator &ev, std::vector &columns, size_t nthreads) + : ev_(ev), columns_(columns), nthreads_(nthreads) { set_num_residuals(static_cast(ev.NumObservations())); mutable_parameter_block_sizes()->push_back(6); } @@ -198,16 +217,24 @@ public: if (!ev_.Residuals(parameters[0], residuals)) return false; if (jacobians != nullptr && jacobians[0] != nullptr) { - std::vector shifted(n); - for (int j = 0; j < 6; j++) { + const double step = ev_.JacobianStep(); + std::vector> shifted(6, std::vector(n)); + std::array ok{}; + ParallelFor(6, nthreads_, [&](int j) { double q[6]; std::copy(parameters[0], parameters[0] + 6, q); - const double step = ev_.JacobianStep(); q[j] += step; - if (!ev_.Residuals(q, shifted.data())) + columns_[j].UseSolvent(ev_.k_sol, ev_.b_sol); + ok[j] = columns_[j].Residuals(q, shifted[j].data()); + }); + // Counted as the serial loop counted them: it stopped at the first column that failed. + for (int j = 0; j < 6; j++) { + ++ev_.evaluations; + ev_.unmatched = columns_[j].unmatched; + if (!ok[j]) return false; for (int i = 0; i < n; i++) - jacobians[0][i * 6 + j] = (shifted[i] - residuals[i]) / step; + jacobians[0][i * 6 + j] = (shifted[j][i] - residuals[i]) / step; } } return true; @@ -215,6 +242,8 @@ public: private: Evaluator &ev_; + std::vector &columns_; + size_t nthreads_; }; // The directions in which this space group's origin is free. Translating the whole cell content along @@ -274,7 +303,8 @@ RigidBodyRefineResult RefineRigidBody(gemmi::Model &model, const gemmi::SpaceGroup &sg, const gemmi::AsuData> &fobs, double d_min, - Logger &logger) { + Logger &logger, + size_t nthreads) { const auto t0 = std::chrono::steady_clock::now(); RigidBodyRefineResult result; @@ -301,7 +331,14 @@ RigidBodyRefineResult RefineRigidBody(gemmi::Model &model, ladder.push_back(d_min); const gemmi::Mat33 gauge = GaugeProjector(sg, cell); - Evaluator ev(model, cell, sg, base, placement); + Evaluator ev(model, cell, sg, base, placement, nthreads); + // The Jacobian's six column evaluators, each over its own copy of the model (see RigidBodyCost). + // The copies only ever hold probe placements; `model` is left at the answer below. + std::vector column_models(6, model); + std::vector columns; + columns.reserve(6); + for (int j = 0; j < 6; j++) + columns.emplace_back(column_models[j], cell, sg, base, placement, 1); double q[6] = {0, 0, 0, 0, 0, 0}; bool any_zone_solved = false; for (double zone : ladder) { @@ -315,9 +352,11 @@ RigidBodyRefineResult RefineRigidBody(gemmi::Model &model, continue; result.zones.push_back(zone); // the ladder WALKED, which a thin zone drops out of ev.SetZone(zone_obs, zone); + for (Evaluator &c : columns) + c.SetZone(zone_obs, zone); ceres::Problem problem; - problem.AddResidualBlock(new RigidBodyCost(ev), nullptr, q); + problem.AddResidualBlock(new RigidBodyCost(ev, columns, nthreads), nullptr, q); ceres::Solver::Options options; options.linear_solver_type = ceres::DENSE_QR; options.max_num_iterations = 15; diff --git a/rugnux/RigidBodyRefine.h b/rugnux/RigidBodyRefine.h index 2f57fdc5d..a17f3489b 100644 --- a/rugnux/RigidBodyRefine.h +++ b/rugnux/RigidBodyRefine.h @@ -3,6 +3,7 @@ #pragma once +#include #include #include "gemmi/asudata.hpp" @@ -41,4 +42,5 @@ RigidBodyRefineResult RefineRigidBody(gemmi::Model &model, const gemmi::SpaceGroup &sg, const gemmi::AsuData> &fobs, double d_min, - Logger &logger); + Logger &logger, + size_t nthreads = 1); diff --git a/tests/ModelValidationTest.cpp b/tests/ModelValidationTest.cpp index 943a559cc..3fcb7bee8 100644 --- a/tests/ModelValidationTest.cpp +++ b/tests/ModelValidationTest.cpp @@ -622,3 +622,43 @@ TEST_CASE("ModelValidation_MapToFPhiMatchesGemmi", "[ModelValidation]") { for (size_t i = 0; i < a.v.size(); i++) CHECK(a.v[i].hkl == b.v[i].hkl); } + +// The Jacobian's six columns are evaluated in parallel; the placement must be the serial one, bit for +// bit. +TEST_CASE("ModelValidation_RigidBodySameOnAnyNumberOfThreads", "[ModelValidation]") { + Logger logger("ModelValidation_RigidBodySameOnAnyNumberOfThreads"); + + const auto path = WriteTemp("rigid_body_threads_test.pdb", ClusterPdb().c_str()); + gemmi::Structure st = gemmi::read_structure_gz(path, gemmi::CoorFormat::Detect); + const gemmi::SpaceGroup *sg = st.find_spacegroup(); + REQUIRE(sg != nullptr); + st.setup_cell_images(); + + const auto ref = ModelReferenceIntensities(path, {}, {}, 3.0, logger); + REQUIRE_FALSE(ref.empty()); + gemmi::AsuData> fobs; + fobs.unit_cell_ = st.cell; + fobs.spacegroup_ = sg; + for (const auto &r : ref) + fobs.v.push_back({{{r.h, r.k, r.l}}, {std::sqrt(r.I), 1.0f}}); + fobs.ensure_sorted(); + + std::vector displaced; + for (const gemmi::Position &p : ModelPositions(st.models[0])) + displaced.emplace_back(p.x + 0.40, p.y - 0.30, p.z + 0.20); + + gemmi::Model serial = st.models[0], parallel = st.models[0]; + SetModelPositions(serial, displaced); + SetModelPositions(parallel, displaced); + const RigidBodyRefineResult r1 = RefineRigidBody(serial, st.cell, *sg, fobs, 3.0, logger, 1); + const RigidBodyRefineResult r6 = RefineRigidBody(parallel, st.cell, *sg, fobs, 3.0, logger, 6); + CHECK(r1.evaluations == r6.evaluations); + CHECK(r1.angle_deg == r6.angle_deg); + CHECK(r1.shift_A == r6.shift_A); + const auto p1 = ModelPositions(serial), p6 = ModelPositions(parallel); + REQUIRE(p1.size() == p6.size()); + for (size_t i = 0; i < p1.size(); i++) + CHECK((p1[i].x == p6[i].x && p1[i].y == p6[i].y && p1[i].z == p6[i].z)); + + std::filesystem::remove(path); +}