485 lines
15 KiB
Go
485 lines
15 KiB
Go
// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
|
||
// SPDX-License-Identifier: MIT
|
||
|
||
package stats
|
||
|
||
import (
|
||
"math"
|
||
"testing"
|
||
|
||
"sourcedock.dev/petrbalvin/tensor/internal/core"
|
||
)
|
||
|
||
// The mixed model against referents. The balanced one-way random
|
||
// effects model has a closed-form REML answer, the analysis of
|
||
// variance estimators, so the sweep's optimum is compared against
|
||
// figures computed from the raw data rather than quoted; the rest
|
||
// pins the recovery of known effects, the refusal surface and the
|
||
// determinism of the whole pipeline.
|
||
|
||
// mixedNoise is a deterministic stand-in for measurement noise: a
|
||
// bounded, aperiodic wiggle no sweep can mistake for structure.
|
||
func mixedNoise(i int) float64 {
|
||
return 0.3*math.Sin(7.3*float64(i)+1.1)*math.Cos(2.1*float64(i)) +
|
||
0.1*math.Sin(0.7*float64(i))
|
||
}
|
||
|
||
// mixedJitter is deterministic white jitter on [−1, 1): the xorshift
|
||
// finaliser of the house generator's mixing constants, run on the row
|
||
// index. Unlike the smooth wiggle above it cannot be absorbed by a
|
||
// within-group linear span, which is what the random-slope fit needs
|
||
// its residual scale to be.
|
||
func mixedJitter(i int) float64 {
|
||
z := uint64(i)*2685821657736338717 + 1
|
||
z ^= z >> 13
|
||
z ^= z << 7
|
||
z ^= z >> 17
|
||
return float64(z>>11)/(1<<52)*2 - 1
|
||
}
|
||
|
||
func mixedVec(t *testing.T, vals []float64, shape ...int) *core.Array {
|
||
t.Helper()
|
||
a, err := core.FromFloats(vals, shape...)
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
return a
|
||
}
|
||
|
||
func TestMixedModelBalancedOneWayANOVAReferent(t *testing.T) {
|
||
// The balanced one-way random effects model: y_ig = μ + b_g + ε_ig
|
||
// with k observations in each of m groups. REML's optimum is the
|
||
// ANOVA answer: σ̂²_e = MSW and σ̂²_b = (MSB − MSW)/k, and the
|
||
// intercept's GLS variance is (k·σ²_b + σ²_e)/(mk).
|
||
const (
|
||
mGroups = 8
|
||
k = 6
|
||
)
|
||
y := make([]float64, 0, mGroups*k)
|
||
for g := range mGroups {
|
||
effect := 2.0 * math.Sin(1.7*float64(g)+0.4) // the drawn b_g
|
||
for i := range k {
|
||
y = append(y, 5+effect+mixedNoise(g*k+i))
|
||
}
|
||
}
|
||
groups := make([]int, mGroups*k)
|
||
for g := range mGroups {
|
||
for i := range k {
|
||
groups[g*k+i] = g
|
||
}
|
||
}
|
||
res, err := LinearMixedModel(
|
||
mixedVec(t, y, len(y)),
|
||
mixedVec(t, ones(len(y)), len(y), 1),
|
||
mixedVec(t, ones(len(y)), len(y), 1),
|
||
groups)
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
if !res.Converged {
|
||
t.Fatalf("the balanced fit did not converge (%d iterations)", res.Iterations)
|
||
}
|
||
// MSW and MSB from the raw data.
|
||
grand := 0.0
|
||
groupMeans := make([]float64, mGroups)
|
||
for g := range mGroups {
|
||
s := 0.0
|
||
for i := range k {
|
||
s += y[g*k+i]
|
||
}
|
||
groupMeans[g] = s / float64(k)
|
||
grand += groupMeans[g]
|
||
}
|
||
grand /= float64(mGroups)
|
||
msw := 0.0
|
||
for g := range mGroups {
|
||
for i := range k {
|
||
d := y[g*k+i] - groupMeans[g]
|
||
msw += d * d
|
||
}
|
||
}
|
||
msw /= float64(mGroups * (k - 1))
|
||
msb := 0.0
|
||
for g := range mGroups {
|
||
d := groupMeans[g] - grand
|
||
msb += d * d
|
||
}
|
||
msb *= float64(k) / float64(mGroups-1)
|
||
wantWithin := math.Max((msb-msw)/float64(k), 0)
|
||
if math.Abs(res.ResidualVariance-msw) > 0.02*msw {
|
||
t.Fatalf("σ̂²_e = %.6f, want the MSW %.6f", res.ResidualVariance, msw)
|
||
}
|
||
if math.Abs(res.RandomCovariance[0]-wantWithin) > 0.05*wantWithin {
|
||
t.Fatalf("σ̂²_b = %.6f, want the ANOVA answer %.6f", res.RandomCovariance[0], wantWithin)
|
||
}
|
||
if math.Abs(res.Coefficients[0]-grand) > 1e-6 {
|
||
t.Fatalf("μ̂ = %.8f, want the grand mean %.8f", res.Coefficients[0], grand)
|
||
}
|
||
wantVar := (float64(k)*res.RandomCovariance[0] + res.ResidualVariance) / float64(mGroups*k)
|
||
if se := res.StandardErrors[0]; math.Abs(se*se-wantVar) > 1e-9*math.Max(1, wantVar) {
|
||
t.Fatalf("SE² = %.10f, want the GLS variance %.10f", se*se, wantVar)
|
||
}
|
||
// The conditional fitted values reproduce the group means plus the
|
||
// shrinkage the model applies; the residuals must complement them
|
||
// to the response.
|
||
for i := range len(y) {
|
||
if math.Abs(res.Fitted[i]+res.Residuals[i]-y[i]) > 1e-9 {
|
||
t.Fatalf("row %d: fitted + residuals = %g, want %g", i, res.Fitted[i]+res.Residuals[i], y[i])
|
||
}
|
||
}
|
||
}
|
||
|
||
// ones returns n constant 1 values, the intercept column.
|
||
func ones(n int) []float64 {
|
||
out := make([]float64, n)
|
||
for i := range out {
|
||
out[i] = 1
|
||
}
|
||
return out
|
||
}
|
||
|
||
func TestMixedModelRandomSlopeRecovery(t *testing.T) {
|
||
// Fixed effects of 1 and 2 with a per-group random slope, built by
|
||
// hand so the truth is known exactly. The random design carries
|
||
// the covariate alone: with the intercept column beside it the
|
||
// random span covers the fixed design and the REML surface loses
|
||
// its interior optimum to a ridge of singular covariance (pinned
|
||
// by the divergence test below). The slope values are centred, so
|
||
// the fixed part of the truth is exactly (1, 2).
|
||
slopeValues := []float64{0.9, -1.1, 1.9, -0.3, -1.8, 0.4}
|
||
groups := make([]int, 0, 36)
|
||
xs := make([]float64, 0, 36)
|
||
y := make([]float64, 0, 36)
|
||
jitter := make([]float64, 0, 36)
|
||
for g := range slopeValues {
|
||
for _, x := range []float64{-1, -0.6, -0.2, 0.2, 0.6, 1} {
|
||
groups = append(groups, g)
|
||
xs = append(xs, x)
|
||
jitter = append(jitter, mixedJitter(len(xs)-1))
|
||
}
|
||
}
|
||
// The jitter is centred on its own sample: the fixed part of the
|
||
// truth must stay exactly (1, 2), and a noise vector with a mean
|
||
// would tilt the intercept instead of testing the recovery.
|
||
mean := 0.0
|
||
for _, j := range jitter {
|
||
mean += j
|
||
}
|
||
mean /= float64(len(jitter))
|
||
for g, s := range slopeValues {
|
||
for k := range 6 {
|
||
i := g*6 + k
|
||
x := xs[i]
|
||
y = append(y, 1+2*x+0.5*s*x+0.1*(jitter[i]-mean))
|
||
}
|
||
}
|
||
design := make([]float64, 0, 2*len(xs))
|
||
for _, x := range xs {
|
||
design = append(design, 1, x)
|
||
}
|
||
res, err := LinearMixedModel(
|
||
mixedVec(t, y, len(y)),
|
||
mixedVec(t, design, len(xs), 2),
|
||
mixedVec(t, xs, len(xs), 1),
|
||
groups)
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
if !res.Converged {
|
||
t.Fatalf("the slope fit did not converge (%d iterations)", res.Iterations)
|
||
}
|
||
if math.Abs(res.Coefficients[0]-1) > 0.05 || math.Abs(res.Coefficients[1]-2) > 0.05 {
|
||
t.Fatalf("β̂ = (%.4f, %.4f), want (1, 2)", res.Coefficients[0], res.Coefficients[1])
|
||
}
|
||
// The random slopes must rank with the true ones, and the group
|
||
// labels must come back in first-appearance order.
|
||
if len(res.GroupLabels) != 6 {
|
||
t.Fatalf("group labels %v, want six groups", res.GroupLabels)
|
||
}
|
||
strongest := 0
|
||
weakest := 0
|
||
for g, s := range slopeValues {
|
||
if s > slopeValues[strongest] {
|
||
strongest = g
|
||
}
|
||
if s < slopeValues[weakest] {
|
||
weakest = g
|
||
}
|
||
}
|
||
if !(res.RandomEffects[weakest][0] < res.RandomEffects[strongest][0]) {
|
||
t.Fatalf("the random slopes do not rank with the truth (%v against %v)",
|
||
res.RandomEffects[weakest][0], slopeValues[strongest])
|
||
}
|
||
for _, se := range res.StandardErrors {
|
||
if !(se > 0) || math.IsInf(se, 0) {
|
||
t.Fatalf("the standard error %g is not finite and positive", se)
|
||
}
|
||
}
|
||
}
|
||
|
||
func TestMixedModelDivergenceRefusal(t *testing.T) {
|
||
// A random design that spans the fixed one under an unstructured
|
||
// covariance: the intercept and slope of every group absorb what
|
||
// the fixed effects name, and the REML surface climbs a ridge of
|
||
// singular Σ without a summit. The fit refuses with the condition
|
||
// named instead of publishing the climb.
|
||
groups := []int{0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2}
|
||
xs := []float64{-1, -0.5, 0.5, 1, -1, -0.5, 0.5, 1, -1, -0.5, 0.5, 1}
|
||
y := make([]float64, len(groups))
|
||
for i, g := range groups {
|
||
y[i] = 1 + 2*xs[i] + 0.4*float64(g)*xs[i] + float64(g) + 0.1*mixedNoise(i)
|
||
}
|
||
design := make([]float64, 0, 2*len(xs))
|
||
for _, x := range xs {
|
||
design = append(design, 1, x)
|
||
}
|
||
_, err := LinearMixedModel(
|
||
mixedVec(t, y, len(y)),
|
||
mixedVec(t, design, len(xs), 2),
|
||
mixedVec(t, design, len(xs), 2),
|
||
groups)
|
||
if err == nil {
|
||
t.Fatal("a saturated random design was accepted")
|
||
}
|
||
}
|
||
|
||
func TestMixedModelMatchesOLSWithoutRandomEffects(t *testing.T) {
|
||
// When the groups carry no shared signal the components collapse
|
||
// towards zero and the fit must land on the plain least squares
|
||
// answer.
|
||
xs := make([]float64, 18)
|
||
y := make([]float64, 18)
|
||
for i := range 18 {
|
||
x := -1 + 2*float64(i)/17
|
||
xs[i] = x
|
||
y[i] = 3 - x + 0.4*mixedNoise(i)
|
||
}
|
||
groups := make([]int, 18)
|
||
for i := range 18 {
|
||
groups[i] = i % 6
|
||
}
|
||
design := make([]float64, 0, 36)
|
||
for _, x := range xs {
|
||
design = append(design, 1, x)
|
||
}
|
||
res, err := LinearMixedModel(
|
||
mixedVec(t, y, len(y)),
|
||
mixedVec(t, design, len(xs), 2),
|
||
mixedVec(t, design, len(xs), 2),
|
||
groups)
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
ref, rerr := LinearRegression(mixedVec(t, design, len(xs), 2), mixedVec(t, y, len(y)))
|
||
if rerr != nil {
|
||
t.Fatal(rerr)
|
||
}
|
||
for j := range 2 {
|
||
if math.Abs(res.Coefficients[j]-ref.Coefficients[j]) > 0.01 {
|
||
t.Fatalf("coefficient %d = %.6f, want the OLS %.6f", j, res.Coefficients[j], ref.Coefficients[j])
|
||
}
|
||
}
|
||
if res.RandomCovariance[0] > 0.01 {
|
||
t.Fatalf("Σ̂ = %g on a group-free sample, want a collapsed component", res.RandomCovariance[0])
|
||
}
|
||
}
|
||
|
||
func TestMixedModelDeterministic(t *testing.T) {
|
||
y := make([]float64, 16)
|
||
groups := make([]int, 16)
|
||
for i := range 16 {
|
||
y[i] = 2 + 0.9*math.Sin(float64(i%4)) + mixedNoise(i)
|
||
groups[i] = i / 4
|
||
}
|
||
run := func() *LinearMixedModelResult {
|
||
res, err := LinearMixedModel(
|
||
mixedVec(t, y, len(y)),
|
||
mixedVec(t, ones(len(y)), len(y), 1),
|
||
mixedVec(t, ones(len(y)), len(y), 1),
|
||
groups)
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
return res
|
||
}
|
||
a, b := run(), run()
|
||
if a.LogLikelihood != b.LogLikelihood || a.ResidualVariance != b.ResidualVariance ||
|
||
a.RandomCovariance[0] != b.RandomCovariance[0] || a.Iterations != b.Iterations {
|
||
t.Fatal("two identical fits disagreed")
|
||
}
|
||
for i := range a.Fitted {
|
||
if a.Fitted[i] != b.Fitted[i] {
|
||
t.Fatalf("row %d: fitted %g against %g", i, a.Fitted[i], b.Fitted[i])
|
||
}
|
||
}
|
||
}
|
||
|
||
func TestMixedModelShuffledLabels(t *testing.T) {
|
||
// Labels out of order and with gaps must canonicalise by first
|
||
// appearance, and permuting the rows with their labels must leave
|
||
// the fitted components where they started.
|
||
y := []float64{1.0, 1.2, 3.0, 3.1, 5.2, 5.1, 1.1, 3.2, 5.0, 1.3}
|
||
groups := []int{7, 7, 3, 3, 5, 5, 7, 3, 5, 7}
|
||
res, err := LinearMixedModel(
|
||
mixedVec(t, y, len(y)),
|
||
mixedVec(t, ones(len(y)), len(y), 1),
|
||
mixedVec(t, ones(len(y)), len(y), 1),
|
||
groups)
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
want := []int{7, 3, 5}
|
||
for i, label := range res.GroupLabels {
|
||
if label != want[i] {
|
||
t.Fatalf("group labels %v, want %v", res.GroupLabels, want)
|
||
}
|
||
}
|
||
if !(res.RandomEffects[0][0] < res.RandomEffects[1][0] && res.RandomEffects[1][0] < res.RandomEffects[2][0]) {
|
||
t.Fatalf("the random effects %v do not rank with the group means", res.RandomEffects)
|
||
}
|
||
}
|
||
|
||
func TestMixedModelRefusals(t *testing.T) {
|
||
y := mixedVec(t, []float64{1, 2, 3, 4}, 4)
|
||
x := mixedVec(t, []float64{1, 1, 1, 1}, 4, 1)
|
||
z := mixedVec(t, []float64{1, 1, 1, 1}, 4, 1)
|
||
groups := []int{0, 0, 1, 1}
|
||
if _, err := LinearMixedModel(mixedVec(t, []float64{1, 2}, 2, 1), x, z, groups); err == nil {
|
||
t.Fatal("a rank-2 response was accepted")
|
||
}
|
||
if _, err := LinearMixedModel(mixedVec(t, []float64{1, 2, 3, 4, 5}, 5), x, z, groups); err == nil {
|
||
t.Fatal("a row count mismatch was accepted")
|
||
}
|
||
if _, err := LinearMixedModel(y, mixedVec(t, []float64{1, 1, 1, 1}, 4, 1), mixedVec(t, []float64{1, 1, 1, 1}, 4, 1), []int{0, 1, 2}); err == nil {
|
||
t.Fatal("a short label vector was accepted")
|
||
}
|
||
if _, err := LinearMixedModel(y, x, z, []int{0, 0, 1, -2}); err == nil {
|
||
t.Fatal("a negative label was accepted")
|
||
}
|
||
// More coefficients than observations.
|
||
big := mixedVec(t, []float64{1, 0, 1, 0, 1, 0, 1, 0, 1, 1, 1, 1}, 4, 3)
|
||
if _, err := LinearMixedModel(y, big, z, groups); err == nil {
|
||
t.Fatal("a saturated design was accepted")
|
||
}
|
||
// A singular fixed design: two identical columns.
|
||
sing := mixedVec(t, []float64{1, 1, 1, 1, 2, 2, 2, 2}, 4, 2)
|
||
if _, err := LinearMixedModel(y, sing, z, groups); err == nil {
|
||
t.Fatal("a collinear fixed design was accepted")
|
||
}
|
||
// A single observation carries no fit.
|
||
if _, err := LinearMixedModel(mixedVec(t, []float64{1}, 1), mixedVec(t, []float64{1}, 1, 1), mixedVec(t, []float64{1}, 1, 1), []int{0}); err == nil {
|
||
t.Fatal("a one-row fit was accepted")
|
||
}
|
||
}
|
||
|
||
// cholSolveInTest factors a symmetric positive definite matrix by its
|
||
// own Cholesky and solves against one right-hand side, an independent
|
||
// route the REML referent below evaluates its pieces through.
|
||
func cholSolveInTest(v []float64, rhs []float64, m int) (logDet float64, solution []float64) {
|
||
l := make([]float64, m*m)
|
||
for i := range m {
|
||
for j := range i + 1 {
|
||
s := v[i*m+j]
|
||
for k := range j {
|
||
s -= l[i*m+k] * l[j*m+k]
|
||
}
|
||
if i == j {
|
||
l[i*m+j] = math.Sqrt(s)
|
||
} else {
|
||
l[i*m+j] = s / l[j*m+j]
|
||
}
|
||
}
|
||
logDet += 2 * math.Log(l[i*m+i])
|
||
}
|
||
x := append([]float64(nil), rhs...)
|
||
for i := range m {
|
||
s := x[i]
|
||
for k := range i {
|
||
s -= l[i*m+k] * x[k]
|
||
}
|
||
x[i] = s / l[i*m+i]
|
||
}
|
||
for i := m - 1; i >= 0; i-- {
|
||
s := x[i]
|
||
for k := i + 1; k < m; k++ {
|
||
s -= l[k*m+i] * x[k]
|
||
}
|
||
x[i] = s / l[i*m+i]
|
||
}
|
||
return logDet, x
|
||
}
|
||
|
||
func TestMixedModelREMLLogLikelihoodReferent(t *testing.T) {
|
||
// The reported LogLikelihood is checked against an independent
|
||
// evaluation of the REML criterion at the fitted components,
|
||
// assembled from first principles in this test:
|
||
// -2·logL = Σ_g log|V_g| + rᵀV⁻¹r + log|XᵀV⁻¹X| + (n−p)·ln 2π,
|
||
// with V_g = Z_gΣZ_gᵀ + σ²I and r the fixed-part residual. The
|
||
// one-way balanced design keeps V_g compound symmetric, so the
|
||
// pieces are small and the route shares no arithmetic with the fit.
|
||
const (
|
||
mGroups = 6
|
||
k = 4
|
||
)
|
||
y := make([]float64, 0, mGroups*k)
|
||
groups := make([]int, mGroups*k)
|
||
for g := range mGroups {
|
||
effect := 1.5 * math.Sin(0.9*float64(g)+0.2)
|
||
for i := range k {
|
||
y = append(y, 3+effect+mixedNoise(g*k+i))
|
||
groups[g*k+i] = g
|
||
}
|
||
}
|
||
res, err := LinearMixedModel(
|
||
mixedVec(t, y, len(y)),
|
||
mixedVec(t, ones(len(y)), len(y), 1),
|
||
mixedVec(t, ones(len(y)), len(y), 1),
|
||
groups)
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
if !res.Converged {
|
||
t.Fatalf("the fit did not converge (%d iterations)", res.Iterations)
|
||
}
|
||
sigma2 := res.ResidualVariance
|
||
tau2 := res.RandomCovariance[0]
|
||
beta := res.Coefficients[0]
|
||
n := len(y)
|
||
p := 1
|
||
logDetV := 0.0
|
||
quad := 0.0
|
||
xtvix := 0.0
|
||
for g := range mGroups {
|
||
m := k
|
||
v := make([]float64, m*m)
|
||
for i := range m {
|
||
for j := range m {
|
||
v[i*m+j] = tau2
|
||
}
|
||
v[i*m+i] += sigma2
|
||
}
|
||
r := make([]float64, m)
|
||
for i := range m {
|
||
r[i] = y[g*k+i] - beta
|
||
}
|
||
ld, u := cholSolveInTest(v, r, m)
|
||
logDetV += ld
|
||
for i := range m {
|
||
quad += r[i] * u[i]
|
||
}
|
||
onesRHS := make([]float64, m)
|
||
for i := range onesRHS {
|
||
onesRHS[i] = 1
|
||
}
|
||
_, w := cholSolveInTest(v, onesRHS, m)
|
||
for i := range m {
|
||
xtvix += w[i]
|
||
}
|
||
}
|
||
want := -0.5 * (logDetV + quad + math.Log(xtvix) + float64(n-p)*math.Log(2*math.Pi))
|
||
if math.Abs(res.LogLikelihood-want) > 1e-8*(1+math.Abs(want)) {
|
||
t.Fatalf("LogLikelihood = %.10f, want the independent REML %.10f (difference %.3e)",
|
||
res.LogLikelihood, want, res.LogLikelihood-want)
|
||
}
|
||
}
|