jepa-fx-risk/cmd/eval/main.go

// cmd/eval — CLI driver for the jepa-fx-risk evaluation harness.
//
//	./bin/eval -metric probe|silhouette|erank [-emb embeddings.json]
//
// embeddings.json format (from train.py EXPORT_EMBEDDINGS=1):
//
//	{
//	  "embeddings":         [[...], ...],  // OOS frozen embeddings
//	  "realized_vol":       [...],         // OOS target (next-day RV)
//	  "hv_label":           [...],         // binary HV label (top-33%)
//	  "train_embeddings":   [[...], ...],  // train-set frozen embeddings
//	  "train_realized_vol": [...]          // train-set RV targets
//	}
//
// eval:probe standardises both sets using train statistics (no leakage).
// Falls back to internal 70/30 split of OOS if train_embeddings absent.
package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"log/slog"
	"math"
	"os"

	"gitea.d-ma.be/mathias/jepa-fx-risk/internal/eval"
)

func main() {
	metric := flag.String("metric", "probe", "probe | silhouette | erank")
	embFile := flag.String("emb", "embeddings.json", "path to embeddings JSON")
	flag.Parse()

	log := slog.New(slog.NewJSONHandler(os.Stdout, nil))

	d, err := readJSON(*embFile)
	if err != nil {
		log.Error("load embeddings", "err", err)
		os.Exit(1)
	}
	log.Info("loaded", "oos", len(d.Embeddings), "dim", len(d.Embeddings[0]),
		"train", len(d.TrainEmbeddings), "metric", *metric)

	switch *metric {
	case "probe":
		var r2 float64
		if len(d.TrainEmbeddings) > 0 {
			// standardise both sets using train statistics to prevent leakage
			trEmb, mu, sd := standardiseCompute(d.TrainEmbeddings)
			oosEmb := applyStandardise(d.Embeddings, mu, sd)
			r2 = eval.LinearProbeTrainTest(trEmb, d.TrainRealizedVol, oosEmb, d.RealizedVol, 1e-3)
			log.Info("probe mode", "fit_on", "train_embeddings", "eval_on", "oos")
		} else {
			// fallback: internal 70/30 split of OOS embeddings
			oosEmb, mu, sd := standardiseCompute(d.Embeddings)
			n70 := int(float64(len(oosEmb)) * 0.7)
			oos70 := applyStandardise(d.Embeddings[n70:], mu, sd)
			r2 = eval.LinearProbeTrainTest(oosEmb[:n70], d.RealizedVol[:n70],
				oos70, d.RealizedVol[n70:], 1e-3)
			log.Info("probe mode", "fit_on", "oos[0:70%]", "eval_on", "oos[70%:]")
		}
		fmt.Printf(`{"metric":"val_vol_r2","value":%.6f}`+"\n", r2)
		log.Info("linear probe", "val_vol_r2", fmt.Sprintf("%.4f", r2))

	case "silhouette":
		if len(d.HVLabel) == 0 {
			log.Error("silhouette requires hv_label in embeddings.json")
			os.Exit(1)
		}
		oosEmb := standardise(d.Embeddings)
		sil, err := eval.Silhouette(oosEmb, d.HVLabel)
		if err != nil {
			log.Error("silhouette", "err", err)
			os.Exit(1)
		}
		fmt.Printf(`{"metric":"silhouette","value":%.6f}`+"\n", sil)
		log.Info("silhouette", "score", fmt.Sprintf("%.4f", sil))

	case "erank":
		oosEmb := standardise(d.Embeddings)
		er := eval.EffectiveRank(oosEmb)
		fmt.Printf(`{"metric":"effective_rank","value":%.6f}`+"\n", er)
		log.Info("effective rank", "erank", fmt.Sprintf("%.2f", er))

	default:
		log.Error("unknown metric", "metric", *metric)
		os.Exit(1)
	}
}

type embJSON struct {
	Embeddings       [][]float64 `json:"embeddings"`
	Dates            []string    `json:"dates"`
	RealizedVol      []float64   `json:"realized_vol"`
	HVLabel          []int       `json:"hv_label"`
	TrainEmbeddings  [][]float64 `json:"train_embeddings"`
	TrainRealizedVol []float64   `json:"train_realized_vol"`
}

func readJSON(path string) (*embJSON, error) {
	f, err := os.Open(path)
	if err != nil {
		return nil, fmt.Errorf("open %s: %w", path, err)
	}
	defer func() { _ = f.Close() }()
	var d embJSON
	if err := json.NewDecoder(f).Decode(&d); err != nil {
		return nil, fmt.Errorf("decode: %w", err)
	}
	if len(d.Embeddings) == 0 {
		return nil, fmt.Errorf("empty embeddings in %s", path)
	}
	return &d, nil
}

// standardise centres + scales to zero mean / unit std; returns normalised rows.
func standardise(rows [][]float64) [][]float64 {
	out, _, _ := standardiseCompute(rows)
	return out
}

// standardiseCompute centres + scales and returns (normalised, mu, sd) for reuse.
func standardiseCompute(rows [][]float64) ([][]float64, []float64, []float64) {
	if len(rows) == 0 {
		return rows, nil, nil
	}
	n, dim := len(rows), len(rows[0])
	mu := make([]float64, dim)
	for _, r := range rows {
		for j, v := range r {
			mu[j] += v
		}
	}
	for j := range mu {
		mu[j] /= float64(n)
	}
	sd := make([]float64, dim)
	for _, r := range rows {
		for j, v := range r {
			diff := v - mu[j]
			sd[j] += diff * diff
		}
	}
	for j := range sd {
		sd[j] = math.Sqrt(sd[j]/float64(n)) + 1e-8
	}
	out := make([][]float64, n)
	for i, r := range rows {
		out[i] = make([]float64, dim)
		for j, v := range r {
			out[i][j] = (v - mu[j]) / sd[j]
		}
	}
	return out, mu, sd
}

// applyStandardise normalises rows using pre-computed mu and sd.
func applyStandardise(rows [][]float64, mu, sd []float64) [][]float64 {
	out := make([][]float64, len(rows))
	for i, r := range rows {
		out[i] = make([]float64, len(r))
		for j, v := range r {
			out[i][j] = (v - mu[j]) / sd[j]
		}
	}
	return out
}