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

rugnux: fit the canonical pass's post-refinement beside its merge chain

Browse Source 
On a canonical pass the post-refinement only measures: nothing in ScaleMergeAndSymmetry reads
result.post_refine or prepass_detector_geometry_ - RunAllPasses does, after RunPipeline returns. So
the pass now gathers the observations where it did (beside the guard's merge and the engine's host
ingest - the gather reads the outcomes, which the merge chain later re-indexes in place) and runs the
fit on a thread of its own, beside the merges. RunPipeline joins it when ScaleMergeAndSymmetry
returns and replays its log lines and notes there. The speculative geometry probe the fit may start
is started from that thread, on a copy of the run taken on the main thread before the fit, because
the merge chain goes on to change the run's members. The probe-only pass keeps the sequential path.

8a1a, GPU build, two interleaved warm runs each: the window after the image pass loses the fit
(post-refinement 7.6 s -> 2.4 s gather on the main thread), but the quality guard's merge of the
pre-pass frames (7.8 s) is what that window waits for, so nothing moves: tail 30.57 -> 30.59 s,
whole run 55.78 -> 55.79 s; 8qaw 104.08 -> 104.06 s (guard 13.8-14.1 s); myob/cytc within noise.
Peak RSS on 8qaw 24.07 -> 24.35 GB. The gain needs the guard off the critical path as well.
p.mtz byte-identical on myob, cytc, thau, kdp, 8a1a and 8qaw, GPU and CPU builds.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01SVmAWnzCmRKAXVUCdc4iNi
This commit is contained in:
leonarski_f
2026-10-08 06:46:26 +02:00
co-authored by Claude Opus 5.5
parent 11122e4924
commit a8ea35ab97
4 changed files with 109 additions and 54 deletions
Download Patch File Download Diff File Expand all files Collapse all files
rugnux/Rugnux.cpp
+10 -3
View File
@@ -433,9 +433,9 @@ void Rugnux::AddFirstPassMemo(const FirstPassMemo &memo) {
first_pass_memo_.push_back(memo);
}
void Rugnux::StartSpeculativeGeometryProbe(const std::array<float, 5> &geometry) {
void Rugnux::StartSpeculativeGeometryProbe(std::shared_ptr<Rugnux> run, 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->prepass_detector_geometry_ = prepass_detector_geometry_;
run->speculative_probe_.reset();
run->speculate_geometry_probe_ = false;
run->experiment_.BeamX_pxl(geometry[0]).BeamY_pxl(geometry[1]).DetectorDistance_mm(geometry[2])
@@ -886,7 +886,14 @@ ProcessResult Rugnux::RunPipeline(RugnuxObserver *observer, bool write_output, b
if (full && !cancelled_ && (result.indexing_rate.value_or(0.0f) > 0 || result.rotation_lattice_found)
&& (config_.run_scaling || !config_.reference_data.empty())) {
TimingMark mark(logger, "scale, merge and symmetry");
if (ScaleMergeAndSymmetry(p))
const bool finished = ScaleMergeAndSymmetry(p);
if (p.postrefine_ahead.valid()) {
p.postrefine_ahead.get();
p.postrefine_log.ReplayInto(logger);
for (const auto &note : p.postrefine_notes)
Note(note);
}
if (finished)
return result;
}
rugnux/Rugnux.h
+4 -1
View File
@@ -796,7 +796,10 @@ class Rugnux {
};
std::shared_ptr<SpeculativeProbe> speculative_probe_;
bool speculate_geometry_probe_ = false;
void StartSpeculativeGeometryProbe(const std::array<float, 5> &geometry);
// run: the copy of this run the probe runs on, taken by the caller on the thread that owns the
// run - the post-refinement that decides whether to start it runs beside the merge, which does
// not leave the run's members alone.
void StartSpeculativeGeometryProbe(std::shared_ptr<Rugnux> run, 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.
rugnux/RugnuxPipeline.h
+5
View File
@@ -166,4 +166,9 @@ struct PipelineLocals {
EndMessage &end_msg;
std::optional<DiffractionExperiment> &unmerged_experiment;
std::future<gemmi::Mtz> &unmerged_mtz;
// The canonical pass's post-refinement fit, run beside its merge chain (ScaleMergeAndSymmetry)
// and joined by RunPipeline when that returns: its log lines and notes wait in these until then.
Logger postrefine_log = Logger::Buffered();
std::vector<std::string> postrefine_notes;
std::future<void> postrefine_ahead; // last, so on an exception it is waited for first
};
rugnux/RugnuxScaleMerge.cpp
+90 -50
View File
@@ -592,9 +592,16 @@ bool Rugnux::ScaleMergeAndSymmetry(PipelineLocals &p) {
std::optional<RotationScaleMerge> rsm;
std::optional<ScaleMergeResult> guard_merge; // what the quality guard reads, where not the search merge
std::optional<PostRefineObservations> prepass_postrefine_obs;
// The rotation geometry post-refinement (see its call sites below for what it is for).
const auto post_refine_geometry = [&] {
TimingMark mark(logger, "post-refinement");
// The rotation geometry post-refinement (see its call sites below for what it is for), in two
// parts. prepare_post_refine reads the pass - the outcomes, the indexer, the scaling engine - and
// returns the fit with everything it needs; fit_post_refine reads only that, so the canonical pass
// can run it beside its merge chain (below).
struct PostRefineJob {
RotationIndexerResult rot;
PostRefineSettings prs;
PostRefineObservations obs;
};
const auto prepare_post_refine = [&]() -> std::optional<PostRefineJob> {
const auto &outcomes = indexer->GetIntegrationOutcome();
if (geometry_prepass) {
prepass_mosaicity_.assign(outcomes.size(), NAN);
@@ -655,48 +662,61 @@ bool Rugnux::ScaleMergeAndSymmetry(PipelineLocals &p) {
config_.nthreads,
/*release_reflections=*/false, logger);
prepass_postrefine_obs.reset();
const auto pr = PostRefineRotationGeometry(std::move(pr_obs), *rot->axis,
rot->geom, rot->lattice, prs, logger);
result.post_refine = pr;
if (pr.ok && pr.detector_refined) {
// Store, don't apply: this pass continues to write its own (header-geometry) output;
// Run() applies the refined geometry before the second pass.
prepass_detector_geometry_ = std::array<float, 5>{
static_cast<float>(pr.beam_x_after_px),
static_cast<float>(pr.beam_y_after_px),
static_cast<float>(pr.distance_after_mm),
rot->geom.GetPoniRot1_rad(),
rot->geom.GetPoniRot2_rad()};
logger.Info("Two-pass: post-refined detector geometry (distance {:.3f} mm, "
"beam {:.2f},{:.2f} px, tilt {:.4f},{:.4f} deg) - {}",
pr.distance_after_mm, pr.beam_x_after_px, pr.beam_y_after_px,
rot->geom.GetPoniRot1_rad() * 180.0 / PI,
rot->geom.GetPoniRot2_rad() * 180.0 / PI,
geometry_prepass ? "it will drive the second integration pass"
: "what the fit still asks for at the geometry "
"this pass ran at");
Note(fmt::format("{}: {:.2f} mm, beam {:.1f}, {:.1f}",
geometry_prepass ? "geometry refined" : "post-refinement asks",
pr.distance_after_mm, pr.beam_x_after_px, pr.beam_y_after_px));
} else {
logger.Info("Two-pass: geometry post-refine committed no detector change{}",
geometry_prepass ? " - the second pass reproduces the first" : "");
Note("geometry post-refinement: no detector change committed");
}
// The goniometer rotation scale in pr is only a fit: RunAllPasses walks it to its
// fixed point and decides it on the validation frames of the whole sweep
// (WalkRotationScale). It does not ride on pr.ok - the crystals that have a stage
// fault are exactly the ones whose cell and detector steps do NOT pass, because the
// angle error is what their residual is made of.
return PostRefineJob{*rot, prs, std::move(pr_obs)};
}
}
return std::nullopt;
};
const auto fit_post_refine = [this, &result = p.result, geometry_prepass = p.geometry_prepass]
(PostRefineJob job, Logger &logger, std::vector<std::string> &notes) {
const auto &rot = job.rot;
const auto pr = PostRefineRotationGeometry(std::move(job.obs), *rot.axis,
rot.geom, rot.lattice, job.prs, logger);
result.post_refine = pr;
if (pr.ok && pr.detector_refined) {
// Store, don't apply: this pass continues to write its own (header-geometry) output;
// Run() applies the refined geometry before the second pass.
prepass_detector_geometry_ = std::array<float, 5>{
static_cast<float>(pr.beam_x_after_px),
static_cast<float>(pr.beam_y_after_px),
static_cast<float>(pr.distance_after_mm),
rot.geom.GetPoniRot1_rad(),
rot.geom.GetPoniRot2_rad()};
logger.Info("Two-pass: post-refined detector geometry (distance {:.3f} mm, "
"beam {:.2f},{:.2f} px, tilt {:.4f},{:.4f} deg) - {}",
pr.distance_after_mm, pr.beam_x_after_px, pr.beam_y_after_px,
rot.geom.GetPoniRot1_rad() * 180.0 / PI,
rot.geom.GetPoniRot2_rad() * 180.0 / PI,
geometry_prepass ? "it will drive the second integration pass"
: "what the fit still asks for at the geometry "
"this pass ran at");
notes.push_back(fmt::format("{}: {:.2f} mm, beam {:.1f}, {:.1f}",
geometry_prepass ? "geometry refined" : "post-refinement asks",
pr.distance_after_mm, pr.beam_x_after_px, pr.beam_y_after_px));
} else {
logger.Info("Two-pass: geometry post-refine committed no detector change{}",
geometry_prepass ? " - the second pass reproduces the first" : "");
notes.push_back("geometry post-refinement: no detector change committed");
}
// The goniometer rotation scale in pr is only a fit: RunAllPasses walks it to its
// fixed point and decides it on the validation frames of the whole sweep
// (WalkRotationScale). It does not ride on pr.ok - the crystals that have a stage
// fault are exactly the ones whose cell and detector steps do NOT pass, because the
// angle error is what their residual is made of.
};
const auto post_refine_geometry = [&] {
TimingMark mark(logger, "post-refinement");
if (auto job = prepare_post_refine()) {
std::vector<std::string> notes;
fit_post_refine(std::move(*job), logger, notes);
for (const auto &note : notes)
Note(note);
}
};
// On a CANONICAL pass the post-refinement only measures (see below), and it reads nothing the
// scaling engine produces: the reflections' hkl, I, sigma and positions are final once the
// images are integrated, and the engine writes back per-frame fields only. So measure it
// before the engine is built rather than beside it - the gathered observations and their
// rocking events are gigabytes on a long axis, and beside the engine's own arrays they set
// the pass's memory high-water mark.
// images are integrated, and the engine writes back per-frame fields only. So it is gathered
// here, beside the engine's ingest, and fitted beside the merges that follow.
// The quality guard judges a refined pass against the pre-pass on the search merge, and a
// pre-pass that integrated only the first part of the sweep merged only that part. Against
// it a merge of the whole sweep holds more reflections, more signal and more of the axial
@@ -736,22 +756,42 @@ bool Rugnux::ScaleMergeAndSymmetry(PipelineLocals &p) {
config_.nthreads, engine_log, config_.observation_dump_path);
rsm->IngestHost();
});
if (postrefine_probe_ && !geometry_prepass && is_rotation) {
if (postrefine_probe_ && !geometry_prepass && is_rotation && !postrefine_probe_only_) {
// Nothing in the merge chain reads the fit (RunAllPasses does, after RunPipeline has joined
// it), so only the gather runs here: it reads the outcomes, which the chain re-indexes in
// place once the space group is settled. The fit and the speculative probe it may start
// run beside the chain; the probe's copy of the run is taken now, as the chain goes on to
// change the run's members.
TimingMark mark(logger, "post-refinement observations");
if (auto job = prepare_post_refine()) {
std::shared_ptr<Rugnux> probe_run;
if (speculate_geometry_probe_)
probe_run = std::make_shared<Rugnux>(*this);
p.postrefine_ahead = std::async(std::launch::async,
[this, &p, fit_post_refine, job = std::move(*job), probe_run]() mutable {
TimingMark fit_mark(p.postrefine_log, "post-refinement");
fit_post_refine(std::move(job), p.postrefine_log, p.postrefine_notes);
fit_mark.Stop();
if (probe_run && !cancelled_ && prepass_detector_geometry_ && p.result.post_refine
&& p.result.post_refine->detector_refined && !p.result.post_refine->large_move)
StartSpeculativeGeometryProbe(probe_run, *prepass_detector_geometry_);
});
}
}
if (postrefine_probe_ && !geometry_prepass && is_rotation && postrefine_probe_only_) {
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_);
StartSpeculativeGeometryProbe(std::make_shared<Rugnux>(*this), *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.
if (postrefine_probe_only_) {
result.processing_time_s = std::chrono::duration<double>(
std::chrono::steady_clock::now() - start_time).count();
logger.Info("Probe pass: {} images integrated and post-refined in {:.2f} s "
"(the comparison reads the post-refinement only, so this pass does not merge)",
result.images_processed, result.processing_time_s);
return true;
}
result.processing_time_s = std::chrono::duration<double>(
std::chrono::steady_clock::now() - start_time).count();
logger.Info("Probe pass: {} images integrated and post-refined in {:.2f} s "
"(the comparison reads the post-refinement only, so this pass does not merge)",
result.images_processed, result.processing_time_s);
return true;
}
// A pass RunAllPasses re-runs at the fixed radius for a starved background ring: its merge,
// search and files are all replaced by the re-run's, and what anything reads from it - the