jepa-fx-risk/internal/eval/eval.go

// Package eval implements the Go evaluation harness for jepa-fx-risk (#4).
// Three diagnostics on frozen embeddings exported from train.py:
//   - LinearProbe  — val_vol_r2: OOS R² of a ridge probe predicting next-day realized vol
//   - Silhouette   — mean silhouette score of embeddings vs a binary label (HV regime)
//   - EffectiveRank — Roy's effective rank: exp(H(σ²)) where H is entropy of normalised singular values
package eval

import (
	"errors"
	"math"
)

// LinearProbeTrainTest fits ridge regression on (trainEmb, trainY) and evaluates
// on (testEmb, testY). Returns OOS R². Use this for proper held-out evaluation.
func LinearProbeTrainTest(trainEmb [][]float64, trainY []float64,
	testEmb [][]float64, testY []float64, lambda float64) float64 {
	n := len(trainEmb)
	if n == 0 || len(testEmb) == 0 {
		return 0
	}
	d := len(trainEmb[0])
	p := d + 1

	A := make([][]float64, n)
	for i, e := range trainEmb {
		row := make([]float64, p)
		copy(row, e)
		row[d] = 1.0
		A[i] = row
	}
	AtA := make([][]float64, p)
	for i := range AtA {
		AtA[i] = make([]float64, p)
	}
	Aty := make([]float64, p)
	for i := 0; i < n; i++ {
		for j := 0; j < p; j++ {
			Aty[j] += A[i][j] * trainY[i]
			for k := 0; k < p; k++ {
				AtA[j][k] += A[i][j] * A[i][k]
			}
		}
	}
	for j := 0; j < p; j++ {
		AtA[j][j] += lambda
	}
	w := solveCholesky(AtA, Aty)

	yMean := mean(testY)
	var ssRes, ssTot float64
	for i, e := range testEmb {
		row := make([]float64, p)
		copy(row, e)
		row[d] = 1.0
		pred := dot(row, w)
		ssRes += (testY[i] - pred) * (testY[i] - pred)
		ssTot += (testY[i] - yMean) * (testY[i] - yMean)
	}
	if ssTot == 0 {
		return 0
	}
	return 1 - ssRes/ssTot
}

// LinearProbe fits a ridge regression (closed-form) on (emb, y) with regularisation λ
// and returns R² on the same data. Call with train embeddings; probe on held-out by
// splitting before calling.
//
// emb[i] is the embedding vector for sample i; y[i] is the scalar target.
func LinearProbe(emb [][]float64, y []float64, lambda float64) float64 {
	n := len(emb)
	if n == 0 {
		return 0
	}
	d := len(emb[0])

	// Build augmented design matrix A = [emb | 1] (n × d+1)
	A := make([][]float64, n)
	for i, e := range emb {
		row := make([]float64, d+1)
		copy(row, e)
		row[d] = 1.0
		A[i] = row
	}

	// Normal equations: (AᵀA + λI) w = Aᵀy  (ridge)
	p := d + 1
	AtA := make([][]float64, p)
	for i := range AtA {
		AtA[i] = make([]float64, p)
	}
	Aty := make([]float64, p)

	for i := 0; i < n; i++ {
		for j := 0; j < p; j++ {
			Aty[j] += A[i][j] * y[i]
			for k := 0; k < p; k++ {
				AtA[j][k] += A[i][j] * A[i][k]
			}
		}
	}
	for j := 0; j < p; j++ {
		AtA[j][j] += lambda
	}

	w := solveCholesky(AtA, Aty)

	// R² = 1 - SS_res / SS_tot
	yMean := mean(y)
	var ssRes, ssTot float64
	for i := 0; i < n; i++ {
		pred := dot(A[i], w)
		ssRes += (y[i] - pred) * (y[i] - pred)
		ssTot += (y[i] - yMean) * (y[i] - yMean)
	}
	if ssTot == 0 {
		return 0
	}
	return 1 - ssRes/ssTot
}

// Silhouette returns the mean silhouette coefficient of the embeddings with respect
// to the given integer labels. Distances are Euclidean. Returns an error if fewer
// than 2 distinct labels are present.
func Silhouette(emb [][]float64, labels []int) (float64, error) {
	n := len(emb)
	if n == 0 {
		return 0, errors.New("eval: empty embeddings")
	}

	// count distinct labels
	labelSet := map[int]struct{}{}
	for _, l := range labels {
		labelSet[l] = struct{}{}
	}
	if len(labelSet) < 2 {
		return 0, errors.New("eval: silhouette requires at least 2 distinct labels")
	}

	// group indices by label
	groups := map[int][]int{}
	for i, l := range labels {
		groups[l] = append(groups[l], i)
	}

	var total float64
	for i := 0; i < n; i++ {
		li := labels[i]

		// a(i) = mean intra-cluster distance
		var aSum float64
		inGroup := groups[li]
		for _, j := range inGroup {
			if j != i {
				aSum += euclidean(emb[i], emb[j])
			}
		}
		var a float64
		if len(inGroup) > 1 {
			a = aSum / float64(len(inGroup)-1)
		}

		// b(i) = min mean inter-cluster distance
		b := math.MaxFloat64
		for l, idxs := range groups {
			if l == li {
				continue
			}
			var dSum float64
			for _, j := range idxs {
				dSum += euclidean(emb[i], emb[j])
			}
			avg := dSum / float64(len(idxs))
			if avg < b {
				b = avg
			}
		}

		s := (b - a) / math.Max(a, b)
		total += s
	}
	return total / float64(n), nil
}

