ScaleAndMerge: Fix partiality model

image_analysis/scale_merge/ScaleAndMerge.cpp

@@ -114,39 +114,30 @@ inline HKLKey CanonicalizeHKLKey(const Reflection& r, const ScaleMergeOptions& o
 }
 
 struct IntensityResidual {
-    IntensityResidual(const Reflection& r, double sigma_obs, double wedge_deg, double s2, bool refine_b,
-        bool partiality_model)
+    IntensityResidual(const Reflection& r, double sigma_obs, double wedge_deg, bool refine_partiality)
         : Iobs_(static_cast<double>(r.I)),
           inv_sigma_(SafeInv(sigma_obs, 1.0)),
           delta_phi_rad_(static_cast<double>(r.delta_phi) * M_PI / 180.0),
           lp_(SafeInv(static_cast<double>(r.rlp), 1.0)),
           half_wedge_rad_(wedge_deg * M_PI / 180.0 / 2.0),
-          c1_(std::sqrt(2.0) * SafeInv(static_cast<double>(r.zeta), 1.0)),
-          s2_(s2),
-          refine_b_(refine_b),
-          partiality_model_(partiality_model) {}
+          c1_(r.zeta / std::sqrt(2.0)),
+          partiality_(r.partiality),
+          refine_partiality_(refine_partiality) {}
 
     template<typename T>
-    bool operator()(const T* const log_k,
-                    const T* const b,
+    bool operator()(const T* const k,
                     const T* const mosaicity_rad,
-                    const T* const log_Itrue,
+                    const T* const Itrue,
                     T* residual) const {
-        const T k = ceres::exp(log_k[0]);
-        const T Itrue = ceres::exp(log_Itrue[0]);
-
-        T partiality = T(1.0);
-        if (partiality_model_) {
-            const T arg_plus  = T(delta_phi_rad_ + half_wedge_rad_) * (T(c1_) * mosaicity_rad[0]);
-            const T arg_minus = T(delta_phi_rad_ - half_wedge_rad_) * (T(c1_) * mosaicity_rad[0]);
+        T partiality;
+        if (refine_partiality_ && mosaicity_rad[0] != 0.0) {
+            const T arg_plus  = T(delta_phi_rad_ + half_wedge_rad_) * T(c1_) / mosaicity_rad[0];
+            const T arg_minus = T(delta_phi_rad_ - half_wedge_rad_) * T(c1_) / mosaicity_rad[0];
             partiality = (ceres::erf(arg_plus) - ceres::erf(arg_minus)) / T(2.0);
-        }
-        T wilson = T(1.0);
-        if (refine_b_) {
-            wilson = ceres::exp(-b[0] * T(s2_));
-        }
+        } else
+            partiality = T(partiality_);
 
-        const T Ipred = k * wilson * partiality * T(lp_) * Itrue;
+        const T Ipred = k[0] * partiality * T(lp_) * Itrue[0];
         residual[0] = (Ipred - T(Iobs_)) * T(inv_sigma_);
         return true;
     }
@@ -157,9 +148,8 @@ struct IntensityResidual {
     double lp_;
     double half_wedge_rad_;
     double c1_;
-    double s2_;
-    bool refine_b_;
-    bool partiality_model_;
+    double partiality_;
+    bool refine_partiality_;
 };
 
 struct ScaleRegularizationResidual {
@@ -167,9 +157,8 @@ struct ScaleRegularizationResidual {
         : inv_sigma_(SafeInv(sigma_k, 1.0)) {}
 
     template <typename T>
-    bool operator()(const T* const log_k, T* residual) const {
-        const T k = ceres::exp(log_k[0]);
-        residual[0] = (k - T(1.0)) * T(inv_sigma_);
+    bool operator()(const T* const k, T* residual) const {
+        residual[0] = (k[0] - T(1.0)) * T(inv_sigma_);
         return true;
     }
 
@@ -187,7 +176,6 @@ ScaleMergeResult ScaleAndMergeReflectionsCeres(const std::vector<Reflection>& ob
         int img_id = 0;
         int img_slot = -1;
         int hkl_slot = -1;
-        double s2 = 0.0;
         double sigma = 0.0;
     };
 
@@ -213,7 +201,6 @@ ScaleMergeResult ScaleAndMergeReflectionsCeres(const std::vector<Reflection>& ob
             continue;
 
         const double sigma = SafeSigma(static_cast<double>(r.sigma), opt.min_sigma);
-        const double s2 = 1.0 / (d * d);
 
         const int img_id = RoundImageId(r.image_number, opt.image_number_rounding);
 
@@ -247,7 +234,6 @@ ScaleMergeResult ScaleAndMergeReflectionsCeres(const std::vector<Reflection>& ob
         o.img_id = img_id;
         o.img_slot = img_slot;
         o.hkl_slot = hkl_slot;
-        o.s2 = s2;
         o.sigma = sigma;
         obs.push_back(o);
     }
@@ -256,16 +242,14 @@ ScaleMergeResult ScaleAndMergeReflectionsCeres(const std::vector<Reflection>& ob
     const int nhkl = static_cast<int>(hklToSlot.size());
 
     out.image_scale_k.assign(nimg, 1.0);
-    out.image_b_factor.assign(nimg, 0.0);
     out.image_ids.assign(nimg, 0);
 
     for (const auto& kv : imgIdToSlot) {
         out.image_ids[kv.second] = kv.first;
     }
 
-    std::vector<double> log_k(nimg, 0.0);
-    std::vector<double> b(nimg, 0.0);
-    std::vector<double> log_Itrue(nhkl, 0.0);
+    std::vector<double> k(nimg, 1.0);
+    std::vector<double> Itrue(nhkl, 0.0);
 
     auto deg2rad = [](double deg) { return deg * M_PI / 180.0; };
 
@@ -281,31 +265,30 @@ ScaleMergeResult ScaleAndMergeReflectionsCeres(const std::vector<Reflection>& ob
         for (int h = 0; h < nhkl; ++h) {
             auto& v = per_hkl_I[h];
             if (v.empty()) {
-                log_Itrue[h] = std::log(std::max(opt.min_sigma, 1e-6));
+                Itrue[h] = std::max(opt.min_sigma, 1e-6);
                 continue;
             }
             std::nth_element(v.begin(), v.begin() + static_cast<long>(v.size() / 2), v.end());
             double med = v[v.size() / 2];
             if (!std::isfinite(med) || med <= opt.min_sigma)
                 med = opt.min_sigma;
-            log_Itrue[h] = std::log(med);
+            Itrue[h] = med;
         }
     }
 
     ceres::Problem problem;
 
-    const bool partiality_model = opt.wedge_deg > 0.0;
+    const bool refine_partiality = opt.wedge_deg > 0.0;
 
     for (const auto& o : obs) {
-        auto* cost = new ceres::AutoDiffCostFunction<IntensityResidual, 1, 1, 1, 1, 1>(
-            new IntensityResidual(*o.r, o.sigma, opt.wedge_deg, o.s2, opt.refine_b_factor, partiality_model));
+        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,   // no loss function
-                                 &log_k[o.img_slot],
-                                 &b[o.img_slot],
+                                 &k[o.img_slot],
                                  &mosaicity_rad[o.img_slot],
-                                 &log_Itrue[o.hkl_slot]);
+                                 &Itrue[o.hkl_slot]);
     }
 
     // Optional Kabsch-like regularization for k
@@ -313,28 +296,14 @@ ScaleMergeResult ScaleAndMergeReflectionsCeres(const std::vector<Reflection>& ob
         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, &log_k[i]);
+            problem.AddResidualBlock(rcost, nullptr, &k[i]);
         }
     }
 
     // Fix gauge freedom
     if (opt.fix_first_image_scale && nimg > 0) {
-        log_k[0] = 0.0;
-        problem.SetParameterBlockConstant(&log_k[0]);
-    }
-
-    if (!opt.refine_b_factor) {
-        for (int i = 0; i < nimg; ++i) {
-            b[i] = 0.0;
-            problem.SetParameterBlockConstant(&b[i]);
-        }
-    } else {
-        for (int i = 0; i < nimg; ++i) {
-            if (opt.b_min)
-                problem.SetParameterLowerBound(&b[i], 0, *opt.b_min);
-            if (opt.b_max)
-                problem.SetParameterUpperBound(&b[i], 0, *opt.b_max);
-        }
+        k[0] = 1.0;
+        problem.SetParameterBlockConstant(&k[0]);
     }
 
     // Mosaicity refinement + bounds
@@ -342,14 +311,22 @@ ScaleMergeResult ScaleAndMergeReflectionsCeres(const std::vector<Reflection>& ob
         for (int i = 0; i < nimg; ++i)
             problem.SetParameterBlockConstant(&mosaicity_rad[i]);
     } else {
-        const std::optional<double> min_rad = opt.mosaicity_min_deg
-            ? std::optional<double>(deg2rad(*opt.mosaicity_min_deg)) : std::nullopt;
-        const std::optional<double> max_rad = opt.mosaicity_max_deg
-            ? std::optional<double>(deg2rad(*opt.mosaicity_max_deg)) : std::nullopt;
+        const double min_rad = deg2rad(opt.mosaicity_min_deg);
+        const double max_rad = deg2rad(opt.mosaicity_max_deg);
 
         for (int i = 0; i < nimg; ++i) {
-            if (min_rad) problem.SetParameterLowerBound(&mosaicity_rad[i], 0, *min_rad);
-            if (max_rad) problem.SetParameterUpperBound(&mosaicity_rad[i], 0, *max_rad);
+            problem.SetParameterLowerBound(&mosaicity_rad[i], 0, min_rad);
+            problem.SetParameterUpperBound(&mosaicity_rad[i], 0, max_rad);
+        }
+    }
+
+    // Force k[i] > 0 for all images (scale factors must be positive)
+    {
+        constexpr double k_floor = 1e-8;
+        for (int i = 0; i < nimg; ++i) {
+            if (!(opt.fix_first_image_scale && i == 0)) {
+                problem.SetParameterLowerBound(&k[i], 0, k_floor);
+            }
         }
     }
 
@@ -364,10 +341,8 @@ ScaleMergeResult ScaleAndMergeReflectionsCeres(const std::vector<Reflection>& ob
     ceres::Solve(options, &problem, &summary);
 
     // --- Export per-image results ---
-    for (int i = 0; i < nimg; ++i) {
-        out.image_scale_k[i] = std::exp(log_k[i]);
-        out.image_b_factor[i] = opt.refine_b_factor ? b[i] : 0.0;
-    }
+    for (int i = 0; i < nimg; ++i)
+        out.image_scale_k[i] = k[i];
 
     out.mosaicity_deg.resize(nimg);
     for (int i = 0; i < nimg; ++i)
@@ -375,13 +350,11 @@ ScaleMergeResult ScaleAndMergeReflectionsCeres(const std::vector<Reflection>& ob
 
     // --- Compute goodness-of-fit (reduced chi-squared) ---
     const int n_obs = static_cast<int>(obs.size());
-    // Count free parameters: nhkl log_Itrue + per-image (log_k + b + mosaicity) minus fixed ones
+    // Count free parameters: nhkl Itrue + per-image (k + mosaicity) minus fixed ones
     int n_params = nhkl;
     for (int i = 0; i < nimg; ++i) {
         if (!(opt.fix_first_image_scale && i == 0))
-            n_params += 1; // log_k
-        if (opt.refine_b_factor)
-            n_params += 1; // b
+            n_params += 1; // k
         if (opt.refine_mosaicity)
             n_params += 1; // mosaicity
     }
@@ -393,22 +366,20 @@ ScaleMergeResult ScaleAndMergeReflectionsCeres(const std::vector<Reflection>& ob
         const int i = o.img_slot;
         const int h = o.hkl_slot;
 
-        const double ki = std::exp(log_k[i]);
-        const double atten = opt.refine_b_factor ? std::exp(-b[i] * o.s2) : 1.0;
-        const double Itrue = std::exp(log_Itrue[h]);
+        const double ki = k[i];
 
-        const double zeta = static_cast<double>(o.r->zeta);
         const double mosa_i = mosaicity_rad[i];
-        const double c1 = std::sqrt(2.0) * (mosa_i / zeta);
+        const double c2 = o.r->zeta / (std::sqrt(2.0) * mosa_i);
 
         const double delta_phi_rad = static_cast<double>(o.r->delta_phi) * M_PI / 180.0;
-        const double partiality = partiality_model
-            ? (std::erf((delta_phi_rad + half_wedge_rad) * c1) - std::erf((delta_phi_rad - half_wedge_rad) * c1)) / 2.0
-            : 1.0;
+        const double partiality = refine_partiality
+            ? (std::erf((delta_phi_rad + half_wedge_rad) * c2)
+                - std::erf((delta_phi_rad - half_wedge_rad) * c2)) / 2.0
+            : o.r->partiality;
 
-        const double lp = SafeInv(static_cast<double>(o.r->rlp), 1.0);
+        const double lp = SafeInv(o.r->rlp, 1.0);
 
-        const double Ipred = ki * atten * partiality * lp * Itrue;
+        const double Ipred = ki * partiality * lp * Itrue[h];
         const double resid = (Ipred - static_cast<double>(o.r->I)) / o.sigma;
         sum_r2 += resid * resid;
     }
@@ -427,7 +398,7 @@ ScaleMergeResult ScaleAndMergeReflectionsCeres(const std::vector<Reflection>& ob
         out.merged[h].h = slotToHKL[h].h;
         out.merged[h].k = slotToHKL[h].k;
         out.merged[h].l = slotToHKL[h].l;
-        out.merged[h].I = std::exp(log_Itrue[h]);
+        out.merged[h].I = Itrue[h];
         out.merged[h].sigma = std::numeric_limits<double>::quiet_NaN();
         out.merged[h].d = 0.0;
     }
@@ -457,17 +428,16 @@ ScaleMergeResult ScaleAndMergeReflectionsCeres(const std::vector<Reflection>& ob
     std::vector<std::pair<const double*, const double*>> covariance_blocks;
     covariance_blocks.reserve(nhkl);
     for (int h = 0; h < nhkl; ++h) {
-        covariance_blocks.emplace_back(&log_Itrue[h], &log_Itrue[h]);
+        covariance_blocks.emplace_back(&Itrue[h], &Itrue[h]);
     }
 
     if (covariance.Compute(covariance_blocks, &problem)) {
         for (int h = 0; h < nhkl; ++h) {
-            double var_log_I = 0.0;
-            covariance.GetCovarianceBlock(&log_Itrue[h], &log_Itrue[h], &var_log_I);
+            double var_I = 0.0;
+            covariance.GetCovarianceBlock(&Itrue[h], &Itrue[h], &var_I);
 
-            // σ(I) = I * sqrt( var(log I) * GoF² )
-            const double Itrue = std::exp(log_Itrue[h]);
-            out.merged[h].sigma = Itrue * std::sqrt(var_log_I * gof2);
+            // σ(I) = I * sqrt( var(I) * GoF² )
+            out.merged[h].sigma = std::sqrt(var_I * gof2);
         }
     }