PowerSums / CrossPowerSums: checked_affine — power sums of an affine image from the sums alone - #344

What

Two methods next to the existing checked_merge, both reached through ndarray::simd:

• PowerSums::checked_affine(a, b) returns the power sums of a·x + b.

• CrossPowerSums::checked_affine([a, b, c, d], [tx, ty]) returns the joint sums of A·(x, y) + t, using

  S' = A·S + n·t
  M' = A·M·Aᵀ + A·S·tᵀ + t·Sᵀ·Aᵀ + n·t·tᵀ
  

Both are exact i128 arithmetic with every step checked. They return None instead of wrapping when an intermediate or a narrowed field (i64 / u128 / i128) overflows.

Chains compose: two calls equal one call with A2·A1 and A2·t1 + t2 whenever both return Some. Failure need not agree, because a step can overflow a field that the composed map never produces. Translating by i64::MAX and then back fails stepwise, but composes to the identity.

Why

The lance-graph algebraic_recipe_probe (D-ART-1, merged in lance-graph #1418) measured three things:

• composed affine chains over these summaries are bitwise equal to materialise-and-fold;
• the summary path is 2.5–4× cheaper than materialising the transformed rows;
• a warm summary transform costs 14 ns, independent of the row count.

Nothing in either repo did this before. I read every impl block on origin/master (36ce111): there are only checked_merge, x() and y(). lance-graph's jc::stats only reads summaries.

Precondition

The group key must not depend on a transformed column. This is documented, and a test pins it: with the key x ≥ 0, shifting x by 20 moves rows between groups, so per-group transforms differ from the truth.

Evidence

• Equivalence: 400 random cases match folding the transformed rows through masked_group_cross_power_sums_i32. Each case has 1–300 rows, coefficients in [−5, 5] (reflections included) and translations in [−100, 100], under random masks and 3 groups. The univariate form is checked against the x marginal.

• Composition: two calls equal the composed call when both succeed. The asymmetric-failure case is pinned. The identity map leaves the sums unchanged, and the zero map leaves only the translation.

• Overflow: refused in every narrowed field.

• Doctests: 2.

• Disable runs, each red then restored:

  • the 2·tx·(A·S)x term dropped;
  • the ty·(A·S)x cross term dropped;
  • n·tx dropped;
  • the univariate 2ab·Σx term dropped;
  • as i64 instead of checked narrowing, separately for sum_x and sum_y.

  The first overflow fixture was vacuous: its translation already overflowed the square, so the i64 narrowing never ran. It was rebuilt so that only the narrowing can refuse.

• Full run: cargo test -p ndarray --lib passed 2547 (32 ignored), with debug 0. clippy -D warnings and fmt are clean.

This is scalar code with no target-specific paths. It is a struct method on a summary type, like checked_merge, not a new SIMD primitive.

CI note: the realization/nightly × x86_64 failure on the first head was a SIGILL inside the cached paste proc-macro. It was not caused by this change; see the comment below for the cause and a proposed workflow fix.

🤖 Generated with Claude Code

https://claude.ai/code/session_01EFw2WdKr1oxvaKCJC2ua2R