Merge branch 'refobs' into rc170

# Conflicts:
#	docs/CHANGELOG.md

docs/CHANGELOG.md

@@ -4,7 +4,7 @@
 ### 1.0.0-rc.170
 
 * Fixed a `jfjoch_broker` crash during indexing: sorting no longer misbehaves on non-finite values, and GPU FFT indexer kernel launches are now error-checked.
-* rugnux needs about a third less peak memory to scale, merge and post-refine rotation data, with identical results.
+* rugnux needs about 40% less peak memory to scale, merge and post-refine rotation data, with identical results.
 * `rugnux --model`: the placed coordinate file carries the space group its own coordinates obey, and says so when that is not the group the reflection files beside it carry.
 * `jfjoch_viewer`: fixes in the dataset plots, inspector and layout; spot markers lose their black outline by default (a checkbox under "Image features" restores it) and the highest-pixel markers are white boxes around the pixel.
 

image_analysis/scale_merge/RotationScaleMerge.cpp

@@ -299,6 +299,15 @@ void RotationScaleMerge::MeasureIncidentFlux(const std::vector<double> &mean_bkg
 
 void RotationScaleMerge::Ingest() {
     n_frames = static_cast<int>(partials_out.size());
+    partials_released = false;
+    resident_ingest = false;
+#ifdef JFJOCH_USE_CUDA
+    // Probed before the observation array is built: which record the build lays down (the narrow
+    // resident one or the full Obs) depends on whether the device pipeline will be active.
+    gpu_ = std::make_unique<RotationScaleMergeGPU>();
+    gpu_active_ = gpu_->Available();
+    resident_ingest = gpu_active_ && observation_dump_path.empty();
+#endif
 
     size_t total = 0;
     for (const auto &o : partials_out) total += o.reflections.size();
@@ -490,23 +499,40 @@ void RotationScaleMerge::Ingest() {
     BuildInRangeObservations(keys);
     rawrun_group.assign(rawrun_start.size(), -1);
     logger.Info("RotationScaleMerge: ingested {} partial observations from {} frames ({} distinct hkl)",
-                partials.size(), n_frames, rawrun_start.size());
+                n_partials_obs, n_frames, rawrun_start.size());
 
     SmoothMosaicityAndPartiality();
 
-    // The corr every pass starts from (see the header).
-    corr_ingested.resize(partials.size());
-    ParallelChunks(static_cast<int>(partials.size()), ThreadsForWork(partials.size(), nthreads),
-                   [&](int lo, int hi) {
-        for (int i = lo; i < hi; ++i) corr_ingested[i] = partials[i].corr;
-    });
+    // The corr every pass starts from (see the header). The resident build filled it already,
+    // straight from the source reflections (and the smoothing between there and here never
+    // touches corr).
+    if (!resident_ingest) {
+        corr_ingested.resize(partials.size());
+        ParallelChunks(static_cast<int>(partials.size()), ThreadsForWork(partials.size(), nthreads),
+                       [&](int lo, int hi) {
+            for (int i = lo; i < hi; ++i) corr_ingested[i] = partials[i].corr;
+        });
+    }
 
 #ifdef JFJOCH_USE_CUDA
-    // Bring the partial-scaling loop onto the GPU when one is present. Upload the immutable per-obs
-    // fields once (corr lives on the device, refreshed each pass); the CPU keeps the sort/keying/combine.
-    gpu_ = std::make_unique<RotationScaleMergeGPU>();
-    gpu_active_ = gpu_->Available();
-    if (gpu_active_) {
+    // Feed the device (probed at the top of Ingest). On the resident path the layout and every
+    // immutable field went up inside BuildInRangeObservations, straight from the source
+    // reflections; what is staged here is the partiality, which the recompute above rewrote.
+    if (resident_ingest) {
+        using F = RotationScaleMergeGPU::ObsField;
+        constexpr int STAGE_SLICE = 8 << 20;
+        const int n = n_partials_obs;
+        const int slice = std::min(std::max(n, 1), STAGE_SLICE);
+        const std::unique_ptr<float[]> buf(new float[slice]);
+        for (int at = 0; at < n; at += slice) {
+            const int cnt = std::min(n - at, slice);
+            ParallelChunks(cnt, nthreads, [&](int lo, int hi) {
+                for (int i = lo; i < hi; ++i) buf[i] = partials_ingest[at + i].partiality;
+            });
+            gpu_->SetObsField(F::Partiality, at, cnt, buf.get());
+        }
+    }
+    if (gpu_active_ && !resident_ingest) {
         const int n = static_cast<int>(partials.size());
         gpu_->SetPartialsLayout(n, n_frames, frame_start.data(), frame_count.data());
         // Staged in bounded slices: fourteen full-length arrays at once are more than half of the
@@ -545,12 +571,23 @@ void RotationScaleMerge::Ingest() {
             gpu_->SetObsFrame(at, cnt, frm.get());
             gpu_->SetObsOnIce(at, cnt, onice.get());
         }
+    }
+    if (gpu_active_) {
         gpu_->SetRawRuns(static_cast<int>(rawrun_start.size()), static_cast<int>(perm.size()), perm.data(),
                          rawrun_start.data(), rawrun_count.data(),
                          rawrun_h.data(), rawrun_k.data(), rawrun_l.data());
         gpu_->SetFrameCellOk(frame_cell_ok.data());
         logger.Info("RotationScaleMerge: GPU scaling + combine + scale-fulls + merge active");
     }
+    // The device now holds every per-obs field the resident pipeline reads, so the host record's
+    // one remaining reader is RockingEventFrames - which on this path sees only ingest-time values
+    // (see the header). Take its answer now and hand the record back. The dump path keeps the full
+    // Obs array: the CPU combine is what writes the dump.
+    if (resident_ingest) {
+        rocking_event_frames_at_ingest = RockingEventFrames();
+        std::vector<ObsIngest>().swap(partials_ingest);
+        partials_released = true;
+    }
 #endif
 }
 
@@ -653,9 +690,99 @@ void RotationScaleMerge::BuildInRangeObservations(std::vector<SortKey> &keys) {
     // The observations themselves, one frame per thread. finite_ok is the per-obs AcceptReflection
     // finiteness (immutable) - it lets ComputeAsuGroups stamp the ASU-group id per obs from a flat
     // 1-byte array instead of re-reading the fat Obs struct for every space group.
+    finite_ok.assign(n_keep, 0);
+#ifdef JFJOCH_USE_CUDA
+    if (resident_ingest) {
+        // The resident path: build the narrow host record, then stage every immutable per-obs field
+        // to the device straight from the source reflections, in frame-aligned bounded slices. The
+        // 80-byte Obs is never materialised here - beside the source reflections it was another
+        // copy of the whole payload - and the slice buffers are bounded whatever the run.
+        partials_ingest.clear();
+        partials_ingest.resize(n_keep);
+        corr_ingested.resize(n_keep);
+        ParallelFor(n_frames, nthreads, [&](int o) {
+            int32_t src = src_start[o];
+            int32_t at = frame_start[o];
+            for (const auto &r : partials_out[o].reflections) {
+                if (!keep[src]) { ++src; continue; }
+                ++src;
+                ObsIngest &obs = partials_ingest[at];
+                obs.h = r.h; obs.k = r.k; obs.l = r.l;
+                obs.frame = o;
+                obs.image_number = r.image_number;
+                obs.d = r.d;
+                obs.zeta = r.zeta;
+                obs.delta_phi = r.delta_phi_deg;
+                obs.partiality = r.partiality;
+                const float prescaling_corr = static_cast<float>(
+                        r.prescaling_corr * r.qe_corr * r.flight_corr / frame_flux[o]);
+                corr_ingested[at] = r.image_scale_corr;
+                finite_ok[at] = (std::isfinite(r.I) && std::isfinite(prescaling_corr)
+                                 && prescaling_corr != 0.0f
+                                 && std::isfinite(r.sigma) && r.sigma > 0.0f) ? 1 : 0;
+                obs.rock_ok = (std::isfinite(r.image_scale_corr) && r.image_scale_corr > 0.0f
+                               && std::isfinite(r.I) && std::isfinite(r.sigma) && r.sigma > 0.0f)
+                              ? 1 : 0;
+                ++at;
+            }
+        });
+
+        gpu_->SetPartialsLayout(n_keep, n_frames, frame_start.data(), frame_count.data());
+        using F = RotationScaleMergeGPU::ObsField;
+        constexpr int STAGE_SLICE = 8 << 20;
+        constexpr int NFLOAT = 11;
+        int max_frame_obs = 0;
+        for (int o = 0; o < n_frames; ++o) max_frame_obs = std::max(max_frame_obs, static_cast<int>(frame_count[o]));
+        const int bufsize = std::max(1, std::max(STAGE_SLICE, max_frame_obs));
+        std::unique_ptr<float[]> fbuf[NFLOAT];
+        for (auto &b : fbuf) b.reset(new float[bufsize]);
+        const std::unique_ptr<int32_t[]> frm(new int32_t[bufsize]);
+        const std::unique_ptr<uint8_t[]> onice(new uint8_t[bufsize]);
+        int f0 = 0;
+        while (f0 < n_frames) {
+            int f1 = f0 + 1;
+            while (f1 < n_frames && frame_start[f1] + frame_count[f1] - frame_start[f0] <= bufsize)
+                ++f1;
+            const int32_t base = frame_start[f0];
+            const int cnt = frame_start[f1 - 1] + frame_count[f1 - 1] - base;
+            if (cnt > 0) {
+                ParallelFor(f1 - f0, nthreads, [&](int fo) {
+                    const int o = f0 + fo;
+                    int32_t src = src_start[o];
+                    int at = frame_start[o] - base;
+                    for (const auto &r : partials_out[o].reflections) {
+                        if (!keep[src]) { ++src; continue; }
+                        ++src;
+                        fbuf[0][at] = r.I;
+                        fbuf[1][at] = r.sigma;
+                        fbuf[2][at] = static_cast<float>(
+                                r.prescaling_corr * r.qe_corr * r.flight_corr / frame_flux[o]);
+                        fbuf[3][at] = r.zeta;
+                        fbuf[4][at] = r.image_scale_corr;
+                        fbuf[5][at] = r.bkg;
+                        fbuf[6][at] = r.var_bkg;
+                        fbuf[7][at] = r.image_number;
+                        fbuf[8][at] = r.d;
+                        fbuf[9][at] = r.predicted_x;
+                        fbuf[10][at] = r.predicted_y;
+                        frm[at] = o;
+                        onice[at] = r.on_ice_ring ? 1 : 0;
+                        ++at;
+                    }
+                });
+                constexpr F fields[NFLOAT] = {F::I, F::Sigma, F::PrescalingCorr, F::Zeta, F::Corr0,
+                                              F::Bkg, F::VarBkg, F::ImageNumber, F::D, F::Px, F::Py};
+                for (int k = 0; k < NFLOAT; ++k)
+                    gpu_->SetObsField(fields[k], base, cnt, fbuf[k].get());
+                gpu_->SetObsFrame(base, cnt, frm.get());
+                gpu_->SetObsOnIce(base, cnt, onice.get());
+            }
+            f0 = f1;
+        }
+    } else {
+#endif
     partials.clear();
     partials.resize(n_keep);
-    finite_ok.assign(n_keep, 0);
     ParallelFor(n_frames, nthreads, [&](int o) {
         int32_t src = src_start[o];
         int32_t at = frame_start[o];
@@ -682,6 +809,11 @@ void RotationScaleMerge::BuildInRangeObservations(std::vector<SortKey> &keys) {
             ++at;
         }
     });
+#ifdef JFJOCH_USE_CUDA
+    }
+#endif
+
+    n_partials_obs = n_keep;
 
     if (n_keep < n_obs)
         logger.Info("RotationScaleMerge: dropped {} of {} observations ({} of {} distinct hkl) outside the "
@@ -749,7 +881,7 @@ namespace {
 // window would over-smooth a fast-slipping crystal and under-smooth a sparse one. The frame's own
 // value is excluded when scoring, so a window is only preferred if it PREDICTS that frame better,
 // which is exactly the noise-versus-drift trade.
-void RotationScaleMerge::SmoothGeometry() {
+template <class T> void RotationScaleMerge::SmoothGeometry(std::vector<T> &obs) {
     const auto gon = x.GetGoniometer();
     if (!gon || n_frames < 5)
         return;
@@ -856,10 +988,10 @@ void RotationScaleMerge::SmoothGeometry() {
     }
 
     std::atomic<int64_t> changed{0};
-    ParallelChunks(static_cast<int>(partials.size()), nthreads, [&](int lo, int hi) {
+    ParallelChunks(static_cast<int>(obs.size()), nthreads, [&](int lo, int hi) {
         int64_t local = 0;
         for (int i = lo; i < hi; ++i) {
-            auto &o = partials[i];
+            auto &o = obs[i];
             if (o.frame < 0 || o.frame >= n_frames || !ok[o.frame])
                 continue;
             const std::array<Coord, 3> &r = frame_recip[o.frame];
@@ -877,11 +1009,16 @@ void RotationScaleMerge::SmoothGeometry() {
 
     logger.Info("Smoothed per-frame geometry over +-{} frames (chosen by cross-validation); "
                 "recomputed delta_phi for {} of {} partials", best_half, changed.load(),
-                partials.size());
+                obs.size());
 }
 
 void RotationScaleMerge::SmoothMosaicityAndPartiality() {
-    SmoothGeometry();
+    if (resident_ingest) SmoothMosaicityAndPartialityOn(partials_ingest);
+    else                 SmoothMosaicityAndPartialityOn(partials);
+}
+
+template <class T> void RotationScaleMerge::SmoothMosaicityAndPartialityOn(std::vector<T> &obs) {
+    SmoothGeometry(obs);
 
     // One rocking event, one exact-Bragg angle. Every partial's delta_phi is solved from its OWN
     // frame's lattice - by the predictor, and again by SmoothGeometry above - so the per-frame
@@ -907,15 +1044,15 @@ void RotationScaleMerge::SmoothMosaicityAndPartiality() {
                 int i = lo;
                 while (i < hi) {
                     int j = i + 1;
-                    while (j < hi && partials[perm[j]].image_number
-                                     - partials[perm[j - 1]].image_number <= max_frame_gap)
+                    while (j < hi && obs[perm[j]].image_number
+                                     - obs[perm[j - 1]].image_number <= max_frame_gap)
                         ++j;
                     if (j - i >= 2) {
-                        const float f0 = partials[perm[i]].image_number;
+                        const float f0 = obs[perm[i]].image_number;
                         double sum = 0.0;
                         int n = 0;
                         for (int m = i; m < j; ++m) {
-                            const auto &o = partials[perm[m]];
+                            const auto &o = obs[perm[m]];
                             if (!std::isfinite(o.delta_phi)) continue;
                             sum += o.delta_phi + increment_deg * (o.image_number - f0);
                             ++n;
@@ -923,7 +1060,7 @@ void RotationScaleMerge::SmoothMosaicityAndPartiality() {
                         if (n >= 2) {
                             const double c = sum / n;
                             for (int m = i; m < j; ++m) {
-                                auto &o = partials[perm[m]];
+                                auto &o = obs[perm[m]];
                                 o.delta_phi = static_cast<float>(c - increment_deg
                                                                     * (o.image_number - f0));
                             }
@@ -998,9 +1135,9 @@ void RotationScaleMerge::SmoothMosaicityAndPartiality() {
     // partiality recomputed here matches the one integration was based on. sin(theta_B) = lambda/(2d).
     const double bandwidth_sigma = x.GetBandwidthFWHM().value_or(0.0f) / 2.3548;
     const double half_wavelength_A = x.GetWavelength_A() / 2.0;
-    ParallelChunks(static_cast<int>(partials.size()), nthreads, [&](int lo, int hi) {
+    ParallelChunks(static_cast<int>(obs.size()), nthreads, [&](int lo, int hi) {
         for (int i = lo; i < hi; ++i) {
-            auto &o = partials[i];
+            auto &o = obs[i];
             float mos = mos_smooth[o.frame];
             if (bandwidth_sigma > 0.0 && o.d > 0.0f) {
                 const double sin_theta = half_wavelength_A / o.d;
@@ -1087,7 +1224,7 @@ int RotationScaleMerge::ComputeAsuGroups(const HKLKeyGenerator &keygen) {
 #endif
 #ifdef JFJOCH_USE_CUDA
     std::vector<int32_t> group_ids;
-    if (gpu_active_) group_ids.resize(partials.size());
+    if (gpu_active_) group_ids.resize(n_partials_obs);
 #endif
     // Parallel over raw-hkl runs: distinct runs own disjoint perm ranges, hence disjoint observations.
     ParallelChunks(n_run, nthreads, [&](int rlo, int rhi) {
@@ -1111,7 +1248,7 @@ int RotationScaleMerge::ComputeAsuGroups(const HKLKeyGenerator &keygen) {
     // by group, parallel via per-chunk histograms over the flat group_ids; bit-identical to a serial fill
     // (chunk order == obs-index order, each chunk fills its groups sequentially).
     if (gpu_active_) {
-        const int n = static_cast<int>(partials.size());
+        const int n = n_partials_obs;
         const int nt = static_cast<int>(std::max<size_t>(1, std::min(nthreads, static_cast<size_t>(std::max(1, n)))));
         const int chunk = (n + nt - 1) / nt;
         const int n_hist = chunk > 0 ? (n + chunk - 1) / chunk : 0;   // chunks with work in them; <= nt
@@ -1990,6 +2127,23 @@ namespace {
 // the finest stretch of the sweep this file is allowed to speak about: the partials of one event are
 // welded into a single full, so two frames closer together than this are not separable observations.
 int RotationScaleMerge::RockingEventFrames() const {
+    if (partials_released)
+        return rocking_event_frames_at_ingest;
+    if (resident_ingest)   // Ingest's own call, before the narrow record is handed back
+        return RockingEventFramesOver(
+                [&](int i) { return partials_ingest[i].rock_ok != 0; },
+                [&](int i) { return partials_ingest[i].image_number; });
+    return RockingEventFramesOver(
+            [&](int i) {
+                const Obs &o = partials[i];
+                return std::isfinite(o.corr) && o.corr > 0.0f && std::isfinite(o.I)
+                       && std::isfinite(o.sigma) && o.sigma > 0.0f;
+            },
+            [&](int i) { return partials[i].image_number; });
+}
+
+template <class UsableFn, class ImgFn>
+int RotationScaleMerge::RockingEventFramesOver(UsableFn usable, ImgFn img) const {
     std::vector<int> hist(n_frames + 2, 0);
     std::vector<float> im;
     int n_event = 0;
@@ -1997,10 +2151,9 @@ int RotationScaleMerge::RockingEventFrames() const {
         const int lo = rawrun_start[r], hi = lo + rawrun_count[r];
         im.clear();
         for (int p = lo; p < hi; ++p) {
-            const Obs &o = partials[perm[p]];
-            if (std::isfinite(o.corr) && o.corr > 0.0f && std::isfinite(o.I)
-                && std::isfinite(o.sigma) && o.sigma > 0.0f)
-                im.push_back(o.image_number);
+            const int i = perm[p];
+            if (usable(i))
+                im.push_back(img(i));
         }
         for (size_t i = 0; i < im.size();) {
             size_t k = i + 1;

image_analysis/scale_merge/RotationScaleMerge.h

@@ -128,6 +128,22 @@ private:
     // re-run over the samples that survive the automatic resolution cutoff. See MergeAndStats.
     struct Sample { double s2, I2, dev2; float d; };
 
+    // The record the ingest keeps on the host when the device pipeline is resident (a GPU and no
+    // observation dump): every immutable per-obs field is uploaded straight from the source
+    // reflections inside BuildInRangeObservations, and the host keeps only what its own remaining
+    // stages touch - the geometry smoothing (delta_phi, then the recomputed partiality), plus the
+    // rocking-event walk's inputs. Forty bytes against the Obs's eighty, and on that path the fat
+    // record is never materialised at all: on a fine-sliced long axis it alone is many gigabytes,
+    // held beside the source reflections it was copied from.
+    struct ObsIngest {
+        int32_t h, k, l, frame;
+        float image_number, d, zeta, delta_phi, partiality;
+        // The rocking-event usability test (RockingEventFrames), evaluated on the ingest-time
+        // values - which is what that walk reads on the resident path, where the host corr is
+        // never refreshed (see partials_released).
+        uint8_t rock_ok;
+    };
+
     // The narrow per-observation record the ingest sort orders: the raw hkl the runs are cut on, the
     // frame position that breaks a tie inside one, the observation's own index (which makes the order
     // total - see the .cpp), and the resolution the range test reads. Twenty-four bytes against the
@@ -190,8 +206,24 @@ private:
     double resolution_cc_target = 0.30;
     int report_shell_count = 9;
 
-    // Flat buffers, allocated once by Ingest() and reused across Run() calls.
+    // Flat buffers, allocated once by Ingest() and reused across Run() calls. Exactly one of the
+    // two observation arrays is built: the narrow resident record when the device pipeline is
+    // active and no observation dump was asked for (resident_ingest), the full Obs otherwise.
     std::vector<Obs> partials;                    // all per-frame partials, grouped by frame
+    std::vector<ObsIngest> partials_ingest;       // the resident-path record (same order)
+    bool resident_ingest = false;
+    // How many observations Ingest built (== partials.size() until any release below), so the
+    // sizes survive when the array itself does not.
+    int n_partials_obs = 0;
+    // With the GPU resident and no observation dump, no host stage reads the 80-byte Obs records
+    // after the device upload: scaling, combine, scale-fulls and merge run on the device,
+    // ComputeAsuGroups stamps the flat group_ids array, every pass restarts corr from
+    // corr_ingested, and the CPU combine only runs for the dump. The one late host reader,
+    // RockingEventFrames, sees only ingest-time values on that path (the host corr is refreshed
+    // only in the dump fallback), so Ingest takes its answer and hands the array - a second
+    // full-size copy of the partials, gigabytes on a fine-sliced long axis - straight back.
+    bool partials_released = false;
+    int rocking_event_frames_at_ingest = 0;
     std::vector<int32_t> frame_start, frame_count; // CSR ranges of `partials` per frame
     std::vector<uint8_t> frame_cell_ok;           // per-frame cell-consistency mask (1 = kept)
     std::vector<uint8_t> finite_ok;               // per-obs AcceptReflection finiteness (immutable; 1 = kept)
@@ -392,8 +424,9 @@ private:
     // per-frame partials of one rocking event tile the curve consistently (they sum toward 1) before the
     // 3D combine. Deterministic (frame order); replaces the old arrival-order mosaicity moving average
     // that prediction applied. SG-independent, so done once in Ingest.
-    void SmoothGeometry();
-    void SmoothMosaicityAndPartiality();
+    template <class T> void SmoothGeometry(std::vector<T> &obs);
+    void SmoothMosaicityAndPartiality();   // dispatch: the resident narrow record or the full Obs
+    template <class T> void SmoothMosaicityAndPartialityOn(std::vector<T> &obs);
 
     void Combine();   // partials -> fulls (CPU)
 
@@ -458,6 +491,8 @@ private:
     // Walked on the partials, because the combine also runs on the GPU and a full keeps only the frame
     // of its peak partial.
     [[nodiscard]] int RockingEventFrames() const;
+    template <class UsableFn, class ImgFn>
+    [[nodiscard]] int RockingEventFramesOver(UsableFn usable, ImgFn img) const;
     // Per-batch delta-CC1/2 on the corrected fulls: measure what keeping each batch of the sweep costs
     // the merged intensities, convict the batches that cost significantly, slide the conviction's edges
     // onto the frames that carry it, and turn the result into the disposition ledger. Fills