From c33366f3d6bd697ec1de270c3ef2c97f2c4780ab Mon Sep 17 00:00:00 2001
From: Filip Leonarski <filip.leonarski@psi.ch>
Date: Fri, 20 Feb 2026 13:18:16 +0100
Subject: [PATCH] IndexAndRefine: Use mosaicity from indexing for scaling

---
 image_analysis/IndexAndRefine.cpp            |  9 ++++++--
 image_analysis/IndexAndRefine.h              |  1 +
 image_analysis/scale_merge/ScaleAndMerge.cpp | 24 ++++++++------------
 3 files changed, 18 insertions(+), 16 deletions(-)

diff --git a/image_analysis/IndexAndRefine.cpp b/image_analysis/IndexAndRefine.cpp
index e5b20981..6311c596 100644
--- a/image_analysis/IndexAndRefine.cpp
+++ b/image_analysis/IndexAndRefine.cpp
@@ -18,6 +18,7 @@ IndexAndRefine::IndexAndRefine(const DiffractionExperiment &x, IndexerThreadPool
     if (indexer && x.IsRotationIndexing())
         rotation_indexer = std::make_unique<RotationIndexer>(x, *indexer);
     reflections.resize(x.GetImageNum());
+    mosaicity.resize(x.GetImageNum(), NAN);
 }
 
 IndexAndRefine::IndexingOutcome IndexAndRefine::DetermineLatticeAndSymmetry(DataMessage &msg) {
@@ -167,8 +168,10 @@ void IndexAndRefine::QuickPredictAndIntegrate(DataMessage &msg,
     if (experiment.GetGoniometer().has_value()) {
         wedge_deg = experiment.GetGoniometer()->GetWedge_deg() / 2.0;
 
-        if (msg.mosaicity_deg)
+        if (msg.mosaicity_deg) {
             mos_deg = msg.mosaicity_deg.value();
+            mosaicity[msg.number] = mos_deg;
+        }
     }
 
     const BraggPredictionSettings settings_prediction{
@@ -255,8 +258,10 @@ std::optional<ScaleMergeResult> IndexAndRefine::ScaleRotationData(const ScaleMer
     ScaleMergeOptions options = opts;
 
     // If the experiment provides a wedge, propagate it
-    if (experiment.GetGoniometer().has_value())
+    if (experiment.GetGoniometer().has_value()) {
         options.wedge_deg = experiment.GetGoniometer()->GetWedge_deg();
+        options.mosaicity_init_deg_vec = mosaicity;
+    }
 
     // If caller left space_group unset, try to pick it from the indexed lattice
     if (!options.space_group.has_value()) {
diff --git a/image_analysis/IndexAndRefine.h b/image_analysis/IndexAndRefine.h
index a04592b8..4b0f893b 100644
--- a/image_analysis/IndexAndRefine.h
+++ b/image_analysis/IndexAndRefine.h
@@ -46,6 +46,7 @@ class IndexAndRefine {
 
     mutable std::mutex reflections_mutex;
     std::vector<std::vector<Reflection>> reflections;
+    std::vector<double> mosaicity;
 
     IndexingOutcome DetermineLatticeAndSymmetry(DataMessage &msg);
     void RefineGeometryIfNeeded(DataMessage &msg, IndexingOutcome &outcome);
diff --git a/image_analysis/scale_merge/ScaleAndMerge.cpp b/image_analysis/scale_merge/ScaleAndMerge.cpp
index 33dee5fb..74f8c667 100644
--- a/image_analysis/scale_merge/ScaleAndMerge.cpp
+++ b/image_analysis/scale_merge/ScaleAndMerge.cpp
@@ -214,8 +214,6 @@ ScaleMergeResult ScaleAndMergeReflectionsCeres(const std::vector<std::vector<Ref
 
     size_t n_image_slots = observations.size() / opt.image_cluster + (observations.size() % opt.image_cluster > 0 ? 1 : 0);
 
-    std::vector<double> g(n_image_slots, 1.0);
-    std::vector<double> mosaicity(n_image_slots, opt.mosaicity_init_deg);
     std::vector<uint8_t> image_slot_used(n_image_slots, 0);
 
     for (int i = 0; i < observations.size(); i++) {
@@ -262,17 +260,20 @@ ScaleMergeResult ScaleAndMergeReflectionsCeres(const std::vector<std::vector<Ref
         }
     }
 
-    const int nhkl = static_cast<int>(hklToSlot.size());
-
-    std::vector<double> Itrue(nhkl, 0.0);
-
+    std::vector<double> g(n_image_slots, 1.0);
+    std::vector<double> mosaicity(n_image_slots, opt.mosaicity_init_deg);
     for (int i = 0; i < n_image_slots; i++) {
         if (!image_slot_used[i]) {
             mosaicity[i] = NAN;
             g[i] = NAN;
+        } else if (opt.mosaicity_init_deg_vec.size() > i && std::isfinite(opt.mosaicity_init_deg_vec[i])) {
+            mosaicity[i] = opt.mosaicity_init_deg_vec[i];
         }
     }
 
+    const int nhkl = static_cast<int>(hklToSlot.size());
+    std::vector<double> Itrue(nhkl, 0.0);
+
     // Initialize Itrue from per-HKL median of observed intensities
     {
         std::vector<std::vector<double> > per_hkl_I(nhkl);
@@ -305,6 +306,7 @@ ScaleMergeResult ScaleAndMergeReflectionsCeres(const std::vector<std::vector<Ref
                                  &Itrue[o.hkl_slot]);
         is_valid_hkl_slot[o.hkl_slot] = true;
     }
+
     for (int i = 0; i < g.size(); ++i) {
         if (image_slot_used[i]) {
             auto *cost = new ceres::AutoDiffCostFunction<ScaleRegularizationResidual, 1, 1>(
@@ -320,10 +322,7 @@ ScaleMergeResult ScaleAndMergeReflectionsCeres(const std::vector<std::vector<Ref
                 auto *cost = new ceres::AutoDiffCostFunction<SmoothnessRegularizationResidual, 1, 1, 1, 1>(
                     new SmoothnessRegularizationResidual(0.05));
 
-                problem.AddResidualBlock(cost, nullptr,
-                                         &g[i],
-                                         &g[i + 1],
-                                         &g[i + 2]);
+                problem.AddResidualBlock(cost, nullptr, &g[i], &g[i + 1], &g[i + 2]);
             }
         }
     }
@@ -334,10 +333,7 @@ ScaleMergeResult ScaleAndMergeReflectionsCeres(const std::vector<std::vector<Ref
                 auto *cost = new ceres::AutoDiffCostFunction<SmoothnessRegularizationResidual, 1, 1, 1, 1>(
                     new SmoothnessRegularizationResidual(0.05));
 
-                problem.AddResidualBlock(cost, nullptr,
-                                         &mosaicity[i],
-                                         &mosaicity[i + 1],
-                                         &mosaicity[i + 2]);
+                problem.AddResidualBlock(cost, nullptr, &mosaicity[i], &mosaicity[i + 1], &mosaicity[i + 2]);
             }
         }
     }