Merge branch 'par-centre' into rc175

rugnux/Rugnux.cpp

@@ -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

@@ -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

@@ -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) {