mx/Jungfraujoch
3
2 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 93 Wiki Activity

RotationScaleMerge: the partials' scaling loop walks only what the means can read, smaller fit records, per-parity surface references, output unchanged

Browse Source 
Three more cuts in the bytes the CPU scaling moves per iteration:

- The partials' working copy is ordered with the partials a group mean can take in first - in a
  group and not below the partiality floor, ReferenceWeight's fixed tests - group by group in the
  CSR's order, and all the others after them. The group means walk only the first part: a partial
  they would skip anyway adds nothing to a sum, so no sum changes, and on a fine-sliced sweep about
  half of the partials are rocking-curve tails below the floor. The others are read by nothing
  during the loop, so their per-iteration corr updates are kept (a few per-frame arrays an
  iteration) and applied to them once, in order, when the loop ends.
- ScalingFit holds the four floats the per-frame fit reads (20 bytes, was 32) and the fit works
  out 1/sigma^2 and partiality/prescaling_corr from them where it uses them: the same expressions,
  so the same doubles.
- The correction surfaces keep their group-ordered reference terms in one copy per frame parity,
  so a cross-validation reference pass reads only its own half; a pass over all the terms walks a
  group's two halves together in term order, which is the order the single copy had.

Samples in the partial loop on a 3600-frame set (perf, cycles, whole scaling-and-merging step):
group means 29.7k -> 15.1k, per-frame fit 22.4k -> 13.7k, the separate corr update pass 11.9k ->
gone. Wall time on the shared machine (load 26-33) is within its noise of the previous commit.

Validated with the battery against the start of the branch (873880bba), --no-gpu: p.hkl and MTZ
data byte-identical on 4nwv, 8qaw, 9v16, lyso_x06da_ref, cytc_x10sa, myob_x10sa and
yag_x10sa_20keV; the same bytes again at -N 11 on all seven; and on the GPU path the same bytes
on lyso_x06da_ref and 9v16.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01PXc1zD1yxHqaR1DGdCTHL4
This commit is contained in:
leonarski_fandClaude Opus 5.5 committed 2026-10-15 22:24:40 +02:00
1 parent c914fa9f33
commit 4a8aff41c2
2 files changed
+85 -40
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

