Statistics and Regression Analysis

// Package special holds the numerical machinery behind the t, chi-square and
// F distributions: the regularized incomplete beta and gamma functions, and
// the distribution functions built on them.
//
// You don't write these. They're numerical-methods code (continued fractions
// and series, after Numerical Recipes), not statistics. Module 4 has you build
// the t distribution on top of RegIncBeta yourself; modules 5–10 use the
// reference copies here, so a stuck module never blocks a later one.
package special

import "math"

const (
	eps     = 1e-15
	tiny    = 1e-300
	maxIter = 500
)

// RegIncBeta is the regularized incomplete beta function I_x(a, b) for
// a, b > 0 and 0 <= x <= 1.
func RegIncBeta(a, b, x float64) float64 {
	switch {
	case x <= 0:
		return 0
	case x >= 1:
		return 1
	}
	la, _ := math.Lgamma(a)
	lb, _ := math.Lgamma(b)
	lab, _ := math.Lgamma(a + b)
	front := math.Exp(lab - la - lb + a*math.Log(x) + b*math.Log1p(-x))
	// The continued fraction converges fast for x < (a+1)/(a+b+2); use the
	// symmetry I_x(a,b) = 1 - I_{1-x}(b,a) on the other side.
	if x < (a+1)/(a+b+2) {
		return front * betaCF(a, b, x) / a
	}
	return 1 - front*betaCF(b, a, 1-x)/b
}

// betaCF evaluates the continued fraction for I_x(a, b) by Lentz's method.
func betaCF(a, b, x float64) float64 {
	qab, qap, qam := a+b, a+1, a-1
	c, d := 1.0, 1-qab*x/qap
	if math.Abs(d) < tiny {
		d = tiny
	}
	d = 1 / d
	h := d
	for m := 1; m <= maxIter; m++ {
		fm := float64(m)
		m2 := 2 * fm
		aa := fm * (b - fm) * x / ((qam + m2) * (a + m2))
		d = 1 + aa*d
		if math.Abs(d) < tiny {
			d = tiny
		}
		c = 1 + aa/c
		if math.Abs(c) < tiny {
			c = tiny
		}
		d = 1 / d
		h *= d * c
		aa = -(a + fm) * (qab + fm) * x / ((a + m2) * (qap + m2))
		d = 1 + aa*d
		if math.Abs(d) < tiny {
			d = tiny
		}
		c = 1 + aa/c
		if math.Abs(c) < tiny {
			c = tiny
		}
		d = 1 / d
		del := d * c
		h *= del
		if math.Abs(del-1) < eps {
			break
		}
	}
	return h
}

// RegIncGamma is the regularized lower incomplete gamma function P(a, x)
// for a > 0 and x >= 0.
func RegIncGamma(a, x float64) float64 {
	if x <= 0 {
		return 0
	}
	lg, _ := math.Lgamma(a)
	if x < a+1 {
		// series: P = e^-x x^a / Γ(a) · Σ x^n / (a(a+1)…(a+n))
		sum, term, ap := 1/a, 1/a, a
		for n := 0; n < maxIter; n++ {
			ap++
			term *= x / ap
			sum += term
			if math.Abs(term) < math.Abs(sum)*eps {
				break
			}
		}
		return sum * math.Exp(-x+a*math.Log(x)-lg)
	}
	// continued fraction for Q = 1 - P, by Lentz's method
	b := x + 1 - a
	c, d := 1/tiny, 1/b
	h := d
	for i := 1; i <= maxIter; i++ {
		an := -float64(i) * (float64(i) - a)
		b += 2
		d = an*d + b
		if math.Abs(d) < tiny {
			d = tiny
		}
		c = b + an/c
		if math.Abs(c) < tiny {
			c = tiny
		}
		d = 1 / d
		del := d * c
		h *= del
		if math.Abs(del-1) < eps {
			break
		}
	}
	return 1 - math.Exp(-x+a*math.Log(x)-lg)*h
}

// NormalCDF is P(Z <= z) for a standard normal Z.
func NormalCDF(z float64) float64 { return 0.5 * math.Erfc(-z/math.Sqrt2) }

