feat: initial release

Assisted-by: GLM 5.3 Flash

stats/mixed_test.go

@@ -0,0 +1,484 @@
+// 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)
+	}
+}