tensor/internal/core/scan_compensation_test.go

// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
// SPDX-License-Identifier: MIT

package core

import (
	"fmt"
	"math"
	"math/big"
	mrand "math/rand/v2"
	"testing"
)

// The compensated-scan evidence file: the sequential float64 CumSum
// chain against the Neumaier-compensated candidate, both measured
// against an exact big.Float referent (256 bits of precision, every
// widening exact). The twins below carry the two walks side by side so
// one binary answers both the accuracy question and the cost question.

// scanLegacyFloats is the sequential chain the Float64 scan arm keeps:
// out[i] = out[i-1] + src[i], one rounding per step.
func scanLegacyFloats(src, dst []float64) {
	acc := src[0]
	dst[0] = acc
	for i := 1; i < len(src); i++ {
		acc += src[i]
		dst[i] = acc
	}
}

// scanNeumaierFloats is the compensated candidate: a running high part
// and a correction, the output the corrected running value. A non-finite
// partial freezes the correction, so an overflow sticks to infinity and
// a NaN poisons the tail exactly as the plain chain does.
func scanNeumaierFloats(src, dst []float64) {
	sum, comp := src[0], 0.0
	dst[0] = sum
	for i := 1; i < len(src); i++ {
		x := src[i]
		t := sum + x
		if !math.IsInf(t, 0) && t == t {
			if math.Abs(sum) >= math.Abs(x) {
				comp += (sum - t) + x
			} else {
				comp += (x - t) + sum
			}
		}
		sum = t
		dst[i] = sum + comp
	}
}

// scanPrefixExact folds the prefixes in big.Float at the given
// precision; every input is a float64, so the widening is exact.
func scanPrefixExact(src []float64, prec uint) []*big.Float {
	out := make([]*big.Float, len(src))
	acc := new(big.Float).SetPrec(prec)
	for i, v := range src {
		acc = new(big.Float).SetPrec(prec).Add(acc, new(big.Float).SetPrec(prec).SetFloat64(v))
		out[i] = acc
	}
	return out
}

// scanError reports the largest deviation of the computed prefixes from
// the exact referent, both absolute and relative to Σ|src|.
func scanError(got []float64, exact []*big.Float) (abs, rel float64) {
	scale := 0.0
	for _, v := range got {
		scale += math.Abs(v)
	}
	for i := range got {
		ref, _ := exact[i].Float64()
		if e := math.Abs(got[i] - ref); e > abs {
			abs = e
		}
	}
	if scale > 0 {
		rel = abs / scale
	}
	return abs, rel
}

// scanAlternatingData builds the cancellation pattern [1, 1e100, 1,
// -1e100] repeated: the plain chain loses every small addend it meets,
// the compensated walk carries them.
func scanAlternatingData(n int) []float64 {
	src := make([]float64, n)
	for i := 0; i < n; i += 4 {
		src[i], src[i+1], src[i+2], src[i+3] = 1, 1e100, 1, -1e100
	}
	return src
}

// scanSpreadData builds a deterministic log-uniform random sample whose
// magnitudes span fifty orders, the shape of data that ages a plain
// chain fastest.
func scanSpreadData(n int, seed uint64) []float64 {
	rng := mrand.New(mrand.NewPCG(seed, seed))
	src := make([]float64, n)
	for i := range src {
		mag := math.Pow(10, -25+50*rng.Float64())
		if rng.Float64() < 0.5 {
			mag = -mag
		}
		src[i] = mag
	}
	return src
}