image_analysis/scale_merge/RotationScaleMerge.cpp
+80 -38
View File
@@ -1696,9 +1696,11 @@ void RotationScaleMerge::FitScalingFrames(const FillVector<ScalingFit> &fit, con
// iterations, which is what RunScalingLoop's tolerance is for.
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;
const double weight = SafeInv(o.sigma, 1.0), w2 = weight * weight;
const double coeff_per_mean = o.partiality * SafeInv(o.prescaling_corr, 1.0);
const double coeff = coeff_per_mean * mean;
num += w2 * coeff * static_cast<double>(o.I);
den += w2 * coeff * coeff;
++n;
}
if (n < min_obs) return;
@@ -3897,18 +3899,25 @@ void RotationScaleMerge::ApplyCellSurface(const std::vector<int32_t> &cell, int
// terms in fulls order, because the counting sort is stable (it walks `term`, which is in fulls
// order). A separate copy rather than a permutation: chasing every term through an index at random
// over a few hundred megabytes is the locality this pass used to stay serial to keep.
struct RefTerm { float I, sigma, corr; int32_t cell; uint8_t parity; };
std::vector<int32_t> g_start;
std::vector<RefTerm> gterm;
// One copy per frame parity: most reference passes are of one half (the cross-validation), and
// each then reads only its own. A pass over all the terms walks a group's two halves together in
// term order (`at`, the term's position), which is the fulls order the single copy held.
struct RefTerm { float I, sigma, corr; int32_t cell, at; };
std::array<std::vector<int32_t>, 2> g_start;
std::array<std::vector<RefTerm>, 2> gterm;
if (!on_gpu) {
g_start.assign(n_groups + 1, 0);
for (const Term &t : term) ++g_start[t.group + 1];
for (int g = 0; g < n_groups; ++g) g_start[g + 1] += g_start[g];
gterm.resize(term.size());
std::vector<int32_t> fill(g_start.begin(), g_start.end() - 1);
for (int par = 0; par < 2; ++par) g_start[par].assign(n_groups + 1, 0);
for (size_t k = 0; k < term.size(); ++k) ++g_start[term_parity[k]][term[k].group + 1];
std::array<std::vector<int32_t>, 2> fill;
for (int par = 0; par < 2; ++par) {
for (int g = 0; g < n_groups; ++g) g_start[par][g + 1] += g_start[par][g];
gterm[par].resize(g_start[par][n_groups]);
fill[par].assign(g_start[par].begin(), g_start[par].end() - 1);
}
for (size_t k = 0; k < term.size(); ++k) {
const Term &t = term[k];
gterm[fill[t.group]++] = {t.I, t.sigma, t.corr, t.cell, term_parity[k]};
const int par = term_parity[k];
gterm[par][fill[par][t.group]++] = {t.I, t.sigma, t.corr, t.cell, static_cast<int32_t>(k)};
}
}
@@ -3927,9 +3936,7 @@ void RotationScaleMerge::ApplyCellSurface(const std::vector<int32_t> &cell, int
ParallelChunks(n_groups, nt, [&](int glo, int ghi) {
for (int g = glo; g < ghi; ++g) {
double s_w = 0.0, s_wI = 0.0;
for (int k = g_start[g]; k < g_start[g + 1]; ++k) {
const RefTerm &t = gterm[k];
if (parity >= 0 && t.parity != parity) continue;
auto add = [&](const RefTerm &t) {
// Keep the intensity's own product together before it is weighted: w * ((I*corr)*a)
// is what the loops this replaced summed, and w * I * corr * a rounds differently.
// The score decides whether a surface is kept, so a last-place difference here can
@@ -3937,6 +3944,14 @@ void RotationScaleMerge::ApplyCellSurface(const std::vector<int32_t> &cell, int
const double a = A[t.cell], Is = static_cast<double>(t.I) * t.corr * a;
const double sc = static_cast<double>(t.sigma) * t.corr * a, w = 1.0 / (sc * sc);
s_w += w; s_wI += w * Is;
};
if (parity >= 0) {
for (int k = g_start[parity][g]; k < g_start[parity][g + 1]; ++k) add(gterm[parity][k]);
} else {
int k0 = g_start[0][g], k1 = g_start[1][g];
const int e0 = g_start[0][g + 1], e1 = g_start[1][g + 1];
while (k0 < e0 || k1 < e1)
add(k1 == e1 || (k0 < e0 && gterm[0][k0].at < gterm[1][k1].at) ? gterm[0][k0++] : gterm[1][k1++]);
}
sw[g] = s_w; swI[g] = s_wI;
iref[g] = s_w > 0.0 ? s_wI / s_w : std::numeric_limits<double>::quiet_NaN();
@@ -6511,18 +6526,43 @@ RotationScaleMerge::Result RotationScaleMerge::RunArm(bool scale_partials, bool
} else {
// 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.
// position q of the copy: first the partials a group mean can take in - in a group and not
// below the partiality floor (ReferenceWeight's fixed tests) - group by group in the CSR's
// order, then all the others in index order. The means walk only the first part: a partial
// they would skip in any case adds nothing to a sum, so leaving it out of the walk changes no
// sum, and on a fine-sliced sweep half of the partials are rocking-curve tails below the floor.
const int n_partials = static_cast<int>(partials.size());
const size_t partial_threads = ThreadsForWork(partials.size(), nthreads);
const int n_grouped = partial_group_start.empty()
? 0 : partial_group_start.back() + partial_group_count.back();
FillVector<int32_t> scaling_index(n_partials);
const std::vector<int32_t> ungrouped = KeptIndices(n_partials, nthreads, [&](int i) {
return partial_group[i] < 0;
});
// Taken once, in index order, so that the walks over the CSR below gather one byte a partial
// rather than its record.
FillVector<uint8_t> in_ref(n_partials);
ParallelChunks(n_partials, partial_threads, [&](int lo, int hi) {
for (int q = lo; q < hi; ++q)
scaling_index[q] = q < n_grouped ? partial_group_perm[q] : ungrouped[q - n_grouped];
for (int i = lo; i < hi; ++i)
in_ref[i] = partial_group[i] >= 0 && !(partials[i].partiality < min_partiality) ? 1 : 0;
});
auto in_reference = [&](int i) { return in_ref[i] != 0; };
std::vector<int32_t> ref_start(n_groups), ref_count(n_groups);
ParallelChunks(n_groups, ThreadsForWork(partials.size(), nthreads), [&](int glo, int ghi) {
for (int g = glo; g < ghi; ++g) {
int32_t c = 0;
for (int p = partial_group_start[g]; p < partial_group_start[g] + partial_group_count[g]; ++p)
c += in_reference(partial_group_perm[p]) ? 1 : 0;
ref_count[g] = c;
}
});
int n_ref = 0;
for (int g = 0; g < n_groups; ++g) { ref_start[g] = n_ref; n_ref += ref_count[g]; }
FillVector<int32_t> scaling_index(n_partials);
ParallelChunks(n_groups, ThreadsForWork(partials.size(), nthreads), [&](int glo, int ghi) {
for (int g = glo; g < ghi; ++g) {
int32_t q = ref_start[g];
for (int p = partial_group_start[g]; p < partial_group_start[g] + partial_group_count[g]; ++p)
if (in_reference(partial_group_perm[p])) scaling_index[q++] = partial_group_perm[p];
}
});
const std::vector<int32_t> rest = KeptIndices(n_partials, nthreads, [&](int i) { return !in_reference(i); });
ParallelChunks(static_cast<int>(rest.size()), ThreadsForWork(rest.size(), nthreads), [&](int lo, int hi) {
for (int j = lo; j < hi; ++j) scaling_index[n_ref + j] = rest[j];
});
FillVector<ScalingRef> scaling_refs(n_partials);
FillVector<float> scaling_corr(n_partials);
@@ -6558,9 +6598,7 @@ RotationScaleMerge::Result RotationScaleMerge::RunArm(bool scale_partials, bool
for (int i = frame_start[f]; i < frame_start[f] + frame_count[f]; ++i) {
if (!fitted_obs(i)) continue;
const Obs &o = partials[i];
const double weight = SafeInv(o.sigma, 1.0);
scaling_fit[k++] = {partial_group[i], o.partiality * SafeInv(o.prescaling_corr, 1.0),
static_cast<double>(o.I), weight * weight};
scaling_fit[k++] = {partial_group[i], o.I, o.sigma, o.partiality, o.prescaling_corr};
}
});
@@ -6568,14 +6606,17 @@ RotationScaleMerge::Result RotationScaleMerge::RunArm(bool scale_partials, bool
// UpdateCorr is folded into the rescale: the G and the fitted frames as the iteration left
// them, then the rescale on top - the same two roundings. Nor is that a pass of its own over
// the copy: the next iteration's group means are the next thing to read corr, so each partial
// in a group is brought up to date there, in the same walk, and only the partials in no group
// (which the means never read) are updated right away. Whatever is still owed when the loop
// stops is applied after it. Each corr goes through the same updates in the same order.
// they walk is brought up to date there, in the same walk, and whatever is still owed when the
// loop stops is applied after it. The partials the means do not walk are read by nothing until
// the loop is over, so their updates are kept - a few per-frame arrays an iteration - and
// applied to them once, in order, at the end. Each corr goes through the same updates in the
// same order.
CorrUpdate update;
bool update_owed = false; // the grouped partials have not had `update` yet
bool update_owed = false; // the walked partials have not had `update` yet
std::vector<CorrUpdate> rest_updates;
partial_loop = RunScalingLoop("partials", g_partial, frame_count, smooth_window,
[&] {
ReduceScalingGroupMeans(n_groups, partial_group_start.data(), partial_group_count.data(),
ReduceScalingGroupMeans(n_groups, ref_start.data(), ref_count.data(),
scaling_refs, scaling_corr, partial_mean, update_owed ? &update : nullptr);
update_owed = false;
FitScalingFrames(scaling_fit, fit_start, MIN_REFLECTIONS, partial_mean, g_partial);
@@ -6585,14 +6626,16 @@ RotationScaleMerge::Result RotationScaleMerge::RunArm(bool scale_partials, bool
update.apply = apply;
update.ratio = ratio;
update_owed = true;
ParallelChunks(n_partials - n_grouped, partial_threads, [&](int lo, int hi) {
for (int q = n_grouped + lo; q < n_grouped + hi; ++q) update.Apply(scaling_refs[q], scaling_corr[q]);
});
rest_updates.push_back(update);
});
if (update_owed)
ParallelChunks(n_grouped, partial_threads, [&](int lo, int hi) {
ParallelChunks(n_ref, partial_threads, [&](int lo, int hi) {
for (int q = lo; q < hi; ++q) update.Apply(scaling_refs[q], scaling_corr[q]);
});
ParallelChunks(n_partials - n_ref, partial_threads, [&](int lo, int hi) {
for (int q = n_ref + lo; q < n_ref + hi; ++q)
for (const CorrUpdate &u : rest_updates) u.Apply(scaling_refs[q], scaling_corr[q]);
});
// Back into the partials' own order - the order the memo keeps - in the compact array, and
// from there onto the partials in one sequential pass.
@@ -6925,8 +6968,7 @@ RotationScaleMerge::Result RotationScaleMerge::RunArm(bool scale_partials, bool
for (int i = fulls_frame_start[f]; i < fulls_frame_start[f] + fulls_frame_count[f]; ++i) {
const Obs &o = fulls[i];
if (!fitted_full(o)) continue;
const double weight = SafeInv(o.sigma, 1.0);
full_fit[k++] = {o.group, 1.0, static_cast<double>(o.I), weight * weight};
full_fit[k++] = {o.group, o.I, o.sigma, 1.0f, 1.0f};
}
});
image_analysis/scale_merge/RotationScaleMerge.h
+5 -2
View File
@@ -666,10 +666,13 @@ private:
int32_t frame;
float I, sigma, partiality, prescaling_corr;
};
// The fit's factors that do not depend on the reference - the coefficient of a mean of 1,
// partiality / prescaling_corr (1 for a full), and the weight 1/sigma^2 - are worked out from these
// floats where they are used: twenty bytes an observation read every iteration instead of thirty-two,
// and the same doubles.
struct ScalingFit {
int32_t group;
double coeff_per_mean; // partiality / prescaling_corr (1 for a full): the fit's coeff for a mean of 1
double Iobs, w2;
float I, sigma, partiality, prescaling_corr;
};
// The observation's inverse-variance weight in the scaling reference, and its scaled intensity;