correlation_test.go

package stats_test

import (
	"fmt"
	"math"
	"testing"

	"github.com/montanaflynn/stats"
)

func ExampleCorrelation() {
	s1 := []float64{1, 2, 3, 4, 5}
	s2 := []float64{1, 2, 3, 5, 6}
	a, _ := stats.Correlation(s1, s2)
	rounded, _ := stats.Round(a, 5)
	fmt.Println(rounded)
	// Output: 0.99124
}

func TestCorrelation(t *testing.T) {
	s1 := []float64{1, 2, 3, 4, 5}
	s2 := []float64{10, -51.2, 8}
	s3 := []float64{1, 2, 3, 5, 6}
	s4 := []float64{}
	s5 := []float64{0, 0, 0}
	testCases := []struct {
		name   string
		input  [][]float64
		output float64
		err    error
	}{
		{"Empty Slice Error", [][]float64{s4, s4}, math.NaN(), stats.EmptyInputErr},
		{"Different Length Error", [][]float64{s1, s2}, math.NaN(), stats.SizeErr},
		{"Correlation Value", [][]float64{s1, s3}, 0.9912407071619302, nil},
		{"Same Input Value", [][]float64{s5, s5}, 0.00, stats.ErrZero},
		{"Single Pair", [][]float64{{1}, {2}}, 0.00, stats.ErrTooFew},
	}
	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			a, err := stats.Correlation(tc.input[0], tc.input[1])
			if err != nil {
				if err != tc.err {
					t.Errorf("Should have returned error %s", tc.err)
				}
			} else if !veryclose(a, tc.output) {
				t.Errorf("Result %.08f should be %.08f", a, tc.output)
			}
			a2, err2 := stats.Pearson(tc.input[0], tc.input[1])
			if err2 != nil {
				if err2 != tc.err {
					t.Errorf("Should have returned error %s", tc.err)
				}
			} else if !veryclose(a2, tc.output) {
				t.Errorf("Result %.08f should be %.08f", a2, tc.output)
			}
		})
	}
}

func TestCorrelationScale(t *testing.T) {
	a := []float64{3.2, 1.1, 4.8, 2.5, 9.0, 7.3, 0.4, 5.5}
	b := []float64{1.5, 2.7, 3.1, 4.9, 5.2, 6.8, 7.4, 8.6}
	// scipy.stats.pearsonr(a, b) at every scale below
	want := 0.22609894303162795
	for _, scale := range []float64{1, 1e200, 1e-200, -1e200, 1e300, 1e-300} {
		t.Run(fmt.Sprintf("scale=%g", scale), func(t *testing.T) {
			sa, sb := make([]float64, len(a)), make([]float64, len(b))
			for i := range a {
				sa[i], sb[i] = a[i]*scale, b[i]*math.Abs(scale)
			}
			got, err := stats.Correlation(sa, sb)
			if err != nil {
				t.Fatal(err)
			}
			if scale < 0 {
				got = -got
			}
			if !tolerance(got, want, 1e-12) {
				t.Errorf("Correlation = %v, want %v", got, want)
			}
		})
	}
}

func TestCorrelationNonFinite(t *testing.T) {
	for _, data := range [][]float64{
		{1, math.NaN(), 3},
		{1, math.Inf(1), 3},
	} {
		got, err := stats.Correlation(data, []float64{1, 2, 3})
		if err != nil || !math.IsNaN(got) {
			t.Errorf("Correlation(%v) = (%v, %v), want (NaN, nil)", data, got, err)
		}
	}
}

func ExampleSpearman() {
	s1 := []float64{1, 2, 3, 4, 5}
	s2 := []float64{5, 6, 7, 8, 7}
	a, _ := stats.Spearman(s1, s2)
	rounded, _ := stats.Round(a, 5)
	fmt.Println(rounded)
	// Output: 0.82078
}

