diff --git a/image_analysis/scale_merge/ScaleAndMerge.cpp b/image_analysis/scale_merge/ScaleAndMerge.cpp index e8bf2f5e..90afbd75 100644 --- a/image_analysis/scale_merge/ScaleAndMerge.cpp +++ b/image_analysis/scale_merge/ScaleAndMerge.cpp @@ -185,16 +185,6 @@ namespace { double inv_sigma_; }; -} // namespace - -ScaleMergeResult ScaleAndMergeReflectionsCeres(const std::vector> &observations, - const ScaleMergeOptions &opt) { - if (opt.image_cluster <= 0) - throw std::invalid_argument("image_cluster must be positive"); - - const bool rotation_crystallography = opt.wedge_deg.has_value(); - - ceres::Problem problem; struct ObsRef { const Reflection *r = nullptr; @@ -203,269 +193,130 @@ ScaleMergeResult ScaleAndMergeReflectionsCeres(const std::vector obs; - obs.reserve(nrefl); - - std::unordered_map hklToSlot; - hklToSlot.reserve(nrefl); - - size_t n_image_slots = observations.size() / opt.image_cluster + (observations.size() % opt.image_cluster > 0 ? 1 : 0); - - std::vector image_slot_used(n_image_slots, 0); - - for (int i = 0; i < observations.size(); i++) { - for (const auto &r: observations[i]) { - const double d = SafeD(r.d); - if (!std::isfinite(d)) - continue; - if (!std::isfinite(r.I)) - continue; - - if (opt.d_min_limit_A > 0.0 && d < opt.d_min_limit_A) - continue; - - if (!std::isfinite(r.zeta) || r.zeta <= 0.0f) - continue; - if (!std::isfinite(r.rlp) || r.rlp == 0.0f) - continue; - - const double sigma = SafeSigma(r.sigma, opt.min_sigma); - - const int img_id = i / opt.image_cluster; - image_slot_used[img_id] = 1; - - int hkl_slot; - try { - const HKLKey key = CanonicalizeHKLKey(r, opt); - auto it = hklToSlot.find(key); - if (it == hklToSlot.end()) { - hkl_slot = static_cast(hklToSlot.size()); - hklToSlot.emplace(key, hkl_slot); - } else { - hkl_slot = it->second; - } - } catch (...) { - continue; - } - - ObsRef o; - o.r = &r; - o.img_id = img_id; - o.hkl_slot = hkl_slot; - o.sigma = sigma; - obs.push_back(o); - } - } - - std::vector g(n_image_slots, 1.0); - std::vector 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(hklToSlot.size()); - std::vector Itrue(nhkl, 0.0); - - // Initialize Itrue from per-HKL median of observed intensities - { - std::vector > per_hkl_I(nhkl); - for (const auto &o: obs) { - per_hkl_I[o.hkl_slot].push_back(static_cast(o.r->I)); - } - for (int h = 0; h < nhkl; ++h) { - auto &v = per_hkl_I[h]; - if (v.empty()) { - Itrue[h] = std::max(opt.min_sigma, 1e-6); - continue; - } - std::nth_element(v.begin(), v.begin() + static_cast(v.size() / 2), v.end()); - double med = v[v.size() / 2]; - if (!std::isfinite(med) || med <= opt.min_sigma) - med = opt.min_sigma; - Itrue[h] = med; - } - } - - std::vector is_valid_hkl_slot(nhkl, false); - - for (const auto &o: obs) { - auto *cost = new ceres::AutoDiffCostFunction( - new IntensityResidual(*o.r, o.sigma, opt.wedge_deg.value_or(0.0), rotation_crystallography)); - problem.AddResidualBlock(cost, - nullptr, - &g[o.img_id], - &mosaicity[o.img_id], - &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( - new ScaleRegularizationResidual(0.05)); - problem.AddResidualBlock(cost, nullptr, &g[i]); - } - } - - if (rotation_crystallography) { - if (opt.smoothen_g) { - for (int i = 0; i < g.size() - 2; ++i) { - if (image_slot_used[i] && image_slot_used[i + 1] && image_slot_used[i + 2]) { - auto *cost = new ceres::AutoDiffCostFunction( - new SmoothnessRegularizationResidual(0.05)); - - problem.AddResidualBlock(cost, nullptr, &g[i], &g[i + 1], &g[i + 2]); - } - } - } - - if (opt.smoothen_mos && opt.refine_mosaicity) { - for (int i = 0; i < mosaicity.size() - 2; ++i) { - if (image_slot_used[i] && image_slot_used[i + 1] && image_slot_used[i + 2]) { - auto *cost = new ceres::AutoDiffCostFunction( - new SmoothnessRegularizationResidual(0.05)); - - problem.AddResidualBlock(cost, nullptr, &mosaicity[i], &mosaicity[i + 1], &mosaicity[i + 2]); - } - } - } - } - - // Scaling factors must be always positive - for (int i = 0; i < g.size(); i++) { - if (image_slot_used[i]) - problem.SetParameterLowerBound(&g[i], 0, 1e-12); - } - - // Mosaicity refinement + bounds - if (!opt.refine_mosaicity) { - for (int i = 0; i < mosaicity.size(); ++i) { - if (image_slot_used[i]) - problem.SetParameterBlockConstant(&mosaicity[i]); - } - } else { - for (int i = 0; i < mosaicity.size(); ++i) { - if (image_slot_used[i]) { - problem.SetParameterLowerBound(&mosaicity[i], 0, opt.mosaicity_min_deg); - problem.SetParameterUpperBound(&mosaicity[i], 0, opt.mosaicity_max_deg); - } - } - } - - // use all available threads - unsigned int hw = std::thread::hardware_concurrency(); - if (hw == 0) - hw = 1; // fallback - - ceres::Solver::Options options; - - options.linear_solver_type = ceres::SPARSE_NORMAL_CHOLESKY; - options.minimizer_progress_to_stdout = true; - options.max_num_iterations = opt.max_num_iterations; - options.max_solver_time_in_seconds = opt.max_solver_time_s; - options.num_threads = static_cast(hw); - - ceres::Solver::Summary summary; - ceres::Solve(options, &problem, &summary); - - ScaleMergeResult out; - - out.image_scale_g.resize(observations.size(), NAN); - out.mosaicity_deg.resize(observations.size(), NAN); - for (int i = 0; i < observations.size(); i++) { - size_t img_slot = i / opt.image_cluster; - if (image_slot_used[img_slot]) { - out.image_scale_g[i] = g[img_slot]; - out.mosaicity_deg[i] = mosaicity[img_slot]; - } - } - - std::vector slotToHKL(nhkl); - for (const auto &kv: hklToSlot) - slotToHKL[kv.second] = kv.first; - - out.merged.resize(nhkl); - for (int h = 0; h < nhkl; ++h) { - 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 = Itrue[h]; - out.merged[h].sigma = 0.0; - out.merged[h].d = 0.0; - } - - // Populate d from median of observations per HKL - { - std::vector > per_hkl_d(nhkl); - for (const auto &o: obs) { - const double d_val = static_cast(o.r->d); - if (std::isfinite(d_val) && d_val > 0.0) - per_hkl_d[o.hkl_slot].push_back(d_val); - } - for (int h = 0; h < nhkl; ++h) { - auto &v = per_hkl_d[h]; - if (!v.empty()) { - std::nth_element(v.begin(), v.begin() + static_cast(v.size() / 2), v.end()); - out.merged[h].d = v[v.size() / 2]; - } - } - } - - std::cout << summary.FullReport() << std::endl; - - // ---- Compute corrected observations once (used for both merging and statistics) ---- struct CorrectedObs { int hkl_slot; double I_corr; double sigma_corr; }; - std::vector corr_obs; - corr_obs.reserve(obs.size()); - { - const double half_wedge = opt.wedge_deg.value_or(0.0) / 2.0; + void scale(const ScaleMergeOptions &opt, + std::vector &g, + std::vector &mosaicity, + const std::vector &image_slot_used, + bool rotation_crystallography, + size_t nhkl, + const std::vector &obs) { + ceres::Problem problem; + + std::vector Itrue(nhkl, 0.0); + + // Initialize Itrue from per-HKL median of observed intensities + { + std::vector > per_hkl_I(nhkl); + for (const auto &o: obs) { + per_hkl_I[o.hkl_slot].push_back(static_cast(o.r->I)); + } + for (int h = 0; h < nhkl; ++h) { + auto &v = per_hkl_I[h]; + if (v.empty()) { + Itrue[h] = std::max(opt.min_sigma, 1e-6); + continue; + } + std::nth_element(v.begin(), v.begin() + static_cast(v.size() / 2), v.end()); + double med = v[v.size() / 2]; + if (!std::isfinite(med) || med <= opt.min_sigma) + med = opt.min_sigma; + Itrue[h] = med; + } + } + + std::vector is_valid_hkl_slot(nhkl, false); for (const auto &o: obs) { - const Reflection &r = *o.r; - const double lp = SafeInv(static_cast(r.rlp), 1.0); - const double G_i = g[o.img_id]; + auto *cost = new ceres::AutoDiffCostFunction( + new IntensityResidual(*o.r, o.sigma, opt.wedge_deg.value_or(0.0), rotation_crystallography)); + problem.AddResidualBlock(cost, + nullptr, + &g[o.img_id], + &mosaicity[o.img_id], + &Itrue[o.hkl_slot]); + is_valid_hkl_slot[o.hkl_slot] = true; + } - // Compute partiality with refined mosaicity - double partiality; - if (rotation_crystallography && mosaicity[o.img_id] > 0.0) { - const double c1 = r.zeta / std::sqrt(2.0); - const double arg_plus = (r.delta_phi_deg + half_wedge) * c1 / mosaicity[o.img_id]; - const double arg_minus = (r.delta_phi_deg - half_wedge) * c1 / mosaicity[o.img_id]; - partiality = (std::erf(arg_plus) - std::erf(arg_minus)) / 2.0; - } else { - partiality = r.partiality; + for (int i = 0; i < g.size(); ++i) { + if (image_slot_used[i]) { + auto *cost = new ceres::AutoDiffCostFunction( + new ScaleRegularizationResidual(0.05)); + problem.AddResidualBlock(cost, nullptr, &g[i]); + } + } + + if (rotation_crystallography) { + if (opt.smoothen_g) { + for (int i = 0; i < g.size() - 2; ++i) { + if (image_slot_used[i] && image_slot_used[i + 1] && image_slot_used[i + 2]) { + auto *cost = new ceres::AutoDiffCostFunction( + new SmoothnessRegularizationResidual(0.05)); + + problem.AddResidualBlock(cost, nullptr, &g[i], &g[i + 1], &g[i + 2]); + } + } } - if (partiality <= opt.min_partiality_for_merge) - continue; - const double correction = G_i * partiality * lp; - if (correction <= 0.0) - continue; + if (opt.smoothen_mos && opt.refine_mosaicity) { + for (int i = 0; i < mosaicity.size() - 2; ++i) { + if (image_slot_used[i] && image_slot_used[i + 1] && image_slot_used[i + 2]) { + auto *cost = new ceres::AutoDiffCostFunction( + new SmoothnessRegularizationResidual(0.05)); - corr_obs.push_back({ - o.hkl_slot, - static_cast(r.I) / correction, - o.sigma / correction - }); + problem.AddResidualBlock(cost, nullptr, &mosaicity[i], &mosaicity[i + 1], &mosaicity[i + 2]); + } + } + } } + + // Scaling factors must be always positive + for (int i = 0; i < g.size(); i++) { + if (image_slot_used[i]) + problem.SetParameterLowerBound(&g[i], 0, 1e-12); + } + + // Mosaicity refinement + bounds + if (!opt.refine_mosaicity) { + for (int i = 0; i < mosaicity.size(); ++i) { + if (image_slot_used[i]) + problem.SetParameterBlockConstant(&mosaicity[i]); + } + } else { + for (int i = 0; i < mosaicity.size(); ++i) { + if (image_slot_used[i]) { + problem.SetParameterLowerBound(&mosaicity[i], 0, opt.mosaicity_min_deg); + problem.SetParameterUpperBound(&mosaicity[i], 0, opt.mosaicity_max_deg); + } + } + } + + // use all available threads + unsigned int hw = std::thread::hardware_concurrency(); + if (hw == 0) + hw = 1; // fallback + + ceres::Solver::Options options; + + options.linear_solver_type = ceres::SPARSE_NORMAL_CHOLESKY; + options.minimizer_progress_to_stdout = true; + options.max_num_iterations = opt.max_num_iterations; + options.max_solver_time_in_seconds = opt.max_solver_time_s; + options.num_threads = static_cast(hw); + + ceres::Solver::Summary summary; + ceres::Solve(options, &problem, &summary); + std::cout << summary.FullReport() << std::endl; } - // ---- Merge (XDS/XSCALE style: inverse-variance weighted mean) ---- - { + void merge(size_t nhkl, ScaleMergeResult &out, const std::vector &corr_obs) { + // ---- Merge (XDS/XSCALE style: inverse-variance weighted mean) ---- + struct HKLAccum { double sum_wI = 0.0; double sum_w = 0.0; @@ -501,6 +352,7 @@ ScaleMergeResult ScaleAndMergeReflectionsCeres(const std::vector &corr_obs) // ---- Compute per-shell merging statistics ---- { constexpr int kStatShells = 10; @@ -657,6 +509,187 @@ ScaleMergeResult ScaleAndMergeReflectionsCeres(const std::vector &g, + std::vector &mosaicity, + bool rotation_crystallography, + const std::vector &obs, + std::vector &corr_obs) { + + // ---- Compute corrected observations once (used for both merging and statistics) ---- + const double half_wedge = opt.wedge_deg.value_or(0.0) / 2.0; + + for (const auto &o: obs) { + const Reflection &r = *o.r; + const double lp = SafeInv(static_cast(r.rlp), 1.0); + const double G_i = g[o.img_id]; + + // Compute partiality with refined mosaicity + double partiality; + if (rotation_crystallography && mosaicity[o.img_id] > 0.0) { + const double c1 = r.zeta / std::sqrt(2.0); + const double arg_plus = (r.delta_phi_deg + half_wedge) * c1 / mosaicity[o.img_id]; + const double arg_minus = (r.delta_phi_deg - half_wedge) * c1 / mosaicity[o.img_id]; + partiality = (std::erf(arg_plus) - std::erf(arg_minus)) / 2.0; + } else { + partiality = r.partiality; + } + + if (partiality <= opt.min_partiality_for_merge) + continue; + const double correction = G_i * partiality * lp; + if (correction <= 0.0) + continue; + + corr_obs.push_back({ + o.hkl_slot, + static_cast(r.I) / correction, + o.sigma / correction + }); + } + } + + void proc_obs(const std::vector > &observations, + const ScaleMergeOptions &opt, + std::vector &image_slot_used, + std::vector &obs, + std::unordered_map &hklToSlot + ) { + for (int i = 0; i < observations.size(); i++) { + for (const auto &r: observations[i]) { + const double d = SafeD(r.d); + if (!std::isfinite(d)) + continue; + if (!std::isfinite(r.I)) + continue; + + if (opt.d_min_limit_A > 0.0 && d < opt.d_min_limit_A) + continue; + + if (!std::isfinite(r.zeta) || r.zeta <= 0.0f) + continue; + if (!std::isfinite(r.rlp) || r.rlp == 0.0f) + continue; + + const double sigma = SafeSigma(r.sigma, opt.min_sigma); + + const int img_id = i / opt.image_cluster; + image_slot_used[img_id] = 1; + + int hkl_slot; + try { + const HKLKey key = CanonicalizeHKLKey(r, opt); + auto it = hklToSlot.find(key); + if (it == hklToSlot.end()) { + hkl_slot = static_cast(hklToSlot.size()); + hklToSlot.emplace(key, hkl_slot); + } else { + hkl_slot = it->second; + } + } catch (...) { + continue; + } + + ObsRef o; + o.r = &r; + o.img_id = img_id; + o.hkl_slot = hkl_slot; + o.sigma = sigma; + obs.push_back(o); + } + } + + } +} // namespace + +ScaleMergeResult ScaleAndMergeReflectionsCeres(const std::vector > &observations, + const ScaleMergeOptions &opt) { + if (opt.image_cluster <= 0) + throw std::invalid_argument("image_cluster must be positive"); + + const bool rotation_crystallography = opt.wedge_deg.has_value(); + + size_t nrefl = 0; + for (const auto &i: observations) + nrefl += i.size(); + + std::vector obs; + std::vector corr_obs; + + obs.reserve(nrefl); + corr_obs.reserve(nrefl); + + std::unordered_map hklToSlot; + hklToSlot.reserve(nrefl); + + size_t n_image_slots = observations.size() / opt.image_cluster + + (observations.size() % opt.image_cluster > 0 ? 1 : 0); + + std::vector image_slot_used(n_image_slots, 0); + + proc_obs(observations, opt, image_slot_used, obs, hklToSlot); + + const int nhkl = static_cast(hklToSlot.size()); + + std::vector g(n_image_slots, 1.0); + std::vector 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]; + } + } + + scale(opt, g, mosaicity, image_slot_used, rotation_crystallography, nhkl, obs); + + ScaleMergeResult out; + + out.image_scale_g.resize(observations.size(), NAN); + out.mosaicity_deg.resize(observations.size(), NAN); + for (int i = 0; i < observations.size(); i++) { + size_t img_slot = i / opt.image_cluster; + if (image_slot_used[img_slot]) { + out.image_scale_g[i] = g[img_slot]; + out.mosaicity_deg[i] = mosaicity[img_slot]; + } + } + + std::vector slotToHKL(nhkl); + for (const auto &kv: hklToSlot) + slotToHKL[kv.second] = kv.first; + + out.merged.resize(nhkl); + for (int h = 0; h < nhkl; ++h) { + 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 = 0.0; + out.merged[h].sigma = 0.0; + out.merged[h].d = 0.0; + } + + // Populate d from median of observations per HKL + { + std::vector > per_hkl_d(nhkl); + for (const auto &o: obs) { + const double d_val = static_cast(o.r->d); + if (std::isfinite(d_val) && d_val > 0.0) + per_hkl_d[o.hkl_slot].push_back(d_val); + } + for (int h = 0; h < nhkl; ++h) { + auto &v = per_hkl_d[h]; + if (!v.empty()) { + std::nth_element(v.begin(), v.begin() + static_cast(v.size() / 2), v.end()); + out.merged[h].d = v[v.size() / 2]; + } + } + } + + calc_obs(opt, g, mosaicity, rotation_crystallography, obs, corr_obs); + merge(nhkl, out, corr_obs); + stats(opt, nhkl, out, corr_obs); return out; }