Skip to content

Commit

Permalink
Implement rolling window Mean and StdDev (#117)
Browse files Browse the repository at this point in the history
* Add series.Rolling (#76) and RollingWindow.Mean

* Add RollingWindow.StdDev
  • Loading branch information
jfussion authored Apr 25, 2021
1 parent 478e4ee commit 218714d
Show file tree
Hide file tree
Showing 2 changed files with 137 additions and 0 deletions.
52 changes: 52 additions & 0 deletions series/rolling_window.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,52 @@
package series

// 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 _, block := range r.getBlocks() {
s.Append(block.Mean())
}

return
}

// StdDev returns the rolling mean.
func (r RollingWindow) StdDev() (s Series) {
s = New([]float64{}, Float, "StdDev")
for _, block := range r.getBlocks() {
s.Append(block.StdDev())
}

return
}

func (r RollingWindow) getBlocks() (blocks []Series) {
for i := 1; i <= r.series.Len(); i++ {
if i < r.window {
blocks = append(blocks, r.series.Empty())
continue
}

index := []int{}
for j := i - r.window; j < i; j++ {
index = append(index, j)
}
blocks = append(blocks, r.series.Subset(index))
}

return
}
85 changes: 85 additions & 0 deletions series/rolling_window_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,85 @@
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,
)
}
}
}
}

func TestSeries_RollingStdDev(t *testing.T) {
tests := []struct {
window int
series Series
expected Series
}{
{
3,
Ints([]int{5, 5, 6, 7, 5, 5, 5}),
Floats([]float64{math.NaN(), math.NaN(), 0.5773502691896257, 1.0, 1.0, 1.1547005383792515, 0.0}),
},
{
2,
Floats([]float64{1.0, 2.0, 3.0}),
Floats([]float64{math.NaN(), 0.7071067811865476, 0.7071067811865476}),
},
{
0,
Floats([]float64{}),
Floats([]float64{}),
},
}

for testnum, test := range tests {
expected := test.expected
received := test.series.Rolling(test.window).StdDev()

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,
)
}
}
}
}

0 comments on commit 218714d

Please sign in to comment.