indicator/linreg: LinReg indicator

pkg/indicator/linreg.go

@@ -0,0 +1,98 @@
+package indicator
+
+import (
+	"github.com/sirupsen/logrus"
+	"time"
+
+	"github.com/c9s/bbgo/pkg/datatype/floats"
+	"github.com/c9s/bbgo/pkg/types"
+)
+
+var log = logrus.WithField("indicator", "supertrend")
+
+// LinReg is Linear Regression baseline
+//go:generate callbackgen -type LinReg
+type LinReg struct {
+	types.SeriesBase
+	types.IntervalWindow
+
+	// Values are the slopes of linear regression baseline
+	Values floats.Slice
+	klines types.KLineWindow
+
+	EndTime         time.Time
+	UpdateCallbacks []func(value float64)
+}
+
+// Last slope of linear regression baseline
+func (lr *LinReg) Last() float64 {
+	if lr.Values.Length() == 0 {
+		return 0.0
+	}
+	return lr.Values.Last()
+}
+
+// Index returns the slope of specified index
+func (lr *LinReg) Index(i int) float64 {
+	if i >= lr.Values.Length() {
+		return 0.0
+	}
+
+	return lr.Values.Index(i)
+}
+
+// Length of the slope values
+func (lr *LinReg) Length() int {
+	return lr.Values.Length()
+}
+
+var _ types.SeriesExtend = &LinReg{}
+
+// Update Linear Regression baseline slope
+func (lr *LinReg) Update(kline types.KLine) {
+	lr.klines.Add(kline)
+	lr.klines.Truncate(lr.Window)
+	if len(lr.klines) < lr.Window {
+		lr.Values.Push(0)
+		return
+	}
+
+	var sumX, sumY, sumXSqr, sumXY float64 = 0, 0, 0, 0
+	end := len(lr.klines) - 1 // The last kline
+	for i := end; i >= end-lr.Window+1; i-- {
+		val := lr.klines[i].GetClose().Float64()
+		per := float64(end - i + 1)
+		sumX += per
+		sumY += val
+		sumXSqr += per * per
+		sumXY += val * per
+	}
+	length := float64(lr.Window)
+	slope := (length*sumXY - sumX*sumY) / (length*sumXSqr - sumX*sumX)
+	average := sumY / length
+	endPrice := average - slope*sumX/length + slope
+	startPrice := endPrice + slope*(length-1)
+	lr.Values.Push((endPrice - startPrice) / (length - 1))
+
+	log.Debugf("linear regression baseline slope: %f", lr.Last())
+}
+
+func (lr *LinReg) BindK(target KLineClosedEmitter, symbol string, interval types.Interval) {
+	target.OnKLineClosed(types.KLineWith(symbol, interval, lr.PushK))
+}
+
+func (lr *LinReg) PushK(k types.KLine) {
+	var zeroTime = time.Time{}
+	if lr.EndTime != zeroTime && k.EndTime.Before(lr.EndTime) {
+		return
+	}
+
+	lr.Update(k)
+	lr.EndTime = k.EndTime.Time()
+}
+
+func (lr *LinReg) LoadK(allKLines []types.KLine) {
+	for _, k := range allKLines {
+		lr.PushK(k)
+	}
+}