mx/Jungfraujoch
3
2 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 93 Wiki Activity

Rugnux: run independent probe passes beside each other, on copies of the run

Browse Source 
The rotation two-pass runs several first-pass indexing probes one after the
other although some of them do not depend on each other:

- The rotation-scale walk always indexes at the stored angles and then at
  the pass-1 fit; both are known before it starts. The probe at the fit now
  runs at the same time, on a copy of the run as it stands before either.
  Its answer is taken only if the probe at the stored angles left behind
  none of the state the next pass reads (the spot-finding settings a
  first-pass ladder rung adopts; each probe puts the experiment back
  itself) - otherwise it is run again in sequence, as before.
- The geometry walk's indexing probe at the canonical pass's post-refined
  geometry depends on nothing that pass does after its post-refinement. It
  is now started there, on a copy, and runs beside the pass's scaling, merge
  and output; its evidence comes back through the first-pass memo (now a
  short list) and the walk's own probe takes it through the usual key check,
  so RUGNUX_VERIFY_FIRST_PASS_MEMO covers it too. It is started only for the
  passes whose post-refinement the walk probes with an indexing pass.

The copies are plain copies of Rugnux: the cancel flag becomes a shared
flag (so cancelling the run cancels them) and everything else was already
copyable. Probe passes run on a copy get no observer.

md5-identical output on two 16M sets; the walk's two probes take ~2.3 s
together instead of ~3 s, and the last probe of the geometry walk is hidden
behind the merge.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01D1G8gJVAy6gp1K5Dz3NE5C
This commit is contained in:
leonarski_fandClaude Opus 5.5 committed 2026-09-26 19:30:01 +02:00
1 parent 59d92a7238
commit 25ea9bf555
2 files changed
+133 -20
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

rugnux/Rugnux.cpp
+106 -18
View File
@@ -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<float>(pass1.post_refine->rotation_scale);
std::shared_ptr<Rugnux> fit_copy;
std::future<ProcessResult> at_fit;
const SpotFindingSettings walk_start_spot_finding = config_.spot_finding;
if (first_fit != 1.0f && !cancelled_) {
fit_copy = std::make_shared<Rugnux>(*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<double> {
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<float, 5> &geometry) {
// What the walk's index_at_geometry sets up on the run before it indexes at this geometry.
auto run = std::make_shared<Rugnux>(*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<SpeculativeProbe>();
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<double> Rugnux::ExperimentKey() const {
std::vector<double> 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<double> first_pass_key = FirstPassInputKey();
const bool verify_memo = std::getenv("RUGNUX_VERIFY_FIRST_PASS_MEMO") != nullptr;
std::optional<ValidationSpotEvidence> 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<SpotEngine> 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<int32_t>(t) % gpus);
pin_gpu(static_cast<int32_t>(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.
rugnux/Rugnux.h
+27 -2
View File
@@ -5,6 +5,8 @@
#include <array>
#include <atomic>
#include <future>
#include <memory>
#include <cmath>
#include <functional>
#include <optional>
@@ -565,7 +567,15 @@ class Rugnux {
DiffractionExperiment experiment_;
PixelMask pixel_mask_;
ProcessConfig config_;
std::atomic<bool> 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<std::atomic<bool>> flag = std::make_shared<std::atomic<bool>>(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<double> key;
ValidationSpotEvidence evidence;
};
std::optional<FirstPassMemo> first_pass_memo_;
std::vector<FirstPassMemo> 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<Rugnux> run;
std::future<void> done;
};
std::shared_ptr<SpeculativeProbe> speculative_probe_;
bool speculate_geometry_probe_ = false;
void StartSpeculativeGeometryProbe(const std::array<float, 5> &geometry);
void CollectSpeculativeProbe();
[[nodiscard]] std::vector<double> FirstPassInputKey() const;
// The part of FirstPassInputKey that is the experiment and the spot-finding settings.
[[nodiscard]] std::vector<double> ExperimentKey() const;