Merge pull request c9s#536 from narumiruna/indicator/atr

pkg/indicator/atr.go

@@ -0,0 +1,107 @@
+package indicator
+
+import (
+	"math"
+	"time"
+
+	"github.com/c9s/bbgo/pkg/types"
+)
+
+//go:generate callbackgen -type ATR
+type ATR struct {
+	types.IntervalWindow
+	Values               types.Float64Slice
+	TrueRanges           types.Float64Slice
+	PercentageVolatility types.Float64Slice
+	PriviousClose        float64
+
+	EndTime         time.Time
+	UpdateCallbacks []func(value float64)
+}
+
+func (inc *ATR) Update(kLine types.KLine) {
+	if inc.Window <= 0 {
+		panic("window must be greater than 0")
+	}
+
+	cloze := kLine.Close.Float64()
+	high := kLine.High.Float64()
+	low := kLine.Low.Float64()
+
+	if inc.PriviousClose == 0 {
+		inc.PriviousClose = kLine.Close.Float64()
+		return
+	}
+
+	// calculate true range
+	trueRange := types.Float64Slice{
+		high - low,
+		math.Abs(high - inc.PriviousClose),
+		math.Abs(low - inc.PriviousClose),
+	}.Max()
+	inc.TrueRanges.Push(trueRange)
+
+	inc.PriviousClose = cloze
+
+	// apply rolling moving average
+	if len(inc.TrueRanges) < inc.Window {
+		return
+	}
+
+	if len(inc.TrueRanges) == inc.Window {
+		atr := inc.TrueRanges.Mean()
+		inc.Values.Push(atr)
+		inc.PercentageVolatility.Push(atr / cloze)
+		return
+	}
+
+	lambda := 1 / float64(inc.Window)
+	atr := inc.Values.Last()*(1-lambda) + inc.TrueRanges.Last()*lambda
+	inc.Values.Push(atr)
+	inc.PercentageVolatility.Push(atr / cloze)
+}
+
+func (inc *ATR) Last() float64 {
+	if len(inc.Values) == 0 {
+		return 0.0
+	}
+	return inc.Values[len(inc.Values)-1]
+}
+
+func (inc *ATR) Index(i int) float64 {
+	length := len(inc.Values)
+	if length == 0 || length-i-1 < 0 {
+		return 0
+	}
+	return inc.Values[length-i-1]
+}
+
+func (inc *ATR) Length() int {
+	return len(inc.Values)
+}
+
+var _ types.Series = &ATR{}
+
+func (inc *ATR) calculateAndUpdate(kLines []types.KLine) {
+	for _, k := range kLines {
+		if inc.EndTime != zeroTime && !k.EndTime.After(inc.EndTime) {
+			continue
+		}
+		inc.Update(k)
+	}
+
+	inc.EmitUpdate(inc.Last())
+	inc.EndTime = kLines[len(kLines)-1].EndTime.Time()
+}
+
+func (inc *ATR) handleKLineWindowUpdate(interval types.Interval, window types.KLineWindow) {
+	if inc.Interval != interval {
+		return
+	}
+
+	inc.calculateAndUpdate(window)
+}
+
+func (inc *ATR) Bind(updater KLineWindowUpdater) {
+	updater.OnKLineWindowUpdate(inc.handleKLineWindowUpdate)
+}

pkg/indicator/atr_test.go

@@ -0,0 +1,53 @@
+package indicator
+
+import (
+	"encoding/json"
+	"math"
+	"testing"
+
+	"github.com/c9s/bbgo/pkg/fixedpoint"
+	"github.com/c9s/bbgo/pkg/types"
+)
+
+func Test_calculateATR(t *testing.T) {
+	var bytes = []byte(`{
+		"high": [40145.0, 40186.36, 40196.39, 40344.6, 40245.48, 40273.24, 40464.0, 40699.0, 40627.48, 40436.31, 40370.0, 40376.8, 40227.03, 40056.52, 39721.7, 39597.94, 39750.15, 39927.0, 40289.02, 40189.0], 
+		"low": [39870.71, 39834.98, 39866.31, 40108.31, 40016.09, 40094.66, 40105.0, 40196.48, 40154.99, 39800.0, 39959.21, 39922.98, 39940.02, 39632.0, 39261.39, 39254.63, 39473.91, 39555.51, 39819.0, 40006.84],
+		"close": [40105.78, 39935.23, 40183.97, 40182.03, 40212.26, 40149.99, 40378.0, 40618.37, 40401.03, 39990.39, 40179.13, 40097.23, 40014.72, 39667.85, 39303.1, 39519.99, 39693.79, 39827.96, 40074.94, 40059.84]
+	}`)
+	buildKLines := func(bytes []byte) (kLines []types.KLine) {
+		var prices map[string][]fixedpoint.Value
+		_ = json.Unmarshal(bytes, &prices)
+		for i, h := range prices["high"] {
+			kLine := types.KLine{High: h, Low: prices["low"][i], Close: prices["close"][i]}
+			kLines = append(kLines, kLine)
+		}
+		return kLines
+	}
+
+	tests := []struct {
+		name   string
+		kLines []types.KLine
+		window int
+		want   float64
+	}{
+		{
+			name:   "test_binance_btcusdt_1h",
+			kLines: buildKLines(bytes),
+			window: 14,
+			want:   364.048648,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			atr := ATR{IntervalWindow: types.IntervalWindow{Window: tt.window}}
+			atr.calculateAndUpdate(tt.kLines)
+			got := atr.Values.Last()
+			diff := math.Trunc((got-tt.want)*100) / 100
+			if diff != 0 {
+				t.Errorf("calculateATR() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}