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

rugnux: index at the measured beam centre beside the file centre's first pass, on a second GPU

Browse Source 
The beam-centre check runs a second first pass at the centre measured from
the background, and reads the first pass's answer only once that pass is
over. With two GPUs or more, its spot finding (on the second card), scheme
feeding and FFT indexing now start beside the file centre's pass, on a
copy of the experiment at the measured centre with workers, a mapping and a
spot list of their own (SpotFindingWorkers). The check takes them only
where the experiment (rotation angles included), SpotFindingKey at the
measured centre and the spot budget are still what the run has there, and
scores the schemes on the run's own indexer as before; otherwise it runs
that pass itself as before. The spots at the two centres are not shared:
the spot list depends on the centre (ring background, resolution, budget).

With one GPU, and in the CPU build, nothing changes: on one card the two
passes share the card and the reader, and side by side they saved nothing.

FindFirstPassSpots is the first pass's per-frame spot finding, now shared;
ExperimentKey and SpotFindingKey take an experiment other than the run's.

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 16:50:35 +02:00
co-authored by Claude Opus 5.5
parent 72c7409327
commit de911fe126
3 changed files with 240 additions and 37 deletions
Download Patch File Download Diff File Expand all files Collapse all files
rugnux/Rugnux.cpp
+14 -13
View File
@@ -483,12 +483,12 @@ void Rugnux::CollectSpeculativeProbe() {
speculative_probe_.reset();
}
std::vector<double> Rugnux::ExperimentKey(bool rotation_angles) const {
std::vector<double> k{experiment_.GetBeamX_pxl(), experiment_.GetBeamY_pxl(),
experiment_.GetDetectorDistance_mm(), experiment_.GetPoniRot1_rad(),
experiment_.GetPoniRot2_rad(), experiment_.GetWavelength_A(),
static_cast<double>(experiment_.GetMaxSpotCount())};
if (const auto g = experiment_.GetGoniometer()) {
std::vector<double> Rugnux::ExperimentKey(const DiffractionExperiment &e, bool rotation_angles) const {
std::vector<double> k{e.GetBeamX_pxl(), e.GetBeamY_pxl(),
e.GetDetectorDistance_mm(), e.GetPoniRot1_rad(),
e.GetPoniRot2_rad(), e.GetWavelength_A(),
static_cast<double>(e.GetMaxSpotCount())};
if (const auto g = e.GetGoniometer()) {
const Coord a = g->GetAxis();
const auto helical = g->GetHelicalStep().value_or(Coord(-1, -1, -1));
k.insert(k.end(), {1.0, a.x, a.y, a.z, helical.x, helical.y, helical.z});
@@ -497,11 +497,11 @@ std::vector<double> Rugnux::ExperimentKey(bool rotation_angles) const {
g->GetScreeningWedge().value_or(-1.0f)});
} else
k.push_back(0.0);
if (const auto uc = experiment_.GetUnitCell())
if (const auto uc = e.GetUnitCell())
k.insert(k.end(), {1.0, uc->a, uc->b, uc->c, uc->alpha, uc->beta, uc->gamma});
else
k.push_back(0.0);
const auto &x = experiment_.GetIndexingSettings();
const auto &x = e.GetIndexingSettings();
k.insert(k.end(), {x.GetTolerance(), x.GetFFT_MinUnitCell_A(), x.GetFFT_MaxUnitCell_A(),
static_cast<double>(x.GetFFT_NumVectors()), x.GetFFT_HighResolution_A(),
x.GetFFT_MinAngle_deg(), x.GetFFT_MaxAngle_deg(), x.GetUnitCellDistTolerance(),
@@ -514,23 +514,24 @@ std::vector<double> Rugnux::ExperimentKey(bool rotation_angles) const {
k.insert(k.end(), {f.signal_to_noise_threshold, static_cast<double>(f.photon_count_threshold),
static_cast<double>(f.min_pix_per_spot.value_or(-1)), static_cast<double>(f.max_pix_per_spot),
f.high_resolution_limit.value_or(-1.0f)});
k.insert(k.end(), {static_cast<double>(experiment_.GetSpaceGroupOrP1().number),
static_cast<double>(experiment_.GetImageNum())});
k.insert(k.end(), {static_cast<double>(e.GetSpaceGroupOrP1().number),
static_cast<double>(e.GetImageNum())});
return k;
}
std::vector<double> Rugnux::SpotFindingKey(const AzimuthalIntegrationMapping &mapping) const {
std::vector<double> Rugnux::SpotFindingKey(const DiffractionExperiment &x,
const AzimuthalIntegrationMapping &mapping) const {
// Without the rotation angles: a frame's spots are found the same under any of them, and the angle
// each spot carries is stamped again where a list is taken from the store (prefetch_spots). The
// probes of the rotation-scale walk differ only there, so they find the spots of a frame once.
std::vector<double> k = ExperimentKey(/*rotation_angles=*/false);
std::vector<double> k = ExperimentKey(x, /*rotation_angles=*/false);
const auto &f = config_.spot_finding;
k.insert(k.end(), {static_cast<double>(f.enable), f.low_resolution_limit.value_or(-1.0f),
f.cutoff_spot_count_low_res, f.high_res_gap_Q_recipA.value_or(-1.0f),
f.ice_ring_width_Q_recipA, static_cast<double>(f.adaptive_threshold),
f.false_pixels_per_frame, static_cast<double>(f.measured_ring_q_recipA.size())});
k.insert(k.end(), f.measured_ring_q_recipA.begin(), f.measured_ring_q_recipA.end());
k.insert(k.end(), {static_cast<double>(experiment_.IsDetectIceRings()),
k.insert(k.end(), {static_cast<double>(x.IsDetectIceRings()),
static_cast<double>(pixel_mask_.GetBinaryMaskChecksum()),
static_cast<double>(mapping.GetPixelToBinChecksum()),
static_cast<double>(mapping.GetCorrectionsChecksum())});
rugnux/Rugnux.h
+11 -2
View File
@@ -787,7 +787,12 @@ class Rugnux {
// not evaluate it again. Shared with the copies the run makes of itself.
std::shared_ptr<AzimuthalIntegrationGeometryCache> azint_geometry_ =
std::make_shared<AzimuthalIntegrationGeometryCache>();
[[nodiscard]] std::vector<double> SpotFindingKey(const AzimuthalIntegrationMapping &mapping) const;
[[nodiscard]] std::vector<double> SpotFindingKey(const AzimuthalIntegrationMapping &mapping) const {
return SpotFindingKey(experiment_, mapping);
}
// The same, for spots found at another experiment (a trial beam centre) against its mapping.
[[nodiscard]] std::vector<double> SpotFindingKey(const DiffractionExperiment &x,
const AzimuthalIntegrationMapping &mapping) const;
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
@@ -809,7 +814,11 @@ class Rugnux {
[[nodiscard]] std::vector<double> FirstPassInputKey() const;
// The part of FirstPassInputKey that is the experiment and the spot-finding settings; without the
// goniometer's angles where rotation_angles is false (SpotFindingKey).
[[nodiscard]] std::vector<double> ExperimentKey(bool rotation_angles = true) const;
[[nodiscard]] std::vector<double> ExperimentKey(bool rotation_angles = true) const {
return ExperimentKey(experiment_, rotation_angles);
}
// The same, of another experiment beside the run's settings.
[[nodiscard]] std::vector<double> ExperimentKey(const DiffractionExperiment &x, bool rotation_angles = true) const;
// The file's detector distance, from before the first pass of the rotation two-pass: a pass whose
// distance is still this one asks the post-refinement to test "the header distance is right"
// (PostRefineSettings::distance_at_header); a pass that has walked off it does not.
rugnux/RugnuxFirstPass.cpp
+215 -22
View File
@@ -88,6 +88,105 @@
using namespace rugnux_internal;
namespace {
// The spots of one first-pass frame. A frame's spots are a pure function of that frame and the
// settings - the engine carries nothing from one image to the next - so it does not matter which
// engine, or which thread, finds them.
std::vector<SpotToSave> FindFirstPassSpots(JFJochReader &reader, const SpotFindingSettings &settings,
const JFJochReaderDataset &dataset,
MXAnalysisWithoutFPGA &analysis,
AzimuthalIntegrationProfile &profile,
JFJochReaderRawImage &img, int ordinal, int image_idx,
Logger &logger) {
std::vector<SpotToSave> spots;
try {
if (reader.ReadRawImage(image_idx, img)) {
DataMessage m{};
m.number = ordinal;
m.original_number = image_idx;
auto first_pass = settings;
first_pass.indexing = false;
first_pass.quick_integration = false;
m.image = img.image;
if (dataset.efficiency.size() > image_idx)
m.image_collection_efficiency = dataset.efficiency[image_idx];
analysis.Analyze(m, profile, first_pass);
spots = std::move(m.spots);
}
} catch (const std::exception &e) {
if (IsFatalResourceError(e)) throw;
logger.Warning("First-pass spot read failed for image {}: {}", image_idx, e.what());
}
return spots;
}
// First-pass spot finding on workers of its own, all on one card (`device`), for work that runs
// beside the first pass. Each worker keeps its engine from one call to the next - building one
// costs far more than a frame - and gives it back on that card.
class SpotFindingWorkers {
public:
SpotFindingWorkers(JFJochReader &reader, const DiffractionExperiment &x,
const AzimuthalIntegrationMapping &mapping, const PixelMask &mask,
const SpotFindingSettings &settings, const JFJochReaderDataset &dataset,
int start_image, int stride, size_t nworkers, int device)
: reader(reader), x(x), mapping(mapping), mask(mask), settings(settings), dataset(dataset),
start_image(start_image), stride(stride), device(device), engines(std::max<size_t>(nworkers, 1)) {}
~SpotFindingWorkers() {
OnWorkers(engines.size(), [&](size_t t) { engines[t] = Engine{}; });
}
// The spots of each ordinal, in the order given.
std::vector<std::vector<SpotToSave>> Find(const std::vector<int> &ordinals) {
std::vector<std::vector<SpotToSave>> found(ordinals.size());
std::atomic<size_t> next{0};
OnWorkers(std::min(engines.size(), ordinals.size()), [&](size_t t) {
Engine &e = engines[t];
if (!e.analysis) {
e.indexer = std::make_unique<IndexAndRefine>(x, nullptr, /*retain_outcomes=*/false);
e.analysis = std::make_unique<MXAnalysisWithoutFPGA>(x, mapping, mask, *e.indexer,
/*enable_fused_adaptive_gpu=*/true);
e.profile = std::make_unique<AzimuthalIntegrationProfile>(mapping);
}
for (size_t i = next.fetch_add(1); i < ordinals.size(); i = next.fetch_add(1))
found[i] = FindFirstPassSpots(reader, settings, dataset, *e.analysis, *e.profile,
e.raw_image, ordinals[i], start_image + ordinals[i] * stride,
logger);
});
return found;
}
private:
struct Engine {
// Before the analysis engine, which page-locks these bytes for its uploads.
JFJochReaderRawImage raw_image;
std::unique_ptr<IndexAndRefine> indexer;
std::unique_ptr<MXAnalysisWithoutFPGA> analysis;
std::unique_ptr<AzimuthalIntegrationProfile> profile;
};
// fn(t) on n threads of its own, on the card.
template <class Fn> void OnWorkers(size_t n, Fn fn) {
std::vector<std::future<void>> futures;
for (size_t t = 0; t < n; t++)
futures.emplace_back(std::async(std::launch::async, [this, &fn, t] {
pin_gpu(device);
fn(t);
}));
for (auto &f : futures)
f.get();
}
JFJochReader &reader;
const DiffractionExperiment &x;
const AzimuthalIntegrationMapping &mapping;
const PixelMask &mask;
const SpotFindingSettings settings;
const JFJochReaderDataset &dataset;
const int start_image, stride, device;
std::vector<Engine> engines;
Logger logger{"Rugnux"};
};
// Field by field, bitwise, for the spot memo's verification.
bool SameSpots(const std::vector<SpotToSave> &a, const std::vector<SpotToSave> &b) {
if (a.size() != b.size())
@@ -248,27 +347,8 @@ bool Rugnux::FirstPassRotationIndexing(PipelineLocals &p) {
// not matter which engine, or which thread, finds them.
const auto find_spots = [&](MXAnalysisWithoutFPGA &analysis, AzimuthalIntegrationProfile &profile,
JFJochReaderRawImage &img, int ordinal) {
const int image_idx = start_image + ordinal * config_.stride;
std::vector<SpotToSave> spots;
try {
if (reader_.ReadRawImage(image_idx, img)) {
DataMessage m{};
m.number = ordinal;
m.original_number = image_idx;
auto first_pass = config_.spot_finding;
first_pass.indexing = false;
first_pass.quick_integration = false;
m.image = img.image;
if (dataset->efficiency.size() > image_idx)
m.image_collection_efficiency = dataset->efficiency[image_idx];
analysis.Analyze(m, profile, first_pass);
spots = std::move(m.spots);
}
} catch (const std::exception &e) {
if (IsFatalResourceError(e)) throw;
logger.Warning("First-pass spot read failed for image {}: {}", image_idx, e.what());
}
return spots;
return FindFirstPassSpots(reader_, config_.spot_finding, *dataset, analysis, profile, img,
ordinal, start_image + ordinal * config_.stride, logger);
};
// Find the spots of every ordinal in the list that is not cached yet, on several workers, and
@@ -1200,7 +1280,97 @@ bool Rugnux::FirstPassRotationIndexing(PipelineLocals &p) {
std::vector<std::future<void>> futures; // last, so it is waited for before the rest goes
};
std::optional<ShortAxisSpeculation> short_spec;
// The beam-centre check below runs a second first pass at the centre measured from the
// background, and it reads the first one's answer only after that pass is over. So the spots at
// the measured centre are found, and both schemes fed and indexed, here beside the file centre's
// pass - on a copy of the experiment at that centre, with workers, a mapping and a spot list of
// their own - and taken at the check's place only if what they read is still what the run has
// there (the experiment with its rotation angles, SpotFindingKey at the measured centre, the spot
// budget); otherwise that pass is run there as before. The schemes are scored there in either
// case, on the run's own indexer.
//
// Only with a second GPU, whose card its spot finding takes: on one card the two passes share the
// card and the reader, and measured, running them side by side saved nothing. There the check
// runs its pass after the first, as it always did.
struct CentreSpeculation {
std::vector<double> experiment_key;
std::vector<double> key;
size_t spots_per_image = 0;
std::unique_ptr<DiffractionExperiment> experiment;
std::unique_ptr<PixelMask> mask;
std::unique_ptr<AzimuthalIntegrationMapping> mapping;
std::map<int, std::vector<SpotToSave>> spots;
SchemeIndexers ris;
std::future<void> done; // last, so it is waited for before the rest goes
};
std::optional<CentreSpeculation> centre_spec;
const auto start_centre_speculation = [&] {
if (get_gpu_count() < 2)
return;
JoinBeamCenterCapture();
if (cancelled_ || !config_.beam_center_check || !background_center_)
return;
auto &cs = centre_spec.emplace();
cs.experiment = std::make_unique<DiffractionExperiment>(experiment_);
cs.experiment->BeamX_pxl(background_center_->beam_x_pxl).BeamY_pxl(background_center_->beam_y_pxl);
cs.mask = std::make_unique<PixelMask>(pixel_mask_);
cs.mapping = std::make_unique<AzimuthalIntegrationMapping>(*cs.experiment, *cs.mask, *azint_geometry_);
cs.experiment_key = ExperimentKey(*cs.experiment);
cs.key = SpotFindingKey(*cs.experiment, *cs.mapping);
cs.spots_per_image = first_pass_spots_per_image;
{
// What an earlier pass already found at this centre under these settings.
std::lock_guard lock(first_pass_spots_->m);
for (const auto &[k, spots] : first_pass_spots_->entries)
if (k == cs.key)
cs.spots = spots;
}
// At this pass's angles, as prefetch_spots takes them.
const DiffractionGeometry spot_geometry = cs.experiment->GetDiffractionGeometry();
for (auto &[ordinal, spots] : cs.spots)
StampSpotAngles(spots, spot_geometry);
cs.done = std::async(std::launch::async, [this, &cs, &schemes, &validation, &indexer_pool, &run_indexing,
settings = config_.spot_finding, dataset, start_image,
nworkers = engines.size()] {
SpotFindingWorkers workers(reader_, *cs.experiment, *cs.mapping, *cs.mask, settings, *dataset,
start_image, config_.stride, nworkers, /*device=*/1);
const auto find = [&](const std::vector<int> &ordinals) {
std::vector<int> wanted;
for (const int o : ordinals)
if (!cs.spots.contains(o))
wanted.push_back(o);
auto found = workers.Find(wanted);
for (size_t i = 0; i < wanted.size(); i++)
cs.spots.emplace(wanted[i], std::move(found[i]));
};
// As feed_scheme, at the measured centre.
const auto gonio = cs.experiment->GetGoniometer();
for (const auto &[name, ordinals] : schemes) {
auto ri = std::make_unique<RotationIndexer>(*cs.experiment, *indexer_pool);
ri->MaxSpotsPerImage(cs.spots_per_image);
ri->VerifyMemo(config_.verify_first_pass_memo);
for (size_t i = 0; i < ordinals.size(); i++) {
if (cancelled_ || ri->AccumulationFull())
break;
if (!cs.spots.contains(ordinals[i]))
find({ordinals.begin() + i,
ordinals.begin() + std::min(ordinals.size(), i + SPOT_PREFETCH_CHUNK)});
std::optional<float> angle;
if (gonio)
angle = gonio->GetAngle_deg(static_cast<float>(ordinals[i])) + gonio->GetWedge_deg() / 2.0f;
ri->ProcessImage(ordinals[i], cs.spots.at(ordinals[i]), angle);
}
cs.ris.push_back(std::move(ri));
}
find(validation);
std::vector<RotationIndexer *> rp;
for (const auto &ri : cs.ris)
rp.push_back(ri.get());
run_indexing(rp);
});
};
FirstPass best = pick_best(*indexer_pool, *indexer, [&] {
start_centre_speculation();
if (cancelled_ || experiment_.GetUnitCell().has_value())
return;
auto &sp = short_spec.emplace();
@@ -1322,7 +1492,30 @@ bool Rugnux::FirstPassRotationIndexing(PipelineLocals &p) {
{
const int majority = static_cast<int>(validation.size()) / 2;
try_beam_center(measured_x, measured_y);
FirstPass alt = pick_best(*indexer_pool, *indexer);
FirstPass alt;
if (centre_spec)
centre_spec->done.get();
if (centre_spec && centre_spec->experiment_key == ExperimentKey()
&& centre_spec->key == SpotFindingKey(*spot_mapping)
&& centre_spec->spots_per_image == first_pass_spots_per_image) {
// What pick_best would have found and fed here: hand the spots to the run's cache and
// its memo as prefetch_spots does, and score the schemes on the run's own indexer.
spot_cache = centre_spec->spots;
{
std::lock_guard lock(first_pass_spots_->m);
auto &entries = first_pass_spots_->entries;
if (std::none_of(entries.begin(), entries.end(),
[&](const auto &e) { return e.first == centre_spec->key; })) {
constexpr size_t MAX_SPOT_KEYS = 8;
if (entries.size() == MAX_SPOT_KEYS)
entries.erase(entries.begin());
entries.emplace_back(centre_spec->key, centre_spec->spots);
}
}
alt = score_schemes(centre_spec->ris, *indexer);
} else
alt = pick_best(*indexer_pool, *indexer);
centre_spec.reset();
const bool header_indexes = best.result.has_value() && best.score > majority;
bool measured_indexes = alt.result.has_value() && alt.score > majority;
if (!header_indexes) {