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

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

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<RotationIndexer>(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<MergedReflection> &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<std::vector<Reflection> > &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<std::vector<Reflection> > &IndexAndRefine::GetReflections() {
-    return reflections;
+    return ScalingResult(integration_outcome);
 }
 
 const std::vector<float> &IndexAndRefine::GetImageCC() const {
@@ -387,4 +391,12 @@ std::optional<UnitCell> IndexAndRefine::GetConsensusUnitCell() const {
     }
 
     return MeanUnitCell(accepted);
-}
+}
+
+std::vector<IntegrationOutcome> &IndexAndRefine::GetIntegrationOutcome() {
+    return integration_outcome;
+}
+
+const std::vector<IntegrationOutcome> &IndexAndRefine::GetIntegrationOutcome() const {
+    return integration_outcome;
+}

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<std::vector<Reflection>> reflections;
+    std::vector<IntegrationOutcome> integration_outcome;
     std::vector<float> mosaicity;
     std::vector<float> scale_cc;
     std::vector<std::optional<UnitCell> > unit_cells;
@@ -74,10 +75,10 @@ public:
     std::optional<UnitCell> GetConsensusUnitCell() const;
 
     // Not thread safe, need to be run after processing is all done
-    const std::vector<std::vector<Reflection>> &GetReflections() const;
-    std::vector<std::vector<Reflection>> &GetReflections();
     const std::vector<float> &GetImageCC() const;
     const std::vector<std::optional<UnitCell> > &GetUnitCells() const;
 
+    std::vector<IntegrationOutcome> &GetIntegrationOutcome();
+    const std::vector<IntegrationOutcome> &GetIntegrationOutcome() const;
 
 };

image_analysis/IntegrationOutcome.h

@@ -0,0 +1,20 @@
+// SPDX-FileCopyrightText: 2025 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
+// 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<Reflection> reflections;
+    std::optional<float> mosaicity_deg;
+    std::optional<float> image_scale_b_factor_Ang2;
+    std::optional<float> image_scale_cc;
+    std::optional<int64_t> image_scale_cc_n;
+    std::optional<float> image_scale_g;
+    std::optional<float> image_scale_wedge_deg;
+};

image_analysis/scale_merge/Merge.cpp

@@ -153,6 +153,22 @@ void MergeOnTheFly::AddAll(const std::vector<std::vector<Reflection> > &reflecti
     }
 }
 