// EffectiveRank computes Roy's effective rank of the embedding matrix:
// exp(H) where H = -∑ pᵢ log(pᵢ) is the Shannon entropy of the normalised
// squared singular values. Returns 1 for a rank-1 matrix and ≈ dim for
// a full-rank isotropic matrix.
func EffectiveRank(emb [][]float64) float64 {
	n := len(emb)
	if n == 0 {
		return 0
	}
	d := len(emb[0])

	// Compute covariance-like matrix CᵀC where C is mean-centered embedding.
	mu := make([]float64, d)
	for _, e := range emb {
		for j, v := range e {
			mu[j] += v
		}
	}
	for j := range mu {
		mu[j] /= float64(n)
	}

	// C = emb - mu  (n × d); compute CᵀC (d × d)
	CtC := make([][]float64, d)
	for i := range CtC {
		CtC[i] = make([]float64, d)
	}
	for _, e := range emb {
		for j := 0; j < d; j++ {
			cj := e[j] - mu[j]
			for k := 0; k < d; k++ {
				CtC[j][k] += cj * (e[k] - mu[k])
			}
		}
	}

	// Eigenvalues of CᵀC via power iteration approximation isn't great;
	// use the Frobenius / trace approach: σᵢ² ∝ eigenvalues of CᵀC.
	// For a pure-Go impl without LAPACK: use the fact that the normalised
	// squared singular values equal normalised eigenvalues of CᵀC.
	// Compute them via Jacobi iteration for small d, or use the analytical
	// formula for 2×2, or use iterative QR for general d.
	eigs := jacobiEigenvalues(CtC)

	// normalise to sum-1 distribution
	var sumEig float64
	for _, v := range eigs {
		if v > 0 {
			sumEig += v
		}
	}
	if sumEig == 0 {
		return 1
	}
	var H float64
	for _, v := range eigs {
		if v > 0 {
			p := v / sumEig
			H -= p * math.Log(p)
		}
	}
	return math.Exp(H)
}

// ── internal helpers ──────────────────────────────────────────────────────────

func euclidean(a, b []float64) float64 {
	var s float64
	for i := range a {
		d := a[i] - b[i]
		s += d * d
	}
	return math.Sqrt(s)
}

func dot(a, b []float64) float64 {
	var s float64
	for i := range a {
		s += a[i] * b[i]
	}
	return s
}

func mean(y []float64) float64 {
	var s float64
	for _, v := range y {
		s += v
	}
	return s / float64(len(y))
}

// solveCholesky solves Ax = b for symmetric positive-definite A via
// Cholesky decomposition. Falls back to pseudo-inverse on failure.
func solveCholesky(A [][]float64, b []float64) []float64 {
	n := len(A)
	// Cholesky decomposition: A = LLᵀ
	L := make([][]float64, n)
	for i := range L {
		L[i] = make([]float64, n)
	}
	for i := 0; i < n; i++ {
		for j := 0; j <= i; j++ {
			s := A[i][j]
			for k := 0; k < j; k++ {
				s -= L[i][k] * L[j][k]
			}
			if i == j {
				if s <= 0 {
					s = 1e-12
				}
				L[i][j] = math.Sqrt(s)
			} else {
				L[i][j] = s / L[j][j]
			}
		}
	}
	// Forward substitution Ly = b
	y := make([]float64, n)
	for i := 0; i < n; i++ {
		s := b[i]
		for k := 0; k < i; k++ {
			s -= L[i][k] * y[k]
		}
		y[i] = s / L[i][i]
	}
	// Back substitution Lᵀx = y
	x := make([]float64, n)
	for i := n - 1; i >= 0; i-- {
		s := y[i]
		for k := i + 1; k < n; k++ {
			s -= L[k][i] * x[k]
		}
		x[i] = s / L[i][i]
	}
	return x
}

// jacobiEigenvalues returns eigenvalues of a symmetric matrix via Jacobi iteration.
func jacobiEigenvalues(A [][]float64) []float64 {
	n := len(A)
	// copy
	a := make([][]float64, n)
	for i := range a {
		a[i] = make([]float64, n)
		copy(a[i], A[i])
	}

	const maxIter = 100
	const tol = 1e-10
	for iter := 0; iter < maxIter; iter++ {
		// find largest off-diagonal element
		p, q, amax := 0, 1, 0.0
		for i := 0; i < n; i++ {
			for j := i + 1; j < n; j++ {
				if v := math.Abs(a[i][j]); v > amax {
					amax = v
					p, q = i, j
				}
			}
		}
		if amax < tol {
			break
		}
		// Jacobi rotation
		theta := 0.5 * math.Atan2(2*a[p][q], a[q][q]-a[p][p])
		c, s := math.Cos(theta), math.Sin(theta)
		// apply rotation
		newA := make([][]float64, n)
		for i := range newA {
			newA[i] = make([]float64, n)
			copy(newA[i], a[i])
		}
		app := c*c*a[p][p] + 2*c*s*a[p][q] + s*s*a[q][q]
		aqq := s*s*a[p][p] - 2*c*s*a[p][q] + c*c*a[q][q]
		apq := 0.0
		newA[p][p] = app
		newA[q][q] = aqq
		newA[p][q] = apq
		newA[q][p] = apq
		for r := 0; r < n; r++ {
			if r == p || r == q {
				continue
			}
			arp := c*a[r][p] + s*a[r][q]
			arq := -s*a[r][p] + c*a[r][q]
			newA[r][p] = arp
			newA[p][r] = arp
			newA[r][q] = arq
			newA[q][r] = arq
		}
		a = newA
	}

	eigs := make([]float64, n)
	for i := range eigs {
		eigs[i] = a[i][i]
	}
	return eigs
}