Eval Metrics — Go source

// Package evalmetrics is the Go twin of CosmoDev's src/lib/evalMetrics.ts
// (dual source: the web lib is TypeScript, the CLI lib is Go — kept in
// lock-step). Pure + deterministic, never panics. The table-driven tests in
// evalmetrics_test.go share vectors with src/lib/evalMetrics.test.ts so the
// two implementations are held to the same contract.
//
// Mirrors the TS lib exactly:
//   - pass@k: the unbiased estimator from the Codex paper (Chen et al. 2021).
//   - precision / recall / F1 over confusion counts, and over label sets.
//   - exact-match rate over paired strings.
//
// Error mapping: every TS function that throws RangeError returns a non-nil
// error here instead (zero-value result alongside it); setMatch never throws
// in TS, so it returns a bare PRF1 here.
package evalmetrics

import (
	"errors"
)

// PRF1 mirrors the PRF1 interface: precision, recall and their harmonic mean.
type PRF1 struct {
	Precision float64
	Recall    float64
	F1        float64
}

// ConfusionCounts mirrors the ConfusionCounts interface. The zero value is the
// TS "all counts omitted/zero" object. TN is optional in TS and ignored there
// (only tp/fp/fn are destructured), so it is likewise carried but unused.
type ConfusionCounts struct {
	TP int
	FP int
	FN int
	TN int // shape parity with TS; never read, exactly like the TS lib
}

// PassAtK is the unbiased pass@k: the probability that at least one of k
// samples drawn without replacement from n (of which c are correct) passes.
// It is the Go twin of passAtK() in src/lib/evalMetrics.ts.
//
//	1                                       when n - c < k  (a wrong draw is impossible)
//	1 - Π_{i=0..k-1} (n - c - i) / (n - i)  otherwise
func PassAtK(n, c, k int) (float64, error) {
	if n <= 0 {
		return 0, errors.New("n must be > 0")
	}
	if c < 0 || c > n {
		return 0, errors.New("c must be in [0, n]")
	}
	if k <= 0 || k > n {
		return 0, errors.New("k must be in [1, n]")
	}
	if n-c < k {
		return 1, nil
	}
	product := 1.0
	for i := 0; i < k; i++ {
		product *= float64(n-c-i) / float64(n-i)
	}
	return 1 - product, nil
}

// PrecisionRecall computes precision/recall/F1 over confusion counts. Zero
// denominators score 0. It is the Go twin of precisionRecall() in
// src/lib/evalMetrics.ts.
func PrecisionRecall(counts ConfusionCounts) (PRF1, error) {
	tp, fp, fn := counts.TP, counts.FP, counts.FN
	if tp < 0 || fp < 0 || fn < 0 {
		return PRF1{}, errors.New("counts must be >= 0")
	}
	var precision, recall float64
	if tp+fp > 0 {
		precision = float64(tp) / float64(tp+fp)
	}
	if tp+fn > 0 {
		recall = float64(tp) / float64(tp+fn)
	}
	var f1 float64
	if precision+recall > 0 {
		f1 = (2 * precision * recall) / (precision + recall)
	}
	return PRF1{Precision: precision, Recall: recall, F1: f1}, nil
}

// MicroAverage pools per-class confusion counts before dividing. It is the Go
// twin of microAverage() in src/lib/evalMetrics.ts: the sums are what get
// validated, so (exactly as in TS) per-class negatives that cancel out do not
// error.
func MicroAverage(perClass []ConfusionCounts) (PRF1, error) {
	sums := ConfusionCounts{}
	for _, c := range perClass {
		sums.TP += c.TP
		sums.FP += c.FP
		sums.FN += c.FN
	}
	return PrecisionRecall(sums)
}

// ExactMatchRate is the exact-match rate over paired predictions/references
// (case-sensitive); empty input scores 0. It is the Go twin of
// exactMatchRate() in src/lib/evalMetrics.ts.
func ExactMatchRate(predictions, references []string) (float64, error) {
	if len(predictions) != len(references) {
		return 0, errors.New("predictions and references must have the same length")
	}
	if len(predictions) == 0 {
		return 0, nil
	}
	hits := 0
	for i := range predictions {
		if predictions[i] == references[i] {
			hits++
		}
	}
	return float64(hits) / float64(len(predictions)), nil
}

// SetMatch computes P/R/F1 over label SETS — the standard multi-label /
// extraction metric. It is the Go twin of setMatch() in
// src/lib/evalMetrics.ts and never errors (its counts are set cardinalities).
func SetMatch(prediction, reference []string) PRF1 {
	p := make(map[string]struct{}, len(prediction))
	for _, label := range prediction {
		p[label] = struct{}{}
	}
	r := make(map[string]struct{}, len(reference))
	for _, label := range reference {
		r[label] = struct{}{}
	}
	tp := 0
	for label := range r {
		if _, ok := p[label]; ok {
			tp++
		}
	}
	fp := 0
	for label := range p {
		if _, ok := r[label]; !ok {
			fp++
		}
	}
	fn := len(r) - tp
	prf1, _ := PrecisionRecall(ConfusionCounts{TP: tp, FP: fp, FN: fn})
	return prf1
}