From b5980c9ba0e781122b6c41b2775d98fe03427a64 Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Fri, 9 Oct 2026 12:10:46 +0200 Subject: [PATCH] rugnux: start the validation in the model's setting on a forecast, beside the first Where a model that fits is written on other axes than the data, the files take its setting and the validation is made a second time on the relabelled data. That second validation depends on the first only through the setting and the indexing the first settles, and both are known long before the first has finished: the change of basis right after the frame scoring, the indexing the probe prefers before the null. So ValidateAgainstModel now reports them (ModelFrameForecast, through ModelValidationSchedule:: on_forecast), and the second validation is started there, on the relabelling AdoptModelFrame and relabel_output would make, applied to a copy of the merge - beside the first one's null, real fit and maps. It is kept only where the first decides exactly what was forecast (the model fits, same change of basis, same indexing; a model asserting the other enantiomorph is not forecast, as the label is decided last, on the anomalous map). Until then its log is held (Logger::Buffered, replayed where the serial run logged it) and its files wait on a gate (ModelValidationSchedule::write_gate) placed before the first map is written; otherwise it is released with false and returns unwritten, and the serial validation runs as before. GPU memory: two validations at once take twice the rigid-body engines. The parallel start is decided up front from sizes, never from what is free: allowed where twice what the first pool asked for (bytes per engine times the engines wanted) fits a quarter of the card's TOTAL memory, the share one validation may take. Threads: both validations submit to the one ParallelFor pool from threads outside it (the second runs on a std::async thread, as the null's replicates do), so no pass runs inline on a pool worker and the pool's size bounds the workers. Measured on the loaded 16-core workstation, TIMING model validation: 8sa8 30.6 -> 19.7 s, 8xtg 19.9 -> 13.9 s, 9ea5 22.1 -> 16.4 s (this and the previous commit together). p.mtz, p.hkl, p.cif, the three maps, p_maps.mtz and p_model.cif/pdb md5-identical to the base on 8sa8, 8xtg, 9ea5, 7qis and myob_x10sa; the validation's log lines identical as a set (the null's replicate lines were already in completion order). Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01SVmAWnzCmRKAXVUCdc4iNi --- .../structure_refinement/ModelValidation.cpp | 24 ++++- .../structure_refinement/ModelValidation.h | 25 ++++- .../structure_refinement/RigidBodyGPU.cpp | 2 + .../structure_refinement/RigidBodyGPU.h | 6 ++ rugnux/RugnuxScaleMerge.cpp | 92 +++++++++++++++---- 5 files changed, 128 insertions(+), 21 deletions(-) diff --git a/image_analysis/structure_refinement/ModelValidation.cpp b/image_analysis/structure_refinement/ModelValidation.cpp index 3bbf9af2d..123ca02b0 100644 --- a/image_analysis/structure_refinement/ModelValidation.cpp +++ b/image_analysis/structure_refinement/ModelValidation.cpp @@ -353,6 +353,7 @@ ModelValidationResult Validate(const std::vector &merged, size_t nthreads, double wavelength_A, const std::vector &report_shell_d_min, + const ModelValidationSchedule &schedule, bool rigid_body_gpu) { ModelValidationResult result; result.model_path = model_path; @@ -807,6 +808,15 @@ ModelValidationResult Validate(const std::vector &merged, const bool decision_pending = result.model_enantiomorph_candidate || !(indexing.op == gemmi::Op::identity()) || !(result.change_of_basis_op == gemmi::Op::identity()); + if (schedule.on_forecast) { + // Two validations at once take twice the engines, and are allowed where twice what this one + // asked for fits the share of the card one validation may take (RigidBodyGPUPool::Create). + bool may_run_beside = true; + if (rigid_body_pool != nullptr) + may_run_beside = 2 * rigid_body_pool->PlannedBytes() <= rigid_body_pool->CardBytes() / 4; + schedule.on_forecast({result.change_of_basis_op, indexing.op, result.model_enantiomorph_candidate, + may_run_beside}); + } std::vector equivalent; // the orientations the crystal cannot tell from the model's double reach_deg = 0; std::vector null_rotation; @@ -1273,6 +1283,13 @@ ModelValidationResult Validate(const std::vector &merged, mapfofc.v.push_back({terms[i].hkl, delfwt[i] * ph}); } + // A validation started before the one it follows had decided writes nothing until that one has, + // and nothing at all where it decided otherwise. + if (schedule.write_gate.valid() && !schedule.write_gate.get()) { + result.failure_reason = "superseded before its files were written"; + return result; + } + // --- write the maps and score the 2mFo-DFc map at atom centres (a real map peaks there) --- // The three maps - this one, the difference map and the anomalous map - share nothing but what // they read, so the other two are made and written on threads of their own beside this one. @@ -1500,14 +1517,15 @@ ModelValidationResult ValidateAgainstModel(const std::vector & bool probe_indexing_ambiguity, size_t nthreads, double wavelength_A, - const std::vector &report_shell_d_min) { + const std::vector &report_shell_d_min, + const ModelValidationSchedule &schedule) { #ifdef JFJOCH_USE_CUDA // A CUDA failure in the rigid body does not end the run: the validation is started again from the // model as read, on the CPU throughout. It is re-runnable, and a dead model check should not take a // finished merge with it. try { return Validate(merged, cell, model_path, output_prefix, logger, data_space_group, - probe_indexing_ambiguity, nthreads, wavelength_A, report_shell_d_min, true); + probe_indexing_ambiguity, nthreads, wavelength_A, report_shell_d_min, schedule, true); } catch (const RigidBodyGPUFailure &e) { cuda_clear_error(); logger.Warning("Model validation: the rigid body failed on the GPU ({}); validating again on the CPU", @@ -1515,7 +1533,7 @@ ModelValidationResult ValidateAgainstModel(const std::vector & } #endif return Validate(merged, cell, model_path, output_prefix, logger, data_space_group, probe_indexing_ambiguity, - nthreads, wavelength_A, report_shell_d_min, false); + nthreads, wavelength_A, report_shell_d_min, schedule, false); } namespace { diff --git a/image_analysis/structure_refinement/ModelValidation.h b/image_analysis/structure_refinement/ModelValidation.h index bcfcad445..a2f24f3cc 100644 --- a/image_analysis/structure_refinement/ModelValidation.h +++ b/image_analysis/structure_refinement/ModelValidation.h @@ -4,6 +4,8 @@ #pragma once #include +#include +#include #include #include #include @@ -226,6 +228,26 @@ struct ModelValidationResult { // cc_model_shells is binned on. Sharing the grid is the point: a reader has to be able to put a // CC(model, data) row beside that shell's CC1/2 and know the two describe the same reflections. // Empty (the default) means no shells were given and none are reported. +// What a validation knows of its outcome before it has one: the setting the model was put into and +// the indexing the probe prefers - which is what it decides where the model turns out to fit and the +// probe's lead turns out to be real. Enough for a caller to start what follows on that forecast and +// keep it only where the outcome agrees. +struct ModelFrameForecast { + gemmi::Op change_of_basis_op = gemmi::Op::identity(); + gemmi::Op indexing_op = gemmi::Op::identity(); + bool enantiomorph_candidate = false; + // Whether a second validation of this size may run beside this one: decided on the sizes of its + // GPU engines against the memory of the cards, never on what happens to be free. + bool may_run_beside = true; +}; + +struct ModelValidationSchedule { + // Called once, before the null, with the forecast above. + std::function on_forecast; + // Where valid, waited on before the first file is written: false returns without writing any. + std::shared_future write_gate; +}; + ModelValidationResult ValidateAgainstModel(const std::vector &merged, const UnitCell &cell, const std::string &model_path, @@ -235,7 +257,8 @@ ModelValidationResult ValidateAgainstModel(const std::vector & bool probe_indexing_ambiguity = true, size_t nthreads = 1, double wavelength_A = 0.0, - const std::vector &report_shell_d_min = {}); + const std::vector &report_shell_d_min = {}, + const ModelValidationSchedule &schedule = {}); // Reindex `merged` into the frame ValidateAgainstModel reported, so the reflection files that are // written describe the same indexing as the R-factors and the maps. Returns the space group they are diff --git a/image_analysis/structure_refinement/RigidBodyGPU.cpp b/image_analysis/structure_refinement/RigidBodyGPU.cpp index d85f696fc..2c74506f6 100644 --- a/image_analysis/structure_refinement/RigidBodyGPU.cpp +++ b/image_analysis/structure_refinement/RigidBodyGPU.cpp @@ -210,6 +210,8 @@ std::unique_ptr RigidBodyGPUPool::Create(const gemmi::Model &m const size_t fit = bytes > 0 ? budget / bytes : 0; const size_t want = std::min(fit, std::max(max_engines, 1)); const int device = RigidBodyGPUEngine::CurrentDevice(); + pool->planned_bytes_ = bytes * std::max(max_engines, 1); + pool->card_bytes_ = total; for (size_t i = 0; i < want; i++) pool->engines_.push_back(std::make_unique(cap, device)); if (pool->engines_.empty()) { diff --git a/image_analysis/structure_refinement/RigidBodyGPU.h b/image_analysis/structure_refinement/RigidBodyGPU.h index 7f2651fb5..6b0db0ee1 100644 --- a/image_analysis/structure_refinement/RigidBodyGPU.h +++ b/image_analysis/structure_refinement/RigidBodyGPU.h @@ -49,6 +49,10 @@ public: ~RigidBodyGPUPool(); size_t Engines() const { return engines_.size(); } + // What max_engines engines take, however many fitted, and the memory of the card: the sizes a + // caller decides by whether a second validation may run beside this one. + size_t PlannedBytes() const { return planned_bytes_; } + size_t CardBytes() const { return card_bytes_; } RigidBodyGPUEngine &Acquire(); void Release(RigidBodyGPUEngine &engine); @@ -64,6 +68,8 @@ public: private: RigidBodyGPUPool() = default; + size_t planned_bytes_ = 0; + size_t card_bytes_ = 0; std::vector> engines_; std::vector idle_; std::mutex m_; diff --git a/rugnux/RugnuxScaleMerge.cpp b/rugnux/RugnuxScaleMerge.cpp index 3197f3765..33793e685 100644 --- a/rugnux/RugnuxScaleMerge.cpp +++ b/rugnux/RugnuxScaleMerge.cpp @@ -3247,19 +3247,62 @@ bool Rugnux::ScaleMergeAndSymmetry(PipelineLocals &p) { std::vector report_shell_d_min; for (const auto &sh : sm.statistics.shells) report_shell_d_min.push_back(sh.d_min); - const auto validate = [&, cell = *result.consensus_cell, - wavelength = experiment_.GetWavelength_A()](Logger &log) { - return ValidateAgainstModel(sm.merged, cell, config_.model_path, - config_.output_prefix, log, - data_sg ? &*data_sg : nullptr, - /*probe_indexing_ambiguity=*/config_.reference_data.empty(), - static_cast(config_.nthreads), - wavelength, report_shell_d_min); + // The validation in the model's setting further down depends on this one only through the + // setting and the indexing it settles, and both are forecast well before this one has finished + // (ModelFrameForecast). So it is started on that forecast, beside this one, and kept only where the decision is the one forecast; its log is held, + // and its files wait, until then. On any other decision it is dropped, unwritten. The relabelling + // it runs on is the one AdoptModelFrame and relabel_output below make, made here on a copy. + // A model asserting the other enantiomorph is not forecast: whether the label is taken is + // only known at the end, from the anomalous map. + struct Ahead { + gemmi::Op change_of_basis_op, indexing_op; + Logger log = Logger::Buffered(); + std::future run; + std::promise write; // destroyed before `run`, so a run left waiting on it is released + } ahead; + ModelValidationSchedule first_schedule; + first_schedule.on_forecast = [&](const ModelFrameForecast &f) { + if (ahead.run.valid() || !config_.reference_data.empty() || !data_sg.has_value() + || !f.may_run_beside || f.enantiomorph_candidate + || f.change_of_basis_op == gemmi::Op::identity()) + return; + ahead.change_of_basis_op = f.change_of_basis_op; + ahead.indexing_op = f.indexing_op; + ahead.run = std::async(std::launch::async, + [&, merged = sm.merged, sg = experiment_.GetSpaceGroupOrP1(), cell = *result.consensus_cell, + friedel = experiment_.GetScalingSettings().GetMergeFriedel(), + rfree_fraction = experiment_.GetScalingSettings().GetRfreeFraction(), + wavelength = experiment_.GetWavelength_A(), + gate = ahead.write.get_future().share()]() mutable { + if (!(ahead.indexing_op == gemmi::Op::identity())) + merged = ReindexMergedIntoAsu(merged, ahead.indexing_op, sg, friedel); + gemmi::Op to_model = ahead.change_of_basis_op; + to_model.tran = {0, 0, 0}; + const gemmi::Op cob = to_model.inverse(); + const gemmi::SpaceGroup *in = &sg; + if (const gemmi::SpaceGroup *moved = SpaceGroupInBasis(sg, cob)) { + merged = ReindexMergedIntoAsu(merged, HklOperator(cob), *moved, friedel); + cell = CellInBasis(cell, cob); + AssignRfreeFlags(merged, *moved, rfree_fraction, 500, cell, config_.nthreads); + in = moved; + } + ModelValidationSchedule schedule; + schedule.write_gate = gate; + return ValidateAgainstModel(merged, cell, config_.model_path, config_.output_prefix, ahead.log, + in, /*probe_indexing_ambiguity=*/false, + static_cast(config_.nthreads), wavelength, + report_shell_d_min, schedule); + }); }; // The ledger's P1 merge (started above) is taken before anything below acts on what the // validation decides: relabelling the outcomes and the group is what it must not see happen. // Every such relabelling reaches the P1 merge afterwards, through merge_to_written. - ModelValidationResult validation = validate(logger); + ModelValidationResult validation = + ValidateAgainstModel(sm.merged, *result.consensus_cell, config_.model_path, + config_.output_prefix, logger, data_sg ? &*data_sg : nullptr, + /*probe_indexing_ambiguity=*/config_.reference_data.empty(), + static_cast(config_.nthreads), experiment_.GetWavelength_A(), + report_shell_d_min, first_schedule); take_ledger_p1(); // A model that was asked for and could not be used has to say so where anyone will see // it. Without this the run ends successfully with no R-free, no maps and nothing in the @@ -3315,18 +3358,33 @@ bool Rugnux::ScaleMergeAndSymmetry(PipelineLocals &p) { // validation is then made again on the relabelled data, so the maps and the placed model // come out on the same axes as the reflections; what it had decided is kept, since the // second one, facing a model already in the data's setting, has nothing left to decide. - if (config_.reference_data.empty() && validation.model_fits - && !(validation.change_of_basis_op == gemmi::Op::identity())) { + const bool remake = config_.reference_data.empty() && validation.model_fits + && !(validation.change_of_basis_op == gemmi::Op::identity()); + const bool forecast_held = remake && ahead.run.valid() && !validation.adopted_model_enantiomorph + && validation.change_of_basis_op == ahead.change_of_basis_op + && validation.indexing_op == ahead.indexing_op; + if (ahead.run.valid() && !forecast_held) { + ahead.write.set_value(false); + ahead.run.wait(); + } + if (remake) { gemmi::Op to_model = validation.change_of_basis_op; to_model.tran = {0, 0, 0}; relabel_output(to_model.inverse(), "MODEL"); const auto sg_now = experiment_.GetGemmiSpaceGroup(); - auto remade = ValidateAgainstModel(sm.merged, *result.consensus_cell, config_.model_path, - config_.output_prefix, logger, - sg_now ? &*sg_now : nullptr, - /*probe_indexing_ambiguity=*/false, - static_cast(config_.nthreads), - experiment_.GetWavelength_A(), report_shell_d_min); + ModelValidationResult remade; + if (forecast_held) { + ahead.write.set_value(true); + remade = ahead.run.get(); + ahead.log.ReplayInto(logger); + } else { + remade = ValidateAgainstModel(sm.merged, *result.consensus_cell, config_.model_path, + config_.output_prefix, logger, + sg_now ? &*sg_now : nullptr, + /*probe_indexing_ambiguity=*/false, + static_cast(config_.nthreads), + experiment_.GetWavelength_A(), report_shell_d_min); + } if (remade.ok) { KeepModelVerdict(remade, validation); validation = std::move(remade);