// NormalQuantile is the z with P(Z <= z) = p, for 0 < p < 1.
func NormalQuantile(p float64) float64 { return math.Sqrt2 * math.Erfinv(2*p-1) }

// TCDF is P(T <= t) for Student's t distribution with df degrees of freedom.
func TCDF(t, df float64) float64 {
	// tail = P(T > |t|) = ½·I_{df/(df+t²)}(df/2, ½). Near t = 0 that x rounds
	// to exactly 1, so use the mirror form ½·(1 − I_{t²/(df+t²)}(½, df/2)).
	var tail float64
	if t*t < df {
		tail = 0.5 * (1 - RegIncBeta(0.5, df/2, t*t/(df+t*t)))
	} else {
		tail = 0.5 * RegIncBeta(df/2, 0.5, df/(df+t*t))
	}
	if t >= 0 {
		return 1 - tail
	}
	return tail
}

// TQuantile is the t with P(T <= t) = p, for 0 < p < 1.
func TQuantile(p, df float64) float64 {
	return invert(func(t float64) float64 { return TCDF(t, df) }, p)
}

// ChiSquareCDF is P(X <= x) for a chi-square distribution with k degrees of
// freedom.
func ChiSquareCDF(x, k float64) float64 { return RegIncGamma(k/2, x/2) }

// FCDF is P(F <= f) for an F distribution with d1 and d2 degrees of freedom.
func FCDF(f, d1, d2 float64) float64 {
	if f <= 0 {
		return 0
	}
	return RegIncBeta(d1/2, d2/2, d1*f/(d1*f+d2))
}

// FSF is the upper tail P(F > f), computed directly so that tiny p-values
// keep their digits (1 − FCDF would round them away).
func FSF(f, d1, d2 float64) float64 {
	if f <= 0 {
		return 1
	}
	return RegIncBeta(d2/2, d1/2, d2/(d2+d1*f))
}

// invert finds x with cdf(x) = p for an increasing cdf on the whole real
// line: widen [lo, hi] until it brackets p, then bisect.
func invert(cdf func(float64) float64, p float64) float64 {
	lo, hi := -1.0, 1.0
	for cdf(lo) > p {
		lo *= 2
	}
	for cdf(hi) < p {
		hi *= 2
	}
	for i := 0; i < 200 && hi-lo > 1e-12*math.Max(1, math.Abs(lo)); i++ {
		mid := (lo + hi) / 2
		if cdf(mid) < p {
			lo = mid
		} else {
			hi = mid
		}
	}
	return (lo + hi) / 2
}

package special

import (
	"math"
	"testing"
)

