experiment with chunk sizes

CMakeLists.txt

@@ -415,6 +415,7 @@ set(SourceFiles
     ${CMAKE_CURRENT_SOURCE_DIR}/src/ROIGeometry.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/src/DetectorGeometry.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/src/hist/PedestalTrackingPixelHistogram.cpp
+    ${CMAKE_CURRENT_SOURCE_DIR}/src/DataSetAggregation.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/src/Dtype.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/src/File.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/src/FilePtr.cpp

include/aare/DataSetAggregation.hpp

@@ -83,19 +83,14 @@ NDArray<FRAME_TYPE, 3> mean_over_dataset(const std::filesystem::path &path,
 
     int chunk_size = num_samples / num_threads;
 
-    for (int i = 0; i < num_threads - 1; ++i) {
-        int end_sample =
-            std::min((i + 1) * static_cast<size_t>(chunk_size), num_samples);
+    for (int i = 0; i < num_threads; ++i) {
+        size_t end_sample =
+            (i == num_threads - 1) ? num_samples : (i + 1) * chunk_size;
 
         futures[i] = std::async(std::launch::async, mean_over_samples,
                                 i * chunk_size, end_sample);
     }
 
-    int end_sample = num_samples;
-    futures[num_threads - 1] =
-        std::async(std::launch::async, mean_over_samples,
-                   (num_threads - 1) * chunk_size, end_sample);
-
     for (auto &f : futures) {
         try {
             f.get();
@@ -108,7 +103,7 @@ NDArray<FRAME_TYPE, 3> mean_over_dataset(const std::filesystem::path &path,
     return mean_per_sample;
 }
 
-/*
+/**
  * @brief Compute the mean current source data from a Jungfrau detector file
  * @param path The path to the Jungfrau detector file.
  * @param mean_over_num_frames The number of frames to mean over for each
@@ -117,126 +112,13 @@ NDArray<FRAME_TYPE, 3> mean_over_dataset(const std::filesystem::path &path,
  * @param num_threads The number of threads to use for parallel computation.
  * @param file_offset The offset in the file to start reading from.
  * @return An NDArray containing the mean current source data for each sample.
+ * @note On Linux the file is read with O_DIRECT in large chunks, bypassing the
+ * page cache, since the datasets are typically larger than the available RAM.
  */
 NDArray<uint16_t, 3> mean_current_source(const std::filesystem::path &path,
                                          const size_t mean_over_num_frames,
                                          const size_t num_samples,
                                          const int num_threads = 16,
-                                         size_t file_offset = 0) {
+                                         size_t file_offset = 0);
 
-    File meta_file(path);
-
-    size_t total_num_frames = meta_file.total_frames();
-
-    const size_t num_rows = meta_file.rows();
-    const size_t num_cols = meta_file.cols();
-    const size_t half_rows = num_rows / 2;
-
-    constexpr size_t columns_per_subcolumn = 64;
-
-    constexpr size_t num_subcolumns = 16;
-
-    if (total_num_frames !=
-        mean_over_num_frames * num_samples * columns_per_subcolumn +
-            file_offset) {
-        throw std::runtime_error(fmt::format(
-            "Frames mismatch: Expected {} frames got {}",
-            mean_over_num_frames * num_samples * columns_per_subcolumn +
-                file_offset,
-            total_num_frames));
-    }
-
-    const double inv_mean_over_num_frames =
-        1.0 / static_cast<double>(mean_over_num_frames);
-
-    NDArray<uint16_t, 3> cs_data({static_cast<ssize_t>(num_samples),
-                                  static_cast<ssize_t>(num_rows),
-                                  static_cast<ssize_t>(num_cols)},
-                                 0.0);
-
-    auto mean_over_samples = [&](size_t start_sample, size_t end_sample) {
-        File file(path);
-
-        file.seek(start_sample * mean_over_num_frames * columns_per_subcolumn +
-                  file_offset);
-
-        NDArray<uint16_t, 2> frame(
-            {static_cast<ssize_t>(num_rows), static_cast<ssize_t>(num_cols)});
-
-        std::vector<double> sum(num_rows * num_cols,
-                                0.0); // need to be double to avoid
-                                      // overflow for uint16_t
-
-        for (size_t i = start_sample; i < end_sample; ++i) {
-
-            std::fill(sum.begin(), sum.end(), 0.0);
-
-            for (size_t acq_idx = 0; acq_idx < mean_over_num_frames;
-                 ++acq_idx) {
-                for (size_t col_idx = 0; col_idx < columns_per_subcolumn;
-                     ++col_idx) {
-                    file.read_into(reinterpret_cast<std::byte *>(frame.data()));
-
-                    const uint16_t *frame_data_ptr = frame.data();
-
-                    for (size_t r = 0; r < half_rows; ++r) {
-                        const size_t row_offset = r * num_cols;
-                        for (size_t subcol_idx = 0; subcol_idx < num_subcolumns;
-                             ++subcol_idx) {
-
-                            const size_t absolute_col_idx =
-                                subcol_idx * columns_per_subcolumn +
-                                col_idx; // TODO: maybe precompute and use a
-                                         // lookup table for this
-                            sum[row_offset + absolute_col_idx] +=
-                                frame_data_ptr[row_offset + absolute_col_idx];
-                        }
-                    }
-
-                    for (size_t r = half_rows; r < num_rows; ++r) {
-                        const size_t row_offset = r * num_cols;
-                        for (size_t subcol_idx = 0; subcol_idx < num_subcolumns;
-                             ++subcol_idx) {
-
-                            const size_t absolute_col_idx =
-                                (subcol_idx + 1) * columns_per_subcolumn -
-                                col_idx;
-                            sum[row_offset + absolute_col_idx] +=
-                                frame_data_ptr[row_offset + absolute_col_idx];
-                        }
-                    }
-                }
-            }
-
-            uint16_t *cs_data_ptr = cs_data.data() + i * num_rows * num_cols;
-            for (size_t idx = 0; idx < num_rows * num_cols; ++idx) {
-                cs_data_ptr[idx] =
-                    static_cast<uint16_t>(sum[idx] * inv_mean_over_num_frames);
-            }
-        }
-    };
-
-    std::vector<std::future<void>> futures(num_threads);
-
-    const size_t chunk_size = num_samples / num_threads;
-
-    for (int i = 0; i < num_threads - 1; ++i) {
-        size_t end_sample = std::min((i + 1) * chunk_size, num_samples);
-
-        futures[i] = std::async(std::launch::async, mean_over_samples,
-                                i * chunk_size, end_sample);
-    }
-
-    int end_sample = num_samples;
-    futures[num_threads - 1] =
-        std::async(std::launch::async, mean_over_samples,
-                   (num_threads - 1) * chunk_size, end_sample);
-
-    for (auto &f : futures) {
-        f.get();
-    }
-
-    return cs_data;
-}
-
-} // namespace aare
+} // namespace aare

src/DataSetAggregation.cpp

@@ -0,0 +1,144 @@
+// SPDX-License-Identifier: MPL-2.0
+#include "aare/DataSetAggregation.hpp"
+#include "aare/JungfrauDataFile.hpp"
+
+#include <atomic>
+#include <cstdint>
+#include <cstdlib>
+#include <cstring>
+#include <fmt/format.h>
+#include <future>
+#include <memory>
+#include <stdexcept>
+#include <vector>
+
+namespace aare {
+
+// TODO: maybe O_DIRECT rea dis better to avoid the page cache
+
+NDArray<uint16_t, 3> mean_current_source(const std::filesystem::path &path,
+                                         const size_t mean_over_num_frames,
+                                         const size_t num_samples,
+                                         const int num_threads,
+                                         size_t file_offset) {
+
+    constexpr size_t columns_per_subcolumn = 64;
+    constexpr size_t frames_per_chunk = 16; // ~16 MiB per read for 512x1024
+
+    size_t num_rows = 0;
+    size_t num_cols = 0;
+    size_t total_num_frames = 0;
+    {
+        JungfrauDataFile meta_file(path);
+        num_rows = meta_file.rows();
+        num_cols = meta_file.cols();
+        total_num_frames = meta_file.total_frames();
+    }
+
+    const size_t half_rows = num_rows / 2;
+    const size_t num_subcolumns = num_cols / columns_per_subcolumn;
+    const size_t frames_per_sample =
+        mean_over_num_frames * columns_per_subcolumn;
+
+    if (total_num_frames != frames_per_sample * num_samples + file_offset) {
+        throw std::runtime_error(fmt::format(
+            "Frames mismatch: Expected {} frames got {}",
+            frames_per_sample * num_samples + file_offset, total_num_frames));
+    }
+    if (num_threads < 1)
+        throw std::invalid_argument("num_threads must be at least 1");
+
+    const double inv_mean_over_num_frames =
+        1.0 / static_cast<double>(mean_over_num_frames);
+
+    NDArray<uint16_t, 3> cs_data({static_cast<ssize_t>(num_samples),
+                                  static_cast<ssize_t>(num_rows),
+                                  static_cast<ssize_t>(num_cols)},
+                                 0);
+
+    auto worker = [&](const size_t sample_begin, const size_t sample_end) {
+        std::vector<uint64_t> sum(
+            num_rows * num_cols); // uint64_t to avoid overflow for uint16_t
+
+        JungfrauDataFile reader(path);
+
+        for (size_t sample = sample_begin; sample < sample_end; ++sample) {
+            std::fill(sum.begin(), sum.end(), 0);
+
+            const size_t first = file_offset + sample * frames_per_sample;
+            const size_t last = first + frames_per_sample;
+
+            reader.seek(first);
+
+            for (size_t frame_chunk = first; frame_chunk < last;
+                 frame_chunk += frames_per_chunk) {
+
+                const size_t n_frames_to_read =
+                    std::min({frames_per_chunk, last - frame_chunk});
+
+                auto frames = reader.read_n(n_frames_to_read);
+
+                for (size_t f = 0; f < n_frames_to_read; ++f) {
+                    // each frame illuminates every 64th column
+                    const size_t illum_col_idx =
+                        (frame_chunk + f - first) % columns_per_subcolumn;
+
+                    const auto *data =
+                        reinterpret_cast<const uint16_t *>(frames[f].data());
+
+                    for (size_t r = 0; r < half_rows; ++r) {
+                        const size_t column_offset =
+                            r * num_cols + illum_col_idx;
+                        for (size_t s = 0; s < num_subcolumns; ++s) {
+                            const size_t idx =
+                                column_offset + s * columns_per_subcolumn;
+                            sum[idx] += data[idx];
+                        }
+                    }
+                    // top half modules are flipped
+                    for (size_t r = half_rows; r < num_rows; ++r) {
+                        const size_t column_offset = r * num_cols +
+                                                     columns_per_subcolumn - 1 -
+                                                     illum_col_idx;
+                        for (size_t s = 0; s < num_subcolumns; ++s) {
+                            const size_t idx =
+                                column_offset + s * columns_per_subcolumn;
+                            sum[idx] += data[idx];
+                        }
+                    }
+                }
+            }
+
+            uint16_t *out = cs_data.data() + sample * num_rows * num_cols;
+            for (size_t idx = 0; idx < num_rows * num_cols; ++idx) {
+                out[idx] = static_cast<uint16_t>(static_cast<double>(sum[idx]) *
+                                                 inv_mean_over_num_frames);
+            }
+        }
+    };
+
+    std::vector<std::future<void>> futures;
+    futures.reserve(static_cast<size_t>(num_threads));
+
+    const size_t samples_per_thread = num_samples / num_threads;
+
+    for (int i = 0; i < num_threads; ++i) {
+        size_t end_sample =
+            (i == num_threads - 1) ? num_samples : (i + 1) * samples_per_thread;
+        futures.push_back(std::async(std::launch::async, worker,
+                                     i * samples_per_thread, end_sample));
+    }
+
+    for (auto &f : futures) {
+        try {
+            f.get();
+        } catch (const std::exception &e) {
+            throw std::runtime_error(std::string("Exception in thread: ") +
+                                     e.what());
+        }
+    }
+
+    return cs_data;
+}
+
+} // namespace aare

src/DataSetAggregation.test.cpp

@@ -54,4 +54,56 @@ TEST_CASE("mean_over_dataset", "[DataSetAggregation]") {
             }
         }
     }
