diff --git a/CMakeLists.txt b/CMakeLists.txt index a9972be3..d756da1b 100644 --- a/CMakeLists.txt +++ b/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 diff --git a/include/aare/DataSetAggregation.hpp b/include/aare/DataSetAggregation.hpp index 6ebcb3ca..d01566ab 100644 --- a/include/aare/DataSetAggregation.hpp +++ b/include/aare/DataSetAggregation.hpp @@ -83,19 +83,14 @@ NDArray 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(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 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 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 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(mean_over_num_frames); - - NDArray cs_data({static_cast(num_samples), - static_cast(num_rows), - static_cast(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 frame( - {static_cast(num_rows), static_cast(num_cols)}); - - std::vector 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(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(sum[idx] * inv_mean_over_num_frames); - } - } - }; - - std::vector> 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 \ No newline at end of file +} // namespace aare diff --git a/src/DataSetAggregation.cpp b/src/DataSetAggregation.cpp new file mode 100644 index 00000000..7a07ce78 --- /dev/null +++ b/src/DataSetAggregation.cpp @@ -0,0 +1,144 @@ +// SPDX-License-Identifier: MPL-2.0 +#include "aare/DataSetAggregation.hpp" +#include "aare/JungfrauDataFile.hpp" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace aare { + +// TODO: maybe O_DIRECT rea dis better to avoid the page cache + +NDArray 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(mean_over_num_frames); + + NDArray cs_data({static_cast(num_samples), + static_cast(num_rows), + static_cast(num_cols)}, + 0); + + auto worker = [&](const size_t sample_begin, const size_t sample_end) { + std::vector 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(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(static_cast(sum[idx]) * + inv_mean_over_num_frames); + } + } + }; + + std::vector> futures; + futures.reserve(static_cast(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 diff --git a/src/DataSetAggregation.test.cpp b/src/DataSetAggregation.test.cpp index 6a4d2f54..8122f55d 100644 --- a/src/DataSetAggregation.test.cpp +++ b/src/DataSetAggregation.test.cpp @@ -54,4 +54,56 @@ TEST_CASE("mean_over_dataset", "[DataSetAggregation]") { } } } -} \ No newline at end of file +} + +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 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( + (c * mean_over_num_frames + sum_mean_over_frames) / + mean_over_num_frames)); + } + } + } +}