func TestSpearman(t *testing.T) {
	testCases := []struct {
		name   string
		data1  []float64
		data2  []float64
		output float64
		err    error
	}{
		// Error cases
		{"Empty Slice Error", []float64{}, []float64{}, math.NaN(), stats.EmptyInputErr},
		{"Different Length Error", []float64{1, 2, 3, 4, 5}, []float64{10, -51.2, 8}, math.NaN(), stats.SizeErr},

		// Perfect correlation
		{"Perfect Positive", []float64{1, 2, 3, 4, 5}, []float64{1, 2, 3, 4, 5}, 1.0, nil},
		{"Perfect Negative", []float64{1, 2, 3, 4, 5}, []float64{5, 4, 3, 2, 1}, -1.0, nil},

		// Verified against SciPy spearmanr and R cor.test(method="spearman")
		{
			"Basic with tie in y (SciPy/R)",
			[]float64{1, 2, 3, 4, 5},
			[]float64{5, 6, 7, 8, 7},
			0.8207826816681233,
			nil,
		},
		{
			"Negative with tie (SciPy mstats)",
			[]float64{5.05, 6.75, 3.21, 2.66},
			[]float64{1.65, 2.64, 2.64, 6.95},
			-0.632455532033676,
			nil,
		},
		{
			"15 elements with ties in both (SciPy/R)",
			[]float64{2.0, 47.4, 42.0, 10.8, 60.1, 1.7, 64.0, 63.1, 1.0, 1.4, 7.9, 0.3, 3.9, 0.3, 6.7},
			[]float64{22.6, 8.3, 44.4, 11.9, 24.6, 0.6, 5.7, 41.6, 0.0, 0.6, 6.7, 3.8, 1.0, 1.2, 1.4},
			0.6887298747763864,
			nil,
		},
		{
			"Adjacent swap (SciPy)",
			[]float64{0, 1, 2, 3, 4, 5},
			[]float64{0, 1, 2, 3, 5, 4},
			0.9428571428571429,
			nil,
		},
		{
			"Heavy ties in both (SciPy)",
			[]float64{1, 1, 1, 2, 2, 2},
			[]float64{1, 1, 2, 2, 3, 3},
			0.8164965809277261,
			nil,
		},

		// Zero standard deviation case (all same values)
		{"Same Input Value", []float64{0, 0, 0}, []float64{0, 0, 0}, 0.00, stats.ErrZero},
		{"Single Pair", []float64{1}, []float64{2}, 0.00, stats.ErrTooFew},
	}
	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			a, err := stats.Spearman(tc.data1, tc.data2)
			if err != nil {
				if err != tc.err {
					t.Errorf("Should have returned error %s, got %s", tc.err, err)
				}
			} else if !veryclose(a, tc.output) {
				t.Errorf("Result %.16f should be %.16f", a, tc.output)
			}

			// Also test the Float64Data method
			a2, err2 := stats.Float64Data(tc.data1).Spearman(tc.data2)
			if err2 != nil {
				if err2 != tc.err {
					t.Errorf("Should have returned error %s, got %s", tc.err, err2)
				}
			} else if !veryclose(a2, tc.output) {
				t.Errorf("Result %.16f should be %.16f", a2, tc.output)
			}
		})
	}
}

func BenchmarkSpearman(b *testing.B) {
	s1 := []float64{2.0, 47.4, 42.0, 10.8, 60.1, 1.7, 64.0, 63.1, 1.0, 1.4, 7.9, 0.3, 3.9, 0.3, 6.7}
	s2 := []float64{22.6, 8.3, 44.4, 11.9, 24.6, 0.6, 5.7, 41.6, 0.0, 0.6, 6.7, 3.8, 1.0, 1.2, 1.4}
	for i := 0; i < b.N; i++ {
		_, _ = stats.Spearman(s1, s2)
	}
}

func ExampleAutoCorrelation() {
	s1 := []float64{1, 2, 3, 4, 5}
	a, _ := stats.AutoCorrelation(s1, 1)
	fmt.Println(a)
	// Output: 0.4
}

func TestAutoCorrelation(t *testing.T) {
	s1 := []float64{1, 2, 3, 4, 5}
	s2 := []float64{}

	a, err := stats.AutoCorrelation(s1, 1)
	if err != nil {
		t.Errorf("Should not have returned an error")
	}
	if a != 0.4 {
		t.Errorf("Should have returned 0.4")
	}

	_, err = stats.AutoCorrelation(s2, 1)
	if err != stats.EmptyInputErr {
		t.Errorf("Should have returned empty input error")
	}

	// Reference values cross-checked against statsmodels.tsa.stattools.acf,
	// R acf(), and MATLAB autocorr (issue #83).
	s3 := []float64{22, 24, 25, 25, 28, 29, 34, 37, 40, 44, 51, 48, 47, 50, 51}
	expected := []float64{
		1.0,
		0.83174224,
		0.65632458,
		0.49105012,
		0.27863962,
		0.03102625,
		-0.16527446,
		-0.30369928,
		-0.40095465,
		-0.45823389,
		-0.45047733,
	}
	for lag, want := range expected {
		got, err := stats.AutoCorrelation(s3, lag)
		if err != nil {
			t.Errorf("AutoCorrelation(s3, %d) returned error: %v", lag, err)
		}
		if math.Abs(got-want) > 1e-8 {
			t.Errorf("AutoCorrelation(s3, %d) = %v, want %v", lag, got, want)
		}
	}

	if _, err := stats.AutoCorrelation(s1, -1); err != stats.BoundsErr {
		t.Errorf("Should have returned bounds error for negative lag")
	}
	if _, err := stats.AutoCorrelation(s1, len(s1)); err != stats.BoundsErr {
		t.Errorf("Should have returned bounds error for lag >= len(data)")
	}

	constant := []float64{3, 3, 3, 3, 3}
	if a, err := stats.AutoCorrelation(constant, 1); err != stats.ErrZero || a != 0 {
		t.Errorf("AutoCorrelation of constant series should be 0 with ErrZero, got %v err=%v", a, err)
	}
}

func TestAutoCorrelationExtremeScales(t *testing.T) {
	base := stats.Float64Data{1, 2, 3, 10, 4}
	want, _ := stats.AutoCorrelation(base, 1)
	for _, s := range []float64{1e200, 1e-200} {
		d := make(stats.Float64Data, len(base))
		for i, v := range base {
			d[i] = v * s
		}
		if got, err := stats.AutoCorrelation(d, 1); err != nil || !close(got, want) {
			t.Errorf("AutoCorrelation(%v, 1) = %v, %v; want %v", d, got, err, want)
		}
	}
}