diff --git a/image_analysis/scale_merge/RotationScaleMerge.cpp b/image_analysis/scale_merge/RotationScaleMerge.cpp index 6840ec09e..08bbd1b51 100644 --- a/image_analysis/scale_merge/RotationScaleMerge.cpp +++ b/image_analysis/scale_merge/RotationScaleMerge.cpp @@ -275,6 +275,12 @@ namespace { // the loop shrank every scale by 10-30% per round for thirty rounds under the robust fit and // settled in fourteen under the plain one. With the weights the reference uses, the plain // alternating fit is exact coordinate descent on one objective and cannot climb. + double ScaleFromSums(double num, double den, double &info) { + const double G = den > 0.0 ? num / den : NAN; + info = den; + return std::isfinite(G) ? std::max(0.0, G) : 1.0; + } + double SolveScale(const std::vector &obs, double &info) { double num = 0.0, den = 0.0; for (const auto &o : obs) { @@ -282,9 +288,7 @@ namespace { num += w2 * o.coeff * o.Iobs; den += w2 * o.coeff * o.coeff; } - const double G = den > 0.0 ? num / den : NAN; - info = den; - return std::isfinite(G) ? std::max(0.0, G) : 1.0; + return ScaleFromSums(num, den, info); } double median_of(std::vector &v) { @@ -1512,18 +1516,6 @@ int RotationScaleMerge::ComputeAsuGroups(const HKLKeyGenerator &keygen) { return n_groups; } -double RotationScaleMerge::ReferenceWeight(int32_t group, float corr, float partiality, float I, float sigma, - double min_partiality, float &I_corr) { - if (group < 0) return 0.0; - if (!(corr > 0.0f) || !std::isfinite(corr)) return 0.0; - if (partiality < min_partiality) return 0.0; - const float Ic = I * corr; - const float sigma_corr = sigma * corr; - if (!std::isfinite(Ic) || !std::isfinite(sigma_corr) || sigma_corr <= 0.0f) return 0.0; - I_corr = Ic; - return 1.0 / (static_cast(sigma_corr) * sigma_corr); -} - void RotationScaleMerge::ReduceFullGroupMeans(int n_groups, std::vector &out_mean) const { // Inverse-variance per-group mean of I*corr = the merge reference (no cell mask: the scaling // reference is built without a reference cell - only the final merge applies it). A group per task @@ -1565,7 +1557,7 @@ void RotationScaleMerge::ReducePartialGroupMeans(int n_groups, std::vector &refs, const std::vector &corr, std::vector &out_mean) const { - out_mean.assign(n_groups, NAN); + out_mean.resize(n_groups); ParallelChunks(n_groups, ThreadsForWork(corr.size(), nthreads), [&](int glo, int ghi) { for (int g = glo; g < ghi; ++g) { double sw = 0.0, swI = 0.0; @@ -1573,13 +1565,12 @@ void RotationScaleMerge::ReduceScalingGroupMeans(int n_groups, const std::vector for (int p = lo; p < hi; ++p) { const ScalingRef &r = refs[p]; float I_corr = 0.0f; - const double w = ReferenceWeight(g, corr[partial_group_perm[p]], r.partiality, r.I, r.sigma, - min_partiality, I_corr); + const double w = ReferenceWeight(g, corr[p], r.partiality, r.I, r.sigma, min_partiality, I_corr); if (w <= 0.0) continue; sw += w; swI += w * I_corr; } - if (sw > 0.0) out_mean[g] = swI / sw; + out_mean[g] = sw > 0.0 ? swI / sw : NAN; } }); } @@ -1709,6 +1700,31 @@ void RotationScaleMerge::FitPerFrameG(const std::vector &obs, const std::vect frame_info_scratch = std::move(info); } +void RotationScaleMerge::FitScalingFrames(const std::vector &fit, const std::vector &fit_start, + const std::vector &group_mean_in, std::vector &g) { + const int n_fr = static_cast(fit_start.size()) - 1; + std::vector scaled(n_fr, 0); + std::vector info(n_fr, 0.0); + ParallelFor(n_fr, nthreads, [&](int f) { + double num = 0.0, den = 0.0; + size_t n = 0; + for (int i = fit_start[f]; i < fit_start[f + 1]; ++i) { + const ScalingFit &o = fit[i]; + const double mean = group_mean_in[o.group]; + if (!std::isfinite(mean)) continue; + const double coeff = o.coeff_per_mean * mean; + num += o.w2 * coeff * o.Iobs; + den += o.w2 * coeff * coeff; + ++n; + } + if (n < MIN_REFLECTIONS) return; + g[f] = ScaleFromSums(num, den, info[f]); + scaled[f] = 1; + }); + frame_scaled_scratch = std::move(scaled); + frame_info_scratch = std::move(info); +} + void RotationScaleMerge::ShrinkToRestrained(const char *what, const std::vector &g_restrained, const std::vector &fitted, const std::vector &frame_obs_count, @@ -6405,28 +6421,49 @@ RotationScaleMerge::Result RotationScaleMerge::RunArm(bool scale_partials, bool logger->Info("Per-frame scaling (partials): the same grouping and settings as an earlier merge, " "so its {} iterations stand", partial_loop.iterations); } else { - // The loop works on its own compact copy of the partials (ScalingObs in the header) and hands - // corr back to them when it is done. + // The loop works on its own compact copy of the partials (ScalingRef and ScalingFit in the + // header) and hands corr back to them when it is done. scaling_index[q] is the partial at + // position q of the copy: the CSR's group order, then the partials in no group. const int n_partials = static_cast(partials.size()); const size_t partial_threads = ThreadsForWork(partials.size(), nthreads); - std::vector scaling_obs(n_partials); + const int n_grouped = partial_group_start.empty() + ? 0 : partial_group_start.back() + partial_group_count.back(); + std::vector scaling_index(partial_group_perm.get(), partial_group_perm.get() + n_grouped); + for (int i = 0; i < n_partials; ++i) + if (partials[i].group < 0) scaling_index.push_back(i); + std::vector scaling_refs(n_partials); std::vector scaling_corr(n_partials); ParallelChunks(n_partials, partial_threads, [&](int lo, int hi) { - for (int i = lo; i < hi; ++i) { - const Obs &o = partials[i]; - scaling_obs[i] = {o.frame, o.group, o.I, o.sigma, o.partiality, o.prescaling_corr, o.zeta, - o.on_ice}; - scaling_corr[i] = o.corr; + for (int q = lo; q < hi; ++q) { + const Obs &o = partials[scaling_index[q]]; + scaling_refs[q] = {o.frame, o.I, o.sigma, o.partiality, o.prescaling_corr}; + scaling_corr[q] = o.corr; } }); - std::vector scaling_refs(partial_group_start.empty() - ? 0 : partial_group_start.back() + partial_group_count.back()); - ParallelChunks(n_groups, partial_threads, [&](int glo, int ghi) { - for (int g = glo; g < ghi; ++g) - for (int p = partial_group_start[g]; p < partial_group_start[g] + partial_group_count[g]; ++p) { - const ScalingObs &o = scaling_obs[partial_group_perm[p]]; - scaling_refs[p] = {o.I, o.sigma, o.partiality}; - } + + // FitPerFrameG's filters, all but the reference mean's, frame by frame. + auto fitted_obs = [&](const Obs &o) { + return o.group >= 0 && !o.on_ice && o.partiality >= min_partiality + && std::isfinite(o.zeta) && o.zeta > 0.0f; + }; + const int n_fit_frames = static_cast(frame_start.size()); + std::vector fit_start(n_fit_frames + 1, 0); + ParallelFor(n_fit_frames, nthreads, [&](int f) { + fit_start[f + 1] = static_cast(std::count_if( + partials.begin() + frame_start[f], partials.begin() + frame_start[f] + frame_count[f], + fitted_obs)); + }); + for (int f = 0; f < n_fit_frames; ++f) fit_start[f + 1] += fit_start[f]; + std::vector scaling_fit(fit_start.back()); + ParallelFor(n_fit_frames, nthreads, [&](int f) { + int32_t k = fit_start[f]; + for (int i = frame_start[f]; i < frame_start[f] + frame_count[f]; ++i) { + const Obs &o = partials[i]; + if (!fitted_obs(o)) continue; + const double weight = SafeInv(o.sigma, 1.0); + scaling_fit[k++] = {o.group, o.partiality * SafeInv(o.prescaling_corr, 1.0), + static_cast(o.I), weight * weight}; + } }); // RunScalingLoop rescales after every iteration before anything reads corr, so the iteration's @@ -6437,14 +6474,14 @@ RotationScaleMerge::Result RotationScaleMerge::RunArm(bool scale_partials, bool partial_loop = RunScalingLoop("partials", g_partial, frame_count, smooth_window, [&] { ReduceScalingGroupMeans(n_groups, scaling_refs, scaling_corr, partial_mean); - FitPerFrameG(scaling_obs, frame_start, frame_count, partial_mean, /*unity=*/false, g_partial); + FitScalingFrames(scaling_fit, fit_start, partial_mean, g_partial); g_fitted = g_partial; frame_fitted = frame_scaled_scratch; }, [&](const std::vector &apply, const std::vector &ratio) { ParallelChunks(n_partials, partial_threads, [&](int lo, int hi) { - for (int i = lo; i < hi; ++i) { - const ScalingObs &o = scaling_obs[i]; - float &corr = scaling_corr[i]; + for (int q = lo; q < hi; ++q) { + const ScalingRef &o = scaling_refs[q]; + float &corr = scaling_corr[q]; if (frame_fitted[o.frame]) { // UpdateCorr const double denom = static_cast(o.partiality) * g_fitted[o.frame]; if (std::isfinite(o.prescaling_corr) && std::isfinite(denom) && denom > 0.0) @@ -6459,7 +6496,11 @@ RotationScaleMerge::Result RotationScaleMerge::RunArm(bool scale_partials, bool }); ParallelChunks(n_partials, partial_threads, [&](int lo, int hi) { - for (int i = lo; i < hi; ++i) partials[i].corr = scaling_corr[i]; + for (int q = lo; q < hi; ++q) partials[scaling_index[q]].corr = scaling_corr[q]; + }); + // The memo keeps corr in the partials' own order. + ParallelChunks(n_partials, partial_threads, [&](int lo, int hi) { + for (int i = lo; i < hi; ++i) scaling_corr[i] = partials[i].corr; }); PartialLoopMemo m{key, partial_loop, frame_scaled_scratch, g_partial, std::move(scaling_corr)}; if (partial_loop_memos.size() == MAX_PARTIAL_LOOP_MEMOS) diff --git a/image_analysis/scale_merge/RotationScaleMerge.h b/image_analysis/scale_merge/RotationScaleMerge.h index 34baebcd9..cb77e956b 100644 --- a/image_analysis/scale_merge/RotationScaleMerge.h +++ b/image_analysis/scale_merge/RotationScaleMerge.h @@ -616,22 +616,35 @@ private: // The CPU partial scaling loop's working copy of the partials. It passes over them several times // per iteration and reads a handful of fields each time, so it works on these rather than on the - // eighty-byte Obs: ScalingObs holds, in the partials' own order, what FitPerFrameG and the corr - // update read; ScalingRef holds, in the CSR's group order, what the group means read besides corr, - // so that only corr is gathered. The same values, so the same arithmetic in the same order. - struct ScalingObs { - int32_t frame, group; - float I, sigma, partiality, prescaling_corr, zeta; - uint8_t on_ice; - }; + // eighty-byte Obs. ScalingRef holds what the group means and the corr update read, in the CSR's + // group order with the partials in no group after them, and corr is kept beside it in the same + // order - so neither pass gathers. ScalingFit holds, frame by frame, the observations + // FitScalingFrames fits: FitPerFrameG's filters, which do not change over the loop, applied once, + // and the factors of its sums that do not depend on the reference worked out once. The same + // values, so the same arithmetic in the same order. struct ScalingRef { - float I, sigma, partiality; + int32_t frame; + float I, sigma, partiality, prescaling_corr; + }; + struct ScalingFit { + int32_t group; + double coeff_per_mean; // partiality / prescaling_corr: FitPerFrameG's coeff for a mean of 1 + double Iobs, w2; }; // The observation's inverse-variance weight in the scaling reference, and its scaled intensity; // 0 when the observation is not in the reference at all. static double ReferenceWeight(int32_t group, float corr, float partiality, float I, float sigma, - double min_partiality, float &I_corr); + double min_partiality, float &I_corr) { + if (group < 0) return 0.0; + if (!(corr > 0.0f) || !std::isfinite(corr)) return 0.0; + if (partiality < min_partiality) return 0.0; + const float Ic = I * corr; + const float sigma_corr = sigma * corr; + if (!std::isfinite(Ic) || !std::isfinite(sigma_corr) || sigma_corr <= 0.0f) return 0.0; + I_corr = Ic; + return 1.0 / (static_cast(sigma_corr) * sigma_corr); + } static double ReferenceWeight(const Obs &o, double min_partiality, float &I_corr) { return ReferenceWeight(o.group, o.corr, o.partiality, o.I, o.sigma, min_partiality, I_corr); } @@ -640,17 +653,20 @@ private: void ReduceFullGroupMeans(int n_groups, std::vector &out_mean) const; // The same over the partials, using the CSR above - the same sums in the same order. void ReducePartialGroupMeans(int n_groups, std::vector &out_mean) const; - // The same again over the scaling loop's copy: refs in group order, corr in the partials' order. + // The same again over the scaling loop's copy, refs and corr both in group order. void ReduceScalingGroupMeans(int n_groups, const std::vector &refs, const std::vector &corr, std::vector &out_mean) const; // Robust per-frame G fit (IRLS, Cauchy k=3), unity=false uses the rotation partiality, unity=true the // scale-fulls (partiality already folded in). Reads out_mean[group] as the reference intensity. - // O is Obs or ScalingObs. template void FitPerFrameG(const std::vector &obs, const std::vector &fstart, const std::vector &fcount, const std::vector &group_mean_in, bool unity, std::vector &g); + // FitPerFrameG(unity = false) over the scaling loop's ScalingFit, frame f's at + // [fit_start[f], fit_start[f + 1]). + void FitScalingFrames(const std::vector &fit, const std::vector &fit_start, + const std::vector &group_mean_in, std::vector &g); // The alternating per-frame scale fit run to convergence: `iterate` fits every frame once against // the current reference (leaving G in `g` and the frames it fitted in frame_scaled_scratch),