// Every expected value below comes from SciPy (scipy.special.betainc and
// gammainc, scipy.stats.t, chi2, f and norm).
func TestAgainstSciPy(t *testing.T) {
	for _, tc := range []struct {
		name      string
		got, want float64
	}{
		{"RegIncBeta(0.5, 0.5, 0.37)", RegIncBeta(0.5, 0.5, 0.37), 0.41627743252839533},
		{"RegIncBeta(2.5, 3, 0.1)", RegIncBeta(2.5, 3, 0.1), 0.021483723853768934},
		{"RegIncBeta(7, 12, 0.82)", RegIncBeta(7, 12, 0.82), 0.9999927548696877},
		{"RegIncBeta(30, 0.5, 0.999)", RegIncBeta(30, 0.5, 0.999), 0.8072373061595369},
		{"RegIncBeta(1, 1, 0.5)", RegIncBeta(1, 1, 0.5), 0.5},
		{"RegIncGamma(0.5, 0.01)", RegIncGamma(0.5, 0.01), 0.11246291601828491},
		{"RegIncGamma(2.5, 3)", RegIncGamma(2.5, 3), 0.6937810815867212},
		{"RegIncGamma(7, 9)", RegIncGamma(7, 9), 0.7932191601400129},
		{"RegIncGamma(30, 40)", RegIncGamma(30, 40), 0.9567713178482644},
		{"TCDF(-40, 1)", TCDF(-40.0, 1.0), 0.007956089912025812},
		{"TCDF(-1.2, 4.5)", TCDF(-1.2, 4.5), 0.14473852742098223},
		{"TCDF(0, 9)", TCDF(0.0, 9.0), 0.5},
		{"TCDF(2.1, 30)", TCDF(2.1, 30.0), 0.9778787643688382},
		{"TCDF(6, 200)", TCDF(6.0, 200.0), 0.9999999954424398},
		{"TQuantile(0.0005, 1)", TQuantile(0.0005, 1.0), -636.6192487687196},
		{"TQuantile(0.025, 9)", TQuantile(0.025, 9.0), -2.2621571627982053},
		{"TQuantile(0.5, 4.5)", TQuantile(0.5, 4.5), 0.0},
		{"TQuantile(0.975, 30)", TQuantile(0.975, 30.0), 2.0422724563012378},
		{"TQuantile(0.9995, 200)", TQuantile(0.9995, 200.0), 3.3398354062757107},
		{"ChiSquareCDF(0.2, 1)", ChiSquareCDF(0.2, 1.0), 0.34527915398142317},
		{"ChiSquareCDF(4, 5)", ChiSquareCDF(4.0, 5.0), 0.4505840486472198},
		{"ChiSquareCDF(30, 13)", ChiSquareCDF(30.0, 13.0), 0.9952902952345285},
		{"FCDF(1, 3, 25)", FCDF(1.0, 3.0, 25.0), 0.590849712985762},
		{"FSF(1, 3, 25)", FSF(1.0, 3.0, 25.0), 0.40915028701423795},
		{"FCDF(2.5, 7, 2)", FCDF(2.5, 7.0, 2.0), 0.6847184505954043},
		{"FSF(2.5, 7, 2)", FSF(2.5, 7.0, 2.0), 0.31528154940459574},
		{"FCDF(444.751, 3, 11)", FCDF(444.75050444838416, 3.0, 11.0), 0.999999999990704},
		{"FSF(444.751, 3, 11)", FSF(444.75050444838416, 3.0, 11.0), 9.296103842548068e-12},
		{"NormalCDF(-9)", NormalCDF(-9.0), 1.1285884059538324e-19},
		{"NormalCDF(-1.96)", NormalCDF(-1.96), 0.024997895148220435},
		{"NormalCDF(0)", NormalCDF(0.0), 0.5},
		{"NormalCDF(1.3)", NormalCDF(1.3), 0.9031995154143897},
		{"NormalQuantile(0.001)", NormalQuantile(0.001), -3.090232306167813},
		{"NormalQuantile(0.5)", NormalQuantile(0.5), 0.0},
		{"NormalQuantile(0.975)", NormalQuantile(0.975), 1.959963984540054},
	} {
		// 1e-9 relative, with a floor for values near zero; TQuantile's
		// bisection stops at about 1e-12, so this has room to spare
		tol := 1e-9 * math.Max(math.Abs(tc.want), 1e-300)
		if tc.want == 0 {
			tol = 1e-9
		}
		if math.Abs(tc.got-tc.want) > tol {
			t.Errorf("%s = %v, want %v", tc.name, tc.got, tc.want)
		}
	}
}

// Package linalg is the small amount of matrix arithmetic that regression
// needs, on plain [][]float64 (row-major: m[i][j] is row i, column j).
//
// Module 8 has you build these operations yourself. Modules 9 and 10 use the
// reference copies here, so they don't depend on your module 8 code.
package linalg

import (
	"errors"
	"math"
)

// ErrSingular is returned by Inverse when the matrix has no inverse — in a
// regression, that means one column is a combination of the others.
var ErrSingular = errors.New("linalg: matrix is singular")

// Transpose returns aᵀ.
func Transpose(a [][]float64) [][]float64 {
	if len(a) == 0 {
		return nil
	}
	t := make([][]float64, len(a[0]))
	for j := range t {
		t[j] = make([]float64, len(a))
		for i := range a {
			t[j][i] = a[i][j]
		}
	}
	return t
}

