433 lines
14 KiB
Go
433 lines
14 KiB
Go
// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
|
||
// SPDX-License-Identifier: MIT
|
||
|
||
package stats
|
||
|
||
import (
|
||
"sourcedock.dev/petrbalvin/tensor/internal/base"
|
||
"sourcedock.dev/petrbalvin/tensor/internal/core"
|
||
)
|
||
|
||
import (
|
||
"math"
|
||
"slices"
|
||
)
|
||
|
||
// Statistical inference on samples: association matrices, group
|
||
// comparison, goodness of fit, distribution comparison and the
|
||
// percentile bootstrap. Everything here is classical frequentist
|
||
// tooling built on the distribution functions of cdf.go, so a
|
||
// p-value travels no further than the library's own incomplete
|
||
// gamma and beta.
|
||
|
||
// CovarianceMatrix returns the sample covariance matrix of the
|
||
// observations in a, an (n, p) array whose rows are observations and
|
||
// columns variables: the (p, p) result has cov(i, j) the sample
|
||
// covariance of columns i and j with the 1/(n−1) normalisation. At
|
||
// least two observations are needed.
|
||
func CovarianceMatrix(a *core.Array) (*core.Array, error) {
|
||
return scatterMatrix("CovarianceMatrix", a, false)
|
||
}
|
||
|
||
// CorrelationMatrix returns the Pearson correlation matrix of the
|
||
// observations in a, the covariance matrix normalised by each
|
||
// column's sample standard deviation; the diagonal is exactly 1.
|
||
func CorrelationMatrix(a *core.Array) (*core.Array, error) {
|
||
return scatterMatrix("CorrelationMatrix", a, true)
|
||
}
|
||
|
||
// scatterMatrix builds the covariance or correlation matrix of an
|
||
// (n, p) observation array.
|
||
func scatterMatrix(name string, a *core.Array, correlate bool) (*core.Array, error) {
|
||
if a.NDim() != 2 {
|
||
return nil, base.Errf("%s: needs a 2-D array of observations, got shape %s",
|
||
name, base.ShapeText(a.Shape()))
|
||
}
|
||
n, p := a.Shape()[0], a.Shape()[1]
|
||
if n < 2 {
|
||
return nil, base.Errf("%s: at least two observations are needed, got %d", name, n)
|
||
}
|
||
if a.Dtype() == core.Complex {
|
||
return nil, base.Errf("%s: complex observations are not supported", name)
|
||
}
|
||
if err := checkFinite(name, "the observations", a); err != nil {
|
||
return nil, err
|
||
}
|
||
// Column means, then centred column copies. A dense float64 payload
|
||
// is read straight from the raw slice, one dispatch-free pass.
|
||
rows := rawFloats(a)
|
||
means := make([]float64, p)
|
||
for j := range p {
|
||
s := 0.0
|
||
if rows != nil {
|
||
for i := range n {
|
||
s += rows[i*p+j]
|
||
}
|
||
} else {
|
||
for i := range n {
|
||
s += a.FloatAt(i*p + j)
|
||
}
|
||
}
|
||
means[j] = s / float64(n)
|
||
}
|
||
cols := make([][]float64, p)
|
||
for j := range p {
|
||
cols[j] = make([]float64, n)
|
||
if rows != nil {
|
||
for i := range n {
|
||
cols[j][i] = rows[i*p+j] - means[j]
|
||
}
|
||
} else {
|
||
for i := range n {
|
||
cols[j][i] = a.FloatAt(i*p+j) - means[j]
|
||
}
|
||
}
|
||
}
|
||
scale := float64(n - 1)
|
||
cov := make([]float64, p*p)
|
||
stds := make([]float64, p)
|
||
for i := range p {
|
||
for j := i; j < p; j++ {
|
||
s := 0.0
|
||
for k := range n {
|
||
s += cols[i][k] * cols[j][k]
|
||
}
|
||
cov[i*p+j] = s / scale
|
||
cov[j*p+i] = cov[i*p+j]
|
||
}
|
||
if correlate {
|
||
stds[i] = math.Sqrt(cov[i*p+i])
|
||
if stds[i] == 0 {
|
||
return nil, base.Errf("CorrelationMatrix: column %d has zero variance", i)
|
||
}
|
||
}
|
||
}
|
||
if correlate {
|
||
for i := range p {
|
||
for j := i; j < p; j++ {
|
||
c := cov[i*p+j] / (stds[i] * stds[j])
|
||
cov[i*p+j] = c
|
||
cov[j*p+i] = c
|
||
}
|
||
cov[i*p+i] = 1
|
||
}
|
||
}
|
||
return floatsToArray(cov, []int{p, p}), nil
|
||
}
|
||
|
||
// WelchTTest compares two independent samples with Welch's t-test,
|
||
// which asks no equal-variance assumption: the statistic is
|
||
// (mean(a) − mean(b)) over the pooled standard error of the two
|
||
// means, and the degrees of freedom follow the Welch-Satterthwaite
|
||
// estimate. Returns the statistic, the (generally fractional) degrees
|
||
// of freedom and the two-sided p-value. Both samples need at least
|
||
// two observations.
|
||
func WelchTTest(a, b *core.Array) (t, df, pValue float64, err error) {
|
||
m1, v1, n1, err := sampleMeanVar(a, "WelchTTest")
|
||
if err != nil {
|
||
return 0, 0, 0, err
|
||
}
|
||
m2, v2, n2, err := sampleMeanVar(b, "WelchTTest")
|
||
if err != nil {
|
||
return 0, 0, 0, err
|
||
}
|
||
se1 := v1 / float64(n1)
|
||
se2 := v2 / float64(n2)
|
||
se := se1 + se2
|
||
if se == 0 {
|
||
return 0, 0, 0, base.Errf("WelchTTest: both samples have zero variance")
|
||
}
|
||
t = (m1 - m2) / math.Sqrt(se)
|
||
df = se * se / (se1*se1/float64(n1-1) + se2*se2/float64(n2-1))
|
||
// Two-sided p = P(T > |t|) + P(T < −|t|) = I_z(df/2, 1/2) with
|
||
// z = df/(df + t²), the Student-t tail in closed form.
|
||
z := df / (df + t*t)
|
||
pValue, err = BetaIncomplete(z, df/2, 0.5)
|
||
if err != nil {
|
||
return 0, 0, 0, base.Errf("WelchTTest: %w", err)
|
||
}
|
||
return t, df, pValue, nil
|
||
}
|
||
|
||
// sampleMeanVar returns the mean and the 1/(n−1) variance of a sample.
|
||
func sampleMeanVar(a *core.Array, name string) (mean, variance float64, n int, err error) {
|
||
if a.Dtype() == core.Complex {
|
||
return 0, 0, 0, base.Errf("%s: complex samples are not supported", name)
|
||
}
|
||
n = a.Len()
|
||
if n < 2 {
|
||
return 0, 0, 0, base.Errf("%s: at least two observations are needed, got %d", name, n)
|
||
}
|
||
// The finiteness refusal belongs here, before the arithmetic: a NaN
|
||
// sample would otherwise surface only when the tail function
|
||
// rejects the NaN statistic it produced, under the tail function's
|
||
// own name.
|
||
if err := checkFinite(name, "the sample", a); err != nil {
|
||
return 0, 0, 0, err
|
||
}
|
||
mean = 0.0
|
||
// The walk is bounded by the sample's own element count: a rebased
|
||
// view's payload may run past its visible elements, and those
|
||
// invisible tail slots are nobody's observations.
|
||
if fs := rawFloats(a); fs != nil {
|
||
fs = fs[:n]
|
||
for _, v := range fs {
|
||
mean += v
|
||
}
|
||
mean /= float64(n)
|
||
variance = sqDeviations(fs, mean)
|
||
} else {
|
||
for i := range n {
|
||
mean += a.FloatAt(i)
|
||
}
|
||
mean /= float64(n)
|
||
variance = sqDeviationsAt(a, mean)
|
||
}
|
||
variance /= float64(n - 1)
|
||
return mean, variance, n, nil
|
||
}
|
||
|
||
// ChiSquareGoodnessOfFit runs Pearson's test of an observed frequency
|
||
// table against expected ones: the statistic is Σ(O−E)²/E over the
|
||
// bins, the degrees of freedom the bin count minus one, and the
|
||
// p-value the upper tail of that χ² distribution. Every expected
|
||
// entry must be positive.
|
||
func ChiSquareGoodnessOfFit(observed, expected *core.Array) (chi2 float64, df int, pValue float64, err error) {
|
||
if observed.Dtype() == core.Complex || expected.Dtype() == core.Complex {
|
||
return 0, 0, 0, base.Errf("ChiSquareGoodnessOfFit: complex arrays are not supported")
|
||
}
|
||
n := observed.Len()
|
||
if expected.Len() != n {
|
||
return 0, 0, 0, base.Errf("ChiSquareGoodnessOfFit: observed has %d bins, expected %d",
|
||
n, expected.Len())
|
||
}
|
||
if n < 2 {
|
||
return 0, 0, 0, base.Errf("ChiSquareGoodnessOfFit: at least two bins are needed, got %d", n)
|
||
}
|
||
// A NaN observation would only surface later, rejected by the
|
||
// tail function under its own name; the scan keeps the error with
|
||
// the entry point the caller actually used.
|
||
if err := checkFinite("ChiSquareGoodnessOfFit", "the observed frequencies", observed); err != nil {
|
||
return 0, 0, 0, err
|
||
}
|
||
// A non-finite expectation passes the positivity gate below (+Inf
|
||
// is positive) and turns the statistic into a NaN the tail function
|
||
// would reject under its own, misleading, name.
|
||
if err := checkFinite("ChiSquareGoodnessOfFit", "the expected frequencies", expected); err != nil {
|
||
return 0, 0, 0, err
|
||
}
|
||
chi2 = 0.0
|
||
// Both walks are bounded by the bin count: a rebased view's payload
|
||
// may run past its visible bins, and indexing one past the other's
|
||
// payload would be a panic before it was a wrong statistic.
|
||
obs := rawFloats(observed)
|
||
exp := rawFloats(expected)
|
||
if obs != nil {
|
||
obs = obs[:n]
|
||
}
|
||
if exp != nil {
|
||
exp = exp[:n]
|
||
}
|
||
if obs != nil && exp != nil {
|
||
for i, e := range exp {
|
||
if !(e > 0) {
|
||
return 0, 0, 0, base.Errf("ChiSquareGoodnessOfFit: expected[%d] = %g must be positive", i, e)
|
||
}
|
||
d := obs[i] - e
|
||
chi2 += d * d / e
|
||
}
|
||
} else {
|
||
for i := range n {
|
||
e := expected.FloatAt(i)
|
||
if !(e > 0) {
|
||
return 0, 0, 0, base.Errf("ChiSquareGoodnessOfFit: expected[%d] = %g must be positive", i, e)
|
||
}
|
||
d := observed.FloatAt(i) - e
|
||
chi2 += d * d / e
|
||
}
|
||
}
|
||
df = n - 1
|
||
pValue, err = GammaUpper(float64(df)/2, chi2/2)
|
||
if err != nil {
|
||
return 0, 0, 0, base.Errf("ChiSquareGoodnessOfFit: %w", err)
|
||
}
|
||
return chi2, df, pValue, nil
|
||
}
|
||
|
||
// KolmogorovSmirnovTest compares two samples by the largest vertical
|
||
// distance d between their empirical distribution functions. The
|
||
// p-value is the asymptotic Kolmogorov distribution evaluated at
|
||
// λ = d·√(nm/(n+m)), accurate for samples of a few dozen and upward;
|
||
// small samples only get an order-of-magnitude answer. Both samples
|
||
// must be finite: a NaN never compares true against the support point
|
||
// the merge walk advances on, so it would leave the walk stuck, and an
|
||
// ±Inf would distort the distance.
|
||
func KolmogorovSmirnovTest(a, b *core.Array) (d, pValue float64, err error) {
|
||
const name = "KolmogorovSmirnovTest"
|
||
if a.Len() == 0 || b.Len() == 0 {
|
||
return 0, 0, base.Errf("%s: both samples must be non-empty", name)
|
||
}
|
||
if a.Dtype() == core.Complex || b.Dtype() == core.Complex {
|
||
return 0, 0, base.Errf("%s: complex samples are not supported", name)
|
||
}
|
||
if err := checkFinite(name, "the first sample", a); err != nil {
|
||
return 0, 0, err
|
||
}
|
||
if err := checkFinite(name, "the second sample", b); err != nil {
|
||
return 0, 0, err
|
||
}
|
||
xs := make([]float64, a.Len())
|
||
if fs := rawFloats(a); fs != nil {
|
||
copy(xs, fs)
|
||
} else {
|
||
for i := range a.Len() {
|
||
xs[i] = a.FloatAt(i)
|
||
}
|
||
}
|
||
ys := make([]float64, b.Len())
|
||
if fs := rawFloats(b); fs != nil {
|
||
copy(ys, fs)
|
||
} else {
|
||
for i := range b.Len() {
|
||
ys[i] = b.FloatAt(i)
|
||
}
|
||
}
|
||
slices.Sort(xs)
|
||
slices.Sort(ys)
|
||
// Walk the merged support, tracking both empirical CDFs.
|
||
i, j := 0, 0
|
||
d = 0.0
|
||
for i < len(xs) || j < len(ys) {
|
||
var x float64
|
||
switch {
|
||
case i < len(xs) && (j >= len(ys) || xs[i] <= ys[j]):
|
||
x = xs[i]
|
||
default:
|
||
x = ys[j]
|
||
}
|
||
for i < len(xs) && xs[i] <= x {
|
||
i++
|
||
}
|
||
for j < len(ys) && ys[j] <= x {
|
||
j++
|
||
}
|
||
if gap := math.Abs(float64(i)/float64(len(xs)) - float64(j)/float64(len(ys))); gap > d {
|
||
d = gap
|
||
}
|
||
}
|
||
en := math.Sqrt(float64(a.Len()) * float64(b.Len()) / float64(a.Len()+b.Len()))
|
||
pValue = kolmogorovTail(d * en)
|
||
return d, pValue, nil
|
||
}
|
||
|
||
// kolmogorovEps is the magnitude at which a term of the Kolmogorov
|
||
// series is below the rounding level of the sum, so the alternating
|
||
// series can be truncated there: the first omitted term bounds the
|
||
// truncation error, and the terms decrease from the first one on.
|
||
const kolmogorovEps = 1e-18
|
||
|
||
// maxKolmogorovTerms caps the series for the smallest lambdas that
|
||
// reach it at all.
|
||
const maxKolmogorovTerms = 100
|
||
|
||
// kolmogorovTermCount returns how many terms of the Kolmogorov series
|
||
// matter at lambda: the first k whose term 2·exp(−2k²λ²) has fallen
|
||
// below the rounding level. The series itself sums exactly that many,
|
||
// so the count is the truncation rule in one place instead of a
|
||
// condition the loop cannot reach.
|
||
func kolmogorovTermCount(lambda float64) int {
|
||
for k := 1; k <= maxKolmogorovTerms; k++ {
|
||
if 2*math.Exp(-2*float64(k)*float64(k)*lambda*lambda) < kolmogorovEps {
|
||
return k
|
||
}
|
||
}
|
||
return maxKolmogorovTerms
|
||
}
|
||
|
||
// kolmogorovTail evaluates the asymptotic Kolmogorov distribution
|
||
// Q(λ) = 2·Σ_{k≥1} (−1)^{k−1}·e^{−2k²λ²}, summing the terms that
|
||
// matter.
|
||
func kolmogorovTail(lambda float64) float64 {
|
||
if lambda < 0.2 {
|
||
return 1
|
||
}
|
||
total := 0.0
|
||
for k := 1; k <= kolmogorovTermCount(lambda); k++ {
|
||
term := 2 * math.Exp(-2*float64(k)*float64(k)*lambda*lambda)
|
||
if k%2 == 0 {
|
||
total -= term
|
||
} else {
|
||
total += term
|
||
}
|
||
}
|
||
return min(1, max(0, total))
|
||
}
|
||
|
||
// BootstrapCI estimates a confidence interval of a statistic by the
|
||
// percentile bootstrap: resamples the data with replacement
|
||
// resamples times from the seeded generator, evaluates the statistic
|
||
// on every resample and returns the alpha/2 and 1−alpha/2 quantiles
|
||
// of the bootstrap distribution, alpha = 1 − level. The statistic
|
||
// receives the resample as a fresh array and may return its own
|
||
// error. The data must be real-valued.
|
||
func BootstrapCI(data *core.Array, statistic func(*core.Array) (float64, error), level float64,
|
||
resamples int, seed int64) (lower, upper float64, err error) {
|
||
if data.Len() == 0 {
|
||
return 0, 0, base.Errf("BootstrapCI: the data must not be empty")
|
||
}
|
||
if data.Dtype() == core.Complex {
|
||
return 0, 0, base.Errf("BootstrapCI: complex data are not supported")
|
||
}
|
||
if !(level > 0 && level < 1) {
|
||
return 0, 0, base.Errf("BootstrapCI: level must lie in (0, 1), got %g", level)
|
||
}
|
||
if resamples < 2 {
|
||
return 0, 0, base.Errf("BootstrapCI: resamples must be ≥ 2, got %d", resamples)
|
||
}
|
||
g := core.NewGenerator(seed)
|
||
n := data.Len()
|
||
// The source values are read once and the index payload hoisted:
|
||
// the resample loop then indexes plain slices.
|
||
dataVals := rawFloats(data)
|
||
values := make([]float64, resamples)
|
||
for r := range resamples {
|
||
idx, ierr := core.Ints(g, n, 0, int64(n))
|
||
if ierr != nil {
|
||
return 0, 0, base.Errf("BootstrapCI: %w", ierr)
|
||
}
|
||
idxVals := idx.RawInts()
|
||
// A fresh slice per resample honours the documented contract: a
|
||
// statistic that retains its argument must not observe the next
|
||
// resample's mutation through the alias, and the resample is
|
||
// written into that slice directly rather than into a scratch the
|
||
// caller then copies.
|
||
sample := make([]float64, n)
|
||
if dataVals != nil {
|
||
for i := range n {
|
||
sample[i] = dataVals[idxVals[i]]
|
||
}
|
||
} else {
|
||
for i := range n {
|
||
sample[i] = data.FloatAt(int(idxVals[i]))
|
||
}
|
||
}
|
||
v, serr := statistic(wrapVector(sample))
|
||
if serr != nil {
|
||
return 0, 0, base.Errf("BootstrapCI: %w", serr)
|
||
}
|
||
values[r] = v
|
||
}
|
||
alpha := (1 - level) / 2
|
||
qs, qerr := Quantile(floatsToArray(values, []int{resamples}), []float64{alpha, 1 - alpha})
|
||
if qerr != nil {
|
||
return 0, 0, base.Errf("BootstrapCI: %w", qerr)
|
||
}
|
||
return qs.FloatAt(0), qs.FloatAt(1), nil
|
||
}
|
||
|
||
// wrapVector views a float64 slice as a rank-1 Array without copying.
|
||
func wrapVector(v []float64) *core.Array {
|
||
a, _ := core.FloatsFromArray(v, len(v))
|
||
return a
|
||
}
|