Add series.Rolling (go-gota#76) and RollingWindow.Mean

series/rolling_window.go

@@ -0,0 +1,39 @@
+package series
+
+import "math"
+
+// RollingWindow is used for rolling window calculations.
+type RollingWindow struct {
+	window int
+	series Series
+}
+
+// Rolling creates new RollingWindow
+func (s Series) Rolling(window int) RollingWindow {
+	return RollingWindow{
+		window: window,
+		series: s,
+	}
+}
+
+// Mean returns the rolling mean.
+func (r RollingWindow) Mean() (s Series) {
+	s = New([]float64{}, Float, "mean")
+
+	for i := 1; i <= r.series.Len(); i++ {
+		if i < r.window {
+			s.Append(math.NaN())
+			continue
+		}
+
+		index := []int{}
+		for j := i - r.window; j < i; j++ {
+			index = append(index, j)
+		}
+
+		mean := r.series.Subset(index).Mean()
+		s.Append(mean)
+	}
+
+	return
+}

series/rolling_window_test.go

@@ -0,0 +1,46 @@
+package series
+
+import (
+	"math"
+	"strings"
+	"testing"
+)
+
+func TestSeries_RollingMean(t *testing.T) {
+	tests := []struct {
+		window   int
+		series   Series
+		expected Series
+	}{
+		{
+			3,
+			Ints([]int{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}),
+			Floats([]float64{math.NaN(), math.NaN(), 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0}),
+		},
+		{
+			2,
+			Floats([]float64{1.0, 2.0, 3.0}),
+			Floats([]float64{math.NaN(), 1.5, 2.5}),
+		},
+		{
+			0,
+			Floats([]float64{}),
+			Floats([]float64{}),
+		},
+	}
+
+	for testnum, test := range tests {
+		expected := test.expected
+		received := test.series.Rolling(test.window).Mean()
+
+		for i := 0; i < expected.Len(); i++ {
+			if strings.Compare(expected.Elem(i).String(),
+				received.Elem(i).String()) != 0 {
+				t.Errorf(
+					"Test:%v\nExpected:\n%v\nReceived:\n%v",
+					testnum, expected, received,
+				)
+			}
+		}
+	}
+}