// TestScanCompensationAccuracy measures both walks against the exact
// referent and pins the compensated result's correctness on the
// cancellation pattern, where the plain chain reports a bare zero.
func TestScanCompensationAccuracy(t *testing.T) {
	// The referent runs at 512 bits: the alternating pattern's values
	// reach 1e100, whose mantissa needs 333 bits, and every prefix must
	// stay exactly representable for the referent to judge ulp errors.
	const prec = 512
	cases := []struct {
		name string
		src  []float64
	}{
		{"alternating 2^20", scanAlternatingData(1 << 20)},
		{"spread 2^20", scanSpreadData(1<<20, 0xC0FFEE)},
		{"spread 2^16", scanSpreadData(1<<16, 7)},
		{"unit 2^20", func() []float64 {
			rng := mrand.New(mrand.NewPCG(42, 42))
			src := make([]float64, 1<<20)
			for i := range src {
				src[i] = rng.Float64()
			}
			return src
		}()},
	}
	for _, tc := range cases {
		exact := scanPrefixExact(tc.src, prec)
		oldGot := make([]float64, len(tc.src))
		newGot := make([]float64, len(tc.src))
		scanLegacyFloats(tc.src, oldGot)
		scanNeumaierFloats(tc.src, newGot)
		oldAbs, oldRel := scanError(oldGot, exact)
		newAbs, newRel := scanError(newGot, exact)
		t.Logf("%s: legacy abs=%.3e rel=%.3e | neumaier abs=%.3e rel=%.3e | improvement %.0fx (abs)",
			tc.name, oldAbs, oldRel, newAbs, newRel, oldAbs/newAbs)
		// A regression guard: the compensated walk is never the worse of
		// the two on these data.
		if newAbs > oldAbs {
			t.Errorf("%s: compensated error %.3e exceeds the chain's %.3e", tc.name, newAbs, oldAbs)
		}
	}
	// The cancellation pin: every complete pattern's prefix returns to
	// 2·k exactly, which the plain chain reports as zero. This pin fails
	// against the uncompenated walk by construction.
	src := []float64{1, 1e100, 1, -1e100, 1, 1e100, 1, -1e100}
	exact := scanPrefixExact(src, prec)
	got := make([]float64, len(src))
	scanNeumaierFloats(src, got)
	for i := range src {
		want, _ := exact[i].Float64()
		if math.Abs(got[i]-want) > 1e-9*math.Max(1, math.Abs(want)) {
			t.Errorf("compensated prefix %d: got %v, want %v", i, got[i], want)
		}
	}
	if last, _ := exact[len(src)-1].Float64(); last != 4 {
		t.Fatalf("referent check: final exact prefix %v, want 4", last)
	}
	// The plain chain's final prefix is 0 here: the document of what the
	// compensation buys.
	legacy := make([]float64, len(src))
	scanLegacyFloats(src, legacy)
	if legacy[len(src)-1] == 4 {
		t.Log("note: the plain chain happened to land exactly on this sample")
	}
}

// TestScanNaNInfSemantics pins the non-finite contract of the
// compensated walk: an overflow sticks to infinity and a NaN poisons
// the tail, exactly as the plain chain reports them.
func TestScanNaNInfSemantics(t *testing.T) {
	cases := [][]float64{
		{math.Inf(1), 1, 2, 3},
		{1, 2, math.Inf(-1), 4},
		{1, math.NaN(), 3, 4},
		{math.MaxFloat64, math.MaxFloat64, 1, 2},
		{1, -1, math.MaxFloat64, math.MaxFloat64},
	}
	for _, src := range cases {
		oldGot := make([]float64, len(src))
		newGot := make([]float64, len(src))
		scanLegacyFloats(src, oldGot)
		scanNeumaierFloats(src, newGot)
		for i := range src {
			wantClass, gotClass := math.Signbit(oldGot[i]), math.Signbit(newGot[i])
			finiteW, finiteG := oldGot[i]-oldGot[i] == 0, newGot[i]-newGot[i] == 0
			if math.IsInf(oldGot[i], 0) != math.IsInf(newGot[i], 0) ||
				math.IsNaN(oldGot[i]) != math.IsNaN(newGot[i]) ||
				(finiteW != finiteG) || (finiteW && wantClass != gotClass) {
				t.Errorf("src %v prefix %d: legacy %v, compensated %v", src, i, oldGot[i], newGot[i])
			}
		}
	}
}

// TestCumSumCompensationProduction pins the shipped entry point on the
// cancellation pattern: the final prefix of two full patterns is 4,
// which the uncompensated chain reports as 0 (the whole sum cancels and
// every small addend is lost on the way). This pin fails against the
// plain-chain walk by construction.
func TestCumSumCompensationProduction(t *testing.T) {
	a, err := FromFloats([]float64{1, 1e100, 1, -1e100, 1, 1e100, 1, -1e100}, 8)
	if err != nil {
		t.Fatal(err)
	}
	got, err := CumSum(a, 0)
	if err != nil {
		t.Fatal(err)
	}
	if last := got.FloatAt(7); math.Abs(last-4) > 1e-9 {
		t.Errorf("compensated CumSum final prefix: %v, want 4", last)
	}
	// The intermediate prefixes carry their compensation too: after the
	// first pattern the running value is 2, not the chain's 0.
	if mid := got.FloatAt(3); math.Abs(mid-2) > 1e-9 {
		t.Errorf("compensated CumSum prefix 3: %v, want 2", mid)
	}
}

// BenchmarkScanCompensated runs the two walks side by side over one
// line of the given length; the sub-benchmarks alternate so one binary
// answers the cost question.
func BenchmarkScanCompensated(b *testing.B) {
	for _, n := range []int{1 << 10, 1 << 16, 1 << 20} {
		src := scanSpreadData(n, 3)
		b.Run(fmt.Sprintf("legacy/n=%d", n), func(b *testing.B) {
			dst := make([]float64, n)
			for b.Loop() {
				scanLegacyFloats(src, dst)
			}
			b.SetBytes(int64(n) * 8)
		})
		b.Run(fmt.Sprintf("neumaier/n=%d", n), func(b *testing.B) {
			dst := make([]float64, n)
			for b.Loop() {
				scanNeumaierFloats(src, dst)
			}
			b.SetBytes(int64(n) * 8)
		})
	}
}