diff --git a/rugnux/Rugnux.cpp b/rugnux/Rugnux.cpp index a36244dda..df88c0f0a 100644 --- a/rugnux/Rugnux.cpp +++ b/rugnux/Rugnux.cpp @@ -2247,27 +2247,63 @@ ProcessResult Rugnux::RunAllPasses(RugnuxObserver *observer) { const auto mosaicity_before_walk = prepass_mosaicity_; const bool searched_before_walk = beam_center_searched_; beam_center_searched_ = true; - const auto probe = [&](float k, bool index_only) { - experiment_ = before_walk; - experiment_.Goniometer(ScaleRotation(*before_walk.GetGoniometer(), k)); - prepass_rotation_scale_ = k; - prepass_mosaicity_.clear(); - indexing_probe_only_ = index_only; - postrefine_probe_ = postrefine_probe_only_ = !index_only; - ProcessResult r; + // A probe pass at scale k, run on r - this run, or a copy of it (see below). + const auto probe_on = [&](Rugnux &r, RugnuxObserver *obs, float k, bool index_only) { + r.experiment_ = before_walk; + r.experiment_.Goniometer(ScaleRotation(*before_walk.GetGoniometer(), k)); + r.prepass_rotation_scale_ = k; + r.prepass_mosaicity_.clear(); + r.indexing_probe_only_ = index_only; + r.postrefine_probe_ = r.postrefine_probe_only_ = !index_only; + ProcessResult res; try { - r = RunPipeline(observer, /*write_output=*/false, /*geometry_prepass=*/false); + res = r.RunPipeline(obs, /*write_output=*/false, /*geometry_prepass=*/false); } catch (const std::exception &e) { if (IsFatalResourceError(e)) throw; // Angles under which nothing indexes have scored nothing, which is what an empty // result reads as. logger.Info("Rotation scale {:.5f}: the probe pass did not complete ({})", k, e.what()); } - indexing_probe_only_ = postrefine_probe_ = postrefine_probe_only_ = false; + r.indexing_probe_only_ = r.postrefine_probe_ = r.postrefine_probe_only_ = false; + return res; + }; + const auto probe = [&](float k, bool index_only) { ++arm_passes; - return r; + return probe_on(*this, observer, k, index_only); + }; + // The walk always indexes at the stored angles first and at the pass-1 fit next (see + // WalkRotationScale). Both are known now and each is its own pass, so the one at the fit + // runs at the same time, on a copy of the run as it stands before either. What a probe + // leaves behind for the next pass - the spot-finding settings a first-pass ladder rung + // adopts; the experiment is put back by the probe itself - is checked when the fit's turn + // comes: if the probe at the stored angles changed it, the copy started from a state the + // sequential walk would not have, and the probe is run again here. + const float first_fit = static_cast(pass1.post_refine->rotation_scale); + std::shared_ptr fit_copy; + std::future at_fit; + const SpotFindingSettings walk_start_spot_finding = config_.spot_finding; + if (first_fit != 1.0f && !cancelled_) { + fit_copy = std::make_shared(*this); + at_fit = std::async(std::launch::async, [&probe_on, fit_copy, first_fit] { + return probe_on(*fit_copy, nullptr, first_fit, true); + }); + } + const auto index_at = [&](float k) { + if (at_fit.valid() && k == first_fit) { + ProcessResult r = at_fit.get(); + const bool same_start = + config_.spot_finding.high_resolution_limit == walk_start_spot_finding.high_resolution_limit + && config_.spot_finding.measured_ring_q_recipA == walk_start_spot_finding.measured_ring_q_recipA; + if (same_start) { + ++arm_passes; + config_.spot_finding = fit_copy->config_.spot_finding; + for (const auto &m : fit_copy->first_pass_memo_) + AddFirstPassMemo(m); + return r.validation_evidence; + } + } + return probe(k, true).validation_evidence; }; - const auto index_at = [&](float k) { return probe(k, true).validation_evidence; }; const auto refit_at = [&](float k) -> std::optional { const ProcessResult r = probe(k, false); if (!r.post_refine) @@ -2490,7 +2526,9 @@ ProcessResult Rugnux::RunAllPasses(RugnuxObserver *observer) { logger.Info("Rotation two-pass geometry post-refinement: second pass (refined geometry, canonical) -> {}_*", base_prefix); + speculate_geometry_probe_ = true; // a walk that follows starts with an indexing probe auto pass2 = RunPipeline(observer, /*write_output=*/true, /*geometry_prepass=*/false); + speculate_geometry_probe_ = false; // Walk the geometry to the fit's fixed point. The post-refinement is fitted to reflections // integrated at the very geometry it is correcting, so one fit reaches only as far as those @@ -2553,6 +2591,7 @@ ProcessResult Rugnux::RunAllPasses(RugnuxObserver *observer) { const bool searched_before_probe = beam_center_searched_; beam_center_searched_ = true; // a probe scores the geometry it is given, nothing else set_geometry(g); + CollectSpeculativeProbe(); indexing_probe_only_ = true; postrefine_probe_ = false; ProcessResult r; @@ -2622,7 +2661,9 @@ ProcessResult Rugnux::RunAllPasses(RugnuxObserver *observer) { "asks for (round {} of at most {})", fit.distance_before_mm, fit.distance_after_mm, geometry_rounds + 1, MAX_GEOMETRY_ROUNDS); set_geometry(g); + speculate_geometry_probe_ = !large_walk; pass2 = RunPipeline(observer, /*write_output=*/true, /*geometry_prepass=*/false); + speculate_geometry_probe_ = false; ++geometry_rounds; ++walk_passes; const double realised = pass2.post_refine ? pass2.post_refine->held_out_before : NAN; @@ -2966,6 +3007,50 @@ namespace { } } +void Rugnux::AddFirstPassMemo(const FirstPassMemo &memo) { + for (const auto &m : first_pass_memo_) + if (m.key == memo.key) + return; + // A handful covers every probe a walk can ask for again; the oldest go first. + constexpr size_t MAX_FIRST_PASS_MEMOS = 8; + if (first_pass_memo_.size() == MAX_FIRST_PASS_MEMOS) + first_pass_memo_.erase(first_pass_memo_.begin()); + first_pass_memo_.push_back(memo); +} + +void Rugnux::StartSpeculativeGeometryProbe(const std::array &geometry) { + // What the walk's index_at_geometry sets up on the run before it indexes at this geometry. + auto run = std::make_shared(*this); + run->speculative_probe_.reset(); + run->speculate_geometry_probe_ = false; + run->experiment_.BeamX_pxl(geometry[0]).BeamY_pxl(geometry[1]).DetectorDistance_mm(geometry[2]) + .PoniRot1_rad(geometry[3]).PoniRot2_rad(geometry[4]); + run->beam_center_searched_ = true; + run->indexing_probe_only_ = true; + run->postrefine_probe_ = false; + run->postrefine_probe_only_ = false; + auto sp = std::make_shared(); + sp->run = run; + // A probe that fails leaves no memo, and the walk's own probe then runs it again and meets the + // failure itself. + sp->done = std::async(std::launch::async, [run] { + try { + run->RunPipeline(nullptr, /*write_output=*/false, /*geometry_prepass=*/false); + } catch (const std::exception &) { + } + }); + speculative_probe_ = std::move(sp); +} + +void Rugnux::CollectSpeculativeProbe() { + if (!speculative_probe_) + return; + speculative_probe_->done.wait(); + for (const auto &m : speculative_probe_->run->first_pass_memo_) + AddFirstPassMemo(m); + speculative_probe_.reset(); +} + std::vector Rugnux::ExperimentKey() const { std::vector k{experiment_.GetBeamX_pxl(), experiment_.GetBeamY_pxl(), experiment_.GetDetectorDistance_mm(), experiment_.GetPoniRot1_rad(), @@ -3278,8 +3363,10 @@ ProcessResult Rugnux::RunPipeline(RugnuxObserver *observer, bool write_output, b const std::vector first_pass_key = FirstPassInputKey(); const bool verify_memo = std::getenv("RUGNUX_VERIFY_FIRST_PASS_MEMO") != nullptr; std::optional memo_hit; - if (indexing_probe_only_ && first_pass_memo_ && first_pass_memo_->key == first_pass_key) { - memo_hit = first_pass_memo_->evidence; + const auto memo = std::find_if(first_pass_memo_.begin(), first_pass_memo_.end(), + [&](const FirstPassMemo &m) { return m.key == first_pass_key; }); + if (indexing_probe_only_ && memo != first_pass_memo_.end()) { + memo_hit = memo->evidence; if (!verify_memo) { logger.Info("Indexing probe: the first pass has already run on these inputs - {}/{} " "validation spots on the lattice, {} at a wrong spindle angle", @@ -3366,7 +3453,7 @@ ProcessResult Rugnux::RunPipeline(RugnuxObserver *observer, bool write_output, b std::vector engines(std::max(spot_workers, 1)); const auto use_gpu_of_worker = [](size_t t) { if (const int32_t gpus = get_gpu_count(); gpus > 0) - set_gpu(static_cast(t) % gpus); + pin_gpu(static_cast(t) % gpus); }; int prefetch_batches = 0; double prefetch_time_s = 0.0; @@ -5083,9 +5170,7 @@ ProcessResult Rugnux::RunPipeline(RugnuxObserver *observer, bool write_output, b // Kept only where the pass ended on the inputs it started from, and did not take the // beam-centre ladder, which a probe (beam_center_searched_) never would. if (!geometry_prepass && !beam_center_ladder_ran && FirstPassInputKey() == first_pass_key) - first_pass_memo_ = FirstPassMemo{first_pass_key, evidence}; - else - first_pass_memo_.reset(); + AddFirstPassMemo(FirstPassMemo{first_pass_key, evidence}); // A pass run only to score a lattice has its score, whatever it is: a lattice that does not // beat chance is a result for the comparison that asked, not a reason to stop the run. if (indexing_probe_only_) { @@ -6266,6 +6351,9 @@ ProcessResult Rugnux::RunPipeline(RugnuxObserver *observer, bool write_output, b // the pass's memory high-water mark. if (postrefine_probe_ && !geometry_prepass && is_rotation) { post_refine_geometry(); + if (speculate_geometry_probe_ && !cancelled_ && prepass_detector_geometry_ && result.post_refine + && result.post_refine->detector_refined && !result.post_refine->large_move) + StartSpeculativeGeometryProbe(*prepass_detector_geometry_); // A pass the lattice arms run to compare two Bravais classes has now produced the only // thing the comparison reads - the held-out residual of the fit just measured. Stop here // rather than scale, search, fit surfaces and merge for a result nothing looks at. diff --git a/rugnux/Rugnux.h b/rugnux/Rugnux.h index c856fcc3a..15210ec3f 100644 --- a/rugnux/Rugnux.h +++ b/rugnux/Rugnux.h @@ -5,6 +5,8 @@ #include #include +#include +#include #include #include #include @@ -565,7 +567,15 @@ class Rugnux { DiffractionExperiment experiment_; PixelMask pixel_mask_; ProcessConfig config_; - std::atomic cancelled_{false}; + // The run's cancel flag, shared with the copies the run makes of itself to run passes side by + // side (see RunAllPasses), so cancelling the run cancels those too. + class CancelFlag { + std::shared_ptr> flag = std::make_shared>(false); + public: + operator bool() const { return flag->load(); } + CancelFlag &operator=(bool v) { flag->store(v); return *this; } + }; + CancelFlag cancelled_; // Per-frame (by image number) mosaicity fitted by RotationScaleMerge in the two-pass geometry pre-pass, // fed to the second pass's Bragg prediction. Empty outside the rotation two-pass. @@ -715,7 +725,22 @@ class Rugnux { std::vector key; ValidationSpotEvidence evidence; }; - std::optional first_pass_memo_; + std::vector first_pass_memo_; + void AddFirstPassMemo(const FirstPassMemo &memo); + // The geometry walk's next indexing probe, started on a copy of the run as soon as the canonical + // pass has post-refined - it depends on nothing that pass does afterwards - so that it runs beside + // the pass's scaling, merge and output. It leaves its evidence in the copy's first-pass memo, and + // the walk's own probe takes it from there through the usual key check (CollectSpeculativeProbe). + // RunAllPasses sets speculate_geometry_probe_ on the passes whose post-refinement the walk will + // probe with an indexing pass. + struct SpeculativeProbe { + std::shared_ptr run; + std::future done; + }; + std::shared_ptr speculative_probe_; + bool speculate_geometry_probe_ = false; + void StartSpeculativeGeometryProbe(const std::array &geometry); + void CollectSpeculativeProbe(); [[nodiscard]] std::vector FirstPassInputKey() const; // The part of FirstPassInputKey that is the experiment and the spot-finding settings. [[nodiscard]] std::vector ExperimentKey() const;