From f5447b347880bd95dafbdeed42a9d58a373c6dea Mon Sep 17 00:00:00 2001
From: leonarski_f <filip.leonarski@psi.ch>
Date: Sun, 17 May 2026 16:23:49 +0200
Subject: [PATCH] Update resolution of reflections based on consensus unit cell

---
 image_analysis/IndexAndRefine.cpp            |  4 ++++
 image_analysis/IndexAndRefine.h              |  6 +++++-
 image_analysis/scale_merge/ScaleOnTheFly.cpp | 18 ++++++++++++++++++
 image_analysis/scale_merge/ScaleOnTheFly.h   |  2 ++
 tools/jfjoch_process.cpp                     |  3 +++
 tools/jfjoch_scale.cpp                       |  9 +++++++++
 6 files changed, 41 insertions(+), 1 deletion(-)

diff --git a/image_analysis/IndexAndRefine.cpp b/image_analysis/IndexAndRefine.cpp
index 0f05f180..5f75ca5f 100644
--- a/image_analysis/IndexAndRefine.cpp
+++ b/image_analysis/IndexAndRefine.cpp
@@ -311,6 +311,10 @@ const std::vector<std::vector<Reflection> > &IndexAndRefine::GetReflections() co
     return reflections;
 }
 
+std::vector<std::vector<Reflection> > &IndexAndRefine::GetReflections() {
+    return reflections;
+}
+
 const std::vector<float> &IndexAndRefine::GetImageCC() const {
     return scale_cc;
 }
diff --git a/image_analysis/IndexAndRefine.h b/image_analysis/IndexAndRefine.h
index 216f25cb..f70f8b58 100644
--- a/image_analysis/IndexAndRefine.h
+++ b/image_analysis/IndexAndRefine.h
@@ -71,9 +71,13 @@ public:
 
     std::optional<RotationIndexerResult> Finalize();
 
+    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;
-    std::optional<UnitCell> GetConsensusUnitCell() const;
     const std::vector<std::optional<UnitCell> > &GetUnitCells() const;
+
+
 };
diff --git a/image_analysis/scale_merge/ScaleOnTheFly.cpp b/image_analysis/scale_merge/ScaleOnTheFly.cpp
index cc3a9c64..02b23b8f 100644
--- a/image_analysis/scale_merge/ScaleOnTheFly.cpp
+++ b/image_analysis/scale_merge/ScaleOnTheFly.cpp
@@ -355,3 +355,21 @@ ScalingResult ScaleOnTheFly::Scale(std::vector<std::vector<Reflection> > &reflec
 
     return result;
 }
+
+void UpdateReflectionResolution(const UnitCell &cell, std::vector<std::vector<Reflection> > &reflections) {
+    CrystalLattice latt(cell);
+    const auto astar = latt.Astar();
+    const auto bstar = latt.Bstar();
+    const auto cstar = latt.Cstar();
+
+    for (auto &image: reflections) {
+        for (auto &r: image) {
+            Coord q = r.h * astar + r.k * bstar + r.l * cstar;
+            auto qlen = q.Length();
+            if (qlen > 1e6)
+                r.d = 1/qlen;
+            else
+                r.d = NAN;
+        }
+    }
+}
diff --git a/image_analysis/scale_merge/ScaleOnTheFly.h b/image_analysis/scale_merge/ScaleOnTheFly.h
index fd214dd3..b42c71e8 100644
--- a/image_analysis/scale_merge/ScaleOnTheFly.h
+++ b/image_analysis/scale_merge/ScaleOnTheFly.h
@@ -22,6 +22,8 @@ struct ScaleOnTheFlyResult {
     bool succesful = false;
 };
 
+void UpdateReflectionResolution(const UnitCell &cell, std::vector<std::vector<Reflection> > &reflections);
+
 class ScaleOnTheFly {
     constexpr static size_t MIN_REFLECTIONS = 20;
 
diff --git a/tools/jfjoch_process.cpp b/tools/jfjoch_process.cpp
index cb46cf3a..300d6ccf 100644
--- a/tools/jfjoch_process.cpp
+++ b/tools/jfjoch_process.cpp
@@ -606,6 +606,7 @@ int main(int argc, char **argv) {
     const auto consensus_cell = indexer.GetConsensusUnitCell();
     auto consensus_end_time = std::chrono::steady_clock::now();
     auto consensus_duration = std::chrono::duration<double>(consensus_end_time - consensus_start_time).count();
+
     if (consensus_cell) {
         logger.Info("Consensus unit cell found in {:.2f} ms", consensus_duration * 1e3);
         logger.Info("UC:         a={:.2f} b={:.2f} c={:.2f} alpha={:.2f} beta={:.2f} gamma={:.2f}",
@@ -615,6 +616,8 @@ int main(int argc, char **argv) {
         auto rejected_uc = CalcMergeMaskUnitCell(experiment, *consensus_cell, indexer.GetUnitCells(), merging_mask_uc);
         if (rejected_uc > 0)
             logger.Info("Rejected {} images for merging due to unit cell being too far from consensus", rejected_uc);
+        UpdateReflectionResolution(consensus_cell.value(), indexer.GetReflections());
+        logger.Info("Reflection resolution updated based on consensus unit cell");
     } else
         logger.Info("Consensus unit cell not found - calculation tool {:.2f} ms", consensus_duration * 1e3);
     end_msg.unit_cell = consensus_cell;
diff --git a/tools/jfjoch_scale.cpp b/tools/jfjoch_scale.cpp
index 3c3d9ed5..8e49fcdb 100644
--- a/tools/jfjoch_scale.cpp
+++ b/tools/jfjoch_scale.cpp
@@ -241,6 +241,15 @@ int main(int argc, char **argv) {
 
     logger.Info("Running scaling (mosaicity refinement) ...");
     auto reflections = reader.ReadReflections(start_image, end_image);
+
+    if (experiment.GetUnitCell()) {
+        UpdateReflectionResolution(experiment.GetUnitCell().value(), reflections);
+        logger.Info("Reflection resolution updated based on experiment unit cell");
+    } else {
+        logger.Error("Experiment unit cell not found, cannot update reflection resolution");
+        exit(EXIT_FAILURE);
+    }
+
     std::vector<float> mosaicity(end_image - start_image + 1);
     for (int i = 0; i < end_image - start_image + 1; i++) {
         mosaicity[i] = dataset->mosaicity_deg[start_image + i];