// Mul returns the matrix product ab. Columns of a must equal rows of b.
func Mul(a, b [][]float64) [][]float64 {
	out := make([][]float64, len(a))
	for i := range a {
		out[i] = make([]float64, len(b[0]))
		for k, aik := range a[i] {
			for j, bkj := range b[k] {
				out[i][j] += aik * bkj
			}
		}
	}
	return out
}

// MulVec returns the matrix-vector product av.
func MulVec(a [][]float64, v []float64) []float64 {
	out := make([]float64, len(a))
	for i, row := range a {
		for j, x := range row {
			out[i] += x * v[j]
		}
	}
	return out
}

// Inverse returns a⁻¹ for a square matrix, by Gauss–Jordan elimination with
// partial pivoting. a is not modified.
func Inverse(a [][]float64) ([][]float64, error) {
	n := len(a)
	// augment [a | I]
	m := make([][]float64, n)
	for i := range a {
		m[i] = make([]float64, 2*n)
		copy(m[i], a[i])
		m[i][n+i] = 1
	}
	for col := 0; col < n; col++ {
		// pick the row with the largest entry in this column
		pivot := col
		for r := col + 1; r < n; r++ {
			if math.Abs(m[r][col]) > math.Abs(m[pivot][col]) {
				pivot = r
			}
		}
		if math.Abs(m[pivot][col]) < 1e-12 {
			return nil, ErrSingular
		}
		m[col], m[pivot] = m[pivot], m[col]
		p := m[col][col]
		for j := range m[col] {
			m[col][j] /= p
		}
		for r := 0; r < n; r++ {
			if r == col || m[r][col] == 0 {
				continue
			}
			f := m[r][col]
			for j := range m[r] {
				m[r][j] -= f * m[col][j]
			}
		}
	}
	inv := make([][]float64, n)
	for i := range m {
		inv[i] = m[i][n:]
	}
	return inv, nil
}

.PHONY: test solutions progress vet

# Test your exercises (every module). Expect failures until you implement them.
test:
	go test ./0* ./10*

# Test the reference solutions — should be all green.
solutions:
	go test -tags solutions ./...

# Show a per-module progress report.
progress:
	go run ./cmd/progress

# Vet both builds.
vet:
	go vet ./...
	go vet -tags solutions ./...

module github.com/bozcode/go-statistics-and-regression

go 1.22

// Command progress runs the course tests and shows how far along you are.
//
//	go run ./cmd/progress                 # your exercises
//	go run ./cmd/progress -tags solutions # the reference solutions (all green)
//
// Any arguments are passed straight through to `go test`.
package main

import (
	"bufio"
	"bytes"
	"encoding/json"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"sort"
	"strings"
)

type module struct {
	dir      string   // e.g. "01-arrays"
	title    string   // e.g. "Arrays"
	tests    []string // top-level test names, in file order
	passed   map[string]bool
	failed   map[string]bool
	buildErr bool
}

var testFunc = regexp.MustCompile(`(?m)^func (Test\w+)\(t \*testing\.T\)`)

func main() {
	root, err := findRoot()
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}

	mods, err := discover(root)
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}

	args := append([]string{"test", "-json", "-count=1"}, os.Args[1:]...)
	for _, m := range mods {
		args = append(args, "./"+m.dir+"/")
	}
	cmd := exec.Command("go", args...)
	cmd.Dir = root
	var out bytes.Buffer
	cmd.Stdout = &out
	cmd.Stderr = &out
	_ = cmd.Run() // a non-zero exit just means some tests fail

	byPkg := map[string]*module{}
	for _, m := range mods {
		byPkg[m.dir] = m
	}
	sc := bufio.NewScanner(&out)
	sc.Buffer(make([]byte, 1024*1024), 1024*1024)
	for sc.Scan() {
		var ev struct {
			Action  string
			Package string
			Test    string
			Output  string
		}
		if json.Unmarshal(sc.Bytes(), &ev) != nil {
			continue
		}
		m := byPkg[filepath.Base(ev.Package)]
		if m == nil {
			continue
		}
		if ev.Test == "" {
			if ev.Action == "output" && strings.Contains(ev.Output, "[build failed]") {
				m.buildErr = true
			}
			continue
		}
		if strings.Contains(ev.Test, "/") {
			continue // subtest
		}
		switch ev.Action {
		case "pass":
			m.passed[ev.Test] = true
		case "fail":
			m.failed[ev.Test] = true
		}
	}

	report(mods)
}

