From 86afb7697e7ba3be84afaf7f7486d1dec5730b72 Mon Sep 17 00:00:00 2001
From: Filip Leonarski <filip.leonarski@psi.ch>
Date: Wed, 11 Feb 2026 16:20:14 +0100
Subject: [PATCH] ScaleAndMerge: Add log-scaling residual

---
 image_analysis/scale_merge/ScaleAndMerge.cpp | 96 +++++++++++++++++---
 image_analysis/scale_merge/ScaleAndMerge.h   |  1 +
 2 files changed, 83 insertions(+), 14 deletions(-)

diff --git a/image_analysis/scale_merge/ScaleAndMerge.cpp b/image_analysis/scale_merge/ScaleAndMerge.cpp
index 5f20d674..42a5a525 100644
--- a/image_analysis/scale_merge/ScaleAndMerge.cpp
+++ b/image_analysis/scale_merge/ScaleAndMerge.cpp
@@ -114,6 +114,58 @@ inline HKLKey CanonicalizeHKLKey(const Reflection& r, const ScaleMergeOptions& o
     return key;
 }
 
+/// CrystFEL-like log-scaling residual
+///
+///   residual = w * [ ln(I_obs) - ln(G) - ln(partiality) - ln(lp) - ln(I_true) ]
+///
+/// Only observations with I_obs > 0 should be fed in (the caller skips the rest).
+/// G and I_true are constrained to be positive via Ceres lower bounds.
+struct LogIntensityResidual {
+    LogIntensityResidual(const Reflection& r, double sigma_obs, double wedge_deg, bool refine_partiality)
+        : log_Iobs_(std::log(std::max(static_cast<double>(r.I), 1e-30))),
+          weight_(SafeInv(sigma_obs / std::max(static_cast<double>(r.I), 1e-30), 1.0)),
+          delta_phi_(r.delta_phi_deg),
+          log_lp_(std::log(std::max(SafeInv(static_cast<double>(r.rlp), 1.0), 1e-30))),
+          half_wedge_(wedge_deg / 2.0),
+          c1_(r.zeta / std::sqrt(2.0)),
+          partiality_(r.partiality),
+          refine_partiality_(refine_partiality) {}
+
+    template<typename T>
+    bool operator()(const T* const G,
+                    const T* const mosaicity,
+                    const T* const Itrue,
+                    T* residual) const {
+        T partiality;
+        if (refine_partiality_ && mosaicity[0] != 0.0) {
+            const T arg_plus  = T(delta_phi_ + half_wedge_) * T(c1_) / mosaicity[0];
+            const T arg_minus = T(delta_phi_ - half_wedge_) * T(c1_) / mosaicity[0];
+            partiality = (ceres::erf(arg_plus) - ceres::erf(arg_minus)) / T(2.0);
+        } else {
+            partiality = T(partiality_);
+        }
+
+        // Clamp partiality away from zero so log is safe
+        const T min_p = T(1e-30);
+        if (partiality < min_p)
+            partiality = min_p;
+
+        // ln(I_pred) = ln(G) + ln(partiality) + ln(lp) + ln(Itrue)
+        const T log_Ipred = ceres::log(G[0]) + ceres::log(partiality) + T(log_lp_) + ceres::log(Itrue[0]);
+        residual[0] = (log_Ipred - T(log_Iobs_)) * T(weight_);
+        return true;
+    }
+
+    double log_Iobs_;
+    double weight_;       // w_i ≈ I_obs / sigma_obs  (relative weight in log-space)
+    double delta_phi_;
+    double log_lp_;
+    double half_wedge_;
+    double c1_;
+    double partiality_;
+    bool refine_partiality_;
+};
+    
 struct IntensityResidual {
     IntensityResidual(const Reflection& r, double sigma_obs, double wedge_deg, bool refine_partiality)
         : Iobs_(static_cast<double>(r.I)),
@@ -280,24 +332,40 @@ ScaleMergeResult ScaleAndMergeReflectionsCeres(const std::vector<Reflection>& ob
     const bool refine_partiality = opt.wedge_deg > 0.0;
 
     for (const auto& o : obs) {
-        auto* cost = new ceres::AutoDiffCostFunction<IntensityResidual, 1, 1, 1, 1>(
-            new IntensityResidual(*o.r, o.sigma, opt.wedge_deg, refine_partiality));
+        if (opt.log_scaling_residual) {
+            // Log residual requires positive I_obs
+            if (o.r->I <= 0.0f)
+                continue;
 
-        problem.AddResidualBlock(cost,
-                                 nullptr,   // no loss function
-                                 &g[o.img_slot],
-                                 &mosaicity[o.img_slot],
-                                 &Itrue[o.hkl_slot]);
+            auto* cost = new ceres::AutoDiffCostFunction<LogIntensityResidual, 1, 1, 1, 1>(
+                new LogIntensityResidual(*o.r, o.sigma, opt.wedge_deg, refine_partiality));
+
+            problem.AddResidualBlock(cost,
+                                     nullptr,
+                                     &g[o.img_slot],
+                                     &mosaicity[o.img_slot],
+                                     &Itrue[o.hkl_slot]);
+        } else {
+            auto* cost = new ceres::AutoDiffCostFunction<IntensityResidual, 1, 1, 1, 1>(
+                new IntensityResidual(*o.r, o.sigma, opt.wedge_deg, refine_partiality));
+
+            problem.AddResidualBlock(cost,
+                                     nullptr,
+                                     &g[o.img_slot],
+                                     &mosaicity[o.img_slot],
+                                     &Itrue[o.hkl_slot]);
+        }
+    }
+
+    // For log residual, G and Itrue must stay positive
+    if (opt.log_scaling_residual) {
+        for (int i = 0; i < nimg; ++i)
+            problem.SetParameterLowerBound(&g[i], 0, 1e-12);
+        for (int h = 0; h < nhkl; ++h)
+            problem.SetParameterLowerBound(&Itrue[h], 0, 1e-12);
     }
 
     // Optional Kabsch-like regularization for k
-    if (opt.regularize_scale_to_one) {
-        for (int i = 0; i < nimg; ++i) {
-            auto* rcost = new ceres::AutoDiffCostFunction<ScaleRegularizationResidual, 1, 1>(
-                new ScaleRegularizationResidual(opt.scale_regularization_sigma));
-            problem.AddResidualBlock(rcost, nullptr, &g[i]);
-        }
-    }
 
     // Mosaicity refinement + bounds
     if (!opt.refine_mosaicity) {
diff --git a/image_analysis/scale_merge/ScaleAndMerge.h b/image_analysis/scale_merge/ScaleAndMerge.h
index 65f5f310..a3f2875d 100644
--- a/image_analysis/scale_merge/ScaleAndMerge.h
+++ b/image_analysis/scale_merge/ScaleAndMerge.h
@@ -40,6 +40,7 @@ struct ScaleMergeOptions {
     // --- Optional: regularize per-image scale k towards 1 (Kabsch-like) ---
     bool regularize_scale_to_one = true;
     double scale_regularization_sigma = 0.05;
+    bool log_scaling_residual = false;
 };
 
 struct MergedReflection {