regression/linear_test.go at master · toolsparty/regression · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
package regression

import (
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func TestNewLinear(t *testing.T) {
	testCases := []*struct {
		Name    string
		X, Y    []float64
		WantErr string
	}{
		{
			Name:    "axis y is empty",
			X:       nil,
			Y:       nil,
			WantErr: "axis y is empty",
		},
		{
			Name:    "axis y is empty",
			X:       nil,
			Y:       []float64{},
			WantErr: "axis y is empty",
		},
		{
			Name:    "axis y is empty",
			X:       []float64{},
			Y:       []float64{},
			WantErr: "axis y is empty",
		},
		{
			Name:    "axis y is empty",
			X:       []float64{1},
			Y:       []float64{},
			WantErr: "axis y is empty",
		},
		{
			Name:    "different lengths",
			X:       []float64{1},
			Y:       []float64{1, 2},
			WantErr: "axes x and y has a different lengths",
		},
		{
			Name:    "success",
			X:       []float64{},
			Y:       []float64{1, 2},
			WantErr: "",
		},
		{
			Name:    "success",
			X:       []float64{},
			Y:       []float64{1, 2},
			WantErr: "",
		},
		{
			Name:    "success",
			X:       []float64{2, 4},
			Y:       []float64{1, 2},
			WantErr: "",
		},
	}

	for _, testCase := range testCases {
		t.Run(testCase.Name, func(t *testing.T) {
			got, err := NewLinear(testCase.X, testCase.Y)
			if testCase.WantErr != "" {
				assert.EqualError(t, err, testCase.WantErr)
				assert.Nil(t, got)
			} else {
				assert.NoError(t, err)
				assert.NotNil(t, got)
			}
		})
	}
}

func TestLinear_GetKB(t *testing.T) {
	y := []int{3, 5, 7, 9, 11}
	reg, err := NewLinear(nil, y)
	require.NoError(t, err)
	require.NotNil(t, reg)

	assert.Equal(t, 2.0, reg.GetK())
	assert.Equal(t, 1.0, reg.GetB())
}

func TestLinear_Predict(t *testing.T) {
	y := []int{3, 5, 7, 9, 11}
	reg, err := NewLinear(nil, y)
	require.NoError(t, err)
	require.NotNil(t, reg)

	got, err := reg.Predict(2)
	assert.NoError(t, nil)
	assert.Equal(t, 5, got)

	got, err = reg.Predict(5)
	assert.NoError(t, nil)
	assert.Equal(t, 11, got)

	got, err = reg.Predict(10)
	assert.NoError(t, nil)
	assert.Equal(t, 21, got)
}

func TestLinear_GetTheta(t *testing.T) {
	y := []int{3, 5, 7, 9, 11}
	reg, err := NewLinear(nil, y)
	require.NoError(t, err)
	require.NotNil(t, reg)

	theta := reg.GetTheta()
	assert.Equal(t, 1.0, theta)

	reg.Append(6, 20)

	assert.Equal(t, 3.0, reg.GetK())
	assert.Equal(t, -1.333333333333334, reg.GetB())
}