-}
+}
+
+TEST_CASE("mean_current_source", "[DataSetAggregation]") {
+
+    constexpr size_t rows = 256;
+    constexpr size_t cols = 256;
+    constexpr size_t columns_per_subcolumn = 64;
+    constexpr size_t mean_over_num_frames = 3;
+    constexpr size_t num_samples = 5;
+    constexpr size_t file_offset = 2 * columns_per_subcolumn;
+    constexpr size_t frames_per_sample =
+        mean_over_num_frames * columns_per_subcolumn;
+    constexpr size_t num_frames = num_samples * frames_per_sample + file_offset;
+
+    // spread over several files to cross file boundaries mid sample
+    TemporaryJungfrauFiles temp_files(
+        num_frames,
+        [&](size_t frame) {
+            NDArray<uint16_t, 2> data({rows, cols});
+            int frame_idx = ((frame - file_offset) / columns_per_subcolumn) %
+                            mean_over_num_frames;
+
+            for (size_t r = 0; r < rows; ++r)
+                for (size_t c = 0; c < cols; ++c)
+                    data(r, c) = c + frame_idx;
+            return data;
+        },
+        3);
+
+    auto result =
+        mean_current_source(temp_files.file_path(), mean_over_num_frames,
+                            num_samples, 4, file_offset);
+
+    REQUIRE(result.shape(0) == num_samples);
+    REQUIRE(result.shape(1) == rows);
+    REQUIRE(result.shape(2) == cols);
+
+    size_t sum_mean_over_frames = (mean_over_num_frames) *
+                                  (mean_over_num_frames - 1) /
+                                  2; // sum of 0,1,...,mean_over_num_frames-1
+
+    for (size_t sample = 0; sample < num_samples; ++sample) {
+        for (size_t r = 0; r < rows; ++r) {
+            for (size_t c = 0; c < cols; ++c) {
+
+                CHECK(result(sample, r, c) ==
+                      static_cast<uint16_t>(
+                          (c * mean_over_num_frames + sum_mean_over_frames) /
+                          mean_over_num_frames));
+            }
+        }
+    }
+}