feat: initial release

Assisted-by: GLM 5.3 Flash

stats/mixed.go

@@ -0,0 +1,826 @@
+// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
+// SPDX-License-Identifier: MIT
+
+package stats
+
+import (
+	"math"
+
+	"sourcedock.dev/petrbalvin/tensor/internal/base"
+	"sourcedock.dev/petrbalvin/tensor/internal/core"
+)
+
+// The linear mixed model: a fixed-effect design shared by every
+// observation plus a random-effect design whose coefficients vary by
+// group, y = X·β + Z·b + ε with b_g ~ N(0, Σ) per group and
+// ε ~ N(0, σ²I). The variance components (Σ and σ²) are estimated by
+// residual maximum likelihood, the REML criterion that compares only
+// the error contrasts and so does not let the fixed effects soak up
+// degrees of freedom the variance estimate needs; the fixed effects
+// follow by general least squares at the fitted components.
+//
+// The maximisation is the expectation-maximisation sweep over the
+// random-effect posterior: every update is a conditional expectation,
+// so the sweep climbs the log likelihood monotonically until the
+// tolerance stops it, and the whole fit is deterministic for a given
+// input. No generator enters: the starting point comes from the
+// data's own ordinary least squares, and the sweep walks from there.
+
+// mixedMaxIterations caps the EM sweeps; a fit that has not settled by
+// then is reported with Converged false rather than pretended to have
+// converged.
+const mixedMaxIterations = 500
+
+// mixedTolerance is the convergence tolerance on the REML log
+// likelihood: the sweeps stop once it moves by less than
+// mixedTolerance scaled by 1 + |log likelihood|, past which no
+// parameter the M step can produce moves the objective materially.
+const mixedTolerance = 1e-10
+
+// LinearMixedModelResult carries the fit of a linear mixed model.
+type LinearMixedModelResult struct {
+	// Coefficients are the fitted fixed effects β̂, one per column of
+	// the fixed design, and StandardErrors their estimated standard
+	// deviations, the square roots of the diagonal of the GLS
+	// covariance (XᵀV⁻¹X)⁻¹ at the fitted variance components.
+	Coefficients   []float64
+	StandardErrors []float64
+	// RandomEffects holds one coefficient vector per group, in the
+	// order GroupLabels names, each of the length of a row of the
+	// random design: the posterior means b̂_g at the fitted components.
+	RandomEffects [][]float64
+	// GroupLabels are the distinct group labels in the order the fit
+	// met them, the order RandomEffects keeps.
+	GroupLabels []int
+	// RandomCovariance is the fitted between-group covariance Σ̂ of the
+	// random effects, row-major, and ResidualVariance the fitted σ̂².
+	RandomCovariance []float64
+	ResidualVariance float64
+	// LogLikelihood is the maximised REML log likelihood, the
+	// criterion the sweep climbed.
+	LogLikelihood float64
+	// Fitted and Residuals align with the rows of the design: the
+	// conditional fit X·β̂ + Z·b̂ per row and its remainder.
+	Fitted    []float64
+	Residuals []float64
+	// Iterations counts the EM sweeps taken; Converged reports whether
+	// the log likelihood settled under mixedTolerance before the
+	// budget ran out.
+	Iterations int
+	Converged  bool
+}
+
+// mixedGroup holds one group's fixed pieces, read once before the
+// sweep: the row indices, both designs restricted to the group, the
+// response restricted to it, and the random design's own scatter
+// ZᵀZ, which no iteration moves.
+type mixedGroup struct {
+	label int
+	rows  []int
+	yg    []float64
+	xg    []float64
+	zg    []float64
+	ztz   []float64
+	xtz   []float64
+	ng    int
+}
+
+// LinearMixedModel fits the linear mixed model over the response y,
+// the fixed design x (n rows, p columns, the intercept included by
+// the caller as a constant column when one is wanted), the random
+// design z (n rows, q columns) and the group label of every row. The
+// labels need no particular order or contiguity: the fit meets them
+// in row order and reports the distinct ones in GroupLabels in the
+// order it met them. The random effects share one unstructured q×q
+// covariance Σ across the groups.
+//
+// The fit runs expectation maximisation over the random-effect
+// posterior to the documented tolerance, within mixedMaxIterations
+// sweeps, and reports the REML log likelihood of the fitted state. A
+// group whose marginal covariance fails to factor under the fitted
+// components, or a fixed design that is singular under them, is an
+// error naming the group or the design.
+//
+// Refuses complex input and any non-finite entry, a response shorter
+// than two rows, a design of fewer than one column, a row count
+// mismatch, more fixed coefficients than observations, and a
+// negative or missing group label.
+func LinearMixedModel(y, x, z *core.Array, groups []int) (*LinearMixedModelResult, error) {
+	const name = "LinearMixedModel"
+	if y.NDim() != 1 {
+		return nil, base.Errf("%s: the response must be rank 1, got shape %s", name, base.ShapeText(y.Shape()))
+	}
+	if x.NDim() != 2 || z.NDim() != 2 {
+		return nil, base.Errf("%s: both designs must be rank 2", name)
+	}
+	if y.Dtype() == core.Complex || x.Dtype() == core.Complex || z.Dtype() == core.Complex {
+		return nil, base.Errf("%s: complex inputs are not supported", name)
+	}
+	n := y.Len()
+	p := x.Shape()[1]
+	q := z.Shape()[1]
+	if x.Shape()[0] != n || z.Shape()[0] != n {
+		return nil, base.Errf("%s: the response holds %d rows, the designs %d and %d", name, n, x.Shape()[0], z.Shape()[0])
+	}
+	if n < 2 {
+		return nil, base.Errf("%s: at least two observations are needed, got %d", name, n)
+	}
+	if p < 1 || q < 1 {
+		return nil, base.Errf("%s: both designs need at least one column, got %d and %d", name, p, q)
+	}
+	if n <= p {
+		return nil, base.Errf("%s: %d observations cannot carry %d fixed coefficients", name, n, p)
+	}
+	if len(groups) != n {
+		return nil, base.Errf("%s: %d group labels for %d observations", name, len(groups), n)
+	}
+	if err := checkFinite(name, "the response", y); err != nil {
+		return nil, err
+	}
+	if err := checkFinite(name, "the fixed design", x); err != nil {
+		return nil, err
+	}
+	if err := checkFinite(name, "the random design", z); err != nil {
+		return nil, err
+	}
+	// The response and the designs are read once into plain slices:
+	// every sweep below indexes them flat.
+	yVals := make([]float64, n)
+	if fs := rawFloats(y); fs != nil {
+		copy(yVals, fs)
+	} else {
+		for i := range n {
+			yVals[i] = y.FloatAt(i)
+		}
+	}
+	xVals := make([]float64, n*p)
+	if fs := rawFloats(x); fs != nil {
+		copy(xVals, fs[:n*p])
+	} else {
+		for i := range xVals {
+			xVals[i] = x.FloatAt(i)
+		}
+	}
+	zVals := make([]float64, n*q)
+	if fs := rawFloats(z); fs != nil {
+		copy(zVals, fs[:n*q])
+	} else {
+		for i := range zVals {
+			zVals[i] = z.FloatAt(i)
+		}
+	}
+	// The groups are canonicalised by first appearance: the labels
+	// stay the caller's own, the fit only needs them distinct and
+	// stable.
+	labelIndex := make(map[int]int)
+	labels := make([]int, 0, 8)
+	member := make([][]int, 0, 8)
+	for i, label := range groups {
+		if label < 0 {
+			return nil, base.Errf("%s: group label %d is negative", name, label)
+		}
+		gi, ok := labelIndex[label]
+		if !ok {
+			gi = len(labels)
+			labelIndex[label] = gi
+			labels = append(labels, label)
+			member = append(member, nil)
+		}
+		member[gi] = append(member[gi], i)
+	}
+	gs := make([]mixedGroup, len(labels))
+	for gi, rows := range member {
+		g := &gs[gi]
+		g.label = labels[gi]
+		g.rows = rows
+		g.ng = len(rows)
+		g.yg = make([]float64, g.ng)
+		g.xg = make([]float64, g.ng*p)
+		g.zg = make([]float64, g.ng*q)
+		g.ztz = make([]float64, q*q)
+		g.xtz = make([]float64, p*q)
+		for li, row := range rows {
+			g.yg[li] = yVals[row]
+			copy(g.xg[li*p:(li+1)*p], xVals[row*p:(row+1)*p])
+			copy(g.zg[li*q:(li+1)*q], zVals[row*q:(row+1)*q])
+		}
+		for a := range q {
+			for b := range q {
+				s := 0.0
+				for li := range g.ng {
+					s += g.zg[li*q+a] * g.zg[li*q+b]
+				}
+				g.ztz[a*q+b] = s
+			}
+		}
+		for j := range p {
+			for b := range q {
+				s := 0.0
+				for li := range g.ng {
+					s += g.xg[li*p+j] * g.zg[li*q+b]
+				}
+				g.xtz[j*q+b] = s
+			}
+		}
+	}
+	// The starting point: ordinary least squares on the fixed design,
+	// the residual variance around it, and a between-group scatter of
+	// its residuals floored above zero so the first sweep can see the
+	// random effects at all.
+	beta, err := mixedOLSStart(name, xVals, yVals, n, p)
+	if err != nil {
+		return nil, err
+	}
+	residual := make([]float64, n)
+	rss := 0.0
+	totalSquare := 0.0
+	for i := range n {
+		fit := 0.0
+		for j := range p {
+			fit += xVals[i*p+j] * beta[j]
+		}
+		residual[i] = yVals[i] - fit
+		rss += residual[i] * residual[i]
+		totalSquare += yVals[i] * yVals[i]
+	}
+	scale := 1 + totalSquare/float64(n)
+	sigma2 := math.Max(rss/float64(n), 1e-12*scale)
+	between := 0.0
+	if len(gs) > 1 {
+		overall := 0.0
+		means := make([]float64, len(gs))
+		for gi, g := range gs {
+			s := 0.0
+			for _, row := range g.rows {
+				s += residual[row]
+			}
+			means[gi] = s / float64(g.ng)
+			overall += means[gi]
+		}
+		overall /= float64(len(gs))
+		for _, m := range means {
+			d := m - overall
+			between += d * d
+		}
+		between /= float64(len(gs) - 1)
+	}
+	comp := math.Max(between, 0.1*sigma2)
+	sigma := make([]float64, q*q)
+	for a := range q {
+		sigma[a*q+a] = comp
+	}
+	// The sweep. The per-group buffers a pass touches are allocated
+	// once here and refilled in place; the fixed-effect solve keeps
+	// its own small p×p workspace, reallocated each pass.
+	chol := make([][][]float64, len(gs))
+	scratch := make([]mixedSweep, len(gs))
+	for gi := range gs {
+		ng := gs[gi].ng
+		chol[gi] = make([][]float64, ng)
+		for i := range ng {
+			chol[gi][i] = make([]float64, ng)
+		}
+		scratch[gi] = newMixedSweep(ng, p, q)
+	}
+	// The fixed scatter XᵀX, which no iteration moves: the REML M
+	// step's expectation of the squared residual carries its trace
+	// against the fixed effects' posterior covariance.
+	xtx := make([]float64, p*p)
+	for i := range n {
+		for j := range p {
+			xj := xVals[i*p+j]
+			for k := range p {
+				xtx[j*p+k] += xj * xVals[i*p+k]
+			}
+		}
+	}
+	xtvix := make([]float64, p*p)
+	xtviy := make([]float64, p)
+	sigmaNew := make([]float64, q*q)
+	logLik := math.Inf(-1)
+	converged := false
+	iterations := mixedMaxIterations
+	for iter := 1; iter <= mixedMaxIterations; iter++ {
+		pieces, perr := mixedPass(gs, chol, scratch, xtx, beta, sigma, sigma2, xtvix, xtviy, sigmaNew)
+		if perr != nil {
+			return nil, base.Errf("%s: %w", name, perr)
+		}
+		next := -0.5 * (pieces.logDetV + pieces.quad + pieces.logDetX + float64(n-p)*math.Log(2*math.Pi))
+		if math.IsNaN(next) {
+			return nil, base.Errf("%s: the fit left the finite domain at iteration %d", name, iter)
+		}
+		move := next - logLik
+		logLik = next
+		if iter > 1 && math.Abs(move) <= mixedTolerance*(1+math.Abs(logLik)) {
+			converged = true
+			iterations = iter
+			break
+		}
+		// The M step: the averaged posterior scatter, the expected
+		// residual variance with its posterior-trace correction, and
+		// the GLS fixed effects at this pass's components.
+		copy(sigma, sigmaNew)
+		for a := range q {
+			for b := a + 1; b < q; b++ {
+				sigma[b*q+a] = sigma[a*q+b]
+			}
+		}
+		sigma2 = math.Max((pieces.sse+pieces.trace)/float64(n), 1e-12*scale)
+		// The divergence fuse: some designs give the REML criterion no
+		// interior optimum, the classical case being a random design
+		// that spans the fixed one under an unstructured covariance,
+		// where the surface climbs a ridge of singular Σ without a
+		// summit. The fit refuses with the condition named instead of
+		// publishing the climb.
+		if sigma2 > 1e12*scale {
+			return nil, base.Errf("%s: the variance components diverged past the data's scale; this design gives the REML criterion no interior optimum", name)
+		}
+		for _, v := range sigma {
+			if math.Abs(v) > 1e12*scale {
+				return nil, base.Errf("%s: the variance components diverged past the data's scale; this design gives the REML criterion no interior optimum", name)
+			}
+		}
+		if err := lltSolveSquare(name, xtvix, xtviy, p); err != nil {
+			return nil, base.Errf("%s: the fixed design is singular under the fitted covariance (%w)", name, err)
+		}
+		copy(beta, xtviy)
+	}
+	// The reporting pass at the fitted parameters: the same sweep once
+	// more at fixed state, then the quantities the caller holds. The
+	// sweep is deterministic, so its totals are the fitted state's
+	// own.
+	pieces, perr := mixedPass(gs, chol, scratch, xtx, beta, sigma, sigma2, xtvix, xtviy, sigmaNew)
+	if perr != nil {
+		return nil, base.Errf("%s: %w", name, perr)
+	}
+	logLik = -0.5*(pieces.logDetV+pieces.quad+pieces.logDetX) - 0.5*float64(n-p)*math.Log(2*math.Pi)
+	// The standard errors: the square roots of the diagonal of the GLS
+	// covariance, from the inverse of the same matrix the log
+	// determinant came from.
+	diag, serr := lltInverseDiagonal(name, xtvix, p)
+	if serr != nil {
+		return nil, base.Errf("%s: %w", name, serr)
+	}
+	se := make([]float64, p)
+	for i := range p {
+		se[i] = math.Sqrt(diag[i])
+	}
+	fitted := make([]float64, n)
+	resids := make([]float64, n)
+	for gi, g := range gs {
+		bhat := scratch[gi].bhat
+		for li, row := range g.rows {
+			fit := 0.0
+			for j := range p {
+				fit += xVals[row*p+j] * beta[j]
+			}
+			for a := range q {
+				fit += g.zg[li*q+a] * bhat[a]
+			}
+			fitted[row] = fit
+			resids[row] = yVals[row] - fit
+		}
+	}
+	effects := make([][]float64, len(gs))
+	for gi := range gs {
+		effects[gi] = append([]float64(nil), scratch[gi].bhat...)
+	}
+	return &LinearMixedModelResult{
+		Coefficients:     append([]float64(nil), beta...),
+		StandardErrors:   se,
+		RandomEffects:    effects,
+		GroupLabels:      labels,
+		RandomCovariance: append([]float64(nil), sigma...),
+		ResidualVariance: sigma2,
+		LogLikelihood:    logLik,
+		Fitted:           fitted,
+		Residuals:        resids,
+		Iterations:       iterations,
+		Converged:        converged,
+	}, nil
+}
+
+// mixedSweep holds one group's reusable pass buffers: the marginal
+// covariance and its factor's solves against the response and both
+// designs, the conditional residual, and the posterior mean and
+// scatter of the group's random effects.
+type mixedSweep struct {
+	u      []float64 // V⁻¹ r
+	vy     []float64 // V⁻¹ y
+	vz     []float64 // V⁻¹ Z, ng×q
+	vx     []float64 // V⁻¹ X, ng×p
+	slope  []float64 // −Σ Zᵀ (V⁻¹ X), the posterior mean's slope in β, q×p
+	bhat   []float64 // Σ Zᵀ u
+	post   []float64 // P + b̂b̂ᵀ, the conditional scatter
+	v      []float64 // the marginal covariance, ng×ng
+	resid  []float64 // y − Xβ per row
+	mzx    []float64 // Σ Zᵀ (V⁻¹ Z), q×q
+	sc     []float64 // slope·C, refilled by the REML correction, q×p
+	bq     []float64 // sc·slopeᵀ, the correction's scatter, q×q
+	colvec []float64 // one design column, refilled per solve
+}
+
+// newMixedSweep allocates one group's buffers.
+func newMixedSweep(ng, p, q int) mixedSweep {
+	return mixedSweep{
+		u:      make([]float64, ng),
+		vy:     make([]float64, ng),
+		vz:     make([]float64, ng*q),
+		vx:     make([]float64, ng*p),
+		slope:  make([]float64, q*p),
+		bhat:   make([]float64, q),
+		post:   make([]float64, q*q),
+		v:      make([]float64, ng*ng),
+		resid:  make([]float64, ng),
+		mzx:    make([]float64, q*q),
+		sc:     make([]float64, q*p),
+		bq:     make([]float64, q*q),
+		colvec: make([]float64, ng),
+	}
+}
+
+// mixedPassTotals gathers what one pass accumulates across the
+// groups.
+type mixedPassTotals struct {
+	logDetV float64 // Σ log|V_g|, the marginal covariances' own logarithms
+	quad    float64 // Σ r_gᵀ V_g⁻¹ r_g
+	sse     float64 // Σ ‖r_g − Z_g b̂_g‖²
+	trace   float64 // Σ tr(Z_gᵀZ_g·P_g)
+	logDetX float64 // log|XᵀV⁻¹X|
+}
+
+// mixedPass runs one expectation pass over the groups at the given
+// parameters: it factors every group's marginal covariance, solves it
+// against the response and both designs, forms the posterior means
+// and scatters, and accumulates the totals the log likelihood and the
+// M step read. The M step's expectations are taken under the flat
+// prior on the fixed effects, so the fixed effects' own posterior
+// covariance (XᵀV⁻¹X)⁻¹ joins both updates through the posterior
+// mean's slope in β; that is the step that makes the fixed point the
+// REML optimum rather than the ML one. The accumulators xtvix, xtviy
+// and sigmaNew are overwritten; the posterior means stay in the
+// group's scratch for the reporting pass to collect.
+func mixedPass(gs []mixedGroup, chol [][][]float64, scratch []mixedSweep, xtx []float64,
+	beta []float64, sigma []float64, sigma2 float64, xtvix, xtviy, sigmaNew []float64) (*mixedPassTotals, error) {
+	const name = "LinearMixedModel"
+	p := len(beta)
+	q := len(scratch[0].bhat)
+	clear(xtvix)
+	clear(xtviy)
+	clear(sigmaNew)
+	totals := &mixedPassTotals{}
+	for gi := range gs {
+		g := &gs[gi]
+		s := &scratch[gi]
+		ng := g.ng
+		// The marginal covariance V = Z Σ Zᵀ + σ²I, both triangles
+		// written so the factorisation's symmetry check sees a mirror
+		// pair with identical bits.
+		for i := range ng {
+			for j := range i + 1 {
+				total := 0.0
+				for a := range q {
+					za := g.zg[i*q+a]
+					if za != 0 {
+						for b := range q {
+							total += za * sigma[a*q+b] * g.zg[j*q+b]
+						}
+					}
+				}
+				if i == j {
+					total += sigma2
+				}
+				s.v[i*ng+j] = total
+				s.v[j*ng+i] = total
+			}
+		}
+		if err := mvnCholeskyFlatInto(name, s.v, chol[gi], ng); err != nil {
+			return nil, base.Errf("group %d's marginal covariance failed to factor (%w)", g.label, err)
+		}
+		l := chol[gi]
+		for i := range ng {
+			// The factor's determinant answers |V_g| = |L|², so the
+			// diagonal logs enter doubled: the criterion the sweep reads
+			// below carries log|V_g| itself, not half of it.
+			totals.logDetV += 2 * math.Log(l[i][i])
+		}
+		// The residual against the fixed part alone, then u = V⁻¹ r.
+		for li := range ng {
+			r := 0.0
+			for j := range p {
+				r += g.xg[li*p+j] * beta[j]
+			}
+			s.resid[li] = g.yg[li] - r
+		}
+		copy(s.u, s.resid)
+		lltSolveInPlace(l, s.u)
+		totals.quad += dot(s.resid, s.u)
+		// V⁻¹ Z and V⁻¹ X, column by column through the same factor.
+		for a := range q {
+			for li := range ng {
+				s.colvec[li] = g.zg[li*q+a]
+			}
+			lltSolveInPlace(l, s.colvec)
+			copy(s.vz[a*ng:a*ng+ng], s.colvec)
+		}
+		for j := range p {
+			for li := range ng {
+				s.colvec[li] = g.xg[li*p+j]
+			}
+			lltSolveInPlace(l, s.colvec)
+			copy(s.vx[j*ng:j*ng+ng], s.colvec)
+		}
+		// V⁻¹ y, free of charge out of the solves already run:
+		// V⁻¹ y = V⁻¹ r + (V⁻¹ X)·β.
+		for li := range ng {
+			total := s.u[li]
+			for j := range p {
+				total += s.vx[j*ng+li] * beta[j]
+			}
+			s.vy[li] = total
+		}
+		// The posterior mean b̂ = Σ Zᵀ u.
+		clear(s.bhat)
+		for a := range q {
+			total := 0.0
+			for li := range ng {
+				zt := 0.0
+				for b := range q {
+					zt += sigma[a*q+b] * g.zg[li*q+b]
+				}
+				total += zt * s.u[li]
+			}
+			s.bhat[a] = total
+		}
+		// mzx = Σ Zᵀ (V⁻¹ Z), then the posterior covariance
+		// P = Σ − mzx·Σ on its upper triangle, joined by the mean's
+		// outer product into the scatter the M step averages.
+		clear(s.mzx)
+		for a := range q {
+			for b := range q {
+				total := 0.0
+				for c := range q {
+					factor := sigma[a*q+c]
+					if factor != 0 {
+						for li := range ng {
+							total += factor * g.zg[li*q+c] * s.vz[li*q+b]
+						}
+					}
+				}
+				s.mzx[a*q+b] = total
+			}
+		}
+		clear(s.post)
+		for a := range q {
+			for b := a; b < q; b++ {
+				pv := sigma[a*q+b]
+				for c := range q {
+					pv -= s.mzx[a*q+c] * sigma[c*q+b]
+				}
+				s.post[a*q+b] = pv + s.bhat[a]*s.bhat[b]
+			}
+		}
+		for a := range q {
+			for b := a; b < q; b++ {
+				sigmaNew[a*q+b] += s.post[a*q+b] / float64(len(gs))
+			}
+		}
+		// The conditional residual and its posterior-trace correction.
+		for li := range ng {
+			r := s.resid[li]
+			for a := range q {
+				r -= g.zg[li*q+a] * s.bhat[a]
+			}
+			s.resid[li] = r
+			totals.sse += r * r
+		}
+		for a := range q {
+			for b := a; b < q; b++ {
+				pv := sigma[a*q+b]
+				for c := range q {
+					pv -= s.mzx[a*q+c] * sigma[c*q+b]
+				}
+				totals.trace += g.ztz[a*q+b] * pv
+				if b != a {
+					totals.trace += g.ztz[b*q+a] * pv
+				}
+			}
+		}
+		// XᵀV⁻¹X and XᵀV⁻¹y accumulate across the groups. The
+		// information is symmetric by construction but its two rounding
+		// paths differ, so only the upper triangle is accumulated and
+		// mirrored: the factorisation's relative symmetry check would
+		// otherwise condemn a matrix whose one side rounded to an exact
+		// zero.
+		for j := range p {
+			for k := j; k < p; k++ {
+				total := 0.0
+				for li := range ng {
+					total += g.xg[li*p+k] * s.vx[j*ng+li]
+				}
+				xtvix[j*p+k] += total
+				if k != j {
+					xtvix[k*p+j] += total
+				}
+			}
+			total := 0.0
+			for li := range ng {
+				total += g.xg[li*p+j] * s.vy[li]
+			}
+			xtviy[j] += total
+		}
+		// The posterior mean's slope in β: E[b_g|β, y] moves as
+		// b̂ − (Σ ZᵀV⁻¹X)(β − β̂), and the slope carries the fixed
+		// effects' uncertainty into the M step's expectations.
+		for a := range q {
+			for j := range p {
+				total := 0.0
+				for b := range q {
+					zt := sigma[a*q+b]
+					if zt != 0 {
+						for li := range ng {
+							total += zt * g.zg[li*q+b] * s.vx[j*ng+li]
+						}
+					}
+				}
+				s.slope[a*p+j] = -total
+			}
+		}
+	}
+	// The fixed information through its own Cholesky: the log
+	// determinant the REML criterion reads, and the posterior
+	// covariance C = (XᵀV⁻¹X)⁻¹ the corrections below walk through.
+	// A non-positive pivot means the design has collapsed under the
+	// fitted covariance, which is an error, not a fit.
+	cholX := make([][]float64, p)
+	for i := range p {
+		cholX[i] = make([]float64, p)
+	}
+	if err := mvnCholeskyFlatInto(name, xtvix, cholX, p); err != nil {
+		return nil, base.Errf("the fixed design is singular under the fitted covariance (%w)", err)
+	}
+	logDetX := 0.0
+	for i := range p {
+		logDetX += math.Log(cholX[i][i])
+	}
+	totals.logDetX = 2 * logDetX
+	inverse := make([]float64, p*p)
+	column := make([]float64, p)
+	for j := range p {
+		clear(column)
+		column[j] = 1
+		lltSolveInPlace(cholX, column)
+		for i := range p {
+			inverse[i*p+j] = column[i]
+		}
+	}
+	// The flat-prior corrections. The expected squared residual gains
+	// the trace of XᵀX against C and, per group, twice the cross term
+	// of XᵀZ against C·slopeᵀ; the expected random-effect scatter
+	// gains slope·C·slopeᵀ. With no random effect the two terms leave
+	// σ² = RSS/(n−p), the REML answer, as the fixed point.
+	traceXX := 0.0
+	for a := range p {
+		for b := range p {
+			traceXX += xtx[a*p+b] * inverse[a*p+b]
+		}
+	}
+	crossTotal := 0.0
+	for gi := range gs {
+		g := &gs[gi]
+		s := &scratch[gi]
+		for a := range q {
+			for c := range p {
+				total := 0.0
+				for j := range p {
+					total += s.slope[a*p+j] * inverse[j*p+c]
+				}
+				s.sc[a*p+c] = total
+			}
+		}
+		clear(s.bq)
+		for a := range q {
+			for b := a; b < q; b++ {
+				total := 0.0
+				for c := range p {
+					total += s.sc[a*p+c] * s.slope[b*p+c]
+				}
+				s.bq[a*q+b] = total
+			}
+		}
+		for a := range q {
+			for b := a; b < q; b++ {
+				sigmaNew[a*q+b] += s.bq[a*q+b] / float64(len(gs))
+				totals.trace += g.ztz[a*q+b] * s.bq[a*q+b]
+				if b != a {
+					totals.trace += g.ztz[b*q+a] * s.bq[a*q+b]
+				}
+			}
+		}
+		for j := range p {
+			for b := range q {
+				crossTotal += g.xtz[j*q+b] * s.sc[b*p+j]
+			}
+		}
+	}
+	totals.sse += traceXX + 2*crossTotal
+	return totals, nil
+}
+
+// dot returns the plain dot product of two equal-length slices.
+func dot(a, b []float64) float64 {
+	total := 0.0
+	for i, v := range a {
+		total += v * b[i]
+	}
+	return total
+}
+
+// lltSolveInPlace solves L·Lᵀ x = x in place: the forward sweep
+// leaves L's solution in x, the backward sweep reads the entries the
+// descending order has already updated.
+func lltSolveInPlace(l [][]float64, x []float64) {
+	n := len(x)
+	for i := range n {
+		total := x[i]
+		for j := range i {
+			total -= l[i][j] * x[j]
+		}
+		x[i] = total / l[i][i]
+	}
+	for i := n - 1; i >= 0; i-- {
+		total := x[i]
+		for j := i + 1; j < n; j++ {
+			total -= l[j][i] * x[j]
+		}
+		x[i] = total / l[i][i]
+	}
+}
+
+// lltSolveSquare solves the symmetric positive definite system a·x =
+// b in place, factoring a through the house Cholesky and leaving b
+// holding x.
+func lltSolveSquare(name string, a, b []float64, d int) error {
+	l := make([][]float64, d)
+	for i := range d {
+		l[i] = make([]float64, d)
+	}
+	if err := mvnCholeskyFlatInto(name, a, l, d); err != nil {
+		return err
+	}
+	lltSolveInPlace(l, b)
+	return nil
+}
+
+// lltInverseDiagonal returns the diagonal of the inverse of the
+// symmetric positive definite a, one solve against each unit vector.
+func lltInverseDiagonal(name string, a []float64, d int) ([]float64, error) {
+	l := make([][]float64, d)
+	for i := range d {
+		l[i] = make([]float64, d)
+	}
+	if err := mvnCholeskyFlatInto(name, a, l, d); err != nil {
+		return nil, err
+	}
+	diag := make([]float64, d)
+	column := make([]float64, d)
+	for i := range d {
+		clear(column)
+		column[i] = 1
+		lltSolveInPlace(l, column)
+		diag[i] = column[i]
+	}
+	return diag, nil
+}
+
+// mixedOLSStart returns the ordinary least squares coefficients of
+// the fixed design, the fixed effects' starting point.
+func mixedOLSStart(name string, xVals, yVals []float64, n, p int) ([]float64, error) {
+	xtx := make([]float64, p*p)
+	xty := make([]float64, p)
+	for i := range n {
+		for j := range p {
+			xj := xVals[i*p+j]
+			xty[j] += xj * yVals[i]
+			for k := range p {
+				xtx[j*p+k] += xj * xVals[i*p+k]
+			}
+		}
+	}
+	rhs := make([][]float64, 1)
+	rhs[0] = xty
+	solved, err := base.SolveSystem(name, toRows(xtx, p), rhs)
+	if err != nil {
+		return nil, base.Errf("%s: the fixed design is singular, the fixed effects cannot be initialised (%w)", name, err)
+	}
+	return append([]float64(nil), solved[0]...), nil
+}
+
+// toRows views a flat row-major matrix as rows for the shared solve.
+func toRows(flat []float64, d int) [][]float64 {
+	rows := make([][]float64, d)
+	for i := range d {
+		rows[i] = flat[i*d : i*d+d]
+	}
+	return rows
+}