+void MergeOnTheFly::AddAll(const std::vector<IntegrationOutcome> &integration_outcome,
+                           const std::vector<uint8_t> &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<MergedReflection> MergeOnTheFly::ExportReflections() {
     std::unique_lock ul(merged_mutex);
 
@@ -233,6 +249,14 @@ std::vector<MergedReflection> MergeAll(const DiffractionExperiment &x,
     return merge.ExportReflections();
 }
 
+std::vector<MergedReflection> MergeAll(const DiffractionExperiment &x,
+                                       const std::vector<IntegrationOutcome > &integration_outcome,
+                                       const std::vector<uint8_t> &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<MergedReflection> &merged,
+                           const std::vector<IntegrationOutcome > &integration_outcome,
+                           const UnitCell &cell,
+                           const std::vector<uint8_t> &merge_mask,
+                           const std::vector<MergedReflection> &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<uint64_t, float> 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<float>::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<ShellAccum> 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<MergedReflection> &merged,
                            const std::vector<std::vector<Reflection> > &reflections,

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<std::vector<Reflection> > &reflections,
                 const std::vector<uint8_t> &merge_mask = {});
 
+    void AddAll(const std::vector<IntegrationOutcome> &integration_outcome,
+        const std::vector<uint8_t> &merge_mask = {});
     void AddImage(const std::vector<Reflection> &v, double image_cc, const UnitCell &cell);
 
     std::vector<MergedReflection> ExportReflections();
@@ -95,9 +98,22 @@ std::vector<MergedReflection> MergeAll(const DiffractionExperiment &x,
                                        const std::vector<std::vector<Reflection> > &reflections,
                                        const std::vector<uint8_t> &merge_mask = {});
 
+
+std::vector<MergedReflection> MergeAll(const DiffractionExperiment &x,
+                                       const std::vector<IntegrationOutcome> &reflections,
+                                       const std::vector<uint8_t> &merge_mask = {});
+
 MergeStatistics MergeStats(const DiffractionExperiment &x,
                            const std::vector<MergedReflection> &merged,
                            const std::vector<std::vector<Reflection> > &reflections,
                            const UnitCell &cell,
                            const std::vector<uint8_t> &merge_mask = {},
                            const std::vector<MergedReflection> &reference = {});
+
+
+MergeStatistics MergeStats(const DiffractionExperiment &x,
+                           const std::vector<MergedReflection> &merged,
+                           const std::vector<IntegrationOutcome> &reflections,
+                           const UnitCell &cell,
+                           const std::vector<uint8_t> &merge_mask = {},
+                           const std::vector<MergedReflection> &reference = {});

image_analysis/scale_merge/ScaleOnTheFly.cpp

@@ -331,7 +331,7 @@ std::pair<double, size_t> ScaleOnTheFly::CalculateGlobalCC(const std::vector<Ref
 }
 
 ScaleOnTheFlyResult ScaleOnTheFly::Scale(std::vector<Reflection> &reflections,
-    std::optional<float> mosaicity_deg) {
+    std::optional<float> mosaicity_deg) const {
     auto start = std::chrono::steady_clock::now();
 
     ceres::Problem problem;
@@ -480,6 +480,54 @@ ScaleOnTheFlyResult ScaleOnTheFly::Scale(std::vector<Reflection> &reflections,
     return result;
 }
 
+void ScaleOnTheFly::Scale(std::vector<IntegrationOutcome> &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<std::future<void>> futures;
+        futures.reserve(local_nthreads);
+        std::atomic<size_t> 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<std::vector<Reflection> > &reflections,
                         const std::vector<float> &mosaicity,
                         size_t nthreads) {

image_analysis/scale_merge/ScaleOnTheFly.h

@@ -7,6 +7,7 @@
 #include "Merge.h"
 #include "../../common/DiffractionExperiment.h"
 #include "ScalingResult.h"
+#include "../IntegrationOutcome.h"
 
 #include <map>
 
@@ -37,10 +38,12 @@ class ScaleOnTheFly {
     [[nodiscard]] std::pair<double, size_t> CalculateGlobalCC(const std::vector<Reflection> &reflections) const;
 public:
     ScaleOnTheFly(const DiffractionExperiment &x, const std::vector<MergedReflection> &ref);
-    ScaleOnTheFlyResult Scale(std::vector<Reflection> &r, std::optional<float> mosaicity_deg);
+    ScaleOnTheFlyResult Scale(std::vector<Reflection> &r, std::optional<float> mosaicity_deg) const;
 
     ScalingResult Scale(std::vector<std::vector<Reflection> > &reflections,
                         const std::vector<float> &mosaicity,
                         size_t nthreads = 0);
+
+    void Scale(std::vector<IntegrationOutcome> &integration_outcome, size_t nthreads = 0) const;
 };
 

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<IntegrationOutcome> &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);

image_analysis/scale_merge/ScalingResult.h

@@ -5,6 +5,7 @@
 
 #include <vector>
 #include <string>
+#include "../IntegrationOutcome.h"
 
 struct ScalingResult {
     std::vector<float> image_scale_g;
@@ -14,5 +15,6 @@ struct ScalingResult {
     std::vector<float> image_cc;
     std::vector<int> image_cc_n;
     explicit ScalingResult(size_t n);
+    explicit ScalingResult(const std::vector<IntegrationOutcome> &v);
     void SaveToFile(const std::string &filename);
 };

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<double>(total_uncompressed_bytes) / (processing_time * 1e6);
     double frame_rate = static_cast<double>(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;