From a85eb4b1373d39be88300bb023b79155d97516ef Mon Sep 17 00:00:00 2001 From: leonarski_f Date: Mon, 25 May 2026 10:09:49 +0200 Subject: [PATCH] Scaling/Merging: Work in progress to improve the information that is fed to scaling. This is a mess at the moment - with redundant data structures, etc. - will clean this up, later. --- image_analysis/CMakeLists.txt | 3 +- image_analysis/IndexAndRefine.cpp | 36 ++-- image_analysis/IndexAndRefine.h | 7 +- image_analysis/IntegrationOutcome.h | 20 ++ image_analysis/scale_merge/Merge.cpp | 191 +++++++++++++++++++ image_analysis/scale_merge/Merge.h | 16 ++ image_analysis/scale_merge/ScaleOnTheFly.cpp | 50 ++++- image_analysis/scale_merge/ScaleOnTheFly.h | 5 +- image_analysis/scale_merge/ScalingResult.cpp | 17 ++ image_analysis/scale_merge/ScalingResult.h | 2 + tools/jfjoch_process.cpp | 7 +- 11 files changed, 332 insertions(+), 22 deletions(-) create mode 100644 image_analysis/IntegrationOutcome.h diff --git a/image_analysis/CMakeLists.txt b/image_analysis/CMakeLists.txt index 66537115..f6a66733 100644 --- a/image_analysis/CMakeLists.txt +++ b/image_analysis/CMakeLists.txt @@ -34,7 +34,8 @@ ADD_LIBRARY(JFJochImageAnalysis STATIC LoadFCalcFromMtz.cpp LoadFCalcFromMtz.h UpdateReflectionResolution.cpp - UpdateReflectionResolution.h) + UpdateReflectionResolution.h + IntegrationOutcome.h) FIND_PACKAGE(Eigen3 3.4 REQUIRED NO_MODULE) # provides Eigen3::Eigen diff --git a/image_analysis/IndexAndRefine.cpp b/image_analysis/IndexAndRefine.cpp index 5f75ca5f..f9008b6c 100644 --- a/image_analysis/IndexAndRefine.cpp +++ b/image_analysis/IndexAndRefine.cpp @@ -18,7 +18,8 @@ IndexAndRefine::IndexAndRefine(const DiffractionExperiment &x, IndexerThreadPool indexer_(indexer) { if (indexer && x.IsRotationIndexing()) rotation_indexer = std::make_unique(x, *indexer); - reflections.resize(x.GetImageNum()); + integration_outcome.resize(x.GetImageNum()); + mosaicity.resize(x.GetImageNum(), NAN); scale_cc.resize(x.GetImageNum(), 0); unit_cells.resize(x.GetImageNum()); @@ -233,7 +234,14 @@ void IndexAndRefine::QuickPredictAndIntegrate(DataMessage &msg, { std::unique_lock ul(reflections_mutex); - reflections[msg.number] = msg.reflections; // Image is not processed twice, so thread-safe in principle, but better safe than sorry :) + integration_outcome[msg.number] = IntegrationOutcome{ + .geom = outcome.experiment.GetDiffractionGeometry(), + .latt = latt, + .reflections = msg.reflections, + .mosaicity_deg = msg.image_scale_mosaicity, + .image_scale_b_factor_Ang2 = msg.image_scale_b_factor, + .image_scale_cc = msg.image_scale_cc, + }; } } @@ -302,17 +310,13 @@ void IndexAndRefine::ScaleImage(DataMessage &msg) { ScalingResult IndexAndRefine::ScaleAllImages(const std::vector &reference, size_t nthreads) { ScaleOnTheFly scaling(experiment, reference); - auto result = scaling.Scale(reflections, mosaicity, nthreads); - scale_cc = result.image_cc; - return result; -} + scaling.Scale(integration_outcome, nthreads); + scale_cc.resize(integration_outcome.size()); -const std::vector > &IndexAndRefine::GetReflections() const { - return reflections; -} + for (int i = 0; i < integration_outcome.size(); i++) + scale_cc.at(i) = integration_outcome[i].image_scale_cc.value_or(NAN); -std::vector > &IndexAndRefine::GetReflections() { - return reflections; + return ScalingResult(integration_outcome); } const std::vector &IndexAndRefine::GetImageCC() const { @@ -387,4 +391,12 @@ std::optional IndexAndRefine::GetConsensusUnitCell() const { } return MeanUnitCell(accepted); -} \ No newline at end of file +} + +std::vector &IndexAndRefine::GetIntegrationOutcome() { + return integration_outcome; +} + +const std::vector &IndexAndRefine::GetIntegrationOutcome() const { + return integration_outcome; +} diff --git a/image_analysis/IndexAndRefine.h b/image_analysis/IndexAndRefine.h index f70f8b58..4e8e1615 100644 --- a/image_analysis/IndexAndRefine.h +++ b/image_analysis/IndexAndRefine.h @@ -16,6 +16,7 @@ #include "RotationParameters.h" #include "scale_merge/ScaleOnTheFly.h" #include "scale_merge/ScalingResult.h" +#include "IntegrationOutcome.h" class IndexAndRefine { const bool index_ice_rings; @@ -46,7 +47,7 @@ class IndexAndRefine { }; mutable std::mutex reflections_mutex; - std::vector> reflections; + std::vector integration_outcome; std::vector mosaicity; std::vector scale_cc; std::vector > unit_cells; @@ -74,10 +75,10 @@ public: std::optional GetConsensusUnitCell() const; // Not thread safe, need to be run after processing is all done - const std::vector> &GetReflections() const; - std::vector> &GetReflections(); const std::vector &GetImageCC() const; const std::vector > &GetUnitCells() const; + std::vector &GetIntegrationOutcome(); + const std::vector &GetIntegrationOutcome() const; }; diff --git a/image_analysis/IntegrationOutcome.h b/image_analysis/IntegrationOutcome.h new file mode 100644 index 00000000..3b9cbec6 --- /dev/null +++ b/image_analysis/IntegrationOutcome.h @@ -0,0 +1,20 @@ +// SPDX-FileCopyrightText: 2025 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#pragma once + +#include "../common/Reflection.h" +#include "../common/CrystalLattice.h" +#include "../common/DiffractionGeometry.h" + +struct IntegrationOutcome { + DiffractionGeometry geom; + CrystalLattice latt; + std::vector reflections; + std::optional mosaicity_deg; + std::optional image_scale_b_factor_Ang2; + std::optional image_scale_cc; + std::optional image_scale_cc_n; + std::optional image_scale_g; + std::optional image_scale_wedge_deg; +}; \ No newline at end of file diff --git a/image_analysis/scale_merge/Merge.cpp b/image_analysis/scale_merge/Merge.cpp index 4a7c5fba..418383e4 100644 --- a/image_analysis/scale_merge/Merge.cpp +++ b/image_analysis/scale_merge/Merge.cpp @@ -153,6 +153,22 @@ void MergeOnTheFly::AddAll(const std::vector > &reflecti } } +void MergeOnTheFly::AddAll(const std::vector &integration_outcome, + const std::vector &merge_mask) { + if (!merge_mask.empty() && merge_mask.size() != integration_outcome.size()) + throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Merge mask size mismatch"); + + std::unique_lock ul(merged_mutex); + for (int i = 0; i < integration_outcome.size(); ++i) { + const int half = half_dist(rng); + if (!merge_mask.empty() && merge_mask[i] == 0) + continue; + if (integration_outcome[i].reflections.empty()) + continue; + ProcessImage_i(integration_outcome[i].reflections, half); + } +} + std::vector MergeOnTheFly::ExportReflections() { std::unique_lock ul(merged_mutex); @@ -233,6 +249,14 @@ std::vector MergeAll(const DiffractionExperiment &x, return merge.ExportReflections(); } +std::vector MergeAll(const DiffractionExperiment &x, + const std::vector &integration_outcome, + const std::vector &merge_mask) { + MergeOnTheFly merge(x); + merge.AddAll(integration_outcome, merge_mask); + return merge.ExportReflections(); +} + struct ShellAccum { int total_obs = 0; int unique = 0; @@ -276,6 +300,173 @@ void CalcPossibleReflections(const DiffractionExperiment &x, } } + +MergeStatistics MergeStats(const DiffractionExperiment &x, + const std::vector &merged, + const std::vector &integration_outcome, + const UnitCell &cell, + const std::vector &merge_mask, + const std::vector &reference) { + if (!merge_mask.empty() && merge_mask.size() != integration_outcome.size()) + throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Merge mask size mismatch"); + + constexpr int n_shells = 10; + + auto min_partiality = x.GetScalingSettings().GetMinPartiality(); + auto d_min_limit_A = x.GetScalingSettings().GetHighResolutionLimit_A(); + auto scaling_settings = x.GetScalingSettings(); + HKLKeyGenerator key_generator(scaling_settings.GetMergeFriedel(), x.GetSpaceGroupNumber().value_or(1)); + + std::unordered_map reference_intensities; + if (!reference.empty()) { + reference_intensities.reserve(reference.size()); + for (const auto &r: reference) { + if (!std::isfinite(r.I)) + continue; + + const auto hkl = key_generator(r); + reference_intensities[hkl.pack()] = r.I; + } + } + + float d_min = std::numeric_limits::max(); + float d_max = 0.0f; + + for (const auto &m: merged) { + if (!std::isfinite(m.d) || m.d <= 0.0f) + continue; + if (d_min_limit_A && m.d < d_min_limit_A) + continue; + + d_min = std::min(d_min, m.d); + d_max = std::max(d_max, m.d); + } + + if (!(d_min < d_max && d_min > 0.0f)) + throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, + "MergeStats: Error in resolution calculation"); + + const float d_min_pad = d_min * 0.999f; + const float d_max_pad = d_max * 1.001f; + + ResolutionShells shells(d_min_pad, d_max_pad, n_shells); + const auto shell_mean_1_d2 = shells.GetShellMeanOneOverResSq(); + const auto shell_min_res = shells.GetShellMinRes(); + + std::vector acc(n_shells); + + CalcPossibleReflections(x, cell, d_min_pad, d_max_pad, shells, acc); + + CorrelationCoefficient cc_half_overall; + CorrelationCoefficient cc_ref_overall; + + for (const auto &m: merged) { + const auto shell = shells.GetShell(m.d); + if (!shell.has_value()) + continue; + + const int s = *shell; + if (s >= 0 && s < n_shells) { + if (std::isfinite(m.I) && std::isfinite(m.sigma) && m.sigma > 0.0) { + acc[s].unique++; + acc[s].sum_i_over_sigma += m.I / m.sigma; + ++acc[s].n_i_over_sigma; + + if (!reference_intensities.empty()) { + const auto hkl = key_generator(m); + const auto ref_it = reference_intensities.find(hkl.pack()); + if (ref_it != reference_intensities.end() && std::isfinite(ref_it->second)) { + acc[s].cc_ref.Add(m.I, ref_it->second); + cc_ref_overall.Add(m.I, ref_it->second); + } + } + + if (std::isfinite(m.I_half[0]) && std::isfinite(m.I_half[1])) { + acc[s].cc_half.Add(m.I_half[0], m.I_half[1]); + cc_half_overall.Add(m.I_half[0], m.I_half[1]); + } + + } + } + } + + for (int i = 0; i < integration_outcome.size(); ++i) { + if (!merge_mask.empty() && merge_mask[i] == 0) + continue; + + for (const auto &r: integration_outcome[i].reflections) { + if (key_generator.IsSystematicallyAbsent(r)) + continue; + if (r.image_scale_corr <= 0.0 || !std::isfinite(r.image_scale_corr)) + continue; + if (!AcceptReflection(r, d_min_limit_A)) + continue; + if (r.partiality < min_partiality) + continue; + + const float I_corr = r.I * r.image_scale_corr; + const float sigma_corr = r.sigma * r.image_scale_corr; + if (!std::isfinite(I_corr) || !std::isfinite(sigma_corr) || sigma_corr <= 0.0f) + continue; + + const auto shell = shells.GetShell(r.d); + if (!shell.has_value()) + continue; + const int s = *shell; + if (s >= 0 && s < n_shells) + acc[s].total_obs++; + } + } + + MergeStatistics out; + out.shells.resize(n_shells); + + for (int s = 0; s < n_shells; ++s) { + const auto &sa = acc[s]; + auto &ss = out.shells[s]; + + ss.mean_one_over_d2 = shell_mean_1_d2[s]; + ss.d_min = shell_min_res[s]; + ss.d_max = s == 0 ? d_max_pad : shell_min_res[s - 1]; + ss.total_observations = sa.total_obs; + ss.unique_reflections = sa.unique; + ss.possible_unique_reflections = sa.possible; + ss.mean_i_over_sigma = sa.n_i_over_sigma > 0 + ? sa.sum_i_over_sigma / sa.n_i_over_sigma + : 0.0; + + ss.cc_half = sa.cc_half.GetCC(); + ss.cc_ref = sa.cc_ref.GetCC(); + } + + auto &overall = out.overall; + overall.d_min = d_min; + overall.d_max = d_max; + + int all_possible = 0; + int all_unique = 0; + double sum_i_over_sigma = 0.0; + int n_i_over_sigma = 0; + + + for (const auto &sa: acc) { + overall.total_observations += sa.total_obs; + all_unique += sa.unique; + all_possible += sa.possible; + sum_i_over_sigma += sa.sum_i_over_sigma; + n_i_over_sigma += sa.n_i_over_sigma; + + } + + overall.possible_unique_reflections = all_possible; + overall.unique_reflections = all_unique; + overall.mean_i_over_sigma = n_i_over_sigma > 0 ? sum_i_over_sigma / n_i_over_sigma : 0.0; + overall.cc_half = cc_half_overall.GetCC(); + overall.cc_ref = cc_ref_overall.GetCC(); + + return out; +} + MergeStatistics MergeStats(const DiffractionExperiment &x, const std::vector &merged, const std::vector > &reflections, diff --git a/image_analysis/scale_merge/Merge.h b/image_analysis/scale_merge/Merge.h index 31a05da3..0297767a 100644 --- a/image_analysis/scale_merge/Merge.h +++ b/image_analysis/scale_merge/Merge.h @@ -10,6 +10,7 @@ #include "../../common/Logger.h" #include "../../common/DiffractionExperiment.h" #include "../../common/Reflection.h" +#include "../IntegrationOutcome.h" #include "HKLKey.h" @@ -77,6 +78,8 @@ public: void AddAll(const std::vector > &reflections, const std::vector &merge_mask = {}); + void AddAll(const std::vector &integration_outcome, + const std::vector &merge_mask = {}); void AddImage(const std::vector &v, double image_cc, const UnitCell &cell); std::vector ExportReflections(); @@ -95,9 +98,22 @@ std::vector MergeAll(const DiffractionExperiment &x, const std::vector > &reflections, const std::vector &merge_mask = {}); + +std::vector MergeAll(const DiffractionExperiment &x, + const std::vector &reflections, + const std::vector &merge_mask = {}); + MergeStatistics MergeStats(const DiffractionExperiment &x, const std::vector &merged, const std::vector > &reflections, const UnitCell &cell, const std::vector &merge_mask = {}, const std::vector &reference = {}); + + +MergeStatistics MergeStats(const DiffractionExperiment &x, + const std::vector &merged, + const std::vector &reflections, + const UnitCell &cell, + const std::vector &merge_mask = {}, + const std::vector &reference = {}); diff --git a/image_analysis/scale_merge/ScaleOnTheFly.cpp b/image_analysis/scale_merge/ScaleOnTheFly.cpp index d49bbcd9..dbd94278 100644 --- a/image_analysis/scale_merge/ScaleOnTheFly.cpp +++ b/image_analysis/scale_merge/ScaleOnTheFly.cpp @@ -331,7 +331,7 @@ std::pair ScaleOnTheFly::CalculateGlobalCC(const std::vector &reflections, - std::optional mosaicity_deg) { + std::optional mosaicity_deg) const { auto start = std::chrono::steady_clock::now(); ceres::Problem problem; @@ -480,6 +480,54 @@ ScaleOnTheFlyResult ScaleOnTheFly::Scale(std::vector &reflections, return result; } +void ScaleOnTheFly::Scale(std::vector &integration, size_t nthreads) const { + ScalingResult result(integration.size()); + + if (nthreads == 0) + nthreads = std::thread::hardware_concurrency(); + + if (nthreads <= 1) { + for (int i = 0; i < integration.size(); i++) { + if (integration[i].reflections.empty()) + continue; + + auto local_result = Scale(integration[i].reflections, integration[i].mosaicity_deg); + integration[i].mosaicity_deg = local_result.mos; + integration[i].image_scale_b_factor_Ang2 = local_result.B; + integration[i].image_scale_g = local_result.G; + integration[i].image_scale_wedge_deg = local_result.wedge; + integration[i].image_scale_cc = local_result.cc; + integration[i].image_scale_cc_n = local_result.cc_n; + } + } else { + auto local_nthreads = std::min(nthreads, integration.size()); + std::vector> futures; + futures.reserve(local_nthreads); + std::atomic curr_image = 0; + + for (size_t t = 0; t < local_nthreads; ++t) + futures.emplace_back(std::async(std::launch::async, [&] { + size_t i = curr_image.fetch_add(1); + while (i < integration.size()) { + if (integration[i].reflections.empty()) + continue; + + auto local_result = Scale(integration[i].reflections, integration[i].mosaicity_deg); + integration[i].mosaicity_deg = local_result.mos; + integration[i].image_scale_b_factor_Ang2 = local_result.B; + integration[i].image_scale_g = local_result.G; + integration[i].image_scale_wedge_deg = local_result.wedge; + integration[i].image_scale_cc = local_result.cc; + integration[i].image_scale_cc_n = local_result.cc_n; + i = curr_image.fetch_add(1); + } + })); + + for (auto &f: futures) + f.get(); + } +} + ScalingResult ScaleOnTheFly::Scale(std::vector > &reflections, const std::vector &mosaicity, size_t nthreads) { diff --git a/image_analysis/scale_merge/ScaleOnTheFly.h b/image_analysis/scale_merge/ScaleOnTheFly.h index b825f054..2216e11f 100644 --- a/image_analysis/scale_merge/ScaleOnTheFly.h +++ b/image_analysis/scale_merge/ScaleOnTheFly.h @@ -7,6 +7,7 @@ #include "Merge.h" #include "../../common/DiffractionExperiment.h" #include "ScalingResult.h" +#include "../IntegrationOutcome.h" #include @@ -37,10 +38,12 @@ class ScaleOnTheFly { [[nodiscard]] std::pair CalculateGlobalCC(const std::vector &reflections) const; public: ScaleOnTheFly(const DiffractionExperiment &x, const std::vector &ref); - ScaleOnTheFlyResult Scale(std::vector &r, std::optional mosaicity_deg); + ScaleOnTheFlyResult Scale(std::vector &r, std::optional mosaicity_deg) const; ScalingResult Scale(std::vector > &reflections, const std::vector &mosaicity, size_t nthreads = 0); + + void Scale(std::vector &integration_outcome, size_t nthreads = 0) const; }; diff --git a/image_analysis/scale_merge/ScalingResult.cpp b/image_analysis/scale_merge/ScalingResult.cpp index a5dc4aae..318948b9 100644 --- a/image_analysis/scale_merge/ScalingResult.cpp +++ b/image_analysis/scale_merge/ScalingResult.cpp @@ -15,6 +15,23 @@ ScalingResult::ScalingResult(size_t n) image_cc(n, NAN), image_cc_n(n, 0) {} +ScalingResult::ScalingResult(const std::vector &v) +: image_scale_g(v.size(), NAN), + mosaicity_deg(v.size(), NAN), + image_bfactor_Ang2(v.size(), NAN), + rotation_wedge_deg(v.size(), NAN), + image_cc(v.size(), NAN), + image_cc_n(v.size(), 0) { + for (int i = 0; i < v.size(); i++) { + image_scale_g[i] = v[i].image_scale_g.value_or(NAN); + mosaicity_deg[i] = v[i].mosaicity_deg.value_or(NAN); + image_bfactor_Ang2[i] = v[i].image_scale_b_factor_Ang2.value_or(NAN); + rotation_wedge_deg[i] = v[i].image_scale_wedge_deg.value_or(NAN); + image_cc[i] = v[i].image_scale_cc.value_or(NAN); + image_cc_n[i] = v[i].image_scale_cc_n.value_or(0); + } +} + void ScalingResult::SaveToFile(const std::string &filename) { const std::string img_path = filename + "_image.dat"; std::ofstream img_file(img_path, std::ofstream::out | std::ofstream::trunc); diff --git a/image_analysis/scale_merge/ScalingResult.h b/image_analysis/scale_merge/ScalingResult.h index 91bc8c2a..7003b5a8 100644 --- a/image_analysis/scale_merge/ScalingResult.h +++ b/image_analysis/scale_merge/ScalingResult.h @@ -5,6 +5,7 @@ #include #include +#include "../IntegrationOutcome.h" struct ScalingResult { std::vector image_scale_g; @@ -14,5 +15,6 @@ struct ScalingResult { std::vector image_cc; std::vector image_cc_n; explicit ScalingResult(size_t n); + explicit ScalingResult(const std::vector &v); void SaveToFile(const std::string &filename); }; diff --git a/tools/jfjoch_process.cpp b/tools/jfjoch_process.cpp index a00acd29..0a8d72b9 100644 --- a/tools/jfjoch_process.cpp +++ b/tools/jfjoch_process.cpp @@ -718,7 +718,7 @@ int main(int argc, char **argv) { auto scale_start = std::chrono::steady_clock::now(); for (int i = 0; i < scaling_iter; i++) { auto iter_start = std::chrono::steady_clock::now(); - auto merge_result = MergeAll(experiment, indexer.GetReflections(), merging_mask_uc); + auto merge_result = MergeAll(experiment, indexer.GetIntegrationOutcome(), merging_mask_uc); scale_result = indexer.ScaleAllImages(merge_result); scale_result.SaveToFile(output_prefix + "_iter" + std::to_string(i) + "_scale.dat"); auto iter_end = std::chrono::steady_clock::now(); @@ -744,8 +744,8 @@ int main(int argc, char **argv) { if (rejected_cc > 0) logger.Info("Rejected {} images for merging due to low CC with reference", rejected_cc); - auto merged_reflections = MergeAll(experiment, indexer.GetReflections(), merging_mask_uc); - auto merged_statistics = MergeStats(experiment, merged_reflections, indexer.GetReflections(), *consensus_cell, merging_mask_uc, + auto merged_reflections = MergeAll(experiment, indexer.GetIntegrationOutcome(), merging_mask_uc); + auto merged_statistics = MergeStats(experiment, merged_reflections, indexer.GetIntegrationOutcome(), *consensus_cell, merging_mask_uc, reference_data); auto merge_end = std::chrono::steady_clock::now(); @@ -789,7 +789,6 @@ int main(int argc, char **argv) { double throughput_MBs = static_cast(total_uncompressed_bytes) / (processing_time * 1e6); double frame_rate = static_cast(images_to_process) / processing_time; - std::cout << fmt::format("Processing time: {:.2f} s", processing_time) << std::endl; std::cout << fmt::format("Frame rate: {:.2f} Hz", frame_rate) << std::endl; std::cout << fmt::format("Total throughput:{:.2f} MB/s", throughput_MBs) << std::endl;