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); +}