From df90848bcd8d69274eacc8705be2fb80b97c82a4 Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Wed, 7 Oct 2026 15:40:27 +0200 Subject: [PATCH] rugnux: move RefineStillsGeometry to RugnuxStills.cpp A verbatim move; the new file starts with Rugnux.cpp's include block so that the same declarations are visible to the moved body. Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01SVmAWnzCmRKAXVUCdc4iNi --- rugnux/CMakeLists.txt | 1 + rugnux/Rugnux.cpp | 180 --------------------------- rugnux/RugnuxStills.cpp | 268 ++++++++++++++++++++++++++++++++++++++++ 3 files changed, 269 insertions(+), 180 deletions(-) create mode 100644 rugnux/RugnuxStills.cpp diff --git a/rugnux/CMakeLists.txt b/rugnux/CMakeLists.txt index b34ed52fb..f2cc51b7e 100644 --- a/rugnux/CMakeLists.txt +++ b/rugnux/CMakeLists.txt @@ -4,6 +4,7 @@ ADD_LIBRARY(JFJochRugnux STATIC Rugnux.cpp RugnuxPipeline.h + RugnuxStills.cpp RugnuxDefaults.cpp RugnuxDefaults.h Rugnux.h diff --git a/rugnux/Rugnux.cpp b/rugnux/Rugnux.cpp index 1824be674..759d08183 100644 --- a/rugnux/Rugnux.cpp +++ b/rugnux/Rugnux.cpp @@ -1482,186 +1482,6 @@ void Rugnux::MaskDefectivePixels(int start_image, const std::vector &sample Note(fmt::format("{} defective pixel{} masked", hot.hot + hot.error, hot.hot + hot.error == 1 ? "" : "s")); } -void Rugnux::RefineStillsGeometry(int start_image, int end_image, int images_to_process, - RugnuxObserver *observer) { - Logger logger("Rugnux"); - - // Rotation has its own two-pass; this is stills only. - if (experiment_.IsRotationIndexing()) { - logger.Warning("--refine-geometry is stills-only; rotation data uses two-pass indexing. Ignoring."); - return; - } - // The bundle anchors a known cell; without one there is nothing to refine against. - const auto cell = experiment_.GetUnitCell(); - if (!cell.has_value()) { - logger.Warning("--refine-geometry needs a reference unit cell (-C / -S / reference MTZ). Skipping."); - return; - } - - const int refine_frames = std::max(1, config_.refine_geometry.value()); - if (observer) - observer->OnPhase("Geometry refinement (first pass)"); - Step("Geometry refinement"); - - const auto dataset = reader_.GetDataset(); - - // Index a spread sample large enough to yield ~refine_frames strong frames even at a low hit rate. - const int sample_budget = std::min(images_to_process, std::max(refine_frames * 50, 8000)); - const std::vector sample = select_equally_spaced_image_ordinals(images_to_process, sample_budget); - - // First-pass engines, built from the current (nominal) geometry: an index-only pass (no - // integration) to obtain each frame's spots + assigned HKL + orientation. - AzimuthalIntegrationMapping mapping(experiment_, pixel_mask_); - IndexerThreadPool pool(experiment_.GetIndexingSettings(), IndexerConstruction::OnFirstUse); - IndexAndRefine indexer(experiment_, &pool, /*retain_outcomes=*/false); - - auto pass_settings = config_.spot_finding; - pass_settings.enable = true; - pass_settings.indexing = true; - pass_settings.quick_integration = false; - - // A frame is worth bundling only if it indexed this many spots; the bundle then takes the - // refine_frames strongest of those. - constexpr int MIN_STRONG_SPOTS = 20; - // Stop sampling once there are comfortably more strong frames than the bundle will use, instead of - // draining the whole budget. The budget is sized for a low-hit-rate dataset (refine_frames * 50, at - // least 8000), so on data that indexes well it meant re-indexing the ENTIRE run to keep 200 frames - - // 99% of this pass was sampling. Four times the bundle leaves the "strongest N" selection a real - // pool to choose from. - const int strong_target = refine_frames * 4; - - // The sample is cut into a FIXED number of interleaved stripes - stripe s is sample positions - // s, s+STRIPES, s+2*STRIPES, ... - and each stripe stops once it has contributed its share of the - // strong frames. Two things follow, both of which the previous shared-cursor-with-a-shared-stop - // arrangement got wrong. Every stripe is spread over the WHOLE run, so stopping early no longer - // means fitting the geometry to the beginning of it. And which frames get examined depends only - // on the data: it is not a race between the workers, and it does not change with -N, because a - // stripe is processed identically whichever worker happens to claim it. - constexpr int STRIPES = 32; - const int stripe_strong_target = std::max(1, (strong_target + STRIPES - 1) / STRIPES); - - std::vector frames; - std::mutex frames_mutex; - std::atomic next_stripe = 0; - std::atomic examined = 0; - - auto worker = [&]() { - pin_gpu(); // round-robin per worker thread; must precede engine construction - // Before the analysis engine, which page-locks these bytes for its uploads and unregisters - // them when it is destroyed: it must not outlive the buffer it was given. - JFJochReaderRawImage img; - MXAnalysisWithoutFPGA analysis(experiment_, mapping, pixel_mask_, indexer, - /*enable_fused_adaptive_gpu=*/true); - AzimuthalIntegrationProfile profile(mapping); - - for (int stripe = next_stripe.fetch_add(1); stripe < STRIPES && !cancelled_; - stripe = next_stripe.fetch_add(1)) { - int stripe_strong = 0; - for (int idx = stripe; idx < static_cast(sample.size()) && !cancelled_ - && stripe_strong < stripe_strong_target; idx += STRIPES) { - examined.fetch_add(1, std::memory_order_relaxed); - const int ordinal = sample[idx]; - const int image_idx = start_image + ordinal * config_.stride; - - bool read = false; - try { - read = reader_.ReadRawImage(image_idx, img); - } catch (const std::exception &e) { - if (IsFatalResourceError(e)) throw; - logger.Warning("Geometry refinement: failed to load image {}: {}", image_idx, e.what()); - continue; - } - if (!read) continue; - - DataMessage msg{}; - msg.image = img.image; - msg.number = ordinal; - msg.original_number = image_idx; - if (dataset->efficiency.size() > image_idx) - msg.image_collection_efficiency = dataset->efficiency[image_idx]; - - try { - analysis.Analyze(msg, profile, pass_settings); - } catch (const std::exception &e) { - if (IsFatalResourceError(e)) throw; - continue; - } - if (!msg.indexing_result.value_or(false) || !msg.indexing_lattice.has_value()) - continue; - - GeomRefineFrame f; - f.lattice = *msg.indexing_lattice; - f.ordinal = ordinal; - for (const auto &s : msg.spots) - if (s.indexed && s.lattice == 0) - f.spots.push_back(GeomRefineSpot{s.x, s.y, - static_cast(s.h), static_cast(s.k), static_cast(s.l)}); - if (f.spots.size() >= 6) { - if (static_cast(f.spots.size()) >= MIN_STRONG_SPOTS) - stripe_strong++; - const std::unique_lock ul(frames_mutex); - frames.push_back(std::move(f)); - } - } - } - }; - - std::vector> futures; - futures.reserve(config_.nthreads); - for (int i = 0; i < config_.nthreads; ++i) - futures.push_back(std::async(std::launch::async, worker)); - for (auto &fut : futures) - fut.get(); - - if (cancelled_) - return; - - // Keep the strongest frames (most indexed spots) for the bundle - the true cell indexes many - // spots per frame, and orientation diversity comes for free from independent serial stills. - std::vector strong; - for (auto &f : frames) - if (static_cast(f.spots.size()) >= MIN_STRONG_SPOTS) - strong.push_back(std::move(f)); - // frames is in the order the workers happened to finish, so the spot count alone does not order - // it: without the ordinal tie-break, equally strong frames would swap places between runs and a - // different bundle would be refined. - std::sort(strong.begin(), strong.end(), - [](const GeomRefineFrame &a, const GeomRefineFrame &b) { - if (a.spots.size() != b.spots.size()) return a.spots.size() > b.spots.size(); - return a.ordinal < b.ordinal; - }); - if (static_cast(strong.size()) > refine_frames) - strong.resize(refine_frames); - - logger.Info("Geometry refinement: indexed {} of {} sampled frames (budget {}), bundling {} strong " - "frames (>= {} spots)", - frames.size(), std::min(examined.load(), static_cast(sample.size())), - sample.size(), strong.size(), MIN_STRONG_SPOTS); - - GeometryRefinerSettings gr_settings; - gr_settings.crystal_system = experiment_.GetCrystalSystem(); - gr_settings.num_threads = config_.nthreads; - - const GeometryRefinerResult r = RefineGlobalGeometry( - experiment_.GetDiffractionGeometry(), *cell, strong, gr_settings); - if (!r.ok) { - logger.Warning("Geometry refinement did not run (too few strong frames / did not converge); " - "keeping the input geometry"); - return; - } - - logger.Info("Geometry refinement: beam ({:.2f}, {:.2f}) -> ({:.2f}, {:.2f}) px, distance {:.4f} -> {:.4f} mm", - experiment_.GetBeamX_pxl(), experiment_.GetBeamY_pxl(), r.beam_x_px, r.beam_y_px, - experiment_.GetDetectorDistance_mm(), r.distance_mm); - logger.Info("Geometry refinement: cell a/b/c {:.3f}/{:.3f}/{:.3f} -> {:.3f}/{:.3f}/{:.3f} A " - "(median residual {:.3f} px, {} frames, {} spots)", - cell->a, cell->b, cell->c, r.cell.a, r.cell.b, r.cell.c, - r.median_residual_px, r.frames_used, r.spots_used); - - experiment_.BeamX_pxl(r.beam_x_px).BeamY_pxl(r.beam_y_px).DetectorDistance_mm(r.distance_mm); - experiment_.SetUnitCell(r.cell); -} - ProcessResult Rugnux::Run(RugnuxObserver *observer) { // Each pass times only itself, so the canonical result's processing_time_s is the last pass alone - // on the rotation two-pass that is under a third of what the run actually took, the pre-scan and the diff --git a/rugnux/RugnuxStills.cpp b/rugnux/RugnuxStills.cpp new file mode 100644 index 000000000..e2e6ab916 --- /dev/null +++ b/rugnux/RugnuxStills.cpp @@ -0,0 +1,268 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#include +#include "Rugnux.h" +#include "../image_analysis/structure_refinement/ModelValidation.h" +#include "../writer/WriteModel.h" +#include "../image_analysis/indexing/SpindleCuspLoss.h" +#include "../image_analysis/bragg_integration/SpotWidth.h" +#include "../image_analysis/bragg_integration/SpotFootprint.h" +#include "../image_analysis/bragg_integration/BraggStencil.h" // BRAGG_FOOTPRINT_NSIGMA +#include "../image_analysis/hot_pixels/HotPixels.h" +#include "../image_analysis/SensorAbsorption.h" +#include "DiagnosticOutput.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "../reader/JFJochHDF5Reader.h" +#include "../common/JFJochMath.h" +#include "../common/ParallelFor.h" +#include "../common/Logger.h" +#include "../common/AzimuthalIntegrationMapping.h" +#include "../common/AzimuthalIntegrationProfile.h" +#include "../common/CUDAWrapper.h" +#include "../common/JFJochException.h" +#include "../common/time_utc.h" +#include "../writer/FileWriter.h" +#include "../image_analysis/MXAnalysisWithoutFPGA.h" +#include "../image_analysis/beam_stop/ShadowFinder.h" +#include "../image_analysis/IndexAndRefine.h" +#include "../image_analysis/indexing/AnalyzeIndexing.h" +#include "../image_analysis/indexing/HarmonicContamination.h" +#include "../image_analysis/geom_refinement/BeamCenterFFT.h" +#include "../image_analysis/geom_refinement/BeamCenterFromBackground.h" +#include "../image_analysis/geom_refinement/GeometryRefiner.h" +#include "../image_analysis/indexing/IndexerThreadPool.h" +#include "../image_analysis/spot_finding/ImageSpotFinderCPU.h" +#include "../image_analysis/spot_finding/SpotUtils.h" +#include "gemmi/twin.hpp" +#include "../image_analysis/azint/AzIntEngineCPU.h" +#include "../image_analysis/image_preprocessing/ImagePreprocessorCPU.h" +#include "../image_analysis/image_preprocessing/ImagePreprocessorBuffer.h" +#ifdef JFJOCH_USE_CUDA +#include "../image_analysis/image_preprocessing/ImagePreprocessorGPU.h" +#include "../image_analysis/image_preprocessing/ImagePreprocessorBufferGPU.h" +#include "../image_analysis/spot_finding/AdaptiveSpotFinderGPU.h" +#endif +#include "../image_analysis/scale_merge/Merge.h" +#include "../image_analysis/scale_merge/RfreeFlags.h" +#include "../image_analysis/scale_merge/RotationScaleMerge.h" +#include "../image_analysis/scale_merge/ResolutionCutoff.h" +#include "../image_analysis/scale_merge/ReindexAmbiguity.h" +#include "../image_analysis/scale_merge/ScalingResult.h" +#include "../image_analysis/scale_merge/SearchSpaceGroup.h" +#include "../image_analysis/geom_refinement/PostRefine.h" +#include "../image_analysis/lattice_search/LatticeSearch.h" +#include "../image_analysis/lattice_search/LePageLattice.h" +#include "../image_analysis/scale_merge/AnisotropyAnalysis.h" +#include "../image_analysis/scale_merge/FrenchWilson.h" +#include "../image_analysis/scale_merge/TwinningAnalysis.h" +#include "../image_analysis/scale_merge/TranslationalNCS.h" +#include "../image_analysis/scale_merge/HKLKey.h" +#include "../image_analysis/scale_merge/ScaleOnTheFly.h" +#include "../image_analysis/scale_merge/StillsPartialityRefine.h" +#include "../image_analysis/WriteReflections.h" +#include "../image_analysis/bragg_integration/CalcISigma.h" +#include "../common/Definitions.h" +#include "../common/CorrelationCoefficient.h" +#include +#include +#include "RugnuxPipeline.h" + +using namespace rugnux_internal; + +void Rugnux::RefineStillsGeometry(int start_image, int end_image, int images_to_process, + RugnuxObserver *observer) { + Logger logger("Rugnux"); + + // Rotation has its own two-pass; this is stills only. + if (experiment_.IsRotationIndexing()) { + logger.Warning("--refine-geometry is stills-only; rotation data uses two-pass indexing. Ignoring."); + return; + } + // The bundle anchors a known cell; without one there is nothing to refine against. + const auto cell = experiment_.GetUnitCell(); + if (!cell.has_value()) { + logger.Warning("--refine-geometry needs a reference unit cell (-C / -S / reference MTZ). Skipping."); + return; + } + + const int refine_frames = std::max(1, config_.refine_geometry.value()); + if (observer) + observer->OnPhase("Geometry refinement (first pass)"); + Step("Geometry refinement"); + + const auto dataset = reader_.GetDataset(); + + // Index a spread sample large enough to yield ~refine_frames strong frames even at a low hit rate. + const int sample_budget = std::min(images_to_process, std::max(refine_frames * 50, 8000)); + const std::vector sample = select_equally_spaced_image_ordinals(images_to_process, sample_budget); + + // First-pass engines, built from the current (nominal) geometry: an index-only pass (no + // integration) to obtain each frame's spots + assigned HKL + orientation. + AzimuthalIntegrationMapping mapping(experiment_, pixel_mask_); + IndexerThreadPool pool(experiment_.GetIndexingSettings(), IndexerConstruction::OnFirstUse); + IndexAndRefine indexer(experiment_, &pool, /*retain_outcomes=*/false); + + auto pass_settings = config_.spot_finding; + pass_settings.enable = true; + pass_settings.indexing = true; + pass_settings.quick_integration = false; + + // A frame is worth bundling only if it indexed this many spots; the bundle then takes the + // refine_frames strongest of those. + constexpr int MIN_STRONG_SPOTS = 20; + // Stop sampling once there are comfortably more strong frames than the bundle will use, instead of + // draining the whole budget. The budget is sized for a low-hit-rate dataset (refine_frames * 50, at + // least 8000), so on data that indexes well it meant re-indexing the ENTIRE run to keep 200 frames - + // 99% of this pass was sampling. Four times the bundle leaves the "strongest N" selection a real + // pool to choose from. + const int strong_target = refine_frames * 4; + + // The sample is cut into a FIXED number of interleaved stripes - stripe s is sample positions + // s, s+STRIPES, s+2*STRIPES, ... - and each stripe stops once it has contributed its share of the + // strong frames. Two things follow, both of which the previous shared-cursor-with-a-shared-stop + // arrangement got wrong. Every stripe is spread over the WHOLE run, so stopping early no longer + // means fitting the geometry to the beginning of it. And which frames get examined depends only + // on the data: it is not a race between the workers, and it does not change with -N, because a + // stripe is processed identically whichever worker happens to claim it. + constexpr int STRIPES = 32; + const int stripe_strong_target = std::max(1, (strong_target + STRIPES - 1) / STRIPES); + + std::vector frames; + std::mutex frames_mutex; + std::atomic next_stripe = 0; + std::atomic examined = 0; + + auto worker = [&]() { + pin_gpu(); // round-robin per worker thread; must precede engine construction + // Before the analysis engine, which page-locks these bytes for its uploads and unregisters + // them when it is destroyed: it must not outlive the buffer it was given. + JFJochReaderRawImage img; + MXAnalysisWithoutFPGA analysis(experiment_, mapping, pixel_mask_, indexer, + /*enable_fused_adaptive_gpu=*/true); + AzimuthalIntegrationProfile profile(mapping); + + for (int stripe = next_stripe.fetch_add(1); stripe < STRIPES && !cancelled_; + stripe = next_stripe.fetch_add(1)) { + int stripe_strong = 0; + for (int idx = stripe; idx < static_cast(sample.size()) && !cancelled_ + && stripe_strong < stripe_strong_target; idx += STRIPES) { + examined.fetch_add(1, std::memory_order_relaxed); + const int ordinal = sample[idx]; + const int image_idx = start_image + ordinal * config_.stride; + + bool read = false; + try { + read = reader_.ReadRawImage(image_idx, img); + } catch (const std::exception &e) { + if (IsFatalResourceError(e)) throw; + logger.Warning("Geometry refinement: failed to load image {}: {}", image_idx, e.what()); + continue; + } + if (!read) continue; + + DataMessage msg{}; + msg.image = img.image; + msg.number = ordinal; + msg.original_number = image_idx; + if (dataset->efficiency.size() > image_idx) + msg.image_collection_efficiency = dataset->efficiency[image_idx]; + + try { + analysis.Analyze(msg, profile, pass_settings); + } catch (const std::exception &e) { + if (IsFatalResourceError(e)) throw; + continue; + } + if (!msg.indexing_result.value_or(false) || !msg.indexing_lattice.has_value()) + continue; + + GeomRefineFrame f; + f.lattice = *msg.indexing_lattice; + f.ordinal = ordinal; + for (const auto &s : msg.spots) + if (s.indexed && s.lattice == 0) + f.spots.push_back(GeomRefineSpot{s.x, s.y, + static_cast(s.h), static_cast(s.k), static_cast(s.l)}); + if (f.spots.size() >= 6) { + if (static_cast(f.spots.size()) >= MIN_STRONG_SPOTS) + stripe_strong++; + const std::unique_lock ul(frames_mutex); + frames.push_back(std::move(f)); + } + } + } + }; + + std::vector> futures; + futures.reserve(config_.nthreads); + for (int i = 0; i < config_.nthreads; ++i) + futures.push_back(std::async(std::launch::async, worker)); + for (auto &fut : futures) + fut.get(); + + if (cancelled_) + return; + + // Keep the strongest frames (most indexed spots) for the bundle - the true cell indexes many + // spots per frame, and orientation diversity comes for free from independent serial stills. + std::vector strong; + for (auto &f : frames) + if (static_cast(f.spots.size()) >= MIN_STRONG_SPOTS) + strong.push_back(std::move(f)); + // frames is in the order the workers happened to finish, so the spot count alone does not order + // it: without the ordinal tie-break, equally strong frames would swap places between runs and a + // different bundle would be refined. + std::sort(strong.begin(), strong.end(), + [](const GeomRefineFrame &a, const GeomRefineFrame &b) { + if (a.spots.size() != b.spots.size()) return a.spots.size() > b.spots.size(); + return a.ordinal < b.ordinal; + }); + if (static_cast(strong.size()) > refine_frames) + strong.resize(refine_frames); + + logger.Info("Geometry refinement: indexed {} of {} sampled frames (budget {}), bundling {} strong " + "frames (>= {} spots)", + frames.size(), std::min(examined.load(), static_cast(sample.size())), + sample.size(), strong.size(), MIN_STRONG_SPOTS); + + GeometryRefinerSettings gr_settings; + gr_settings.crystal_system = experiment_.GetCrystalSystem(); + gr_settings.num_threads = config_.nthreads; + + const GeometryRefinerResult r = RefineGlobalGeometry( + experiment_.GetDiffractionGeometry(), *cell, strong, gr_settings); + if (!r.ok) { + logger.Warning("Geometry refinement did not run (too few strong frames / did not converge); " + "keeping the input geometry"); + return; + } + + logger.Info("Geometry refinement: beam ({:.2f}, {:.2f}) -> ({:.2f}, {:.2f}) px, distance {:.4f} -> {:.4f} mm", + experiment_.GetBeamX_pxl(), experiment_.GetBeamY_pxl(), r.beam_x_px, r.beam_y_px, + experiment_.GetDetectorDistance_mm(), r.distance_mm); + logger.Info("Geometry refinement: cell a/b/c {:.3f}/{:.3f}/{:.3f} -> {:.3f}/{:.3f}/{:.3f} A " + "(median residual {:.3f} px, {} frames, {} spots)", + cell->a, cell->b, cell->c, r.cell.a, r.cell.b, r.cell.c, + r.median_residual_px, r.frames_used, r.spots_used); + + experiment_.BeamX_pxl(r.beam_x_px).BeamY_pxl(r.beam_y_px).DetectorDistance_mm(r.distance_mm); + experiment_.SetUnitCell(r.cell); +}