func report(mods []*module) {
	fmt.Println()
	fmt.Println("  Statistics and Regression Analysis in Go — progress")
	fmt.Println()
	width := 0 // pad titles to the longest one so the bars line up
	for _, m := range mods {
		width = max(width, len(m.title))
	}
	total, done := 0, 0
	var nextMod *module
	var nextTest string
	for i, m := range mods {
		n := len(m.tests)
		p := 0
		for _, t := range m.tests {
			if m.passed[t] {
				p++
			} else if nextMod == nil {
				nextMod, nextTest = m, t
			}
		}
		total += n
		done += p

		status := ""
		switch {
		case m.buildErr:
			status = "  ✗ doesn't compile — run: go vet ./" + m.dir + "/"
		case p == n:
			status = "  ✓ complete"
		}
		fmt.Printf("  %2d  %-*s  %s  %2d/%-2d%s\n", i+1, width, m.title, bar(p, n, 12), p, n, status)
	}
	fmt.Println()
	fmt.Printf("  Overall: %d/%d tests passing (%d%%)\n", done, total, pct(done, total))
	if nextMod != nil {
		fmt.Printf("  Next up: %s → %s\n", nextMod.title, nextTest)
		fmt.Printf("           go test -run '^%s$' -v ./%s/\n", nextTest, nextMod.dir)
	} else {
		fmt.Println("  Every test passes. Nice work.")
	}
	fmt.Println()
}

func bar(p, n, width int) string {
	if n == 0 {
		return strings.Repeat("·", width)
	}
	filled := p * width / n
	return strings.Repeat("█", filled) + strings.Repeat("░", width-filled)
}

func pct(a, b int) int {
	if b == 0 {
		return 0
	}
	return a * 100 / b
}

// discover finds the numbered module directories and their test names.
func discover(root string) ([]*module, error) {
	entries, err := os.ReadDir(root)
	if err != nil {
		return nil, err
	}
	numbered := regexp.MustCompile(`^\d\d-`)
	var mods []*module
	for _, e := range entries {
		if !e.IsDir() || !numbered.MatchString(e.Name()) {
			continue
		}
		m := &module{
			dir:    e.Name(),
			title:  titleFor(e.Name()),
			passed: map[string]bool{},
			failed: map[string]bool{},
		}
		files, _ := filepath.Glob(filepath.Join(root, e.Name(), "*_test.go"))
		sort.Strings(files)
		for _, f := range files {
			src, err := os.ReadFile(f)
			if err != nil {
				return nil, err
			}
			for _, match := range testFunc.FindAllSubmatch(src, -1) {
				m.tests = append(m.tests, string(match[1]))
			}
		}
		mods = append(mods, m)
	}
	sort.Slice(mods, func(i, j int) bool { return mods[i].dir < mods[j].dir })
	if len(mods) == 0 {
		return nil, fmt.Errorf("no module directories found under %s", root)
	}
	return mods, nil
}

// titleFor turns "03-linked-lists" into "Linked Lists".
func titleFor(dir string) string {
	words := strings.Split(dir[3:], "-")
	for i, w := range words {
		words[i] = strings.ToUpper(w[:1]) + w[1:]
	}
	return strings.Join(words, " ")
}

// findRoot walks up from the working directory to the directory holding go.mod.
func findRoot() (string, error) {
	dir, err := os.Getwd()
	if err != nil {
		return "", err
	}
	for {
		if _, err := os.Stat(filepath.Join(dir, "go.mod")); err == nil {
			return dir, nil
		}
		parent := filepath.Dir(dir)
		if parent == dir {
			return "", fmt.Errorf("go.mod not found — run this from inside the course folder")
		}
		dir = parent
	}
}